From 8902901e61fcc514244bd5cd2ec0109bc56f6571 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 12 Aug 2026 22:55:17 +0000 Subject: [PATCH 01/42] feat(qwen35): complete concurrent speculative serving --- server/CMakeLists.txt | 37 +- server/deps/llama.cpp/ggml/include/ggml.h | 19 +- .../ggml/src/ggml-cuda/paged-attn.cu | 192 +++- server/deps/llama.cpp/ggml/src/ggml.c | 54 +- .../src/common/concurrency/paged_kv_pool.cpp | 110 ++- server/src/common/concurrency/paged_kv_pool.h | 48 + .../common/concurrency/paged_kv_residency.cpp | 929 ++++++++++++++++++ .../common/concurrency/paged_kv_residency.h | 260 +++++ .../concurrency/qwen_paged_kv_transfer.cpp | 405 ++++++++ .../concurrency/qwen_paged_kv_transfer.h | 87 ++ .../qwen_paged_kv_transfer_layout.cpp | 119 +++ server/src/common/concurrency/seq_engine.h | 65 +- server/src/common/ddtree.cpp | 10 + server/src/common/ddtree.h | 10 + server/src/common/feature_gate.cpp | 56 +- server/src/common/gpu_runtime_compat.h | 1 + server/src/common/step_graph.h | 10 +- server/src/internal.h | 36 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 655 +++++++++++- .../qwen35/concurrency/qwen35_seq_engine.h | 36 +- .../concurrency/qwen35_slot_manager.cpp | 251 ++++- .../qwen35/concurrency/qwen35_slot_manager.h | 53 +- server/src/qwen35/graph_builders.cpp | 156 +++ server/src/qwen35/graph_builders.h | 40 + server/src/qwen35/qwen35_backend.cpp | 99 +- server/src/qwen35/qwen35_backend.h | 6 +- server/src/qwen35/qwen35_target_graph.cpp | 187 +++- server/src/server/http_server.cpp | 27 +- server/src/server/http_server.h | 50 + server/src/server/scheduler.cpp | 124 ++- server/test/bench_paged_attention.cpp | 3 +- server/test/seq_engine_contract.h | 13 + server/test/test_ddtree_path.cpp | 46 + server/test/test_feature_gate.cpp | 63 ++ server/test/test_paged_attention.cpp | 192 +++- server/test/test_paged_kv_pool.cpp | 77 ++ server/test/test_paged_kv_residency.cpp | 559 +++++++++++ .../test_qwen_paged_kv_transfer_layout.cpp | 86 ++ server/test/test_recurrent_snapshot.cpp | 33 + server/test/test_seq_batch_plan.cpp | 37 + server/test/test_seq_engine_contract.cpp | 8 + server/test/test_seq_slot_manager.cpp | 58 +- server/test/test_server_unit.cpp | 71 ++ 43 files changed, 5168 insertions(+), 210 deletions(-) create mode 100644 server/src/common/concurrency/paged_kv_residency.cpp create mode 100644 server/src/common/concurrency/paged_kv_residency.h create mode 100644 server/src/common/concurrency/qwen_paged_kv_transfer.cpp create mode 100644 server/src/common/concurrency/qwen_paged_kv_transfer.h create mode 100644 server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp create mode 100644 server/test/test_ddtree_path.cpp create mode 100644 server/test/test_paged_kv_residency.cpp create mode 100644 server/test/test_qwen_paged_kv_transfer_layout.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 46ecfda81..ad6ef874f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -450,6 +450,9 @@ add_library(dflash_common STATIC src/common/dflash_draft_kv.cpp src/common/dflash_spec_decode.cpp src/common/concurrency/paged_kv_pool.cpp + src/common/concurrency/paged_kv_residency.cpp + src/common/concurrency/qwen_paged_kv_transfer_layout.cpp + src/common/concurrency/qwen_paged_kv_transfer.cpp src/qwen35/concurrency/qwen35_slot_manager.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp @@ -1391,12 +1394,33 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/src) list(APPEND _raw_unit_test_targets test_paged_kv_pool) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_paged_kv_residency.cpp") + # Pure host-side multi-sequence residency policy + mock DMA test. + add_executable(test_paged_kv_residency + test/test_paged_kv_residency.cpp + src/common/concurrency/paged_kv_pool.cpp + src/common/concurrency/paged_kv_residency.cpp) + target_include_directories(test_paged_kv_residency PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_paged_kv_residency) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_qwen_paged_kv_transfer_layout.cpp") + # Pure host-side validation of packed K/V block byte/stride math. + add_executable(test_qwen_paged_kv_transfer_layout + test/test_qwen_paged_kv_transfer_layout.cpp + src/common/concurrency/qwen_paged_kv_transfer_layout.cpp) + target_include_directories(test_qwen_paged_kv_transfer_layout PRIVATE + ${DFLASH27B_SRC_INCLUDE_DIRS} + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + list(APPEND _raw_unit_test_targets test_qwen_paged_kv_transfer_layout) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_seq_slot_manager.cpp") # Host-side slot bookkeeping test (concurrent serving): no GPU. add_executable(test_seq_slot_manager test/test_seq_slot_manager.cpp src/qwen35/concurrency/qwen35_slot_manager.cpp - src/common/concurrency/paged_kv_pool.cpp) + src/common/concurrency/paged_kv_pool.cpp + src/common/concurrency/paged_kv_residency.cpp) target_include_directories(test_seq_slot_manager PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) list(APPEND _raw_unit_test_targets test_seq_slot_manager) @@ -1409,6 +1433,15 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/test) list(APPEND _raw_unit_test_targets test_seq_engine_contract) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_ddtree_path.cpp") + # Pure host-side accepted-path/pending-token contract tests. + add_executable(test_ddtree_path + test/test_ddtree_path.cpp + src/common/ddtree.cpp) + target_include_directories(test_ddtree_path PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_ddtree_path) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_seq_batch_plan.cpp") # Pure-host tests for model-neutral token-budget/FIFO planning. add_executable(test_seq_batch_plan test/test_seq_batch_plan.cpp) @@ -1791,6 +1824,8 @@ if(DFLASH27B_TESTS) set(_unit_ctest_name recurrent_snapshot) elseif(_unit_target STREQUAL "test_paged_kv_pool") set(_unit_ctest_name paged_kv_pool) + elseif(_unit_target STREQUAL "test_paged_kv_residency") + set(_unit_ctest_name paged_kv_residency) endif() add_test(NAME "${_unit_ctest_name}" COMMAND ${_unit_target}) set_tests_properties("${_unit_ctest_name}" PROPERTIES SKIP_RETURN_CODE 77) diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index acac40c1f..beb6fc817 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2502,6 +2502,18 @@ extern "C" { // prefill chunks can attend the paged pool causally. A negative position // marks a padding row. NULL keeps the decode semantics (full cached // length per row). + // + // parent_ids/tree_sizes optionally enable packed tree verification. + // Queries are flattened sequence-major: tree sequence s occupies rows + // [s*tree_width, (s+1)*tree_width). parent_ids is contiguous I32 + // [tree_width, n_tree_seq] (root parent -1), and tree_sizes is contiguous + // I32 [n_tree_seq]. active_slot_ids is required and remains per query row; + // it selects the physical block-table column and scratch slab. Each live + // query attends its complete committed prefix from the block table plus + // its own candidate node and ancestors from physical K/V rows + // tree_scratch_base + slot*tree_scratch_stride + node. Siblings and rows + // at or beyond tree_sizes[s] are excluded. query_positions must be NULL + // in tree mode. Pass NULL/NULL/0/0/0 to retain standard paged attention. GGML_API struct ggml_tensor * ggml_paged_attn_ext( struct ggml_context * ctx, struct ggml_tensor * q, @@ -2513,7 +2525,12 @@ extern "C" { struct ggml_tensor * query_positions, float scale, int block_size, - int max_kv_seq_len); + int max_kv_seq_len, + struct ggml_tensor * parent_ids, + struct ggml_tensor * tree_sizes, + int tree_width, + int tree_scratch_base, + int tree_scratch_stride); // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse. // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse). diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu index c76d97b36..a881b93d3 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu @@ -29,6 +29,42 @@ static __host__ __device__ __forceinline__ int32_t paged_attn_partitions( return requested < available ? requested : available; } +// parent_ids is a sequence-major [tree_width, n_tree_seq] table. Walk only +// from the current query node toward the root; a candidate is visible iff it +// appears on that chain. The bounded walk also turns malformed cycles or +// out-of-range parents into invisible edges instead of an unsafe read. +static __device__ __forceinline__ bool paged_attn_tree_visible( + const char * __restrict__ parent_ids, + int64_t parent_nb0, + int64_t parent_nb1, + int32_t tree_seq, + int32_t query_node, + int32_t candidate, + int32_t tree_size) { + if (candidate < 0 || candidate >= tree_size || + query_node < 0 || query_node >= tree_size) { + return false; + } + + int32_t current = query_node; + for (int32_t depth = 0; depth < tree_size; ++depth) { + if (current == candidate) { + return true; + } + if (current < 0 || current >= tree_size) { + return false; + } + const int32_t parent = *(const int32_t *) ( + parent_ids + (int64_t) current * parent_nb0 + + (int64_t) tree_seq * parent_nb1); + if (parent == current) { + return false; + } + current = parent; + } + return false; +} + // All scores are computed in the log2 domain: log2(e) is folded into the same // Q prescale that already carries the 1/sqrt(D) attention scale, so every // softmax exponential uses the fast exp2f SFU path. @@ -180,6 +216,8 @@ static __global__ void paged_attn_decode( const char * __restrict__ kv_seq_lens, const char * __restrict__ active_slot_ids, const char * __restrict__ query_positions, + const char * __restrict__ parent_ids, + const char * __restrict__ tree_sizes, char * __restrict__ dst, half * __restrict__ partial_acc, float2 * __restrict__ partial_meta, @@ -189,6 +227,7 @@ static __global__ void paged_attn_decode( int64_t bt_nb0, int64_t bt_nb1, int64_t ksl_nb0, int64_t asi_nb0, int64_t qpos_nb0, + int64_t parent_nb0, int64_t parent_nb1, int64_t tree_size_nb0, int64_t dst_nb1, int64_t dst_nb2, int32_t n_table_seq, int32_t n_head, @@ -197,6 +236,9 @@ static __global__ void paged_attn_decode( int32_t max_blocks, int32_t block_size, int32_t min_partitions, + int32_t tree_width, + int32_t tree_scratch_base, + int32_t tree_scratch_stride, float scale) { constexpr int nthreads = WARP_SIZE; constexpr int values_per_load = 4; @@ -222,6 +264,15 @@ static __global__ void paged_attn_decode( const int n_seq = gridDim.y; const int n_partitions = gridDim.z; + const bool tree_mode = parent_ids != nullptr; + const int32_t tree_seq = tree_mode ? seq / tree_width : 0; + const int32_t query_node = + tree_mode ? seq - tree_seq * tree_width : -1; + const int32_t tree_size = tree_mode + ? *(const int32_t *) ( + tree_sizes + (int64_t) tree_seq * tree_size_nb0) + : 0; + const int32_t physical_seq_raw = active_slot_ids ? *(const int32_t *) (active_slot_ids + (int64_t) seq * asi_nb0) : seq; @@ -229,21 +280,22 @@ static __global__ void paged_attn_decode( ? *(const int32_t *) (query_positions + (int64_t) seq * qpos_nb0) : -1; // A row is live when its slot id selects a real block-table column and, - // for ragged batches, its causal position is non-negative. Dead rows are - // pinned to column 0 with kv_seq_len forced to 0, which routes every - // partition through the existing zero-output early path; the block table - // is then never read for them. + // for ragged batches, its causal position is non-negative. Tree padding + // rows are validated by tree_sizes. Dead rows are pinned to column 0 with + // an empty virtual context, so the block table and scratch are never read. const bool valid_query = physical_seq_raw >= 0 && physical_seq_raw < n_table_seq && - (!query_positions || query_pos >= 0); + (!query_positions || query_pos >= 0) && + (!tree_mode || + (tree_size >= 0 && tree_size <= tree_width && + query_node < tree_size)); const int32_t physical_seq = valid_query ? physical_seq_raw : 0; int32_t kv_seq_len_raw = valid_query ? *(const int32_t *) (kv_seq_lens + (int64_t) physical_seq * ksl_nb0) : 0; - // The inclusive clamp IS the causal mask: this row attends tokens - // [0, pos] only, and every downstream bound (partition count, token loop - // extents) already derives from kv_seq_len. + // The inclusive clamp IS the causal mask for non-tree ragged rows. Tree + // rows always read the whole committed prefix carried by kv_seq_lens. if (query_positions && query_pos < kv_seq_len_raw) { kv_seq_len_raw = query_pos + 1; } @@ -254,8 +306,14 @@ static __global__ void paged_attn_decode( : (kv_seq_len_raw < table_capacity ? kv_seq_len_raw : (int32_t) table_capacity); + // Treat the candidate slab as a virtual tail of tree_width tokens. The + // normal partition split then covers prefix and tree candidates in one + // stable softmax; invisible siblings/padding resolve to no physical row. + const int32_t virtual_tokens = valid_query + ? kv_seq_len + (tree_mode ? tree_width : 0) + : 0; const int32_t n_logical_blocks = - (kv_seq_len + block_size - 1) / block_size; + (virtual_tokens + block_size - 1) / block_size; const int32_t active_partitions = paged_attn_partitions(n_logical_blocks, min_partitions, n_partitions); @@ -289,7 +347,7 @@ static __global__ void paged_attn_decode( const int32_t token_begin = logical_block_begin * block_size; const int32_t token_end_blocks = logical_block_end * block_size; const int32_t token_end = - kv_seq_len < token_end_blocks ? kv_seq_len : token_end_blocks; + virtual_tokens < token_end_blocks ? virtual_tokens : token_end_blocks; constexpr bool quantize_q = type_K != GGML_TYPE_F16; constexpr int q_registers = (D / 2) / nthreads; @@ -363,7 +421,12 @@ static __global__ void paged_attn_decode( qk_sum[h] = 0.0f; } - const int32_t n_physical_blocks = pool_tokens / block_size; + // In tree mode the committed block table may address only the prefix + // pool before tree_scratch_base. Candidate rows are addressed directly + // below, keeping uncommitted nodes out of every sequence block table. + const int32_t prefix_pool_tokens = + tree_mode ? tree_scratch_base : pool_tokens; + const int32_t n_physical_blocks = prefix_pool_tokens / block_size; for (int32_t tile_begin = token_begin; tile_begin < token_end; @@ -380,7 +443,7 @@ static __global__ void paged_attn_decode( // read; their tokens contribute nothing, mirroring the CPU reference. int32_t phys_mine = -1; const int32_t my_token = tile_begin + lane; - if (my_token < token_end) { + if (my_token < token_end && my_token < kv_seq_len) { const int32_t logical_block = my_token / block_size; const int32_t physical_block = *(const int32_t *) (block_table + @@ -390,6 +453,19 @@ static __global__ void paged_attn_decode( phys_mine = physical_block * block_size + my_token % block_size; } + } else if (tree_mode && my_token < token_end) { + const int32_t candidate = my_token - kv_seq_len; + if (paged_attn_tree_visible( + parent_ids, parent_nb0, parent_nb1, + tree_seq, query_node, candidate, tree_size)) { + const int64_t physical = + (int64_t) tree_scratch_base + + (int64_t) physical_seq * tree_scratch_stride + + candidate; + if (physical >= 0 && physical < pool_tokens) { + phys_mine = (int32_t) physical; + } + } } float score_mine[n_batch_heads]; @@ -619,6 +695,8 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { const ggml_tensor * kv_seq_lens = dst->src[4]; const ggml_tensor * active_slot_ids = dst->src[5]; const ggml_tensor * query_positions = dst->src[6]; + const ggml_tensor * parent_ids = dst->src[7]; + const ggml_tensor * tree_sizes = dst->src[8]; if (!q || !k || !v || !block_table || !kv_seq_lens) { return false; @@ -628,6 +706,11 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { if (query_positions && !active_slot_ids) { return false; } + const bool tree_mode = parent_ids || tree_sizes; + if ((parent_ids == nullptr) != (tree_sizes == nullptr) || + (tree_mode && (!active_slot_ids || query_positions))) { + return false; + } if (dst->type != GGML_TYPE_F32 || q->type != GGML_TYPE_F32 || @@ -636,7 +719,9 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { block_table->type != GGML_TYPE_I32 || kv_seq_lens->type != GGML_TYPE_I32 || (active_slot_ids && active_slot_ids->type != GGML_TYPE_I32) || - (query_positions && query_positions->type != GGML_TYPE_I32)) { + (query_positions && query_positions->type != GGML_TYPE_I32) || + (parent_ids && parent_ids->type != GGML_TYPE_I32) || + (tree_sizes && tree_sizes->type != GGML_TYPE_I32)) { return false; } @@ -647,6 +732,8 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { kv_seq_lens->nb[0] != sizeof(int32_t) || (active_slot_ids && active_slot_ids->nb[0] != sizeof(int32_t)) || (query_positions && query_positions->nb[0] != sizeof(int32_t)) || + (parent_ids && parent_ids->nb[0] != sizeof(int32_t)) || + (tree_sizes && tree_sizes->nb[0] != sizeof(int32_t)) || dst->nb[0] != sizeof(float)) { return false; } @@ -701,10 +788,47 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { const int32_t block_size = ggml_get_op_params_i32(dst, 1); const int32_t max_kv_seq_len = ggml_get_op_params_i32(dst, 2); - return block_size > 0 && - max_kv_seq_len > 0 && - max_kv_seq_len <= k->ne[1] && - k->ne[1] % block_size == 0; + const int32_t tree_width = ggml_get_op_params_i32(dst, 3); + const int32_t tree_scratch_base = ggml_get_op_params_i32(dst, 4); + const int32_t tree_scratch_stride = ggml_get_op_params_i32(dst, 5); + if (block_size <= 0 || + max_kv_seq_len <= 0 || + (int64_t) max_kv_seq_len + tree_width > INT32_MAX || + k->ne[1] % block_size != 0) { + return false; + } + + if (!tree_mode) { + return tree_width == 0 && + tree_scratch_base == 0 && + tree_scratch_stride == 0; + } + + if (tree_width <= 0 || + tree_scratch_base <= 0 || + tree_scratch_base % block_size != 0 || + tree_scratch_stride < tree_width || + !ggml_is_contiguous(parent_ids) || + !ggml_is_contiguous(tree_sizes) || + parent_ids->ne[0] != tree_width || + parent_ids->ne[1] <= 0 || + parent_ids->ne[1] != tree_sizes->ne[0] || + parent_ids->ne[2] != 1 || + parent_ids->ne[3] != 1 || + tree_sizes->ne[1] != 1 || + tree_sizes->ne[2] != 1 || + tree_sizes->ne[3] != 1 || + parent_ids->ne[1] > INT64_MAX / tree_width || + q->ne[1] != parent_ids->ne[1] * tree_width || + (int64_t) max_kv_seq_len + tree_width > INT32_MAX) { + return false; + } + + const int64_t scratch_end = + (int64_t) tree_scratch_base + + (block_table->ne[1] - 1) * (int64_t) tree_scratch_stride + + tree_width; + return scratch_end <= k->ne[1]; } // Cached max resident blocks/SM for this instantiation at the given block @@ -764,6 +888,12 @@ static bool try_launch_paged_attn( const ggml_tensor * kv_seq_lens = dst->src[4]; const ggml_tensor * active_slot_ids = dst->src[5]; const ggml_tensor * query_positions = dst->src[6]; + const ggml_tensor * parent_ids = dst->src[7]; + const ggml_tensor * tree_sizes = dst->src[8]; + + const int32_t tree_width = ggml_get_op_params_i32(dst, 3); + const int32_t tree_scratch_base = ggml_get_op_params_i32(dst, 4); + const int32_t tree_scratch_stride = ggml_get_op_params_i32(dst, 5); const int32_t n_head = (int32_t) q->ne[2]; const int32_t n_head_kv = (int32_t) k->ne[2]; @@ -837,15 +967,21 @@ static bool try_launch_paged_attn( if (min_partitions > partition_limit) { min_partitions = partition_limit; } - if (min_partitions > block_table->ne[0]) { - min_partitions = (int32_t) block_table->ne[0]; + const int32_t tree_blocks = + (tree_width + block_size - 1) / block_size; + const int64_t partitionable_blocks = + block_table->ne[0] + (parent_ids ? tree_blocks : 0); + if (min_partitions > partitionable_blocks) { + min_partitions = (int32_t) partitionable_blocks; } - // Size the launch from the live maximum sequence length carried in the - // graph op, not the block-table capacity. Ragged sequences still clamp - // their own active partition count from kv_seq_lens on device. + // Size the launch from the live maximum committed prefix plus the virtual + // tree tail. Ragged/tree rows still clamp their own active partition count + // from device metadata. + const int32_t live_tokens = + max_kv_seq_len + (parent_ids ? tree_width : 0); const int32_t live_blocks = - (max_kv_seq_len + block_size - 1) / block_size; + (live_tokens + block_size - 1) / block_size; int32_t n_partitions = paged_attn_partitions( live_blocks, min_partitions, PAGED_ATTN_MAX_PARTITIONS); @@ -861,7 +997,7 @@ static bool try_launch_paged_attn( }(); if (forced_partitions >= 1 && forced_partitions <= PAGED_ATTN_MAX_PARTITIONS && - forced_partitions <= block_table->ne[0]) { + forced_partitions <= partitionable_blocks) { min_partitions = forced_partitions; n_partitions = forced_partitions; } @@ -914,6 +1050,8 @@ static bool try_launch_paged_attn( (const char *) kv_seq_lens->data, active_slot_ids ? (const char *) active_slot_ids->data : nullptr, query_positions ? (const char *) query_positions->data : nullptr, + parent_ids ? (const char *) parent_ids->data : nullptr, + tree_sizes ? (const char *) tree_sizes->data : nullptr, (char *) dst->data, partial_acc, partial_meta, @@ -924,6 +1062,9 @@ static bool try_launch_paged_attn( kv_seq_lens->nb[0], active_slot_ids ? active_slot_ids->nb[0] : 0, query_positions ? query_positions->nb[0] : 0, + parent_ids ? parent_ids->nb[0] : 0, + parent_ids ? parent_ids->nb[1] : 0, + tree_sizes ? tree_sizes->nb[0] : 0, dst->nb[1], dst->nb[2], (int32_t) block_table->ne[1], n_head, @@ -932,6 +1073,9 @@ static bool try_launch_paged_attn( (int32_t) block_table->ne[0], block_size, min_partitions, + tree_width, + tree_scratch_base, + tree_scratch_stride, scale); if (n_partitions > 1) { diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 86e4520d6..77a590cca 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5708,7 +5708,12 @@ struct ggml_tensor * ggml_paged_attn_ext( struct ggml_tensor * query_positions, float scale, int block_size, - int max_kv_seq_len) { + int max_kv_seq_len, + struct ggml_tensor * parent_ids, + struct ggml_tensor * tree_sizes, + int tree_width, + int tree_scratch_base, + int tree_scratch_stride) { GGML_ASSERT(q->type == GGML_TYPE_F32); GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_Q4_0 || k->type == GGML_TYPE_Q8_0); GGML_ASSERT(v->type == GGML_TYPE_F16 || v->type == GGML_TYPE_Q4_0 || v->type == GGML_TYPE_Q8_0); @@ -5720,6 +5725,13 @@ struct ggml_tensor * ggml_paged_attn_ext( GGML_ASSERT(query_positions == NULL || active_slot_ids != NULL); GGML_ASSERT(query_positions == NULL || query_positions->type == GGML_TYPE_I32); + const bool tree_mode = parent_ids != NULL || tree_sizes != NULL; + GGML_ASSERT((parent_ids == NULL) == (tree_sizes == NULL)); + GGML_ASSERT(!tree_mode || active_slot_ids != NULL); + GGML_ASSERT(!tree_mode || query_positions == NULL); + GGML_ASSERT(!tree_mode || parent_ids->type == GGML_TYPE_I32); + GGML_ASSERT(!tree_mode || tree_sizes->type == GGML_TYPE_I32); + GGML_ASSERT(q->ne[0] == k->ne[0] && q->ne[0] == v->ne[0]); GGML_ASSERT(k->ne[1] == v->ne[1]); GGML_ASSERT(k->ne[2] > 0); @@ -5749,13 +5761,49 @@ struct ggml_tensor * ggml_paged_attn_ext( GGML_ASSERT(block_size > 0); GGML_ASSERT(k->ne[1] % block_size == 0); GGML_ASSERT(max_kv_seq_len > 0); - GGML_ASSERT(max_kv_seq_len <= k->ne[1]); + // This is a padded logical launch bound, not a physical-cache extent. + // Each row clamps its actual sequence length to the block-table capacity + // and validates every resolved physical block before dereferencing K/V. + GGML_ASSERT((int64_t) max_kv_seq_len + tree_width <= INT32_MAX); + + if (tree_mode) { + GGML_ASSERT(tree_width > 0); + GGML_ASSERT(tree_scratch_base > 0); + GGML_ASSERT(tree_scratch_base % block_size == 0); + GGML_ASSERT(tree_scratch_stride >= tree_width); + GGML_ASSERT(ggml_is_contiguous(parent_ids)); + GGML_ASSERT(ggml_is_contiguous(tree_sizes)); + GGML_ASSERT(parent_ids->ne[0] == tree_width); + GGML_ASSERT(parent_ids->ne[1] == tree_sizes->ne[0]); + GGML_ASSERT(parent_ids->ne[2] == 1 && parent_ids->ne[3] == 1); + GGML_ASSERT(tree_sizes->ne[1] == 1 && tree_sizes->ne[2] == 1 && tree_sizes->ne[3] == 1); + GGML_ASSERT(parent_ids->ne[1] > 0); + GGML_ASSERT(parent_ids->ne[1] <= INT64_MAX / tree_width); + GGML_ASSERT(q->ne[1] == parent_ids->ne[1] * tree_width); + + // Every physical sequence slot owns one non-overlapping scratch slab. + // Bound the largest address with int64 arithmetic before the GPU sees + // the int32 op parameters. + const int64_t scratch_end = + (int64_t) tree_scratch_base + + (block_table->ne[1] - 1) * (int64_t) tree_scratch_stride + + tree_width; + GGML_ASSERT(scratch_end <= k->ne[1]); + GGML_ASSERT((int64_t) max_kv_seq_len + tree_width <= INT32_MAX); + } else { + GGML_ASSERT(tree_width == 0); + GGML_ASSERT(tree_scratch_base == 0); + GGML_ASSERT(tree_scratch_stride == 0); + } struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, q->ne); ggml_set_op_params_f32(result, 0, scale); ggml_set_op_params_i32(result, 1, block_size); ggml_set_op_params_i32(result, 2, max_kv_seq_len); + ggml_set_op_params_i32(result, 3, tree_width); + ggml_set_op_params_i32(result, 4, tree_scratch_base); + ggml_set_op_params_i32(result, 5, tree_scratch_stride); result->op = GGML_OP_PAGED_ATTN; result->src[0] = q; @@ -5765,6 +5813,8 @@ struct ggml_tensor * ggml_paged_attn_ext( result->src[4] = kv_seq_lens; result->src[5] = active_slot_ids; result->src[6] = query_positions; + result->src[7] = parent_ids; + result->src[8] = tree_sizes; return result; } diff --git a/server/src/common/concurrency/paged_kv_pool.cpp b/server/src/common/concurrency/paged_kv_pool.cpp index 6bdf0f3fe..741ace600 100644 --- a/server/src/common/concurrency/paged_kv_pool.cpp +++ b/server/src/common/concurrency/paged_kv_pool.cpp @@ -20,6 +20,12 @@ const char * paged_kv_status_string(PagedKvStatus status) { return "physical blocks exhausted"; case PagedKvStatus::StaleHandle: return "stale sequence handle"; + case PagedKvStatus::LogicalBlockOutOfRange: + return "logical block out of range"; + case PagedKvStatus::BlockNotResident: + return "logical block is not resident"; + case PagedKvStatus::BlockAlreadyResident: + return "logical block is already resident"; } return "unknown paged KV status"; } @@ -138,8 +144,35 @@ PagedKvAppendResult PagedKvPool::append(PagedKvSequenceHandle handle, const uint32_t old_kv_seq_len = sequence.kv_seq_len; const uint32_t new_kv_seq_len = old_kv_seq_len + token_count; - result.status = - extend_block_table(sequence, blocks_for_tokens(new_kv_seq_len)); + const uint32_t required_blocks = blocks_for_tokens(new_kv_seq_len); + const uint32_t additional_blocks = + required_blocks - static_cast(sequence.block_table.size()); + + // Appending into a partially-filled cold head needs one physical remap in + // addition to any newly opened logical blocks. Preflight the aggregate so + // append remains all-or-nothing on BlocksExhausted. + const bool remap_head = + old_kv_seq_len % block_size_ != 0 && + sequence.block_table[old_kv_seq_len / block_size_] == + PAGED_KV_COLD_BLOCK; + const uint32_t allocations = additional_blocks + (remap_head ? 1u : 0u); + const uint64_t available = + static_cast(sequence.reserved_blocks.size()) + + free_blocks_.size(); + if (allocations > available) { + result.status = PagedKvStatus::BlocksExhausted; + return result; + } + + sequence.block_table.reserve(required_blocks); + if (remap_head) { + const uint32_t logical_block = old_kv_seq_len / block_size_; + const uint32_t physical_block = take_append_block(sequence); + sequence.block_table[logical_block] = physical_block; + result.remapped_cold_blocks.push_back( + {logical_block, physical_block}); + } + result.status = extend_block_table(sequence, required_blocks); if (result.status != PagedKvStatus::Ok) return result; const auto make_slot = [&](uint32_t logical_position) { @@ -176,7 +209,9 @@ PagedKvStatus PagedKvPool::release(PagedKvSequenceHandle handle) { SequenceState & sequence = sequences_[handle.slot]; request_to_slot_.erase(sequence.request_id); for (uint32_t block : sequence.block_table) { - give_back(free_blocks_, block); + if (block != PAGED_KV_COLD_BLOCK) { + give_back(free_blocks_, block); + } } for (uint32_t block : sequence.reserved_blocks) { give_back(free_blocks_, block); @@ -226,11 +261,70 @@ PagedKvStatus PagedKvPool::owned_block_count( if (status != PagedKvStatus::Ok) return status; const SequenceState & sequence = sequences_[handle.slot]; + // Cold logical blocks remain appended sequence capacity even though they + // no longer own a physical page. Preserve the pre-residency API contract: + // appended logical blocks plus blocks reserved for future append. out_count = static_cast( sequence.block_table.size() + sequence.reserved_blocks.size()); return PagedKvStatus::Ok; } +PagedKvStatus PagedKvPool::page_out_block( + PagedKvSequenceHandle handle, uint32_t logical_block, + uint32_t & out_physical_block) { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + SequenceState & sequence = sequences_[handle.slot]; + if (logical_block >= sequence.block_table.size()) { + return PagedKvStatus::LogicalBlockOutOfRange; + } + const uint32_t physical_block = sequence.block_table[logical_block]; + if (physical_block == PAGED_KV_COLD_BLOCK) { + return PagedKvStatus::BlockNotResident; + } + + sequence.block_table[logical_block] = PAGED_KV_COLD_BLOCK; + give_back(free_blocks_, physical_block); + out_physical_block = physical_block; + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::page_in_block( + PagedKvSequenceHandle handle, uint32_t logical_block, + uint32_t & out_physical_block) { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + SequenceState & sequence = sequences_[handle.slot]; + if (logical_block >= sequence.block_table.size()) { + return PagedKvStatus::LogicalBlockOutOfRange; + } + if (sequence.block_table[logical_block] != PAGED_KV_COLD_BLOCK) { + return PagedKvStatus::BlockAlreadyResident; + } + if (free_blocks_.empty()) { + return PagedKvStatus::BlocksExhausted; + } + + const uint32_t physical_block = take_lowest(free_blocks_); + sequence.block_table[logical_block] = physical_block; + out_physical_block = physical_block; + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::resident_block_count( + PagedKvSequenceHandle handle, uint32_t & out_count) const { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + const SequenceState & sequence = sequences_[handle.slot]; + out_count = static_cast(std::count_if( + sequence.block_table.begin(), sequence.block_table.end(), + [](uint32_t block) { return block != PAGED_KV_COLD_BLOCK; })); + return PagedKvStatus::Ok; +} + uint32_t PagedKvPool::blocks_for_tokens(uint32_t token_count) const { if (token_count == 0) return 0; return 1 + (token_count - 1) / block_size_; @@ -262,13 +356,17 @@ PagedKvStatus PagedKvPool::extend_block_table(SequenceState & sequence, sequence.block_table.reserve(required_blocks); for (uint32_t i = 0; i < additional_blocks; ++i) { - std::vector & source = sequence.reserved_blocks.empty() - ? free_blocks_ : sequence.reserved_blocks; - sequence.block_table.push_back(take_lowest(source)); + sequence.block_table.push_back(take_append_block(sequence)); } return PagedKvStatus::Ok; } +uint32_t PagedKvPool::take_append_block(SequenceState & sequence) { + std::vector & source = sequence.reserved_blocks.empty() + ? free_blocks_ : sequence.reserved_blocks; + return take_lowest(source); +} + void PagedKvPool::take_reserved_blocks(SequenceState & sequence, uint32_t additional_blocks) { sequence.reserved_blocks.reserve( diff --git a/server/src/common/concurrency/paged_kv_pool.h b/server/src/common/concurrency/paged_kv_pool.h index 9d745739f..3cf90964a 100644 --- a/server/src/common/concurrency/paged_kv_pool.h +++ b/server/src/common/concurrency/paged_kv_pool.h @@ -33,6 +33,9 @@ enum class PagedKvStatus : uint8_t { SequenceSlotsExhausted, // all max_sequences slots are in use BlocksExhausted, // not enough free physical blocks for the growth StaleHandle, // handle refers to a released or reused slot + LogicalBlockOutOfRange, // logical block is not materialized by the sequence + BlockNotResident, // page_out_block() targeted an already-cold block + BlockAlreadyResident, // page_in_block() targeted a resident block }; // Human-readable status name for logs and error messages. @@ -57,6 +60,23 @@ struct PagedKvWriteSlot { uint64_t physical_token_index = 0; }; +// Sentinel in PagedKvSequenceSnapshot::block_table for a logical block whose +// full-attention K/V bytes are host-backed by PagedKvResidencyManager. It is +// never returned as a PagedKvWriteSlot::physical_block. +inline constexpr uint32_t PAGED_KV_COLD_BLOCK = + std::numeric_limits::max(); + +// A formerly-cold append-head block that append() had to remap. The caller +// must restore the block's host-backed bytes before consuming or overwriting +// any row in the returned physical block. PagedKvResidencyManager::append() +// and prepare_append() do that before returning to the engine; this record +// primarily makes a direct pool append fail-obvious instead of losing the +// remap information. +struct PagedKvBlockRemap { + uint32_t logical_block = 0; + uint32_t physical_block = 0; +}; + // Outcome of append(). On success, `token_count` is the number of appended // tokens. By default, `write_slots` holds one entry per token in logical // order. With `only_first_last_slots`, it is empty and `first` and `last` @@ -68,6 +88,7 @@ struct PagedKvAppendResult { std::vector write_slots; PagedKvWriteSlot first; PagedKvWriteSlot last; + std::vector remapped_cold_blocks; explicit operator bool() const { return status == PagedKvStatus::Ok; } }; @@ -75,6 +96,8 @@ struct PagedKvAppendResult { // Copy of one sequence's bookkeeping state, as returned by sequence(). struct PagedKvSequenceSnapshot { uint32_t kv_seq_len = 0; + // Entries are physical block indices or PAGED_KV_COLD_BLOCK. Logical + // length and block-table length do not shrink when a block is paged out. std::vector block_table; // Physical blocks held for future append() calls by this sequence. They // are not visible in block_table until append consumes them. @@ -155,6 +178,26 @@ class PagedKvPool { PagedKvStatus owned_block_count(PagedKvSequenceHandle handle, uint32_t & out_count) const; + // Relinquish one materialized physical block while preserving its logical + // block-table position as PAGED_KV_COLD_BLOCK. The caller must first copy + // the complete block to host backing. On success, out_physical_block is + // the returned pool block and may be reused immediately. + PagedKvStatus page_out_block(PagedKvSequenceHandle handle, + uint32_t logical_block, + uint32_t & out_physical_block); + + // Allocate a physical block for one cold logical entry. The caller must + // restore its complete host-backed bytes before attention or append reads + // it. On failure, the cold mapping and output argument are unchanged. + PagedKvStatus page_in_block(PagedKvSequenceHandle handle, + uint32_t logical_block, + uint32_t & out_physical_block); + + // Number of materialized logical blocks that currently own physical + // storage. Reserved append capacity is deliberately excluded. + PagedKvStatus resident_block_count(PagedKvSequenceHandle handle, + uint32_t & out_count) const; + private: // Bookkeeping for one sequence slot. `generation` survives release so // the next acquire on this slot invalidates old handles. @@ -182,6 +225,11 @@ class PagedKvPool { PagedKvStatus extend_block_table(SequenceState & sequence, uint32_t required_blocks); + // Allocate a physical block from a sequence's reservation first, then the + // global free list. Used only by append(), where consuming promised future + // capacity is correct. + uint32_t take_append_block(SequenceState & sequence); + // Move exactly `additional_blocks` globally free blocks into a sequence's // private reservation. Caller must preflight availability. void take_reserved_blocks(SequenceState & sequence, diff --git a/server/src/common/concurrency/paged_kv_residency.cpp b/server/src/common/concurrency/paged_kv_residency.cpp new file mode 100644 index 000000000..43c8443b0 --- /dev/null +++ b/server/src/common/concurrency/paged_kv_residency.cpp @@ -0,0 +1,929 @@ +#include "paged_kv_residency.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +PagedKvResidencyStatus from_pool_status(PagedKvStatus status) { + switch (status) { + case PagedKvStatus::Ok: + return PagedKvResidencyStatus::Ok; + case PagedKvStatus::StaleHandle: + return PagedKvResidencyStatus::StaleHandle; + case PagedKvStatus::BlocksExhausted: + return PagedKvResidencyStatus::PoolExhausted; + case PagedKvStatus::InvalidArgument: + case PagedKvStatus::LogicalBlockOutOfRange: + case PagedKvStatus::BlockNotResident: + case PagedKvStatus::BlockAlreadyResident: + return PagedKvResidencyStatus::InvalidArgument; + case PagedKvStatus::DuplicateRequest: + case PagedKvStatus::SequenceSlotsExhausted: + return PagedKvResidencyStatus::InconsistentPoolState; + } + return PagedKvResidencyStatus::InconsistentPoolState; +} + +} // namespace + +const char * paged_kv_residency_status_string( + PagedKvResidencyStatus status) { + switch (status) { + case PagedKvResidencyStatus::Ok: + return "ok"; + case PagedKvResidencyStatus::InvalidArgument: + return "invalid argument"; + case PagedKvResidencyStatus::SequenceNotRegistered: + return "sequence not registered"; + case PagedKvResidencyStatus::StaleHandle: + return "stale sequence handle"; + case PagedKvResidencyStatus::PoolExhausted: + return "physical pool exhausted"; + case PagedKvResidencyStatus::NoEvictableBlock: + return "no evictable block"; + case PagedKvResidencyStatus::HostAllocationFailed: + return "pinned host allocation failed"; + case PagedKvResidencyStatus::TransferFailed: + return "K/V transfer failed"; + case PagedKvResidencyStatus::HostCopyMissing: + return "cold block has no valid host copy"; + case PagedKvResidencyStatus::InconsistentPoolState: + return "residency state disagrees with paged pool"; + } + return "unknown paged K/V residency status"; +} + +PagedKvResidencyManager::PagedKvResidencyManager( + PagedKvPool & pool, PagedKvResidencyConfig config, + PagedKvResidencyTransferOps transfers) + : pool_(pool), config_(config), transfers_(std::move(transfers)) { + if (config_.block_bytes == 0 || !transfers_.allocate_pinned || + !transfers_.free_pinned || !transfers_.copy_out_async || + !transfers_.copy_in_async || !transfers_.synchronize) { + throw std::invalid_argument("invalid paged K/V residency callbacks"); + } + if (config_.resident_budget_blocks == 0) { + config_.resident_budget_blocks = pool_.physical_block_count(); + } + if (config_.resident_budget_blocks > pool_.physical_block_count()) { + throw std::invalid_argument("resident budget exceeds paged K/V pool"); + } + sequences_.resize(pool_.max_sequences()); +} + +PagedKvResidencyManager::~PagedKvResidencyManager() { + try { + reset(); + } catch (...) { + // Destructors must not propagate callback failures. Production pinned + // allocators/free functions are non-throwing; this catch protects test + // and plugin callbacks from terminating teardown. + } +} + +PagedKvResidencyStatus PagedKvResidencyManager::validate_registered( + PagedKvSequenceHandle handle) const { + if (transfer_barrier_failed_) { + return PagedKvResidencyStatus::TransferFailed; + } + if (handle.slot >= sequences_.size()) { + return PagedKvResidencyStatus::StaleHandle; + } + const SequenceState & state = sequences_[handle.slot]; + if (!state.active) { + return PagedKvResidencyStatus::SequenceNotRegistered; + } + if (state.generation != handle.generation) { + return PagedKvResidencyStatus::StaleHandle; + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::register_sequence( + PagedKvSequenceHandle handle) { + if (transfer_barrier_failed_) { + return PagedKvResidencyStatus::TransferFailed; + } + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) { + return from_pool_status(pool_status); + } + if (handle.slot >= sequences_.size()) { + return PagedKvResidencyStatus::StaleHandle; + } + if (std::find(snapshot.block_table.begin(), snapshot.block_table.end(), + PAGED_KV_COLD_BLOCK) != snapshot.block_table.end()) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + + SequenceState & state = sequences_[handle.slot]; + if (state.active && state.generation == handle.generation) { + return PagedKvResidencyStatus::Ok; + } + if (state.active) { + const auto synced = synchronize_before_read(); + if (synced != PagedKvResidencyStatus::Ok) return synced; + free_sequence_buffers(state); + } + state.active = true; + state.generation = handle.generation; + state.blocks.resize(snapshot.block_table.size()); + for (BlockState & block : state.blocks) block.last_use = ++clock_; + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::forget_sequence( + PagedKvSequenceHandle handle) { + // Teardown is also the recovery path after a failed copy-stream barrier. + // validate_registered() deliberately rejects ordinary operations while a + // transfer is quarantined, so validate the generation directly here and + // allow synchronize_before_read() to retry the barrier. + if (handle.slot >= sequences_.size()) { + return PagedKvResidencyStatus::StaleHandle; + } + const SequenceState & state = sequences_[handle.slot]; + if (!state.active) { + return PagedKvResidencyStatus::SequenceNotRegistered; + } + if (state.generation != handle.generation) { + return PagedKvResidencyStatus::StaleHandle; + } + const auto synced = synchronize_before_read(); + if (synced != PagedKvResidencyStatus::Ok) return synced; + free_sequence_buffers(sequences_[handle.slot]); + return PagedKvResidencyStatus::Ok; +} + +void PagedKvResidencyManager::reset() { + // A failed barrier should not let teardown free memory still referenced by + // an async transfer. Leak those buffers rather than creating a use-after- + // free; the process/backend is already unhealthy in this case. + if (synchronize_before_read() != PagedKvResidencyStatus::Ok) return; + for (SequenceState & state : sequences_) free_sequence_buffers(state); + stats_ = {}; + clock_ = 0; +} + +void PagedKvResidencyManager::free_sequence_buffers( + SequenceState & state) noexcept { + for (BlockState & block : state.blocks) { + if (!block.host) continue; + try { + transfers_.free_pinned(block.host); + } catch (...) { + // The callback contract is non-throwing. Continue freeing other + // pages if a third-party implementation violates it. + } + block.host = nullptr; + if (stats_.host_bytes >= config_.block_bytes) { + stats_.host_bytes -= config_.block_bytes; + } + } + state = {}; +} + +PagedKvResidencyStatus PagedKvResidencyManager::refresh_sequence( + PagedKvSequenceHandle handle, bool appended_rows, + uint32_t append_first, uint32_t append_count) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + + SequenceState & state = sequences_[handle.slot]; + if (snapshot.block_table.size() < state.blocks.size()) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + state.blocks.resize(snapshot.block_table.size()); + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (snapshot.block_table[logical] == PAGED_KV_COLD_BLOCK && + !state.blocks[logical].host_valid) { + return PagedKvResidencyStatus::HostCopyMissing; + } + } + + if (appended_rows && append_count > 0) { + const uint64_t last = static_cast(append_first) + + append_count - 1; + const uint32_t first_block = append_first / pool_.block_size(); + const uint32_t last_block = + static_cast(last / pool_.block_size()); + if (last_block >= state.blocks.size()) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + for (uint32_t logical = first_block; logical <= last_block; ++logical) { + // The target graph has not written the returned rows yet. Retain a + // restored partial page's old host image until commit, but make the + // physical page ineligible for eviction in every policy path. + state.blocks[logical].write_pending = true; + state.blocks[logical].last_use = ++clock_; + } + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::prepare_append( + PagedKvSequenceHandle handle, uint32_t token_count) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (token_count > std::numeric_limits::max() - + snapshot.kv_seq_len) { + return PagedKvResidencyStatus::InvalidArgument; + } + if (token_count == 0) return finish_transfers(PagedKvResidencyStatus::Ok); + + bool restore_partial_head = false; + uint32_t partial_head = 0; + if (snapshot.kv_seq_len % pool_.block_size() != 0) { + partial_head = snapshot.kv_seq_len / pool_.block_size(); + restore_partial_head = + snapshot.block_table[partial_head] == PAGED_KV_COLD_BLOCK; + } + + const uint32_t new_length = snapshot.kv_seq_len + token_count; + const uint32_t required_blocks = + 1 + (new_length - 1) / pool_.block_size(); + const uint32_t additional_blocks = required_blocks - + static_cast(snapshot.block_table.size()); + const uint32_t globally_needed = additional_blocks > + snapshot.reserved_block_count + ? additional_blocks - snapshot.reserved_block_count : 0; + // Reserve the partial-head restoration and every new physical page as one + // transaction. Restoring first and making room later can otherwise select + // the just-restored append head as the next eviction victim. + const auto room = make_room( + handle, globally_needed + (restore_partial_head ? 1u : 0u), + additional_blocks + (restore_partial_head ? 1u : 0u)); + if (room != PagedKvResidencyStatus::Ok) return finish_transfers(room); + if (restore_partial_head) { + const auto restored = restore_block_async( + handle, partial_head); + if (restored != PagedKvResidencyStatus::Ok) { + return finish_transfers(restored); + } + } + return finish_transfers(PagedKvResidencyStatus::Ok); +} + +PagedKvResidentAppendResult PagedKvResidencyManager::append( + PagedKvSequenceHandle handle, uint32_t token_count, + bool only_first_last_slots) { + PagedKvResidentAppendResult result; + result.status = prepare_append(handle, token_count); + if (result.status != PagedKvResidencyStatus::Ok) return result; + + result.pool_result = + pool_.append(handle, token_count, only_first_last_slots); + if (result.pool_result.status != PagedKvStatus::Ok) { + result.status = from_pool_status(result.pool_result.status); + return result; + } + result.status = observe_append(handle, result.pool_result); + return result; +} + +PagedKvResidencyStatus PagedKvResidencyManager::observe_append( + PagedKvSequenceHandle handle, const PagedKvAppendResult & result) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + if (result.status != PagedKvStatus::Ok) return from_pool_status(result.status); + + SequenceState & state = sequences_[handle.slot]; + for (const PagedKvBlockRemap & remap : result.remapped_cold_blocks) { + if (remap.logical_block >= state.blocks.size() || + !state.blocks[remap.logical_block].host_valid) { + return finish_transfers(PagedKvResidencyStatus::HostCopyMissing); + } + bool queued = false; + try { + queued = transfers_.copy_in_async( + handle, remap.logical_block, remap.physical_block, + state.blocks[remap.logical_block].host, config_.block_bytes); + } catch (...) { + queued = false; + } + if (!queued) return finish_transfers(PagedKvResidencyStatus::TransferFailed); + transfers_pending_ = true; + stats_.page_ins++; + stats_.moved_bytes += config_.block_bytes; + } + const auto synced = finish_transfers(PagedKvResidencyStatus::Ok); + if (synced != PagedKvResidencyStatus::Ok) return synced; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (result.token_count > snapshot.kv_seq_len) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + return refresh_sequence( + handle, /*appended_rows=*/true, + snapshot.kv_seq_len - result.token_count, result.token_count); +} + +PagedKvResidencyStatus PagedKvResidencyManager::commit_pending_writes( + PagedKvSequenceHandle handle) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + + // The caller supplies the target-compute -> pager dependency by + // synchronizing the target backend before this call. Any old host image is + // stale only now, after the device rows have actually been overwritten. + SequenceState & state = sequences_[handle.slot]; + for (BlockState & block : state.blocks) { + if (!block.write_pending) continue; + block.write_pending = false; + block.host_valid = false; + block.last_use = ++clock_; + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::finish_transfers( + PagedKvResidencyStatus status) { + const auto synced = synchronize_before_read(); + return synced == PagedKvResidencyStatus::Ok ? status : synced; +} + +PagedKvResidencyStatus PagedKvResidencyManager::synchronize_before_read() { + if (!transfers_pending_) { + return transfer_barrier_failed_ + ? PagedKvResidencyStatus::TransferFailed + : PagedKvResidencyStatus::Ok; + } + bool ok = false; + try { + ok = transfers_.synchronize(); + } catch (...) { + ok = false; + } + if (!ok) { + // A failed barrier does not prove that the stream stopped using its + // source and destination blocks. Keep every transfer pending and every + // H2D destination mapped (therefore quarantined from the free list), + // and reject all manager operations until an explicit retry confirms + // that the stream has drained. + transfer_barrier_failed_ = true; + return PagedKvResidencyStatus::TransferFailed; + } + + transfers_pending_ = false; + transfer_barrier_failed_ = false; + + PagedKvResidencyStatus result = PagedKvResidencyStatus::Ok; + for (const PendingTransfer & transfer : pending_page_outs_) { + if (validate_registered(transfer.handle) != + PagedKvResidencyStatus::Ok) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + BlockState & block = + sequences_[transfer.handle.slot].blocks[transfer.logical_block]; + block.page_out_pending = false; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(transfer.handle, snapshot) != PagedKvStatus::Ok || + transfer.logical_block >= snapshot.block_table.size() || + snapshot.block_table[transfer.logical_block] != + transfer.physical_block) { + block.host_valid = false; + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + uint32_t released = PAGED_KV_COLD_BLOCK; + if (pool_.page_out_block( + transfer.handle, transfer.logical_block, released) != + PagedKvStatus::Ok || released != transfer.physical_block) { + block.host_valid = false; + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + block.host_valid = true; + stats_.page_outs++; + stats_.moved_bytes += config_.block_bytes; + } + for (const PendingTransfer & transfer : pending_page_ins_) { + if (validate_registered(transfer.handle) != + PagedKvResidencyStatus::Ok) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + BlockState & block = + sequences_[transfer.handle.slot].blocks[transfer.logical_block]; + block.page_in_pending = false; + stats_.page_ins++; + stats_.moved_bytes += config_.block_bytes; + } + pending_page_outs_.clear(); + pending_page_ins_.clear(); + if (result != PagedKvResidencyStatus::Ok) return result; + return PagedKvResidencyStatus::Ok; +} + +bool PagedKvResidencyManager::is_protected( + PagedKvSequenceHandle handle, uint32_t logical_block) const { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(handle, snapshot) != PagedKvStatus::Ok || + logical_block >= snapshot.block_table.size()) { + return true; + } + if (logical_block < config_.sink_blocks) return true; + const uint32_t tail_begin = snapshot.block_table.size() > config_.tail_blocks + ? static_cast(snapshot.block_table.size()) - + config_.tail_blocks + : 0; + return logical_block >= tail_begin; +} + +uint32_t PagedKvResidencyManager::sequence_resident_count( + PagedKvSequenceHandle handle) const { + uint32_t count = 0; + return pool_.resident_block_count(handle, count) == PagedKvStatus::Ok + ? count : 0; +} + +uint32_t PagedKvResidencyManager::sequence_pending_page_out_count( + PagedKvSequenceHandle handle) const { + if (handle.slot >= sequences_.size()) return 0; + const SequenceState & state = sequences_[handle.slot]; + if (!state.active || state.generation != handle.generation) return 0; + return static_cast(std::count_if( + state.blocks.begin(), state.blocks.end(), + [](const BlockState & block) { return block.page_out_pending; })); +} + +uint32_t PagedKvResidencyManager::sequence_protected_count( + PagedKvSequenceHandle handle) const { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(handle, snapshot) != PagedKvStatus::Ok) return 0; + uint32_t count = 0; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (snapshot.block_table[logical] != PAGED_KV_COLD_BLOCK && + is_protected(handle, logical)) { + ++count; + } + } + return count; +} + +uint32_t PagedKvResidencyManager::total_resident_count() const { + uint32_t total = 0; + for (uint32_t slot = 0; slot < sequences_.size(); ++slot) { + const SequenceState & state = sequences_[slot]; + if (!state.active) continue; + total += sequence_resident_count({slot, state.generation}); + } + return total; +} + +uint32_t PagedKvResidencyManager::quota_for_slot(uint32_t slot) const { + uint32_t active = 0; + uint32_t rank = 0; + for (uint32_t i = 0; i < sequences_.size(); ++i) { + if (!sequences_[i].active) continue; + if (i < slot) ++rank; + ++active; + } + if (active == 0 || slot >= sequences_.size() || + !sequences_[slot].active) { + return 0; + } + const uint32_t base = config_.resident_budget_blocks / active; + const uint32_t remainder = config_.resident_budget_blocks % active; + const uint32_t fair = base + (rank < remainder ? 1u : 0u); + const PagedKvSequenceHandle handle{slot, sequences_[slot].generation}; + return std::max(fair, sequence_protected_count(handle)); +} + +uint32_t PagedKvResidencyManager::fair_quota( + PagedKvSequenceHandle handle) const { + return validate_registered(handle) == PagedKvResidencyStatus::Ok + ? quota_for_slot(handle.slot) : 0; +} + +PagedKvResidencyManager::Victim PagedKvResidencyManager::choose_victim( + PagedKvSequenceHandle requester) const { + Victim best; + for (uint32_t slot = 0; slot < sequences_.size(); ++slot) { + const SequenceState & state = sequences_[slot]; + if (!state.active) continue; + const PagedKvSequenceHandle owner{slot, state.generation}; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(owner, snapshot) != PagedKvStatus::Ok) continue; + const uint32_t resident = sequence_resident_count(owner); + const uint32_t pending_page_outs = + sequence_pending_page_out_count(owner); + const uint32_t resident_after_pending = + resident > pending_page_outs ? resident - pending_page_outs : 0; + const uint32_t quota = quota_for_slot(slot); + + int class_rank = 2; + if (resident_after_pending > quota) class_rank = 0; + else if (slot == requester.slot && + state.generation == requester.generation) class_rank = 1; + + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (snapshot.block_table[logical] == PAGED_KV_COLD_BLOCK || + state.blocks[logical].write_pending || + state.blocks[logical].page_out_pending || + state.blocks[logical].page_in_pending || + state.blocks[logical].reservation_pending || + is_protected(owner, logical)) { + continue; + } + const BlockState & block = state.blocks[logical]; + // An unscored page is the first eviction candidate once relevance + // scoring is active for only part of a sequence. + const float score = block.score_valid + ? block.score : -std::numeric_limits::infinity(); + const bool better = !best.found || + class_rank < best.class_rank || + (class_rank == best.class_rank && score < best.score) || + (class_rank == best.class_rank && score == best.score && + block.last_use < best.last_use) || + (class_rank == best.class_rank && score == best.score && + block.last_use == best.last_use && + (slot < best.handle.slot || + (slot == best.handle.slot && + logical < best.logical_block))); + if (better) { + best.found = true; + best.handle = owner; + best.logical_block = logical; + best.class_rank = class_rank; + best.score = score; + best.last_use = block.last_use; + } + } + } + return best; +} + +PagedKvResidencyStatus PagedKvResidencyManager::make_room( + PagedKvSequenceHandle requester, uint32_t pool_blocks_needed, + uint32_t future_resident_blocks) { + const uint32_t free = pool_.free_block_count(); + const uint32_t need_for_pool = pool_blocks_needed > free + ? pool_blocks_needed - free : 0; + const uint32_t resident = total_resident_count(); + const uint64_t projected = + static_cast(resident) + future_resident_blocks; + const uint32_t need_for_budget = + projected > config_.resident_budget_blocks + ? static_cast(projected - config_.resident_budget_blocks) + : 0; + const uint32_t evictions = std::max(need_for_pool, need_for_budget); + for (uint32_t i = 0; i < evictions; ++i) { + const Victim victim = choose_victim(requester); + if (!victim.found) { + return finish_transfers(PagedKvResidencyStatus::NoEvictableBlock); + } + const auto status = evict_block_async( + victim.handle, victim.logical_block, + /*allow_protected=*/false); + if (status != PagedKvResidencyStatus::Ok) return finish_transfers(status); + } + // A page is recyclable only after the complete D2H batch succeeds. + return finish_transfers(PagedKvResidencyStatus::Ok); +} + +PagedKvResidencyStatus PagedKvResidencyManager::evict_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block, + bool allow_protected) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (logical_block >= snapshot.block_table.size()) { + return PagedKvResidencyStatus::InvalidArgument; + } + if (snapshot.block_table[logical_block] == PAGED_KV_COLD_BLOCK) { + return PagedKvResidencyStatus::InvalidArgument; + } + + BlockState & block = sequences_[handle.slot].blocks[logical_block]; + if (block.write_pending || block.page_out_pending || + block.page_in_pending || + (!allow_protected && is_protected(handle, logical_block))) { + return PagedKvResidencyStatus::NoEvictableBlock; + } + if (!block.host) { + try { + block.host = transfers_.allocate_pinned(config_.block_bytes); + } catch (...) { + block.host = nullptr; + } + if (!block.host) return PagedKvResidencyStatus::HostAllocationFailed; + stats_.host_bytes += config_.block_bytes; + } + + const uint32_t physical = snapshot.block_table[logical_block]; + bool queued = false; + try { + queued = transfers_.copy_out_async( + handle, logical_block, physical, block.host, + config_.block_bytes); + } catch (...) { + queued = false; + } + if (!queued) return PagedKvResidencyStatus::TransferFailed; + transfers_pending_ = true; + block.host_valid = false; + block.page_out_pending = true; + pending_page_outs_.push_back({handle, logical_block, physical}); + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::restore_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (logical_block >= snapshot.block_table.size()) { + return PagedKvResidencyStatus::InvalidArgument; + } + if (snapshot.block_table[logical_block] != PAGED_KV_COLD_BLOCK) { + sequences_[handle.slot].blocks[logical_block].last_use = ++clock_; + return PagedKvResidencyStatus::Ok; + } + BlockState & block = sequences_[handle.slot].blocks[logical_block]; + if (block.page_out_pending || block.page_in_pending) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + if (!block.host || !block.host_valid) { + return PagedKvResidencyStatus::HostCopyMissing; + } + + uint32_t physical = PAGED_KV_COLD_BLOCK; + const PagedKvStatus page_status = + pool_.page_in_block(handle, logical_block, physical); + if (page_status != PagedKvStatus::Ok) return from_pool_status(page_status); + + bool queued = false; + try { + queued = transfers_.copy_in_async( + handle, logical_block, physical, block.host, + config_.block_bytes); + } catch (...) { + queued = false; + } + if (!queued) { + uint32_t released = PAGED_KV_COLD_BLOCK; + if (pool_.page_out_block(handle, logical_block, released) != + PagedKvStatus::Ok) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + return PagedKvResidencyStatus::TransferFailed; + } + transfers_pending_ = true; + block.page_in_pending = true; + pending_page_ins_.push_back({handle, logical_block, physical}); + block.last_use = ++clock_; + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::ensure_resident( + PagedKvSequenceHandle handle, + const std::vector & logical_blocks) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + std::vector cold; + std::vector seen(snapshot.block_table.size(), 0); + // Validate the complete request before reserving any member. Otherwise a + // later bad index leaves earlier blocks permanently ineligible as victims. + if (std::any_of(logical_blocks.begin(), logical_blocks.end(), + [&](uint32_t logical) { + return logical >= snapshot.block_table.size(); + })) { + return PagedKvResidencyStatus::InvalidArgument; + } + for (uint32_t logical : logical_blocks) { + if (seen[logical]) continue; + seen[logical] = 1; + sequences_[handle.slot].blocks[logical].reservation_pending = true; + if (snapshot.block_table[logical] == PAGED_KV_COLD_BLOCK) { + cold.push_back(logical); + } else { + sequences_[handle.slot].blocks[logical].last_use = ++clock_; + } + } + const auto clear_reservations = [&] { + for (uint32_t logical = 0; logical < seen.size(); ++logical) { + if (seen[logical]) { + sequences_[handle.slot].blocks[logical].reservation_pending = false; + } + } + }; + const auto room = make_room( + handle, static_cast(cold.size()), + static_cast(cold.size())); + if (room != PagedKvResidencyStatus::Ok) { + clear_reservations(); + return room; + } + + // Restore the requested set from one reserved capacity pool. No member can + // become the victim of a later member in this same operation. + for (uint32_t logical : cold) { + const auto status = restore_block_async( + handle, logical); + if (status != PagedKvResidencyStatus::Ok) { + const auto finished = finish_transfers(status); + clear_reservations(); + return finished; + } + } + const auto finished = finish_transfers(PagedKvResidencyStatus::Ok); + clear_reservations(); + return finished; +} + +PagedKvResidencyStatus PagedKvResidencyManager::evict_block( + PagedKvSequenceHandle handle, uint32_t logical_block, + bool allow_protected) { + return finish_transfers( + evict_block_async(handle, logical_block, allow_protected)); +} + +PagedKvResidencyStatus PagedKvResidencyManager::touch( + PagedKvSequenceHandle handle, uint32_t logical_block) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + if (logical_block >= sequences_[handle.slot].blocks.size()) { + return PagedKvResidencyStatus::InvalidArgument; + } + sequences_[handle.slot].blocks[logical_block].last_use = ++clock_; + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::set_scores( + PagedKvSequenceHandle handle, const std::vector & scores) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + SequenceState & state = sequences_[handle.slot]; + if (scores.empty()) { + for (BlockState & block : state.blocks) block.score_valid = false; + return PagedKvResidencyStatus::Ok; + } + if (scores.size() != state.blocks.size() || + std::any_of(scores.begin(), scores.end(), + [](float score) { return !std::isfinite(score); })) { + return PagedKvResidencyStatus::InvalidArgument; + } + for (uint32_t i = 0; i < scores.size(); ++i) { + state.blocks[i].score = scores[i]; + state.blocks[i].score_valid = true; + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::reselect( + PagedKvSequenceHandle handle) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + stats_.reselects++; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + SequenceState & state = sequences_[handle.slot]; + + std::vector wanted(snapshot.block_table.size(), 0); + uint32_t wanted_count = 0; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (state.blocks[logical].write_pending || + is_protected(handle, logical)) { + wanted[logical] = 1; + ++wanted_count; + } + } + const uint32_t target = std::max( + wanted_count, + std::min(quota_for_slot(handle.slot), + snapshot.block_table.size())); + + std::vector candidates; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (!wanted[logical]) candidates.push_back(logical); + } + std::sort(candidates.begin(), candidates.end(), + [&](uint32_t a, uint32_t b) { + const BlockState & lhs = state.blocks[a]; + const BlockState & rhs = state.blocks[b]; + if (lhs.score_valid != rhs.score_valid) return lhs.score_valid; + if (lhs.score_valid && lhs.score != rhs.score) return lhs.score > rhs.score; + if (lhs.last_use != rhs.last_use) return lhs.last_use > rhs.last_use; + return a < b; + }); + for (uint32_t logical : candidates) { + if (wanted_count >= target) break; + wanted[logical] = 1; + ++wanted_count; + } + + for (uint32_t logical = 0; logical < wanted.size(); ++logical) { + if (wanted[logical]) state.blocks[logical].reservation_pending = true; + } + const auto clear_reservations = [&] { + for (uint32_t logical = 0; logical < wanted.size(); ++logical) { + if (wanted[logical]) state.blocks[logical].reservation_pending = false; + } + }; + + // Out first, making one batch of free physical pages before recalls. + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (!wanted[logical] && + snapshot.block_table[logical] != PAGED_KV_COLD_BLOCK) { + const auto status = evict_block_async( + handle, logical, /*allow_protected=*/false); + if (status != PagedKvResidencyStatus::Ok) { + const auto finished = finish_transfers(status); + clear_reservations(); + return finished; + } + } + } + auto finished = finish_transfers(PagedKvResidencyStatus::Ok); + if (finished != PagedKvResidencyStatus::Ok) { + clear_reservations(); + return finished; + } + + std::vector cold_wanted; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (wanted[logical] && !is_resident(handle, logical)) { + cold_wanted.push_back(logical); + } + } + const auto room = make_room( + handle, static_cast(cold_wanted.size()), + static_cast(cold_wanted.size())); + if (room != PagedKvResidencyStatus::Ok) { + clear_reservations(); + return room; + } + for (uint32_t logical : cold_wanted) { + const auto status = restore_block_async( + handle, logical); + if (status != PagedKvResidencyStatus::Ok) { + finished = finish_transfers(status); + clear_reservations(); + return finished; + } + } + finished = finish_transfers(PagedKvResidencyStatus::Ok); + clear_reservations(); + return finished; +} + +PagedKvResidencyStatus PagedKvResidencyManager::rebalance() { + while (true) { + // The dummy requester prevents any real sequence from getting the + // requester-swap class. A class-0 victim is necessarily over quota. + const PagedKvSequenceHandle none{ + std::numeric_limits::max(), 0}; + const Victim victim = choose_victim(none); + if (!victim.found || victim.class_rank != 0) break; + const auto status = evict_block_async( + victim.handle, victim.logical_block, + /*allow_protected=*/false); + if (status != PagedKvResidencyStatus::Ok) { + return finish_transfers(status); + } + } + return finish_transfers(PagedKvResidencyStatus::Ok); +} + +bool PagedKvResidencyManager::is_resident( + PagedKvSequenceHandle handle, uint32_t logical_block) const { + if (validate_registered(handle) != PagedKvResidencyStatus::Ok) return false; + PagedKvSequenceSnapshot snapshot; + return pool_.sequence(handle, snapshot) == PagedKvStatus::Ok && + logical_block < snapshot.block_table.size() && + snapshot.block_table[logical_block] != PAGED_KV_COLD_BLOCK; +} + +PagedKvResidencyStats PagedKvResidencyManager::stats() const { + PagedKvResidencyStats out = stats_; + out.resident_blocks = total_resident_count(); + return out; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/paged_kv_residency.h b/server/src/common/concurrency/paged_kv_residency.h new file mode 100644 index 000000000..dc6f67cf5 --- /dev/null +++ b/server/src/common/concurrency/paged_kv_residency.h @@ -0,0 +1,260 @@ +// Multi-sequence host residency for PagedKvPool full-attention K/V blocks. +// +// The policy and bookkeeping are backend-neutral. The Qwen engine supplies a +// pinned allocator plus async full-block D2H/H2D callbacks that capture its +// dedicated copy stream and know how to concatenate every full-attention K/V +// tensor into one host block. The manager batches transfers and synchronizes +// once before returning any operation that permits device K/V to be read or a +// recycled physical block to be written. + +#pragma once + +#include "paged_kv_pool.h" + +#include +#include +#include +#include + +namespace dflash::common { + +enum class PagedKvResidencyStatus : uint8_t { + Ok = 0, + InvalidArgument, + SequenceNotRegistered, + StaleHandle, + PoolExhausted, + NoEvictableBlock, + HostAllocationFailed, + TransferFailed, + HostCopyMissing, + InconsistentPoolState, +}; + +const char * paged_kv_residency_status_string(PagedKvResidencyStatus status); + +struct PagedKvResidencyConfig { + // Bytes copied for one physical pool block across all full-attention K/V + // tensors. Recurrent/DeltaNet state is intentionally outside this pager. + size_t block_bytes = 0; + + // Hard bound for materialized logical blocks. Zero uses the complete + // physical PagedKvPool. Reserved-but-unappended pool blocks still consume + // physical capacity and may temporarily make the effective bound smaller. + uint32_t resident_budget_blocks = 0; + + // Attention sinks and the trailing local window are never automatic + // eviction victims. Explicit evict_block(..., allow_protected=true) is + // available for teardown/tests only. + uint32_t sink_blocks = 1; + uint32_t tail_blocks = 4; +}; + +struct PagedKvResidencyTransferOps { + // Production implementations should use cudaMallocHost/hipHostMalloc (or + // the equivalent backend pinned allocator). Every allocation is exactly + // config.block_bytes and lives until forget_sequence()/reset(). + std::function allocate_pinned; + std::function free_pinned; + + // Queue a complete physical-block copy on one dedicated copy stream. + // handle/logical_block are supplied for diagnostics only; physical_block + // identifies the source/destination in the Qwen paged K/V tensors. + std::function copy_out_async; + std::function copy_in_async; + + // Synchronize that copy stream. The manager calls it once after a batch + // and before attention reads or append writes can observe remapped blocks. + std::function synchronize; +}; + +struct PagedKvResidencyStats { + uint64_t page_ins = 0; + uint64_t page_outs = 0; + uint64_t resident_blocks = 0; + uint64_t reselects = 0; + uint64_t host_bytes = 0; + uint64_t moved_bytes = 0; +}; + +struct PagedKvResidentAppendResult { + PagedKvResidencyStatus status = PagedKvResidencyStatus::Ok; + PagedKvAppendResult pool_result; + + explicit operator bool() const { + return status == PagedKvResidencyStatus::Ok && + pool_result.status == PagedKvStatus::Ok; + } +}; + +class PagedKvResidencyManager { +public: + // Throws std::invalid_argument for an empty transfer callback, zero block + // bytes, or a resident budget larger than the physical pool. + PagedKvResidencyManager(PagedKvPool & pool, + PagedKvResidencyConfig config, + PagedKvResidencyTransferOps transfers); + ~PagedKvResidencyManager(); + + PagedKvResidencyManager(const PagedKvResidencyManager &) = delete; + PagedKvResidencyManager & operator=(const PagedKvResidencyManager &) = delete; + + // Call immediately after PagedKvPool::acquire[_reserved](). Registration + // rejects pre-existing cold entries because their host backing is unknown. + PagedKvResidencyStatus register_sequence(PagedKvSequenceHandle handle); + + // Free this sequence's pinned backing. Call before pool.release, and do + // not release the pool handle unless this returns Ok: a failed transfer + // barrier keeps its physical pages and pinned buffers quarantined until + // forget_sequence is retried. The handle must still match the manager's + // registered generation. + PagedKvResidencyStatus forget_sequence(PagedKvSequenceHandle handle); + void reset(); + + // Recommended append integration: restores a cold partial append head, + // fairly evicts enough blocks, synchronizes the copy stream, appends in + // PagedKvPool, and marks every touched block pending until the engine has + // finished the target graph that writes the returned physical rows. + PagedKvResidentAppendResult append(PagedKvSequenceHandle handle, + uint32_t token_count, + bool only_first_last_slots = false); + + // Split integration for callers whose slot manager owns pool.append(): + // prepare_append(handle, n); pool.append(handle, n); observe_append(...) + // prepare_append always synchronizes before returning Ok. observe_append + // must be called before the next residency operation. + PagedKvResidencyStatus prepare_append(PagedKvSequenceHandle handle, + uint32_t token_count); + PagedKvResidencyStatus observe_append( + PagedKvSequenceHandle handle, const PagedKvAppendResult & result); + + // Clear all blocks staged by append()/observe_append() for this sequence. + // Call only after synchronizing the target compute that wrote every returned + // physical row. Pending blocks cannot be evicted or deselected, preventing a + // later slot staged in the same packed step from recycling an unwritten page. + PagedKvResidencyStatus commit_pending_writes( + PagedKvSequenceHandle handle); + + // Restore the requested host-backed logical blocks as one transfer batch. + // The operation synchronizes before returning Ok, so attention may read + // the returned pool block table immediately. + PagedKvResidencyStatus ensure_resident( + PagedKvSequenceHandle handle, + const std::vector & logical_blocks); + + // Explicit single-block eviction. Automatic policy never evicts sink/tail + // blocks; protected eviction requires an explicit opt-in. Pending target + // writes are never evictable, including with allow_protected=true. + PagedKvResidencyStatus evict_block(PagedKvSequenceHandle handle, + uint32_t logical_block, + bool allow_protected = false); + + // Update recency after a block participates in attention. + PagedKvResidencyStatus touch(PagedKvSequenceHandle handle, + uint32_t logical_block); + + // Optional relevance array, one score per materialized logical block. + // Higher values are retained. Without scores, reselection uses LRU. + PagedKvResidencyStatus set_scores(PagedKvSequenceHandle handle, + const std::vector & scores); + PagedKvResidencyStatus reselect(PagedKvSequenceHandle handle); + + // Evict unprotected blocks above each active sequence's deterministic + // fair share. Spare capacity remains borrowable; it is reclaimed first + // from borrowers when another sequence needs a block. + PagedKvResidencyStatus rebalance(); + + // Explicit barrier for engine paths that directly issue operations and + // then read the paged K/V tensors. Normally append/ensure/reselect/evict + // already provide this barrier. + PagedKvResidencyStatus synchronize_before_read(); + + uint32_t fair_quota(PagedKvSequenceHandle handle) const; + bool is_resident(PagedKvSequenceHandle handle, + uint32_t logical_block) const; + PagedKvResidencyStats stats() const; + +private: + struct BlockState { + void * host = nullptr; + bool host_valid = false; + bool write_pending = false; + bool page_out_pending = false; + bool page_in_pending = false; + bool reservation_pending = false; + bool score_valid = false; + float score = 0.0f; + uint64_t last_use = 0; + }; + + struct SequenceState { + bool active = false; + uint64_t generation = 0; + std::vector blocks; + }; + + struct Victim { + bool found = false; + PagedKvSequenceHandle handle; + uint32_t logical_block = 0; + int class_rank = 0; + float score = 0.0f; + uint64_t last_use = 0; + }; + + struct PendingTransfer { + PagedKvSequenceHandle handle; + uint32_t logical_block = 0; + uint32_t physical_block = PAGED_KV_COLD_BLOCK; + }; + + PagedKvResidencyStatus validate_registered( + PagedKvSequenceHandle handle) const; + PagedKvResidencyStatus refresh_sequence( + PagedKvSequenceHandle handle, bool appended_rows = false, + uint32_t append_first = 0, uint32_t append_count = 0); + PagedKvResidencyStatus finish_transfers( + PagedKvResidencyStatus status); + + bool is_protected(PagedKvSequenceHandle handle, + uint32_t logical_block) const; + uint32_t sequence_resident_count(PagedKvSequenceHandle handle) const; + uint32_t sequence_pending_page_out_count( + PagedKvSequenceHandle handle) const; + uint32_t sequence_protected_count(PagedKvSequenceHandle handle) const; + uint32_t total_resident_count() const; + uint32_t quota_for_slot(uint32_t slot) const; + + Victim choose_victim(PagedKvSequenceHandle requester) const; + PagedKvResidencyStatus make_room(PagedKvSequenceHandle requester, + uint32_t pool_blocks_needed, + uint32_t future_resident_blocks); + PagedKvResidencyStatus evict_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block, + bool allow_protected); + PagedKvResidencyStatus restore_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block); + + void free_sequence_buffers(SequenceState & state) noexcept; + + PagedKvPool & pool_; + PagedKvResidencyConfig config_; + PagedKvResidencyTransferOps transfers_; + std::vector sequences_; + PagedKvResidencyStats stats_; + uint64_t clock_ = 0; + bool transfers_pending_ = false; + bool transfer_barrier_failed_ = false; + std::vector pending_page_outs_; + std::vector pending_page_ins_; +}; + +} // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer.cpp b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp new file mode 100644 index 000000000..111e6a36f --- /dev/null +++ b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp @@ -0,0 +1,405 @@ +#include "qwen_paged_kv_transfer.h" + +#include "common/gpu_runtime_compat.h" +#include "internal.h" + +#include "ggml.h" + +#include +#include +#include + +namespace dflash::common { +namespace { + +void set_error(std::string * error, const std::string & message) { + if (error) *error = message; +} + +bool runtime_pointer_is_device(const void * pointer, int device) { + cudaPointerAttributes attributes{}; + const cudaError_t status = cudaPointerGetAttributes(&attributes, pointer); + if (status != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + return attributes.type == hipMemoryTypeDevice && + attributes.device == device; +#else + return attributes.type == cudaMemoryTypeDevice && + attributes.device == device; +#endif +} + +bool tensor_is_device_backed(const ggml_tensor * tensor, int device, + std::string * error) { + if (!tensor || !tensor->buffer || !tensor->data) { + set_error(error, "paged K/V tensor is null or unallocated"); + return false; + } + const ggml_backend_buffer_type_t buft = + ggml_backend_buffer_get_type(tensor->buffer); + if (!buft || ggml_backend_buft_is_meta(buft)) { + set_error(error, + "meta/tensor-parallel paged K/V buffers are unsupported"); + return false; + } + if (ggml_backend_buft_is_host(buft)) { + set_error(error, "host-backed paged K/V buffers are unsupported"); + return false; + } + const ggml_backend_dev_t tensor_device = + ggml_backend_buft_get_device(buft); + if (!tensor_device) { + set_error(error, "paged K/V buffer has no backend device"); + return false; + } + const enum ggml_backend_dev_type tensor_device_type = + ggml_backend_dev_type(tensor_device); + if (tensor_device_type != GGML_BACKEND_DEVICE_TYPE_GPU && + tensor_device_type != GGML_BACKEND_DEVICE_TYPE_IGPU) { + set_error(error, "paged K/V buffer is not GPU device memory"); + return false; + } + if (!runtime_pointer_is_device(tensor->data, device)) { + set_error(error, + "paged K/V tensor is not on the requested HIP/CUDA device"); + return false; + } + return true; +} + +} // namespace + +struct QwenPagedKvResidencyTransfer::State { + struct TensorCopy { + uint8_t * device_data = nullptr; + QwenPagedKvTensorLayout layout; + size_t host_offset = 0; + size_t host_head_bytes = 0; + }; + + int device = -1; + uint32_t block_size = 0; + size_t block_bytes = 0; + uint32_t physical_block_count = 0; + cudaStream_t stream = nullptr; + bool stream_may_reference_host = false; + std::vector tensors; + std::vector pinned_allocations; + + ~State() { + if (device >= 0 && cudaSetDevice(device) == cudaSuccess) { + if (stream && cudaStreamSynchronize(stream) != cudaSuccess) { + // Match PagedKvResidencyManager's fail-safe teardown rule: + // leak stream/backing rather than free host memory that a + // failed runtime may still reference. + pinned_allocations.clear(); + stream = nullptr; + return; + } + for (void * pointer : pinned_allocations) { + if (pointer) (void)cudaFreeHost(pointer); + } + pinned_allocations.clear(); + if (stream) (void)cudaStreamDestroy(stream); + } + stream = nullptr; + } + + bool select_device() const { + return device >= 0 && cudaSetDevice(device) == cudaSuccess; + } + + void * allocate_pinned(size_t bytes) { + if (bytes != block_bytes || !select_device()) return nullptr; + void * pointer = nullptr; + if (cudaMallocHost(&pointer, bytes) != cudaSuccess || !pointer) { + (void)cudaGetLastError(); + return nullptr; + } + try { + pinned_allocations.push_back(pointer); + } catch (...) { + (void)cudaFreeHost(pointer); + return nullptr; + } + return pointer; + } + + void free_pinned(void * pointer) { + if (!pointer) return; + const auto found = std::find( + pinned_allocations.begin(), pinned_allocations.end(), pointer); + if (found == pinned_allocations.end()) return; + if (stream_may_reference_host) { + if (!stream || !select_device() || + cudaStreamSynchronize(stream) != cudaSuccess) { + return; + } + stream_may_reference_host = false; + } + if (select_device() && cudaFreeHost(pointer) == cudaSuccess) { + pinned_allocations.erase(found); + } + } + + bool synchronize() { + if (!stream || !select_device()) return false; + if (cudaStreamSynchronize(stream) != cudaSuccess) { + (void)cudaGetLastError(); + stream_may_reference_host = true; + return false; + } + stream_may_reference_host = false; + return true; + } + + bool queue_copy(uint32_t physical_block, void * host, size_t bytes, + bool copy_out) { + if (!host || bytes != block_bytes || !stream || + physical_block >= physical_block_count || !select_device()) { + return false; + } + + const size_t first_row = + static_cast(physical_block) * block_size; + uint8_t * const host_bytes = static_cast(host); + for (const TensorCopy & tensor : tensors) { + uint8_t * const device_block = tensor.device_data + + first_row * tensor.layout.row_stride; + uint8_t * const host_tensor = + host_bytes + tensor.host_offset; + + cudaError_t status = cudaSuccess; + if (tensor.layout.row_stride == tensor.layout.row_bytes) { + // Common Qwen cache layout: block rows are contiguous inside + // each head plane, so one 2D copy covers every head. + status = copy_out + ? cudaMemcpy2DAsync( + host_tensor, tensor.host_head_bytes, + device_block, tensor.layout.head_stride, + tensor.host_head_bytes, tensor.layout.heads, + cudaMemcpyDeviceToHost, stream) + : cudaMemcpy2DAsync( + device_block, tensor.layout.head_stride, + host_tensor, tensor.host_head_bytes, + tensor.host_head_bytes, tensor.layout.heads, + cudaMemcpyHostToDevice, stream); + } else { + // Preserve uncommon row padding with one pitched async copy + // per head. Host images remain tightly packed payload bytes. + for (size_t head = 0; + head < static_cast(tensor.layout.heads); + ++head) { + uint8_t * const device_head = + device_block + head * tensor.layout.head_stride; + uint8_t * const host_head = + host_tensor + head * tensor.host_head_bytes; + status = copy_out + ? cudaMemcpy2DAsync( + host_head, tensor.layout.row_bytes, + device_head, tensor.layout.row_stride, + tensor.layout.row_bytes, block_size, + cudaMemcpyDeviceToHost, stream) + : cudaMemcpy2DAsync( + device_head, tensor.layout.row_stride, + host_head, tensor.layout.row_bytes, + tensor.layout.row_bytes, block_size, + cudaMemcpyHostToDevice, stream); + if (status != cudaSuccess) break; + } + } + if (status != cudaSuccess) { + // A callback that returns false is not marked pending by the + // residency manager. Drain any prefix already queued here so + // its pinned buffer can still be released safely. + (void)cudaGetLastError(); + if (cudaStreamSynchronize(stream) != cudaSuccess) { + stream_may_reference_host = true; + } + return false; + } + } + return true; + } +}; + +QwenPagedKvResidencyTransfer::QwenPagedKvResidencyTransfer( + std::shared_ptr state) + : state_(std::move(state)) {} + +QwenPagedKvResidencyTransfer::~QwenPagedKvResidencyTransfer() = default; + +std::unique_ptr +QwenPagedKvResidencyTransfer::create( + const TargetCache & cache, + ggml_backend_t backend, + int device, + uint32_t block_size, + std::string * error) { + if (error) error->clear(); + if (!backend || device < 0 || block_size == 0) { + set_error(error, "invalid paged K/V transfer backend/device/block size"); + return nullptr; + } + const ggml_backend_dev_t backend_device = ggml_backend_get_device(backend); + if (!backend_device) { + set_error(error, "paged K/V transfer backend has no device"); + return nullptr; + } + const enum ggml_backend_dev_type backend_type = + ggml_backend_dev_type(backend_device); + if (backend_type == GGML_BACKEND_DEVICE_TYPE_META || + ggml_backend_buft_is_meta( + ggml_backend_get_default_buffer_type(backend))) { + set_error(error, + "meta/tensor-parallel paged K/V transfer is unsupported"); + return nullptr; + } + if (backend_type != GGML_BACKEND_DEVICE_TYPE_GPU && + backend_type != GGML_BACKEND_DEVICE_TYPE_IGPU) { + set_error(error, "paged K/V transfer requires a GPU backend"); + return nullptr; + } + if (cache.backend && cache.backend != backend) { + set_error(error, "paged K/V cache belongs to a different backend"); + return nullptr; + } + if (cache.attn_k.empty() || + cache.attn_k.size() != cache.attn_v.size()) { + set_error(error, "paged K/V cache has no complete attention layers"); + return nullptr; + } + + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || + device >= device_count || cudaSetDevice(device) != cudaSuccess) { + (void)cudaGetLastError(); + set_error(error, "invalid HIP/CUDA device for paged K/V transfer"); + return nullptr; + } + + std::vector cache_tensors; + std::vector layouts; + try { + cache_tensors.reserve(cache.attn_k.size() * 2); + layouts.reserve(cache.attn_k.size() * 2); + } catch (...) { + set_error(error, "paged K/V transfer layout allocation failed"); + return nullptr; + } + + for (size_t layer = 0; layer < cache.attn_k.size(); ++layer) { + const ggml_tensor * pair[] = { + cache.attn_k[layer], cache.attn_v[layer], + }; + if (!pair[0] || !pair[1]) { + set_error(error, + "partial/tensor-parallel paged K/V cache is unsupported"); + return nullptr; + } + if (pair[0]->ne[1] != pair[1]->ne[1] || + pair[0]->ne[2] != pair[1]->ne[2]) { + set_error(error, "K/V cache pair dimensions do not match"); + return nullptr; + } + for (const ggml_tensor * tensor : pair) { + if (!tensor_is_device_backed(tensor, device, error)) return nullptr; + if (tensor->ne[0] <= 0 || tensor->ne[1] <= 0 || + tensor->ne[2] <= 0 || tensor->ne[3] != 1 || + tensor->ne[0] % ggml_blck_size(tensor->type) != 0) { + set_error(error, "unsupported paged K/V tensor dimensions"); + return nullptr; + } + const size_t row_bytes = + ggml_row_size(tensor->type, tensor->ne[0]); + cache_tensors.push_back(tensor); + layouts.push_back({ + row_bytes, + tensor->nb[1], + tensor->nb[2], + ggml_nbytes(tensor), + static_cast(tensor->ne[1]), + static_cast(tensor->ne[2]), + }); + } + } + + QwenPagedKvBlockLayout plan; + if (!plan_qwen_paged_kv_block_layout( + layouts, block_size, plan, error)) { + return nullptr; + } + + std::shared_ptr state; + try { + state = std::make_shared(); + state->device = device; + state->block_size = block_size; + state->block_bytes = plan.block_bytes; + state->physical_block_count = plan.physical_block_count; + state->tensors.reserve(cache_tensors.size()); + for (size_t i = 0; i < cache_tensors.size(); ++i) { + state->tensors.push_back({ + static_cast(cache_tensors[i]->data), + layouts[i], + plan.tensor_offsets[i], + plan.tensor_head_bytes[i], + }); + } + } catch (...) { + set_error(error, "paged K/V transfer state allocation failed"); + return nullptr; + } + + if (cudaStreamCreateWithFlags( + &state->stream, cudaStreamNonBlocking) != cudaSuccess) { + (void)cudaGetLastError(); + set_error(error, "failed to create paged K/V transfer stream"); + return nullptr; + } + + try { + return std::unique_ptr( + new QwenPagedKvResidencyTransfer(std::move(state))); + } catch (...) { + set_error(error, "paged K/V transfer wrapper allocation failed"); + return nullptr; + } +} + +size_t QwenPagedKvResidencyTransfer::block_bytes() const noexcept { + return state_ ? state_->block_bytes : 0; +} + +PagedKvResidencyTransferOps +QwenPagedKvResidencyTransfer::callbacks() const { + PagedKvResidencyTransferOps result; + const std::shared_ptr state = state_; + if (!state) return result; + result.allocate_pinned = [state](size_t bytes) { + return state->allocate_pinned(bytes); + }; + result.free_pinned = [state](void * pointer) { + state->free_pinned(pointer); + }; + result.copy_out_async = + [state](PagedKvSequenceHandle, uint32_t, uint32_t physical_block, + void * host, size_t bytes) { + return state->queue_copy( + physical_block, host, bytes, /*copy_out=*/true); + }; + result.copy_in_async = + [state](PagedKvSequenceHandle, uint32_t, uint32_t physical_block, + const void * host, size_t bytes) { + return state->queue_copy( + physical_block, const_cast(host), bytes, + /*copy_out=*/false); + }; + result.synchronize = [state] { return state->synchronize(); }; + return result; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer.h b/server/src/common/concurrency/qwen_paged_kv_transfer.h new file mode 100644 index 000000000..2aee297a4 --- /dev/null +++ b/server/src/common/concurrency/qwen_paged_kv_transfer.h @@ -0,0 +1,87 @@ +// Production HIP/CUDA block transfers for Qwen paged K/V residency. +// +// A host image packs the payload bytes for one physical block from every +// full-attention K and V tensor. Device row/head strides are retained in the +// copy plan, so quantized K/V types and padded tensor layouts are copied +// without reinterpretation. + +#pragma once + +#include "paged_kv_residency.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct TargetCache; + +// GPU-independent input/output for block-layout validation. storage_bytes is +// the accessible span starting at the tensor data pointer (ggml_nbytes for a +// normal cache tensor). +struct QwenPagedKvTensorLayout { + size_t row_bytes = 0; + size_t row_stride = 0; + size_t head_stride = 0; + size_t storage_bytes = 0; + uint64_t physical_rows = 0; + uint64_t heads = 0; +}; + +struct QwenPagedKvBlockLayout { + size_t block_bytes = 0; + uint32_t physical_block_count = 0; + // One entry per input tensor. Tensor payloads are packed consecutively; + // each head occupies tensor_head_bytes[i] bytes in the host image. + std::vector tensor_offsets; + std::vector tensor_head_bytes; +}; + +// Validates all dimensions/strides and computes the packed host block image. +// Every tensor must cover the same physical row count and head count. Row +// padding is supported; padding bytes are not copied into the host image. +bool plan_qwen_paged_kv_block_layout( + const std::vector & tensors, + uint32_t block_size, + QwenPagedKvBlockLayout & out, + std::string * error = nullptr); + +// Owns the dedicated nonblocking transfer stream. callbacks() retains shared +// ownership of the stream/layout state, so this wrapper may be destroyed as +// soon as callbacks are handed to PagedKvResidencyManager. TargetCache and +// its K/V buffers must outlive those callbacks. Before asking the residency +// manager to evict, the engine must have completed the compute work that last +// wrote the source block; the manager's synchronize callback supplies the +// opposite copy-stream -> later-attention/append barrier. +class QwenPagedKvResidencyTransfer { +public: + static std::unique_ptr create( + const TargetCache & cache, + ggml_backend_t backend, + int device, + uint32_t block_size, + std::string * error = nullptr); + + ~QwenPagedKvResidencyTransfer(); + + QwenPagedKvResidencyTransfer( + const QwenPagedKvResidencyTransfer &) = delete; + QwenPagedKvResidencyTransfer & operator=( + const QwenPagedKvResidencyTransfer &) = delete; + + size_t block_bytes() const noexcept; + PagedKvResidencyTransferOps callbacks() const; + +private: + struct State; + explicit QwenPagedKvResidencyTransfer(std::shared_ptr state); + + std::shared_ptr state_; +}; + +} // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp b/server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp new file mode 100644 index 000000000..d75a3f864 --- /dev/null +++ b/server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp @@ -0,0 +1,119 @@ +#include "qwen_paged_kv_transfer.h" + +#include + +namespace dflash::common { +namespace { + +bool fail(QwenPagedKvBlockLayout & out, std::string * error, + const char * message) { + out = {}; + if (error) *error = message; + return false; +} + +bool checked_mul(size_t lhs, size_t rhs, size_t & out) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return false; + } + out = lhs * rhs; + return true; +} + +bool checked_add(size_t lhs, size_t rhs, size_t & out) { + if (rhs > std::numeric_limits::max() - lhs) return false; + out = lhs + rhs; + return true; +} + +} // namespace + +bool plan_qwen_paged_kv_block_layout( + const std::vector & tensors, + uint32_t block_size, + QwenPagedKvBlockLayout & out, + std::string * error) { + out = {}; + if (tensors.empty() || block_size == 0) { + return fail(out, error, "empty paged K/V layout or zero block size"); + } + + const uint64_t physical_rows = tensors.front().physical_rows; + const uint64_t heads = tensors.front().heads; + if (physical_rows == 0 || heads == 0 || + physical_rows % block_size != 0 || + physical_rows / block_size > std::numeric_limits::max()) { + return fail(out, error, "invalid paged K/V physical dimensions"); + } + out.physical_block_count = + static_cast(physical_rows / block_size); + + try { + out.tensor_offsets.reserve(tensors.size()); + out.tensor_head_bytes.reserve(tensors.size()); + } catch (...) { + return fail(out, error, "paged K/V layout allocation failed"); + } + + size_t host_offset = 0; + for (const QwenPagedKvTensorLayout & tensor : tensors) { + if (tensor.row_bytes == 0 || tensor.row_stride < tensor.row_bytes || + tensor.head_stride == 0 || tensor.storage_bytes == 0 || + tensor.physical_rows != physical_rows || tensor.heads != heads) { + return fail(out, error, "inconsistent paged K/V tensor layout"); + } + if (physical_rows > std::numeric_limits::max() || + heads > std::numeric_limits::max()) { + return fail(out, error, "paged K/V dimensions exceed host size_t"); + } + + size_t last_row_offset = 0; + size_t last_head_offset = 0; + size_t used_end = 0; + if (!checked_mul(static_cast(physical_rows - 1), + tensor.row_stride, last_row_offset) || + !checked_mul(static_cast(heads - 1), + tensor.head_stride, last_head_offset) || + !checked_add(last_head_offset, last_row_offset, used_end) || + !checked_add(used_end, tensor.row_bytes, used_end) || + used_end > tensor.storage_bytes) { + return fail(out, error, "paged K/V tensor strides exceed storage"); + } + + // Heads must not overlap. Larger head strides (alignment/padding) are + // fine because device copies retain the original pitch. + size_t plane_span = 0; + if (!checked_mul(static_cast(physical_rows - 1), + tensor.row_stride, plane_span) || + !checked_add(plane_span, tensor.row_bytes, plane_span) || + tensor.head_stride < plane_span) { + return fail(out, error, "overlapping paged K/V head planes"); + } + + size_t head_bytes = 0; + size_t tensor_bytes = 0; + size_t next_offset = 0; + if (!checked_mul(tensor.row_bytes, block_size, head_bytes) || + !checked_mul(head_bytes, static_cast(heads), + tensor_bytes) || + !checked_add(host_offset, tensor_bytes, next_offset)) { + return fail(out, error, "paged K/V host block size overflow"); + } + try { + out.tensor_offsets.push_back(host_offset); + out.tensor_head_bytes.push_back(head_bytes); + } catch (...) { + return fail(out, error, "paged K/V layout allocation failed"); + } + host_offset = next_offset; + } + + if (host_offset == 0) { + return fail(out, error, "zero-byte paged K/V host block"); + } + out.block_bytes = host_offset; + if (error) error->clear(); + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index fa4dbba35..fa2909da3 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -8,9 +8,10 @@ // interface. admit() claims a slot and queues its prompt without compute. // Each step() then advances a scheduler-selected cohort of prompt slices // alongside the complete live decode batch. Once a prefill completes, the -// scheduler advances that slot one token per step(), feeding each sampled -// token back as the next step's input — which is what lets it override a token -// (thinking-budget force-close) before it is committed to the cache. +// scheduler feeds the final sampled token back as the next step's input. An +// engine may also return speculative children it already committed before that +// final pending token. The scheduler can disable that burst path per slot when +// it must retain authority to substitute a thinking-budget close token. // // The split of duties is deliberate and is the reason this interface exists // apart from ModelBackend: @@ -181,14 +182,31 @@ class SeqEngine { struct StepInput { int slot = -1; int32_t token = -1; // token to commit at this slot's next position + // False when scheduler-side policy may replace the sampled token + // before it is committed (currently the thinking-budget close hook). + bool allow_speculation = true; }; struct DecodeOutput { int slot = -1; - int32_t token = -1; // newly sampled token (pending until next step) + // Final newly sampled token, pending until the scheduler feeds it into + // the next step. `committed_tokens`, when non-empty, precede this token + // and are already present in backend state. + int32_t token = -1; bool failed = false; // Present when failed=true so the scheduler can report an honest // per-request error instead of silently truncating generation. std::string error; + std::vector committed_tokens; + + // Per-slot deltas for this engine step. The scheduler aggregates them + // until retirement and emits one machine-readable proof record. + uint64_t ddtree_steps = 0; + uint64_t ddtree_accepted_tokens = 0; + uint64_t target_forwards = 0; + uint64_t kvflash_page_ins = 0; + uint64_t kvflash_page_outs = 0; + uint64_t kvflash_resident_blocks = 0; + uint64_t kvflash_reselects = 0; }; struct PrefillOutput { @@ -250,6 +268,19 @@ class SeqEngine { virtual bool token_is_eos(int32_t token) const = 0; }; +// Deliver a successful decode result in wire order. The visitor returns false +// after a stop/EOS/output-cap decision; in that case later committed children +// and the final pending token are intentionally hidden and the slot is retired. +template +inline bool consume_decode_output_tokens( + const SeqEngine::DecodeOutput & output, Advance advance) { + if (output.failed) return false; + for (int32_t token : output.committed_tokens) { + if (!advance(token)) return false; + } + return advance(output.token); +} + // Validate the model-neutral step protocol before the scheduler consumes any // output. Malformed row ownership is fatal because re-feeding a token after an // omitted output would silently corrupt that sequence. @@ -265,6 +296,7 @@ inline std::string validate_step_result( } std::vector decode_planned((size_t)slot_count, 0); + std::vector speculation_allowed((size_t)slot_count, 0); std::vector prefill_planned((size_t)slot_count, 0); for (const SeqEngine::StepInput & input : plan.decode) { if (input.slot < 0 || input.slot >= slot_count || input.token < 0) @@ -272,6 +304,8 @@ inline std::string validate_step_result( if (decode_planned[(size_t)input.slot]) return "decode plan contains a duplicate slot"; decode_planned[(size_t)input.slot] = 1; + speculation_allowed[(size_t)input.slot] = + input.allow_speculation ? 1 : 0; } for (const PrefillSlice & slice : plan.prefills) { if (slice.slot < 0 || slice.slot >= slot_count || @@ -290,10 +324,25 @@ inline std::string validate_step_result( return "decode output names an unplanned slot"; if (decode_seen[(size_t)output.slot]) return "step returned duplicate decode outputs"; - if (output.failed && (output.token >= 0 || output.error.empty())) - return "failed decode has invalid payload"; - if (!output.failed && (output.token < 0 || !output.error.empty())) - return "successful decode has invalid payload"; + if (output.failed) { + if (output.error.empty()) + return "failed decode has no diagnostic"; + if (output.token >= 0 || !output.committed_tokens.empty()) + return "failed decode exposes token payload"; + } else { + if (output.token < 0) + return "successful decode has no pending token"; + if (!output.error.empty()) + return "successful decode carries an error diagnostic"; + if (!speculation_allowed[(size_t)output.slot] && + !output.committed_tokens.empty()) + return "decode output burst violates disabled speculation"; + if (std::any_of( + output.committed_tokens.begin(), + output.committed_tokens.end(), + [](int32_t token) { return token < 0; })) + return "decode output burst contains an invalid token"; + } decode_seen[(size_t)output.slot] = 1; } diff --git a/server/src/common/ddtree.cpp b/server/src/common/ddtree.cpp index 08ca33464..54402c2cc 100644 --- a/server/src/common/ddtree.cpp +++ b/server/src/common/ddtree.cpp @@ -223,4 +223,14 @@ std::vector follow_verified_tree(const DDTree & tree, return accepted; } +bool truncate_verified_path(std::vector & accepted, + std::size_t max_committed, + const int32_t * posterior, + int & out_next_token) { + if (accepted.size() <= max_committed) return false; + accepted.resize(max_committed); + out_next_token = accepted.empty() ? -1 : posterior[accepted.back()]; + return true; +} + } // namespace dflash::common diff --git a/server/src/common/ddtree.h b/server/src/common/ddtree.h index afe22f226..718c05b57 100644 --- a/server/src/common/ddtree.h +++ b/server/src/common/ddtree.h @@ -8,6 +8,7 @@ #pragma once +#include #include #include #include @@ -61,4 +62,13 @@ std::vector follow_verified_tree(const DDTree & tree, int & out_next_token, int * out_node_idx = nullptr); +// Bound a verified path to the number of tokens that can actually be +// committed. When truncation removes the old tip, the pending token must be +// recomputed from the posterior at the new tip; otherwise it describes model +// state that was never committed. +bool truncate_verified_path(std::vector & accepted, + std::size_t max_committed, + const int32_t * posterior, + int & out_next_token); + } // namespace dflash::common diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index f655f5d4f..63099832b 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -163,6 +163,26 @@ std::string check_feature_compatibility( "' does not support PFlash compression"; } + const bool concurrent_paged_qwen = + arch == "qwen35" && args.paged_attention && + args.max_concurrency > 1 && + !args.device.is_layer_split() && + !args.remote_target_shard.enabled() && + args.fa_window == 0; + const bool concurrent_local_paged_qwen = + concurrent_paged_qwen && !args.remote_draft.enabled() && + target_backend == draft_backend && + args.device.gpu == args.draft_device.gpu && + !args.device.is_tensor_parallel(); + const bool concurrent_local_ddtree = + concurrent_local_paged_qwen && args.draft_path != nullptr && + args.ddtree_mode; + + if (args.ddtree_mode && + (args.ddtree_budget < 1 || args.ddtree_budget > 255)) { + return "--ddtree-budget must be in [1, 255]"; + } + // ── --paged-attention × architecture, placement, and decode features // Paged decode swaps the contiguous K/V cache for a block table owned by // the monolithic qwen35 backend, so every rule below is about reaching @@ -180,20 +200,24 @@ std::string check_feature_compatibility( args.remote_target_shard.enabled()) { return "--paged-attention requires one local target device"; } - if (args.draft_path != nullptr || args.remote_draft.enabled() || - args.ddtree_mode) { - return "--paged-attention requires autoregressive decode without a " - "draft or DDTree"; + if (args.remote_draft.enabled()) { + return "concurrent paged DDTree requires a local draft on the target device"; + } + if ((args.draft_path != nullptr || args.ddtree_mode) && + !concurrent_local_ddtree) { + return "paged draft decode is supported only as concurrent local DDTree " + "on one target/draft device"; } if (args.fa_window != 0) { return "--paged-attention requires full attention (--fa-window 0)"; } - if (features.pflash_enabled) { - return "--paged-attention cannot be combined with PFlash prefill " - "compression"; + if (features.pflash_enabled && !concurrent_local_paged_qwen) { + return "paged PFlash prefill compression requires concurrent local " + "Qwen3.5/Qwen3.6 serving on one target/draft device"; } - if (features.kvflash_enabled) { - return "--paged-attention cannot be combined with KVFlash"; + if (features.kvflash_enabled && !concurrent_paged_qwen) { + return "paged KVFlash requires concurrent local Qwen3.5/Qwen3.6 " + "serving with full attention"; } // The pool rounds max_ctx up to a whole number of blocks, so the top // of the range is what can be rounded without overflowing int. @@ -232,10 +256,16 @@ std::string check_feature_compatibility( if (args.max_concurrency <= 1) { return "--kv-pool-tokens requires --max-concurrency greater than 1"; } - // The cache appends one scratch block after the physical pool, and - // the requested pool itself is rounded up to a whole block. Cap the - // request at the largest aligned pool that leaves room for scratch. - const int64_t max_pool_tokens = paged_kv_address_cap(); + // Reserve the dead-row block plus one rounded DDTree candidate slab + // per slot, exactly matching Qwen35Backend's cache allocation. + const int64_t tree_scratch = concurrent_local_ddtree + ? (int64_t)args.max_concurrency * + paged_token_capacity(args.ddtree_budget + 1) + : 0; + const int64_t scratch_tokens = PAGED_BLOCK_SIZE + tree_scratch; + const int64_t max_pool_tokens = + ((int64_t)INT32_MAX - scratch_tokens) / PAGED_BLOCK_SIZE * + PAGED_BLOCK_SIZE; if (args.kv_pool_tokens < PAGED_BLOCK_SIZE || args.kv_pool_tokens > max_pool_tokens) { return "--kv-pool-tokens must be in [" + diff --git a/server/src/common/gpu_runtime_compat.h b/server/src/common/gpu_runtime_compat.h index dba8eaa7c..8a21fe025 100644 --- a/server/src/common/gpu_runtime_compat.h +++ b/server/src/common/gpu_runtime_compat.h @@ -60,6 +60,7 @@ #define cudaPointerAttributes hipPointerAttribute_t #define cudaPointerGetAttributes hipPointerGetAttributes #define cudaStreamCreate hipStreamCreate +#define cudaStreamCreateWithFlags hipStreamCreateWithFlags #define cudaStreamDefault hipStreamDefault #define cudaStreamDestroy hipStreamDestroy #define cudaStreamNonBlocking hipStreamNonBlocking diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index cdd4d0210..5e1814a72 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -35,7 +35,8 @@ struct StepGraph { ggml_tensor * inp_embed = nullptr; ggml_tensor * positions = nullptr; ggml_tensor * attn_mask = nullptr; // may be null - ggml_tensor * parent_ids = nullptr; // DDTree tree-mode; null for chain mode + ggml_tensor * parent_ids = nullptr; // DDTree [tree_width,n_tree_seqs] + ggml_tensor * tree_sizes = nullptr; // DDTree [n_tree_seqs], 0 = padding ggml_tensor * target_hidden_cat = nullptr; // draft only ggml_tensor * positions_k = nullptr; // draft only ggml_tensor * pad_mask_full = nullptr; // draft only; padded-ctx mask @@ -55,11 +56,16 @@ struct StepGraph { // state_slot_ids has the same shape but maps padding to a safe readable // slot for graph-level conv-state gathers. ggml_tensor * active_slot_ids = nullptr; + // Recurrent gather rows. Unlike active/paged IDs, padding must name a + // valid harmless slot (normally 0): ggml_get_rows does not mask -1. ggml_tensor * state_slot_ids = nullptr; // Ragged paged read (concurrent prefill): per-row block-table column and // inclusive causal position, [n_tokens] i32 each. Padding rows carry -1. ggml_tensor * paged_query_seq_ids = nullptr; ggml_tensor * paged_query_positions = nullptr; + // DFlash target-feature destination rows. Multi-slot replay maps each + // token to its slot-local ring; padding maps to the cache's dead row. + ggml_tensor * target_feat_rows = nullptr; // Multi-prompt steps: i32 row indices gathered from the final norm // before the LM head (committing rows + decode rows). ggml_tensor * logits_row_indices = nullptr; @@ -92,11 +98,13 @@ inline void step_graph_free(StepGraph & sg) { sg.built_view = false; sg.hidden_input = nullptr; sg.parent_ids = nullptr; + sg.tree_sizes = nullptr; sg.kv_write_rows = nullptr; sg.active_slot_ids = nullptr; sg.state_slot_ids = nullptr; sg.paged_query_seq_ids = nullptr; sg.paged_query_positions = nullptr; + sg.target_feat_rows = nullptr; sg.logits_row_indices = nullptr; sg.logits = nullptr; sg.hidden_states = nullptr; diff --git a/server/src/internal.h b/server/src/internal.h index fffda41ac..9cb2a03b6 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -414,13 +414,12 @@ struct TargetCache { std::vector conv_input_cache; // size = n_delta (48) // Rolling target layer features captured during target forward passes. - // Shape [5 * hidden, target_feat_cap] bf16. target_feat_cap is typically - // << max_ctx (e.g. 4096) so the buffer stays small at 128K context. The - // graph writes to slot `(kv_start + i) % target_feat_cap` so positions - // beyond the cap wrap and overwrite older entries. Readers (draft) only - // need the last DRAFT_CTX_MAX positions, so wrap is invisible in - // practice. Fed into the draft graph's fc projection after a bf16→f32 - // cast (ggml_get_to_fp32_cuda). + // Single-sequence shape: [5 * hidden, target_feat_cap] bf16. A multi-slot + // cache owns one ring per physical sequence slot and one final dead row: + // [5 * hidden, target_feat_cap * n_seq_slots + 1]. Live row P in slot S + // maps to S*target_feat_cap + P%target_feat_cap; bucket padding maps to the + // dead final row because ggml_set_rows does not accept a negative index. + // target_feat_cap remains the per-sequence ring width. ggml_tensor * target_feat = nullptr; int target_feat_cap = 0; @@ -559,6 +558,10 @@ bool restore_target_cache_chain(const PrefixSnapshot * thick, // decode is AR-only). With // n_seq_slots > 1 the attention K/V tensors are sized by ctx_alloc (the shared // pool capacity plus one scratch block) rather than one sequence's max_ctx. +// `concurrent_tree` declares that a paged multi-slot caller will build packed +// DDTree verification graphs. Those graphs are deliberately side-effect-free +// for recurrent state and commit accepted paths through a later replay, so no +// rollback snapshots/intermediates are allocated. bool create_target_cache(const TargetWeights & w, int max_ctx, int max_verify_tokens, @@ -567,7 +570,8 @@ bool create_target_cache(const TargetWeights & w, bool prefill_only = false, int ctx_alloc = 0, bool paged_attention = false, - int n_seq_slots = 1); + int n_seq_slots = 1, + bool concurrent_tree = false); // `f32_ssm_intermediates` enables exact per-token checkpoints for the opt-in // layer-split fast rollback path. The default preserves the established Q8_0 @@ -584,7 +588,8 @@ bool create_target_cache_partial(const TargetWeights & w, int ctx_alloc = 0, bool f32_ssm_intermediates = false, bool paged_attention = false, - int n_seq_slots = 1); + int n_seq_slots = 1, + bool concurrent_tree = false); void free_target_cache(TargetCache & c); @@ -657,7 +662,8 @@ struct QwenGraphInputs { bool capture_moe_router = false; // if true, expose selected expert ids for MoE layers int fa_window = 0; // sliding window for FA layers: 0 = full attention int logits_tail_rows = 0; // compute logits only for last n rows; 0 = all - ggml_tensor * parent_ids = nullptr; // [n_tokens] i32; tree mode when non-null + ggml_tensor * parent_ids = nullptr; // tree: [tree_width,n_tree_seqs] i32 + ggml_tensor * tree_sizes = nullptr; // tree: [n_tree_seqs] i32; 0 = padding tree // [n_tokens,n_head_kv] i64 physical destination rows for the // ggml_set_rows KV write; step-invariant. ggml_tensor * kv_write_rows = nullptr; @@ -683,6 +689,10 @@ struct QwenGraphInputs { // last row plus the decode rows), which a tail view cannot express. // Non-null overrides logits_tail_rows. ggml_tensor * logits_row_indices = nullptr; + // Optional replay-stable DFlash capture destinations. When present, all + // captured layers are concatenated once and written with ggml_set_rows. + // Multi-slot callers provide per-slot ring rows (padding uses dead row). + ggml_tensor * target_feat_rows = nullptr; // [n_tokens] i32 // Prefill segments on the leading token axis (see QwenPrefillSegment). // n_prefill_tokens is their total row count. seq_slot is ignored when // segments are present. @@ -717,6 +727,12 @@ struct QwenGraphInputs { int seq_slot = 0; int paged_max_kv_len = 0; int n_prefill_tokens = 0; + // Packed paged-tree metadata. Tokens are flattened sequence-major: + // row = sequence*tree_width + node. tree_scratch_* describe the physical + // KV scratch slab owned by each physical sequence slot. + int tree_width = 0; + int tree_scratch_base = 0; + int tree_scratch_stride = 0; // Capture the LAST token's post-RoPE/post-rotation Q per full-attention // layer into cache.q_cap (KVFlash target-QK scorer). Step-invariant: // node properties depend only on n_tokens and the layer index. diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 0a54761f9..545e7d9bf 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -12,6 +12,8 @@ #include "attn_masks.h" #include "prefill_helpers.h" #include "common/sampler.h" +#include "common/ddtree.h" +#include "common/geometric_draft_topk_cuda.h" #include "internal.h" #include @@ -35,6 +37,506 @@ int decode_bucket_width(int live_count) { } // namespace +Qwen35SeqEngine::Qwen35SeqEngine( + Qwen35Backend & backend, PagedKvPool & pool, int max_ctx, + int64_t scratch_row, int tree_width, int tree_scratch_base, + int tree_scratch_stride, int max_prefills, + int mixed_prefill_tokens, int long_mixed_prefill_tokens, + int long_prefill_threshold, int idle_prefill_tokens, + int prefill_quantum) + : max_prefills_(std::max(1, max_prefills)), + mixed_prefill_tokens_(std::max(1, mixed_prefill_tokens)), + long_mixed_prefill_tokens_(std::max(1, long_mixed_prefill_tokens)), + long_prefill_threshold_(std::max(1, long_prefill_threshold)), + idle_prefill_tokens_(std::max(1, idle_prefill_tokens)), + prefill_quantum_(std::max(1, prefill_quantum)), b_(backend), + slots_(pool, max_ctx, std::max(1, tree_width), + backend.paged_kv_residency_.get()), + scratch_row_(scratch_row), tree_width_(tree_width), + tree_scratch_base_(tree_scratch_base), + tree_scratch_stride_(tree_scratch_stride) { + const int n_slots = slots_.slot_count(); + slot_draft_kv_.resize((size_t)n_slots); + + // The concurrent DDTree stack is gated to a local same-device drafter. + // Build metadata-only BF16 views over each slot's disjoint target feature + // ring; draft_kv_begin_step converts only newly committed rows to its F32 + // append input instead of syncing the entire 200 MiB ring per round. + capture_features_ = tree_width_ > 0 && b_.cache_.target_feat && + b_.cache_.target_feat_cap > 0 && b_.cfg_.draft_path && + !b_.cfg_.remote_draft.enabled() && !b_.split_gpus_ && + b_.cfg_.draft_gpu == b_.cfg_.device.gpu; + if (!capture_features_) return; + + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * (size_t)(n_slots + 1); + ip.no_alloc = true; + feature_view_ctx_ = ggml_init(ip); + if (!feature_view_ctx_) { + capture_features_ = false; + return; + } + + const int cap = b_.cache_.target_feat_cap; + const int64_t fc_in = + (int64_t)b_.w_.n_capture_layers * b_.w_.n_embd; + slot_feature_mirrors_.resize((size_t)n_slots); + for (int slot = 0; slot < n_slots; ++slot) { + DraftFeatureMirror & mirror = slot_feature_mirrors_[(size_t)slot]; + mirror.target_feat = ggml_view_2d( + feature_view_ctx_, b_.cache_.target_feat, fc_in, cap, + b_.cache_.target_feat->nb[1], + (size_t)slot * (size_t)cap * b_.cache_.target_feat->nb[1]); + mirror.device = b_.cfg_.draft_gpu; + mirror.target_device = b_.cfg_.device.gpu; + mirror.cap = cap; + mirror.n_target_layers = b_.w_.n_capture_layers; + mirror.hidden_size = b_.w_.n_embd; + mirror.storage_type = b_.cache_.target_feat->type; + } +} + +Qwen35SeqEngine::~Qwen35SeqEngine() { + for (std::unique_ptr & state : slot_draft_kv_) { + if (state) draft_kv_free(*state); + } + slot_draft_kv_.clear(); + for (DraftFeatureMirror & mirror : slot_feature_mirrors_) { + draft_feature_mirror_free(mirror); + } + slot_feature_mirrors_.clear(); + if (feature_view_ctx_) { + ggml_free(feature_view_ctx_); + feature_view_ctx_ = nullptr; + } +} + +DraftFeatureMirror * Qwen35SeqEngine::slot_feature_mirror(int slot) { + if (!capture_features_ || slot < 0 || + slot >= (int)slot_feature_mirrors_.size()) { + return nullptr; + } + return &slot_feature_mirrors_[(size_t)slot]; +} + +DraftKvState * Qwen35SeqEngine::ensure_slot_draft_kv(int slot) { + DraftFeatureMirror * mirror = slot_feature_mirror(slot); + if (!mirror || slot < 0 || slot >= (int)slot_draft_kv_.size()) { + return nullptr; + } + std::unique_ptr & state = slot_draft_kv_[(size_t)slot]; + if (state && state->gf && state->built_for == (const void *)&b_.dw_) { + return state.get(); + } + if (state) draft_kv_free(*state); + state = std::make_unique(); + const int cap = std::min( + mirror->cap, std::max(1, b_.cfg_.draft_ctx_max)); + if (!draft_kv_init(*state, b_.dw_, b_.draft_backend_, cap, nullptr)) { + draft_kv_free(*state); + state.reset(); + return nullptr; + } + return state.get(); +} + +bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { + if (tree_width_ <= 1 || !capture_features_ || !plan.prefills.empty() || + plan.decode.empty() || b_.dw_.block_size <= 1 || + b_.cfg_.ddtree_budget + 1 != tree_width_) { + return false; + } + const int min_floor = []() { + const char * value = std::getenv("DFLASH_MIN_TOKENS"); + return value ? std::max(0, std::atoi(value)) : 0; + }(); + for (const StepInput & in : plan.decode) { + if (!in.allow_speculation || in.slot < 0 || + in.slot >= slots_.slot_count() || + !slots_.slot(in.slot).decoding() || + slots_.slot(in.slot).sampler.needs_logit_processing() || + slots_.slot(in.slot).cur_pos < 1 || + slots_.slot(in.slot).cur_pos >= slots_.max_context()) { + return false; + } + const Qwen35Slot & seq = slots_.slot(in.slot); + const int generated = seq.generated_tokens(); + if (generated < min_floor) return false; + } + return true; +} + +std::optional Qwen35SeqEngine::step_ddtree( + const StepPlan & plan) { + StepResult result; + const int active = (int)plan.decode.size(); + const int bucket = decode_bucket_width(active); + const int T = tree_width_; + const int q_len = b_.dw_.block_size; + const int hidden = b_.w_.n_embd; + const int n_head_kv = b_.w_.n_head_kv; + const int n_slots = slots_.slot_count(); + const int K = b_.cfg_.ddtree_budget > q_len - 1 ? 8 : 1; + + struct Proposal { + int slot = -1; + int32_t root = -1; + DDTree tree; + std::vector flat; + std::vector accepted; + int32_t bonus = -1; + }; + std::vector proposals; + proposals.reserve((size_t)active); + std::vector noise((size_t)q_len, b_.w_.mask_token_id); + std::vector noise_embed((size_t)hidden * q_len); + std::vector logits((size_t)b_.w_.n_vocab * q_len); + std::vector top_lp((size_t)q_len * K); + std::vector top_ids((size_t)q_len * K); + + auto proposal_fallback = [&]() -> std::optional { + // begin_step updates persistent drafter bookkeeping before compute. + // A failed proposal graph may therefore leave only part of that + // cache valid. Reset all participating draft rings so a later + // speculative round rebuilds them from committed target features. + for (const StepInput & in : plan.decode) { + if (in.slot >= 0 && in.slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)in.slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)in.slot]); + } + } + return std::nullopt; + }; + + if (!build_lm_head_projection_step( + b_.proj_sg_, b_.w_, b_.target_backend_, q_len)) { + return proposal_fallback(); + } + + // Proposal is sequential by slot: immutable draft weights are shared, + // while each slot owns an independent persistent context-KV ring. + for (const StepInput & in : plan.decode) { + DraftKvState * draft = ensure_slot_draft_kv(in.slot); + DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); + if (!draft || !mirror) return proposal_fallback(); + noise[0] = in.token; + std::fill(noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), q_len, noise_embed.data()) || + !draft_kv_begin_step(*draft, b_.dw_, b_.draft_backend_, + *mirror, slots_.slot(in.slot).cur_pos)) { + return proposal_fallback(); + } + ggml_backend_tensor_set( + draft->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + if (ggml_backend_graph_compute(b_.draft_backend_, draft->gf) != + GGML_STATUS_SUCCESS) { + return proposal_fallback(); + } + ggml_backend_tensor_copy( + draft->hidden_states, b_.proj_sg_.hidden_input); + if (ggml_backend_graph_compute( + b_.target_backend_, b_.proj_sg_.gf) != GGML_STATUS_SUCCESS) { + return proposal_fallback(); + } + bool topk_ready = false; +#ifdef DFLASH27B_HAVE_DRAFT_TOPK + topk_ready = geometric_extract_draft_topk_cuda( + b_.proj_sg_.logits->data, q_len, b_.w_.n_vocab, K, + top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); +#endif + if (!topk_ready) { + ggml_backend_tensor_get( + b_.proj_sg_.logits, logits.data(), 0, + sizeof(float) * logits.size()); + extract_draft_topk( + logits.data(), q_len, b_.w_.n_vocab, K, + top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); + } + + Proposal p; + p.slot = in.slot; + p.root = in.token; + p.tree = build_ddtree( + top_lp.data() + K, top_ids.data() + K, + q_len - 1, K, b_.cfg_.ddtree_budget, + b_.cfg_.ddtree_chain_seed); + p.flat.assign((size_t)T, 0); + p.flat[0] = in.token; + for (int node = 0; node < p.tree.n_nodes; ++node) { + p.flat[(size_t)node + 1] = p.tree.token_ids[(size_t)node]; + } + proposals.push_back(std::move(p)); + } + + StepGraph & tree_sg = b_.sg_; + int max_prefix = 1; + for (const Proposal & p : proposals) { + max_prefix = std::max(max_prefix, slots_.slot(p.slot).cur_pos); + } + if (!build_target_step_paged_tree( + tree_sg, b_.w_, b_.cache_, b_.target_backend_, T, bucket, + max_prefix, tree_scratch_base_, tree_scratch_stride_, + b_.cfg_.kq_stride_pad)) { + result.error = "packed DDTree verify graph build failed"; + return result; + } + + const int total_tree = T * bucket; + std::vector flat_tokens((size_t)total_tree, 0); + std::vector parents((size_t)total_tree, -1); + std::vector sizes((size_t)bucket, 0); + // Negative active IDs identify bucket padding. Recurrent gathers cannot + // index a negative slab, so padded trees use slot 0 only for their + // read-only base-state gather; tree_size=0/query_slot=-1 keeps all of + // their attention/output rows inactive and tree mode never persists it. + std::vector tree_slots((size_t)bucket, -1); + std::vector tree_state_slots((size_t)bucket, 0); + std::vector query_slots((size_t)total_tree, -1); + std::vector tree_rows( + (size_t)total_tree * n_head_kv, scratch_row_); + std::vector tree_pos((size_t)4 * total_tree, 0); + std::vector tree_embed((size_t)hidden * total_tree, 0.0f); + seq_lens_.assign((size_t)n_slots, 0); + + for (int s = 0; s < active; ++s) { + const Proposal & p = proposals[(size_t)s]; + const int base = s * T; + sizes[(size_t)s] = p.tree.n_nodes + 1; + tree_slots[(size_t)s] = p.slot; + tree_state_slots[(size_t)s] = p.slot; + seq_lens_[(size_t)p.slot] = slots_.slot(p.slot).cur_pos; + for (int node = 0; node < sizes[(size_t)s]; ++node) { + const int row = base + node; + flat_tokens[(size_t)row] = p.flat[(size_t)node]; + parents[(size_t)row] = node == 0 ? -1 : + p.tree.parents[(size_t)node]; + query_slots[(size_t)row] = p.slot; + const int depth = node == 0 ? 0 : + p.tree.depths[(size_t)node - 1]; + const int pos = slots_.slot(p.slot).cur_pos + depth; + tree_pos[(size_t)0 * total_tree + row] = pos; + tree_pos[(size_t)1 * total_tree + row] = pos; + tree_pos[(size_t)2 * total_tree + row] = pos; + for (int h = 0; h < n_head_kv; ++h) { + tree_rows[(size_t)h * total_tree + row] = + (int64_t)tree_scratch_base_ + + (int64_t)p.slot * tree_scratch_stride_ + node; + } + } + } + if (!b_.w_.embedder.embed( + flat_tokens.data(), total_tree, tree_embed.data())) { + result.error = "packed DDTree embedding failed"; + return result; + } + ggml_backend_tensor_set(tree_sg.inp_embed, tree_embed.data(), 0, + sizeof(float) * tree_embed.size()); + ggml_backend_tensor_set(tree_sg.positions, tree_pos.data(), 0, + sizeof(int32_t) * tree_pos.size()); + ggml_backend_tensor_set(tree_sg.parent_ids, parents.data(), 0, + sizeof(int32_t) * parents.size()); + ggml_backend_tensor_set(tree_sg.tree_sizes, sizes.data(), 0, + sizeof(int32_t) * sizes.size()); + ggml_backend_tensor_set(tree_sg.active_slot_ids, tree_slots.data(), 0, + sizeof(int32_t) * tree_slots.size()); + ggml_backend_tensor_set(tree_sg.state_slot_ids, tree_state_slots.data(), 0, + sizeof(int32_t) * tree_state_slots.size()); + ggml_backend_tensor_set(tree_sg.paged_query_seq_ids, query_slots.data(), 0, + sizeof(int32_t) * query_slots.size()); + ggml_backend_tensor_set(tree_sg.kv_write_rows, tree_rows.data(), 0, + sizeof(int64_t) * tree_rows.size()); + ggml_backend_tensor_set(b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + if (ggml_backend_graph_compute(b_.target_backend_, tree_sg.gf) != + GGML_STATUS_SUCCESS) { + result.error = "packed DDTree verify compute failed"; + return result; + } + std::vector posterior((size_t)total_tree, -1); + ggml_backend_tensor_get(tree_sg.argmax_tokens, posterior.data(), 0, + sizeof(int32_t) * posterior.size()); + + int replay_total = 0; + for (int s = 0; s < active; ++s) { + Proposal & p = proposals[(size_t)s]; + p.accepted = follow_verified_tree( + p.tree, posterior.data() + (size_t)s * T, p.bonus); + const int room = slots_.max_context() - slots_.slot(p.slot).cur_pos; + truncate_verified_path( + p.accepted, (size_t)std::max(0, room), + posterior.data() + (size_t)s * T, p.bonus); + if (p.accepted.empty()) { + result.error = "DDTree accepted path has no context headroom"; + return result; + } + replay_total += (int)p.accepted.size(); + } + + std::vector replay_segments; + std::vector replay_tokens; + std::vector replay_slots; + std::vector replay_positions; + std::vector replay_rows; + std::vector replay_logits_rows; + replay_segments.reserve((size_t)active); + replay_tokens.reserve((size_t)replay_total); + replay_slots.reserve((size_t)replay_total); + replay_positions.reserve((size_t)replay_total); + replay_rows.assign((size_t)replay_total * n_head_kv, scratch_row_); + replay_logits_rows.reserve((size_t)active); + seq_lens_.assign((size_t)n_slots, 0); + + int replay_offset = 0; + for (Proposal & p : proposals) { + std::vector path; + path.reserve(p.accepted.size()); + for (int dfs : p.accepted) { + path.push_back(dfs == 0 ? p.root : + p.tree.token_ids[(size_t)dfs - 1]); + } + const Qwen35SlotManager::StepAppend app = slots_.append_tokens( + p.slot, path.data(), (int)path.size()); + const bool table_ok = slots_.residency_active() || + upload_block_table_delta(p.slot, app.first_new_block, + app.new_blocks.data(), app.new_blocks.size()); + if (!app.ok || app.physical_rows.size() != path.size() || + !table_ok) { + result.error = app.busy + ? "paged KV pool exhausted during DDTree replay" + : "DDTree replay K/V append failed"; + return result; + } + replay_segments.push_back( + {replay_offset, (int)path.size(), p.slot}); + for (size_t i = 0; i < path.size(); ++i) { + replay_tokens.push_back(path[i]); + replay_slots.push_back(p.slot); + replay_positions.push_back(app.position + (int)i); + for (int h = 0; h < n_head_kv; ++h) { + replay_rows[(size_t)h * replay_total + replay_offset + i] = + app.physical_rows[i]; + } + } + replay_offset += (int)path.size(); + replay_logits_rows.push_back(replay_offset - 1); + seq_lens_[(size_t)p.slot] = app.position + (int)path.size(); + } + + if (!upload_all_active_block_tables()) { + result.error = "DDTree replay block-table refresh failed"; + return result; + } + + StepGraph & replay_sg = b_.sg_; + if (!build_target_step( + replay_sg, b_.w_, b_.cache_, b_.target_backend_, + 0, replay_total, false, true, false, 0, 0, + b_.cfg_.kq_stride_pad, false, false, false, true, + 1, 0, *std::max_element(seq_lens_.begin(), seq_lens_.end()), + replay_total, replay_segments.data(), + (int)replay_segments.size(), active, false) || + !replay_sg.target_feat_rows || !replay_sg.paged_query_seq_ids || + !replay_sg.paged_query_positions || !replay_sg.logits_row_indices || + !replay_sg.argmax_tokens) { + result.error = "DDTree accepted-path replay graph build failed"; + return result; + } + embed_buf_.resize((size_t)hidden * replay_total); + if (!b_.w_.embedder.embed( + replay_tokens.data(), replay_total, embed_buf_.data())) { + result.error = "DDTree replay embedding failed"; + return result; + } + pos_buf_.assign((size_t)4 * replay_total, 0); + feature_rows_.resize((size_t)replay_total); + const int cap = b_.cache_.target_feat_cap; + for (int row = 0; row < replay_total; ++row) { + const int pos = replay_positions[(size_t)row]; + pos_buf_[(size_t)0 * replay_total + row] = pos; + pos_buf_[(size_t)1 * replay_total + row] = pos; + pos_buf_[(size_t)2 * replay_total + row] = pos; + feature_rows_[(size_t)row] = + replay_slots[(size_t)row] * cap + pos % cap; + } + ggml_backend_tensor_set(replay_sg.inp_embed, embed_buf_.data(), 0, + sizeof(float) * embed_buf_.size()); + ggml_backend_tensor_set(replay_sg.positions, pos_buf_.data(), 0, + sizeof(int32_t) * pos_buf_.size()); + ggml_backend_tensor_set(replay_sg.kv_write_rows, replay_rows.data(), 0, + sizeof(int64_t) * replay_rows.size()); + ggml_backend_tensor_set(replay_sg.paged_query_seq_ids, + replay_slots.data(), 0, + sizeof(int32_t) * replay_slots.size()); + ggml_backend_tensor_set(replay_sg.paged_query_positions, + replay_positions.data(), 0, + sizeof(int32_t) * replay_positions.size()); + ggml_backend_tensor_set(replay_sg.logits_row_indices, + replay_logits_rows.data(), 0, + sizeof(int32_t) * replay_logits_rows.size()); + ggml_backend_tensor_set(replay_sg.target_feat_rows, + feature_rows_.data(), 0, + sizeof(int32_t) * feature_rows_.size()); + ggml_backend_tensor_set(b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + if (ggml_backend_graph_compute(b_.target_backend_, replay_sg.gf) != + GGML_STATUS_SUCCESS) { + result.error = "DDTree accepted-path replay compute failed"; + return result; + } + std::vector replay_write_slots; + replay_write_slots.reserve(proposals.size()); + for (const Proposal & p : proposals) replay_write_slots.push_back(p.slot); + if (!commit_residency_writes(replay_write_slots)) { + result.error = "DDTree replay KV write commit failed"; + return result; + } + + // The replay is the durable target forward: its recurrent/KV/feature + // state is what the next step consumes. Use its posterior rather than + // the tree-verify posterior so the pending scalar remains exact even if + // the two graph shapes differ numerically. + std::vector replay_next((size_t)active, -1); + ggml_backend_tensor_get( + replay_sg.argmax_tokens, replay_next.data(), 0, + sizeof(int32_t) * replay_next.size()); + for (int s = 0; s < active; ++s) { + if (replay_next[(size_t)s] < 0) { + result.error = "DDTree replay produced an invalid pending token"; + return result; + } + proposals[(size_t)s].bonus = replay_next[(size_t)s]; + } + + for (Proposal & p : proposals) { + slots_.commit_step(p.slot); + std::string reselect_error; + if (!maybe_reselect_residency(p.slot, reselect_error)) { + result.error = reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error; + return result; + } + } + result.decode.reserve((size_t)active); + for (Proposal & p : proposals) { + DecodeOutput out; + out.slot = p.slot; + out.token = p.bonus; + out.ddtree_steps = 1; + out.ddtree_accepted_tokens = p.accepted.size() - 1; + out.target_forwards = 2; + for (size_t i = 1; i < p.accepted.size(); ++i) { + const int dfs = p.accepted[i]; + out.committed_tokens.push_back( + p.tree.token_ids[(size_t)dfs - 1]); + } + attach_residency_telemetry(out); + result.decode.push_back(std::move(out)); + } + return result; +} + bool Qwen35SeqEngine::token_is_eos(int32_t token) const { return b_.token_is_eos(token); } @@ -46,6 +548,14 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( AdmitResult result = slots_.admit(request_id, prompt, sampler); if (result.status == AdmitResult::Status::admitted) { reset_recurrent_slot(b_.cache_, result.slot); + if (slots_.residency_active()) { + slots_.slot(result.slot).kvflash_last_reselect_generated = + -std::max(1, b_.kvflash_tau_); + } + if (result.slot >= 0 && result.slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)result.slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)result.slot]); + } } return result; } @@ -102,6 +612,77 @@ bool Qwen35SeqEngine::upload_block_table_delta( return true; } +bool Qwen35SeqEngine::upload_all_active_block_tables() { + if (!slots_.residency_active()) return true; + ggml_tensor * table = b_.cache_.paged_block_table; + if (!table) return false; + std::vector column((size_t)table->ne[0], -1); + std::vector snapshot; + for (int slot = 0; slot < slots_.slot_count(); ++slot) { + if (!slots_.is_active(slot)) continue; + std::fill(column.begin(), column.end(), -1); + if (!slots_.block_table_snapshot(slot, snapshot) || + snapshot.size() > column.size()) { + return false; + } + std::copy(snapshot.begin(), snapshot.end(), column.begin()); + ggml_backend_tensor_set( + table, column.data(), (size_t)slot * table->nb[1], + sizeof(int32_t) * column.size()); + } + return true; +} + +bool Qwen35SeqEngine::commit_residency_writes( + const std::vector & slots) { + if (!slots_.residency_active()) return true; + // Graph completion is not a host barrier for every backend. Pending pages + // become evictable only after all target writes are device-complete. + ggml_backend_synchronize(b_.target_backend_); + for (int slot : slots) { + if (!slots_.commit_residency_writes(slot)) return false; + } + return true; +} + +bool Qwen35SeqEngine::maybe_reselect_residency( + int slot, std::string & error) { + if (!slots_.residency_active()) return true; + Qwen35Slot & seq = slots_.slot(slot); + const int generated = seq.generated_tokens(); + const int tau = std::max( + b_.kvflash_tau_, (int)(seq.sample_history.size() / 45)); + if (generated - seq.kvflash_last_reselect_generated < tau) return true; + + b_.kvflash_ensure_scorer(); + std::vector scores; + const std::vector * score_ptr = nullptr; + if (b_.kvflash_scorer_) { + if (!b_.kvflash_scorer_->score_chunks( + seq.sample_history, PAGED_BLOCK_SIZE, scores)) { + // Short histories and recoverable drafter failures are expected + // scorer outcomes. Preserve service with the pager's explicit + // recency/LRU policy; only residency or transfer errors below are + // fatal to the request. + std::fprintf(stderr, + "[parallel-kvflash] scorer unavailable for slot %d; using LRU\n", + slot); + } else { + const size_t blocks = (seq.sample_history.size() + + PAGED_BLOCK_SIZE - 1) / PAGED_BLOCK_SIZE; + scores.resize(blocks, scores.empty() ? 0.0f : scores.back()); + score_ptr = &scores; + } + } + if (!slots_.reselect_residency(slot, score_ptr, &error)) return false; + seq.kvflash_last_reselect_generated = generated; + return upload_all_active_block_tables(); +} + +void Qwen35SeqEngine::attach_residency_telemetry(DecodeOutput & out) { + slots_.take_residency_telemetry(out.slot, out); +} + void Qwen35SeqEngine::fail_prefill( int slot, std::vector & prefill_outputs, const char * log_message, const char * client_message) { @@ -135,9 +716,10 @@ Qwen35SeqEngine::PrefillStage Qwen35SeqEngine::stage_prefill_chunk( "prefill K/V allocation failed"); return PrefillStage{}; } - if (!upload_block_table_delta( - slot, chunk.first_new_block, chunk.new_blocks.data(), - chunk.new_blocks.size())) { + const bool table_ok = slots_.residency_active() || upload_block_table_delta( + slot, chunk.first_new_block, chunk.new_blocks.data(), + chunk.new_blocks.size()); + if (!table_ok) { fail_prefill( slot, prefill_outputs, "prefill block-table delta exceeds device capacity", "prefill block-table update failed"); @@ -207,6 +789,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (inputs.empty() && plan.prefills.empty()) return result; + if (ddtree_eligible(plan)) { + std::optional speculative = step_ddtree(plan); + if (speculative) return std::move(*speculative); + // Proposal setup failed before target/cache mutation. Preserve service + // by taking the existing packed AR path for this iteration. + } + const TargetWeights & w = b_.w_; StepGraph & sg = b_.sg_; const int hidden = w.n_embd; @@ -242,9 +831,10 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { output_rows_.push_back(compact_row); continue; } - if (app.new_block >= 0 && - !upload_block_table_delta( - in.slot, app.new_block_index, &app.new_block, 1)) { + const bool table_ok = slots_.residency_active() || app.new_block < 0 || + upload_block_table_delta(in.slot, app.new_block_index, + &app.new_block, 1); + if (!table_ok) { out.error = "decode block-table entry exceeds device capacity"; decode_outputs.push_back(std::move(out)); output_rows_.push_back(compact_row); @@ -278,6 +868,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } prefills.push_back(std::move(prefill)); } + if (!upload_all_active_block_tables()) { + return fail_step("active KVFlash block-table refresh failed"); + } const int live_count = (int)live_tokens_.size(); const bool with_decode = live_count > 0; @@ -329,7 +922,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { built = build_target_step( sg, w, b_.cache_, b_.target_backend_, /*kv_start=*/0, /*n_tokens=*/n_total, - /*with_mask=*/false, /*capture=*/false, + /*with_mask=*/false, /*capture=*/capture_features_, /*capture_delta_intermediate=*/false, /*fa_window=*/0, /*logits_tail_rows=*/0, b_.cfg_.kq_stride_pad, @@ -347,7 +940,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { built = build_target_step( sg, w, b_.cache_, b_.target_backend_, /*kv_start=*/0, /*n_tokens=*/decode_bucket, - /*with_mask=*/false, /*capture=*/false, + /*with_mask=*/false, /*capture=*/capture_features_, /*capture_delta_intermediate=*/false, /*fa_window=*/0, /*logits_tail_rows=*/0, b_.cfg_.kq_stride_pad, @@ -365,6 +958,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { /*compact_slots=*/true); } if (!built || !sg.kv_write_rows || + (capture_features_ && !sg.target_feat_rows) || (with_prefill && (!sg.paged_query_seq_ids || !sg.paged_query_positions || !sg.logits_row_indices))) { @@ -428,6 +1022,30 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { b_.target_backend_, sg.kv_write_rows, rows_buf_.data(), 0, sizeof(int64_t) * rows_buf_.size()); + if (capture_features_) { + const int cap = b_.cache_.target_feat_cap; + const int dead_row = cap * n_slots; + feature_rows_.assign((size_t)n_total, dead_row); + int feature_offset = 0; + for (size_t i = 0; i < prefills.size(); ++i) { + const PrefillStage & prefill = prefills[i]; + const int slot = plan.prefills[i].slot; + for (int row = 0; row < prefill.chunk; ++row) { + feature_rows_[(size_t)(feature_offset + row)] = + slot * cap + (prefill.kv_pos + row) % cap; + } + feature_offset += prefill.chunk; + } + for (int row = 0; row < live_count; ++row) { + feature_rows_[(size_t)(n_prefill + row)] = + live_slot_ids_[(size_t)row] * cap + + live_positions_[(size_t)row] % cap; + } + ggml_backend_tensor_set_async( + b_.target_backend_, sg.target_feat_rows, feature_rows_.data(), 0, + sizeof(int32_t) * feature_rows_.size()); + } + if (with_prefill) { query_slot_ids_.assign((size_t)n_total, -1); query_positions_.assign((size_t)n_total, -1); @@ -507,6 +1125,16 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { "qwen35.argmax_readback", roctx_metadata); ggml_backend_synchronize(b_.target_backend_); } + std::vector write_slots; + write_slots.reserve(live_slot_ids_.size() + prefills.size()); + write_slots.insert(write_slots.end(), + live_slot_ids_.begin(), live_slot_ids_.end()); + for (size_t i = 0; i < prefills.size(); ++i) { + write_slots.push_back(plan.prefills[i].slot); + } + if (!commit_residency_writes(write_slots)) { + return fail_step("KVFlash pending write commit failed"); + } for (size_t oi = 0; oi < inputs.size(); ++oi) { DecodeOutput & out = decode_outputs[oi]; @@ -517,6 +1145,17 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { out.slot, row, &argmax_buf_[(size_t)row], &logits_buf_); } + for (DecodeOutput & out : decode_outputs) { + if (out.failed) continue; + std::string reselect_error; + if (!maybe_reselect_residency(out.slot, reselect_error)) { + return fail_step(reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error); + } + out.target_forwards = 1; + attach_residency_telemetry(out); + } + int commit_row = 0; for (size_t i = 0; i < prefills.size(); ++i) { const int slot = plan.prefills[i].slot; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index d9391784c..75275df39 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -22,10 +22,15 @@ #pragma once #include "common/concurrency/seq_engine.h" +#include "common/dflash_draft_kv.h" +#include "common/dflash_feature_ring.h" +#include "common/ddtree.h" #include "qwen35_slot_manager.h" #include #include +#include +#include #include namespace dflash::common { @@ -40,19 +45,14 @@ class Qwen35SeqEngine final : public SeqEngine { // `max_prefills` bounds scheduler-selected prompt slices per traversal. Qwen35SeqEngine(Qwen35Backend & backend, PagedKvPool & pool, int max_ctx, int64_t scratch_row, - int max_prefills = 8, + int tree_width = 0, int tree_scratch_base = 0, + int tree_scratch_stride = 0, int max_prefills = 8, int mixed_prefill_tokens = 2048, int long_mixed_prefill_tokens = 4096, int long_prefill_threshold = 768, int idle_prefill_tokens = 4096, - int prefill_quantum = 512) - : max_prefills_(std::max(1, max_prefills)), - mixed_prefill_tokens_(std::max(1, mixed_prefill_tokens)), - long_mixed_prefill_tokens_(std::max(1, long_mixed_prefill_tokens)), - long_prefill_threshold_(std::max(1, long_prefill_threshold)), - idle_prefill_tokens_(std::max(1, idle_prefill_tokens)), - prefill_quantum_(std::max(1, prefill_quantum)), b_(backend), - slots_(pool, max_ctx), scratch_row_(scratch_row) {} + int prefill_quantum = 512); + ~Qwen35SeqEngine() override; int slot_count() const override { return slots_.slot_count(); } int max_context() const override { return slots_.max_context(); } @@ -104,6 +104,10 @@ class Qwen35SeqEngine final : public SeqEngine { bool upload_block_table_delta(int slot, int first_block, const int32_t * blocks, size_t count); + bool upload_all_active_block_tables(); + bool commit_residency_writes(const std::vector & slots); + bool maybe_reselect_residency(int slot, std::string & error); + void attach_residency_telemetry(DecodeOutput & out); void fail_prefill(int slot, std::vector & outputs, const char * log_message, const char * client_message); @@ -112,10 +116,23 @@ class Qwen35SeqEngine final : public SeqEngine { int32_t sample_graph_row(int slot, int logits_row, const int32_t * cached_argmax = nullptr, std::vector * logits_scratch = nullptr); + DraftFeatureMirror * slot_feature_mirror(int slot); + DraftKvState * ensure_slot_draft_kv(int slot); + bool ddtree_eligible(const StepPlan & plan) const; + // nullopt means proposal setup failed before target/cache mutation and the + // caller may safely use the ordinary packed AR path for this iteration. + std::optional step_ddtree(const StepPlan & plan); Qwen35Backend & b_; Qwen35SlotManager slots_; int64_t scratch_row_ = 0; + int tree_width_ = 0; + int tree_scratch_base_ = 0; + int tree_scratch_stride_ = 0; + bool capture_features_ = false; + ggml_context * feature_view_ctx_ = nullptr; + std::vector slot_feature_mirrors_; + std::vector> slot_draft_kv_; // Hoisted per-step buffers (reused across step() calls). std::vector output_rows_; @@ -131,6 +148,7 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector query_slot_ids_; std::vector query_positions_; std::vector logits_rows_; + std::vector feature_rows_; std::vector embed_buf_; std::vector pos_buf_; std::vector rows_buf_; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index b8dcc36ec..f740a283b 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -5,8 +5,12 @@ namespace dflash::common { -Qwen35SlotManager::Qwen35SlotManager(PagedKvPool & pool, int max_ctx) - : pool_(pool), max_ctx_(max_ctx) { +Qwen35SlotManager::Qwen35SlotManager( + PagedKvPool & pool, int max_ctx, int speculative_headroom, + PagedKvResidencyManager * residency) + : pool_(pool), max_ctx_(max_ctx), + headroom_tokens_(std::max(pool.block_size(), speculative_headroom)), + residency_(residency) { slots_.assign(pool.max_sequences(), Qwen35Slot{}); } @@ -21,7 +25,7 @@ int Qwen35SlotManager::decoding_count() const { uint32_t Qwen35SlotManager::decode_headroom_capacity(int logical_tokens) const { const uint64_t extended = static_cast(std::max(0, logical_tokens)) + - pool_.block_size(); + static_cast(headroom_tokens_); return static_cast(std::min( static_cast(max_ctx_), extended)); } @@ -83,6 +87,20 @@ bool Qwen35SlotManager::is_active(int slot) const { bool Qwen35SlotManager::is_prefilling(int slot) const { return is_active(slot) && slots_[(size_t)slot].prefilling(); } +void Qwen35SlotManager::accumulate_residency_delta( + Qwen35Slot & slot, const PagedKvResidencyStats & before) { + if (!residency_) return; + const PagedKvResidencyStats after = residency_->stats(); + if (after.page_ins >= before.page_ins) { + slot.kvflash_page_ins_pending += after.page_ins - before.page_ins; + } + if (after.page_outs >= before.page_outs) { + slot.kvflash_page_outs_pending += after.page_outs - before.page_outs; + } + if (after.reselects >= before.reselects) { + slot.kvflash_reselects_pending += after.reselects - before.reselects; + } +} bool Qwen35SlotManager::has_prefill_prompt_at_least(int tokens) const { if (tokens <= 0) return true; @@ -112,7 +130,7 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( // fail anyway. Hard-fail it up front instead of reporting busy. const uint64_t pool_capacity = (uint64_t)pool_.physical_block_count() * pool_.block_size(); - if ((uint64_t)prompt_len > pool_capacity) { + if (!residency_ && (uint64_t)prompt_len > pool_capacity) { r.error = "prompt needs " + std::to_string(prompt_len) + " KV tokens but the pool holds " + std::to_string(pool_capacity) + @@ -132,7 +150,8 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( // A newly freed block belongs to any older decoder missing its rolling // next-page reserve before it can belong to this admission. - const PagedKvStatus headroom_status = protect_decode_headroom(); + const PagedKvStatus headroom_status = residency_ + ? PagedKvStatus::Ok : protect_decode_headroom(); if (headroom_status != PagedKvStatus::Ok) { r.status = headroom_status == PagedKvStatus::BlocksExhausted ? AdmitStatus::busy : AdmitStatus::failed; @@ -143,7 +162,7 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( } PagedKvSequenceHandle handle; - uint32_t reservation_capacity = + uint32_t reservation_capacity = residency_ ? 0 : decode_headroom_capacity(prompt_len); if (!capacity_fits_pool(reservation_capacity)) { // The prompt itself fits, but this physical pool can never hold its @@ -151,8 +170,9 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( // decode exhaustion later if the sequence reaches that boundary. reservation_capacity = static_cast(prompt_len); } - const PagedKvStatus status = pool_.acquire_reserved( - request_id, reservation_capacity, handle); + const PagedKvStatus status = residency_ + ? pool_.acquire(request_id, handle) + : pool_.acquire_reserved(request_id, reservation_capacity, handle); if (status != PagedKvStatus::Ok) { r.status = status == PagedKvStatus::SequenceSlotsExhausted || status == PagedKvStatus::BlocksExhausted @@ -163,6 +183,17 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } + if (residency_) { + const PagedKvResidencyStatus resident_status = + residency_->register_sequence(handle); + if (resident_status != PagedKvResidencyStatus::Ok) { + (void)pool_.release(handle); + r.error = std::string("KV residency registration failed: ") + + paged_kv_residency_status_string(resident_status); + return r; + } + } + Qwen35Slot & s = slots_[(size_t)slot]; s.phase = Qwen35SlotPhase::prefill; s.handle = handle; @@ -194,7 +225,22 @@ Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( return out; } - PagedKvAppendResult app = pool_.append(s.handle, (uint32_t)n_tokens); + PagedKvAppendResult app; + if (residency_) { + const PagedKvResidencyStats before = residency_->stats(); + PagedKvResidentAppendResult resident = residency_->append( + s.handle, (uint32_t)n_tokens); + accumulate_residency_delta(s, before); + if (!resident) { + std::fprintf(stderr, + "[parallel-kvflash] prefill append failed for slot %d: %s\n", + slot, paged_kv_residency_status_string(resident.status)); + return out; + } + app = std::move(resident.pool_result); + } else { + app = pool_.append(s.handle, (uint32_t)n_tokens); + } if (!app) { // Admission reserved the whole prompt. Treat exhaustion here as a // broken invariant, not a retryable condition: retrying a batch of @@ -218,6 +264,17 @@ Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( out.new_blocks.push_back((int32_t)write.physical_block); } } + if (residency_) { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(s.handle, snapshot) != PagedKvStatus::Ok) { + return out; + } + out.full_block_table.reserve(snapshot.block_table.size()); + for (uint32_t block : snapshot.block_table) { + out.full_block_table.push_back( + block == PAGED_KV_COLD_BLOCK ? -1 : (int32_t)block); + } + } s.cur_pos += n_tokens; out.ok = true; return out; @@ -230,43 +287,189 @@ void Qwen35SlotManager::commit_prefill(int slot) { s.phase = Qwen35SlotPhase::decode; } -Qwen35SlotManager::StepAppend Qwen35SlotManager::append_token(int slot, - int32_t fed_token) { +Qwen35SlotManager::StepAppend Qwen35SlotManager::append_tokens( + int slot, const int32_t * fed_tokens, int n_tokens) { StepAppend out; - if (!is_active(slot) || !slots_[(size_t)slot].decoding()) return out; + if (!is_active(slot) || !slots_[(size_t)slot].decoding() || + !fed_tokens || n_tokens < 1) { + return out; + } Qwen35Slot & s = slots_[(size_t)slot]; - if (s.cur_pos >= max_ctx_) { - // No context left; the scheduler should have stopped this slot. + if (!s.staged_tokens.empty() || s.cur_pos > max_ctx_ || + n_tokens > max_ctx_ - s.cur_pos) { return out; } - PagedKvAppendResult app = pool_.append( - s.handle, 1, /*only_first_last_slots=*/true); - if (!app || app.token_count != 1 || - app.last.logical_position != (uint32_t)s.cur_pos) { + + PagedKvAppendResult app; + if (residency_) { + const PagedKvResidencyStats before = residency_->stats(); + PagedKvResidentAppendResult resident = residency_->append( + s.handle, static_cast(n_tokens)); + accumulate_residency_delta(s, before); + if (!resident) { + out.busy = resident.status == + PagedKvResidencyStatus::PoolExhausted || + resident.status == PagedKvResidencyStatus::NoEvictableBlock; + return out; + } + app = std::move(resident.pool_result); + } else { + app = pool_.append(s.handle, static_cast(n_tokens)); + } + if (!app || app.token_count != static_cast(n_tokens)) { out.busy = app.status == PagedKvStatus::BlocksExhausted; return out; } - s.sample_history.push_back(fed_token); + if (app.write_slots.size() != static_cast(n_tokens) || + app.write_slots.front().logical_position != + static_cast(s.cur_pos) || + app.write_slots.back().logical_position != + static_cast(s.cur_pos + n_tokens - 1)) { + // Pool success guarantees this shape. Retain staged ownership so a + // fatal caller retires the sequence rather than double-appending. + s.staged_tokens.assign(fed_tokens, fed_tokens + n_tokens); + return out; + } + + out.physical_rows.reserve(app.write_slots.size()); + for (const PagedKvWriteSlot & write : app.write_slots) { + out.physical_rows.push_back( + static_cast(write.physical_token_index)); + if (write.block_offset == 0) { + if (out.first_new_block < 0) { + out.first_new_block = static_cast( + write.logical_position / pool_.block_size()); + } + out.new_blocks.push_back( + static_cast(write.physical_block)); + } + } + s.staged_tokens.assign(fed_tokens, fed_tokens + n_tokens); + if (residency_) { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(s.handle, snapshot) != PagedKvStatus::Ok) { + return out; + } + out.full_block_table.reserve(snapshot.block_table.size()); + for (uint32_t block : snapshot.block_table) { + out.full_block_table.push_back( + block == PAGED_KV_COLD_BLOCK ? -1 : (int32_t)block); + } + } out.ok = true; - out.physical_row = (int64_t)app.last.physical_token_index; + out.count = n_tokens; out.position = s.cur_pos; - if ((uint32_t)s.cur_pos % pool_.block_size() == 0) { - out.new_block = (int32_t)app.last.physical_block; - out.new_block_index = s.cur_pos / (int)pool_.block_size(); + if (n_tokens == 1) { + out.physical_row = out.physical_rows.front(); + if (!out.new_blocks.empty()) { + out.new_block = out.new_blocks.front(); + out.new_block_index = out.first_new_block; + } } return out; } +Qwen35SlotManager::StepAppend Qwen35SlotManager::append_token( + int slot, int32_t fed_token) { + return append_tokens(slot, &fed_token, 1); +} + void Qwen35SlotManager::commit_step(int slot) { if (!is_active(slot)) return; - slots_[(size_t)slot].cur_pos += 1; + Qwen35Slot & s = slots_[(size_t)slot]; + if (s.staged_tokens.empty()) return; + s.sample_history.insert( + s.sample_history.end(), s.staged_tokens.begin(), + s.staged_tokens.end()); + s.cur_pos += static_cast(s.staged_tokens.size()); + s.staged_tokens.clear(); +} + +bool Qwen35SlotManager::block_table_snapshot( + int slot, std::vector & out) const { + out.clear(); + if (!is_active(slot)) return false; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(slots_[(size_t)slot].handle, snapshot) != + PagedKvStatus::Ok) { + return false; + } + out.reserve(snapshot.block_table.size()); + for (uint32_t block : snapshot.block_table) { + out.push_back(block == PAGED_KV_COLD_BLOCK ? -1 : (int32_t)block); + } + return true; +} + +bool Qwen35SlotManager::commit_residency_writes(int slot) { + if (!residency_) return true; + if (!is_active(slot)) return false; + return residency_->commit_pending_writes(slots_[(size_t)slot].handle) == + PagedKvResidencyStatus::Ok; +} + +bool Qwen35SlotManager::reselect_residency( + int slot, const std::vector * scores, std::string * error) { + if (!residency_) return true; + if (!is_active(slot)) { + if (error) *error = "inactive KVFlash slot"; + return false; + } + Qwen35Slot & s = slots_[(size_t)slot]; + const PagedKvResidencyStats before = residency_->stats(); + const std::vector no_scores; + PagedKvResidencyStatus status = residency_->set_scores( + s.handle, scores ? *scores : no_scores); + if (status == PagedKvResidencyStatus::Ok) { + status = residency_->reselect(s.handle); + } + accumulate_residency_delta(s, before); + if (status != PagedKvResidencyStatus::Ok) { + if (error) { + *error = std::string("KVFlash reselect failed: ") + + paged_kv_residency_status_string(status); + } + return false; + } + return true; +} + +void Qwen35SlotManager::take_residency_telemetry( + int slot, SeqEngine::DecodeOutput & out) { + if (!residency_ || !is_active(slot)) return; + Qwen35Slot & s = slots_[(size_t)slot]; + out.kvflash_page_ins = s.kvflash_page_ins_pending; + out.kvflash_page_outs = s.kvflash_page_outs_pending; + out.kvflash_reselects = s.kvflash_reselects_pending; + uint32_t resident = 0; + if (pool_.resident_block_count(s.handle, resident) == PagedKvStatus::Ok) { + out.kvflash_resident_blocks = resident; + } + s.kvflash_page_ins_pending = 0; + s.kvflash_page_outs_pending = 0; + s.kvflash_reselects_pending = 0; } void Qwen35SlotManager::retire(int slot) { if (slot < 0 || slot >= (int)slots_.size()) return; Qwen35Slot & s = slots_[(size_t)slot]; if (!s.active()) return; + if (residency_) { + const PagedKvResidencyStatus resident_status = + residency_->forget_sequence(s.handle); + if (resident_status != PagedKvResidencyStatus::Ok && + resident_status != PagedKvResidencyStatus::StaleHandle && + resident_status != PagedKvResidencyStatus::SequenceNotRegistered) { + std::fprintf(stderr, + "[parallel-kvflash] slot %d residency release failed: %s\n", + slot, paged_kv_residency_status_string(resident_status)); + // A failed copy-stream barrier leaves physical pages in flight. + // Keep the slot and its pool handle intact so a later retirement + // can retry forget_sequence without recycling those pages. + return; + } + } const PagedKvStatus status = pool_.release(s.handle); if (status != PagedKvStatus::Ok && status != PagedKvStatus::StaleHandle) { std::fprintf(stderr, "[parallel] slot %d release failed: %s\n", diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index 1da009f69..dd077496c 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -17,6 +17,7 @@ #pragma once #include "common/concurrency/paged_kv_pool.h" +#include "common/concurrency/paged_kv_residency.h" #include "common/sampler.h" #include "common/concurrency/seq_engine.h" @@ -45,6 +46,17 @@ struct Qwen35Slot { // Penalty history is recorded as fed rather than sampled: the scheduler // may override a sample before the model consumes it. std::vector sample_history; + // Decode rows allocated from the pool but not yet made durable by a + // successful target forward. A step stages one contiguous token range + // per slot, then commit_step() publishes all of it atomically. + std::vector staged_tokens; + + // Residency operation deltas are attributed to the requesting slot and + // held across prefill until a DecodeOutput can carry them upstream. + uint64_t kvflash_page_ins_pending = 0; + uint64_t kvflash_page_outs_pending = 0; + uint64_t kvflash_reselects_pending = 0; + int kvflash_last_reselect_generated = 0; int generated_tokens() const { return sample_history.size() > (size_t)prompt_len @@ -61,7 +73,13 @@ class Qwen35SlotManager { public: // `max_ctx` is the per-sequence logical bound; slot count comes from the // pool's max_sequences. The pool must outlive the manager. - Qwen35SlotManager(PagedKvPool & pool, int max_ctx); + Qwen35SlotManager(PagedKvPool & pool, int max_ctx, + int speculative_headroom = 1, + PagedKvResidencyManager * residency = nullptr); + Qwen35SlotManager(const Qwen35SlotManager &) = delete; + Qwen35SlotManager & operator=(const Qwen35SlotManager &) = delete; + Qwen35SlotManager(Qwen35SlotManager &&) = delete; + Qwen35SlotManager & operator=(Qwen35SlotManager &&) = delete; // Claim a free slot and atomically reserve all K/V blocks needed by the // known prompt plus its next logical decode page when that page can exist @@ -80,6 +98,7 @@ class Qwen35SlotManager { // Delta to patch into the slot's device block-table column. std::vector new_blocks; int first_new_block = -1; + std::vector full_block_table; }; // Append `n_tokens` more prompt rows for a prefilling slot. Physical block @@ -93,17 +112,37 @@ class Qwen35SlotManager { struct StepAppend { bool ok = false; bool busy = false; // no physical block available right now + std::vector physical_rows; + std::vector new_blocks; + int first_new_block = -1; + int count = 0; + // Compatibility fields for the common one-token append. int64_t physical_row = -1; - int position = -1; // logical position the fed token is written at + int position = -1; // logical position of the first staged token int32_t new_block = -1; int new_block_index = -1; + std::vector full_block_table; }; - // Allocate the next decode token's cache row, report any new block-table - // entry, and log it to sample_history. cur_pos waits for commit_step(). + // Atomically allocate and stage a contiguous accepted path. The pool + // append is all-or-nothing; sample_history and cur_pos remain unchanged + // until commit_step(). A slot may have only one staged range at a time. + StepAppend append_tokens(int slot, const int32_t * fed_tokens, + int n_tokens); + + // Complete residency snapshots encode cold logical pages as -1. The + // engine pads each device column before upload. + bool residency_active() const { return residency_ != nullptr; } + bool block_table_snapshot(int slot, std::vector & out) const; + bool commit_residency_writes(int slot); + bool reselect_residency(int slot, const std::vector * scores, + std::string * error = nullptr); + void take_residency_telemetry(int slot, SeqEngine::DecodeOutput & out); + + // One-token compatibility wrapper used by ordinary autoregressive decode. StepAppend append_token(int slot, int32_t fed_token); - // The batched step's compute succeeded: cur_pos++. + // Publish every staged token after a successful target forward. void commit_step(int slot); // Release the slot's blocks and clear its state. Safe on inactive slots @@ -111,6 +150,8 @@ class Qwen35SlotManager { void retire(int slot); int slot_count() const { return (int)slots_.size(); } + void accumulate_residency_delta(Qwen35Slot & slot, + const PagedKvResidencyStats & before); int max_context() const { return max_ctx_; } int decoding_count() const; bool is_active(int slot) const; @@ -131,6 +172,8 @@ class Qwen35SlotManager { PagedKvPool & pool_; int max_ctx_ = 0; + int headroom_tokens_; + PagedKvResidencyManager * residency_ = nullptr; std::vector slots_; }; diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index ad3d56b58..0ca65ecdb 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -9,6 +9,44 @@ namespace dflash::common { +bool detail::validate_target_paged_tree_layout( + const TargetCache & cache, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride) { + static constexpr int tree_buckets[] = { + 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, + }; + if (tree_width < 1 || + std::find(std::begin(tree_buckets), std::end(tree_buckets), + n_tree_seqs) == std::end(tree_buckets) || + cache.n_seq_slots <= 1 || !cache.paged_block_table || + !cache.paged_kv_seq_lens || paged_max_kv_len < 1 || + tree_scratch_base <= 0 || + tree_scratch_base % PAGED_BLOCK_SIZE != 0 || + tree_scratch_stride < tree_width) { + return false; + } + + int physical_kv_rows = 0; + for (ggml_tensor * tensor : cache.attn_k) { + if (tensor) { + physical_kv_rows = (int)tensor->ne[1]; + break; + } + } + if (physical_kv_rows < 1) return false; + + const int64_t scratch_end = + (int64_t)tree_scratch_base + + (int64_t)(cache.n_seq_slots - 1) * tree_scratch_stride + + tree_width; + return scratch_end <= physical_kv_rows && + (int64_t)tree_width * n_tree_seqs <= INT32_MAX; +} + // ── build_layer_step ──────────────────────────────────────────── bool build_layer_step( @@ -490,6 +528,12 @@ bool build_target_step( ggml_set_name(sg.logits_row_indices, "logits_row_indices"); ggml_set_input(sg.logits_row_indices); } + if (capture && paged_attention && cache.target_feat) { + sg.target_feat_rows = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + ggml_set_name(sg.target_feat_rows, "target_feat_rows"); + ggml_set_input(sg.target_feat_rows); + } sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); @@ -539,6 +583,7 @@ bool build_target_step( gi.paged_query_seq_ids = sg.paged_query_seq_ids; gi.paged_query_positions = sg.paged_query_positions; gi.logits_row_indices = sg.logits_row_indices; + gi.target_feat_rows = sg.target_feat_rows; gi.prefill_segments = prefill_segments; gi.n_prefill_segments = n_prefill_segments; @@ -630,6 +675,117 @@ bool build_target_step_tree( return ggml_gallocr_alloc_graph(sg.alloc, sg.gf); } +// ── build_target_step_paged_tree ──────────────────────────────── + +bool build_target_step_paged_tree( + StepGraph & sg, + const TargetWeights & w, + TargetCache & cache, + ggml_backend_t backend, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride, + int kq_stride_pad) { + (void)kq_stride_pad; + step_graph_free(sg); + + if (!detail::validate_target_paged_tree_layout( + cache, tree_width, n_tree_seqs, paged_max_kv_len, + tree_scratch_base, tree_scratch_stride)) { + return false; + } + const int n_tokens = tree_width * n_tree_seqs; + + ggml_init_params ip{}; + ip.mem_size = 512 * 1024 * 1024; + static thread_local std::vector g_tree_arena; + if (g_tree_arena.size() < ip.mem_size) g_tree_arena.resize(ip.mem_size); + ip.mem_buffer = g_tree_arena.data(); + ip.no_alloc = true; + sg.ctx = ggml_init(ip); + if (!sg.ctx) return false; + + // Salt graph addresses by the stable bucket shape so captured graphs for + // different T/S buckets never alias in ggml-cuda's topology cache. + for (int i = 0; i < tree_width + n_tree_seqs; ++i) { + (void)ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 1); + } + + sg.inp_embed = ggml_new_tensor_3d( + sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens, 1); + sg.positions = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 4 * n_tokens); + sg.parent_ids = ggml_new_tensor_2d( + sg.ctx, GGML_TYPE_I32, tree_width, n_tree_seqs); + sg.tree_sizes = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); + sg.active_slot_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); + sg.state_slot_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); + sg.paged_query_seq_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + sg.kv_write_rows = ggml_new_tensor_2d( + sg.ctx, GGML_TYPE_I64, n_tokens, w.n_head_kv); + + const struct NamedInput { + ggml_tensor * tensor; + const char * name; + } inputs[] = { + {sg.inp_embed, "inp_embed"}, + {sg.positions, "positions"}, + {sg.parent_ids, "parent_ids"}, + {sg.tree_sizes, "tree_sizes"}, + {sg.active_slot_ids, "active_slot_ids"}, + {sg.state_slot_ids, "state_slot_ids"}, + {sg.paged_query_seq_ids, "paged_query_seq_ids"}, + {sg.kv_write_rows, "kv_write_rows"}, + }; + for (const NamedInput & input : inputs) { + ggml_set_name(input.tensor, input.name); + ggml_set_input(input.tensor); + } + + sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); + QwenGraphInputs gi{}; + gi.inp_embed = sg.inp_embed; + gi.positions = sg.positions; + gi.n_tokens = n_tokens; + gi.kv_start = 0; + gi.capture_layers = false; + gi.capture_delta_intermediate = false; + gi.parent_ids = sg.parent_ids; + gi.tree_sizes = sg.tree_sizes; + gi.kv_write_rows = sg.kv_write_rows; + gi.paged_block_table = cache.paged_block_table; + gi.paged_kv_seq_lens = cache.paged_kv_seq_lens; + gi.active_slot_ids = sg.active_slot_ids; + gi.state_slot_ids = sg.state_slot_ids; + gi.paged_query_seq_ids = sg.paged_query_seq_ids; + gi.n_seqs = n_tree_seqs; + gi.paged_max_kv_len = paged_max_kv_len; + gi.tree_width = tree_width; + gi.tree_scratch_base = tree_scratch_base; + gi.tree_scratch_stride = tree_scratch_stride; + + QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); + if (!go.logits) return false; + sg.logits = go.logits; + ggml_set_output(sg.logits); + sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); + ggml_set_name(sg.argmax_tokens, "paged_tree_verify_argmax"); + ggml_set_output(sg.argmax_tokens); + ggml_build_forward_expand(sg.gf, sg.argmax_tokens); + + if (!sg.alloc) { + sg.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + } + return ggml_gallocr_alloc_graph(sg.alloc, sg.gf); +} + // ── build_lm_head_projection_step ─────────────────────────────── diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index cdbaf75ed..fac406863 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -23,6 +23,22 @@ namespace dflash::common { +namespace detail { + +// Model-free validation shared by the packed-tree builder and its shape +// tests. paged_max_kv_len is a logical launch bound and may exceed the +// bounded physical K/V pool; only the per-slot scratch slabs must fit in the +// physical tensor rows. +bool validate_target_paged_tree_layout( + const TargetCache & cache, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride); + +} // namespace detail + // Layer-segmented prefill: process one target layer for chunk_start..chunk_start+n_tokens. bool build_layer_step( StepGraph & sg, @@ -109,6 +125,10 @@ bool build_hybrid_full_layer_step( // overrides logits_tail_rows. Multi-prompt steps need it because // committing rows are scattered. 0 keeps the tail-view behavior. // `logits_tail_rows` — logits/argmax only for the last n rows (0 = all). +// When `capture && paged_attention`, sg.target_feat_rows is an I32 graph +// input mapping every token to its slot-local feature-ring destination. This +// keeps accepted-path replay graph-stable and leaves legacy offset capture +// unchanged for callers that do not use paged serving. bool build_target_step( StepGraph & sg, const TargetWeights & w, @@ -146,6 +166,26 @@ bool build_target_step_tree( int fa_window = 0, int kq_stride_pad = KQ_MASK_PAD); +// Packed concurrent DDTree verify over a paged multi-slot cache. Tokens are +// flattened sequence-major as [tree_width*n_tree_seqs]. n_tree_seqs is a +// stable graph-bucket width; inactive trees use tree_size=0 and dead/safe row +// mappings. In particular state_slot_ids padding must map to a valid harmless +// slot (normally 0), while active/paged sequence IDs may use -1. The graph +// writes candidate K/V into per-slot scratch slabs but +// does not mutate persistent recurrent state or target features. Accepted +// paths are committed by a later row-indexed replay through build_target_step. +bool build_target_step_paged_tree( + StepGraph & sg, + const TargetWeights & w, + TargetCache & cache, + ggml_backend_t backend, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride, + int kq_stride_pad = KQ_MASK_PAD); + // LM-head projection: project draft hidden states through the target output matrix. bool build_lm_head_projection_step( StepGraph & sg, diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 490f86d3b..c640d778f 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -173,7 +173,7 @@ static FILE * open_dflash_floor_log() { // staging K/V or staging recurrent slab to reserve. static int64_t concurrent_fixed_cache_bytes( const TargetWeights & w, int max_ctx, int n_slots, - int64_t kv_bytes_per_token) { + int64_t kv_bytes_per_token, int64_t scratch_tokens) { const int64_t n_full_attn = w.n_layer / w.full_attention_interval; const int64_t n_delta = w.n_layer - n_full_attn; @@ -188,7 +188,8 @@ static int64_t concurrent_fixed_cache_bytes( state_per_layer * n_delta * (int64_t)n_slots; const int64_t target_feat = (int64_t)w.n_capture_layers * w.n_embd * - std::min(max_ctx, 4096) * (int64_t)sizeof(uint16_t); + ((int64_t)std::min(max_ctx, 4096) * n_slots + 1) * + (int64_t)sizeof(uint16_t); const int64_t q_capture = (int64_t)w.n_embd_head_k * w.n_head * n_full_attn * (int64_t)sizeof(float); @@ -196,7 +197,7 @@ static int64_t concurrent_fixed_cache_bytes( ((int64_t)paged_block_count(max_ctx) * n_slots + n_slots) * (int64_t)sizeof(int32_t); const int64_t scratch = - kv_bytes_per_token * PAGED_BLOCK_SIZE; + kv_bytes_per_token * scratch_tokens; return recurrent + target_feat + q_capture + paged_metadata + scratch; } @@ -331,10 +332,20 @@ bool Qwen35Backend::init() { } } + // Feature-gate validation normally catches this before construction, but + // keep backend arithmetic safe for direct/test callers too. + if (cfg_.ddtree_mode && + (cfg_.ddtree_budget < 1 || cfg_.ddtree_budget > 255)) { + set_last_error("--ddtree-budget must be in [1, 255]"); + return false; + } + // Create KV cache const int max_verify_tokens = cfg_.ddtree_mode ? std::max(dw_.block_size, cfg_.ddtree_budget + 1) : dw_.block_size; + const int n_slots = concurrent_slots(); + // kvflash (bounded residency): pool size from the env, rounded/floored/ // clamped by the shared reader (256-stride keeps FA vec-kernel // eligibility; the floor keeps eviction from deadlocking). @@ -366,16 +377,16 @@ bool Qwen35Backend::init() { if (!post_kvflash_init_gate()) return false; // KVFlash is resolved from env at init; this is the authoritative // paged×KVFlash compatibility check. - if (cfg_.paged_attention && kvflash_active()) { + if (cfg_.paged_attention && kvflash_active() && n_slots <= 1) { std::fprintf(stderr, - "[paged-attention] cannot be combined with KVFlash " - "(resident pool %d tokens)\n", kvflash_tokens_); - set_last_error("paged attention cannot be combined with KVFlash"); + "[paged-attention] single-sequence paged decode cannot be combined " + "with KVFlash (resident pool %d tokens)\n", kvflash_tokens_); + set_last_error( + "paged KVFlash requires --max-concurrency greater than 1"); return false; } // Paged mode sizes the KV cache to whole blocks; otherwise KVFlash // decides the allocation (0 = full max_ctx). - const int n_slots = concurrent_slots(); const int max_concurrent_prefills = n_slots > 1 ? std::clamp( env_int_or_default("DFLASH_MAX_CONCURRENT_PREFILLS", 8), @@ -395,6 +406,17 @@ bool Qwen35Backend::init() { set_last_error("--max-concurrency requires --paged-attention"); return false; } + const bool concurrent_local_ddtree = + n_slots > 1 && cfg_.ddtree_mode && cfg_.draft_path && + !use_remote_draft && !tensor_parallel && !split_gpus_ && + target_backend_ == draft_backend_; + const int tree_width = concurrent_local_ddtree + ? cfg_.ddtree_budget + 1 : 0; + const int tree_stride = concurrent_local_ddtree + ? paged_token_capacity(tree_width) : 0; + const int64_t concurrent_scratch_tokens = + (int64_t)n_slots * tree_stride + PAGED_BLOCK_SIZE; + // Concurrent slots share one physical pool. An explicit // --kv-pool-tokens is rounded up to a whole block; otherwise capacity is // derived from device-free memory after subtracting fixed concurrent cache @@ -404,10 +426,17 @@ bool Qwen35Backend::init() { // pool's index space) as the write target of dead decode-batch rows. int64_t pool_tokens = 0; if (n_slots > 1) { - if (cfg_.kv_pool_tokens > 0) { + if (kvflash_active()) { + pool_tokens = paged_token_capacity(kvflash_tokens_); + std::fprintf(stderr, + "[parallel-kvflash] physical resident pool %lld tokens; " + "logical per-slot cap %d across %d slots " + "(--kv-pool-tokens does not expand resident VRAM)\n", + (long long)pool_tokens, cfg_.device.max_ctx, n_slots); + } else if (cfg_.kv_pool_tokens > 0) { pool_tokens = (int64_t)paged_token_capacity( (int)std::min( - cfg_.kv_pool_tokens, INT32_MAX - PAGED_BLOCK_SIZE)); + cfg_.kv_pool_tokens, INT32_MAX - concurrent_scratch_tokens)); } else { PagedKvAutoBudget budget; // TODO: Size tensor-parallel pools from each device's free memory @@ -416,7 +445,8 @@ bool Qwen35Backend::init() { budget.bytes_per_token = kvf_budget.bytes_per_token; budget.reserve_bytes = kvf_budget.reserve_bytes; budget.fixed_cache_bytes = concurrent_fixed_cache_bytes( - w_, cfg_.device.max_ctx, n_slots, budget.bytes_per_token); + w_, cfg_.device.max_ctx, n_slots, budget.bytes_per_token, + concurrent_scratch_tokens); pool_tokens = paged_kv_auto_pool_tokens( cfg_.device.max_ctx, n_slots, budget); const int64_t one_context = @@ -438,19 +468,19 @@ bool Qwen35Backend::init() { return false; } } - if (pool_tokens + PAGED_BLOCK_SIZE > INT32_MAX) { + if (pool_tokens + concurrent_scratch_tokens > INT32_MAX) { set_last_error("paged KV pool exceeds INT32_MAX tokens"); return false; } } const int ctx_alloc = n_slots > 1 - ? (int)(pool_tokens + PAGED_BLOCK_SIZE) + ? (int)(pool_tokens + concurrent_scratch_tokens) : (cfg_.paged_attention ? paged_token_capacity(cfg_.device.max_ctx) : kvflash_tokens_); if (!create_target_cache(w_, cfg_.device.max_ctx, max_verify_tokens, target_backend_, cache_, /*prefill_only=*/true, ctx_alloc, - cfg_.paged_attention, n_slots)) { + cfg_.paged_attention, n_slots, concurrent_local_ddtree)) { std::fprintf(stderr, "cache: %s\n", dflash27b_last_error()); return false; } @@ -481,10 +511,40 @@ bool Qwen35Backend::init() { e.what()); return false; } + if (n_slots > 1 && kvflash_active()) { + std::string transfer_error; + paged_kv_transfer_ = QwenPagedKvResidencyTransfer::create( + cache_, target_backend_, cfg_.device.gpu, PAGED_BLOCK_SIZE, + &transfer_error); + if (!paged_kv_transfer_) { + set_last_error("concurrent KVFlash transfer init failed: " + + transfer_error); + return false; + } + PagedKvResidencyConfig residency_cfg; + residency_cfg.block_bytes = paged_kv_transfer_->block_bytes(); + residency_cfg.resident_budget_blocks = + paged_kv_pool_->physical_block_count(); + residency_cfg.sink_blocks = 1; + residency_cfg.tail_blocks = 4; + try { + paged_kv_residency_ = + std::make_unique( + *paged_kv_pool_, residency_cfg, + paged_kv_transfer_->callbacks()); + } catch (const std::exception & e) { + set_last_error(std::string( + "concurrent KVFlash residency init failed: ") + e.what()); + return false; + } + } if (n_slots > 1) { + const int tree_scratch_base = (int)pool_tokens; + const int64_t dead_scratch_row = + pool_tokens + (int64_t)n_slots * tree_stride; seq_engine_ = std::make_unique( *this, *paged_kv_pool_, cfg_.device.max_ctx, - /*scratch_row=*/pool_tokens, + dead_scratch_row, tree_width, tree_scratch_base, tree_stride, max_concurrent_prefills, mixed_prefill_tokens, long_mixed_prefill_tokens, long_prefill_threshold, idle_prefill_tokens, prefill_quantum); @@ -509,7 +569,7 @@ bool Qwen35Backend::init() { cfg_.device.max_ctx); std::fflush(stdout); } - if (kvflash_active()) { + if (kvflash_active() && !(cfg_.paged_attention && n_slots > 1)) { KvFlashConfig pc; pc.pool_tokens = kvflash_tokens_; if (!kvflash_pager_.attach(pc, cache_.attn_k, cache_.attn_v)) { @@ -540,7 +600,8 @@ bool Qwen35Backend::init() { // Init feature mirror when draft model is available (needed for spec decode). // On single-GPU, this is an F32 conversion buffer; on split-GPU, a cross-device mirror. - if (cfg_.draft_path && !use_remote_draft) { + if (cfg_.draft_path && !use_remote_draft && + !(n_slots > 1 && concurrent_local_ddtree)) { const int mirror_cap = std::min({cfg_.draft_ctx_max, cfg_.device.max_ctx, cache_.target_feat_cap > 0 ? cache_.target_feat_cap : cfg_.device.max_ctx}); if (!draft_feature_mirror_init(feature_mirror_, draft_backend_, @@ -1184,7 +1245,11 @@ DFlashTarget * Qwen35Backend::dflash_target() { void Qwen35Backend::shutdown() { const bool use_remote_draft = cfg_.remote_draft.enabled(); + seq_engine_.reset(); end_paged_sequence(); + paged_kv_residency_.reset(); + paged_kv_transfer_.reset(); + paged_kv_pool_.reset(); free_drafter(); step_graph_destroy(sg_); step_graph_destroy(draft_sg_); diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index a6a508f73..62fee8524 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -24,7 +24,9 @@ #include "concurrency/qwen35_seq_engine.h" #include "internal.h" // TargetWeights, TargetCache, DraftWeights, PrefixSnapshot #include "qwen3/qwen3_drafter.h" // DrafterContext, load_drafter, free_drafter, drafter_score_and_compress -#include "kvflash_pager.h" // bounded KV residency pool +#include "kvflash_pager.h" +#include "common/concurrency/paged_kv_residency.h" +#include "common/concurrency/qwen_paged_kv_transfer.h" #include "kvflash_scorer.h" // chunk-relevance policy interface #include "kvflash_qk.h" // target-QK scorer (pooled keys + query) @@ -295,6 +297,8 @@ class Qwen35Backend : public ModelBackend { // Page size comes from PAGED_BLOCK_SIZE (paged_attention_config.h), // shared with the graph builder and the cache's block-aligned sizing. std::unique_ptr paged_kv_pool_; + std::unique_ptr paged_kv_transfer_; + std::unique_ptr paged_kv_residency_; std::optional paged_sequence_; PagedKvRequestId paged_request_id_ = 0; diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 1082df6f5..5f10dcdc9 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -80,12 +80,14 @@ bool create_target_cache(const TargetWeights & w, bool prefill_only, int ctx_alloc, bool paged_attention, - int n_seq_slots) { + int n_seq_slots, + bool concurrent_tree) { return create_target_cache_partial(w, max_ctx, max_verify_tokens, backend, out, prefill_only, 0, w.n_layer, true, ctx_alloc, /*f32_ssm_intermediates=*/false, - paged_attention, n_seq_slots); + paged_attention, n_seq_slots, + concurrent_tree); } // concurrent_fixed_cache_bytes() in qwen35_backend.cpp mirrors this @@ -103,7 +105,8 @@ bool create_target_cache_partial(const TargetWeights & w, int ctx_alloc, bool f32_ssm_intermediates, bool paged_attention, - int n_seq_slots) { + int n_seq_slots, + bool concurrent_tree) { if (layer_begin < 0) layer_begin = 0; if (layer_end < 0 || layer_end > w.n_layer) layer_end = w.n_layer; if (layer_begin > layer_end) { @@ -115,6 +118,11 @@ bool create_target_cache_partial(const TargetWeights & w, set_last_error("multi-slot target cache requires paged attention"); return false; } + if (concurrent_tree && (!paged_attention || n_seq_slots <= 1)) { + set_last_error( + "concurrent tree cache requires paged multi-slot serving"); + return false; + } out.backend = backend; out.max_ctx = max_ctx; out.cur_pos = 0; @@ -227,7 +235,14 @@ bool create_target_cache_partial(const TargetWeights & w, out.target_feat_cap = std::min(max_ctx, TARGET_FEAT_CAP_DEFAULT); if (allocate_target_feat) { const int fc_in = w.n_capture_layers * w.n_embd; - out.target_feat = ggml_new_tensor_2d(out.base_ctx, GGML_TYPE_BF16, fc_in, out.target_feat_cap); + // Concurrent slots own disjoint feature rings. The final row is + // dead scratch for padded bucket rows because set_rows does not + // accept negative destination indices. + const int feat_rows = multi_slot + ? out.target_feat_cap * n_seq_slots + 1 + : out.target_feat_cap; + out.target_feat = ggml_new_tensor_2d( + out.base_ctx, GGML_TYPE_BF16, fc_in, feat_rows); ggml_set_name(out.target_feat, "target_feat"); } else { out.target_feat = nullptr; @@ -266,9 +281,10 @@ bool create_target_cache_partial(const TargetWeights & w, } // ── Rollback context: snapshots + intermediates ─────────────────── - // Multi-slot caches skip these entirely: concurrent serving is paged and - // therefore AR-only (no spec-decode rollback), and the tensors are the - // single largest optional allocation (~0.8 GB at 48 delta layers). + // Multi-slot caches skip these entirely. Packed tree verification gathers + // the selected slots' base state without mutating it, then a bounded + // replay commits accepted paths. T*S recurrent captures would be tens of + // GiB at useful Strix concurrency and are intentionally not allocated. if (!prefill_only && !multi_slot) { const int rb_tensors = 4 * n_delta; ggml_init_params ip{}; @@ -722,7 +738,16 @@ static ggml_tensor * build_full_attn_block( int paged_max_kv_len = 0, // Compact decode row -> physical block-table column. Negative ids are // graph-bucket padding rows. - ggml_tensor * active_slot_ids = nullptr + ggml_tensor * active_slot_ids = nullptr, + // Packed paged-tree verification. Query rows are flattened + // sequence-major; row mappings are supplied through + // paged_query_seq_ids, while parent/tree metadata describes each tree. + ggml_tensor * paged_tree_parent_ids = nullptr, + ggml_tensor * paged_tree_sizes = nullptr, + int tree_width = 0, + int tree_scratch_base = 0, + int tree_scratch_stride = 0, + int paged_logical_max_ctx = 0 ) { const int head_dim = w.n_embd_head_k; const int n_head = w.n_head; @@ -810,9 +835,14 @@ static ggml_tensor * build_full_attn_block( Kcur_T = ggml_turbo_wht(ctx, Kcur_T, 0); } + const bool paged_tree = paged_tree_parent_ids || paged_tree_sizes; + GGML_ASSERT((paged_tree_parent_ids == nullptr) == + (paged_tree_sizes == nullptr)); const bool ragged = paged_query_seq_ids != nullptr; - GGML_ASSERT(!ragged || (paged_block_table && paged_query_positions && - kv_write_rows)); + GGML_ASSERT(!ragged || (paged_block_table && kv_write_rows)); + GGML_ASSERT(!ragged || paged_tree || paged_query_positions); + GGML_ASSERT(!paged_tree || + (ragged && !paged_query_positions && tree_width > 0)); if (kv_write_rows) { // Step-invariant: the destination tensor stays fixed while the input // indices carry contiguous, KVFlash, or paged physical rows. @@ -878,12 +908,31 @@ static ggml_tensor * build_full_attn_block( ggml_tensor * row_seq_ids, ggml_tensor * row_positions, bool dense_token_layout) { - const int padded = ((std::max(1, launch_kv_len) + 255) / 256) * 256; - const int launch_len = std::min(padded, (int)cache_k->ne[1]); + // max_kv_seq_len sizes the logical partition grid. In bounded + // KVFlash mode the block table maps that logical range onto a much + // smaller physical K/V pool, so cache_k->ne[1] is not a valid clamp. + // Bound against both sources of logical capacity instead, doing the + // 256-window rounding in i64 to avoid signed overflow at large + // configured contexts. Per-row kv_seq_lens remains the exact runtime + // bound, and the paged kernel bounds every resolved physical row. + GGML_ASSERT(paged_block_table && cache_k && cache_v); + const int64_t table_capacity = + (int64_t)paged_block_table->ne[0] * PAGED_BLOCK_SIZE; + const int64_t logical_capacity = + std::min(paged_logical_max_ctx, table_capacity); + GGML_ASSERT(logical_capacity > 0 && logical_capacity <= INT32_MAX); + const int64_t requested = + std::min(std::max(1, launch_kv_len), + logical_capacity); + const int64_t padded = ((requested + 255) / 256) * 256; + const int launch_len = + (int)std::min(padded, logical_capacity); ggml_tensor * out = ggml_paged_attn_ext( ctx, q, cache_k, cache_v, paged_block_table, paged_kv_seq_lens, row_seq_ids, row_positions, kq_scale, - PAGED_BLOCK_SIZE, launch_len); + PAGED_BLOCK_SIZE, launch_len, + paged_tree_parent_ids, paged_tree_sizes, + tree_width, tree_scratch_base, tree_scratch_stride); if (dense_token_layout) { out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); } @@ -891,7 +940,19 @@ static ggml_tensor * build_full_attn_block( }; ggml_tensor * attn = nullptr; - if (ragged) { + if (paged_tree) { + // ── Packed concurrent tree verify. Every query row selects its + // physical sequence/scratch slab. The paged kernel combines the + // committed block-table prefix with only this node's ancestor chain; + // query_positions is intentionally absent in tree mode. + ggml_tensor * Qfa = q_segment(0, n_tokens); + if (q_fa_out) *q_fa_out = Qfa; + const int launch_kv_len = paged_max_kv_len > 0 + ? paged_max_kv_len : kv_start + n_tokens; + attn = paged_read(Qfa, launch_kv_len, + paged_query_seq_ids, /*row_positions=*/nullptr, + /*dense_token_layout=*/n_tokens > 1); + } else if (ragged) { // ── Ragged concurrent step: prefill chunk rows and decode rows all // read the pool through one call, each row clamped to its own // inclusive position. This step's chunk rows are visible to their @@ -899,6 +960,7 @@ static ggml_tensor * build_full_attn_block( // attention in the graph; cross-sequence isolation is structural // (each row's seq id selects its own block-table column). ggml_tensor * Qfa = q_segment(0, n_tokens); + if (q_fa_out) *q_fa_out = Qfa; const int launch_kv_len = paged_max_kv_len > 0 ? paged_max_kv_len : kv_start + n_tokens; attn = paged_read(Qfa, launch_kv_len, @@ -920,8 +982,8 @@ static ggml_tensor * build_full_attn_block( // bound only over-sizes the partition grid, and partitions past the // real length exit with a zero-weight sentinel. // Batched decode: kv_len (kv_start + n_tokens) describes one sequence; - // the launch bound must cover the longest live slot instead. Clamped - // because ggml_paged_attn asserts max_kv_seq_len <= k->ne[1]. + // the launch bound must cover the longest live slot instead. Bounded + // paged pools may be physically smaller than this logical span. const int launch_kv_len = paged_max_kv_len > 0 ? paged_max_kv_len : kv_len; attn = paged_read( Qfa, launch_kv_len, active_slot_ids, /*row_positions=*/nullptr, @@ -1032,8 +1094,13 @@ static ggml_tensor * build_delta_net_block( prefill_total += prefill_segments[i].n_tokens; } GGML_ASSERT((active_slot_ids == nullptr) == (state_slot_ids == nullptr)); + const bool mapped_tree = active_slot_ids && parent_ids; + GGML_ASSERT(!active_slot_ids || !cap); GGML_ASSERT(!active_slot_ids || - (!cap && !parent_ids && prefill_total + n_seqs == n_tokens)); + (mapped_tree + ? (!ragged && prefill_total == 0 && + n_tokens % n_seqs == 0) + : (prefill_total + n_seqs == n_tokens))); if (!active_slot_ids) { GGML_ASSERT(n_seqs == 1); GGML_ASSERT(prefill_total == 0 || prefill_total == n_tokens); @@ -1065,6 +1132,7 @@ static ggml_tensor * build_delta_net_block( int T; // timesteps per sequence int S; // sequences bool active; // compact decode segment (slot-mapped) + bool tree; // mapped tree: gather-only, no persistence ggml_tensor * conv_st; ggml_tensor * ssm_st; }; @@ -1082,14 +1150,17 @@ static ggml_tensor * build_delta_net_block( ssm_state->ne[0], ssm_state->ne[1], ssm_state->ne[2], 1, ssm_state->nb[1], ssm_state->nb[2], ssm_state->nb[3], (size_t)pf.seq_slot * ssm_state->nb[3]); - segs.push_back({pf.token_offset, pf.n_tokens, 1, false, c, s}); + segs.push_back({pf.token_offset, pf.n_tokens, 1, + false, false, c, s}); } if (active_slot_ids) { - segs.push_back({prefill_total, 1, n_seqs, true, - conv_state, ssm_state}); + const int tree_tokens = mapped_tree ? n_tokens / n_seqs : 1; + segs.push_back({prefill_total, tree_tokens, n_seqs, true, + mapped_tree, conv_state, ssm_state}); } else if (segs.empty()) { // No general [timesteps x sequences] mode: one multi-token sequence. - segs.push_back({0, n_tokens, n_seqs, false, conv_state, ssm_state}); + segs.push_back({0, n_tokens, n_seqs, false, false, + conv_state, ssm_state}); } const int n_segs = (int)segs.size(); @@ -1109,13 +1180,14 @@ static ggml_tensor * build_delta_net_block( const int seg_seqs = seg.S; const int seg_tokens = seg.T * seg.S; const bool seg_active = seg.active; + const bool seg_tree = seg.tree; // Plain one-token decode has no in-graph consumer of the updated state: // the next graph evaluation is the first read. Write the final state // directly into its persistent slab and avoid materializing/copying a // second S_v x S_v x H_v state. The active-aware path also updates each // mapped physical slab directly; only its negative bucket-padding rows // use the result tensor's retained scratch state region. - const bool inplace_state = seg_active || + const bool inplace_state = (seg_active && !seg_tree) || (allow_inplace_state && can_skip_gdn_intermediate && !ragged && n_seq_tokens == 1); @@ -1183,7 +1255,7 @@ static ggml_tensor * build_delta_net_block( w.ssm_d_conv - 1, conv_channels, seg_seqs, conv_input->nb[1], conv_input->nb[2], (conv_input->ne[0] - (w.ssm_d_conv - 1)) * ggml_element_size(conv_input)); - if (seg_active) { + if (seg_active && !seg_tree) { const int64_t slab = (int64_t)(w.ssm_d_conv - 1) * conv_channels; ggml_tensor * compact_last = ggml_reshape_2d( @@ -1193,7 +1265,7 @@ static ggml_tensor * build_delta_net_block( ggml_build_forward_expand( gf, ggml_set_rows_masked( ctx, all_conv, compact_last, active_slot_ids)); - } else { + } else if (!seg_tree) { ggml_build_forward_expand(gf, ggml_cpy(ctx, last_conv, seg.conv_st)); } @@ -1246,10 +1318,25 @@ static ggml_tensor * build_delta_net_block( } // ── SSM state (recurrent): reshape to [S_v, S_v, H_v, n_seqs] - ggml_tensor * s = seg_active - ? seg.ssm_st - : ggml_reshape_4d(ctx, seg.ssm_st, + ggml_tensor * s = nullptr; + if (seg_tree) { + // Packed tree verification starts each tree from the owning slot's + // base state. Gather compact slabs, then leave the persistent tensor + // untouched; accepted paths are committed by a later replay. + const int64_t slab = + (int64_t)head_v_dim * head_v_dim * num_v_heads; + ggml_tensor * all_ssm = ggml_reshape_2d( + ctx, seg.ssm_st, slab, seg.ssm_st->ne[3]); + ggml_tensor * gathered = + ggml_get_rows(ctx, all_ssm, state_slot_ids); + s = ggml_reshape_4d(ctx, gathered, head_v_dim, head_v_dim, num_v_heads, seg_seqs); + } else { + s = seg_active + ? seg.ssm_st + : ggml_reshape_4d(ctx, seg.ssm_st, + head_v_dim, head_v_dim, num_v_heads, seg_seqs); + } // ── Fused Gated DeltaNet op — returns packed (output | new_state [| intermediates]). // In tree mode, the kernel uses parent_ids to reload state at DFS @@ -1302,7 +1389,7 @@ static ggml_tensor * build_delta_net_block( ggml_build_forward_expand(gf, ggml_cpy(ctx, r.new_state, s)); } else { ggml_tensor * result; - if (seg_active) { + if (seg_active && !seg_tree) { result = ggml_gated_delta_net_active_inplace( ctx, q_c, k_c, v_c, g_tensor, beta, s, active_slot_ids); } else if (parent_ids) { @@ -1337,7 +1424,7 @@ static ggml_tensor * build_delta_net_block( S_v * H_v * r_elt, S_v * H_v * n_seq_tokens * r_elt, 0); - if (!inplace_state) { + if (!inplace_state && !seg_tree) { ggml_tensor * new_state = ggml_view_4d(ctx, result, S_v, S_v, H_v, seg_seqs, S_v * r_elt, @@ -1345,8 +1432,8 @@ static ggml_tensor * build_delta_net_block( S_v * S_v * H_v * r_elt, S_v * H_v * n_seq_tokens * seg_seqs * r_elt); - // Persist new_state back to cache. Both compact active decode and the - // plain in-place AR path write state from the GDN kernel directly. + // Persist new_state back to cache. Mapped trees deliberately skip + // this branch: their gathered base state is read-only. ggml_build_forward_expand(gf, ggml_cpy(ctx, new_state, seg.ssm_st)); } @@ -1555,6 +1642,12 @@ QwenGraphOutputs build_qwen35_graph( const int hidden = w.n_embd; const float eps = w.rms_eps; + const bool capture_with_rows = + in.capture_layers && cache.target_feat && in.target_feat_rows; + std::vector capture_slices; + if (capture_with_rows) { + capture_slices.assign((size_t)N_CAPTURE, nullptr); + } for (int il = 0; il < w.n_layer; il++) { const TargetLayer & L = w.layers[il]; @@ -1584,7 +1677,13 @@ QwenGraphOutputs build_qwen35_graph( in.paged_query_seq_ids, in.paged_query_positions, in.paged_max_kv_len, - in.active_slot_ids); + in.active_slot_ids, + in.parent_ids, + in.tree_sizes, + in.tree_width, + in.tree_scratch_base, + in.tree_scratch_stride, + cache.max_ctx); if (want_q_cap && q_fa) { // Last token's Q, all heads: src [head_dim, 1, n_head] view of // [head_dim, n_tokens, n_head]; dst = q_cap plane fa_idx @@ -1676,6 +1775,13 @@ QwenGraphOutputs build_qwen35_graph( if (CAPTURE_LAYERS[k] == il) { capture_idx = k; break; } } if (capture_idx >= 0) { + ggml_tensor * cur_2d = + ggml_reshape_2d(ctx, cur, hidden, n_tokens); + if (capture_with_rows) { + capture_slices[(size_t)capture_idx] = cur_2d; + inpL = cur; + continue; + } const size_t elt = ggml_element_size(cache.target_feat); const size_t col_stride = cache.target_feat->nb[1]; const int cap = cache.target_feat_cap; @@ -1683,8 +1789,6 @@ QwenGraphOutputs build_qwen35_graph( const int pre_n = std::min(n_tokens, cap - slot_start); const int post_n = n_tokens - pre_n; - ggml_tensor * cur_2d = ggml_reshape_2d(ctx, cur, hidden, n_tokens); - // First slice: [slot_start..slot_start+pre_n) in the ring. { const size_t offset = @@ -1714,6 +1818,21 @@ QwenGraphOutputs build_qwen35_graph( inpL = cur; } + if (capture_with_rows) { + GGML_ASSERT(!capture_slices.empty()); + ggml_tensor * feat_cat = capture_slices[0]; + GGML_ASSERT(feat_cat); + for (int k = 1; k < (int)capture_slices.size(); ++k) { + GGML_ASSERT(capture_slices[(size_t)k]); + feat_cat = ggml_concat( + ctx, feat_cat, capture_slices[(size_t)k], 0); + } + feat_cat = ggml_cont(ctx, feat_cat); + ggml_build_forward_expand( + gf, ggml_set_rows( + ctx, cache.target_feat, feat_cat, in.target_feat_rows)); + } + // 2. Final norm ggml_tensor * out = rms_norm_mul(ctx, inpL, w.out_norm, w.rms_eps); diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index e92f495bf..52a4a845c 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -1379,11 +1379,27 @@ int HttpServer::run() { std::fprintf(stderr, "[server] listening on http://%s:%d\n", config_.host.c_str(), config_.port); - // A backend-provided sequence engine replaces the one-request worker - // with the concurrent scheduler. Upstream forwarding stays on the - // classic path even when the local backend exposes an engine. - if (SeqEngine * engine = backend_.seq_engine(); - engine && config_.pflash_upstream_base.empty()) { + // A backend-provided sequence engine replaces the one-request worker. + // Local PFlash stays on this path too: scheduler admission prepares each + // prompt exactly once, then admits the effective tokens. Persistent + // residency is the explicit safety contract that lets compression run + // without parking model state owned by other live slots. + SeqEngine * engine = backend_.seq_engine(); + if (engine && config_.pflash_upstream_base.empty()) { + const ConcurrentPflashPlan pflash_plan = + resolve_concurrent_pflash_plan(config_, drafter_tokenizer_ != nullptr); + if (!pflash_plan.ok()) { + std::fprintf(stderr, "[server] %s\n", pflash_plan.error.c_str()); + socket_close(listen_fd_); + listen_fd_ = kInvalidSocket; + return 2; + } + if (pflash_plan.force_skip_park && !config_.pflash_skip_park) { + config_.pflash_skip_park = true; + std::fprintf(stderr, + "[server] concurrent PFlash: persistent residency enables " + "skip-park for live sequence safety\n"); + } worker_thread_ = std::thread([this, engine]() { scheduler_loop(*engine); }); } else { @@ -4231,6 +4247,7 @@ std::string HttpServer::format_http_response( case 400: reason = "Bad Request"; break; case 404: reason = "Not Found"; break; case 405: reason = "Method Not Allowed"; break; + case 409: reason = "Conflict"; break; case 413: reason = "Payload Too Large"; break; case 500: reason = "Internal Server Error"; break; case 503: reason = "Service Unavailable"; break; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 90e02412e..0fa0ef718 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -250,6 +250,49 @@ bool should_clamp_flowkv_disk_cache( bool flowkv, const DiskPrefixCachePolicy & policy); } // namespace http_detail +// Resolve the prompt-compression contract for a backend-provided sequence +// engine. Concurrent PFlash must keep both models resident: parking either +// model while another slot owns live device state invalidates that slot. The +// explicit persistent policy is therefore the opt-in that also implies the +// effective skip-park behavior; callers do not need a second, redundant CLI +// flag on large-memory concurrent hosts. +struct ConcurrentPflashPlan { + bool enabled = false; + bool force_skip_park = false; + std::string error; + + bool ok() const { return error.empty(); } +}; + +inline ConcurrentPflashPlan resolve_concurrent_pflash_plan( + const ServerConfig & config, bool drafter_tokenizer_available) { + ConcurrentPflashPlan plan; + if (config.pflash_mode == ServerConfig::PflashMode::OFF || + !config.pflash_upstream_base.empty()) { + return plan; + } + plan.enabled = true; + if (!drafter_tokenizer_available) { + plan.error = + "concurrent PFlash requires a loaded --prefill-drafter tokenizer"; + return plan; + } + if (config.draft_residency != DraftResidencyPolicy::Persistent) { + plan.error = + "concurrent PFlash requires --draft-residency persistent so " + "prompt compression cannot park live target/draft state"; + return plan; + } + if (config.prefix_cache_cap > 0 || config.prefill_cache_cap > 0 || + !config.disk_cache_dir.empty()) { + plan.error = + "concurrent paged PFlash does not support prefix/prefill " + "snapshots; disable the snapshot caches"; + return plan; + } + plan.force_skip_park = true; + return plan; +} // ─── Parsed request ───────────────────────────────────────────────────── @@ -605,6 +648,13 @@ struct ServerJob { // server-side prefill/elapsed telemetry does not erase queueing delay. std::chrono::steady_clock::time_point parallel_started_at{}; std::unique_ptr emitter; + // Prompt preparation (FlowKV/PFlash) is expensive and may load a resident + // drafter. Cache its result on the job so a pool-full retry never runs it + // twice. The original request tokens remain untouched for API accounting; + // this vector is the effective prompt admitted to the sequence engine. + bool parallel_prompt_prepared = false; + bool parallel_prompt_compressed = false; + std::vector parallel_prompt_tokens; }; // ─── Parse session_id from a chat-completion JSON body ────────────────── diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 3c0b8f7c3..c34c71fb6 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -36,6 +36,16 @@ struct SchedSlot { double prefill_s = 0.0; int n_gen_cap = 0; int completion_tokens = 0; + int effective_prompt_tokens = 0; + bool prompt_compressed = false; + uint64_t engine_request_id = 0; + uint64_t ddtree_steps = 0; + uint64_t ddtree_accepted_tokens = 0; + uint64_t target_forwards = 0; + uint64_t kvflash_page_ins = 0; + uint64_t kvflash_page_outs = 0; + uint64_t kvflash_resident_blocks = 0; + uint64_t kvflash_reselects = 0; bool client_disconnected = false; bool failed = false; std::string error; @@ -238,7 +248,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { stop_job_stream(s.job, &s.send_buffer); const double decode_s = std::chrono::duration( std::chrono::steady_clock::now() - s.decode_started_at).count(); - const int prompt_tokens = (int)req.prompt_tokens.size(); + const int prompt_tokens = s.effective_prompt_tokens; GenTimings gen_timings{ s.prefill_s, decode_s, @@ -253,9 +263,10 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { perf.prompt_tokens = (int)req.prompt_tokens.size(); perf.completion_tokens = s.completion_tokens; perf.prefill_tok_s = s.prefill_s > 0.0 - ? (double)req.prompt_tokens.size() / s.prefill_s : 0.0; + ? (double)prompt_tokens / s.prefill_s : 0.0; perf.decode_tok_s = decode_s > 0.0 ? (double)s.completion_tokens / decode_s : 0.0; + perf.pflash = s.prompt_compressed; status_.record_perf(perf); } @@ -292,17 +303,38 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { std::chrono::steady_clock::now() - s.started_at).count(); const int out_tokens = (int)s.gen_tokens.size(); std::fprintf(stderr, - "[server] chat DONE %s ok=%s in=%zu out=%d %.1fs %.1f tok/s " + "[server] chat DONE %s ok=%s in=%zu effective_in=%d out=%d %.1fs %.1f tok/s " "finish=%s slot=%d prefill=%.1fs decode=%.1fs(%.1ftok/s) parallel\n", req.response_id.c_str(), (!s.failed && backend_ok) ? "true" : "false", - req.prompt_tokens.size(), out_tokens, elapsed_s, + req.prompt_tokens.size(), prompt_tokens, out_tokens, elapsed_s, elapsed_s > 0.0 ? out_tokens / elapsed_s : 0.0, s.client_disconnected ? "client_disconnect" : s.emitter->finish_reason().c_str(), idx, s.prefill_s, decode_s, decode_s > 0.0 ? out_tokens / decode_s : 0.0); + const json concurrency_metrics = { + {"request_id", req.response_id}, + {"response_id", req.response_id}, + {"engine_request_id", s.engine_request_id}, + {"raw_prompt_tokens", req.prompt_tokens.size()}, + {"effective_prompt_tokens", prompt_tokens}, + {"output_tokens", out_tokens}, + {"pflash_applied", s.prompt_compressed}, + {"pflash_input_tokens", req.prompt_tokens.size()}, + {"pflash_output_tokens", prompt_tokens}, + {"ddtree_steps", s.ddtree_steps}, + {"ddtree_accepted_tokens", s.ddtree_accepted_tokens}, + {"target_forwards", s.target_forwards}, + {"kvflash_page_ins", s.kvflash_page_ins}, + {"kvflash_page_outs", s.kvflash_page_outs}, + {"kvflash_resident_blocks", s.kvflash_resident_blocks}, + {"kvflash_reselects", s.kvflash_reselects}, + }; + std::fprintf(stderr, "[concurrency-metrics] %s\n", + concurrency_metrics.dump().c_str()); + engine.retire(idx); // A retirement may have released the blocks the head job needs. deferred_retry_at = {}; @@ -431,9 +463,65 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { start_job_stream(job); } + // Apply the same FlowKV/PFlash precedence and overflow checks as the + // classic worker. Keep the prepared prompt on the job because an + // atomically-busy admission is retried at the FIFO head later. + if (!job->parallel_prompt_prepared) { + PreparedPrompt prepared = prepare_prompt(req); + if (prepared.error_status != 0) { + const std::string message = prepared.error.empty() + ? "prompt preparation failed" + : prepared.error; + std::fprintf(stderr, + "[server] concurrent prompt preparation failed: %s\n", + message.c_str()); + if (req.stream && job->sse_started) { + stop_job_stream(job); + for (const std::string & chunk : + sse_error_close_chunks(message)) { + send_job_bytes(job, chunk.data(), chunk.size()); + } + } else { + send_error(job->fd, prepared.error_status, message); + } + finish_job(job); + return AdmissionDisposition::Retired; + } + // Paged sequence engines cannot restore the classic snapshot + // format. Startup normally disables those caches; keep this + // check as a hard guard for embedded/non-CLI callers. + if (prepared.full_cache_hit_slot >= 0 || + prepared.full_cache_served_tokens >= 0) { + const std::string message = + "concurrent paged serving cannot restore a prefix snapshot"; + if (req.stream && job->sse_started) { + stop_job_stream(job); + for (const std::string & chunk : + sse_error_close_chunks(message)) { + send_job_bytes(job, chunk.data(), chunk.size()); + } + } else { + send_error(job->fd, 409, message); + } + finish_job(job); + return AdmissionDisposition::Retired; + } + job->parallel_prompt_tokens = std::move(prepared.tokens); + job->parallel_prompt_compressed = prepared.compressed; + job->parallel_prompt_prepared = true; + std::fprintf(stderr, + "[server] concurrent prompt READY %s raw=%zu effective=%zu " + "pflash=%s\n", + req.response_id.c_str(), req.prompt_tokens.size(), + job->parallel_prompt_tokens.size(), + prepared.compressed ? "true" : "false"); + } + const auto & effective_prompt = job->parallel_prompt_tokens; + // Admission only claims the slot and queues the prompt. Prefill // advances one chunk per engine step alongside live decode. - auto ar = engine.admit(next_request_id, req.prompt_tokens, + const uint64_t engine_request_id = next_request_id; + auto ar = engine.admit(engine_request_id, effective_prompt, req.sampler); if (ar.status == SeqEngine::AdmitResult::Status::busy) return AdmissionDisposition::Deferred; @@ -464,9 +552,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.admission_order = next_admission_order++; s.started_at = started_at; s.decode_started_at = started_at; // sane on prefill failure + s.effective_prompt_tokens = (int)effective_prompt.size(); + s.prompt_compressed = job->parallel_prompt_compressed; + s.engine_request_id = engine_request_id; s.n_gen_cap = std::min( n_gen_cap, - engine.max_context() - (int)req.prompt_tokens.size() + 1); + engine.max_context() - s.effective_prompt_tokens + 1); s.emitter = std::move(job->emitter); s.send_buffer.mark_progress(std::chrono::steady_clock::now()); if (budget_active && !config_.think_close_token_ids.empty() && @@ -606,8 +697,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { prefill_candidates.clear(); for (int i = 0; i < n_slots; i++) { if (slots[(size_t)i].job && !slots[(size_t)i].prefilling) { - step_plan.decode.push_back( - {i, slots[(size_t)i].pending_tok}); + SeqEngine::StepInput input; + input.slot = i; + input.token = slots[(size_t)i].pending_tok; + input.allow_speculation = + slots[(size_t)i].hook.close_token_ids.empty(); + step_plan.decode.push_back(input); } else if (slots[(size_t)i].job) { prefill_candidates.push_back( {i, slots[(size_t)i].admission_order}); @@ -655,7 +750,18 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.finished = true; continue; } - advance_slot(s, out.token); + s.ddtree_steps += out.ddtree_steps; + s.ddtree_accepted_tokens += out.ddtree_accepted_tokens; + s.target_forwards += out.target_forwards; + s.kvflash_page_ins += out.kvflash_page_ins; + s.kvflash_page_outs += out.kvflash_page_outs; + s.kvflash_resident_blocks = std::max( + s.kvflash_resident_blocks, out.kvflash_resident_blocks); + s.kvflash_reselects += out.kvflash_reselects; + consume_decode_output_tokens(out, [&](int32_t token) { + advance_slot(s, token); + return !s.finished; + }); } using PrefillStatus = SeqEngine::PrefillOutput::Status; for (const auto & out : step_result.prefills) { diff --git a/server/test/bench_paged_attention.cpp b/server/test/bench_paged_attention.cpp index 49fc00960..1e5afe9c8 100644 --- a/server/test/bench_paged_attention.cpp +++ b/server/test/bench_paged_attention.cpp @@ -725,7 +725,8 @@ bool run_case( ggml_tensor * paged_output = ggml_paged_attn_ext( ctx.value, q_paged, k_paged, v_paged, table, kv_seq_lens_tensor, nullptr, nullptr, 1.0f / std::sqrt(static_cast(D)), - BLOCK_SIZE, max_context); + BLOCK_SIZE, max_context, + nullptr, nullptr, 0, 0, 0); ggml_tensor * contiguous_output = ggml_flash_attn_ext( ctx.value, q_contiguous, k_contiguous, v_contiguous, padding_mask, 1.0f / std::sqrt(static_cast(D)), diff --git a/server/test/seq_engine_contract.h b/server/test/seq_engine_contract.h index 63ca035ab..0361a893c 100644 --- a/server/test/seq_engine_contract.h +++ b/server/test/seq_engine_contract.h @@ -258,6 +258,19 @@ inline std::vector check_seq_engine_contract(SeqEngine & engine) { return violations; } + // Scheduler policy can retain commit authority for selected slots. A + // conforming engine may still use its ordinary AR implementation, but it + // must not return already-committed children for a disabled input. + SeqEngine::StepPlan no_speculation; + no_speculation.decode = decode_inputs(); + for (SeqEngine::StepInput & input : no_speculation.decode) { + input.allow_speculation = false; + } + if (!execute(no_speculation)) { + retire_all(); + return violations; + } + // A full engine is retryable admission pressure, not a request error. if (n_slots == 2) { const SeqEngine::AdmitResult full = engine.admit(3, {31}, greedy); diff --git a/server/test/test_ddtree_path.cpp b/server/test/test_ddtree_path.cpp new file mode 100644 index 000000000..e7d2bbf3a --- /dev/null +++ b/server/test/test_ddtree_path.cpp @@ -0,0 +1,46 @@ +#include "common/ddtree.h" + +#include +#include +#include +#include + +using dflash::common::DDTree; +using dflash::common::follow_verified_tree; +using dflash::common::truncate_verified_path; + +int main() { + DDTree tree; + tree.n_nodes = 2; + tree.token_ids = {11, 22}; + tree.depths = {1, 2}; + tree.parents = {-1, 0, 1}; + tree.child_maps.resize(3); + tree.child_maps[0][11] = 1; + tree.child_maps[1][22] = 2; + + const int32_t posterior[] = {11, 22, 33}; + int pending = -1; + std::vector accepted = + follow_verified_tree(tree, posterior, pending); + assert((accepted == std::vector{0, 1, 2})); + assert(pending == 33); + + // Truncating after node 1 means node 2's token becomes pending. Keeping + // the old value (33) would skip token 22 and describe uncommitted state. + assert(truncate_verified_path(accepted, 2, posterior, pending)); + assert((accepted == std::vector{0, 1})); + assert(pending == 22); + + // An unchanged path preserves the already-computed pending token. + assert(!truncate_verified_path(accepted, 2, posterior, pending)); + assert(pending == 22); + + // No headroom is represented explicitly and never dereferences a tip. + assert(truncate_verified_path(accepted, 0, posterior, pending)); + assert(accepted.empty()); + assert(pending == -1); + + std::puts("ddtree path tests passed"); + return 0; +} diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index d01b9da04..fb1a5571a 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -381,6 +381,53 @@ void test_feature_gate_paged_attention_requires_plain_ar_decode() { ddtree.ddtree_mode = true; CHECK(!gate_result(ddtree, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs concurrent_ddtree = base; + concurrent_ddtree.max_concurrency = 16; + concurrent_ddtree.draft_path = "/nonexistent/draft.gguf"; + concurrent_ddtree.ddtree_mode = true; + concurrent_ddtree.ddtree_budget = 22; + TEST_ASSERT(gate_accepts( + concurrent_ddtree, "qwen35", PlacementBackend::Cuda)); + TEST_ASSERT(gate_accepts( + concurrent_ddtree, "qwen35", PlacementBackend::Hip)); + + BackendArgs tensor_ddtree = concurrent_ddtree; + TEST_ASSERT(parse_placement_device_list( + "cuda:0,cuda:1", tensor_ddtree.device)); + tensor_ddtree.device.split_mode = TargetSplitMode::Tensor; + TEST_ASSERT(!gate_accepts( + tensor_ddtree, "qwen35", PlacementBackend::Cuda)); + + + BackendFeatureConfig concurrent_pflash; + concurrent_pflash.pflash_enabled = true; + concurrent_pflash.pflash_drafter_configured = true; + TEST_ASSERT(gate_accepts(concurrent_ddtree, "qwen35", + PlacementBackend::Hip, concurrent_pflash)); + + BackendArgs concurrent_plain = base; + concurrent_plain.max_concurrency = 16; + TEST_ASSERT(gate_accepts(concurrent_plain, "qwen35", + PlacementBackend::Hip, concurrent_pflash)); + BackendFeatureConfig concurrent_kvflash; + concurrent_kvflash.kvflash_enabled = true; + TEST_ASSERT(gate_accepts(concurrent_plain, "qwen35", + PlacementBackend::Hip, concurrent_kvflash)); + + BackendArgs bad_budget = concurrent_ddtree; + for (int value : {0, -1, 256, INT_MAX}) { + bad_budget.ddtree_budget = value; + TEST_ASSERT(!gate_accepts( + bad_budget, "qwen35", PlacementBackend::Hip)); + } + + BackendArgs remote_ddtree = concurrent_ddtree; + remote_ddtree.remote_draft.ipc_bin = "/usr/bin/draft-ipc"; + remote_ddtree.draft_device.backend = PlacementBackend::Cuda; + remote_ddtree.device.backend = PlacementBackend::Hip; + TEST_ASSERT(!gate_accepts( + remote_ddtree, "qwen35", PlacementBackend::Hip)); + BackendArgs windowed = base; windowed.fa_window = 4096; CHECK(!gate_result( @@ -481,6 +528,22 @@ void test_feature_gate_parallel_and_kv_pool_rules() { pool.kv_pool_tokens = max_pool_tokens; CHECK(gate_result(pool, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs tree_pool = paged; + tree_pool.max_concurrency = 16; + tree_pool.draft_path = "/nonexistent/draft.gguf"; + tree_pool.ddtree_mode = true; + tree_pool.ddtree_budget = 22; + const long long tree_scratch = + (long long)tree_pool.max_concurrency * + paged_token_capacity(tree_pool.ddtree_budget + 1); + const long long max_tree_pool_tokens = + ((long long)INT_MAX - PAGED_BLOCK_SIZE - tree_scratch) / + PAGED_BLOCK_SIZE * PAGED_BLOCK_SIZE; + tree_pool.kv_pool_tokens = max_tree_pool_tokens; + TEST_ASSERT(gate_accepts(tree_pool, "qwen35", PlacementBackend::Cuda)); + tree_pool.kv_pool_tokens = max_tree_pool_tokens + PAGED_BLOCK_SIZE; + TEST_ASSERT(!gate_accepts(tree_pool, "qwen35", PlacementBackend::Cuda)); + // The automatic pool is memory-derived, so a logical slot/context product // larger than the physical tensor address space is legal. BackendArgs overflow = paged; diff --git a/server/test/test_paged_attention.cpp b/server/test/test_paged_attention.cpp index 78338f717..3e550c1de 100644 --- a/server/test/test_paged_attention.cpp +++ b/server/test/test_paged_attention.cpp @@ -30,6 +30,13 @@ struct TestCase { bool corrupt_blocks; }; +struct TreeMetadata { + int width; + int scratch_stride; + std::vector parent_ids; + std::vector tree_sizes; +}; + int clamped_seq_len(const TestCase & test_case, int seq) { return std::max( 0, std::min( @@ -51,6 +58,36 @@ bool block_is_valid(int32_t block, int physical_blocks) { return block >= 0 && block < physical_blocks; } +bool tree_visible( + const TreeMetadata & tree, + int tree_seq, + int query_node, + int candidate) { + const int tree_size = tree.tree_sizes[tree_seq]; + if (tree_size < 0 || tree_size > tree.width || + query_node < 0 || query_node >= tree_size || + candidate < 0 || candidate >= tree_size) { + return false; + } + + int current = query_node; + for (int depth = 0; depth < tree_size; ++depth) { + if (current == candidate) { + return true; + } + if (current < 0 || current >= tree_size) { + return false; + } + const int parent = + tree.parent_ids[tree_seq * tree.width + current]; + if (parent == current) { + return false; + } + current = parent; + } + return false; +} + std::vector make_block_table( const TestCase & test_case, int physical_blocks) { @@ -127,7 +164,9 @@ std::vector reference_attention( const std::vector & k, const std::vector & v, const std::vector * active_slot_ids = nullptr, - const std::vector * query_positions = nullptr) { + const std::vector * query_positions = nullptr, + const TreeMetadata * tree = nullptr, + int tree_scratch_base = 0) { std::vector output(q.size(), 0.0f); const float scale = 1.0f / std::sqrt(static_cast(D)); const int q_per_kv = N_HEAD / N_HEAD_KV; @@ -147,50 +186,74 @@ std::vector reference_attention( // cached tokens [0, position]. kv_seq_len = (*query_positions)[seq] + 1; } + const int tree_seq = tree ? seq / tree->width : 0; + const int query_node = tree ? seq % tree->width : -1; + const int tree_size = tree ? tree->tree_sizes[tree_seq] : 0; + if (tree && + (tree_size < 0 || tree_size > tree->width || + query_node >= tree_size)) { + continue; + } + + std::vector physical_rows; + physical_rows.reserve(kv_seq_len + (tree ? tree->width : 0)); + for (int token = 0; token < kv_seq_len; ++token) { + const int block = + block_table[ + physical_seq * test_case.max_blocks + + token / BLOCK_SIZE]; + physical_rows.push_back( + block_is_valid(block, physical_blocks) + ? block * BLOCK_SIZE + token % BLOCK_SIZE + : -1); + } + if (tree) { + for (int candidate = 0; candidate < tree->width; ++candidate) { + physical_rows.push_back( + tree_visible(*tree, tree_seq, query_node, candidate) + ? tree_scratch_base + + physical_seq * tree->scratch_stride + candidate + : -1); + } + } for (int head = 0; head < N_HEAD; ++head) { const int kv_head = head / q_per_kv; const float * q_row = q.data() + (static_cast(head) * n_seq + seq) * D; - std::vector scores(kv_seq_len); + std::vector scores(physical_rows.size(), -INFINITY); float max_score = -INFINITY; - for (int token = 0; token < kv_seq_len; ++token) { - const int block = - block_table[ - physical_seq * test_case.max_blocks + token / BLOCK_SIZE]; - if (!block_is_valid(block, physical_blocks)) { - // Mirrors the kernel: invalid blocks contribute nothing. - scores[token] = -INFINITY; - continue; - } - const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; + for (size_t row = 0; row < physical_rows.size(); ++row) { + const int physical = physical_rows[row]; + if (physical < 0) continue; const float * k_row = k.data() + (static_cast(kv_head) * pool_tokens + physical) * D; float dot = 0.0f; for (int d = 0; d < D; ++d) dot += q_row[d] * k_row[d]; - scores[token] = dot * scale; - max_score = std::max(max_score, scores[token]); + scores[row] = dot * scale; + max_score = std::max(max_score, scores[row]); } float denominator = 0.0f; for (float & score : scores) { + if (!std::isfinite(score)) { + score = 0.0f; + continue; + } score = std::exp(score - max_score); denominator += score; } float * out_row = output.data() + (static_cast(head) * n_seq + seq) * D; - for (int token = 0; token < kv_seq_len; ++token) { - const int block = - block_table[ - physical_seq * test_case.max_blocks + token / BLOCK_SIZE]; - if (!block_is_valid(block, physical_blocks)) continue; - const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; + for (size_t row = 0; row < physical_rows.size(); ++row) { + const int physical = physical_rows[row]; + if (physical < 0 || denominator == 0.0f) continue; const float * v_row = v.data() + (static_cast(kv_head) * pool_tokens + physical) * D; - const float probability = scores[token] / denominator; + const float probability = scores[row] / denominator; for (int d = 0; d < D; ++d) { out_row[d] += probability * v_row[d]; } @@ -205,7 +268,8 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type, const std::vector * active_slot_ids = nullptr, - const std::vector * query_positions = nullptr) { + const std::vector * query_positions = nullptr, + const TreeMetadata * tree = nullptr) { const int physical_n_seq = static_cast(test_case.kv_seq_lens.size()); const int n_seq = active_slot_ids ? static_cast(active_slot_ids->size()) @@ -214,8 +278,23 @@ bool run_case(ggml_backend_t backend, GGML_ASSERT(!query_positions || (active_slot_ids && query_positions->size() == active_slot_ids->size())); + GGML_ASSERT(!tree || (active_slot_ids && !query_positions)); + if (tree) { + GGML_ASSERT(tree->width > 0); + GGML_ASSERT(tree->scratch_stride >= tree->width); + GGML_ASSERT( + tree->parent_ids.size() == + static_cast(tree->width) * tree->tree_sizes.size()); + GGML_ASSERT( + n_seq == + tree->width * static_cast(tree->tree_sizes.size())); + } const int physical_blocks = count_physical_blocks(test_case); - const int pool_tokens = physical_blocks * BLOCK_SIZE; + const int tree_scratch_base = physical_blocks * BLOCK_SIZE; + const int pool_tokens = tree + ? tree_scratch_base + physical_n_seq * tree->scratch_stride + : tree_scratch_base; + GGML_ASSERT(pool_tokens % BLOCK_SIZE == 0); const std::vector block_table = make_block_table(test_case, physical_blocks); @@ -250,13 +329,26 @@ bool run_case(ggml_backend_t backend, positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seq); ggml_set_input(positions); } + ggml_tensor * parents = nullptr; + ggml_tensor * sizes = nullptr; + if (tree) { + parents = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, tree->width, tree->tree_sizes.size()); + sizes = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, tree->tree_sizes.size()); + ggml_set_input(parents); + ggml_set_input(sizes); + } const float scale = 1.0f / std::sqrt(static_cast(D)); const int max_kv_seq_len = *std::max_element( test_case.kv_seq_lens.begin(), test_case.kv_seq_lens.end()); ggml_tensor * output = ggml_paged_attn_ext( ctx, q, k, v, table, kv_seq_lens, active, positions, - scale, BLOCK_SIZE, max_kv_seq_len); + scale, BLOCK_SIZE, max_kv_seq_len, parents, sizes, + tree ? tree->width : 0, + tree ? tree_scratch_base : 0, + tree ? tree->scratch_stride : 0); ggml_set_output(output); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, output); @@ -315,6 +407,14 @@ bool run_case(ggml_backend_t backend, positions, query_positions->data(), 0, query_positions->size() * sizeof((*query_positions)[0])); } + if (tree) { + ggml_backend_tensor_set( + parents, tree->parent_ids.data(), 0, + tree->parent_ids.size() * sizeof(tree->parent_ids[0])); + ggml_backend_tensor_set( + sizes, tree->tree_sizes.data(), 0, + tree->tree_sizes.size() * sizeof(tree->tree_sizes[0])); + } ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; } @@ -327,7 +427,7 @@ bool run_case(ggml_backend_t backend, reference_attention( test_case, block_table, pool_tokens, physical_blocks, q_data, k_reference, v_reference, active_slot_ids, - query_positions); + query_positions, tree, tree ? tree_scratch_base : 0); max_abs_error = 0.0f; for (size_t i = 0; i < actual.size(); ++i) { if (!std::isfinite(actual[i])) { @@ -340,10 +440,11 @@ bool run_case(ggml_backend_t backend, ok = ok && max_abs_error < MAX_ABS_ERROR; } - std::printf("paged attention %-11s K=%-4s V=%-4s active=%s pos=%s max_abs=%.6g %s\n", + std::printf("paged attention %-11s K=%-4s V=%-4s active=%s pos=%s tree=%s max_abs=%.6g %s\n", test_case.name, ggml_type_name(k_type), ggml_type_name(v_type), active_slot_ids ? "yes" : "no", query_positions ? "yes" : "no", + tree ? "yes" : "no", max_abs_error, ok ? "PASS" : "FAIL"); ggml_gallocr_free(allocator); ggml_free(ctx); @@ -373,7 +474,8 @@ bool rejects_unlaunchable_gqa(ggml_backend_t backend) { ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_tensor * output = ggml_paged_attn_ext( ctx, q, k, v, table, kv_seq_lens, nullptr, nullptr, - 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE, 1); + 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE, 1, + nullptr, nullptr, 0, 0, 0); const bool rejected = !ggml_backend_supports_op(backend, output); std::printf("paged attention unlaunchable GQA support %s\n", @@ -404,6 +506,38 @@ void run_paged_attention_case(const TestCase & test_case) { ggml_backend_free(backend); } +void run_tree_case() { + ggml_backend_t backend = ggml_backend_cuda_init(0); + REQUIRE_NOT_NULL(backend); + const TestCase tree_case{"tree", 65, {1025, 17, 257}, false}; + const TreeMetadata tree_metadata{ + 22, 32, + { + -1, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, + 5, 5, 6, 6, 7, 8, 9, 10, 11, 12, 13, + -1, 0, 0, 1, 1, 2, 2, 3, 4, + -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, + }, + {22, 9}, + }; + const std::vector tree_slot_ids{ + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, + 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, + }; + CHECK(run_case(backend, tree_case, GGML_TYPE_F16, GGML_TYPE_F16, + &tree_slot_ids, nullptr, &tree_metadata)); + CHECK(run_case(backend, tree_case, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, + &tree_slot_ids, nullptr, &tree_metadata)); + CHECK(run_case(backend, tree_case, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, + &tree_slot_ids, nullptr, &tree_metadata)); + ggml_backend_free(backend); +} + void run_active_slot_case(const TestCase & test_case, const std::vector & active_slot_ids) { ggml_backend_t backend = ggml_backend_cuda_init(0); @@ -466,6 +600,10 @@ TEST_CASE(PagedAttention, CompactThreeSlotBucketMatchesReference) { }, {2, 1, 0, -1}); } +TEST_CASE(PagedAttention, PackedTreesMatchReference) { + run_tree_case(); +} + TEST_CASE(PagedAttention, RaggedCausalPositionsMatchReference) { // Interleaved query rows from two sequences attend the paged pool causally // through per-row positions. Sequence 1 spans 65 logical blocks, so its diff --git a/server/test/test_paged_kv_pool.cpp b/server/test/test_paged_kv_pool.cpp index e7fd60c07..a91d5775e 100644 --- a/server/test/test_paged_kv_pool.cpp +++ b/server/test/test_paged_kv_pool.cpp @@ -401,6 +401,83 @@ TEST_CASE(PagedKvPoolFixture, invalid_arguments) { })); } + +TEST_CASE(PagedKvPoolFixture, cold_block_roundtrip_and_release) { + PagedKvPool pool(/*physical_block_count=*/4, + /*max_sequences=*/2, /*block_size=*/16); + const auto first = acquire(pool, 1); + CHECK(pool.append(first, 33).status == PagedKvStatus::Ok); + CHECK(equals(sequence(pool, first).block_table, {0, 1, 2})); + CHECK(pool.free_block_count() == 1); + + uint32_t released = 99; + CHECK(pool.page_out_block(first, 1, released) == PagedKvStatus::Ok); + CHECK(released == 1); + CHECK(pool.free_block_count() == 2); + CHECK(sequence(pool, first).block_table[1] == PAGED_KV_COLD_BLOCK); + uint32_t resident = 0; + uint32_t owned = 0; + CHECK(pool.resident_block_count(first, resident) == PagedKvStatus::Ok); + CHECK(resident == 2); + CHECK(pool.owned_block_count(first, owned) == PagedKvStatus::Ok); + CHECK(owned == 3); // logical appended capacity includes the cold block + + uint32_t unchanged = 123; + CHECK(pool.page_out_block(first, 1, unchanged) == + PagedKvStatus::BlockNotResident); + CHECK(unchanged == 123); + CHECK(pool.page_in_block(first, 99, unchanged) == + PagedKvStatus::LogicalBlockOutOfRange); + CHECK(unchanged == 123); + + uint32_t restored = 99; + CHECK(pool.page_in_block(first, 1, restored) == PagedKvStatus::Ok); + CHECK(restored == 1); + CHECK(equals(sequence(pool, first).block_table, {0, 1, 2})); + CHECK(pool.page_in_block(first, 1, unchanged) == + PagedKvStatus::BlockAlreadyResident); + + CHECK(pool.page_out_block(first, 1, released) == PagedKvStatus::Ok); + // release() skips the sentinel and returns only resident blocks. + CHECK(pool.release(first) == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 4); +} + +TEST_CASE(PagedKvPoolFixture, append_remaps_cold_partial_head_atomically) { + PagedKvPool pool(/*physical_block_count=*/2, + /*max_sequences=*/2, /*block_size=*/16); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pool.append(first, 8).status == PagedKvStatus::Ok); + CHECK(pool.append(second, 16).status == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 0); + + uint32_t released = 99; + CHECK(pool.page_out_block(first, 0, released) == PagedKvStatus::Ok); + CHECK(released == 0); + auto appended = pool.append(first, 4); + CHECK(appended.status == PagedKvStatus::Ok); + CHECK(appended.remapped_cold_blocks.size() == 1); + CHECK(appended.remapped_cold_blocks[0].logical_block == 0); + CHECK(appended.remapped_cold_blocks[0].physical_block == 0); + CHECK(appended.write_slots.front().physical_block == 0); + CHECK(appended.write_slots.front().logical_position == 8); + + // A cold head plus a newly-opened block needs two physical allocations. + CHECK(pool.page_out_block(first, 0, released) == PagedKvStatus::Ok); + const auto before = sequence(pool, first); + appended = pool.append(first, 8); // positions 12..19 + CHECK(appended.status == PagedKvStatus::BlocksExhausted); + const auto after = sequence(pool, first); + CHECK(after.kv_seq_len == before.kv_seq_len); + CHECK(after.block_table == before.block_table); + CHECK(after.block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(appended.remapped_cold_blocks.empty()); + + pool.reset(); + CHECK(pool.free_block_count() == 2); +} + TEST_CASE(PagedKvPoolFixture, auto_pool_sizing) { PagedKvAutoBudget budget; budget.free_bytes = 10'000; diff --git a/server/test/test_paged_kv_residency.cpp b/server/test/test_paged_kv_residency.cpp new file mode 100644 index 000000000..d39a5a0a8 --- /dev/null +++ b/server/test/test_paged_kv_residency.cpp @@ -0,0 +1,559 @@ +// Pure-host tests for multi-sequence paged K/V residency. Transfers use a +// deterministic mock device array but preserve async queue/barrier semantics. + +#define GENERATE_UNIT_TEST_MAIN +#include "CppUnitTestFramework.hpp" +#include "common/concurrency/paged_kv_residency.h" + +#include +#include +#include +#include +#include +#include +#include + +using namespace dflash::common; + +namespace { +struct PagedKvResidencyFixture {}; + +PagedKvSequenceHandle acquire(PagedKvPool & pool, uint64_t request) { + PagedKvSequenceHandle handle; + if (pool.acquire(request, handle) != PagedKvStatus::Ok) { + throw std::runtime_error("acquire failed"); + } + return handle; +} + +PagedKvSequenceSnapshot snapshot(PagedKvPool & pool, + PagedKvSequenceHandle handle) { + PagedKvSequenceSnapshot out; + if (pool.sequence(handle, out) != PagedKvStatus::Ok) { + throw std::runtime_error("sequence failed"); + } + return out; +} + +struct MockTransfers { + struct Pending { + std::function apply; + }; + + explicit MockTransfers(uint32_t blocks, size_t bytes) + : block_bytes(bytes), device((size_t)blocks * bytes, 0) {} + + PagedKvResidencyTransferOps callbacks() { + return { + [this](size_t bytes) -> void * { + if (fail_alloc || bytes != block_bytes) return nullptr; + allocations++; + return new uint8_t[bytes]; + }, + [this](void * ptr) { + frees++; + delete[] static_cast(ptr); + }, + [this](PagedKvSequenceHandle, uint32_t, uint32_t physical, + void * host, size_t bytes) { + if (fail_copy_out || bytes != block_bytes) return false; + pending.push_back({[this, physical, host, bytes] { + std::memcpy(host, &device[(size_t)physical * block_bytes], bytes); + }}); + return true; + }, + [this](PagedKvSequenceHandle, uint32_t, uint32_t physical, + const void * host, size_t bytes) { + if (fail_copy_in || bytes != block_bytes) return false; + pending.push_back({[this, physical, host, bytes] { + std::memcpy(&device[(size_t)physical * block_bytes], host, bytes); + }}); + return true; + }, + [this] { + syncs++; + if (fail_sync) return false; + for (Pending & op : pending) op.apply(); + pending.clear(); + return true; + }, + }; + } + + void fill(uint32_t physical, uint8_t value) { + std::fill_n(&device[(size_t)physical * block_bytes], block_bytes, value); + } + + bool block_is(uint32_t physical, uint8_t value) const { + const auto begin = device.begin() + (size_t)physical * block_bytes; + return std::all_of(begin, begin + block_bytes, + [value](uint8_t byte) { return byte == value; }); + } + + size_t block_bytes; + std::vector device; + std::vector pending; + int allocations = 0; + int frees = 0; + int syncs = 0; + bool fail_alloc = false; + bool fail_copy_out = false; + bool fail_copy_in = false; + bool fail_sync = false; +}; + +PagedKvResidencyConfig config(size_t block_bytes, uint32_t budget, + uint32_t sink = 0, uint32_t tail = 0) { + return {block_bytes, budget, sink, tail}; +} + +} // namespace + +TEST_CASE(PagedKvResidencyFixture, page_roundtrip_is_bit_exact_and_barriered) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/1, /*block size=*/4); + MockTransfers io(3, 32); + PagedKvResidencyManager pager(pool, config(32, 3), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 12)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + const auto table = snapshot(pool, handle).block_table; + io.fill(table[0], 0x11); + io.fill(table[1], 0x22); + io.fill(table[2], 0x33); + + CHECK(pager.evict_block(handle, 1) == PagedKvResidencyStatus::Ok); + CHECK(io.syncs == 1); + CHECK(snapshot(pool, handle).block_table[1] == PAGED_KV_COLD_BLOCK); + io.fill(table[1], 0xEE); // recycled device bytes must not affect backing + + CHECK(pager.ensure_resident(handle, {1}) == PagedKvResidencyStatus::Ok); + CHECK(io.syncs == 2); + const uint32_t restored = snapshot(pool, handle).block_table[1]; + CHECK(io.block_is(restored, 0x22)); + CHECK(pager.stats().page_outs == 1); + CHECK(pager.stats().page_ins == 1); + CHECK(pager.stats().resident_blocks == 3); + CHECK(pager.stats().host_bytes == 32); + CHECK(pager.stats().moved_bytes == 64); +} + +TEST_CASE(PagedKvResidencyFixture, fair_share_reclaims_borrowed_pages) { + PagedKvPool pool(/*physical blocks=*/6, /*sequences=*/2, /*block size=*/4); + MockTransfers io(6, 16); + PagedKvResidencyManager pager(pool, config(16, 6), io.callbacks()); + const auto first = acquire(pool, 1); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 20)); // borrows five of six pages while alone + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(first) == 6); + + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(first) == 3); + CHECK(pager.fair_quota(second) == 3); + CHECK(pager.append(second, 12)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().page_outs == 2); + CHECK(pager.stats().resident_blocks == 6); + + uint32_t first_resident = 0; + uint32_t second_resident = 0; + CHECK(pool.resident_block_count(first, first_resident) == PagedKvStatus::Ok); + CHECK(pool.resident_block_count(second, second_resident) == PagedKvStatus::Ok); + CHECK(first_resident == 3); + CHECK(second_resident == 3); +} + +TEST_CASE(PagedKvResidencyFixture, sink_and_tail_are_never_auto_evicted) { + PagedKvPool pool(/*physical blocks=*/5, /*sequences=*/2, /*block size=*/4); + MockTransfers io(5, 16); + PagedKvResidencyManager pager( + pool, config(16, 5, /*sink=*/1, /*tail=*/1), io.callbacks()); + const auto first = acquire(pool, 1); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 16)); // logical blocks 0..3 + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 8)); // needs two pages; only one was free + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + const auto first_table = snapshot(pool, first).block_table; + CHECK(first_table[0] != PAGED_KV_COLD_BLOCK); // sink + CHECK(first_table[3] != PAGED_KV_COLD_BLOCK); // tail + CHECK(first_table[1] == PAGED_KV_COLD_BLOCK || + first_table[2] == PAGED_KV_COLD_BLOCK); + CHECK(pager.evict_block(first, 0) == + PagedKvResidencyStatus::NoEvictableBlock); +} + +TEST_CASE(PagedKvResidencyFixture, append_restores_cold_partial_head) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 2)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + const uint32_t head = snapshot(pool, handle).block_table[0]; + io.fill(head, 0xA5); + CHECK(pager.evict_block(handle, 0) == PagedKvResidencyStatus::Ok); + io.fill(head, 0x00); + + const auto append = pager.append(handle, 1); + CHECK(append); + CHECK(append.pool_result.remapped_cold_blocks.empty()); + CHECK(append.pool_result.write_slots[0].logical_position == 2); + CHECK(io.block_is(append.pool_result.write_slots[0].physical_block, 0xA5)); + CHECK(pager.stats().page_ins == 1); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, scores_drive_reselection) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); + MockTransfers io(3, 16); + PagedKvResidencyManager pager(pool, config(16, 3), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 12)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 2) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().resident_blocks == 1); + const auto peer = acquire(pool, 2); + CHECK(pager.register_sequence(peer) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(handle) == 2); + CHECK(pager.set_scores(handle, {1.0f, 2.0f, 9.0f}) == + PagedKvResidencyStatus::Ok); + CHECK(pager.reselect(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().reselects == 1); + CHECK(pager.is_resident(handle, 2)); + CHECK(pager.is_resident(handle, 1)); + CHECK(!pager.is_resident(handle, 0)); +} + +TEST_CASE(PagedKvResidencyFixture, allocation_and_copy_failures_are_explicit) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 8)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + + io.fail_alloc = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::HostAllocationFailed); + CHECK(pager.is_resident(handle, 0)); + io.fail_alloc = false; + io.fail_copy_out = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(pager.is_resident(handle, 0)); +} + +TEST_CASE(PagedKvResidencyFixture, forget_frees_host_backing_and_stale_is_rejected) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto old_handle = acquire(pool, 1); + CHECK(pager.register_sequence(old_handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(old_handle, 8)); + CHECK(pager.commit_pending_writes(old_handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(old_handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(io.allocations == 1); + CHECK(pager.forget_sequence(old_handle) == PagedKvResidencyStatus::Ok); + CHECK(io.frees == 1); + CHECK(pool.release(old_handle) == PagedKvStatus::Ok); + + const auto replacement = acquire(pool, 2); + CHECK(replacement.generation != old_handle.generation); + CHECK(pager.register_sequence(replacement) == PagedKvResidencyStatus::Ok); + CHECK(pager.touch(old_handle, 0) == PagedKvResidencyStatus::StaleHandle); +} + +TEST_CASE(PagedKvResidencyFixture, + forget_retries_failed_barrier_before_releasing_backing) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + + io.fail_sync = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(pool.free_block_count() == 0); + CHECK(io.frees == 0); + CHECK(pager.forget_sequence(handle) == + PagedKvResidencyStatus::TransferFailed); + CHECK(pool.free_block_count() == 0); + CHECK(io.frees == 0); + + // Teardown retries the quarantined barrier instead of being rejected by + // the manager-wide failure latch. + io.fail_sync = false; + CHECK(pager.forget_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(io.frees == 1); + CHECK(pool.release(handle) == PagedKvStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, + invalid_restore_request_does_not_leak_reservations) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/2, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 4)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + + CHECK(pager.ensure_resident(first, {0, 1}) == + PagedKvResidencyStatus::InvalidArgument); + + // The valid prefix remains evictable, allowing the peer to make progress. + CHECK(pager.append(second, 4)); + CHECK(snapshot(pool, first).block_table[0] == PAGED_KV_COLD_BLOCK); +} + +TEST_CASE(PagedKvResidencyFixture, + staged_writes_are_never_recycled_across_slots) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); + MockTransfers io(3, 16); + PagedKvResidencyManager pager(pool, config(16, 3), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + + CHECK(pager.append(first, 8)); + CHECK(pager.append(second, 4)); + const auto first_staged = snapshot(pool, first).block_table; + const auto second_staged = snapshot(pool, second).block_table; + CHECK(first_staged.size() == 2); + CHECK(second_staged.size() == 1); + + const auto blocked = pager.append(first, 4); + CHECK(blocked.status == PagedKvResidencyStatus::NoEvictableBlock); + CHECK(pager.stats().page_outs == 0); + CHECK(pager.evict_block(second, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::NoEvictableBlock); + CHECK(pager.reselect(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().page_outs == 0); + CHECK(snapshot(pool, first).block_table == first_staged); + CHECK(snapshot(pool, second).block_table == second_staged); + + // Simulate one synchronized packed target graph, then make room. The still + // pending peer page cannot be the victim even though another slot grows. + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const auto grown = pager.append(first, 4); + CHECK(grown); + CHECK(pager.stats().page_outs == 1); + CHECK(snapshot(pool, second).block_table == second_staged); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, + cross_slot_eviction_is_visible_in_full_victim_snapshot) { + PagedKvPool pool(/*physical blocks=*/4, /*sequences=*/2, /*block size=*/4); + MockTransfers io(4, 16); + PagedKvResidencyManager pager(pool, config(16, 4), io.callbacks()); + const auto first = acquire(pool, 1); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 12)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const auto before = snapshot(pool, first).block_table; + + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 8)); + const auto victim = snapshot(pool, first).block_table; + const auto requester = snapshot(pool, second).block_table; + CHECK(victim.size() == before.size()); + CHECK(requester.size() == 2); + + size_t cold = victim.size(); + for (size_t logical = 0; logical < victim.size(); ++logical) { + if (victim[logical] == PAGED_KV_COLD_BLOCK) { + CHECK(cold == victim.size()); + cold = logical; + } + } + CHECK(cold < victim.size()); + CHECK(std::find(requester.begin(), requester.end(), before[cold]) != + requester.end()); + CHECK(pager.stats().page_outs == 1); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, + sixteen_slots_progress_with_adaptive_fair_quotas) { + constexpr uint32_t kSlots = 16; + constexpr uint32_t kPoolBlocks = 96; + PagedKvPool pool(kPoolBlocks, kSlots, /*block size=*/4); + MockTransfers io(kPoolBlocks, 16); + PagedKvResidencyManager pager( + pool, config(16, kPoolBlocks, /*sink=*/1, /*tail=*/4), + io.callbacks()); + + std::vector handles; + handles.reserve(kSlots); + for (uint32_t slot = 0; slot < kSlots; ++slot) { + handles.push_back(acquire(pool, slot + 1)); + CHECK(pager.register_sequence(handles.back()) == + PagedKvResidencyStatus::Ok); + } + for (const auto handle : handles) { + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == + PagedKvResidencyStatus::Ok); + } + for (const auto handle : handles) { + CHECK(pager.append(handle, 28)); // eight logical blocks total + CHECK(pager.commit_pending_writes(handle) == + PagedKvResidencyStatus::Ok); + } + + CHECK(pager.stats().resident_blocks == kPoolBlocks); + CHECK(pager.stats().page_outs == kSlots * 2); + for (const auto handle : handles) { + CHECK(pager.fair_quota(handle) == kPoolBlocks / kSlots); + const auto table = snapshot(pool, handle).block_table; + CHECK(table.size() == 8); + CHECK(table[0] != PAGED_KV_COLD_BLOCK); + for (size_t logical = 4; logical < 8; ++logical) { + CHECK(table[logical] != PAGED_KV_COLD_BLOCK); + } + } +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_out_barrier_retains_the_device_mapping) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 8)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + const auto before = snapshot(pool, handle).block_table; + + io.fail_sync = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(snapshot(pool, handle).block_table == before); + CHECK(pool.free_block_count() == 0); + CHECK(pager.stats().page_outs == 0); +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_in_barrier_quarantines_mapping_until_stream_drains) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + + io.fail_sync = true; + CHECK(pager.ensure_resident(handle, {0}) == + PagedKvResidencyStatus::TransferFailed); + const uint32_t quarantined = snapshot(pool, handle).block_table[0]; + CHECK(quarantined != PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 0); + CHECK(pager.stats().page_ins == 0); + + // No operation may reuse or mutate the quarantined destination while the + // failed stream is not known to be drained. + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(io.pending.size() == 1); + + io.fail_sync = false; + CHECK(pager.synchronize_before_read() == PagedKvResidencyStatus::Ok); + CHECK(snapshot(pool, handle).block_table[0] == quarantined); + CHECK(pool.free_block_count() == 0); + CHECK(pager.stats().page_ins == 1); +} + +TEST_CASE(PagedKvResidencyFixture, + rebalance_evicts_only_blocks_above_fair_quota) { + PagedKvPool pool(/*physical blocks=*/6, /*sequences=*/2, /*block size=*/4); + MockTransfers io(6, 16); + PagedKvResidencyManager pager(pool, config(16, 6), io.callbacks()); + const auto borrower = acquire(pool, 1); + CHECK(pager.register_sequence(borrower) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(borrower, 20)); + CHECK(pager.commit_pending_writes(borrower) == + PagedKvResidencyStatus::Ok); + + const auto peer = acquire(pool, 2); + CHECK(pager.register_sequence(peer) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(borrower) == 3); + CHECK(pager.rebalance() == PagedKvResidencyStatus::Ok); + + uint32_t borrower_resident = 0; + CHECK(pool.resident_block_count(borrower, borrower_resident) == + PagedKvStatus::Ok); + CHECK(borrower_resident == 3); + CHECK(pager.stats().page_outs == 2); +} + +TEST_CASE(PagedKvResidencyFixture, + append_reserves_partial_head_and_new_block_together) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/2, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 2)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(first, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 4)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + const auto grown = pager.append(first, 3); + CHECK(grown); + const auto table = snapshot(pool, first).block_table; + CHECK(table.size() == 2); + CHECK(table[0] != PAGED_KV_COLD_BLOCK); + CHECK(table[1] != PAGED_KV_COLD_BLOCK); + CHECK(snapshot(pool, second).block_table[0] == PAGED_KV_COLD_BLOCK); +} + +TEST_CASE(PagedKvResidencyFixture, + requested_restore_set_is_protected_as_one_batch) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); + MockTransfers io(3, 16); + PagedKvResidencyManager pager(pool, config(16, 3), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 8)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(first, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(first, 1, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 8)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + CHECK(pager.ensure_resident(first, {0, 1}) == + PagedKvResidencyStatus::Ok); + CHECK(pager.is_resident(first, 0)); + CHECK(pager.is_resident(first, 1)); +} diff --git a/server/test/test_qwen_paged_kv_transfer_layout.cpp b/server/test/test_qwen_paged_kv_transfer_layout.cpp new file mode 100644 index 000000000..59c9d0e7a --- /dev/null +++ b/server/test/test_qwen_paged_kv_transfer_layout.cpp @@ -0,0 +1,86 @@ +#define GENERATE_UNIT_TEST_MAIN +#include "CppUnitTestFramework.hpp" +#include "common/concurrency/qwen_paged_kv_transfer.h" + +#include +#include +#include + +using namespace dflash::common; + +namespace { +struct QwenPagedKvTransferLayoutFixture {}; + +QwenPagedKvTensorLayout packed(size_t row_bytes, uint64_t rows, + uint64_t heads) { + return { + row_bytes, + row_bytes, + row_bytes * static_cast(rows), + row_bytes * static_cast(rows) * + static_cast(heads), + rows, + heads, + }; +} +} // namespace + +TEST_CASE(QwenPagedKvTransferLayoutFixture, + packs_mixed_kv_types_and_all_heads) { + const std::vector tensors = { + packed(8, 64, 2), + packed(4, 64, 2), + packed(6, 64, 2), + packed(2, 64, 2), + }; + QwenPagedKvBlockLayout plan; + std::string error; + CHECK(plan_qwen_paged_kv_block_layout( + tensors, /*block_size=*/16, plan, &error)); + CHECK(error.empty()); + CHECK(plan.physical_block_count == 4); + CHECK(plan.tensor_offsets == + std::vector({0, 256, 384, 576})); + CHECK(plan.tensor_head_bytes == + std::vector({128, 64, 96, 32})); + CHECK(plan.block_bytes == 640); +} + +TEST_CASE(QwenPagedKvTransferLayoutFixture, + accepts_row_and_head_padding_without_storing_padding) { + QwenPagedKvTensorLayout tensor; + tensor.row_bytes = 6; + tensor.row_stride = 8; + tensor.head_stride = 520; + tensor.physical_rows = 64; + tensor.heads = 2; + tensor.storage_bytes = 520 + 63 * 8 + 6; + + QwenPagedKvBlockLayout plan; + CHECK(plan_qwen_paged_kv_block_layout( + {tensor}, /*block_size=*/16, plan)); + CHECK(plan.tensor_head_bytes == std::vector({96})); + CHECK(plan.block_bytes == 192); +} + +TEST_CASE(QwenPagedKvTransferLayoutFixture, + rejects_mismatched_rows_short_storage_and_overflow) { + QwenPagedKvBlockLayout plan; + std::string error; + CHECK(!plan_qwen_paged_kv_block_layout( + {packed(8, 64, 2), packed(8, 32, 2)}, 16, plan, &error)); + CHECK(!error.empty()); + + QwenPagedKvTensorLayout short_tensor = packed(8, 64, 2); + short_tensor.storage_bytes--; + CHECK(!plan_qwen_paged_kv_block_layout( + {short_tensor}, 16, plan, &error)); + + QwenPagedKvTensorLayout overflow = packed(8, 64, 2); + overflow.row_bytes = std::numeric_limits::max(); + overflow.row_stride = overflow.row_bytes; + overflow.head_stride = overflow.row_bytes; + overflow.storage_bytes = overflow.row_bytes; + CHECK(!plan_qwen_paged_kv_block_layout( + {overflow}, 16, plan, &error)); +} diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index 6fb5de4d6..05e8a59cb 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -1,5 +1,6 @@ #include "CppUnitTestFramework.hpp" #include "internal.h" +#include "qwen35/graph_builders.h" #include "ggml-backend.h" #include "ggml-cpu.h" @@ -32,6 +33,38 @@ static std::vector get_tensor(const ggml_tensor * tensor) { return values; } +TEST_CASE(RecurrentSnapshotFixture, validates_paged_tree_layout) { + // The packed-tree launch length is logical. KVFlash may keep a much + // smaller physical resident pool, provided every tree scratch slab still + // fits within that pool. + { + ggml_init_params shape_params{}; + shape_params.mem_size = 8 * ggml_tensor_overhead(); + shape_params.no_alloc = true; + ggml_context * shape_ctx = ggml_init(shape_params); + CHECK(shape_ctx != nullptr); + if (shape_ctx) { + TargetCache shape_cache; + shape_cache.n_seq_slots = 2; + shape_cache.paged_block_table = + ggml_new_tensor_2d(shape_ctx, GGML_TYPE_I32, 4, 2); + shape_cache.paged_kv_seq_lens = + ggml_new_tensor_1d(shape_ctx, GGML_TYPE_I32, 2); + shape_cache.attn_k = { + ggml_new_tensor_4d(shape_ctx, GGML_TYPE_F16, 4, 64, 1, 1), + }; + CHECK(dflash::common::detail::validate_target_paged_tree_layout( + shape_cache, 8, 2, 4096, 32, 16)); + CHECK(!dflash::common::detail::validate_target_paged_tree_layout( + shape_cache, 8, 2, 4096, 48, 16)); + CHECK(!dflash::common::detail::validate_target_paged_tree_layout( + shape_cache, 8, 5, 4096, 32, 16)); + ggml_free(shape_ctx); + } + } + +} + TEST_CASE(RecurrentSnapshotFixture, snapshot_and_restore_recurrent_state) { ggml_backend_t backend = ggml_backend_cpu_init(); CHECK(backend != nullptr); diff --git a/server/test/test_seq_batch_plan.cpp b/server/test/test_seq_batch_plan.cpp index 8f3b87178..a771d160b 100644 --- a/server/test/test_seq_batch_plan.cpp +++ b/server/test/test_seq_batch_plan.cpp @@ -119,6 +119,7 @@ int main() { SeqEngine::StepPlan work; work.decode = {{0, 7}}; work.prefills = {{1, 4}}; + CHECK(work.decode[0].allow_speculation); SeqEngine::StepResult good; good.decode.push_back({0, 11, false, {}}); @@ -126,6 +127,35 @@ int main() { 1, SeqEngine::PrefillOutput::Status::advanced, -1, {}}); CHECK(validate_step_result(work, good, 2).empty()); + SeqEngine::StepResult burst = good; + burst.decode[0].committed_tokens = {8, 9, 10}; + burst.decode[0].ddtree_steps = 1; + burst.decode[0].ddtree_accepted_tokens = 3; + burst.decode[0].target_forwards = 1; + CHECK(validate_step_result(work, burst, 2).empty()); + + // A scheduler stop in the committed prefix must hide the remaining burst + // and final pending token. The backend state is discarded at retirement. + std::vector delivered; + const bool delivered_all = consume_decode_output_tokens( + burst.decode[0], [&](int32_t token) { + delivered.push_back(token); + return token != 9; + }); + CHECK(!delivered_all); + CHECK((delivered == std::vector{8, 9})); + + SeqEngine::StepResult malformed_burst = burst; + malformed_burst.decode[0].committed_tokens = {8, -1}; + CHECK(!validate_step_result(work, malformed_burst, 2).empty()); + + SeqEngine::StepPlan speculation_disabled = work; + speculation_disabled.decode[0].allow_speculation = false; + CHECK(!validate_step_result( + speculation_disabled, burst, 2).empty()); + CHECK(validate_step_result( + speculation_disabled, good, 2).empty()); + SeqEngine::StepResult complete = good; complete.prefills[0] = { 1, SeqEngine::PrefillOutput::Status::completed, 12, {}}; @@ -171,6 +201,13 @@ int main() { success_with_error.decode[0].error = "contradictory diagnostic"; CHECK(!validate_step_result(work, success_with_error, 2).empty()); + SeqEngine::StepResult failed_burst = good; + failed_burst.decode[0].failed = true; + failed_burst.decode[0].token = -1; + failed_burst.decode[0].error = "decode failed"; + failed_burst.decode[0].committed_tokens = {8}; + CHECK(!validate_step_result(work, failed_burst, 2).empty()); + SeqEngine::StepResult failed; failed.error = "device compute failed"; CHECK(validate_step_result(work, failed, 2).empty()); diff --git a/server/test/test_seq_engine_contract.cpp b/server/test/test_seq_engine_contract.cpp index be3f031b4..63a62d18a 100644 --- a/server/test/test_seq_engine_contract.cpp +++ b/server/test/test_seq_engine_contract.cpp @@ -23,6 +23,7 @@ struct Faults { bool overconsume_prefill = false; bool drop_second_completion = false; bool retire_leaks = false; + bool burst_when_speculation_disabled = false; }; struct FakeCapabilities { @@ -111,6 +112,10 @@ class FakeSeqEngine final : public SeqEngine { 100 + input.slot + (int32_t)slot.fed.size(), false, {}, }); + if (faults_.burst_when_speculation_disabled && + !input.allow_speculation) { + result.decode.back().committed_tokens.push_back(91); + } } std::vector completed_this_step; @@ -302,6 +307,9 @@ int main() { "omitted an output"}, {"retire-leak", &Faults::retire_leaks, "succeed while a slot is free"}, + {"ignore-speculation-gate", + &Faults::burst_when_speculation_disabled, + "disabled speculation"}, }; for (const Case & test : cases) { diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index a64329e5f..46a5287d7 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -8,6 +8,7 @@ #include "qwen35/concurrency/qwen35_slot_manager.h" #include "host_check.h" +#include #include #include @@ -85,18 +86,22 @@ int main() { mgr.commit_prefill(0); CHECK(mgr.slot(0).cur_pos == 20); - // Decode appends: row allocation + sample_history; cur_pos advances - // separately after the step's compute. + // Decode append stages row allocation and fed-token history; both + // history and cur_pos publish only after the target compute succeeds. auto st = mgr.append_token(0, /*fed_token=*/42); CHECK(st.ok); CHECK(st.position == 20); CHECK(st.physical_row == 20); // tail of the prompt's last block CHECK(st.new_block < 0 && st.new_block_index < 0); CHECK(mgr.slot(0).cur_pos == 20); - CHECK(mgr.slot(0).sample_history.size() == 21 && - mgr.slot(0).sample_history.back() == 42); + CHECK(mgr.slot(0).sample_history.size() == 20); + CHECK(mgr.slot(0).staged_tokens.size() == 1 && + mgr.slot(0).staged_tokens.back() == 42); mgr.commit_step(0); CHECK(mgr.slot(0).cur_pos == 21); + CHECK(mgr.slot(0).sample_history.size() == 21 && + mgr.slot(0).sample_history.back() == 42); + CHECK(mgr.slot(0).staged_tokens.empty()); // Second admission lands in slot 1 with non-identity rows. auto b = admit(mgr, 2, prompt_tokens(20), greedy_sampler()); @@ -333,6 +338,51 @@ int main() { CHECK(pool.free_block_count() == 0); } + // Accepted-path replay stages multiple rows and publishes history and + // logical position only after the target compute succeeds. + { + PagedKvPool pool(8, 1, /*block_size=*/4); + Qwen35SlotManager mgr(pool, /*max_ctx=*/32, + /*speculative_headroom=*/7); + auto a = admit(mgr, 1, prompt_tokens(3), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(mgr.append_prefill(a.slot, 3).ok); + mgr.commit_prefill(a.slot); + + const int32_t accepted[] = {41, 42, 43, 44, 45, 46}; + auto staged = mgr.append_tokens(a.slot, accepted, 6); + CHECK(staged.ok && !staged.busy && staged.count == 6); + CHECK(staged.position == 3 && staged.physical_rows.size() == 6); + CHECK(staged.first_new_block == 1); + CHECK(staged.new_blocks.size() == 2); + CHECK(mgr.slot(a.slot).cur_pos == 3); + CHECK(mgr.slot(a.slot).sample_history.size() == 3); + CHECK(mgr.slot(a.slot).staged_tokens == + std::vector(accepted, accepted + 6)); + CHECK(!mgr.append_token(a.slot, 99).ok); + + // A failed target step never publishes staged history. Scheduler + // retirement releases both materialized rows and the staged host range. + const uint32_t free_while_staged = pool.free_block_count(); + mgr.retire(a.slot); + CHECK(!mgr.is_active(a.slot)); + CHECK(pool.free_block_count() > free_while_staged); + + a = admit(mgr, 2, prompt_tokens(3), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(mgr.append_prefill(a.slot, 3).ok); + mgr.commit_prefill(a.slot); + staged = mgr.append_tokens(a.slot, accepted, 6); + CHECK(staged.ok); + + mgr.commit_step(a.slot); + CHECK(mgr.slot(a.slot).cur_pos == 9); + CHECK(mgr.slot(a.slot).staged_tokens.empty()); + CHECK(mgr.slot(a.slot).sample_history.size() == 9); + CHECK(std::equal(accepted, accepted + 6, + mgr.slot(a.slot).sample_history.end() - 6)); + } + // Context exhaustion: append_token refuses past max_ctx. { PagedKvPool pool(4, 1, /*block_size=*/16); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 8373859dc..2ff471a3d 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -21,6 +21,7 @@ #include "server/http_server.h" #include "server/chat_template.h" #include "common/sampler.h" +#include "common/concurrency/seq_engine.h" #include "common/backend_precision.h" #include "common/backend_ipc.h" #include "common/moe_hybrid_ffn_eval.h" @@ -1947,6 +1948,33 @@ TEST_CASE(ServerUnitFixture, test_stop_sequence_holdback_extends) { TEST_ASSERT(em.accumulated_text().find("suffix") == std::string::npos); } +TEST_CASE(ServerUnitFixture, + test_concurrent_scheduler_burst_stops_at_eos) { + SeqEngine::DecodeOutput burst; + burst.slot = 0; + burst.committed_tokens = {101, 2, 103}; + burst.token = 104; + + std::vector emitted; + int completion_tokens = 0; + const bool consumed_all = consume_decode_output_tokens( + burst, [&](int32_t token) { + emitted.push_back(token); + ++completion_tokens; + // Mirrors scheduler.cpp: advance_slot marks the slot finished on + // EOS and its callback immediately stops the rest of the burst. + return token != 2; + }); + + TEST_ASSERT(!consumed_all); + TEST_ASSERT((emitted == std::vector{101, 2})); + TEST_ASSERT(completion_tokens == 2); + TEST_ASSERT(std::find(emitted.begin(), emitted.end(), 103) == + emitted.end()); + TEST_ASSERT(std::find(emitted.begin(), emitted.end(), 104) == + emitted.end()); +} + // ═══════════════════════════════════════════════════════════════════════ // Prefix cache hash tests (model-free) // ═══════════════════════════════════════════════════════════════════════ @@ -2369,6 +2397,49 @@ TEST_CASE(ServerUnitFixture, test_pflash_config_modes) { TEST_ASSERT(cfg.pflash_mode != ServerConfig::PflashMode::AUTO); } +TEST_CASE(ServerUnitFixture, test_concurrent_pflash_persistent_forces_skip_park) { + ServerConfig cfg; + cfg.pflash_mode = ServerConfig::PflashMode::AUTO; + cfg.draft_residency = DraftResidencyPolicy::Persistent; + cfg.prefix_cache_cap = 0; + + const auto plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/true); + TEST_ASSERT(plan.ok()); + TEST_ASSERT(plan.enabled); + TEST_ASSERT(plan.force_skip_park); +} + +TEST_CASE(ServerUnitFixture, test_concurrent_pflash_rejects_unsafe_residency) { + ServerConfig cfg; + cfg.pflash_mode = ServerConfig::PflashMode::ALWAYS; + cfg.draft_residency = DraftResidencyPolicy::Auto; + cfg.prefix_cache_cap = 0; + + auto plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/true); + TEST_ASSERT(!plan.ok()); + TEST_ASSERT(plan.error.find("--draft-residency persistent") != + std::string::npos); + + cfg.draft_residency = DraftResidencyPolicy::Persistent; + plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/false); + TEST_ASSERT(!plan.ok()); + TEST_ASSERT(plan.error.find("--prefill-drafter") != std::string::npos); +} + +TEST_CASE(ServerUnitFixture, test_concurrent_pflash_rejects_snapshot_caches) { + ServerConfig cfg; + cfg.pflash_mode = ServerConfig::PflashMode::AUTO; + cfg.draft_residency = DraftResidencyPolicy::Persistent; + // Default prefix_cache_cap is intentionally non-zero. + const auto plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/true); + TEST_ASSERT(!plan.ok()); + TEST_ASSERT(plan.error.find("snapshots") != std::string::npos); +} + TEST_CASE(ServerUnitFixture, test_pflash_compress_request_struct) { ModelBackend::CompressRequest req; req.input_ids = {1, 2, 3, 4, 5}; From 3ad0796906fb085156a472e9066a49ed65e5b5d8 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 12 Aug 2026 22:55:43 +0000 Subject: [PATCH 02/42] bench(qwen36): verify concurrent feature matrix --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 113 ++++ harness/benchmarks/concurrency/README.md | 70 +++ .../concurrency/concurrent_benchmark.py | 348 +++++++++++ .../feature_concurrent_benchmark.py | 255 ++++++++ .../concurrency/generate_feature_prompts.py | 57 ++ .../concurrency/generate_ragged_prompts.py | 111 ++++ .../concurrency/record_feature_runtime.py | 119 ++++ .../concurrency/run_qwen36_concurrency.sh | 194 ++++++ .../concurrency/run_qwen36_feature_matrix.sh | 362 ++++++++++++ .../concurrency/summarize_concurrency.py | 219 +++++++ .../concurrency/summarize_feature_matrix.py | 208 +++++++ .../concurrency/test_concurrency_tools.py | 218 +++++++ .../concurrency/test_concurrent_benchmark.py | 131 +++++ .../test_feature_concurrent_benchmark.py | 123 ++++ .../concurrency/test_feature_metadata.py | 153 +++++ .../concurrency/test_feature_tools.py | 556 ++++++++++++++++++ .../concurrency/verify_feature_metrics.py | 322 ++++++++++ .../concurrency/write_feature_metadata.py | 164 ++++++ 18 files changed, 3723 insertions(+) create mode 100644 harness/benchmarks/concurrency/FEATURE_MATRIX.md create mode 100644 harness/benchmarks/concurrency/README.md create mode 100755 harness/benchmarks/concurrency/concurrent_benchmark.py create mode 100755 harness/benchmarks/concurrency/feature_concurrent_benchmark.py create mode 100644 harness/benchmarks/concurrency/generate_feature_prompts.py create mode 100755 harness/benchmarks/concurrency/generate_ragged_prompts.py create mode 100644 harness/benchmarks/concurrency/record_feature_runtime.py create mode 100755 harness/benchmarks/concurrency/run_qwen36_concurrency.sh create mode 100755 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh create mode 100755 harness/benchmarks/concurrency/summarize_concurrency.py create mode 100755 harness/benchmarks/concurrency/summarize_feature_matrix.py create mode 100644 harness/benchmarks/concurrency/test_concurrency_tools.py create mode 100644 harness/benchmarks/concurrency/test_concurrent_benchmark.py create mode 100644 harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py create mode 100644 harness/benchmarks/concurrency/test_feature_metadata.py create mode 100644 harness/benchmarks/concurrency/test_feature_tools.py create mode 100644 harness/benchmarks/concurrency/verify_feature_metrics.py create mode 100644 harness/benchmarks/concurrency/write_feature_metadata.py diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md new file mode 100644 index 000000000..7a1885f18 --- /dev/null +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -0,0 +1,113 @@ +# Qwen3.6 concurrent feature matrix + +`run_qwen36_feature_matrix.sh` extends the PR #596 protocol with feature +ablations for the complete Strix Halo configuration: + +- `ar`: concurrent paged autoregressive control. +- `ddtree`: adds the decode draft, DDTree, and the recorded budget. +- `pflash`: adds auto prefill compression, its drafter, and persistent draft + residency. +- `kvflash`: adds bounded KV residency in auto mode and explicitly supplies + the hashed prefill drafter for relevance-scored page selection; prefill + compression remains off in this ablation. +- `full`: enables DDTree, PFlash, and KVFlash together with both devices on + `hip:0`. +- `llama`: optional external comparison; its binary is required only when this + variant is explicitly requested. + +Run the default bounded C4 screening repeat (seven applicable fresh-server cases): + +```bash +MODEL=/opt/models/Qwen3.6-27B-Q4_K_M.gguf \ +DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q4_k_m.gguf \ +PREFILL_DRAFTER=/opt/models/Qwen3-0.6B-BF16.gguf \ +REPEATS=1 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +On a 128 GiB Strix Halo host, budget roughly 45–90 minutes for this smoke run; +the long-context AR controls dominate and actual time depends on the build. +Every row remains independently selectable through `VARIANTS`. For example: + +```bash +WORKLOADS=short CLIENTS=4 VARIANTS=ar,ddtree MAX_TOKENS=256 REPEATS=1 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +Use five fresh-process, paired repeats for published measurements: + +```bash +WORKLOADS=short,compression CLIENTS=1,4,8,16 \ +VARIANTS=ar,ddtree,pflash,kvflash,full MAX_TOKENS=256 REPEATS=5 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +That full matrix can take roughly 15–30 hours on Strix Halo; keep the generated +case directory so interrupted or suspect rows can be diagnosed rather than +quoted. + +## Activation workloads + +Auto features cannot be validated with the original 400–4,000 word prompts. +The extension adds two deterministic, disjoint 29-prompt cohorts: + +- `compression`: 34K–40K words, chosen after observing 38,130–44,856 + tokens with the development Qwen GGUF tokenizer. +- `kv-pressure`: 12K–18K words, which produced 13,463–20,190 tokens with + that tokenizer against the runner's explicitly recorded 8K pool cap. + +The observed counts are a fixture sanity check, not a claim derived from word +count and not publication evidence. Each row records the model hash; the proof +cross-checks logged raw PFlash input against `usage.prompt_tokens`, requires +it to meet the recorded auto threshold, and uses server-reported effective +tokens plus actual page traffic for KVFlash. + +The bounded default uses `short,compression` at C4. It runs PFlash and the +full configuration only on `compression`; KVFlash can also be selected on +`kv-pressure`. Inapplicable pairs are printed as skips and never appear as +successful rows. AR controls use the same prompts, so feature deltas remain +paired. + +## Fail-closed feature proof + +The server must write one JSON object per completed request with this prefix: + +```text +[concurrency-metrics] {"request_id":"...", ...} +``` + +Required fields are `effective_prompt_tokens`, `ddtree_steps`, +`ddtree_accepted_tokens`, `target_forwards`, `kvflash_page_ins`, +`kvflash_page_outs`, `kvflash_resident_blocks`, `kvflash_reselects`, +`pflash_applied`, `pflash_input_tokens`, and `pflash_output_tokens`. + +The proof tool correlates log objects with measured SSE request IDs and also +checks the log's effective token count against +`usage.timings.effective_prompt_tokens`. A requested feature fails the case +unless: + +- DDTree has positive step and target-forward counts. Acceptance may be zero. +- PFlash reports `pflash_applied=true`, a smaller output prompt, and (in auto + mode) an input token count at or above the recorded activation threshold. +- KVFlash always records an explicit hashed scorer drafter, reports its + startup-observed physical pool and enabled metadata, and has a positive + resident-block count. `kvflash`-only and `kv-pressure` rows must also show page-in or + page-out traffic. For a `full` row, traffic is required only when a + server-reported `effective_prompt_tokens` value exceeds that observed pool + token limit; zero traffic is valid when PFlash compression fits in the pool. + +After health succeeds, the runner fail-closed parses the backend's +`[parallel-kvflash] physical resident pool ...` and `[paged-attention] ...` +startup markers into `runtime_observed`. This distinguishes the actual resident +pool from both the requested auto cap and `--kv-pool-tokens`, which concurrent +KVFlash intentionally does not use to expand VRAM. + +Each case retains the exact shell-escaped command, controlled launch +environment, literal client process arguments and client-script hash, +binary/shared-library/target/draft/PFlash-and-KV-scorer hashes, the ordered +`literal_screenshot_flags` array, all feature values, raw request report, +server log, and `feature-proof.json`. The summary refuses to +include a Lucebox row whose proof is missing or invalid. + + + diff --git a/harness/benchmarks/concurrency/README.md b/harness/benchmarks/concurrency/README.md new file mode 100644 index 000000000..b39db289e --- /dev/null +++ b/harness/benchmarks/concurrency/README.md @@ -0,0 +1,70 @@ +# Qwen3.6 concurrency benchmark + +This protocol measures the serving behavior targeted by packed continuous +prefill and concurrent decode. It is intentionally small: one streaming client, +one fresh-process runner, one deterministic prompt generator, and one summary +script. + +Run a quick screening repeat: + +```bash +MODEL=/path/to/Qwen3.6-27B-Q4_K_M.gguf \ +LUCE_SERVER_BIN=server/build-hip/dflash_server \ +LLAMA_SERVER_BIN=/path/to/llama-server \ +harness/benchmarks/concurrency/run_qwen36_concurrency.sh +``` + +Run a decode-heavy comparison with the same harness: + +```bash +MODEL=/path/to/Qwen3.6-27B-Q4_K_M.gguf \ +LUCE_SERVER_BIN=server/build-hip/dflash_server \ +LLAMA_SERVER_BIN=/path/to/llama-server \ +WORKLOADS=short MAX_TOKENS=256 VARIANTS=luce-k8,llama REPEATS=3 \ +harness/benchmarks/concurrency/run_qwen36_concurrency.sh +``` + +The short ragged prompts keep admission realistic while 256 forced output +tokens make generation dominate the measured window. Use `REPEATS=5` for +publication. Every measured case starts a fresh server and first runs a +discarded warmup at the same concurrency. The variants are: + +- `luce-k8`: packed prefill with up to eight concurrent prefills. +- `luce-k1`: the same binary/configuration with packing width limited to one. +- `llama`: llama.cpp continuous batching with fixed `-b 2048 -ub 512`. + +The 29 generated prompts are disjoint cohorts for C1/C4/C8/C16. C4 and above +contain four substantial length strata while holding the mean target length +constant. The default short, medium, and long profiles target mean lengths of +400, 1,000, and 3,000 words per request. Those generator targets are not token +counts; reports retain the exact server-observed token counts for the selected +model and tokenizer. The client refuses to wrap or reuse a prompt. + +The headline metric is aggregate output goodput: exact server-reported +completion tokens divided by level wall time. It includes queueing, prefill, +and decode and must not be called decode throughput. + +`Output-window tok/s` divides exact completion tokens by the interval from the +earliest observed first output to the final request completion. It removes the +initial all-prefill interval and is decode-facing, but it can still contain +staggered prefill while later requests await their first token. +`Request decode tok/s` is the median per-request estimate +`(completion_tokens - 1) / (end - first_output)`; it assumes the first +observed streaming event accounts for one token. Neither metric is pure kernel +decode throughput. + +`Prompt tok/s to first` is the sum of server-reported prompt tokens divided by +the latest first-token arrival; it is a useful prefill-facing metric but still +includes admission, queueing, and transport. Report TTFT median/max alongside +all throughput metrics. + +The K8-vs-K1 comparison is the causal packing ablation. The K8-vs-llama +comparison is the product comparison. Five paired repeats, the exact command +and hashes recorded in each case, zero failures, and a fixed declared output +length are required before using results in a post. The standard prefill-facing +protocol uses 64 output tokens; the decode-heavy protocol above uses 256. +Variant gains are computed as the median of same-repeat ratios, not as a ratio +of independently aggregated medians. The summarizer rejects mismatched repeat +sets. It also marks whether each variant produced the same ordered output +hashes across at least two repeats; a one-repeat screen reports stability as +`n/a`, and an unstable result is a correctness warning, not a performance win. diff --git a/harness/benchmarks/concurrency/concurrent_benchmark.py b/harness/benchmarks/concurrency/concurrent_benchmark.py new file mode 100755 index 000000000..773260393 --- /dev/null +++ b/harness/benchmarks/concurrency/concurrent_benchmark.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +"""Measure end-to-end output goodput and TTFT under concurrent streaming load.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import sys +import threading +import time +import urllib.request +from pathlib import Path +from typing import Any, Iterable + + +def sha256_text(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def load_prompts(path: Path) -> list[str]: + prompts = [] + for line_no, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + line = raw.strip() + if not line: + continue + if line.startswith("{"): + value = json.loads(line).get("prompt") + if not isinstance(value, str) or not value: + raise ValueError(f"{path}:{line_no}: missing string 'prompt'") + prompts.append(value) + else: + prompts.append(line) + if not prompts: + raise ValueError(f"{path}: no prompts") + return prompts + + +def request_prompts(prompts: list[str], count: int, offset: int) -> list[str]: + if offset < 0: + raise ValueError("--prompt-offset must be >= 0") + if offset + count > len(prompts): + raise ValueError( + f"need prompts [{offset}, {offset + count}), but only " + f"{len(prompts)} were supplied; refusing to reuse prompts" + ) + return prompts[offset:offset + count] + + +def iter_sse_data(lines: Iterable[bytes]) -> Iterable[str]: + data: list[str] = [] + for raw in lines: + line = raw.decode("utf-8", errors="replace").rstrip("\r\n") + if not line: + if data: + yield "\n".join(data) + data.clear() + elif line.startswith("data:"): + data.append(line[5:].lstrip()) + if data: + yield "\n".join(data) + + +def stream_request(args: argparse.Namespace, prompt: str) -> dict[str, Any]: + started = time.perf_counter() + first = None + content: list[str] = [] + reasoning: list[str] = [] + completion_tokens = None + prompt_tokens = None + finish_reason = None + done_received = False + error = None + payload = { + "model": args.model, + "messages": [{"role": "user", "content": prompt}], + "stream": True, + "stream_options": {"include_usage": True}, + "max_tokens": args.max_tokens, + "temperature": args.temperature, + "seed": args.seed, + } + if args.ignore_eos: + payload["ignore_eos"] = True + headers = {"Content-Type": "application/json"} + if args.api_key: + headers["Authorization"] = f"Bearer {args.api_key}" + request = urllib.request.Request( + args.base_url.rstrip("/") + "/chat/completions", + data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=args.timeout) as response: + for data in iter_sse_data(response): + if data == "[DONE]": + done_received = True + break + event = json.loads(data) + usage = event.get("usage") or {} + if isinstance(usage.get("completion_tokens"), int): + completion_tokens = usage["completion_tokens"] + if isinstance(usage.get("prompt_tokens"), int): + prompt_tokens = usage["prompt_tokens"] + for choice in event.get("choices") or []: + if choice.get("finish_reason") is not None: + finish_reason = choice["finish_reason"] + delta = choice.get("delta") or {} + piece = delta.get("content") + thought = delta.get("reasoning_content") + if isinstance(piece, str) and piece: + first = first or time.perf_counter() + content.append(piece) + if isinstance(thought, str) and thought: + first = first or time.perf_counter() + reasoning.append(thought) + except Exception as exc: # preserve partial timing/output for diagnosis + error = f"{type(exc).__name__}: {exc}" + if error is None and not done_received: + error = "ProtocolError: stream ended before [DONE]" + elif error is None and finish_reason is None: + error = "ProtocolError: stream ended without a terminal finish_reason" + ended = time.perf_counter() + output = "".join(content) + reasoning_output = "".join(reasoning) + decode_duration = ended - first if first is not None and ended > first else None + request_decode_tok_s = ( + (completion_tokens - 1) / decode_duration + if isinstance(completion_tokens, int) and completion_tokens > 0 + and decode_duration is not None else None + ) + return { + "t_start": started, "t_first": first, "t_end": ended, + "duration_s": ended - started, + "ttft_s": first - started if first is not None else None, + "decode_duration_s": decode_duration, + "completion_tokens": completion_tokens, "prompt_tokens": prompt_tokens, + "finish_reason": finish_reason, "done_received": done_received, "error": error, + "content_sha256": sha256_text(output), + "reasoning_content_sha256": sha256_text(reasoning_output), + "content_chars": len(output), "reasoning_content_chars": len(reasoning_output), + "request_output_tok_s": ( + completion_tokens / (ended - started) + if completion_tokens is not None and ended > started else None + ), + "request_decode_tok_s": request_decode_tok_s, + } + + +def run_level( + clients: int, args: argparse.Namespace, prompts: list[str], offset: int, +) -> dict[str, Any]: + selected = request_prompts(prompts, clients, offset) + barrier = threading.Barrier(clients) + records: list[dict[str, Any] | None] = [None] * clients + + def worker(index: int) -> None: + barrier.wait() + record = stream_request(args, selected[index]) + record["prompt_index"] = offset + index + record["prompt_sha256"] = sha256_text(selected[index]) + records[index] = record + + threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(clients)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(args.timeout + 30) + completed = [record for record in records if record is not None] + hung = sum(thread.is_alive() for thread in threads) + failures = hung + sum(record["error"] is not None for record in completed) + ok = [record for record in completed if record["error"] is None] + starts = [record["t_start"] for record in completed] + ends = [record["t_end"] for record in completed] + level_start = min(starts) if starts else time.perf_counter() + wall = max(ends) - level_start if ends else 0.0 + for record in completed: + record["start_offset_s"] = record["t_start"] - level_start + + completion_counts = [r["completion_tokens"] for r in ok] + prompt_counts = [r["prompt_tokens"] for r in ok] + completion_complete = bool(ok) and all(isinstance(v, int) for v in completion_counts) + prompt_complete = bool(ok) and all(isinstance(v, int) for v in prompt_counts) + ttfts = [r["ttft_s"] for r in ok if r["ttft_s"] is not None] + first_window = ( + max(r["start_offset_s"] + r["ttft_s"] for r in ok) + if len(ttfts) == len(ok) and ok else None + ) + first_times = [r["t_first"] for r in ok if r["t_first"] is not None] + output_window = ( + max(r["t_end"] for r in ok) - min(first_times) + if len(first_times) == len(ok) and ok else None + ) + request_decode_rates = [ + r["request_decode_tok_s"] for r in ok + if r.get("request_decode_tok_s") is not None + ] + fixed_valid = ( + failures == 0 and len(ok) == clients + and completion_complete + and all(v == args.max_tokens for v in completion_counts) + ) if args.ignore_eos else None + prompt_hashes = [r["prompt_sha256"] for r in ok] + output_hashes = [ + [r["content_sha256"], r["reasoning_content_sha256"]] for r in ok + ] + digest = lambda value: sha256_text(json.dumps(value, separators=(",", ":"))) + return { + "clients": clients, "requests": clients, "requests_ok": len(ok), + "failures": failures, "wall_s": wall, + "start_skew_s": max(starts) - min(starts) if starts else None, + "completion_tokens_total": sum(completion_counts) if completion_complete else None, + "token_count_complete": completion_complete, + "fixed_token_workload_valid": fixed_valid, + "aggregate_tok_s": ( + sum(completion_counts) / wall if completion_complete and wall > 0 else None + ), + "aggregate_metric": "completion_tokens_per_level_wall_second", + "output_window_s": output_window, + "output_window_tok_s": ( + sum(completion_counts) / output_window + if completion_complete and output_window is not None and output_window > 0 + else None + ), + "output_window_metric": "completion_tokens_per_first_output_to_final_completion_second", + "request_decode_tok_s_median": ( + statistics.median(request_decode_rates) + if len(request_decode_rates) == len(ok) and ok else None + ), + "prompt_tokens_total": sum(prompt_counts) if prompt_complete else None, + "prompt_tokens_min": min(prompt_counts) if prompt_complete else None, + "prompt_tokens_max": max(prompt_counts) if prompt_complete else None, + "prompt_tokens_distinct": len(set(prompt_counts)) if prompt_complete else None, + "prompt_token_count_complete": prompt_complete, + "prompt_to_first_token_s": first_window, + "prompt_tokens_per_s_to_first_token": ( + sum(prompt_counts) / first_window + if prompt_complete and first_window is not None and first_window > 0 else None + ), + "ttft_median_s": statistics.median(ttfts) if ttfts else None, + "ttft_max_s": max(ttfts) if ttfts else None, + "selected_prompt_set_sha256": digest(prompt_hashes), + "selected_output_set_sha256": digest(output_hashes), + "requests_detail": completed, + } + + +def fmt(value: Any, spec: str = ".2f") -> str: + return format(value, spec) if isinstance(value, (int, float)) else "n/a" + + +def markdown(report: dict[str, Any]) -> str: + lines = [ + f"# Concurrent benchmark — {report['label']}", "", + "| C | Ok | Output goodput tok/s | Output-window tok/s | " + "Request decode tok/s | Prompt tok/s to first | Prompt range | " + "TTFT median s | TTFT max s | Wall s |", + "| ---: | ---: | ---: | ---: | ---: | ---: | :--- | ---: | ---: | ---: |", + ] + for level in report["levels"]: + lines.append( + f"| {level['clients']} | {level['requests_ok']}/{level['requests']} | " + f"{fmt(level['aggregate_tok_s'])} | " + f"{fmt(level['output_window_tok_s'])} | " + f"{fmt(level['request_decode_tok_s_median'])} | " + f"{fmt(level['prompt_tokens_per_s_to_first_token'])} | " + f"{fmt(level['prompt_tokens_min'], '.0f')}–{fmt(level['prompt_tokens_max'], '.0f')} | " + f"{fmt(level['ttft_median_s'], '.3f')} | {fmt(level['ttft_max_s'], '.3f')} | " + f"{fmt(level['wall_s'])} |" + ) + return "\n".join(lines) + "\n" + + +def level_failed(level: dict[str, Any], ignore_eos: bool) -> bool: + return bool( + level["failures"] + or not level["token_count_complete"] + or not level["prompt_token_count_complete"] + or (ignore_eos and level["fixed_token_workload_valid"] is not True) + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18080/v1") + parser.add_argument("--api-key", default="") + parser.add_argument("--model", default="luce-dflash") + parser.add_argument("--clients", type=int, action="append", dest="client_levels") + parser.add_argument("--prompt-file", type=Path, required=True) + parser.add_argument("--prompt-offset", type=int, default=0) + parser.add_argument("--require-distinct-prompts", action="store_true", + help="Compatibility flag; this client always refuses reuse") + parser.add_argument("--max-tokens", type=int, default=64) + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--seed", type=int, default=1) + parser.add_argument("--ignore-eos", action="store_true") + parser.add_argument("--timeout", type=float, default=1200.0) + parser.add_argument("--cooldown", type=float, default=0.0) + parser.add_argument("--server-metadata-json", type=Path) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--label", default="") + return parser + + +def run(args: argparse.Namespace) -> int: + levels = args.client_levels or [1, 4, 8, 16] + if any(level < 1 for level in levels): + raise ValueError("--clients must be positive") + if args.prompt_offset < 0 or args.max_tokens < 1 or args.timeout <= 0: + raise ValueError("invalid offset, max-tokens, or timeout") + prompts = load_prompts(args.prompt_file) + results = [] + offset = args.prompt_offset + for index, clients in enumerate(levels): + if index and args.cooldown > 0: + time.sleep(args.cooldown) + print(f"[bench] C={clients} max_tokens={args.max_tokens}", flush=True) + results.append(run_level(clients, args, prompts, offset)) + offset += clients + metadata = ( + json.loads(args.server_metadata_json.read_text(encoding="utf-8")) + if args.server_metadata_json else {} + ) + report = { + "schema_version": 2, "label": args.label, "base_url": args.base_url, + "model": args.model, "max_tokens": args.max_tokens, + "temperature": args.temperature, "seed": args.seed, + "ignore_eos": args.ignore_eos, "prompt_offset": args.prompt_offset, + "prompt_file_sha256": hashlib.sha256(args.prompt_file.read_bytes()).hexdigest(), + "server_metadata": metadata, "levels": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(markdown(report), end="") + bad = any(level_failed(level, args.ignore_eos) for level in results) + return 1 if bad else 0 + + +def main() -> int: + try: + return run(build_parser().parse_args()) + except Exception as exc: + print(f"[bench] error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/feature_concurrent_benchmark.py new file mode 100755 index 000000000..bc509c590 --- /dev/null +++ b/harness/benchmarks/concurrency/feature_concurrent_benchmark.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python3 +"""Concurrency client with request IDs and effective-prompt telemetry.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import sys +import time +import urllib.request +from pathlib import Path +from typing import Any + +import concurrent_benchmark as base + + +CLIENT_SCRIPT = Path(__file__).resolve() + + +def client_provenance(argv: list[str] | None = None) -> dict[str, Any]: + """Return the literal process argv and the exact client source digest.""" + process_argv = list(sys.orig_argv if argv is None else argv) + if not process_argv or not all(isinstance(value, str) for value in process_argv): + raise ValueError("client process argv must be a non-empty string array") + return { + "client_argv": process_argv, + "client_script": str(CLIENT_SCRIPT), + "client_script_sha256": hashlib.sha256(CLIENT_SCRIPT.read_bytes()).hexdigest(), + } + + +def stream_request(args: argparse.Namespace, prompt: str) -> dict[str, Any]: + started = time.perf_counter() + first = None + request_id = None + content: list[str] = [] + reasoning: list[str] = [] + completion_tokens = None + prompt_tokens = None + finish_reason = None + done_received = False + timings: dict[str, Any] = {} + wire_metrics: dict[str, Any] = {} + error = None + payload = { + "model": args.model, + "messages": [{"role": "user", "content": prompt}], + "stream": True, + "stream_options": {"include_usage": True}, + "max_tokens": args.max_tokens, + "temperature": args.temperature, + "seed": args.seed, + } + if args.ignore_eos: + payload["ignore_eos"] = True + headers = {"Content-Type": "application/json"} + if args.api_key: + headers["Authorization"] = f"Bearer {args.api_key}" + request = urllib.request.Request( + args.base_url.rstrip("/") + "/chat/completions", + data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=args.timeout) as response: + for data in base.iter_sse_data(response): + if data == "[DONE]": + done_received = True + break + event = json.loads(data) + if isinstance(event.get("id"), str): + request_id = event["id"] + usage = event.get("usage") or {} + if type(usage.get("completion_tokens")) is int: + completion_tokens = usage["completion_tokens"] + if type(usage.get("prompt_tokens")) is int: + prompt_tokens = usage["prompt_tokens"] + if isinstance(usage.get("timings"), dict): + timings = dict(usage["timings"]) + # Log telemetry is authoritative, but retaining a future wire + # copy makes reports forward-compatible without weakening proof. + if isinstance(usage.get("concurrency_metrics"), dict): + wire_metrics = dict(usage["concurrency_metrics"]) + for choice in event.get("choices") or []: + if choice.get("finish_reason") is not None: + finish_reason = choice["finish_reason"] + delta = choice.get("delta") or {} + piece = delta.get("content") + thought = delta.get("reasoning_content") + if isinstance(piece, str) and piece: + first = first or time.perf_counter() + content.append(piece) + if isinstance(thought, str) and thought: + first = first or time.perf_counter() + reasoning.append(thought) + except Exception as exc: # retain partial data for diagnosis + error = f"{type(exc).__name__}: {exc}" + if error is None and not done_received: + error = "ProtocolError: stream ended before [DONE]" + elif error is None and finish_reason is None: + error = "ProtocolError: stream ended without a terminal finish_reason" + ended = time.perf_counter() + output = "".join(content) + reasoning_output = "".join(reasoning) + decode_duration = ended - first if first is not None and ended > first else None + request_decode_tok_s = ( + (completion_tokens - 1) / decode_duration + if type(completion_tokens) is int and completion_tokens > 0 + and decode_duration is not None else None + ) + return { + "request_id": request_id, + "t_start": started, "t_first": first, "t_end": ended, + "duration_s": ended - started, + "ttft_s": first - started if first is not None else None, + "decode_duration_s": decode_duration, + "completion_tokens": completion_tokens, "prompt_tokens": prompt_tokens, + "effective_prompt_tokens": timings.get("effective_prompt_tokens"), + "prefilled_tokens": timings.get("prefilled_tokens"), + "cached_prefix_tokens": timings.get("cached_prefix_tokens"), + "cache_hit": timings.get("cache_hit"), + "server_prefill_ms": timings.get("prefill_ms"), + "server_decode_ms": timings.get("decode_ms"), + "server_decode_tokens_per_sec": timings.get("decode_tokens_per_sec"), + "server_timings": timings, + "wire_concurrency_metrics": wire_metrics, + "finish_reason": finish_reason, "done_received": done_received, "error": error, + "content_sha256": base.sha256_text(output), + "reasoning_content_sha256": base.sha256_text(reasoning_output), + "content_chars": len(output), "reasoning_content_chars": len(reasoning_output), + "request_output_tok_s": ( + completion_tokens / (ended - started) + if completion_tokens is not None and ended > started else None + ), + "request_decode_tok_s": request_decode_tok_s, + } + + +# The base client owns the concurrency/barrier/accounting implementation. Its +# module-global hook is intentional: this process runs one benchmark at a time. +base.stream_request = stream_request + + +def enrich_level(level: dict[str, Any]) -> None: + ok = [r for r in level["requests_detail"] if r.get("error") is None] + effective = [r.get("effective_prompt_tokens") for r in ok] + effective_complete = bool(ok) and all(type(v) is int for v in effective) + request_ids = [r.get("request_id") for r in ok] + request_ids_complete = ( + bool(ok) and all(isinstance(v, str) and v for v in request_ids) + and len(set(request_ids)) == len(request_ids) + ) + level.update({ + "request_ids_complete": request_ids_complete, + "effective_prompt_token_count_complete": effective_complete, + "effective_prompt_tokens_total": sum(effective) if effective_complete else None, + "effective_prompt_tokens_min": min(effective) if effective_complete else None, + "effective_prompt_tokens_max": max(effective) if effective_complete else None, + "effective_to_wire_prompt_ratio": ( + sum(effective) / level["prompt_tokens_total"] + if effective_complete and level.get("prompt_tokens_total") else None + ), + }) + for key in ("server_prefill_ms", "server_decode_ms", "server_decode_tokens_per_sec"): + values = [r.get(key) for r in ok if type(r.get(key)) in (int, float)] + level[f"{key}_median"] = statistics.median(values) if len(values) == len(ok) and ok else None + + +def markdown(report: dict[str, Any]) -> str: + lines = [ + f"# Concurrent feature benchmark — {report['label']}", "", + "| C | Ok | Output goodput tok/s | Output-window tok/s | " + "Request decode tok/s | Wire prompt range | Effective prompt range | " + "Effective/wire | TTFT max s |", + "| ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | ---: |", + ] + for level in report["levels"]: + lines.append( + f"| {level['clients']} | {level['requests_ok']}/{level['requests']} | " + f"{base.fmt(level['aggregate_tok_s'])} | " + f"{base.fmt(level['output_window_tok_s'])} | " + f"{base.fmt(level['request_decode_tok_s_median'])} | " + f"{base.fmt(level['prompt_tokens_min'], '.0f')}–{base.fmt(level['prompt_tokens_max'], '.0f')} | " + f"{base.fmt(level['effective_prompt_tokens_min'], '.0f')}–" + f"{base.fmt(level['effective_prompt_tokens_max'], '.0f')} | " + f"{base.fmt(level['effective_to_wire_prompt_ratio'], '.3f')} | " + f"{base.fmt(level['ttft_max_s'], '.3f')} |" + ) + return "\n".join(lines) + "\n" + + +def build_parser() -> argparse.ArgumentParser: + parser = base.build_parser() + parser.description = __doc__ + parser.add_argument( + "--require-effective-prompt-telemetry", action="store_true", + help="fail when usage.timings.effective_prompt_tokens is absent", + ) + return parser + + +def run(args: argparse.Namespace) -> int: + levels = args.client_levels or [1, 4, 8, 16] + if any(level < 1 for level in levels): + raise ValueError("--clients must be positive") + if args.prompt_offset < 0 or args.max_tokens < 1 or args.timeout <= 0: + raise ValueError("invalid offset, max-tokens, or timeout") + prompts = base.load_prompts(args.prompt_file) + results = [] + offset = args.prompt_offset + for index, clients in enumerate(levels): + if index and args.cooldown > 0: + time.sleep(args.cooldown) + print(f"[bench] C={clients} max_tokens={args.max_tokens}", flush=True) + level = base.run_level(clients, args, prompts, offset) + enrich_level(level) + results.append(level) + offset += clients + metadata = ( + json.loads(args.server_metadata_json.read_text(encoding="utf-8")) + if args.server_metadata_json else {} + ) + report = { + "schema_version": 3, "label": args.label, "base_url": args.base_url, + "model": args.model, "max_tokens": args.max_tokens, + "temperature": args.temperature, "seed": args.seed, + "ignore_eos": args.ignore_eos, "prompt_offset": args.prompt_offset, + "prompt_file_sha256": hashlib.sha256(args.prompt_file.read_bytes()).hexdigest(), + "server_metadata": metadata, "levels": results, + **client_provenance(), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(markdown(report), end="") + bad = any( + base.level_failed(level, args.ignore_eos) + or not level["request_ids_complete"] + or (args.require_effective_prompt_telemetry + and not level["effective_prompt_token_count_complete"]) + for level in results + ) + return 1 if bad else 0 + + +def main() -> int: + try: + return run(build_parser().parse_args()) + except Exception as exc: + print(f"[bench] error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/generate_feature_prompts.py b/harness/benchmarks/concurrency/generate_feature_prompts.py new file mode 100644 index 000000000..d642b138e --- /dev/null +++ b/harness/benchmarks/concurrency/generate_feature_prompts.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +"""Generate deterministic long-context cohorts that force feature activation.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from generate_ragged_prompts import ( + PROFILES as BASE_PROFILES, + build_profile_records, + write_records, +) + + +PROFILES = { + **BASE_PROFILES, + # These word counts were chosen after observing 38,130--44,856 tokens + # with the development Qwen GGUF tokenizer. Word count is never treated as + # activation proof: runtime wire/log telemetry is checked against the + # recorded PFlash threshold. The stable word bank keeps hashes distinct. + "compression": (34000, 36000, 38000, 40000), + # The development Qwen GGUF tokenizer produced 13,463--20,190 tokens here. + # Runtime effective-token and paging telemetry, not this estimate, proves + # KVFlash pressure. + "kv-pressure": (12000, 14000, 16000, 18000), +} + + +def build_records(profile: str) -> list[dict[str, object]]: + activation_target = ( + "pflash-auto" if profile == "compression" + else "kvflash-pressure" if profile == "kv-pressure" + else "none" + ) + return build_profile_records( + profile, PROFILES, + lambda _profile: {"activation_target": activation_target}, + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--profile", choices=sorted(PROFILES), required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + records = build_records(args.profile) + try: + write_records(args.out, records) + except FileExistsError as exc: + parser.error(str(exc)) + print(f"wrote {len(records)} prompts to {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/generate_ragged_prompts.py b/harness/benchmarks/concurrency/generate_ragged_prompts.py new file mode 100755 index 000000000..ca0b180a7 --- /dev/null +++ b/harness/benchmarks/concurrency/generate_ragged_prompts.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""Generate a small deterministic ragged-prompt manifest for concurrency runs.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Callable, Mapping + + +PROFILES = { + "short": (250, 350, 450, 550), + "medium": (650, 850, 1150, 1350), + "long": (2000, 2600, 3400, 4000), +} + +WORD_BANK = ( + "systems engineers compare latency throughput scheduling memory kernels queues " + "batches requests tokens caches pages attention arithmetic bandwidth occupancy " + "profiling measurement fairness reproducibility workloads concurrency admission " + "prefill decoding evidence tradeoffs implementation validation production service" +).split() + + +def prompt_text(profile: str, cohort: str, index: int, target_words: int) -> str: + prefix = ( + f"Ragged benchmark {profile} cohort {cohort} request {index}. " + "Write a structured engineering analysis of the following observations, " + "including assumptions, likely bottlenecks, and a concise conclusion." + ).split() + words = list(prefix) + cursor = (index * 7 + target_words) % len(WORD_BANK) + while len(words) < target_words: + words.append(WORD_BANK[cursor % len(WORD_BANK)]) + cursor += 1 + return " ".join(words[:target_words]) + + +ExtraFields = Callable[[str], Mapping[str, object]] + + +def build_profile_records( + profile: str, + profiles: Mapping[str, tuple[int, ...]], + extra_fields: ExtraFields | None = None, +) -> list[dict[str, object]]: + """Build the standard disjoint C1/C4/C8/C16 cohort layout.""" + strata = profiles[profile] + if not strata: + raise ValueError(f"profile {profile!r} has no length strata") + layout = [ + ("c1", [sum(strata) // len(strata)]), + ("c4", list(strata)), + ("c8", list(strata) * 2), + ("c16", list(strata) * 4), + ] + records: list[dict[str, object]] = [] + for cohort, targets in layout: + for target in targets: + index = len(records) + record: dict[str, object] = { + "id": f"{profile}-{index:02d}", + "cohort": cohort, + "stratum": strata.index(target) if target in strata else "mean", + "target_words": target, + "prompt": prompt_text(profile, cohort, index, target), + } + if extra_fields is not None: + additions = dict(extra_fields(profile)) + overlap = record.keys() & additions.keys() + if overlap: + raise ValueError( + "extra profile fields may not replace standard fields: " + f"{sorted(overlap)}" + ) + record.update(additions) + records.append(record) + return records + + +def build_records(profile: str) -> list[dict[str, object]]: + return build_profile_records(profile, PROFILES) + + +def write_records(path: Path, records: list[dict[str, object]]) -> None: + if path.exists(): + raise FileExistsError(f"refusing to overwrite {path}") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in records), + encoding="utf-8", + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--profile", choices=sorted(PROFILES), required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + records = build_records(args.profile) + try: + write_records(args.out, records) + except FileExistsError as exc: + parser.error(str(exc)) + print(f"wrote {len(records)} prompts to {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/record_feature_runtime.py b/harness/benchmarks/concurrency/record_feature_runtime.py new file mode 100644 index 000000000..5665a4961 --- /dev/null +++ b/harness/benchmarks/concurrency/record_feature_runtime.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Add startup-observed pool dimensions to one feature-case metadata file.""" + +from __future__ import annotations + +import argparse +import json +import re +import tempfile +from pathlib import Path +from typing import Any + + +KVFLASH_POOL_RE = re.compile( + r"\[parallel-kvflash\] physical resident pool (?P\d+) tokens; " + r"logical per-slot cap (?P\d+) across (?P\d+) slots" +) +PAGED_POOL_RE = re.compile( + r"\[paged-attention\] (?P\d+) physical blocks x " + r"(?P\d+) tokens \((?P\d+) pool tokens, " + r"per-sequence max_ctx (?P\d+)\)" +) + + +def _one_consistent(matches: list[dict[str, int]], label: str) -> dict[str, int] | None: + if not matches: + return None + first = matches[0] + if any(row != first for row in matches[1:]): + raise ValueError(f"conflicting {label} startup markers: {matches}") + return first + + +def observe_startup(log_text: str) -> dict[str, Any]: + kvflash = _one_consistent( + [ + {key: int(value) for key, value in match.groupdict().items()} + for match in KVFLASH_POOL_RE.finditer(log_text) + ], + "KVFlash pool", + ) + paged = _one_consistent( + [ + {key: int(value) for key, value in match.groupdict().items()} + for match in PAGED_POOL_RE.finditer(log_text) + ], + "paged pool", + ) + if paged and paged["blocks"] * paged["block_size"] != paged["tokens"]: + raise ValueError("paged-attention startup marker has inconsistent dimensions") + if kvflash and paged: + if kvflash["tokens"] != paged["tokens"]: + raise ValueError("KVFlash and paged-attention startup pool sizes disagree") + if kvflash["max_ctx"] != paged["max_ctx"]: + raise ValueError("KVFlash and paged-attention logical max_ctx values disagree") + + return { + "kvflash_active": kvflash is not None, + "physical_kv_pool_tokens": ( + kvflash["tokens"] if kvflash else paged["tokens"] if paged else None + ), + "physical_kv_pool_blocks": paged["blocks"] if paged else None, + "kv_block_size_tokens": paged["block_size"] if paged else None, + "logical_per_slot_max_ctx": ( + kvflash["max_ctx"] if kvflash else paged["max_ctx"] if paged else None + ), + "configured_slots": kvflash["slots"] if kvflash else None, + "proof_sources": { + "kvflash_pool_startup_marker": kvflash is not None, + "paged_pool_startup_marker": paged is not None, + }, + } + + +def update_metadata(metadata: dict[str, Any], log_text: str) -> dict[str, Any]: + observed = observe_startup(log_text) + feature_config = metadata.get("feature_config") or {} + kvflash_mode = feature_config.get("kvflash") + kvflash_requested = isinstance(kvflash_mode, str) and kvflash_mode not in ( + "", "off", "0", + ) + if kvflash_requested and not observed["kvflash_active"]: + raise ValueError( + "KVFlash metadata is enabled but its physical-pool startup marker is missing" + ) + if observed["physical_kv_pool_tokens"] is None: + raise ValueError("paged physical-pool startup marker is missing") + + result = dict(metadata) + result["schema_version"] = max(3, int(result.get("schema_version", 0))) + result["runtime_observed"] = observed + return result + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--metadata", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + args = parser.parse_args() + + metadata = json.loads(args.metadata.read_text(encoding="utf-8")) + updated = update_metadata( + metadata, + args.server_log.read_text(encoding="utf-8", errors="replace"), + ) + # Replace atomically so a killed run never leaves half-written metadata. + with tempfile.NamedTemporaryFile( + "w", encoding="utf-8", dir=args.metadata.parent, + prefix=f".{args.metadata.name}.", delete=False, + ) as handle: + json.dump(updated, handle, indent=2, sort_keys=True) + handle.write("\n") + temporary = Path(handle.name) + temporary.replace(args.metadata) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh new file mode 100755 index 000000000..b6f76fe95 --- /dev/null +++ b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh @@ -0,0 +1,194 @@ +#!/usr/bin/env bash +# Paired, fresh-process Qwen3.6 concurrency benchmark for Lucebox and llama.cpp. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +REPO="${REPO:-$(cd -- "$SCRIPT_DIR/../../.." && pwd -P)}" +CLIENT="${CLIENT:-$SCRIPT_DIR/concurrent_benchmark.py}" +GENERATOR="${GENERATOR:-$SCRIPT_DIR/generate_ragged_prompts.py}" +SUMMARIZER="${SUMMARIZER:-$SCRIPT_DIR/summarize_concurrency.py}" + +MODEL="${MODEL:-}" +LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" +LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-$(command -v llama-server 2>/dev/null || true)}" +OUT="${OUT:-$REPO/.harness-runs/qwen36-concurrency-$(date -u +%Y%m%dT%H%M%SZ)}" +REPEATS="${REPEATS:-1}" +WORKLOADS="${WORKLOADS:-short,medium,long}" +VARIANTS="${VARIANTS:-luce-k8,luce-k1,llama}" +CLIENTS="${CLIENTS:-1,4,8,16}" +PORT="${PORT:-18114}" +COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-3}" +HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-600}" +MAX_TOKENS="${MAX_TOKENS:-64}" +WARMUP_TOKENS="${WARMUP_TOKENS:-8}" +SLOTS=16 + +usage() { + cat <<'EOF' +Usage: MODEL=/path/model.gguf [REPEATS=5] run_qwen36_concurrency.sh + +Runs fresh-server, same-concurrency warmup + measurement cases for luce-k8, +luce-k1, and llama at C=1/4/8/16. Defaults to one repeat for screening; use at +least five paired repeats for publication. For a decode-heavy comparison, set +WORKLOADS=short MAX_TOKENS=256 VARIANTS=luce-k8,llama. OUT must not already +exist. +EOF +} + +if [[ "${1:-}" == "--help" ]]; then usage; exit 0; fi +if [[ $# -ne 0 ]]; then usage >&2; exit 2; fi +for cmd in python3 curl sha256sum; do command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; }; done +[[ -r "$MODEL" ]] || { echo "set MODEL to a readable GGUF" >&2; exit 2; } +[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +[[ -x "$LLAMA_SERVER_BIN" ]] || { echo "missing llama.cpp server: $LLAMA_SERVER_BIN" >&2; exit 2; } +[[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } +[[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } +ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ + | grep -v '^LUCE_SERVER_BIN=' || true)" +if [[ -n "$ambient_tuning" ]]; then + echo "refusing ambient GPU/backend tuning variables:" >&2 + echo "$ambient_tuning" >&2 + exit 2 +fi +MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" + +IFS=, read -r -a workload_list <<< "$WORKLOADS" +IFS=, read -r -a variant_list <<< "$VARIANTS" +IFS=, read -r -a client_list <<< "$CLIENTS" +declare -A prompt_offsets=([1]=0 [4]=1 [8]=5 [16]=13) +for c in "${client_list[@]}"; do + [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } +done +for v in "${variant_list[@]}"; do + [[ "$v" == luce-k8 || "$v" == luce-k1 || "$v" == llama ]] || { echo "unknown variant $v" >&2; exit 2; } +done + +mkdir -p "$OUT/prompts" +for workload in "${workload_list[@]}"; do + python3 "$GENERATOR" --profile "$workload" --out "$OUT/prompts/$workload.jsonl" +done + +server_pid="" +stop_server() { + if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then + kill "$server_pid" 2>/dev/null || true + for _ in $(seq 1 30); do + kill -0 "$server_pid" 2>/dev/null || break + sleep 1 + done + kill -9 "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true + fi + server_pid="" +} +trap stop_server EXIT INT TERM + +wait_health() { + local deadline=$((SECONDS + HEALTH_TIMEOUT_SECONDS)) + while (( SECONDS < deadline )); do + kill -0 "$server_pid" 2>/dev/null || return 1 + curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && return 0 + sleep 1 + done + return 1 +} + +write_metadata() { + local path="$1" variant="$2" workload="$3" clients="$4" repeat="$5" binary="$6" max_prefills="$7" command_file="$8" + python3 -c 'import hashlib,json,pathlib,subprocess,sys +p,variant,workload,clients,repeat,binary,max_prefills,cmd_file,model_sha,prompts,repo=sys.argv[1:] +digest=lambda x: hashlib.sha256(pathlib.Path(x).read_bytes()).hexdigest() +libs={} +for line in subprocess.run(["ldd",binary],text=True,capture_output=True).stdout.splitlines(): + fields=line.replace("=>"," ").split() + paths=[x for x in fields if x.startswith("/") and pathlib.Path(x).is_file()] + for lib in paths: libs[str(pathlib.Path(lib).resolve())]=digest(lib) +lucebox_git_head=subprocess.run(["git","-C",repo,"rev-parse","HEAD"],text=True,capture_output=True).stdout.strip() or None +server_version=None +if variant == "llama": + version=subprocess.run([binary,"--version"],text=True,capture_output=True,timeout=30) + server_version="\n".join(x.strip() for x in (version.stdout,version.stderr) if x.strip()) or None + if version.returncode != 0 or server_version is None: + raise RuntimeError(f"cannot identify llama.cpp source version from {binary} --version") +obj={"variant":variant,"workload":workload,"clients":int(clients),"repeat":int(repeat), + "max_concurrent_prefills":int(max_prefills),"server_binary":str(pathlib.Path(binary).resolve()), + "server_binary_sha256":digest(binary),"model_sha256":model_sha, + "prompt_file_sha256":digest(prompts),"server_command":pathlib.Path(cmd_file).read_text().strip(), + "resolved_shared_library_sha256":libs, + "lucebox_git_head":lucebox_git_head if variant != "llama" else None, + "server_version":server_version} +pathlib.Path(p).write_text(json.dumps(obj,indent=2,sort_keys=True)+"\n")' \ + "$path" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$command_file" "$MODEL_SHA256" "$OUT/prompts/$workload.jsonl" "$REPO" +} + +run_case() { + local repeat="$1" workload="$2" clients="$3" variant="$4" + local max_ctx timeout capacity max_prefills binary model_id + if [[ "$workload" == long ]]; then + max_ctx=8192; timeout=1800 + else + max_ctx=4096; timeout=1200 + fi + capacity=$((SLOTS * max_ctx)) + local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" + mkdir -p "$case_dir" + local -a command + if [[ "$variant" == llama ]]; then + binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 + command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" + -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on + -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") + else + binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce + [[ "$variant" == luce-k8 ]] && max_prefills=8 || max_prefills=1 + command=("$binary" "$MODEL" --target-device hip:0 --paged-attention + --max-concurrency "$SLOTS" --kv-pool-tokens "$capacity" --max-ctx "$max_ctx" + --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 + --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 + --host 127.0.0.1 --port "$PORT" --model-name "$model_id") + fi + printf '%q ' "${command[@]}" > "$case_dir/server-command.txt"; printf '\n' >> "$case_dir/server-command.txt" + write_metadata "$case_dir/server-metadata.json" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$case_dir/server-command.txt" + + echo "[run] $workload C=$clients repeat=$repeat variant=$variant" + if [[ "$variant" == llama ]]; then + "${command[@]}" > "$case_dir/server.log" 2>&1 & + else + env DFLASH_MIN_TOKENS="$WARMUP_TOKENS" DFLASH_MAX_CONCURRENT_PREFILLS="$max_prefills" \ + "${command[@]}" > "$case_dir/server.log" 2>&1 & + fi + server_pid=$! + if ! wait_health; then tail -n 80 "$case_dir/server.log" >&2 || true; return 1; fi + + local offset="${prompt_offsets[$clients]}" prompts="$OUT/prompts/$workload.jsonl" + python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ + --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" \ + --require-distinct-prompts --max-tokens "$WARMUP_TOKENS" --temperature 0 \ + --ignore-eos --timeout "$timeout" --cooldown 0 --out "$case_dir/warmup.json" \ + --label "$variant $workload C=$clients warmup" > "$case_dir/warmup.txt" + python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ + --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" \ + --require-distinct-prompts --max-tokens "$MAX_TOKENS" --temperature 0 \ + --ignore-eos --timeout "$timeout" --cooldown 0 \ + --server-metadata-json "$case_dir/server-metadata.json" --out "$case_dir/bench.json" \ + --label "$variant $workload C=$clients repeat=$repeat" | tee "$case_dir/bench.txt" + stop_server + sleep "$COOLDOWN_SECONDS" +} + +for ((repeat=1; repeat<=REPEATS; repeat++)); do + for workload in "${workload_list[@]}"; do + for c_index in "${!client_list[@]}"; do + clients="${client_list[$c_index]}" + # Rotate start variant by case so one engine is not always hot or cold. + shift_by=$(((repeat + c_index) % ${#variant_list[@]})) + for ((i=0; i<${#variant_list[@]}; i++)); do + variant="${variant_list[$(((i + shift_by) % ${#variant_list[@]}))]}" + run_case "$repeat" "$workload" "$clients" "$variant" + done + done + done +done + +python3 "$SUMMARIZER" "$OUT" --out "$OUT/summary.md" +echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh new file mode 100755 index 000000000..e4b23bca1 --- /dev/null +++ b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh @@ -0,0 +1,362 @@ +#!/usr/bin/env bash +# Fresh-process Qwen3.6 concurrent feature ablations with fail-closed proof. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +REPO="${REPO:-$(cd -- "$SCRIPT_DIR/../../.." && pwd -P)}" +CLIENT="${CLIENT:-$SCRIPT_DIR/feature_concurrent_benchmark.py}" +GENERATOR="${GENERATOR:-$SCRIPT_DIR/generate_feature_prompts.py}" +SUMMARIZER="${SUMMARIZER:-$SCRIPT_DIR/summarize_feature_matrix.py}" +PROOF_TOOL="${PROOF_TOOL:-$SCRIPT_DIR/verify_feature_metrics.py}" +METADATA_TOOL="${METADATA_TOOL:-$SCRIPT_DIR/write_feature_metadata.py}" +RUNTIME_METADATA_TOOL="${RUNTIME_METADATA_TOOL:-$SCRIPT_DIR/record_feature_runtime.py}" + +MODEL="${MODEL:-}" +DRAFT_MODEL="${DRAFT_MODEL:-}" +PREFILL_DRAFTER="${PREFILL_DRAFTER:-}" +LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" +LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-$(command -v llama-server 2>/dev/null || true)}" +OUT="${OUT:-$REPO/.harness-runs/qwen36-feature-matrix-$(date -u +%Y%m%dT%H%M%SZ)}" +REPEATS="${REPEATS:-1}" +WORKLOADS="${WORKLOADS:-short,compression}" +VARIANTS="${VARIANTS:-ar,ddtree,pflash,kvflash,full}" +CLIENTS="${CLIENTS:-4}" +PORT="${PORT:-18114}" +COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-3}" +HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-900}" +MAX_TOKENS="${MAX_TOKENS:-64}" +WARMUP_TOKENS="${WARMUP_TOKENS:-8}" +SLOTS="${SLOTS:-16}" +MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" + +# The requested Strix Halo configuration. Every value is serialized into case +# metadata; no performance-affecting DFLASH variable is inherited implicitly. +TARGET_DEVICE="${TARGET_DEVICE:-hip:0}" +DRAFT_DEVICE="${DRAFT_DEVICE:-hip:0}" +DDTREE_BUDGET="${DDTREE_BUDGET:-22}" +DRAFT_RESIDENCY="${DRAFT_RESIDENCY:-persistent}" +PREFILL_COMPRESSION="${PREFILL_COMPRESSION:-auto}" +PREFILL_THRESHOLD="${PREFILL_THRESHOLD:-32000}" +PREFILL_KEEP_RATIO="${PREFILL_KEEP_RATIO:-0.05}" +KVFLASH_MODE="${KVFLASH_MODE:-auto}" +KVFLASH_MAX_POOL_TOKENS="${KVFLASH_MAX_POOL_TOKENS:-8192}" + +usage() { + cat <<'EOF' +Usage: + MODEL=/path/Qwen3.6-27B-Q4_K_M.gguf \ + DRAFT_MODEL=/path/dflash-draft-3.6-q4_k_m.gguf \ + PREFILL_DRAFTER=/path/Qwen3-0.6B-BF16.gguf \ + harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +The default is a bounded C4 smoke matrix with independently selectable +ar, ddtree, pflash, kvflash, and full rows. The full row is the requested +Strix Halo configuration: target/draft hip:0, DDTree +budget 22, persistent PFlash auto, and KVFlash auto. The long-context profiles +are intended to cross the recorded 32K PFlash and 8K KV-residency thresholds; +word count is not treated as proof. Per-request wire/log token counts and +activation telemetry fail the case if an "auto" feature did not execute. + +llama is optional: include it explicitly with VARIANTS=ar,ddtree,llama and set +LLAMA_SERVER_BIN. For publication, set CLIENTS=1,4,8,16 and REPEATS=5. +OUT must not already exist. +EOF +} + +if [[ "${1:-}" == "--help" ]]; then usage; exit 0; fi +if [[ $# -ne 0 ]]; then usage >&2; exit 2; fi +for cmd in python3 curl sha256sum awk; do + command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; } +done +[[ -r "$MODEL" ]] || { echo "set MODEL to a readable target GGUF" >&2; exit 2; } +[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +[[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } +[[ "$SLOTS" =~ ^[1-9][0-9]*$ ]] || { echo "SLOTS must be positive" >&2; exit 2; } +[[ "$MAX_CONCURRENT_PREFILLS" =~ ^[1-9][0-9]*$ ]] || { echo "MAX_CONCURRENT_PREFILLS must be positive" >&2; exit 2; } +[[ "$DDTREE_BUDGET" =~ ^[1-9][0-9]*$ ]] || { echo "DDTREE_BUDGET must be positive" >&2; exit 2; } +[[ "$PREFILL_THRESHOLD" =~ ^[1-9][0-9]*$ ]] || { echo "PREFILL_THRESHOLD must be positive" >&2; exit 2; } +[[ "$KVFLASH_MAX_POOL_TOKENS" =~ ^[1-9][0-9]*$ ]] || { echo "KVFLASH_MAX_POOL_TOKENS must be positive" >&2; exit 2; } +[[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } + +ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ + | grep -v '^LUCE_SERVER_BIN=' || true)" +if [[ -n "$ambient_tuning" ]]; then + echo "refusing ambient GPU/backend tuning variables:" >&2 + echo "$ambient_tuning" >&2 + exit 2 +fi + +IFS=, read -r -a workload_list <<< "$WORKLOADS" +IFS=, read -r -a variant_list <<< "$VARIANTS" +IFS=, read -r -a client_list <<< "$CLIENTS" +reject_duplicates() { + local list_name="$1" value + shift + local -A seen=() + for value in "$@"; do + if [[ -n "${seen[$value]+yes}" ]]; then + echo "$list_name contains duplicate entry: $value" >&2 + return 1 + fi + seen["$value"]=1 + done +} +reject_duplicates CLIENTS "${client_list[@]}" || exit 2 +reject_duplicates VARIANTS "${variant_list[@]}" || exit 2 +declare -A prompt_offsets=([1]=0 [4]=1 [8]=5 [16]=13) +for c in "${client_list[@]}"; do + [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } + (( c <= SLOTS )) || { echo "CLIENTS=$c exceeds SLOTS=$SLOTS" >&2; exit 2; } +done +for v in "${variant_list[@]}"; do + case "$v" in + ar|ddtree|pflash|kvflash|full|llama) ;; + *) echo "unknown variant $v" >&2; exit 2 ;; + esac +done + +contains_variant() { + local needle="$1" value + for value in "${variant_list[@]}"; do [[ "$value" == "$needle" ]] && return 0; done + return 1 +} +if contains_variant ddtree || contains_variant full; then + [[ -r "$DRAFT_MODEL" ]] || { echo "DRAFT_MODEL is required for DDTree/full" >&2; exit 2; } +fi +if contains_variant pflash || contains_variant kvflash || contains_variant full; then + [[ -r "$PREFILL_DRAFTER" ]] || { + echo "PREFILL_DRAFTER is required for PFlash and drafter-scored KVFlash" >&2 + exit 2 + } +fi +if contains_variant llama; then + [[ -x "$LLAMA_SERVER_BIN" ]] || { echo "llama requested but LLAMA_SERVER_BIN is missing" >&2; exit 2; } +fi + +MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" +DRAFT_MODEL_SHA256="" +PREFILL_DRAFTER_SHA256="" +[[ -n "$DRAFT_MODEL" ]] && DRAFT_MODEL_SHA256="$(sha256sum "$DRAFT_MODEL" | awk '{print $1}')" +[[ -n "$PREFILL_DRAFTER" ]] && PREFILL_DRAFTER_SHA256="$(sha256sum "$PREFILL_DRAFTER" | awk '{print $1}')" + +mkdir -p "$OUT/prompts" +for workload in "${workload_list[@]}"; do + python3 "$GENERATOR" --profile "$workload" --out "$OUT/prompts/$workload.jsonl" +done + +server_pid="" +stop_server() { + if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then + kill "$server_pid" 2>/dev/null || true + for _ in $(seq 1 30); do + kill -0 "$server_pid" 2>/dev/null || break + sleep 1 + done + kill -9 "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true + fi + server_pid="" +} +trap stop_server EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +wait_health() { + local deadline=$((SECONDS + HEALTH_TIMEOUT_SECONDS)) + while (( SECONDS < deadline )); do + kill -0 "$server_pid" 2>/dev/null || return 1 + curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && return 0 + sleep 1 + done + return 1 +} + +port_is_available() { + python3 - "$PORT" <<'PY' +import socket +import sys + +port = int(sys.argv[1]) +with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind(("127.0.0.1", port)) + except OSError as exc: + print(f"PORT {port} is unavailable: {exc}", file=sys.stderr) + sys.exit(1) +PY +} + +# Auto modes cannot be proven on sub-threshold inputs. These skips are listed +# explicitly and are not emitted as successful benchmark rows. +case_applicable() { + local variant="$1" workload="$2" + case "$variant" in + pflash|full) [[ "$workload" == compression ]] ;; + kvflash) [[ "$workload" == compression || "$workload" == kv-pressure ]] ;; + *) return 0 ;; + esac +} + +workload_limits() { + case "$1" in + short) echo "4096 1200" ;; + medium) echo "8192 1800" ;; + long) echo "16384 2400" ;; + kv-pressure) echo "32768 3600" ;; + compression) echo "65536 5400" ;; + *) echo "unknown workload $1" >&2; return 1 ;; + esac +} + +run_case() { + local repeat="$1" workload="$2" clients="$3" variant="$4" + local max_ctx timeout + read -r max_ctx timeout <<< "$(workload_limits "$workload")" + local capacity=$((SLOTS * max_ctx)) + local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" + mkdir -p "$case_dir" + + local binary model_id max_prefills + local -a command launch_env metadata expected + if [[ "$variant" == llama ]]; then + binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 + command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" + -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on + -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") + launch_env=() + else + binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce; max_prefills="$MAX_CONCURRENT_PREFILLS" + command=("$binary" "$MODEL" --target-device "$TARGET_DEVICE" --paged-attention + --max-concurrency "$SLOTS" --kv-pool-tokens "$capacity" --max-ctx "$max_ctx" + --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 + --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 + --host 127.0.0.1 --port "$PORT" --model-name "$model_id") + launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills") + if [[ "$variant" == ddtree || "$variant" == full ]]; then + command+=(--draft "$DRAFT_MODEL" --draft-device "$DRAFT_DEVICE" + --ddtree --ddtree-budget "$DDTREE_BUDGET") + expected+=(--expect ddtree) + fi + if [[ "$variant" == pflash || "$variant" == full ]]; then + if [[ "$variant" == pflash ]]; then command+=(--draft-device "$DRAFT_DEVICE"); fi + command+=(--prefill-compression "$PREFILL_COMPRESSION" + --prefill-threshold "$PREFILL_THRESHOLD" + --prefill-keep-ratio "$PREFILL_KEEP_RATIO" + --prefill-drafter "$PREFILL_DRAFTER" + --draft-residency "$DRAFT_RESIDENCY") + expected+=(--expect pflash) + fi + if [[ "$variant" == kvflash || "$variant" == full ]]; then + # --prefill-drafter also selects the default drafter-scored KVFlash + # residency policy; it does not enable PFlash when compression is off. + if [[ "$variant" == kvflash ]]; then + command+=(--prefill-drafter "$PREFILL_DRAFTER") + fi + command+=(--kvflash "$KVFLASH_MODE") + launch_env+=("DFLASH_KVFLASH_MAX_POOL=$KVFLASH_MAX_POOL_TOKENS") + expected+=(--expect kvflash) + fi + fi + + if ((${#launch_env[@]})); then + printf 'env ' > "$case_dir/server-command.txt" + printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$case_dir/server-command.txt" + else + printf '%q ' "${command[@]}" > "$case_dir/server-command.txt" + fi + printf '\n' >> "$case_dir/server-command.txt" + + metadata=(python3 "$METADATA_TOOL" --out "$case_dir/server-metadata.json" + --variant "$variant" --workload "$workload" --clients "$clients" --repeat "$repeat" + --binary "$binary" --model "$MODEL" --model-sha256 "$MODEL_SHA256" + --prompt-file "$OUT/prompts/$workload.jsonl" + --command-file "$case_dir/server-command.txt" --repo "$REPO" + --max-concurrent-prefills "$max_prefills") + if [[ "$variant" != llama ]]; then + metadata+=(--target-device "$TARGET_DEVICE") + local item + for item in "${launch_env[@]}"; do metadata+=(--launch-env "$item"); done + fi + if [[ "$variant" == ddtree || "$variant" == full ]]; then + metadata+=(--draft-device "$DRAFT_DEVICE" --draft-model "$DRAFT_MODEL" + --draft-model-sha256 "$DRAFT_MODEL_SHA256" --ddtree --ddtree-budget "$DDTREE_BUDGET") + fi + if [[ "$variant" == pflash || "$variant" == full ]]; then + metadata+=(--draft-device "$DRAFT_DEVICE" --prefill-compression "$PREFILL_COMPRESSION" + --prefill-threshold "$PREFILL_THRESHOLD" --prefill-keep-ratio "$PREFILL_KEEP_RATIO" + --prefill-drafter "$PREFILL_DRAFTER" + --prefill-drafter-sha256 "$PREFILL_DRAFTER_SHA256" + --draft-residency "$DRAFT_RESIDENCY") + fi + if [[ "$variant" == kvflash || "$variant" == full ]]; then + metadata+=(--kvflash "$KVFLASH_MODE" + --kvflash-max-pool-tokens "$KVFLASH_MAX_POOL_TOKENS" + --kvflash-scorer-drafter "$PREFILL_DRAFTER" + --kvflash-scorer-drafter-sha256 "$PREFILL_DRAFTER_SHA256") + if [[ "$variant" == kvflash ]]; then + # Record the literal server flag too, while keeping compression off. + metadata+=(--prefill-drafter "$PREFILL_DRAFTER" + --prefill-drafter-sha256 "$PREFILL_DRAFTER_SHA256") + fi + fi + "${metadata[@]}" + + echo "[run] $workload C=$clients repeat=$repeat variant=$variant" + port_is_available || return 1 + if ((${#launch_env[@]})); then + env "${launch_env[@]}" "${command[@]}" > "$case_dir/server.log" 2>&1 & + else + "${command[@]}" > "$case_dir/server.log" 2>&1 & + fi + server_pid=$! + if ! wait_health; then tail -n 120 "$case_dir/server.log" >&2 || true; return 1; fi + if [[ "$variant" != llama ]]; then + python3 "$RUNTIME_METADATA_TOOL" --metadata "$case_dir/server-metadata.json" \ + --server-log "$case_dir/server.log" + fi + + local offset="${prompt_offsets[$clients]}" prompts="$OUT/prompts/$workload.jsonl" + local -a common_client=(--base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" + --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" + --require-distinct-prompts --temperature 0 --ignore-eos --timeout "$timeout" --cooldown 0) + local -a telemetry_arg=() + [[ "$variant" != llama ]] && telemetry_arg+=(--require-effective-prompt-telemetry) + python3 "$CLIENT" "${common_client[@]}" --max-tokens "$WARMUP_TOKENS" + "${telemetry_arg[@]}" --out "$case_dir/warmup.json" + --label "$variant $workload C=$clients warmup" > "$case_dir/warmup.txt" + python3 "$CLIENT" "${common_client[@]}" --max-tokens "$MAX_TOKENS" + "${telemetry_arg[@]}" --server-metadata-json "$case_dir/server-metadata.json" + --out "$case_dir/bench.json" --label "$variant $workload C=$clients repeat=$repeat" \ + | tee "$case_dir/bench.txt" + stop_server + + if [[ "$variant" != llama ]]; then + python3 "$PROOF_TOOL" --bench "$case_dir/bench.json" --server-log "$case_dir/server.log" + "${expected[@]}" --out "$case_dir/feature-proof.json" + fi + sleep "$COOLDOWN_SECONDS" +} + +active_cases=0 +for ((repeat=1; repeat<=REPEATS; repeat++)); do + for workload in "${workload_list[@]}"; do + for c_index in "${!client_list[@]}"; do + clients="${client_list[$c_index]}" + shift_by=$(((repeat + c_index) % ${#variant_list[@]})) + for ((i=0; i<${#variant_list[@]}; i++)); do + variant="${variant_list[$(((i + shift_by) % ${#variant_list[@]}))]}" + if ! case_applicable "$variant" "$workload"; then + echo "[skip] $variant requires an activation workload; workload=$workload" + continue + fi + active_cases=$((active_cases + 1)) + run_case "$repeat" "$workload" "$clients" "$variant" + done + done + done +done +(( active_cases > 0 )) || { echo "no applicable benchmark cases" >&2; exit 2; } + +python3 "$SUMMARIZER" "$OUT" --out "$OUT/summary.md" +echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/summarize_concurrency.py b/harness/benchmarks/concurrency/summarize_concurrency.py new file mode 100755 index 000000000..6f5130f41 --- /dev/null +++ b/harness/benchmarks/concurrency/summarize_concurrency.py @@ -0,0 +1,219 @@ +#!/usr/bin/env python3 +"""Summarize paired Lucebox/llama.cpp concurrency benchmark reports.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections import defaultdict +from pathlib import Path + + +def load_reports(root: Path) -> list[dict]: + reports = [] + for path in sorted(root.rglob("bench.json")): + report = json.loads(path.read_text(encoding="utf-8")) + meta = report.get("server_metadata") or {} + if len(report.get("levels", [])) != 1: + raise ValueError(f"{path}: expected exactly one client level") + level = report["levels"][0] + if ( + level.get("failures") + or not level.get("token_count_complete") + or not level.get("prompt_token_count_complete") + ): + raise ValueError(f"{path}: failed or incomplete token accounting") + if report.get("ignore_eos") and level.get("fixed_token_workload_valid") is not True: + raise ValueError(f"{path}: fixed-token validation failed") + reports.append({"path": path, "report": report, "level": level, "meta": meta}) + if not reports: + raise ValueError(f"{root}: no bench.json files found") + return reports + + +def median(values: list[float]) -> float: + return statistics.median(values) + + +def complete_median(values: list[float | None]) -> float | None: + """Return a median only when every repeat measured the metric.""" + if not values or any(value is None for value in values): + return None + return median([value for value in values if value is not None]) + + +def run_signature(item: dict) -> tuple[object, ...]: + report, meta = item["report"], item["meta"] + max_tokens = report.get("max_tokens") + ignore_eos = report.get("ignore_eos") + temperature = report.get("temperature") + seed = report.get("seed") + model_sha256 = meta.get("model_sha256") + if ( + type(max_tokens) is not int or max_tokens <= 0 + or not isinstance(ignore_eos, bool) + or type(temperature) not in (int, float) + or type(seed) is not int + or not isinstance(model_sha256, str) or not model_sha256 + ): + raise ValueError("incomplete run metadata") + return max_tokens, ignore_eos, temperature, seed, model_sha256 + + +def report_key(item: dict) -> tuple[str, int, str]: + meta, level = item["meta"], item["level"] + workload = meta.get("workload") + variant = meta.get("variant") + clients = level.get("clients") + if not isinstance(workload, str) or not workload: + raise ValueError("incomplete report metadata: missing workload") + if not isinstance(variant, str) or not variant: + raise ValueError("incomplete report metadata: missing variant") + if type(clients) is not int or clients <= 0: + raise ValueError("incomplete report metadata: invalid clients") + return workload, clients, variant + + +def summarize(reports: list[dict]) -> str: + grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) + for item in reports: + grouped[report_key(item)].append(item) + for key, items in grouped.items(): + repeats = [item["meta"].get("repeat") for item in items] + if any(type(repeat) is not int or repeat <= 0 for repeat in repeats): + raise ValueError(f"{key}: invalid or missing repeat") + if len(repeats) != len(set(repeats)): + raise ValueError(f"{key}: duplicate repeat") + if len({run_signature(item) for item in items}) != 1: + raise ValueError(f"{key}: incompatible run metadata") + + lines = [ + "# Concurrency benchmark summary", "", + "Aggregate output goodput includes queueing, prefill, and decode. " + "Output-window goodput starts at the first observed output and is decode-facing, " + "but it can include staggered prefill. Prompt tok/s to first token includes " + "admission and TTFT.", "", + "| Workload | C | Variant | Repeats | Output goodput tok/s | " + "Output-window tok/s | Request decode tok/s | Prompt tok/s to first | " + "TTFT max s | Stable output | vs llama | Decode vs llama | K8 vs K1 |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " + ":---: | ---: | ---: | ---: |", + ] + for workload, clients, variant in sorted(grouped): + items = grouped[(workload, clients, variant)] + prompt_digests = [ + item["level"].get("selected_prompt_set_sha256") for item in items + ] + if not all(isinstance(value, str) and value for value in prompt_digests): + raise ValueError( + f"{workload} C={clients} {variant}: missing selected prompt set hash" + ) + hashes = set(prompt_digests) + if len(hashes) != 1: + raise ValueError(f"{workload} C={clients} {variant}: prompt sets differ") + goodput_values = [item["level"].get("aggregate_tok_s") for item in items] + if any(type(value) not in (int, float) for value in goodput_values): + raise ValueError( + f"{workload} C={clients} {variant}: missing aggregate token rate" + ) + goodput = median(goodput_values) + output_window = complete_median([ + item["level"].get("output_window_tok_s") for item in items + ]) + request_decode = complete_median([ + item["level"].get("request_decode_tok_s_median") for item in items + ]) + prompt_rate = complete_median([ + item["level"].get("prompt_tokens_per_s_to_first_token") for item in items + ]) + ttft = complete_median([ + item["level"].get("ttft_max_s") for item in items + ]) + output_digests = [ + item["level"].get("selected_output_set_sha256") for item in items + ] + output_digests_complete = all( + isinstance(value, str) and bool(value) for value in output_digests + ) + output_hashes = { + value for value in output_digests if isinstance(value, str) + } + stable = ( + "n/a" if len(items) < 2 or not output_digests_complete + else "yes" if len(output_hashes) == 1 + else "NO" + ) + + def delta(other: str, metric: str) -> str: + peers = grouped.get((workload, clients, other), []) + if not peers: + return "n/a" + peer_hashes = {p["level"]["selected_prompt_set_sha256"] for p in peers} + if peer_hashes != hashes: + raise ValueError(f"{workload} C={clients}: {variant}/{other} prompts differ") + if {run_signature(item) for item in items} != { + run_signature(peer) for peer in peers + }: + raise ValueError( + f"{workload} C={clients}: {variant}/{other} run metadata differs" + ) + by_repeat = {int(item["meta"]["repeat"]): item for item in items} + peers_by_repeat = {int(item["meta"]["repeat"]): item for item in peers} + if by_repeat.keys() != peers_by_repeat.keys(): + raise ValueError( + f"{workload} C={clients}: {variant}/{other} repeat sets differ" + ) + ratios = [] + for repeat in sorted(by_repeat): + value = by_repeat[repeat]["level"].get(metric) + base = peers_by_repeat[repeat]["level"].get(metric) + if value is None or base is None: + return "n/a" + if base <= 0: + raise ValueError( + f"{workload} C={clients} repeat={repeat}: " + f"non-positive {other} {metric}" + ) + ratios.append(value / base - 1.0) + return f"{median(ratios) * 100:+.1f}%" + + vs_llama = ( + delta("llama", "aggregate_tok_s") + if variant == "luce-k8" else "—" + ) + decode_vs_llama = ( + delta("llama", "output_window_tok_s") + if variant == "luce-k8" else "—" + ) + vs_k1 = ( + delta("luce-k1", "aggregate_tok_s") + if variant == "luce-k8" else "—" + ) + output_window_text = f"{output_window:.2f}" if output_window is not None else "n/a" + request_decode_text = f"{request_decode:.2f}" if request_decode is not None else "n/a" + prompt_rate_text = f"{prompt_rate:.2f}" if prompt_rate is not None else "n/a" + ttft_text = f"{ttft:.3f}" if ttft is not None else "n/a" + lines.append( + f"| {workload} | {clients} | {variant} | {len(items)} | {goodput:.2f} | " + f"{output_window_text} | {request_decode_text} | {prompt_rate_text} | " + f"{ttft_text} | {stable} | {vs_llama} | {decode_vs_llama} | {vs_k1} |" + ) + lines.append("") + return "\n".join(lines) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("root", type=Path) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + text = summarize(load_reports(args.root)) + if args.out: + args.out.write_text(text + "\n", encoding="utf-8") + print(text) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/summarize_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py new file mode 100755 index 000000000..b5b29c01e --- /dev/null +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +"""Summarize Qwen3.6 concurrent feature ablations and activation proof.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections import defaultdict +from pathlib import Path + + +def load_reports(root: Path) -> list[dict]: + reports = [] + for path in sorted(root.rglob("bench.json")): + report = json.loads(path.read_text(encoding="utf-8")) + meta = report.get("server_metadata") or {} + if len(report.get("levels", [])) != 1: + raise ValueError(f"{path}: expected exactly one client level") + level = report["levels"][0] + variant = str(meta.get("variant", "")) + if ( + level.get("failures") + or not level.get("token_count_complete") + or not level.get("prompt_token_count_complete") + or (variant != "llama" + and not level.get("effective_prompt_token_count_complete")) + ): + raise ValueError(f"{path}: failed or incomplete token accounting") + if report.get("ignore_eos") and level.get("fixed_token_workload_valid") is not True: + raise ValueError(f"{path}: fixed-token validation failed") + proof_path = path.with_name("feature-proof.json") + proof = None + if variant != "llama": + if not proof_path.is_file(): + raise ValueError(f"{path}: missing feature-proof.json") + proof = json.loads(proof_path.read_text(encoding="utf-8")) + if proof.get("valid") is not True: + raise ValueError(f"{proof_path}: activation proof failed") + expected_by_variant = { + "ar": [], "ddtree": ["ddtree"], "pflash": ["pflash"], + "kvflash": ["kvflash"], "full": ["ddtree", "kvflash", "pflash"], + } + if variant not in expected_by_variant: + raise ValueError(f"{path}: unknown Lucebox variant {variant!r}") + if proof.get("expected_features") != expected_by_variant[variant]: + raise ValueError( + f"{proof_path}: expected_features does not match variant {variant}" + ) + reports.append({ + "path": path, "report": report, "level": level, + "meta": meta, "proof": proof, + }) + if not reports: + raise ValueError(f"{root}: no bench.json files found") + return reports + + +def median(values: list[float]) -> float: + return statistics.median(values) + + +def fmt(value: float | None, digits: int = 2) -> str: + return f"{value:.{digits}f}" if value is not None else "n/a" + + +def complete_median(values: list[float | None]) -> float | None: + """Return a median only when every repeat measured the metric.""" + if not values or any(value is None for value in values): + return None + return median([value for value in values if value is not None]) + + +def run_signature(item: dict) -> tuple[object, ...]: + report, meta = item["report"], item["meta"] + max_tokens = report.get("max_tokens") + ignore_eos = report.get("ignore_eos") + model_sha256 = meta.get("model_sha256") + if ( + type(max_tokens) is not int or max_tokens <= 0 + or not isinstance(ignore_eos, bool) + or not isinstance(model_sha256, str) or not model_sha256 + ): + raise ValueError("incomplete run metadata") + return max_tokens, ignore_eos, model_sha256 + + +def output_stability(items: list[dict]) -> str: + output_digests = [ + item["level"].get("selected_output_set_sha256") for item in items + ] + complete = all(isinstance(value, str) and bool(value) for value in output_digests) + hashes = {value for value in output_digests if isinstance(value, str)} + return ( + "n/a" if len(items) < 2 or not complete + else "yes" if len(hashes) == 1 + else "NO" + ) + + +def summarize(reports: list[dict]) -> str: + grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) + for item in reports: + meta, level = item["meta"], item["level"] + key = (str(meta["workload"]), int(level["clients"]), str(meta["variant"])) + grouped[key].append(item) + for key, items in grouped.items(): + repeats = [int(item["meta"]["repeat"]) for item in items] + if len(repeats) != len(set(repeats)): + raise ValueError(f"{key}: duplicate repeat") + if len({run_signature(item) for item in items}) != 1: + raise ValueError(f"{key}: incompatible run metadata") + + lines = [ + "# Qwen3.6 concurrent feature matrix", "", + "Every Lucebox row is included only after request-correlated server telemetry " + "proves its requested features executed. Throughput is the median across fresh-process repeats.", + "", + "| Workload | C | Variant | N | Output goodput | Output-window | vs AR | " + "Effective/wire | DDTree accepted/step | Target forwards | KV in/out | " + "PFlash requests | TTFT max s | Stable output |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | " + ":--- | ---: | ---: | :---: |", + ] + for workload, clients, variant in sorted(grouped): + items = grouped[(workload, clients, variant)] + prompt_hashes = {item["level"]["selected_prompt_set_sha256"] for item in items} + if len(prompt_hashes) != 1: + raise ValueError(f"{workload} C={clients} {variant}: prompt sets differ") + goodput = median([item["level"]["aggregate_tok_s"] for item in items]) + window = complete_median([ + item["level"].get("output_window_tok_s") for item in items + ]) + ratio = complete_median([ + item["level"].get("effective_to_wire_prompt_ratio") for item in items + ]) + ttft = complete_median([ + item["level"].get("ttft_max_s") for item in items + ]) + stable = output_stability(items) + + peers = grouped.get((workload, clients, "ar"), []) + vs_ar = "—" if variant == "ar" else "n/a" + if variant != "ar" and peers: + peer_hashes = {p["level"]["selected_prompt_set_sha256"] for p in peers} + if peer_hashes != prompt_hashes: + raise ValueError(f"{workload} C={clients}: {variant}/ar prompts differ") + if {run_signature(item) for item in items} != { + run_signature(peer) for peer in peers + }: + raise ValueError( + f"{workload} C={clients}: {variant}/ar run metadata differs" + ) + by_repeat = {int(item["meta"]["repeat"]): item for item in items} + peers_by_repeat = {int(item["meta"]["repeat"]): item for item in peers} + if by_repeat.keys() != peers_by_repeat.keys(): + raise ValueError( + f"{workload} C={clients}: {variant}/ar repeat sets differ" + ) + if stable != "NO" and output_stability(peers) != "NO": + ratios = [] + for repeat in sorted(by_repeat): + value = by_repeat[repeat]["level"].get("aggregate_tok_s") + base = peers_by_repeat[repeat]["level"].get("aggregate_tok_s") + if value is None or base is None or base <= 0: + raise ValueError( + f"{workload} C={clients} repeat={repeat}: invalid AR goodput" + ) + ratios.append(value / base - 1.0) + vs_ar = f"{median(ratios) * 100:+.1f}%" + + proofs = [item["proof"] for item in items if item["proof"] is not None] + aggregates = [p["aggregate"] for p in proofs] + steps = sum(a["ddtree_steps"] for a in aggregates) + accepted = sum(a["ddtree_accepted_tokens"] for a in aggregates) + accepted_per_step = accepted / steps if steps else None + target_forwards = median([a["target_forwards"] for a in aggregates]) if aggregates else None + page_ins = median([a["kvflash_page_ins"] for a in aggregates]) if aggregates else None + page_outs = median([a["kvflash_page_outs"] for a in aggregates]) if aggregates else None + pflash_requests = median([a["pflash_applied_requests"] for a in aggregates]) if aggregates else None + kv_text = ( + f"{page_ins:.0f}/{page_outs:.0f}" + if page_ins is not None and page_outs is not None else "n/a" + ) + lines.append( + f"| {workload} | {clients} | {variant} | {len(items)} | {goodput:.2f} | " + f"{fmt(window)} | {vs_ar} | {fmt(ratio, 3)} | {fmt(accepted_per_step)} | " + f"{fmt(target_forwards, 0)} | {kv_text} | {fmt(pflash_requests, 0)} | " + f"{fmt(ttft, 3)} | {stable} |" + ) + lines.append("") + return "\n".join(lines) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("root", type=Path) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + text = summarize(load_reports(args.root)) + if args.out: + args.out.write_text(text + "\n", encoding="utf-8") + print(text) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/test_concurrency_tools.py b/harness/benchmarks/concurrency/test_concurrency_tools.py new file mode 100644 index 000000000..1aa8dc2d8 --- /dev/null +++ b/harness/benchmarks/concurrency/test_concurrency_tools.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +"""Unit tests for the deterministic prompt generator and compact summarizer.""" + +from __future__ import annotations + +import importlib.util +import json +import tempfile +import unittest +from pathlib import Path + + +HERE = Path(__file__).parent + + +def load(name: str): + path = HERE / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +generator = load("generate_ragged_prompts") +summarizer = load("summarize_concurrency") + + +class PromptGeneratorTests(unittest.TestCase): + def test_cohorts_are_disjoint_ragged_and_mean_matched(self) -> None: + records = generator.build_records("short") + self.assertEqual(len(records), 29) + self.assertEqual( + [row["cohort"] for row in records], + ["c1"] + ["c4"] * 4 + ["c8"] * 8 + ["c16"] * 16, + ) + self.assertEqual(len({row["prompt"] for row in records}), 29) + by_cohort = { + cohort: [row for row in records if row["cohort"] == cohort] + for cohort in ("c1", "c4", "c8", "c16") + } + means = { + cohort: sum(row["target_words"] for row in rows) / len(rows) + for cohort, rows in by_cohort.items() + } + self.assertEqual(len(set(means.values())), 1) + for cohort in ("c4", "c8", "c16"): + self.assertEqual(len({row["target_words"] for row in by_cohort[cohort]}), 4) + for row in records: + self.assertEqual(len(row["prompt"].split()), row["target_words"]) + + +class SummarizerTests(unittest.TestCase): + @staticmethod + def item( + variant: str, + goodput: float, + output_window: float | None = None, + *, + repeat: int = 1, + output_hash: str | None = "same-outputs", + ) -> dict: + return { + "report": { + "max_tokens": 256, "ignore_eos": True, + "temperature": 0.0, "seed": 1, + }, + "meta": { + "workload": "short", "variant": variant, "repeat": repeat, + "model_sha256": "a" * 64, + }, + "level": { + "clients": 8, + "aggregate_tok_s": goodput, + "output_window_tok_s": output_window if output_window is not None else goodput, + "request_decode_tok_s_median": goodput / 8, + "prompt_tokens_per_s_to_first_token": 100.0, + "ttft_max_s": 2.0, + "selected_prompt_set_sha256": "same-prompts", + "selected_output_set_sha256": output_hash, + }, + } + + def test_summary_reports_product_and_packing_deltas(self) -> None: + text = summarizer.summarize([ + self.item("luce-k8", 20.0), + self.item("luce-k1", 10.0), + self.item("llama", 8.0), + ]) + self.assertIn("+150.0%", text) + self.assertIn("+100.0%", text) + self.assertIn("Decode vs llama", text) + + def test_summary_uses_same_repeat_ratios(self) -> None: + reports = [] + for repeat, luce, llama in ( + (1, 10.0, 1.0), + (2, 20.0, 90.0), + (3, 100.0, 50.0), + ): + reports.extend([ + self.item("luce-k8", luce, repeat=repeat), + self.item("llama", llama, repeat=repeat), + ]) + text = summarizer.summarize(reports) + luce_row = next(line for line in text.splitlines() if "| luce-k8 |" in line) + self.assertIn("+100.0%", luce_row) + self.assertNotIn("-60.0%", luce_row) + + def test_summary_rejects_mismatched_repeat_sets(self) -> None: + reports = [ + self.item("luce-k8", 20.0, repeat=1), + self.item("luce-k8", 22.0, repeat=2), + self.item("llama", 10.0, repeat=1), + ] + with self.assertRaisesRegex(ValueError, "repeat sets differ"): + summarizer.summarize(reports) + + def test_single_repeat_does_not_claim_stability(self) -> None: + text = summarizer.summarize([self.item("llama", 8.0)]) + row = next(line for line in text.splitlines() if "| llama |" in line) + self.assertEqual(row.split("|")[10].strip(), "n/a") + + def test_multiple_repeats_report_output_stability(self) -> None: + stable = summarizer.summarize([ + self.item("llama", 8.0, repeat=1), + self.item("llama", 9.0, repeat=2), + ]) + stable_row = next(line for line in stable.splitlines() if "| llama |" in line) + self.assertEqual(stable_row.split("|")[10].strip(), "yes") + + unstable = summarizer.summarize([ + self.item("llama", 8.0, repeat=1, output_hash="first"), + self.item("llama", 9.0, repeat=2, output_hash="second"), + ]) + unstable_row = next(line for line in unstable.splitlines() if "| llama |" in line) + self.assertEqual(unstable_row.split("|")[10].strip(), "NO") + + def test_missing_output_digest_does_not_claim_stability(self) -> None: + reports = [ + self.item("llama", 8.0, repeat=1, output_hash=None), + self.item("llama", 9.0, repeat=2, output_hash=None), + ] + text = summarizer.summarize(reports) + row = next(line for line in text.splitlines() if "| llama |" in line) + self.assertEqual(row.split("|")[10].strip(), "n/a") + + def test_incomplete_repeat_metrics_are_reported_as_na(self) -> None: + first = self.item("llama", 8.0, repeat=1) + second = self.item("llama", 9.0, repeat=2) + for metric in ( + "output_window_tok_s", "request_decode_tok_s_median", + "prompt_tokens_per_s_to_first_token", "ttft_max_s", + ): + second["level"][metric] = None + text = summarizer.summarize([first, second]) + row = next(line for line in text.splitlines() if "| llama |" in line) + self.assertEqual( + [row.split("|")[column].strip() for column in range(6, 10)], + ["n/a"] * 4, + ) + + def test_comparison_rejects_incompatible_run_metadata(self) -> None: + fields = ( + ("report", "max_tokens", 64), + ("report", "ignore_eos", False), + ("report", "temperature", 0.5), + ("report", "seed", 2), + ("meta", "model_sha256", "b" * 64), + ) + for container, key, value in fields: + with self.subTest(key=key): + luce = self.item("luce-k8", 20.0) + llama = self.item("llama", 10.0) + luce[container][key] = value + with self.assertRaisesRegex(ValueError, "run metadata differs"): + summarizer.summarize([luce, llama]) + + def test_summary_rejects_incomplete_run_metadata(self) -> None: + item = self.item("llama", 10.0) + del item["report"]["max_tokens"] + with self.assertRaisesRegex(ValueError, "incomplete run metadata"): + summarizer.summarize([item]) + + def test_summary_reports_descriptive_errors_for_missing_fields(self) -> None: + fields = ( + ("meta", "workload", "workload"), + ("level", "clients", "clients"), + ("level", "aggregate_tok_s", "aggregate token rate"), + ("level", "selected_prompt_set_sha256", "prompt set hash"), + ) + for container, key, message in fields: + with self.subTest(key=key): + item = self.item("llama", 10.0) + del item[container][key] + with self.assertRaisesRegex(ValueError, message): + summarizer.summarize([item]) + + def test_load_reports_rejects_missing_prompt_usage(self) -> None: + report = { + "ignore_eos": True, + "server_metadata": {"workload": "short", "variant": "llama", "repeat": 1}, + "levels": [{ + "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": False, + "fixed_token_workload_valid": True, + }], + } + with tempfile.TemporaryDirectory() as root: + path = Path(root) / "bench.json" + path.write_text(json.dumps(report), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "incomplete token accounting"): + summarizer.load_reports(Path(root)) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_concurrent_benchmark.py new file mode 100644 index 000000000..c0455edc3 --- /dev/null +++ b/harness/benchmarks/concurrency/test_concurrent_benchmark.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Focused tests for concurrent_benchmark.py.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import unittest +from pathlib import Path +from unittest import mock + + +SCRIPT = Path(__file__).with_name("concurrent_benchmark.py") +SPEC = importlib.util.spec_from_file_location("concurrent_benchmark", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(benchmark) + + +class BenchmarkTests(unittest.TestCase): + def test_sse_parser_handles_events_and_done(self) -> None: + lines = [ + b'data: {"choices":[{"delta":{"content":"hi"}}]}\n', b"\n", + b"data: [DONE]\n", b"\n", + ] + self.assertEqual( + list(benchmark.iter_sse_data(lines)), + ['{"choices":[{"delta":{"content":"hi"}}]}', "[DONE]"], + ) + + def test_prompt_selection_never_wraps(self) -> None: + self.assertEqual(benchmark.request_prompts(["a", "b", "c"], 2, 1), ["b", "c"]) + with self.assertRaisesRegex(ValueError, "refusing to reuse"): + benchmark.request_prompts(["a", "b"], 2, 1) + + def test_level_uses_exact_usage_and_first_token_window(self) -> None: + def fake_request(_args: argparse.Namespace, prompt: str) -> dict: + start, first, end, prompt_tokens = { + "first": (10.0, 12.0, 14.0, 10), + "second": (10.25, 11.25, 15.0, 30), + }[prompt] + return { + "t_start": start, "t_first": first, "t_end": end, + "duration_s": end - start, "ttft_s": first - start, + "decode_duration_s": end - first, + "completion_tokens": 8, "prompt_tokens": prompt_tokens, + "finish_reason": "length", "error": None, + "content_sha256": benchmark.sha256_text(prompt + " output"), + "reasoning_content_sha256": benchmark.sha256_text(""), + "content_chars": 6, "reasoning_content_chars": 0, + "request_output_tok_s": 8 / (end - start), + "request_decode_tok_s": 7 / (end - first), + } + + args = argparse.Namespace(max_tokens=8, ignore_eos=True, timeout=2.0) + with mock.patch.object(benchmark, "stream_request", side_effect=fake_request): + level = benchmark.run_level(2, args, ["first", "second"], 0) + self.assertEqual(level["completion_tokens_total"], 16) + self.assertEqual(level["prompt_tokens_total"], 40) + self.assertTrue(level["fixed_token_workload_valid"]) + # Assert independently known windows before their derived rates. + self.assertEqual(level["wall_s"], 5.0) + self.assertEqual(level["output_window_s"], 3.75) + self.assertEqual(level["prompt_to_first_token_s"], 2.0) + self.assertAlmostEqual(level["aggregate_tok_s"], 3.2) + self.assertAlmostEqual( + level["output_window_tok_s"], 16 / 3.75, + ) + self.assertAlmostEqual( + level["request_decode_tok_s_median"], (3.5 + 7 / 3.75) / 2, + ) + self.assertAlmostEqual(level["prompt_tokens_per_s_to_first_token"], 20.0) + + def test_stream_request_keeps_usage_separate_from_sse_chunks(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"choices":[{"delta":{"content":"one chunk"}}]}\n', b"\n", + b'data: {"choices":[{"delta":{},"finish_reason":"length"}]}\n', b"\n", + b'data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":64}}\n', b"\n", + b"data: [DONE]\n", b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=64, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + record = benchmark.stream_request(args, "prompt") + self.assertEqual(record["completion_tokens"], 64) + self.assertEqual(record["prompt_tokens"], 12) + self.assertTrue(record["done_received"]) + self.assertIsNone(record["error"]) + self.assertIsNotNone(record["request_decode_tok_s"]) + self.assertEqual(record["content_sha256"], benchmark.sha256_text("one chunk")) + + def test_stream_request_rejects_clean_eof_without_done(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"choices":[{"delta":{"content":"partial"}}]}\n', b"\n", + b'data: {"choices":[{"delta":{},"finish_reason":"length"}]}\n', b"\n", + b'data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":64}}\n', b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=64, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + record = benchmark.stream_request(args, "prompt") + self.assertFalse(record["done_received"]) + self.assertIn("before [DONE]", record["error"]) + + def test_missing_prompt_usage_fails_level(self) -> None: + level = { + "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": False, + "fixed_token_workload_valid": True, + } + self.assertTrue(benchmark.level_failed(level, ignore_eos=True)) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py new file mode 100644 index 000000000..d582f6e32 --- /dev/null +++ b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Focused tests for request-correlated feature benchmark telemetry.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import sys +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +SCRIPT = HERE / "feature_concurrent_benchmark.py" +SPEC = importlib.util.spec_from_file_location("feature_concurrent_benchmark", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(benchmark) + + +class FeatureBenchmarkTests(unittest.TestCase): + def test_stream_request_captures_id_and_effective_prompt(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"id":"chatcmpl-42","choices":[{"delta":{"content":"x"}}]}\n', b"\n", + b'data: {"id":"chatcmpl-42","choices":[{"delta":{},' + b'"finish_reason":"length"}]}\n', b"\n", + b'data: {"id":"chatcmpl-42","choices":[],"usage":' + b'{"prompt_tokens":40000,"completion_tokens":8,"timings":' + b'{"effective_prompt_tokens":2000,"prefilled_tokens":2000,' + b'"cached_prefix_tokens":0,"cache_hit":false,"prefill_ms":12.5,' + b'"decode_ms":20.0,"decode_tokens_per_sec":400.0}}}\n', b"\n", + b"data: [DONE]\n", b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=8, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + row = benchmark.stream_request(args, "prompt") + self.assertEqual(row["request_id"], "chatcmpl-42") + self.assertEqual(row["prompt_tokens"], 40000) + self.assertEqual(row["effective_prompt_tokens"], 2000) + self.assertEqual(row["server_prefill_ms"], 12.5) + self.assertTrue(row["done_received"]) + self.assertIsNone(row["error"]) + + def test_clean_eof_without_done_is_rejected(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"id":"chatcmpl-cut","choices":[{"delta":{"content":"x"},' + b'"finish_reason":"length"}]}\n', b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=8, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + row = benchmark.stream_request(args, "prompt") + self.assertIn("before [DONE]", row["error"]) + + def test_enrich_level_reports_compression_ratio(self) -> None: + level = { + "prompt_tokens_total": 40000, + "requests_detail": [{ + "request_id": "r1", "error": None, + "effective_prompt_tokens": 2000, + "server_prefill_ms": 10.0, "server_decode_ms": 20.0, + "server_decode_tokens_per_sec": 100.0, + }], + } + benchmark.enrich_level(level) + self.assertTrue(level["request_ids_complete"]) + self.assertTrue(level["effective_prompt_token_count_complete"]) + self.assertEqual(level["effective_prompt_tokens_total"], 2000) + self.assertEqual(level["effective_to_wire_prompt_ratio"], 0.05) + + def test_duplicate_request_ids_are_not_complete(self) -> None: + level = { + "prompt_tokens_total": 20, + "requests_detail": [ + {"request_id": "same", "error": None, "effective_prompt_tokens": 10}, + {"request_id": "same", "error": None, "effective_prompt_tokens": 10}, + ], + } + benchmark.enrich_level(level) + self.assertFalse(level["request_ids_complete"]) + + def test_boolean_wire_counts_are_not_accepted_as_integers(self) -> None: + level = { + "prompt_tokens_total": 1, + "requests_detail": [{ + "request_id": "r1", "error": None, + "effective_prompt_tokens": True, + }], + } + benchmark.enrich_level(level) + self.assertFalse(level["effective_prompt_token_count_complete"]) + + def test_client_provenance_records_exact_argv_and_source_digest(self) -> None: + argv = ["python3", "feature_concurrent_benchmark.py", "--clients", "4"] + result = benchmark.client_provenance(argv) + self.assertEqual(result["client_argv"], argv) + self.assertEqual(result["client_script"], str(SCRIPT.resolve())) + self.assertEqual( + result["client_script_sha256"], + hashlib.sha256(SCRIPT.read_bytes()).hexdigest(), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_metadata.py b/harness/benchmarks/concurrency/test_feature_metadata.py new file mode 100644 index 000000000..7349fa21a --- /dev/null +++ b/harness/benchmarks/concurrency/test_feature_metadata.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""Tests for literal feature flags and reproducibility metadata.""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +SCRIPT = HERE / "write_feature_metadata.py" +SPEC = importlib.util.spec_from_file_location("write_feature_metadata", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +metadata = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(metadata) +RUNTIME_SCRIPT = HERE / "record_feature_runtime.py" +RUNTIME_SPEC = importlib.util.spec_from_file_location( + "record_feature_runtime", RUNTIME_SCRIPT, +) +assert RUNTIME_SPEC is not None and RUNTIME_SPEC.loader is not None +runtime_metadata = importlib.util.module_from_spec(RUNTIME_SPEC) +RUNTIME_SPEC.loader.exec_module(runtime_metadata) + + +class FeatureMetadataTests(unittest.TestCase): + def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + target = root / "target.gguf" + draft = root / "draft.gguf" + prefill = root / "prefill.gguf" + prompts = root / "prompts.jsonl" + command = root / "command.txt" + out = root / "metadata.json" + target.write_bytes(b"target") + draft.write_bytes(b"draft") + prefill.write_bytes(b"prefill") + prompts.write_text('{"prompt":"p"}\n', encoding="utf-8") + command.write_text("server --target-device hip:0\n", encoding="utf-8") + argv = [ + str(SCRIPT), "--out", str(out), "--variant", "full", + "--workload", "compression", "--clients", "4", "--repeat", "1", + "--binary", "/bin/true", "--model", str(target), + "--model-sha256", hashlib.sha256(target.read_bytes()).hexdigest(), + "--prompt-file", str(prompts), "--command-file", str(command), + "--repo", str(HERE.parents[2]), "--max-concurrent-prefills", "8", + "--target-device", "hip:0", "--draft-device", "hip:0", + "--draft-model", str(draft), "--draft-model-sha256", "draft-sha", + "--ddtree", "--ddtree-budget", "22", + "--prefill-compression", "auto", "--prefill-threshold", "32000", + "--prefill-keep-ratio", "0.05", "--prefill-drafter", str(prefill), + "--prefill-drafter-sha256", "prefill-sha", + "--draft-residency", "persistent", "--kvflash", "auto", + "--kvflash-max-pool-tokens", "8192", + "--kvflash-scorer-drafter", str(prefill), + "--kvflash-scorer-drafter-sha256", "prefill-sha", + ] + with mock.patch.object(sys, "argv", argv): + self.assertEqual(metadata.main(), 0) + result = json.loads(out.read_text(encoding="utf-8")) + self.assertEqual(result["model_sha256"], hashlib.sha256(b"target").hexdigest()) + self.assertTrue(result["server_binary_sha256"]) + self.assertTrue(result["git_head"]) + self.assertEqual(result["feature_config"]["draft_model_sha256"], "draft-sha") + self.assertEqual(result["feature_config"]["prefill_drafter_sha256"], "prefill-sha") + self.assertEqual( + result["feature_config"]["kvflash_scorer_drafter"], + str(prefill.resolve()), + ) + self.assertEqual( + result["feature_config"]["kvflash_scorer_drafter_sha256"], + "prefill-sha", + ) + self.assertEqual(result["literal_screenshot_flags"], [ + "--target-device", "hip:0", + "--draft-device", "hip:0", + "--ddtree", + "--ddtree-budget", "22", + "--draft-residency", "persistent", + "--prefill-compression", "auto", + "--prefill-drafter", str(prefill.resolve()), + "--kvflash", "auto", + ]) + self.assertIsNone(result["runtime_observed"]) + + def test_ldd_failure_is_fatal(self) -> None: + failed = mock.Mock(returncode=1, stdout="", stderr="not a dynamic executable") + with mock.patch.object(metadata.subprocess, "run", return_value=failed): + with self.assertRaisesRegex(RuntimeError, "ldd failed"): + metadata.resolved_libraries(Path("/tmp/server")) + + def test_unresolved_shared_library_is_fatal(self) -> None: + unresolved = mock.Mock( + returncode=0, stdout="libmissing.so => not found\n", stderr="", + ) + with mock.patch.object(metadata.subprocess, "run", return_value=unresolved): + with self.assertRaisesRegex(RuntimeError, "unresolved libraries"): + metadata.resolved_libraries(Path("/tmp/server")) + + def test_git_revision_failure_and_empty_output_are_fatal(self) -> None: + for result, message in ( + (mock.Mock(returncode=128, stdout="", stderr="not a repository"), + "git rev-parse failed"), + (mock.Mock(returncode=0, stdout="\n", stderr=""), + "empty revision"), + ): + with self.subTest(message=message): + with mock.patch.object(metadata.subprocess, "run", return_value=result): + with self.assertRaisesRegex(RuntimeError, message): + metadata.repository_head(Path("/tmp/repo")) + + def test_runtime_records_actual_kvflash_pool_from_startup(self) -> None: + original = { + "schema_version": 3, + "feature_config": { + "kvflash": "auto", + "kvflash_max_pool_tokens": 16384, + }, + } + log = "\n".join(( + "[parallel-kvflash] physical resident pool 8192 tokens; " + "logical per-slot cap 65536 across 16 slots " + "(--kv-pool-tokens does not expand resident VRAM)", + "[paged-attention] 512 physical blocks x 16 tokens " + "(8192 pool tokens, per-sequence max_ctx 65536)", + )) + result = runtime_metadata.update_metadata(original, log) + observed = result["runtime_observed"] + self.assertTrue(observed["kvflash_active"]) + self.assertEqual(observed["physical_kv_pool_tokens"], 8192) + self.assertEqual(observed["physical_kv_pool_blocks"], 512) + self.assertEqual(observed["kv_block_size_tokens"], 16) + self.assertEqual(observed["logical_per_slot_max_ctx"], 65536) + self.assertEqual(observed["configured_slots"], 16) + + def test_runtime_rejects_enabled_kvflash_without_marker(self) -> None: + original = {"feature_config": {"kvflash": "auto"}} + with self.assertRaisesRegex(ValueError, "startup marker is missing"): + runtime_metadata.update_metadata( + original, + "[paged-attention] 512 physical blocks x 16 tokens " + "(8192 pool tokens, per-sequence max_ctx 65536)", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py new file mode 100644 index 000000000..85c2294e6 --- /dev/null +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -0,0 +1,556 @@ +#!/usr/bin/env python3 +"""Tests for pressure prompts, activation proof, and feature summaries.""" + +from __future__ import annotations + +import importlib.util +import json +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) + + +def load(name: str): + path = HERE / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +generator = load("generate_feature_prompts") +proof = load("verify_feature_metrics") +summary = load("summarize_feature_matrix") + + +def report( + variant: str = "full", workload: str = "compression", + effective: tuple[int, int] = (2000, 2100), +) -> dict: + return { + "server_metadata": { + "variant": variant, + "workload": workload, + "feature_config": { + "prefill_compression": "auto", + "prefill_threshold": 32000, + "kvflash": "auto", + "kvflash_max_pool_tokens": 8192, + "kvflash_scorer_drafter": "/models/Qwen3-0.6B-BF16.gguf", + "kvflash_scorer_drafter_sha256": "a" * 64, + }, + "runtime_observed": { + "kvflash_active": True, + "physical_kv_pool_tokens": 8192, + }, + }, + "levels": [{ + "requests_detail": [ + {"request_id": "r1", "error": None, + "prompt_tokens": 40000, + "effective_prompt_tokens": effective[0]}, + {"request_id": "r2", "error": None, + "prompt_tokens": 40000, + "effective_prompt_tokens": effective[1]}, + ], + }], + } + + +def metric(request_id: str, effective: int, page_outs: int = 1) -> dict: + return { + "request_id": request_id, + "effective_prompt_tokens": effective, + "ddtree_steps": 3, + "ddtree_suspensions": 0, + # Zero acceptance is legitimate and must not invalidate execution proof. + "ddtree_accepted_tokens": 0, + "target_forwards": 3, + "kvflash_page_ins": 0, + "kvflash_page_outs": page_outs, + "kvflash_resident_blocks": 8, + "kvflash_reselects": 1, + "pflash_applied": True, + "pflash_input_tokens": 40000, + "pflash_output_tokens": effective, + } + + +class FeaturePromptTests(unittest.TestCase): + def test_activation_profiles_are_disjoint_and_above_thresholds(self) -> None: + compression = generator.build_records("compression") + pressure = generator.build_records("kv-pressure") + self.assertEqual(len(compression), 29) + self.assertEqual(len(pressure), 29) + self.assertEqual(len({row["prompt"] for row in compression}), 29) + self.assertEqual(len({row["prompt"] for row in pressure}), 29) + self.assertTrue( + {row["prompt"] for row in compression}.isdisjoint( + row["prompt"] for row in pressure + ) + ) + self.assertGreaterEqual(min(row["target_words"] for row in compression), 34000) + self.assertGreaterEqual(min(row["target_words"] for row in pressure), 12000) + self.assertTrue(all(row["activation_target"] == "pflash-auto" for row in compression)) + + +class FeatureRunnerShellTests(unittest.TestCase): + def run_invalid_matrix( + self, tmp: str, **overrides: str, + ) -> subprocess.CompletedProcess[str]: + model = Path(tmp) / "model.gguf" + model.touch() + env = { + key: value for key, value in os.environ.items() + if not key.startswith(("GGML_", "DFLASH_", "LUCE_", "HIP_", "ROCR_", "HSA_")) + and key not in ("LD_PRELOAD", "LD_LIBRARY_PATH") + } + env.update({ + "MODEL": str(model), + "LUCE_SERVER_BIN": "/bin/true", + "OUT": str(Path(tmp) / "out"), + "VARIANTS": "ar", + **overrides, + }) + return subprocess.run( + ["bash", str(HERE / "run_qwen36_feature_matrix.sh")], + env=env, capture_output=True, text=True, check=False, + ) + + def test_client_and_proof_invocations_are_array_backed(self) -> None: + runner = (HERE / "run_qwen36_feature_matrix.sh").read_text( + encoding="utf-8", + ) + client_lines = [ + line.strip() for line in runner.splitlines() + if 'python3 "$CLIENT"' in line + ] + self.assertEqual(client_lines, [ + 'python3 "$CLIENT" "${common_client[@]}"', + 'python3 "$CLIENT" "${common_client[@]}"', + ]) + self.assertIn('local -a warmup_cmd=(', runner) + self.assertIn('local -a benchmark_cmd=(', runner) + self.assertIn( + '"${warmup_cmd[@]}" > "$case_dir/warmup.txt"', + runner, + ) + self.assertIn( + '"${benchmark_cmd[@]}" | tee "$case_dir/bench.txt"', + runner, + ) + + proof_lines = [ + line.strip() for line in runner.splitlines() + if 'python3 "$PROOF_TOOL"' in line + ] + self.assertEqual(proof_lines, ['python3 "$PROOF_TOOL"']) + self.assertIn('local -a proof_cmd=(', runner) + self.assertIn('"${proof_cmd[@]}"', runner) + + def test_signals_exit_and_launch_environment_is_recorded(self) -> None: + runner = (HERE / "run_qwen36_feature_matrix.sh").read_text(encoding="utf-8") + self.assertIn("trap stop_server EXIT", runner) + self.assertIn("trap 'exit 130' INT", runner) + self.assertIn("trap 'exit 143' TERM", runner) + self.assertNotIn("trap stop_server EXIT INT TERM", runner) + self.assertIn("'env ' > \"$case_dir/server-command.txt\"", runner) + self.assertIn('"${launch_env[@]}" "${command[@]}"', runner) + + def test_duplicate_clients_are_rejected_before_artifacts_are_created(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self.run_invalid_matrix(tmp, CLIENTS="4,4") + self.assertEqual(result.returncode, 2, result.stderr) + self.assertIn("CLIENTS contains duplicate entry: 4", result.stderr) + self.assertFalse((Path(tmp) / "out").exists()) + + def test_duplicate_variants_are_rejected_before_artifacts_are_created(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self.run_invalid_matrix(tmp, VARIANTS="ar,ar") + self.assertEqual(result.returncode, 2, result.stderr) + self.assertIn("VARIANTS contains duplicate entry: ar", result.stderr) + self.assertFalse((Path(tmp) / "out").exists()) + + +class FeatureProofTests(unittest.TestCase): + def test_full_below_pool_passes_without_page_traffic(self) -> None: + rows = [ + metric("warmup", 50), + metric("r1", 2000, page_outs=0), + metric("r2", 2100, page_outs=0), + ] + result = proof.verify( + report(), rows, {"ddtree", "pflash", "kvflash"}, + ) + self.assertTrue(result["valid"], result["errors"]) + self.assertFalse(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "compressed-prompt-fits-pool", + ) + self.assertEqual(result["aggregate"]["ddtree_accepted_tokens"], 0) + self.assertEqual(result["aggregate"]["ddtree_suspensions"], 0) + self.assertEqual(result["matched_metric_count"], 2) + + def test_ddtree_suspensions_required_and_aggregated(self) -> None: + missing = metric("r1", 2000) + del missing["ddtree_suspensions"] + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + proof.PREFIX + json.dumps(missing) + "\n", + encoding="utf-8", + ) + with self.assertRaisesRegex( + ValueError, "missing telemetry keys.*ddtree_suspensions", + ): + proof.parse_markers(path) + + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[1]["ddtree_suspensions"] = 1 + result = proof.verify( + report(), rows, {"ddtree", "pflash", "kvflash"}, + ) + self.assertTrue(result["valid"], result["errors"]) + self.assertEqual(result["aggregate"]["ddtree_suspensions"], 1) + self.assertEqual( + [row["ddtree_suspensions"] for row in result["requests"]], + [0, 1], + ) + + def test_ddtree_suspensions_must_be_binary_per_request(self) -> None: + raw_two = metric("r1", 2000) + raw_two["ddtree_suspensions"] = 2 + with self.assertRaisesRegex( + ValueError, "ddtree_suspensions must be 0 or 1 per request", + ): + proof.aggregate_rows([raw_two]) + + def test_duplicate_terminal_telemetry_is_rejected(self) -> None: + row = metric("r1", 2000) + with self.assertRaisesRegex(ValueError, "duplicate telemetry request ID r1"): + proof.aggregate_rows([row, dict(row)]) + + def test_boolean_telemetry_counts_are_rejected(self) -> None: + integer_fields = (*proof.COUNTERS, "effective_prompt_tokens", + "kvflash_resident_blocks", "pflash_input_tokens", + "pflash_output_tokens") + for key in integer_fields: + with self.subTest(key=key): + row = metric("r1", 2000) + row[key] = True + with self.assertRaisesRegex(ValueError, key): + proof.aggregate_rows([row]) + + def test_boolean_wire_and_metadata_counts_do_not_prove_features(self) -> None: + rows = [metric("r1", 2000), metric("r2", 2100)] + wire = report() + wire["levels"][0]["requests_detail"][0]["effective_prompt_tokens"] = True + result = proof.verify(wire, rows, set()) + self.assertFalse(result["valid"]) + self.assertIn("missing usage.timings.effective_prompt_tokens", result["errors"][0]) + + threshold = report() + threshold["server_metadata"]["feature_config"]["prefill_threshold"] = True + result = proof.verify(threshold, rows, {"pflash"}) + self.assertFalse(result["valid"]) + self.assertIn("positive recorded token threshold", "\n".join(result["errors"])) + + pool = report() + pool["server_metadata"]["runtime_observed"]["physical_kv_pool_tokens"] = True + result = proof.verify(pool, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + self.assertIn("positive startup-observed physical pool", "\n".join(result["errors"])) + + requested = report() + requested["server_metadata"]["feature_config"]["kvflash_max_pool_tokens"] = True + result = proof.verify(requested, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + self.assertIn("recorded pool-token cap", "\n".join(result["errors"])) + + def test_full_above_pool_requires_page_traffic(self) -> None: + rows = [ + metric("r1", 9000, page_outs=0), + metric("r2", 9100, page_outs=0), + ] + result = proof.verify( + report(effective=(9000, 9100)), rows, + {"ddtree", "pflash", "kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertIn( + "no page-in/page-out", "\n".join(result["errors"]) + ) + + def test_full_aggregate_effective_demand_requires_page_traffic(self) -> None: + rows = [ + metric("r1", 5000, page_outs=0), + metric("r2", 5000, page_outs=0), + ] + result = proof.verify( + report(effective=(5000, 5000)), rows, + {"ddtree", "pflash", "kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "effective-prompt-exceeds-pool", + ) + self.assertIn("no page-in/page-out", "\n".join(result["errors"])) + + def test_full_effective_demand_is_evaluated_per_client_level(self) -> None: + input_report = report(effective=(5000, 5000)) + requests = input_report["levels"][0]["requests_detail"] + input_report["levels"] = [ + {"requests_detail": [requests[0]]}, + {"requests_detail": [requests[1]]}, + ] + rows = [ + metric("r1", 5000, page_outs=0), + metric("r2", 5000, page_outs=0), + ] + result = proof.verify( + input_report, rows, {"ddtree", "pflash", "kvflash"}, + ) + self.assertTrue(result["valid"], result["errors"]) + self.assertFalse(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "compressed-prompt-fits-pool", + ) + + def test_unknown_kvflash_variant_fails_closed(self) -> None: + rows = [ + metric("r1", 2000, page_outs=0), + metric("r2", 2100, page_outs=0), + ] + result = proof.verify( + report(variant="typo"), rows, {"kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertEqual(result["kvflash_page_traffic_reason"], "unknown-variant") + self.assertIn("no page-in/page-out", "\n".join(result["errors"])) + + def test_pflash_auto_requires_measured_input_above_threshold(self) -> None: + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[0]["pflash_input_tokens"] = 31999 + result = proof.verify(report(), rows, {"pflash"}) + self.assertFalse(result["valid"]) + self.assertIn( + "did not reach its recorded token threshold", + "\n".join(result["errors"]), + ) + + def test_kvflash_only_cannot_succeed_without_page_traffic(self) -> None: + rows = [ + metric("r1", 2000, page_outs=0), + metric("r2", 2100, page_outs=0), + ] + result = proof.verify( + report(variant="kvflash"), rows, {"kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "kvflash-only-ablation", + ) + self.assertIn( + "no page-in/page-out", "\n".join(result["errors"]) + ) + + def test_kvflash_requires_explicit_hashed_scorer(self) -> None: + input_report = report(variant="kvflash") + config = input_report["server_metadata"]["feature_config"] + config["kvflash_scorer_drafter"] = None + config["kvflash_scorer_drafter_sha256"] = None + rows = [metric("r1", 9000), metric("r2", 9100)] + result = proof.verify(input_report, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + text = "\n".join(result["errors"]) + self.assertIn("no explicit scorer drafter", text) + self.assertIn("not a valid SHA-256 digest", text) + + def test_kvflash_rejects_malformed_scorer_hash(self) -> None: + rows = [metric("r1", 9000), metric("r2", 9100)] + for digest in ("not-a-hash", "g" * 64, "a" * 63, "a" * 65): + with self.subTest(digest=digest): + input_report = report(variant="kvflash") + input_report["server_metadata"]["feature_config"][ + "kvflash_scorer_drafter_sha256" + ] = digest + result = proof.verify(input_report, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + self.assertIn( + "not a valid SHA-256 digest", "\n".join(result["errors"]) + ) + + def test_kvflash_requires_startup_observed_pool(self) -> None: + input_report = report(variant="kvflash") + input_report["server_metadata"]["runtime_observed"] = None + rows = [metric("r1", 9000), metric("r2", 9100)] + result = proof.verify(input_report, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + text = "\n".join(result["errors"]) + self.assertIn("startup marker was not recorded", text) + self.assertIn("no positive startup-observed physical pool", text) + + def test_requested_features_cannot_succeed_silently(self) -> None: + rows = [metric("r1", 9000, page_outs=0), metric("r2", 9100, page_outs=0)] + rows[0]["ddtree_steps"] = 0 + rows[1]["pflash_applied"] = False + result = proof.verify( + report(effective=(9000, 9100)), rows, + {"ddtree", "pflash", "kvflash"}, + ) + self.assertFalse(result["valid"]) + text = "\n".join(result["errors"]) + self.assertIn("ddtree_steps is zero", text) + self.assertIn("pflash_applied is false", text) + self.assertIn("no page-in/page-out", text) + + def test_pflash_input_mismatch_with_wire_tokens_fails(self) -> None: + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[0]["pflash_input_tokens"] = 39999 + result = proof.verify(report(), rows, {"pflash"}) + self.assertFalse(result["valid"]) + self.assertIn( + "pflash_input_tokens=39999 does not match wire value 40000", + "\n".join(result["errors"]), + ) + + def test_effective_prompt_mismatch_fails(self) -> None: + rows = [metric("r1", 1999), metric("r2", 2100)] + result = proof.verify(report(), rows, set()) + self.assertFalse(result["valid"]) + self.assertIn("does not match wire value", result["errors"][0]) + + def test_measured_request_requires_explicit_error_status(self) -> None: + input_report = report() + del input_report["levels"][0]["requests_detail"][0]["error"] + with self.assertRaisesRegex(ValueError, "explicit error status"): + proof.measured_requests(input_report) + + +class FeatureSummaryTests(unittest.TestCase): + @staticmethod + def item( + variant: str, goodput: float, *, repeat: int = 1, + output_hash: str | None = "same-output", + ) -> dict: + return { + "report": {"max_tokens": 256, "ignore_eos": True}, + "meta": { + "workload": "compression", "variant": variant, "repeat": repeat, + "model_sha256": "a" * 64, + }, + "level": { + "clients": 8, + "aggregate_tok_s": goodput, + "output_window_tok_s": goodput, + "effective_to_wire_prompt_ratio": 0.05 if variant == "full" else 1.0, + "ttft_max_s": 2.0, + "selected_prompt_set_sha256": "same-prompts", + "selected_output_set_sha256": output_hash, + }, + "proof": { + "aggregate": { + "ddtree_steps": 4 if variant == "full" else 0, + "ddtree_suspensions": 1 if variant == "full" else 0, + "ddtree_accepted_tokens": 8 if variant == "full" else 0, + "target_forwards": 4, + "kvflash_page_ins": 2 if variant == "full" else 0, + "kvflash_page_outs": 3 if variant == "full" else 0, + "pflash_applied_requests": 8 if variant == "full" else 0, + }, + }, + } + + def test_summary_compares_feature_row_to_ar(self) -> None: + text = summary.summarize([self.item("ar", 10.0), self.item("full", 12.0)]) + self.assertIn("+20.0%", text) + self.assertIn("2.00", text) + self.assertIn("DDTree steps/susp.", text) + self.assertIn("| 4/1 | 4 | 2/3 |", text) + + def test_feature_row_without_ar_control_reports_na(self) -> None: + text = summary.summarize([self.item("full", 12.0)]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[7].strip(), "n/a") + + def test_missing_output_digest_does_not_claim_stability(self) -> None: + first = self.item("full", 12.0, output_hash=None) + second = self.item("full", 13.0, repeat=2, output_hash=None) + text = summary.summarize([first, second]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[15].strip(), "n/a") + + def test_incomplete_repeat_metrics_are_reported_as_na(self) -> None: + first = self.item("full", 12.0) + second = self.item("full", 13.0, repeat=2) + second["level"]["output_window_tok_s"] = None + second["level"]["effective_to_wire_prompt_ratio"] = None + text = summary.summarize([first, second]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[6].strip(), "n/a") + self.assertEqual(row.split("|")[8].strip(), "n/a") + + def test_incomplete_repeat_ttft_is_reported_as_na(self) -> None: + first = self.item("full", 12.0) + second = self.item("full", 13.0, repeat=2) + second["level"]["ttft_max_s"] = None + text = summary.summarize([first, second]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[14].strip(), "n/a") + + def test_summary_rejects_incompatible_repeat_metadata(self) -> None: + first = self.item("full", 12.0) + second = self.item("full", 13.0, repeat=2) + second["report"]["max_tokens"] = 128 + with self.assertRaisesRegex(ValueError, "incompatible run metadata"): + summary.summarize([first, second]) + + def test_summary_rejects_incompatible_ar_control_metadata(self) -> None: + ar = self.item("ar", 10.0) + feature = self.item("full", 12.0) + feature["meta"]["model_sha256"] = "b" * 64 + with self.assertRaisesRegex(ValueError, "run metadata differs"): + summary.summarize([ar, feature]) + + def test_unstable_feature_row_suppresses_ar_delta(self) -> None: + reports = [ + self.item("ar", 10.0, repeat=1), + self.item("ar", 10.0, repeat=2), + self.item("full", 12.0, repeat=1, output_hash="first"), + self.item("full", 13.0, repeat=2, output_hash="second"), + ] + text = summary.summarize(reports) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[7].strip(), "n/a") + self.assertEqual(row.split("|")[15].strip(), "NO") + + def test_unstable_ar_control_suppresses_feature_delta(self) -> None: + reports = [ + self.item("ar", 10.0, repeat=1, output_hash="first"), + self.item("ar", 10.0, repeat=2, output_hash="second"), + self.item("full", 12.0, repeat=1), + self.item("full", 13.0, repeat=2), + ] + text = summary.summarize(reports) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[7].strip(), "n/a") + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/verify_feature_metrics.py b/harness/benchmarks/concurrency/verify_feature_metrics.py new file mode 100644 index 000000000..b8faf043b --- /dev/null +++ b/harness/benchmarks/concurrency/verify_feature_metrics.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +"""Fail closed unless server telemetry proves requested features executed.""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path +from typing import Any + + +PREFIX = "[concurrency-metrics] " +COUNTERS = ( + "ddtree_steps", "ddtree_suspensions", "ddtree_accepted_tokens", + "target_forwards", "kvflash_page_ins", "kvflash_page_outs", + "kvflash_reselects", +) +REQUIRED_KEYS = ( + "request_id", "effective_prompt_tokens", *COUNTERS, + "kvflash_resident_blocks", "pflash_applied", "pflash_input_tokens", + "pflash_output_tokens", +) + + +def parse_markers(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for line_no, line in enumerate(path.read_text(encoding="utf-8", errors="replace").splitlines(), 1): + marker = line.find(PREFIX) + if marker < 0: + continue + raw = line[marker + len(PREFIX):] + try: + row = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"{path}:{line_no}: invalid concurrency metric JSON: {exc}") from exc + if not isinstance(row, dict): + raise ValueError(f"{path}:{line_no}: concurrency metric must be an object") + missing = [key for key in REQUIRED_KEYS if key not in row] + if missing: + raise ValueError(f"{path}:{line_no}: missing telemetry keys {missing}") + rows.append(row) + return rows + + +def measured_requests(report: dict[str, Any]) -> dict[str, dict[str, Any]]: + requests: dict[str, dict[str, Any]] = {} + for level in report.get("levels") or []: + for row in level.get("requests_detail") or []: + if "error" not in row: + raise ValueError("bench report request lacks an explicit error status") + if row["error"] is not None: + continue + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("bench report lacks a request ID for a successful request") + if request_id in requests: + raise ValueError(f"duplicate measured request ID {request_id}") + requests[request_id] = row + if not requests: + raise ValueError("bench report has no successful measured requests") + return requests + + +def aggregate_rows(rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]: + aggregate: dict[str, dict[str, Any]] = {} + for row in rows: + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("telemetry request_id must be a non-empty string") + if request_id in aggregate: + raise ValueError(f"duplicate telemetry request ID {request_id}") + dst = { + "request_id": request_id, + "effective_prompt_tokens": row["effective_prompt_tokens"], + "kvflash_resident_blocks": row["kvflash_resident_blocks"], + "pflash_applied": False, + "pflash_input_tokens": row["pflash_input_tokens"], + "pflash_output_tokens": row["pflash_output_tokens"], + **{key: 0 for key in COUNTERS}, + } + aggregate[request_id] = dst + for key in COUNTERS: + value = row[key] + if key == "ddtree_suspensions": + if type(value) is not int or value not in (0, 1): + raise ValueError( + f"{request_id}: ddtree_suspensions must be 0 or 1 " + "per request" + ) + elif type(value) is not int or value < 0: + raise ValueError(f"{request_id}: {key} must be a non-negative integer") + dst[key] += value + for key in ( + "effective_prompt_tokens", "kvflash_resident_blocks", + "pflash_input_tokens", "pflash_output_tokens", + ): + value = row[key] + if type(value) is not int or value < 0: + raise ValueError(f"{request_id}: {key} must be a non-negative integer") + dst[key] = value + if not isinstance(row["pflash_applied"], bool): + raise ValueError(f"{request_id}: pflash_applied must be boolean") + dst["pflash_applied"] = dst["pflash_applied"] or row["pflash_applied"] + for request_id, row in aggregate.items(): + if row["ddtree_suspensions"] not in (0, 1): + raise ValueError( + f"{request_id}: ddtree_suspensions must be 0 or 1 per request" + ) + return aggregate + + +def verify( + report: dict[str, Any], markers: list[dict[str, Any]], expected: set[str], +) -> dict[str, Any]: + measured = measured_requests(report) + all_rows = aggregate_rows(markers) + rows = {request_id: all_rows[request_id] for request_id in measured if request_id in all_rows} + errors: list[str] = [] + missing = sorted(set(measured) - set(rows)) + if missing: + errors.append(f"missing concurrency telemetry for {len(missing)} measured request(s): {missing}") + + for request_id, measured_row in measured.items(): + metric = rows.get(request_id) + if metric is None: + continue + wire_effective = measured_row.get("effective_prompt_tokens") + if type(wire_effective) is not int or wire_effective < 0: + errors.append(f"{request_id}: missing usage.timings.effective_prompt_tokens") + elif metric["effective_prompt_tokens"] != wire_effective: + errors.append( + f"{request_id}: log effective_prompt_tokens={metric['effective_prompt_tokens']} " + f"does not match wire value {wire_effective}" + ) + if "pflash" in expected: + wire_input = measured_row.get("prompt_tokens") + if type(wire_input) is not int or wire_input < 0: + errors.append(f"{request_id}: missing usage.prompt_tokens") + elif metric["pflash_input_tokens"] != wire_input: + errors.append( + f"{request_id}: log pflash_input_tokens=" + f"{metric['pflash_input_tokens']} does not match wire value " + f"{wire_input}" + ) + if "ddtree" in expected: + if metric["ddtree_steps"] <= 0: + errors.append(f"{request_id}: DDTree requested but ddtree_steps is zero") + if metric["target_forwards"] <= 0: + errors.append(f"{request_id}: DDTree requested but target_forwards is zero") + if "pflash" in expected: + if metric["pflash_applied"] is not True: + errors.append(f"{request_id}: PFlash requested but pflash_applied is false") + if not (0 < metric["pflash_output_tokens"] < metric["pflash_input_tokens"]): + errors.append( + f"{request_id}: PFlash did not reduce prompt tokens " + f"({metric['pflash_input_tokens']} -> {metric['pflash_output_tokens']})" + ) + + totals = {key: sum(row[key] for row in rows.values()) for key in COUNTERS} + resident = [row["kvflash_resident_blocks"] for row in rows.values()] + metadata = report.get("server_metadata") or {} + feature_config = metadata.get("feature_config") or {} + variant = str(metadata.get("variant") or "") + workload = str(metadata.get("workload") or "") + pflash_mode = feature_config.get("prefill_compression") + pflash_threshold = feature_config.get("prefill_threshold") + if "pflash" in expected: + if not isinstance(pflash_mode, str) or pflash_mode in ("", "off", "0"): + errors.append("PFlash requested but metadata does not prove it was enabled") + if pflash_mode == "auto": + if type(pflash_threshold) is not int or pflash_threshold <= 0: + errors.append( + "PFlash auto row lacks a positive recorded token threshold" + ) + else: + below_threshold = sorted( + request_id for request_id, row in rows.items() + if row["pflash_input_tokens"] < pflash_threshold + ) + if below_threshold: + errors.append( + "PFlash auto input did not reach its recorded token threshold " + f"for request(s): {below_threshold}" + ) + kvflash_mode = feature_config.get("kvflash") + requested_pool_tokens = feature_config.get("kvflash_max_pool_tokens") + runtime_observed = metadata.get("runtime_observed") or {} + pool_tokens = runtime_observed.get("physical_kv_pool_tokens") + kvflash_page_traffic_required = False + kvflash_page_traffic_reason = "not-requested" + if "kvflash" in expected: + if not isinstance(kvflash_mode, str) or kvflash_mode in ("", "off", "0"): + errors.append("KVFlash requested but metadata does not prove it was enabled") + if (type(requested_pool_tokens) is not int + or requested_pool_tokens <= 0): + errors.append( + "KVFlash requested but its recorded pool-token cap is not a positive integer") + scorer_drafter = feature_config.get("kvflash_scorer_drafter") + scorer_sha256 = feature_config.get("kvflash_scorer_drafter_sha256") + if not isinstance(scorer_drafter, str) or not scorer_drafter: + errors.append( + "KVFlash requested but no explicit scorer drafter was recorded" + ) + if ( + not isinstance(scorer_sha256, str) + or re.fullmatch(r"[0-9a-fA-F]{64}", scorer_sha256) is None + ): + errors.append( + "KVFlash requested but the scorer drafter hash is not a valid SHA-256 digest" + ) + if runtime_observed.get("kvflash_active") is not True: + errors.append( + "KVFlash requested but its physical-pool startup marker was not recorded" + ) + if type(pool_tokens) is not int or pool_tokens <= 0: + errors.append( + "KVFlash requested but no positive startup-observed physical pool was recorded" + ) + if not resident or max(resident) <= 0: + errors.append("KVFlash requested but resident block count never became positive") + + effective_demand_by_level = [] + for level in report.get("levels") or []: + demand = 0 + for request in level.get("requests_detail") or []: + request_id = request.get("request_id") + if request.get("error") is None and request_id in rows: + demand += rows[request_id]["effective_prompt_tokens"] + effective_demand_by_level.append(demand) + if variant not in ("kvflash", "full"): + # A report whose declared variant is inconsistent with a KVFlash + # proof must not pass merely because its current prompts fit. + kvflash_page_traffic_required = True + kvflash_page_traffic_reason = "unknown-variant" + elif variant == "kvflash": + kvflash_page_traffic_required = True + kvflash_page_traffic_reason = "kvflash-only-ablation" + elif workload == "kv-pressure": + kvflash_page_traffic_required = True + kvflash_page_traffic_reason = "kv-pressure-workload" + elif variant == "full": + if type(pool_tokens) is not int or pool_tokens <= 0: + errors.append( + "full KVFlash row lacks a positive recorded pool-token limit" + ) + kvflash_page_traffic_reason = "missing-pool-limit" + else: + kvflash_page_traffic_required = any( + demand > pool_tokens + for demand in effective_demand_by_level + ) + kvflash_page_traffic_reason = ( + "effective-prompt-exceeds-pool" + if kvflash_page_traffic_required + else "compressed-prompt-fits-pool" + ) + + if ( + kvflash_page_traffic_required + and totals["kvflash_page_ins"] + totals["kvflash_page_outs"] <= 0 + ): + errors.append( + "KVFlash paging was required but no page-in/page-out was observed" + ) + + return { + "schema_version": 3, + "expected_features": sorted(expected), + "valid": not errors, + "errors": errors, + "measured_request_count": len(measured), + "matched_metric_count": len(rows), + "ignored_marker_count": len(markers) - len(rows), + "kvflash_page_traffic_required": kvflash_page_traffic_required, + "kvflash_page_traffic_reason": kvflash_page_traffic_reason, + "kvflash_pool_tokens": pool_tokens if type(pool_tokens) is int else None, + "kvflash_requested_max_pool_tokens": ( + requested_pool_tokens if type(requested_pool_tokens) is int else None + ), + "pflash_threshold_tokens": ( + pflash_threshold if type(pflash_threshold) is int else None + ), + "aggregate": { + **totals, + "kvflash_resident_blocks_max": max(resident) if resident else None, + "pflash_applied_requests": sum(row["pflash_applied"] is True for row in rows.values()), + "pflash_input_tokens": sum(row["pflash_input_tokens"] for row in rows.values()), + "pflash_output_tokens": sum(row["pflash_output_tokens"] for row in rows.values()), + }, + "requests": [rows[key] for key in sorted(rows)], + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--bench", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + parser.add_argument("--expect", action="append", choices=("ddtree", "pflash", "kvflash"), default=[]) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + try: + report = json.loads(args.bench.read_text(encoding="utf-8")) + result = verify(report, parse_markers(args.server_log), set(args.expect)) + except Exception as exc: + print(f"[proof] error: {exc}", file=sys.stderr) + return 2 + args.out.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") + if not result["valid"]: + for error in result["errors"]: + print(f"[proof] {error}", file=sys.stderr) + return 1 + print( + f"[proof] valid features={','.join(result['expected_features']) or 'ar'} " + f"requests={result['matched_metric_count']}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/write_feature_metadata.py b/harness/benchmarks/concurrency/write_feature_metadata.py new file mode 100644 index 000000000..9500c41b9 --- /dev/null +++ b/harness/benchmarks/concurrency/write_feature_metadata.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Write reproducible server and feature configuration metadata for one case.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import pathlib +import subprocess + + +def digest(path: pathlib.Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def resolved_libraries(binary: pathlib.Path) -> dict[str, str]: + result = subprocess.run( + ["ldd", str(binary)], text=True, capture_output=True, + ) + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() or "no diagnostic" + raise RuntimeError(f"ldd failed for {binary}: {detail}") + unresolved = [ + line.strip() for line in result.stdout.splitlines() + if "=> not found" in line + ] + if unresolved: + raise RuntimeError(f"ldd found unresolved libraries for {binary}: {unresolved}") + libraries: dict[str, str] = {} + for line in result.stdout.splitlines(): + fields = line.replace("=>", " ").split() + paths = [pathlib.Path(value) for value in fields if value.startswith("/")] + for path in paths: + if path.is_file(): + libraries[str(path.resolve())] = digest(path) + return libraries + + +def repository_head(repo: pathlib.Path) -> str: + result = subprocess.run( + ["git", "-C", str(repo), "rev-parse", "HEAD"], + text=True, capture_output=True, + ) + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() or "no diagnostic" + raise RuntimeError(f"git rev-parse failed for {repo}: {detail}") + head = result.stdout.strip() + if not head: + raise RuntimeError(f"git rev-parse returned an empty revision for {repo}") + return head + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=pathlib.Path, required=True) + parser.add_argument("--variant", required=True) + parser.add_argument("--workload", required=True) + parser.add_argument("--clients", type=int, required=True) + parser.add_argument("--repeat", type=int, required=True) + parser.add_argument("--binary", type=pathlib.Path, required=True) + parser.add_argument("--model", type=pathlib.Path, required=True) + parser.add_argument("--model-sha256", required=True) + parser.add_argument("--prompt-file", type=pathlib.Path, required=True) + parser.add_argument("--command-file", type=pathlib.Path, required=True) + parser.add_argument("--repo", type=pathlib.Path, required=True) + parser.add_argument("--max-concurrent-prefills", type=int, required=True) + parser.add_argument("--target-device", default=None) + parser.add_argument("--draft-device", default=None) + parser.add_argument("--draft-model", type=pathlib.Path) + parser.add_argument("--draft-model-sha256") + parser.add_argument("--ddtree", action="store_true") + parser.add_argument("--ddtree-budget", type=int) + parser.add_argument("--prefill-compression", default="off") + parser.add_argument("--prefill-threshold", type=int) + parser.add_argument("--prefill-keep-ratio", type=float) + parser.add_argument("--prefill-drafter", type=pathlib.Path) + parser.add_argument("--prefill-drafter-sha256") + parser.add_argument("--draft-residency", default=None) + parser.add_argument("--kvflash", default="off") + parser.add_argument("--kvflash-max-pool-tokens", type=int) + parser.add_argument("--kvflash-scorer-drafter", type=pathlib.Path) + parser.add_argument("--kvflash-scorer-drafter-sha256") + parser.add_argument("--launch-env", action="append", default=[]) + args = parser.parse_args() + + launch_env: dict[str, str] = {} + for item in args.launch_env: + key, sep, value = item.partition("=") + if not sep or not key: + parser.error(f"bad --launch-env {item!r}; expected KEY=VALUE") + launch_env[key] = value + + libraries = resolved_libraries(args.binary) + git_head = repository_head(args.repo) + literal_flags: list[str] = [] + if args.target_device: + literal_flags += ["--target-device", args.target_device] + if args.draft_device: + literal_flags += ["--draft-device", args.draft_device] + if args.ddtree: + literal_flags += ["--ddtree"] + if args.ddtree_budget is not None: + literal_flags += ["--ddtree-budget", str(args.ddtree_budget)] + if args.draft_residency: + literal_flags += ["--draft-residency", args.draft_residency] + if args.prefill_compression != "off": + literal_flags += ["--prefill-compression", args.prefill_compression] + if args.prefill_drafter: + literal_flags += ["--prefill-drafter", str(args.prefill_drafter.resolve())] + if args.kvflash != "off": + literal_flags += ["--kvflash", args.kvflash] + + obj = { + "schema_version": 3, + "variant": args.variant, + "workload": args.workload, + "clients": args.clients, + "repeat": args.repeat, + "max_concurrent_prefills": args.max_concurrent_prefills, + "server_binary": str(args.binary.resolve()), + "server_binary_sha256": digest(args.binary), + "model": str(args.model.resolve()), + "model_sha256": args.model_sha256, + "prompt_file_sha256": digest(args.prompt_file), + "server_command": args.command_file.read_text(encoding="utf-8").strip(), + "launch_environment": launch_env, + "resolved_shared_library_sha256": libraries, + "git_head": git_head, + "literal_screenshot_flags": literal_flags, + # Populated from fail-closed startup markers after the server is healthy. + "runtime_observed": None, + "feature_config": { + "target_device": args.target_device, + "draft_device": args.draft_device, + "draft_model": str(args.draft_model.resolve()) if args.draft_model else None, + "draft_model_sha256": args.draft_model_sha256, + "ddtree": args.ddtree, + "ddtree_budget": args.ddtree_budget, + "prefill_compression": args.prefill_compression, + "prefill_threshold": args.prefill_threshold, + "prefill_keep_ratio": args.prefill_keep_ratio, + "prefill_drafter": ( + str(args.prefill_drafter.resolve()) if args.prefill_drafter else None + ), + "prefill_drafter_sha256": args.prefill_drafter_sha256, + "draft_residency": args.draft_residency, + "kvflash": args.kvflash, + "kvflash_max_pool_tokens": args.kvflash_max_pool_tokens, + "kvflash_scorer_drafter": ( + str(args.kvflash_scorer_drafter.resolve()) + if args.kvflash_scorer_drafter else None + ), + "kvflash_scorer_drafter_sha256": ( + args.kvflash_scorer_drafter_sha256 + ), + }, + } + args.out.write_text(json.dumps(obj, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 631a9e176cba00fe3df8993e3219b039899cdbb8 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 12 Aug 2026 23:09:18 +0000 Subject: [PATCH 03/42] fix(qwen36): harden feature smoke paths --- .../concurrency/run_qwen36_feature_matrix.sh | 33 +++-- .../qwen35/concurrency/qwen35_seq_engine.cpp | 10 +- server/src/qwen35/graph_builders.cpp | 75 +++++++++-- server/src/qwen35/graph_builders.h | 36 ++++++ server/test/test_recurrent_snapshot.cpp | 116 ++++++++++++++++++ 5 files changed, 248 insertions(+), 22 deletions(-) diff --git a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh index e4b23bca1..d0ea1786a 100755 --- a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh @@ -322,18 +322,33 @@ run_case() { --require-distinct-prompts --temperature 0 --ignore-eos --timeout "$timeout" --cooldown 0) local -a telemetry_arg=() [[ "$variant" != llama ]] && telemetry_arg+=(--require-effective-prompt-telemetry) - python3 "$CLIENT" "${common_client[@]}" --max-tokens "$WARMUP_TOKENS" - "${telemetry_arg[@]}" --out "$case_dir/warmup.json" - --label "$variant $workload C=$clients warmup" > "$case_dir/warmup.txt" - python3 "$CLIENT" "${common_client[@]}" --max-tokens "$MAX_TOKENS" - "${telemetry_arg[@]}" --server-metadata-json "$case_dir/server-metadata.json" - --out "$case_dir/bench.json" --label "$variant $workload C=$clients repeat=$repeat" \ - | tee "$case_dir/bench.txt" + local -a warmup_cmd=( + python3 "$CLIENT" "${common_client[@]}" + --max-tokens "$WARMUP_TOKENS" "${telemetry_arg[@]}" + --out "$case_dir/warmup.json" + --label "$variant $workload C=$clients warmup" + ) + "${warmup_cmd[@]}" > "$case_dir/warmup.txt" + + local -a benchmark_cmd=( + python3 "$CLIENT" "${common_client[@]}" + --max-tokens "$MAX_TOKENS" "${telemetry_arg[@]}" + --server-metadata-json "$case_dir/server-metadata.json" + --out "$case_dir/bench.json" + --label "$variant $workload C=$clients repeat=$repeat" + ) + "${benchmark_cmd[@]}" | tee "$case_dir/bench.txt" stop_server if [[ "$variant" != llama ]]; then - python3 "$PROOF_TOOL" --bench "$case_dir/bench.json" --server-log "$case_dir/server.log" - "${expected[@]}" --out "$case_dir/feature-proof.json" + local -a proof_cmd=( + python3 "$PROOF_TOOL" + --bench "$case_dir/bench.json" + --server-log "$case_dir/server.log" + "${expected[@]}" + --out "$case_dir/feature-proof.json" + ) + "${proof_cmd[@]}" fi sleep "$COOLDOWN_SECONDS" } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 545e7d9bf..10e50450a 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -339,8 +339,14 @@ std::optional Qwen35SeqEngine::step_ddtree( sizeof(int32_t) * parents.size()); ggml_backend_tensor_set(tree_sg.tree_sizes, sizes.data(), 0, sizeof(int32_t) * sizes.size()); - ggml_backend_tensor_set(tree_sg.active_slot_ids, tree_slots.data(), 0, - sizeof(int32_t) * tree_slots.size()); + // Mapped-tree DeltaNet uses active_slot_ids only as a topology marker; + // gallocr may therefore optimize away its backing buffer. The actual + // state/attention mappings below are live graph inputs and remain + // mandatory. Upload the marker only if a future topology consumes it. + if (detail::target_paged_tree_active_slots_need_upload(tree_sg)) { + ggml_backend_tensor_set(tree_sg.active_slot_ids, tree_slots.data(), 0, + sizeof(int32_t) * tree_slots.size()); + } ggml_backend_tensor_set(tree_sg.state_slot_ids, tree_state_slots.data(), 0, sizeof(int32_t) * tree_state_slots.size()); ggml_backend_tensor_set(tree_sg.paged_query_seq_ids, query_slots.data(), 0, diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 0ca65ecdb..f96349af3 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -5,10 +5,56 @@ #include #include #include +#include #include namespace dflash::common { +bool detail::target_graph_capacity_for_parallel_segments( + int n_parallel_segments, + size_t & capacity) { + static constexpr int k_max_parallel_segments = 64; + static constexpr size_t k_base_capacity = 16384; + static constexpr int k_segments_per_capacity = 8; + static constexpr size_t k_max_capacity = + k_base_capacity * + (k_max_parallel_segments / k_segments_per_capacity); + + if (n_parallel_segments < 0 || + n_parallel_segments > k_max_parallel_segments) { + return false; + } + const int64_t scale = std::max( + 1, ((int64_t)n_parallel_segments + + k_segments_per_capacity - 1) / + k_segments_per_capacity); + if ((uint64_t)scale > + std::numeric_limits::max() / k_base_capacity) { + return false; + } + const size_t computed = k_base_capacity * (size_t)scale; + if (computed > k_max_capacity) return false; + capacity = computed; + return true; +} + +bool detail::target_paged_tree_graph_capacity( + int tree_width, + int n_tree_seqs, + size_t & capacity) { + static constexpr int tree_buckets[] = { + 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, + }; + if (tree_width < 1 || tree_width > 256 || + std::find(std::begin(tree_buckets), std::end(tree_buckets), + n_tree_seqs) == std::end(tree_buckets) || + (int64_t)tree_width * n_tree_seqs > INT32_MAX) { + return false; + } + return target_graph_capacity_for_parallel_segments( + n_tree_seqs, capacity); +} + bool detail::validate_target_paged_tree_layout( const TargetCache & cache, int tree_width, @@ -16,12 +62,9 @@ bool detail::validate_target_paged_tree_layout( int paged_max_kv_len, int tree_scratch_base, int tree_scratch_stride) { - static constexpr int tree_buckets[] = { - 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, - }; - if (tree_width < 1 || - std::find(std::begin(tree_buckets), std::end(tree_buckets), - n_tree_seqs) == std::end(tree_buckets) || + size_t graph_capacity = 0; + if (!target_paged_tree_graph_capacity( + tree_width, n_tree_seqs, graph_capacity) || cache.n_seq_slots <= 1 || !cache.paged_block_table || !cache.paged_kv_seq_lens || paged_max_kv_len < 1 || tree_scratch_base <= 0 || @@ -43,8 +86,7 @@ bool detail::validate_target_paged_tree_layout( (int64_t)tree_scratch_base + (int64_t)(cache.n_seq_slots - 1) * tree_scratch_stride + tree_width; - return scratch_end <= physical_kv_rows && - (int64_t)tree_width * n_tree_seqs <= INT32_MAX; + return scratch_end <= physical_kv_rows; } // ── build_layer_step ──────────────────────────────────────────── @@ -370,6 +412,11 @@ bool build_target_step( } if (segment_total != n_prefill_tokens) return false; if (n_logits_rows > 0 && n_prefill_tokens == 0) return false; + size_t graph_capacity = 0; + if (!detail::target_graph_capacity_for_parallel_segments( + n_prefill_segments, graph_capacity)) { + return false; + } // Persistent thread_local arena: rebuilt step graphs land at identical // addresses, keeping the ggml-cuda CUDA-graph cache key (nodes[0]) and @@ -535,7 +582,7 @@ bool build_target_step( ggml_set_input(sg.target_feat_rows); } - sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); + sg.gf = ggml_new_graph_custom(sg.ctx, graph_capacity, false); // Step-invariant KV write: only when topology can't vary per step. // DFLASH_QWEN35_NO_KVPAD=1 restores the legacy cpy append + exact-length @@ -696,6 +743,11 @@ bool build_target_step_paged_tree( tree_scratch_base, tree_scratch_stride)) { return false; } + size_t graph_capacity = 0; + if (!detail::target_paged_tree_graph_capacity( + tree_width, n_tree_seqs, graph_capacity)) { + return false; + } const int n_tokens = tree_width * n_tree_seqs; ggml_init_params ip{}; @@ -748,7 +800,7 @@ bool build_target_step_paged_tree( ggml_set_input(input.tensor); } - sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); + sg.gf = ggml_new_graph_custom(sg.ctx, graph_capacity, false); QwenGraphInputs gi{}; gi.inp_embed = sg.inp_embed; gi.positions = sg.positions; @@ -783,7 +835,8 @@ bool build_target_step_paged_tree( sg.alloc = ggml_gallocr_new( ggml_backend_get_default_buffer_type(backend)); } - return ggml_gallocr_alloc_graph(sg.alloc, sg.gf); + return ggml_gallocr_alloc_graph(sg.alloc, sg.gf) && + detail::target_paged_tree_uploads_ready(sg); } diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index fac406863..cbe58b786 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -25,6 +25,20 @@ namespace dflash::common { namespace detail { +// Qwen's recurrent graph duplicates one small subgraph per ragged sequence. +// Return a graph capacity that covers every supported concurrent bucket while +// keeping the legacy allocation for the common <= 8-sequence case. +bool target_graph_capacity_for_parallel_segments( + int n_parallel_segments, + size_t & capacity); + +// Checked packed-tree shape/capacity contract. The public DDTree budget allows +// at most 255 children plus the root, and concurrent serving at most 64 slots. +bool target_paged_tree_graph_capacity( + int tree_width, + int n_tree_seqs, + size_t & capacity); + // Model-free validation shared by the packed-tree builder and its shape // tests. paged_max_kv_len is a logical launch bound and may exceed the // bounded physical K/V pool; only the per-slot scratch slabs must fit in the @@ -37,6 +51,28 @@ bool validate_target_paged_tree_layout( int tree_scratch_base, int tree_scratch_stride); +// `active_slot_ids` is a topology marker in mapped-tree graphs. It may be +// optimized out by gallocr because the actual recurrent and attention row +// mappings are carried by state_slot_ids and paged_query_seq_ids. Every other +// tensor listed here is read by a graph node and must have backend storage +// before the engine uploads metadata. +inline bool target_paged_tree_uploads_ready(const StepGraph & sg) { + const auto allocated = [](const ggml_tensor * tensor) { + return tensor && tensor->buffer; + }; + return sg.active_slot_ids && + allocated(sg.inp_embed) && allocated(sg.positions) && + allocated(sg.parent_ids) && allocated(sg.tree_sizes) && + allocated(sg.state_slot_ids) && + allocated(sg.paged_query_seq_ids) && + allocated(sg.kv_write_rows); +} + +inline bool target_paged_tree_active_slots_need_upload( + const StepGraph & sg) { + return sg.active_slot_ids && sg.active_slot_ids->buffer; +} + } // namespace detail // Layer-segmented prefill: process one target layer for chunk_start..chunk_start+n_tokens. diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index 05e8a59cb..b0bb0475b 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -12,6 +12,7 @@ using namespace CppUnitTestFramework; using dflash::common::TargetCache; +using dflash::common::StepGraph; using dflash::common::restore_ssm_state; using dflash::common::snapshot_ssm_state; @@ -33,6 +34,90 @@ static std::vector get_tensor(const ggml_tensor * tensor) { return values; } +TEST_CASE(RecurrentSnapshotFixture, hardens_feature_smoke_paths) { + size_t graph_capacity = 0; + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 0, graph_capacity) && graph_capacity == 16384); + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 8, graph_capacity) && graph_capacity == 16384); + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 16, graph_capacity) && graph_capacity == 32768); + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 64, graph_capacity) && graph_capacity == 131072); + CHECK(!dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 65, graph_capacity)); + CHECK(dflash::common::detail::target_paged_tree_graph_capacity( + 23, 16, graph_capacity) && graph_capacity == 32768); + CHECK(!dflash::common::detail::target_paged_tree_graph_capacity( + 257, 16, graph_capacity)); + + // Mapped-tree active_slot_ids is a topology marker and can legitimately + // be left without gallocr storage. Required state/query/write metadata + // remains allocated and uploadable. + { + ggml_backend_t tree_backend = ggml_backend_cpu_init(); + CHECK(tree_backend != nullptr); + ggml_init_params marker_params{}; + marker_params.mem_size = 4 * ggml_tensor_overhead(); + marker_params.no_alloc = true; + ggml_context * marker_ctx = ggml_init(marker_params); + ggml_init_params live_params{}; + live_params.mem_size = 16 * ggml_tensor_overhead(); + live_params.no_alloc = true; + ggml_context * live_ctx = ggml_init(live_params); + CHECK(marker_ctx != nullptr); + CHECK(live_ctx != nullptr); + if (tree_backend && marker_ctx && live_ctx) { + StepGraph tree; + tree.active_slot_ids = + ggml_new_tensor_1d(marker_ctx, GGML_TYPE_I32, 2); + ggml_tensor * unallocated_state_ids = + ggml_new_tensor_1d(marker_ctx, GGML_TYPE_I32, 2); + tree.inp_embed = ggml_new_tensor_2d( + live_ctx, GGML_TYPE_F32, 4, 4); + tree.positions = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 16); + tree.parent_ids = + ggml_new_tensor_2d(live_ctx, GGML_TYPE_I32, 2, 2); + tree.tree_sizes = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 2); + tree.state_slot_ids = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 2); + tree.paged_query_seq_ids = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 4); + tree.kv_write_rows = + ggml_new_tensor_2d(live_ctx, GGML_TYPE_I64, 4, 1); + ggml_backend_buffer_t live_buffer = + ggml_backend_alloc_ctx_tensors(live_ctx, tree_backend); + CHECK(live_buffer != nullptr); + if (live_buffer) { + CHECK(tree.active_slot_ids->buffer == nullptr); + CHECK(dflash::common::detail:: + target_paged_tree_uploads_ready(tree)); + CHECK(!dflash::common::detail:: + target_paged_tree_active_slots_need_upload(tree)); + + const int32_t state_ids[] = {0, 1}; + ggml_backend_tensor_set(tree.state_slot_ids, state_ids, 0, + sizeof(state_ids)); + tree.state_slot_ids = unallocated_state_ids; + CHECK(!dflash::common::detail:: + target_paged_tree_uploads_ready(tree)); + ggml_backend_buffer_free(live_buffer); + } + } + if (live_ctx) ggml_free(live_ctx); + if (marker_ctx) ggml_free(marker_ctx); + if (tree_backend) ggml_backend_free(tree_backend); + } + +} + TEST_CASE(RecurrentSnapshotFixture, validates_paged_tree_layout) { // The packed-tree launch length is logical. KVFlash may keep a much // smaller physical resident pool, provided every tree scratch slab still @@ -70,6 +155,37 @@ TEST_CASE(RecurrentSnapshotFixture, snapshot_and_restore_recurrent_state) { CHECK(backend != nullptr); if (!backend) SKIP("CPU backend is unavailable"); + // C16 x width-23 selects a 32K graph. Prove that graph traversal and + // gallocr can cross the old 16K hard ceiling without asserting. + { + CHECK(dflash::common::detail::target_paged_tree_graph_capacity( + 23, 16, graph_capacity)); + ggml_init_params graph_params{}; + graph_params.mem_size = 32 * 1024 * 1024; + graph_params.no_alloc = true; + ggml_context * graph_ctx = ggml_init(graph_params); + CHECK(graph_ctx != nullptr); + if (graph_ctx) { + ggml_tensor * input = + ggml_new_tensor_1d(graph_ctx, GGML_TYPE_F32, 1); + ggml_set_input(input); + ggml_cgraph * graph = ggml_new_graph_custom( + graph_ctx, graph_capacity, false); + for (int i = 0; i < 16385; ++i) { + ggml_build_forward_expand( + graph, ggml_dup(graph_ctx, input)); + } + CHECK(ggml_graph_n_nodes(graph) == 16385); + ggml_gallocr_t graph_alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + CHECK(graph_alloc != nullptr); + CHECK(graph_alloc && + ggml_gallocr_alloc_graph(graph_alloc, graph)); + if (graph_alloc) ggml_gallocr_free(graph_alloc); + ggml_free(graph_ctx); + } + } + ggml_init_params params{}; params.mem_size = 8 * ggml_tensor_overhead(); params.no_alloc = true; From 5f7c5df392bffe9316e0e0a56244c7466e0b616d Mon Sep 17 00:00:00 2001 From: Graffioh Date: Thu, 13 Aug 2026 00:04:36 +0000 Subject: [PATCH 04/42] feat(qwen36): tune concurrent serving for Strix Halo --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 30 +++- .../concurrency/STRIX_HALO_RESULTS.md | 134 ++++++++++++++++++ .../concurrency/concurrent_benchmark.py | 11 +- .../feature_concurrent_benchmark.py | 8 +- .../concurrency/record_feature_runtime.py | 2 +- .../concurrency/run_qwen36_concurrency.sh | 78 ++++++++-- .../concurrency/run_qwen36_feature_matrix.sh | 32 ++++- .../concurrency/summarize_concurrency.py | 30 ++-- .../concurrency/summarize_feature_matrix.py | 39 +++-- .../concurrency/test_concurrency_tools.py | 48 +++++++ .../concurrency/test_concurrent_benchmark.py | 16 +++ .../test_feature_concurrent_benchmark.py | 26 ++++ .../concurrency/test_feature_metadata.py | 45 +++++- .../concurrency/test_feature_tools.py | 36 ++++- .../concurrency/write_feature_metadata.py | 53 ++++++- .../common/concurrency/paged_kv_residency.cpp | 81 ++++++++--- .../common/concurrency/paged_kv_residency.h | 1 + .../concurrency/qwen_paged_kv_transfer.cpp | 6 +- server/src/common/concurrency/seq_engine.h | 5 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 42 +++++- .../concurrency/qwen35_slot_manager.cpp | 42 +++++- .../qwen35/concurrency/qwen35_slot_manager.h | 26 +++- server/src/qwen35/qwen35_backend.cpp | 7 + server/src/server/scheduler.cpp | 3 + server/test/test_ddtree_path.cpp | 24 ++-- server/test/test_paged_kv_residency.cpp | 105 +++++++++++++- server/test/test_seq_batch_plan.cpp | 13 ++ server/test/test_seq_slot_manager.cpp | 114 +++++++++++++++ 28 files changed, 955 insertions(+), 102 deletions(-) create mode 100644 harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 7a1885f18..62fce6a46 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -1,5 +1,8 @@ # Qwen3.6 concurrent feature matrix +The bounded Strix Halo measurements collected for the draft implementation are +recorded in [`STRIX_HALO_RESULTS.md`](STRIX_HALO_RESULTS.md). + `run_qwen36_feature_matrix.sh` extends the PR #596 protocol with feature ablations for the complete Strix Halo configuration: @@ -19,12 +22,25 @@ Run the default bounded C4 screening repeat (seven applicable fresh-server cases ```bash MODEL=/opt/models/Qwen3.6-27B-Q4_K_M.gguf \ -DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q4_k_m.gguf \ +DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q8_0.gguf \ PREFILL_DRAFTER=/opt/models/Qwen3-0.6B-BF16.gguf \ REPEATS=1 \ harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh ``` +For the canonical AMD Strix Halo recipe, use the published Q8_0 3.6 drafter and make the tuning explicit: + +```bash +MODEL=/opt/models/Qwen3.6-27B-Q4_K_M.gguf \ +DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q8_0.gguf \ +PREFILL_DRAFTER=/opt/models/Qwen3-0.6B-BF16.gguf \ +DRAFT_SWA=2048 PREFILL_UBATCH=512 DDTREE_ADAPTIVE=0 \ +VARIANTS=ddtree,pflash,kvflash,full CLIENTS=1,4,8,16 REPEATS=5 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +`DDTREE_ADAPTIVE=0` matches the blog's continuous DDTree probe policy; leave it at the default `1` when measuring the concurrent engine's adaptive fallback policy. The concurrent path now records a startup `[parallel-ddtree]` marker and per-request `ddtree_steps`; these are the proof that DDTree actually ran. + On a 128 GiB Strix Halo host, budget roughly 45–90 minutes for this smoke run; the long-context AR controls dominate and actual time depends on the build. Every row remains independently selectable through `VARIANTS`. For example: @@ -77,8 +93,9 @@ The server must write one JSON object per completed request with this prefix: ``` Required fields are `effective_prompt_tokens`, `ddtree_steps`, -`ddtree_accepted_tokens`, `target_forwards`, `kvflash_page_ins`, -`kvflash_page_outs`, `kvflash_resident_blocks`, `kvflash_reselects`, +`ddtree_suspensions`, `ddtree_accepted_tokens`, `target_forwards`, +`kvflash_page_ins`, `kvflash_page_outs`, `kvflash_resident_blocks`, +`kvflash_reselects`, `pflash_applied`, `pflash_input_tokens`, and `pflash_output_tokens`. The proof tool correlates log objects with measured SSE request IDs and also @@ -87,6 +104,10 @@ checks the log's effective token count against unless: - DDTree has positive step and target-forward counts. Acceptance may be zero. + The required per-request suspension counter must be either zero or one. It + records adaptive fallback activation, but does not prove whether AR work ran + before or after the suspension; temporal claims require direct ordered-log + evidence. - PFlash reports `pflash_applied=true`, a smaller output prompt, and (in auto mode) an input token count at or above the recorded activation threshold. - KVFlash always records an explicit hashed scorer drafter, reports its @@ -108,6 +129,3 @@ binary/shared-library/target/draft/PFlash-and-KV-scorer hashes, the ordered `literal_screenshot_flags` array, all feature values, raw request report, server log, and `feature-proof.json`. The summary refuses to include a Lucebox row whose proof is missing or invalid. - - - diff --git a/harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md b/harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md new file mode 100644 index 000000000..c42eea4f9 --- /dev/null +++ b/harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md @@ -0,0 +1,134 @@ +# Qwen3.6 concurrent feature results — Strix Halo + +- Date: 2026-08-13 +- Implementation: `568fbac03b498d53d6efc0b2ab5893044543a321` +- Stack base: PR #595 head `a90ffe45c1d4ad58f5f73c4107571d3cf6c51bfd` + +These are bounded engineering measurements for the draft PR, not the +five-repeat publication matrix described in `FEATURE_MATRIX.md`. The paired +AR/DDTree screen has three fresh-process repeats. The long-context activation +rows have one fresh-process repeat per concurrency level because they are much +more expensive; treat their throughput as screening data. + +## System and artifacts + +- AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`), 128 GiB unified memory. +- ROCm runtime 7.2.4. +- Release HIP build for `gfx1151` with + `DFLASH27B_HIP_SM80_EQUIV=ON`. +- Server SHA-256: + `c77b4d2c7d1505fcc751600a6603cd65e51514b685bc66c4f7d33cd64a87c8a6`. +- Target SHA-256: + `5ed60d0af4650a854b1755bd392f9aef4872643dc25a254bc68043fa638392a0`. +- Decode draft SHA-256: + `e2500e90165a0f8e7b52c9882c29ed1fa391c60b300ff11b817bf10e31fa092e`. +- PFlash/KV scorer drafter SHA-256: + `f9c9f1d3c1e21755b82d4e165f88dbbbd4355646d632fb5d6cef7c66ed4ee04e`. + +Every case started a fresh server, discarded an 8-token same-concurrency +warmup, then requested exactly 64 output tokens per request with temperature +zero, seed one, and EOS ignored. Prompts were deterministic, disjoint across +concurrency levels, and identical between paired variants. The runner rotated +variant order across repeats. + +`Output-window` counts all completion tokens from the earliest first output to +the last completion. `Goodput` counts completion tokens over the whole level, +including TTFT. Every reported row passed exact token accounting and the +request-ID-correlated feature proof. + +The retained screening artifacts contain maximum TTFT but not median TTFT. +Their max-only columns below are an explicit screening exception, not a +protocol-complete publication result; a publication rerun must report both. + +## Paired AR and adaptive DDTree + +The DDTree configuration adds the local decode draft, budget 22, and target and +draft placement on `hip:0`. Values are medians over three fresh-process +repeats. + +| C | Variant | N | Goodput tok/s | Output-window tok/s | vs AR goodput | Accepted/step | Steps/suspensions | Max TTFT s | Output hashes stable | +| ---: | :--- | ---: | ---: | ---: | ---: | ---: | :--- | ---: | :---: | +| 1 | AR | 3 | 9.41 | 12.57 | — | — | 0/0 | 1.707 | yes | +| 1 | DDTree | 3 | 9.00 | 11.85 | -4.4% | 1.00 | 1/1 | 1.715 | yes | +| 4 | AR | 3 | 20.46 | 36.21 | — | — | 0/0 | 5.508 | no | +| 4 | DDTree | 3 | 19.43 | 33.61 | n/a | 3.08 | 4/4 | 5.533 | no | +| 8 | AR | 3 | 27.44 | 65.93 | — | — | 0/0 | 11.008 | no | +| 8 | DDTree | 3 | 25.63 | 56.32 | n/a | 2.79 | 8/8 | 11.076 | no | +| 16 | AR | 3 | 31.82 | 58.79 | — | — | 0/0 | 22.487 | no | +| 16 | DDTree | 3 | 29.36 | 54.63 | n/a | 2.00 | 16/16 | 22.514 | no | + +The supplied draft had weak acceptance on this cohort. The adaptive policy +sampled one real packed-tree step, then suspended the whole cohort because its +aggregate emitted yield was below six tokens per request. At C4 and above, +the raw timings are retained only to diagnose this fallback behavior; unstable +outputs do not support a performance comparison with AR. + +At C4 and above, greedy text hashes varied across fresh repeats in both the AR +control and DDTree. C1 was byte-stable. These measurements therefore establish +exact token accounting and feature execution, but do not claim bitwise text +reproducibility for concurrent batches. + +## Full screenshot configuration + +These rows enable the complete requested product configuration: + +```text +--target-device hip:0 +--draft-device hip:0 +--ddtree +--ddtree-budget 22 +--draft-residency persistent +--prefill-compression auto +--prefill-drafter /opt/models/Qwen3-0.6B-BF16.gguf +--kvflash auto +``` + +The controlled runner sets the auto PFlash threshold to 32K tokens, the keep +ratio to 0.05, and the KVFlash resident cap to 8,192 tokens. Startup telemetry +confirmed 512 physical blocks of 16 tokens, 16 configured slots, and a 65,536 +logical-token bound per slot. + +| C | N | Goodput tok/s | Output-window tok/s | Request decode tok/s | Raw prompt range | Effective prompt range | Max TTFT s | DDTree steps/susp. | KV page in/out | PFlash requests | +| ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | :--- | :--- | ---: | +| 1 | 1 | 2.61 | 12.39 | 12.19 | 41,504 | 2,021 | 19.324 | 1/1 | 0/0 | 1 | +| 4 | 1 | 3.00 | 31.61 | 7.95 | 38,142–44,866 | 1,870–2,235 | 77.495 | 4/4 | 1/18 | 4 | +| 8 | 1 | 3.13 | 52.60 | 6.56 | 38,141–44,870 | 1,869–2,237 | 153.782 | 8/8 | 245/792 | 8 | +| 16 | 1 | 3.19 | 18.73 | 2.34 | 38,140–44,872 | 1,867–2,238 | 307.360 | 16/16 | 293/1,880 | 16 | + +All four rows proved DDTree, PFlash, and KVFlash active. PFlash retained about +4.9% of raw prompt tokens. Output-window throughput scaled through C8, then +dropped at C16 while roughly 32K effective prompt tokens shared the 8K resident +pool; the concurrent page traffic rose accordingly. + +## Feature ablations + +| Workload | C | Variant | N | Goodput tok/s | Output-window tok/s | Request decode tok/s | Effective/raw | Max TTFT s | Activation evidence | +| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | :--- | +| compression | 4 | PFlash | 1 | 3.15 | 35.54 | 8.89 | 0.049 | 74.261 | 4/4 prompts compressed, 166,016 -> 8,179 tokens | +| kv-pressure | 4 | KVFlash | 1 | 1.24 | 8.56 | 5.32 | 1.000 | 197.859 | 129 resident blocks max, 0 page-ins / 3,714 page-outs | + +The PFlash-only row uses the same C4 prompts as the full row; adding DDTree and +KVFlash reduced output-window throughput from 35.54 to 31.61 tok/s in this +single screening repeat. The KVFlash-only row deliberately disables PFlash and +uses 13,474–20,203-token histories against the 8K pool. It is an activation and +pressure test, not a recommended latency configuration. + +## Reproduction + +The exact per-case command, controlled environment, startup-observed pool, +binary/shared-library/model hashes, raw request report, server log, and +`feature-proof.json` are retained by the runner. The principal invocations were: + +```bash +WORKLOADS=short CLIENTS=1,4,8,16 VARIANTS=ar,ddtree MAX_TOKENS=64 REPEATS=3 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +WORKLOADS=compression CLIENTS=1,4,8,16 VARIANTS=full MAX_TOKENS=64 REPEATS=1 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +WORKLOADS=compression CLIENTS=4 VARIANTS=pflash MAX_TOKENS=64 REPEATS=1 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +WORKLOADS=kv-pressure CLIENTS=4 VARIANTS=kvflash MAX_TOKENS=64 REPEATS=1 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +Set `MODEL`, `DRAFT_MODEL`, `PREFILL_DRAFTER`, and `LUCE_SERVER_BIN` as shown +in `FEATURE_MATRIX.md`. For publication-quality claims, rerun the documented +five-repeat 256-token matrix. diff --git a/harness/benchmarks/concurrency/concurrent_benchmark.py b/harness/benchmarks/concurrency/concurrent_benchmark.py index 773260393..4c3e85774 100755 --- a/harness/benchmarks/concurrency/concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/concurrent_benchmark.py @@ -164,11 +164,16 @@ def worker(index: int) -> None: threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(clients)] for thread in threads: thread.start() + deadline = time.monotonic() + args.timeout + 30 for thread in threads: - thread.join(args.timeout + 30) - completed = [record for record in records if record is not None] + thread.join(max(0.0, deadline - time.monotonic())) hung = sum(thread.is_alive() for thread in threads) - failures = hung + sum(record["error"] is not None for record in completed) + if hung: + raise TimeoutError( + f"{hung} request worker(s) exceeded the level deadline" + ) + completed = [record for record in records if record is not None] + failures = sum(record["error"] is not None for record in completed) ok = [record for record in completed if record["error"] is None] starts = [record["t_start"] for record in completed] ends = [record["t_end"] for record in completed] diff --git a/harness/benchmarks/concurrency/feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/feature_concurrent_benchmark.py index bc509c590..c9df122e1 100755 --- a/harness/benchmarks/concurrency/feature_concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/feature_concurrent_benchmark.py @@ -171,9 +171,9 @@ def markdown(report: dict[str, Any]) -> str: lines = [ f"# Concurrent feature benchmark — {report['label']}", "", "| C | Ok | Output goodput tok/s | Output-window tok/s | " - "Request decode tok/s | Wire prompt range | Effective prompt range | " - "Effective/wire | TTFT max s |", - "| ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | ---: |", + "Request decode tok/s | Prompt tok/s to first | Wire prompt range | " + "Effective prompt range | Effective/wire | TTFT median/max s |", + "| ---: | ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | :--- |", ] for level in report["levels"]: lines.append( @@ -181,10 +181,12 @@ def markdown(report: dict[str, Any]) -> str: f"{base.fmt(level['aggregate_tok_s'])} | " f"{base.fmt(level['output_window_tok_s'])} | " f"{base.fmt(level['request_decode_tok_s_median'])} | " + f"{base.fmt(level['prompt_tokens_per_s_to_first_token'])} | " f"{base.fmt(level['prompt_tokens_min'], '.0f')}–{base.fmt(level['prompt_tokens_max'], '.0f')} | " f"{base.fmt(level['effective_prompt_tokens_min'], '.0f')}–" f"{base.fmt(level['effective_prompt_tokens_max'], '.0f')} | " f"{base.fmt(level['effective_to_wire_prompt_ratio'], '.3f')} | " + f"{base.fmt(level['ttft_median_s'], '.3f')}/" f"{base.fmt(level['ttft_max_s'], '.3f')} |" ) return "\n".join(lines) + "\n" diff --git a/harness/benchmarks/concurrency/record_feature_runtime.py b/harness/benchmarks/concurrency/record_feature_runtime.py index 5665a4961..afb4fea6e 100644 --- a/harness/benchmarks/concurrency/record_feature_runtime.py +++ b/harness/benchmarks/concurrency/record_feature_runtime.py @@ -83,7 +83,7 @@ def update_metadata(metadata: dict[str, Any], log_text: str) -> dict[str, Any]: raise ValueError( "KVFlash metadata is enabled but its physical-pool startup marker is missing" ) - if observed["physical_kv_pool_tokens"] is None: + if not observed["proof_sources"]["paged_pool_startup_marker"]: raise ValueError("paged physical-pool startup marker is missing") result = dict(metadata) diff --git a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh index b6f76fe95..62c1e513e 100755 --- a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh +++ b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh @@ -55,6 +55,20 @@ MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" IFS=, read -r -a workload_list <<< "$WORKLOADS" IFS=, read -r -a variant_list <<< "$VARIANTS" IFS=, read -r -a client_list <<< "$CLIENTS" +reject_duplicates() { + local list_name="$1" value + shift + local -A seen=() + for value in "$@"; do + if [[ -n "${seen[$value]+yes}" ]]; then + echo "$list_name contains duplicate entry: $value" >&2 + return 1 + fi + seen["$value"]=1 + done +} +reject_duplicates CLIENTS "${client_list[@]}" || exit 2 +reject_duplicates VARIANTS "${variant_list[@]}" || exit 2 declare -A prompt_offsets=([1]=0 [4]=1 [8]=5 [16]=13) for c in "${client_list[@]}"; do [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } @@ -83,16 +97,55 @@ stop_server() { } trap stop_server EXIT INT TERM +served_model_matches() { + python3 - "$PORT" "$1" <<'PY' +import json +import sys +import urllib.request + +port, expected = sys.argv[1:] +with urllib.request.urlopen( + f"http://127.0.0.1:{port}/v1/models", timeout=2, +) as response: + payload = json.load(response) +matches = any( + isinstance(row, dict) and row.get("id") == expected + for row in payload.get("data", []) +) +raise SystemExit(0 if matches else 1) +PY +} + wait_health() { + local expected_model="$1" local deadline=$((SECONDS + HEALTH_TIMEOUT_SECONDS)) while (( SECONDS < deadline )); do kill -0 "$server_pid" 2>/dev/null || return 1 - curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && return 0 + if curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && + served_model_matches "$expected_model" >/dev/null 2>&1; then + return 0 + fi sleep 1 done return 1 } +port_is_available() { + python3 - "$PORT" <<'PY' +import socket +import sys + +port = int(sys.argv[1]) +with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind(("127.0.0.1", port)) + except OSError as exc: + print(f"PORT {port} is unavailable: {exc}", file=sys.stderr) + raise SystemExit(1) +PY +} + write_metadata() { local path="$1" variant="$2" workload="$3" clients="$4" repeat="$5" binary="$6" max_prefills="$7" command_file="$8" python3 -c 'import hashlib,json,pathlib,subprocess,sys @@ -132,12 +185,13 @@ run_case() { capacity=$((SLOTS * max_ctx)) local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" mkdir -p "$case_dir" - local -a command + local -a command launch_env if [[ "$variant" == llama ]]; then binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") + launch_env=() else binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce [[ "$variant" == luce-k8 ]] && max_prefills=8 || max_prefills=1 @@ -146,19 +200,27 @@ run_case() { --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 --host 127.0.0.1 --port "$PORT" --model-name "$model_id") + launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" + "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills") + fi + if ((${#launch_env[@]})); then + printf 'env ' > "$case_dir/server-command.txt" + printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$case_dir/server-command.txt" + else + printf '%q ' "${command[@]}" > "$case_dir/server-command.txt" fi - printf '%q ' "${command[@]}" > "$case_dir/server-command.txt"; printf '\n' >> "$case_dir/server-command.txt" + printf '\n' >> "$case_dir/server-command.txt" write_metadata "$case_dir/server-metadata.json" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$case_dir/server-command.txt" echo "[run] $workload C=$clients repeat=$repeat variant=$variant" - if [[ "$variant" == llama ]]; then - "${command[@]}" > "$case_dir/server.log" 2>&1 & + port_is_available || return 1 + if ((${#launch_env[@]})); then + env "${launch_env[@]}" "${command[@]}" > "$case_dir/server.log" 2>&1 & else - env DFLASH_MIN_TOKENS="$WARMUP_TOKENS" DFLASH_MAX_CONCURRENT_PREFILLS="$max_prefills" \ - "${command[@]}" > "$case_dir/server.log" 2>&1 & + "${command[@]}" > "$case_dir/server.log" 2>&1 & fi server_pid=$! - if ! wait_health; then tail -n 80 "$case_dir/server.log" >&2 || true; return 1; fi + if ! wait_health "$model_id"; then tail -n 80 "$case_dir/server.log" >&2 || true; return 1; fi local offset="${prompt_offsets[$clients]}" prompts="$OUT/prompts/$workload.jsonl" python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ diff --git a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh index d0ea1786a..c7a113916 100755 --- a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh @@ -28,6 +28,9 @@ MAX_TOKENS="${MAX_TOKENS:-64}" WARMUP_TOKENS="${WARMUP_TOKENS:-8}" SLOTS="${SLOTS:-16}" MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" +DRAFT_SWA="${DRAFT_SWA:-2048}" +PREFILL_UBATCH="${PREFILL_UBATCH:-512}" +DDTREE_ADAPTIVE="${DDTREE_ADAPTIVE:-1}" # The requested Strix Halo configuration. Every value is serialized into case # metadata; no performance-affecting DFLASH variable is inherited implicitly. @@ -45,13 +48,14 @@ usage() { cat <<'EOF' Usage: MODEL=/path/Qwen3.6-27B-Q4_K_M.gguf \ - DRAFT_MODEL=/path/dflash-draft-3.6-q4_k_m.gguf \ + DRAFT_MODEL=/path/dflash-draft-3.6-q8_0.gguf \ PREFILL_DRAFTER=/path/Qwen3-0.6B-BF16.gguf \ + DRAFT_SWA=2048 PREFILL_UBATCH=512 \ harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh The default is a bounded C4 smoke matrix with independently selectable ar, ddtree, pflash, kvflash, and full rows. The full row is the requested -Strix Halo configuration: target/draft hip:0, DDTree +Strix Halo configuration: target/draft hip:0, Q8 DFlash SWA=2048, DDTree budget 22, persistent PFlash auto, and KVFlash auto. The long-context profiles are intended to cross the recorded 32K PFlash and 8K KV-residency thresholds; word count is not treated as proof. Per-request wire/log token counts and @@ -59,6 +63,7 @@ activation telemetry fail the case if an "auto" feature did not execute. llama is optional: include it explicitly with VARIANTS=ar,ddtree,llama and set LLAMA_SERVER_BIN. For publication, set CLIENTS=1,4,8,16 and REPEATS=5. +For the AMD blog recipe, use the Q8_0 3.6 drafter, DRAFT_SWA=2048, PREFILL_UBATCH=512, and set DDTREE_ADAPTIVE=0 to keep DDTree active for every eligible step. OUT must not already exist. EOF } @@ -69,13 +74,15 @@ for cmd in python3 curl sha256sum awk; do command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; } done [[ -r "$MODEL" ]] || { echo "set MODEL to a readable target GGUF" >&2; exit 2; } -[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } [[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } [[ "$SLOTS" =~ ^[1-9][0-9]*$ ]] || { echo "SLOTS must be positive" >&2; exit 2; } [[ "$MAX_CONCURRENT_PREFILLS" =~ ^[1-9][0-9]*$ ]] || { echo "MAX_CONCURRENT_PREFILLS must be positive" >&2; exit 2; } [[ "$DDTREE_BUDGET" =~ ^[1-9][0-9]*$ ]] || { echo "DDTREE_BUDGET must be positive" >&2; exit 2; } [[ "$PREFILL_THRESHOLD" =~ ^[1-9][0-9]*$ ]] || { echo "PREFILL_THRESHOLD must be positive" >&2; exit 2; } [[ "$KVFLASH_MAX_POOL_TOKENS" =~ ^[1-9][0-9]*$ ]] || { echo "KVFLASH_MAX_POOL_TOKENS must be positive" >&2; exit 2; } +[[ "$DRAFT_SWA" =~ ^[0-9]+$ ]] || { echo "DRAFT_SWA must be a non-negative integer" >&2; exit 2; } +[[ "$PREFILL_UBATCH" =~ ^[1-9][0-9]*$ ]] || { echo "PREFILL_UBATCH must be positive" >&2; exit 2; } +[[ "$DDTREE_ADAPTIVE" == 0 || "$DDTREE_ADAPTIVE" == 1 ]] || { echo "DDTREE_ADAPTIVE must be 0 or 1" >&2; exit 2; } [[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ @@ -108,12 +115,17 @@ for c in "${client_list[@]}"; do [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } (( c <= SLOTS )) || { echo "CLIENTS=$c exceeds SLOTS=$SLOTS" >&2; exit 2; } done +luce_requested=0 for v in "${variant_list[@]}"; do case "$v" in - ar|ddtree|pflash|kvflash|full|llama) ;; + llama) ;; + ar|ddtree|pflash|kvflash|full) luce_requested=1 ;; *) echo "unknown variant $v" >&2; exit 2 ;; esac done +if (( luce_requested )); then + [[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +fi contains_variant() { local needle="$1" value @@ -232,11 +244,16 @@ run_case() { --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 --host 127.0.0.1 --port "$PORT" --model-name "$model_id") - launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills") + launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills" + "DFLASH27B_DRAFT_SWA=$DRAFT_SWA" "DFLASH27B_PREFILL_UBATCH=$PREFILL_UBATCH") + if [[ "$DDTREE_ADAPTIVE" == 0 ]]; then + launch_env+=("DFLASH_DDTREE_ADAPTIVE=0") + fi if [[ "$variant" == ddtree || "$variant" == full ]]; then command+=(--draft "$DRAFT_MODEL" --draft-device "$DRAFT_DEVICE" - --ddtree --ddtree-budget "$DDTREE_BUDGET") + --ddtree --ddtree-budget "$DDTREE_BUDGET" --fast-rollback) expected+=(--expect ddtree) + if [[ "$variant" == ddtree ]]; then command+=(--draft-residency "$DRAFT_RESIDENCY"); fi fi if [[ "$variant" == pflash || "$variant" == full ]]; then if [[ "$variant" == pflash ]]; then command+=(--draft-device "$DRAFT_DEVICE"); fi @@ -280,7 +297,8 @@ run_case() { fi if [[ "$variant" == ddtree || "$variant" == full ]]; then metadata+=(--draft-device "$DRAFT_DEVICE" --draft-model "$DRAFT_MODEL" - --draft-model-sha256 "$DRAFT_MODEL_SHA256" --ddtree --ddtree-budget "$DDTREE_BUDGET") + --draft-model-sha256 "$DRAFT_MODEL_SHA256" --ddtree --ddtree-budget "$DDTREE_BUDGET" --fast-rollback) + if [[ "$variant" == ddtree ]]; then metadata+=(--draft-residency "$DRAFT_RESIDENCY"); fi fi if [[ "$variant" == pflash || "$variant" == full ]]; then metadata+=(--draft-device "$DRAFT_DEVICE" --prefill-compression "$PREFILL_COMPRESSION" diff --git a/harness/benchmarks/concurrency/summarize_concurrency.py b/harness/benchmarks/concurrency/summarize_concurrency.py index 6f5130f41..0f4374bcb 100755 --- a/harness/benchmarks/concurrency/summarize_concurrency.py +++ b/harness/benchmarks/concurrency/summarize_concurrency.py @@ -43,6 +43,19 @@ def complete_median(values: list[float | None]) -> float | None: return median([value for value in values if value is not None]) +def output_stability(items: list[dict]) -> str: + output_digests = [ + item["level"].get("selected_output_set_sha256") for item in items + ] + complete = all(isinstance(value, str) and bool(value) for value in output_digests) + hashes = {value for value in output_digests if isinstance(value, str)} + return ( + "n/a" if len(items) < 2 or not complete + else "yes" if len(hashes) == 1 + else "NO" + ) + + def run_signature(item: dict) -> tuple[object, ...]: report, meta = item["report"], item["meta"] max_tokens = report.get("max_tokens") @@ -130,20 +143,7 @@ def summarize(reports: list[dict]) -> str: ttft = complete_median([ item["level"].get("ttft_max_s") for item in items ]) - output_digests = [ - item["level"].get("selected_output_set_sha256") for item in items - ] - output_digests_complete = all( - isinstance(value, str) and bool(value) for value in output_digests - ) - output_hashes = { - value for value in output_digests if isinstance(value, str) - } - stable = ( - "n/a" if len(items) < 2 or not output_digests_complete - else "yes" if len(output_hashes) == 1 - else "NO" - ) + stable = output_stability(items) def delta(other: str, metric: str) -> str: peers = grouped.get((workload, clients, other), []) @@ -164,6 +164,8 @@ def delta(other: str, metric: str) -> str: raise ValueError( f"{workload} C={clients}: {variant}/{other} repeat sets differ" ) + if stable == "NO" or output_stability(peers) == "NO": + return "n/a" ratios = [] for repeat in sorted(by_repeat): value = by_repeat[repeat]["level"].get(metric) diff --git a/harness/benchmarks/concurrency/summarize_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py index b5b29c01e..b789a5fbd 100755 --- a/harness/benchmarks/concurrency/summarize_feature_matrix.py +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -75,14 +75,18 @@ def run_signature(item: dict) -> tuple[object, ...]: report, meta = item["report"], item["meta"] max_tokens = report.get("max_tokens") ignore_eos = report.get("ignore_eos") + temperature = report.get("temperature") + seed = report.get("seed") model_sha256 = meta.get("model_sha256") if ( type(max_tokens) is not int or max_tokens <= 0 or not isinstance(ignore_eos, bool) + or type(temperature) not in (int, float) + or type(seed) is not int or not isinstance(model_sha256, str) or not model_sha256 ): raise ValueError("incomplete run metadata") - return max_tokens, ignore_eos, model_sha256 + return max_tokens, ignore_eos, temperature, seed, model_sha256 def output_stability(items: list[dict]) -> str: @@ -117,14 +121,22 @@ def summarize(reports: list[dict]) -> str: "proves its requested features executed. Throughput is the median across fresh-process repeats.", "", "| Workload | C | Variant | N | Output goodput | Output-window | vs AR | " - "Effective/wire | DDTree accepted/step | Target forwards | KV in/out | " - "PFlash requests | TTFT max s | Stable output |", - "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | " - ":--- | ---: | ---: | :---: |", + "Effective/wire | DDTree accepted/step | DDTree steps/susp. | " + "Target forwards | KV in/out | PFlash requests | TTFT max s | " + "Stable output |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " + ":--- | ---: | :--- | ---: | ---: | :---: |", ] for workload, clients, variant in sorted(grouped): items = grouped[(workload, clients, variant)] - prompt_hashes = {item["level"]["selected_prompt_set_sha256"] for item in items} + prompt_digests = [ + item["level"].get("selected_prompt_set_sha256") for item in items + ] + if not all(isinstance(value, str) and value for value in prompt_digests): + raise ValueError( + f"{workload} C={clients} {variant}: missing selected prompt set hash" + ) + prompt_hashes = set(prompt_digests) if len(prompt_hashes) != 1: raise ValueError(f"{workload} C={clients} {variant}: prompt sets differ") goodput = median([item["level"]["aggregate_tok_s"] for item in items]) @@ -174,6 +186,17 @@ def summarize(reports: list[dict]) -> str: steps = sum(a["ddtree_steps"] for a in aggregates) accepted = sum(a["ddtree_accepted_tokens"] for a in aggregates) accepted_per_step = accepted / steps if steps else None + median_steps = median( + [a["ddtree_steps"] for a in aggregates] + ) if aggregates else None + median_suspensions = median( + [a["ddtree_suspensions"] for a in aggregates] + ) if aggregates else None + ddtree_activity = ( + f"{median_steps:.0f}/{median_suspensions:.0f}" + if median_steps is not None and median_suspensions is not None + else "n/a" + ) target_forwards = median([a["target_forwards"] for a in aggregates]) if aggregates else None page_ins = median([a["kvflash_page_ins"] for a in aggregates]) if aggregates else None page_outs = median([a["kvflash_page_outs"] for a in aggregates]) if aggregates else None @@ -185,8 +208,8 @@ def summarize(reports: list[dict]) -> str: lines.append( f"| {workload} | {clients} | {variant} | {len(items)} | {goodput:.2f} | " f"{fmt(window)} | {vs_ar} | {fmt(ratio, 3)} | {fmt(accepted_per_step)} | " - f"{fmt(target_forwards, 0)} | {kv_text} | {fmt(pflash_requests, 0)} | " - f"{fmt(ttft, 3)} | {stable} |" + f"{ddtree_activity} | {fmt(target_forwards, 0)} | {kv_text} | " + f"{fmt(pflash_requests, 0)} | {fmt(ttft, 3)} | {stable} |" ) lines.append("") return "\n".join(lines) diff --git a/harness/benchmarks/concurrency/test_concurrency_tools.py b/harness/benchmarks/concurrency/test_concurrency_tools.py index 1aa8dc2d8..60bc925a9 100644 --- a/harness/benchmarks/concurrency/test_concurrency_tools.py +++ b/harness/benchmarks/concurrency/test_concurrency_tools.py @@ -50,6 +50,23 @@ def test_cohorts_are_disjoint_ragged_and_mean_matched(self) -> None: self.assertEqual(len(row["prompt"].split()), row["target_words"]) +class RunnerShellTests(unittest.TestCase): + def test_runner_guards_case_identity_and_records_launch_environment(self) -> None: + runner = (HERE / "run_qwen36_concurrency.sh").read_text( + encoding="utf-8", + ) + self.assertIn( + 'reject_duplicates CLIENTS "${client_list[@]}"', runner, + ) + self.assertIn( + 'reject_duplicates VARIANTS "${variant_list[@]}"', runner, + ) + self.assertIn('printf \'env \' > "$case_dir/server-command.txt"', runner) + self.assertIn('"${launch_env[@]}" "${command[@]}"', runner) + self.assertIn("port_is_available || return 1", runner) + self.assertIn('wait_health "$model_id"', runner) + + class SummarizerTests(unittest.TestCase): @staticmethod def item( @@ -136,6 +153,37 @@ def test_multiple_repeats_report_output_stability(self) -> None: unstable_row = next(line for line in unstable.splitlines() if "| llama |" in line) self.assertEqual(unstable_row.split("|")[10].strip(), "NO") + def test_unstable_current_variant_suppresses_deltas(self) -> None: + reports = [ + self.item("llama", 10.0, repeat=1), + self.item("llama", 10.0, repeat=2), + self.item("luce-k8", 20.0, repeat=1, output_hash="first"), + self.item("luce-k8", 22.0, repeat=2, output_hash="second"), + ] + text = summarizer.summarize(reports) + row = next(line for line in text.splitlines() if "| luce-k8 |" in line) + self.assertEqual(row.split("|")[10].strip(), "NO") + self.assertEqual(row.split("|")[11].strip(), "n/a") + self.assertEqual(row.split("|")[12].strip(), "n/a") + + def test_unstable_peer_variant_suppresses_deltas(self) -> None: + reports = [ + self.item( + "llama", 10.0, repeat=1, output_hash="first", + ), + self.item( + "llama", 10.0, repeat=2, output_hash="second", + ), + self.item("luce-k8", 20.0, repeat=1), + self.item("luce-k8", 22.0, repeat=2), + ] + text = summarizer.summarize(reports) + row = next(line for line in text.splitlines() if "| luce-k8 |" in line) + self.assertEqual(row.split("|")[10].strip(), "yes") + self.assertEqual(row.split("|")[11].strip(), "n/a") + self.assertEqual(row.split("|")[12].strip(), "n/a") + + def test_missing_output_digest_does_not_claim_stability(self) -> None: reports = [ self.item("llama", 8.0, repeat=1, output_hash=None), diff --git a/harness/benchmarks/concurrency/test_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_concurrent_benchmark.py index c0455edc3..da1251c79 100644 --- a/harness/benchmarks/concurrency/test_concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/test_concurrent_benchmark.py @@ -126,6 +126,22 @@ def test_missing_prompt_usage_fails_level(self) -> None: } self.assertTrue(benchmark.level_failed(level, ignore_eos=True)) + def test_hung_worker_aborts_level_instead_of_overlapping_next(self) -> None: + thread = mock.Mock() + thread.is_alive.return_value = True + args = argparse.Namespace(timeout=1.0) + with ( + mock.patch.object(benchmark.threading, "Thread", return_value=thread), + mock.patch.object( + benchmark.time, "monotonic", side_effect=(10.0, 50.0), + ), + ): + with self.assertRaisesRegex(TimeoutError, "exceeded the level deadline"): + benchmark.run_level(1, args, ["prompt"], 0) + thread.start.assert_called_once_with() + thread.join.assert_called_once_with(0.0) + + if __name__ == "__main__": unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py index d582f6e32..232057dcc 100644 --- a/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py @@ -118,6 +118,32 @@ def test_client_provenance_records_exact_argv_and_source_digest(self) -> None: hashlib.sha256(SCRIPT.read_bytes()).hexdigest(), ) + def test_markdown_reports_prompt_rate_and_ttft_median_max(self) -> None: + text = benchmark.markdown({ + "label": "feature", + "levels": [{ + "clients": 4, + "requests_ok": 4, + "requests": 4, + "aggregate_tok_s": 12.5, + "output_window_tok_s": 20.0, + "request_decode_tok_s_median": 5.0, + "prompt_tokens_per_s_to_first_token": 123.4, + "prompt_tokens_min": 100, + "prompt_tokens_max": 400, + "effective_prompt_tokens_min": 50, + "effective_prompt_tokens_max": 200, + "effective_to_wire_prompt_ratio": 0.5, + "ttft_median_s": 1.25, + "ttft_max_s": 2.5, + }], + }) + self.assertIn("Prompt tok/s to first", text) + self.assertIn("TTFT median/max s", text) + self.assertIn("| 4 | 4/4 | 12.50 | 20.00 | 5.00 | 123.40 |", text) + self.assertIn("| 0.500 | 1.250/2.500 |", text) + + if __name__ == "__main__": unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_metadata.py b/harness/benchmarks/concurrency/test_feature_metadata.py index 7349fa21a..229685030 100644 --- a/harness/benchmarks/concurrency/test_feature_metadata.py +++ b/harness/benchmarks/concurrency/test_feature_metadata.py @@ -41,6 +41,8 @@ def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: target.write_bytes(b"target") draft.write_bytes(b"draft") prefill.write_bytes(b"prefill") + draft_sha = hashlib.sha256(draft.read_bytes()).hexdigest() + prefill_sha = hashlib.sha256(prefill.read_bytes()).hexdigest() prompts.write_text('{"prompt":"p"}\n', encoding="utf-8") command.write_text("server --target-device hip:0\n", encoding="utf-8") argv = [ @@ -51,15 +53,15 @@ def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: "--prompt-file", str(prompts), "--command-file", str(command), "--repo", str(HERE.parents[2]), "--max-concurrent-prefills", "8", "--target-device", "hip:0", "--draft-device", "hip:0", - "--draft-model", str(draft), "--draft-model-sha256", "draft-sha", - "--ddtree", "--ddtree-budget", "22", + "--draft-model", str(draft), "--draft-model-sha256", draft_sha, + "--ddtree", "--ddtree-budget", "22", "--fast-rollback", "--prefill-compression", "auto", "--prefill-threshold", "32000", "--prefill-keep-ratio", "0.05", "--prefill-drafter", str(prefill), - "--prefill-drafter-sha256", "prefill-sha", + "--prefill-drafter-sha256", prefill_sha, "--draft-residency", "persistent", "--kvflash", "auto", "--kvflash-max-pool-tokens", "8192", "--kvflash-scorer-drafter", str(prefill), - "--kvflash-scorer-drafter-sha256", "prefill-sha", + "--kvflash-scorer-drafter-sha256", prefill_sha, ] with mock.patch.object(sys, "argv", argv): self.assertEqual(metadata.main(), 0) @@ -67,21 +69,22 @@ def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: self.assertEqual(result["model_sha256"], hashlib.sha256(b"target").hexdigest()) self.assertTrue(result["server_binary_sha256"]) self.assertTrue(result["git_head"]) - self.assertEqual(result["feature_config"]["draft_model_sha256"], "draft-sha") - self.assertEqual(result["feature_config"]["prefill_drafter_sha256"], "prefill-sha") + self.assertEqual(result["feature_config"]["draft_model_sha256"], draft_sha) + self.assertEqual(result["feature_config"]["prefill_drafter_sha256"], prefill_sha) self.assertEqual( result["feature_config"]["kvflash_scorer_drafter"], str(prefill.resolve()), ) self.assertEqual( result["feature_config"]["kvflash_scorer_drafter_sha256"], - "prefill-sha", + prefill_sha, ) self.assertEqual(result["literal_screenshot_flags"], [ "--target-device", "hip:0", "--draft-device", "hip:0", "--ddtree", "--ddtree-budget", "22", + "--fast-rollback", "--draft-residency", "persistent", "--prefill-compression", "auto", "--prefill-drafter", str(prefill.resolve()), @@ -89,6 +92,24 @@ def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: ]) self.assertIsNone(result["runtime_observed"]) + def test_model_digest_claim_must_match_referenced_file(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + model = Path(tmp) / "model.gguf" + model.write_bytes(b"model") + expected = hashlib.sha256(b"model").hexdigest() + cache: dict[Path, str] = {} + self.assertEqual( + metadata.validated_digest(model, expected, "model", cache), + expected, + ) + with self.assertRaisesRegex(ValueError, "does not match"): + metadata.validated_digest(model, "stale", "model", cache) + with self.assertRaisesRegex(ValueError, "is required"): + metadata.validated_digest(model, None, "model", cache) + with self.assertRaisesRegex(ValueError, "without a model file"): + metadata.validated_digest(None, expected, "model", cache) + + def test_ldd_failure_is_fatal(self) -> None: failed = mock.Mock(returncode=1, stdout="", stderr="not a dynamic executable") with mock.patch.object(metadata.subprocess, "run", return_value=failed): @@ -148,6 +169,16 @@ def test_runtime_rejects_enabled_kvflash_without_marker(self) -> None: "(8192 pool tokens, per-sequence max_ctx 65536)", ) + def test_runtime_rejects_kvflash_marker_without_paged_marker(self) -> None: + original = {"feature_config": {"kvflash": "auto"}} + log = ( + "[parallel-kvflash] physical resident pool 8192 tokens; " + "logical per-slot cap 65536 across 16 slots " + "(--kv-pool-tokens does not expand resident VRAM)" + ) + with self.assertRaisesRegex(ValueError, "paged physical-pool"): + runtime_metadata.update_metadata(original, log) + if __name__ == "__main__": unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index 85c2294e6..5e5645f7d 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -179,6 +179,17 @@ def test_duplicate_variants_are_rejected_before_artifacts_are_created(self) -> N self.assertIn("VARIANTS contains duplicate entry: ar", result.stderr) self.assertFalse((Path(tmp) / "out").exists()) + def test_llama_only_does_not_require_lucebox_binary(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self.run_invalid_matrix( + tmp, VARIANTS="llama", LUCE_SERVER_BIN="/does/not/exist", + LLAMA_SERVER_BIN="/bin/true", REPEATS="0", + ) + self.assertEqual(result.returncode, 2, result.stderr) + self.assertIn("REPEATS must be positive", result.stderr) + self.assertNotIn("missing Lucebox server", result.stderr) + + class FeatureProofTests(unittest.TestCase): def test_full_below_pool_passes_without_page_traffic(self) -> None: @@ -450,7 +461,10 @@ def item( output_hash: str | None = "same-output", ) -> dict: return { - "report": {"max_tokens": 256, "ignore_eos": True}, + "report": { + "max_tokens": 256, "ignore_eos": True, + "temperature": 0.0, "seed": 1, + }, "meta": { "workload": "compression", "variant": variant, "repeat": repeat, "model_sha256": "a" * 64, @@ -528,6 +542,26 @@ def test_summary_rejects_incompatible_ar_control_metadata(self) -> None: with self.assertRaisesRegex(ValueError, "run metadata differs"): summary.summarize([ar, feature]) + def test_summary_rejects_sampling_mismatch(self) -> None: + for field, value in (("temperature", 0.5), ("seed", 2)): + with self.subTest(field=field): + ar = self.item("ar", 10.0) + feature = self.item("full", 12.0) + feature["report"][field] = value + with self.assertRaisesRegex(ValueError, "run metadata differs"): + summary.summarize([ar, feature]) + + def test_summary_rejects_missing_prompt_hash(self) -> None: + for value in (None, ""): + with self.subTest(value=value): + feature = self.item("full", 12.0) + feature["level"]["selected_prompt_set_sha256"] = value + with self.assertRaisesRegex( + ValueError, "missing selected prompt set hash", + ): + summary.summarize([feature]) + + def test_unstable_feature_row_suppresses_ar_delta(self) -> None: reports = [ self.item("ar", 10.0, repeat=1), diff --git a/harness/benchmarks/concurrency/write_feature_metadata.py b/harness/benchmarks/concurrency/write_feature_metadata.py index 9500c41b9..01e40146f 100644 --- a/harness/benchmarks/concurrency/write_feature_metadata.py +++ b/harness/benchmarks/concurrency/write_feature_metadata.py @@ -14,6 +14,27 @@ def digest(path: pathlib.Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() +def validated_digest( + path: pathlib.Path | None, + claimed: str | None, + label: str, + cache: dict[pathlib.Path, str], +) -> str | None: + if path is None: + if claimed is not None: + raise ValueError(f"{label} SHA-256 supplied without a model file") + return None + if not isinstance(claimed, str) or not claimed: + raise ValueError(f"{label} SHA-256 is required") + resolved = path.resolve() + if resolved not in cache: + cache[resolved] = digest(resolved) + actual = cache[resolved] + if claimed != actual: + raise ValueError(f"{label} SHA-256 does not match {resolved}") + return actual + + def resolved_libraries(binary: pathlib.Path) -> dict[str, str]: result = subprocess.run( ["ldd", str(binary)], text=True, capture_output=True, @@ -70,6 +91,7 @@ def main() -> int: parser.add_argument("--draft-model", type=pathlib.Path) parser.add_argument("--draft-model-sha256") parser.add_argument("--ddtree", action="store_true") + parser.add_argument("--fast-rollback", action="store_true") parser.add_argument("--ddtree-budget", type=int) parser.add_argument("--prefill-compression", default="off") parser.add_argument("--prefill-threshold", type=int) @@ -91,6 +113,24 @@ def main() -> int: parser.error(f"bad --launch-env {item!r}; expected KEY=VALUE") launch_env[key] = value + digest_cache: dict[pathlib.Path, str] = {} + model_sha256 = validated_digest( + args.model, args.model_sha256, "target model", digest_cache, + ) + draft_model_sha256 = validated_digest( + args.draft_model, args.draft_model_sha256, "draft model", digest_cache, + ) + prefill_drafter_sha256 = validated_digest( + args.prefill_drafter, args.prefill_drafter_sha256, + "prefill drafter", digest_cache, + ) + kvflash_scorer_drafter_sha256 = validated_digest( + args.kvflash_scorer_drafter, + args.kvflash_scorer_drafter_sha256, + "KVFlash scorer drafter", + digest_cache, + ) + libraries = resolved_libraries(args.binary) git_head = repository_head(args.repo) literal_flags: list[str] = [] @@ -102,6 +142,8 @@ def main() -> int: literal_flags += ["--ddtree"] if args.ddtree_budget is not None: literal_flags += ["--ddtree-budget", str(args.ddtree_budget)] + if args.fast_rollback: + literal_flags += ["--fast-rollback"] if args.draft_residency: literal_flags += ["--draft-residency", args.draft_residency] if args.prefill_compression != "off": @@ -121,7 +163,7 @@ def main() -> int: "server_binary": str(args.binary.resolve()), "server_binary_sha256": digest(args.binary), "model": str(args.model.resolve()), - "model_sha256": args.model_sha256, + "model_sha256": model_sha256, "prompt_file_sha256": digest(args.prompt_file), "server_command": args.command_file.read_text(encoding="utf-8").strip(), "launch_environment": launch_env, @@ -134,8 +176,9 @@ def main() -> int: "target_device": args.target_device, "draft_device": args.draft_device, "draft_model": str(args.draft_model.resolve()) if args.draft_model else None, - "draft_model_sha256": args.draft_model_sha256, + "draft_model_sha256": draft_model_sha256, "ddtree": args.ddtree, + "fast_rollback": args.fast_rollback, "ddtree_budget": args.ddtree_budget, "prefill_compression": args.prefill_compression, "prefill_threshold": args.prefill_threshold, @@ -143,7 +186,7 @@ def main() -> int: "prefill_drafter": ( str(args.prefill_drafter.resolve()) if args.prefill_drafter else None ), - "prefill_drafter_sha256": args.prefill_drafter_sha256, + "prefill_drafter_sha256": prefill_drafter_sha256, "draft_residency": args.draft_residency, "kvflash": args.kvflash, "kvflash_max_pool_tokens": args.kvflash_max_pool_tokens, @@ -151,9 +194,7 @@ def main() -> int: str(args.kvflash_scorer_drafter.resolve()) if args.kvflash_scorer_drafter else None ), - "kvflash_scorer_drafter_sha256": ( - args.kvflash_scorer_drafter_sha256 - ), + "kvflash_scorer_drafter_sha256": kvflash_scorer_drafter_sha256, }, } args.out.write_text(json.dumps(obj, indent=2, sort_keys=True) + "\n", encoding="utf-8") diff --git a/server/src/common/concurrency/paged_kv_residency.cpp b/server/src/common/concurrency/paged_kv_residency.cpp index 43c8443b0..c0af73b12 100644 --- a/server/src/common/concurrency/paged_kv_residency.cpp +++ b/server/src/common/concurrency/paged_kv_residency.cpp @@ -243,10 +243,13 @@ PagedKvResidencyStatus PagedKvResidencyManager::prepare_append( } if (token_count == 0) return finish_transfers(PagedKvResidencyStatus::Ok); + BlockState * append_head = nullptr; bool restore_partial_head = false; uint32_t partial_head = 0; if (snapshot.kv_seq_len % pool_.block_size() != 0) { partial_head = snapshot.kv_seq_len / pool_.block_size(); + append_head = &sequences_[handle.slot].blocks[partial_head]; + append_head->reservation_pending = true; restore_partial_head = snapshot.block_table[partial_head] == PAGED_KV_COLD_BLOCK; } @@ -265,15 +268,22 @@ PagedKvResidencyStatus PagedKvResidencyManager::prepare_append( const auto room = make_room( handle, globally_needed + (restore_partial_head ? 1u : 0u), additional_blocks + (restore_partial_head ? 1u : 0u)); - if (room != PagedKvResidencyStatus::Ok) return finish_transfers(room); + if (room != PagedKvResidencyStatus::Ok) { + if (append_head) append_head->reservation_pending = false; + return room; + } if (restore_partial_head) { const auto restored = restore_block_async( handle, partial_head); if (restored != PagedKvResidencyStatus::Ok) { - return finish_transfers(restored); + const auto finished = finish_transfers(restored); + if (append_head) append_head->reservation_pending = false; + return finished; } } - return finish_transfers(PagedKvResidencyStatus::Ok); + const auto finished = finish_transfers(PagedKvResidencyStatus::Ok); + if (append_head) append_head->reservation_pending = false; + return finished; } PagedKvResidentAppendResult PagedKvResidencyManager::append( @@ -305,18 +315,26 @@ PagedKvResidencyStatus PagedKvResidencyManager::observe_append( !state.blocks[remap.logical_block].host_valid) { return finish_transfers(PagedKvResidencyStatus::HostCopyMissing); } + BlockState & block = state.blocks[remap.logical_block]; + // A callback may reject only after queuing a prefix of a multi-tensor + // copy. Quarantine the destination before invoking it either way. + transfers_pending_ = true; + block.page_in_pending = true; bool queued = false; try { queued = transfers_.copy_in_async( handle, remap.logical_block, remap.physical_block, - state.blocks[remap.logical_block].host, config_.block_bytes); + block.host, config_.block_bytes); } catch (...) { queued = false; } - if (!queued) return finish_transfers(PagedKvResidencyStatus::TransferFailed); - transfers_pending_ = true; - stats_.page_ins++; - stats_.moved_bytes += config_.block_bytes; + const PendingTransfer transfer{ + handle, remap.logical_block, remap.physical_block}; + if (!queued) { + pending_page_in_rollbacks_.push_back(transfer); + return finish_transfers(PagedKvResidencyStatus::TransferFailed); + } + pending_page_ins_.push_back(transfer); } const auto synced = finish_transfers(PagedKvResidencyStatus::Ok); if (synced != PagedKvResidencyStatus::Ok) return synced; @@ -424,8 +442,34 @@ PagedKvResidencyStatus PagedKvResidencyManager::synchronize_before_read() { stats_.page_ins++; stats_.moved_bytes += config_.block_bytes; } + for (const PendingTransfer & transfer : pending_page_in_rollbacks_) { + if (validate_registered(transfer.handle) != + PagedKvResidencyStatus::Ok) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + BlockState & block = + sequences_[transfer.handle.slot].blocks[transfer.logical_block]; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(transfer.handle, snapshot) != PagedKvStatus::Ok || + transfer.logical_block >= snapshot.block_table.size() || + snapshot.block_table[transfer.logical_block] != + transfer.physical_block) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + uint32_t released = PAGED_KV_COLD_BLOCK; + if (pool_.page_out_block( + transfer.handle, transfer.logical_block, released) != + PagedKvStatus::Ok || released != transfer.physical_block) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + block.page_in_pending = false; + } pending_page_outs_.clear(); pending_page_ins_.clear(); + pending_page_in_rollbacks_.clear(); if (result != PagedKvResidencyStatus::Ok) return result; return PagedKvResidencyStatus::Ok; } @@ -630,6 +674,10 @@ PagedKvResidencyStatus PagedKvResidencyManager::evict_block_async( const uint32_t physical = snapshot.block_table[logical_block]; bool queued = false; + // False may mean a multi-tensor callback queued only a prefix. Invalidate + // the old host image and require a barrier before any later operation. + transfers_pending_ = true; + block.host_valid = false; try { queued = transfers_.copy_out_async( handle, logical_block, physical, block.host, @@ -638,8 +686,6 @@ PagedKvResidencyStatus PagedKvResidencyManager::evict_block_async( queued = false; } if (!queued) return PagedKvResidencyStatus::TransferFailed; - transfers_pending_ = true; - block.host_valid = false; block.page_out_pending = true; pending_page_outs_.push_back({handle, logical_block, physical}); return PagedKvResidencyStatus::Ok; @@ -673,6 +719,8 @@ PagedKvResidencyStatus PagedKvResidencyManager::restore_block_async( if (page_status != PagedKvStatus::Ok) return from_pool_status(page_status); bool queued = false; + transfers_pending_ = true; + block.page_in_pending = true; try { queued = transfers_.copy_in_async( handle, logical_block, physical, block.host, @@ -680,17 +728,14 @@ PagedKvResidencyStatus PagedKvResidencyManager::restore_block_async( } catch (...) { queued = false; } + const PendingTransfer transfer{handle, logical_block, physical}; if (!queued) { - uint32_t released = PAGED_KV_COLD_BLOCK; - if (pool_.page_out_block(handle, logical_block, released) != - PagedKvStatus::Ok) { - return PagedKvResidencyStatus::InconsistentPoolState; - } + // The callback may have queued a prefix. Keep the physical mapping + // quarantined until finish_transfers() proves the stream is drained. + pending_page_in_rollbacks_.push_back(transfer); return PagedKvResidencyStatus::TransferFailed; } - transfers_pending_ = true; - block.page_in_pending = true; - pending_page_ins_.push_back({handle, logical_block, physical}); + pending_page_ins_.push_back(transfer); block.last_use = ++clock_; return PagedKvResidencyStatus::Ok; } diff --git a/server/src/common/concurrency/paged_kv_residency.h b/server/src/common/concurrency/paged_kv_residency.h index dc6f67cf5..ac2d982b0 100644 --- a/server/src/common/concurrency/paged_kv_residency.h +++ b/server/src/common/concurrency/paged_kv_residency.h @@ -255,6 +255,7 @@ class PagedKvResidencyManager { bool transfer_barrier_failed_ = false; std::vector pending_page_outs_; std::vector pending_page_ins_; + std::vector pending_page_in_rollbacks_; }; } // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer.cpp b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp index 111e6a36f..d9a2b74cf 100644 --- a/server/src/common/concurrency/qwen_paged_kv_transfer.cpp +++ b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp @@ -212,9 +212,9 @@ struct QwenPagedKvResidencyTransfer::State { } } if (status != cudaSuccess) { - // A callback that returns false is not marked pending by the - // residency manager. Drain any prefix already queued here so - // its pinned buffer can still be released safely. + // Drain any prefix already queued here. The residency manager + // also treats a false return as pending, so a failed drain is + // quarantined and retried before ownership can be released. (void)cudaGetLastError(); if (cudaStreamSynchronize(stream) != cudaSuccess) { stream_may_reference_host = true; diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index fa2909da3..6371e2552 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -202,6 +202,7 @@ class SeqEngine { // until retirement and emits one machine-readable proof record. uint64_t ddtree_steps = 0; uint64_t ddtree_accepted_tokens = 0; + uint64_t ddtree_suspensions = 0; uint64_t target_forwards = 0; uint64_t kvflash_page_ins = 0; uint64_t kvflash_page_outs = 0; @@ -329,6 +330,8 @@ inline std::string validate_step_result( return "failed decode has no diagnostic"; if (output.token >= 0 || !output.committed_tokens.empty()) return "failed decode exposes token payload"; + if (output.ddtree_suspensions != 0) + return "failed decode carries DDTree suspension telemetry"; } else { if (output.token < 0) return "successful decode has no pending token"; @@ -342,6 +345,8 @@ inline std::string validate_step_result( output.committed_tokens.end(), [](int32_t token) { return token < 0; })) return "decode output burst contains an invalid token"; + if (output.ddtree_suspensions > output.ddtree_steps) + return "DDTree suspension has no successful DDTree step"; } decode_seen[(size_t)output.slot] = 1; } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 10e50450a..52949acfc 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -154,6 +154,7 @@ bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { if (!in.allow_speculation || in.slot < 0 || in.slot >= slots_.slot_count() || !slots_.slot(in.slot).decoding() || + !slots_.ddtree_speculation_allowed(in.slot) || slots_.slot(in.slot).sampler.needs_logit_processing() || slots_.slot(in.slot).cur_pos < 1 || slots_.slot(in.slot).cur_pos >= slots_.max_context()) { @@ -234,6 +235,13 @@ std::optional Qwen35SeqEngine::step_ddtree( GGML_STATUS_SUCCESS) { return proposal_fallback(); } + // The draft and target backends own separate HIP streams even when + // both are placed on hip:0. Projection consumes the draft hidden + // state on the target stream, so establish the producer/consumer + // ordering explicitly before the cross-backend tensor copy. Without + // this barrier the first tree proposal can race stale hidden rows and + // collapse acceptance to the one-token fallback. + ggml_backend_synchronize(b_.draft_backend_); ggml_backend_tensor_copy( draft->hidden_states, b_.proj_sg_.hidden_input); if (ggml_backend_graph_compute( @@ -524,19 +532,51 @@ std::optional Qwen35SeqEngine::step_ddtree( return result; } } + + uint64_t cohort_emitted = 0; + for (const Proposal & p : proposals) { + // accepted contains the replay root plus accepted children. The + // output emits those children plus one separately computed pending + // scalar, so accepted.size() is this request's emitted yield. + cohort_emitted += (uint64_t)p.accepted.size(); + } + const bool suspend_cohort = + Qwen35SlotManager::ddtree_cohort_should_suspend( + cohort_emitted, active); + std::vector newly_suspended((size_t)slots_.slot_count(), false); + for (const Proposal & p : proposals) { + newly_suspended[(size_t)p.slot] = + slots_.record_ddtree_sample(p.slot, suspend_cohort); + } + result.decode.reserve((size_t)active); for (Proposal & p : proposals) { DecodeOutput out; out.slot = p.slot; out.token = p.bonus; out.ddtree_steps = 1; - out.ddtree_accepted_tokens = p.accepted.size() - 1; + const int accepted_children = (int)p.accepted.size() - 1; + out.ddtree_accepted_tokens = (uint64_t)accepted_children; out.target_forwards = 2; for (size_t i = 1; i < p.accepted.size(); ++i) { const int dfs = p.accepted[i]; out.committed_tokens.push_back( p.tree.token_ids[(size_t)dfs - 1]); } + if (newly_suspended[(size_t)p.slot]) { + out.ddtree_suspensions = 1; + const Qwen35Slot & seq = slots_.slot(p.slot); + std::fprintf(stderr, + "[parallel-ddtree] adaptive suspend request=%llu slot=%d " + "sample=%llu emitted=%d accepted_children=%d " + "cohort_emitted=%llu cohort_size=%d target_forwards=2 " + "floor=%d\n", + (unsigned long long)seq.request_id, p.slot, + (unsigned long long)seq.ddtree_sampled_steps, + accepted_children + 1, accepted_children, + (unsigned long long)cohort_emitted, active, + Qwen35SlotManager::kDdtreeMinEmittedTokens); + } attach_residency_telemetry(out); result.decode.push_back(std::move(out)); } diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index f740a283b..e0e64170d 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -2,6 +2,7 @@ #include #include +#include namespace dflash::common { @@ -138,9 +139,15 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } + // Retirement keeps failed copy-stream ownership quarantined. Retry those + // barriers before deciding whether a sequence slot is actually available. + for (int i = 0; i < (int)slots_.size(); ++i) { + if (slots_[(size_t)i].retiring()) retire(i); + } + int slot = -1; for (int i = 0; i < (int)slots_.size(); i++) { - if (!slots_[(size_t)i].active()) { slot = i; break; } + if (slots_[(size_t)i].phase == Qwen35SlotPhase::free) { slot = i; break; } } if (slot < 0) { r.status = AdmitStatus::busy; @@ -196,6 +203,9 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( Qwen35Slot & s = slots_[(size_t)slot]; s.phase = Qwen35SlotPhase::prefill; + s.request_id = request_id; + s.ddtree_suspended = false; + s.ddtree_sampled_steps = 0; s.handle = handle; s.cur_pos = 0; s.prompt_len = prompt_len; @@ -214,6 +224,33 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } +bool Qwen35SlotManager::ddtree_speculation_allowed(int slot) const { + return is_active(slot) && !slots_[(size_t)slot].ddtree_suspended; +} + +bool Qwen35SlotManager::ddtree_cohort_should_suspend( + uint64_t total_emitted, int active) { + const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); + if (adaptive && std::atoi(adaptive) == 0) { + return false; + } + if (active <= 0) return false; + return total_emitted < + (uint64_t)active * (uint64_t)kDdtreeMinEmittedTokens; +} + +bool Qwen35SlotManager::record_ddtree_sample( + int slot, bool suspend_cohort) { + if (!is_active(slot)) return false; + Qwen35Slot & s = slots_[(size_t)slot]; + // A suspended request must never pay for another probe. + if (s.ddtree_suspended) return false; + ++s.ddtree_sampled_steps; + if (!suspend_cohort) return false; + s.ddtree_suspended = true; + return true; +} + Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( int slot, int n_tokens) { PrefillChunk out; @@ -454,7 +491,7 @@ void Qwen35SlotManager::take_residency_telemetry( void Qwen35SlotManager::retire(int slot) { if (slot < 0 || slot >= (int)slots_.size()) return; Qwen35Slot & s = slots_[(size_t)slot]; - if (!s.active()) return; + if (s.phase == Qwen35SlotPhase::free) return; if (residency_) { const PagedKvResidencyStatus resident_status = residency_->forget_sequence(s.handle); @@ -467,6 +504,7 @@ void Qwen35SlotManager::retire(int slot) { // A failed copy-stream barrier leaves physical pages in flight. // Keep the slot and its pool handle intact so a later retirement // can retry forget_sequence without recycling those pages. + s.phase = Qwen35SlotPhase::retiring; return; } } diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index dd077496c..9a84fb3be 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -32,10 +32,12 @@ enum class Qwen35SlotPhase { free, prefill, decode, + retiring, }; struct Qwen35Slot { Qwen35SlotPhase phase = Qwen35SlotPhase::free; + uint64_t request_id = 0; PagedKvSequenceHandle handle; // Prompt tokens are the immutable prefix of sample_history. Decode tokens // append to the same allocation, avoiding a second full prompt copy. @@ -64,13 +66,28 @@ struct Qwen35Slot { : 0; } - bool active() const { return phase != Qwen35SlotPhase::free; } + // Real packed-tree samples are counted per request. A low aggregate-yield + // cohort sample suspends every participating request; ordinary AR keeps + // every target cache and feature-ring row current. + bool ddtree_suspended = false; + uint64_t ddtree_sampled_steps = 0; + + bool active() const { + return phase == Qwen35SlotPhase::prefill || + phase == Qwen35SlotPhase::decode; + } bool prefilling() const { return phase == Qwen35SlotPhase::prefill; } bool decoding() const { return phase == Qwen35SlotPhase::decode; } + bool retiring() const { return phase == Qwen35SlotPhase::retiring; } }; class Qwen35SlotManager { public: + // Packed DDTree pays for verify + accepted-path replay. Requiring six + // emitted tokens makes continuation earn at least three tokens per target + // forward before accounting for its additional draft/tree work. + static constexpr int kDdtreeMinEmittedTokens = 6; + // `max_ctx` is the per-sequence logical bound; slot count comes from the // pool's max_sequences. The pool must outlive the manager. Qwen35SlotManager(PagedKvPool & pool, int max_ctx, @@ -139,6 +156,13 @@ class Qwen35SlotManager { std::string * error = nullptr); void take_residency_telemetry(int slot, SeqEngine::DecodeOutput & out); + bool ddtree_speculation_allowed(int slot) const; + // Compare aggregate emitted yield (accepted children plus one replay + // bonus per request) against the cohort continuation floor. + static bool ddtree_cohort_should_suspend(uint64_t total_emitted, int active); + // Returns true exactly once, when this sample newly suspends the request. + bool record_ddtree_sample(int slot, bool suspend_cohort); + // One-token compatibility wrapper used by ordinary autoregressive decode. StepAppend append_token(int slot, int32_t fed_token); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index c640d778f..cb3cf7d0a 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -548,6 +548,13 @@ bool Qwen35Backend::init() { max_concurrent_prefills, mixed_prefill_tokens, long_mixed_prefill_tokens, long_prefill_threshold, idle_prefill_tokens, prefill_quantum); + if (concurrent_local_ddtree) { + const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); + std::fprintf(stderr, + "[parallel-ddtree] enabled budget=%d width=%d mode=packed-verify-replay adaptive=%s\n", + cfg_.ddtree_budget, tree_width, + adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); + } std::printf("[parallel] %d decode slots, up to %d packed prefills " "(mixed short/long %d/%d at >=%d tokens, " "idle %d, quantum %d), " diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index c34c71fb6..045786cb2 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -41,6 +41,7 @@ struct SchedSlot { uint64_t engine_request_id = 0; uint64_t ddtree_steps = 0; uint64_t ddtree_accepted_tokens = 0; + uint64_t ddtree_suspensions = 0; uint64_t target_forwards = 0; uint64_t kvflash_page_ins = 0; uint64_t kvflash_page_outs = 0; @@ -326,6 +327,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { {"pflash_output_tokens", prompt_tokens}, {"ddtree_steps", s.ddtree_steps}, {"ddtree_accepted_tokens", s.ddtree_accepted_tokens}, + {"ddtree_suspensions", s.ddtree_suspensions}, {"target_forwards", s.target_forwards}, {"kvflash_page_ins", s.kvflash_page_ins}, {"kvflash_page_outs", s.kvflash_page_outs}, @@ -752,6 +754,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } s.ddtree_steps += out.ddtree_steps; s.ddtree_accepted_tokens += out.ddtree_accepted_tokens; + s.ddtree_suspensions += out.ddtree_suspensions; s.target_forwards += out.target_forwards; s.kvflash_page_ins += out.kvflash_page_ins; s.kvflash_page_outs += out.kvflash_page_outs; diff --git a/server/test/test_ddtree_path.cpp b/server/test/test_ddtree_path.cpp index e7d2bbf3a..a4118ce1e 100644 --- a/server/test/test_ddtree_path.cpp +++ b/server/test/test_ddtree_path.cpp @@ -1,6 +1,6 @@ #include "common/ddtree.h" +#include "host_check.h" -#include #include #include #include @@ -9,6 +9,8 @@ using dflash::common::DDTree; using dflash::common::follow_verified_tree; using dflash::common::truncate_verified_path; +static int g_checks = 0; + int main() { DDTree tree; tree.n_nodes = 2; @@ -23,23 +25,23 @@ int main() { int pending = -1; std::vector accepted = follow_verified_tree(tree, posterior, pending); - assert((accepted == std::vector{0, 1, 2})); - assert(pending == 33); + CHECK((accepted == std::vector{0, 1, 2})); + CHECK(pending == 33); // Truncating after node 1 means node 2's token becomes pending. Keeping // the old value (33) would skip token 22 and describe uncommitted state. - assert(truncate_verified_path(accepted, 2, posterior, pending)); - assert((accepted == std::vector{0, 1})); - assert(pending == 22); + CHECK(truncate_verified_path(accepted, 2, posterior, pending)); + CHECK((accepted == std::vector{0, 1})); + CHECK(pending == 22); // An unchanged path preserves the already-computed pending token. - assert(!truncate_verified_path(accepted, 2, posterior, pending)); - assert(pending == 22); + CHECK(!truncate_verified_path(accepted, 2, posterior, pending)); + CHECK(pending == 22); // No headroom is represented explicitly and never dereferences a tip. - assert(truncate_verified_path(accepted, 0, posterior, pending)); - assert(accepted.empty()); - assert(pending == -1); + CHECK(truncate_verified_path(accepted, 0, posterior, pending)); + CHECK(accepted.empty()); + CHECK(pending == -1); std::puts("ddtree path tests passed"); return 0; diff --git a/server/test/test_paged_kv_residency.cpp b/server/test/test_paged_kv_residency.cpp index d39a5a0a8..274c7f1e2 100644 --- a/server/test/test_paged_kv_residency.cpp +++ b/server/test/test_paged_kv_residency.cpp @@ -60,7 +60,7 @@ struct MockTransfers { pending.push_back({[this, physical, host, bytes] { std::memcpy(host, &device[(size_t)physical * block_bytes], bytes); }}); - return true; + return !fail_copy_out_after_queue; }, [this](PagedKvSequenceHandle, uint32_t, uint32_t physical, const void * host, size_t bytes) { @@ -68,7 +68,7 @@ struct MockTransfers { pending.push_back({[this, physical, host, bytes] { std::memcpy(&device[(size_t)physical * block_bytes], host, bytes); }}); - return true; + return !fail_copy_in_after_queue; }, [this] { syncs++; @@ -99,6 +99,8 @@ struct MockTransfers { bool fail_alloc = false; bool fail_copy_out = false; bool fail_copy_in = false; + bool fail_copy_out_after_queue = false; + bool fail_copy_in_after_queue = false; bool fail_sync = false; }; @@ -253,6 +255,79 @@ TEST_CASE(PagedKvResidencyFixture, allocation_and_copy_failures_are_explicit) { CHECK(pager.is_resident(handle, 0)); } +TEST_CASE(PagedKvResidencyFixture, + failed_copy_out_prefix_is_barriered_before_retry) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + + io.fail_copy_out_after_queue = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(io.pending.empty()); + CHECK(pager.is_resident(handle, 0)); + CHECK(pager.stats().page_outs == 0); + + io.fail_copy_out_after_queue = false; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(!pager.is_resident(handle, 0)); +} + +TEST_CASE(PagedKvResidencyFixture, + rejected_append_remap_is_rolled_back_after_barrier) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 2)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + + io.fail_copy_in = true; + const auto append = pool.append(handle, 1); + CHECK(append.status == PagedKvStatus::Ok); + CHECK(append.remapped_cold_blocks.size() == 1); + CHECK(pager.observe_append(handle, append) == + PagedKvResidencyStatus::TransferFailed); + CHECK(snapshot(pool, handle).block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 1); + CHECK(pager.stats().page_ins == 0); +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_in_prefix_quarantines_mapping_until_barrier) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + + io.fail_copy_in_after_queue = true; + io.fail_sync = true; + CHECK(pager.ensure_resident(handle, {0}) == + PagedKvResidencyStatus::TransferFailed); + CHECK(snapshot(pool, handle).block_table[0] != PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 0); + CHECK(io.pending.size() == 1); + + io.fail_sync = false; + CHECK(pager.synchronize_before_read() == PagedKvResidencyStatus::Ok); + CHECK(snapshot(pool, handle).block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 1); + CHECK(pager.stats().page_ins == 0); +} + TEST_CASE(PagedKvResidencyFixture, forget_frees_host_backing_and_stale_is_rejected) { PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); MockTransfers io(2, 16); @@ -534,6 +609,32 @@ TEST_CASE(PagedKvResidencyFixture, CHECK(snapshot(pool, second).block_table[0] == PAGED_KV_COLD_BLOCK); } +TEST_CASE(PagedKvResidencyFixture, + resident_partial_head_stays_within_budget_during_growth) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/2, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager( + pool, config(16, /*budget=*/2, /*sink=*/0, /*tail=*/0), + io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 2)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const uint32_t append_head = snapshot(pool, first).block_table[0]; + CHECK(pager.append(second, 4)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + const auto grown = pager.append(first, 3); + CHECK(grown); + const auto table = snapshot(pool, first).block_table; + CHECK(table.size() == 2); + CHECK(table[0] == append_head); + CHECK(table[1] != PAGED_KV_COLD_BLOCK); + CHECK(snapshot(pool, second).block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(pager.stats().resident_blocks == 2); +} TEST_CASE(PagedKvResidencyFixture, requested_restore_set_is_protected_as_one_batch) { PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); diff --git a/server/test/test_seq_batch_plan.cpp b/server/test/test_seq_batch_plan.cpp index a771d160b..e1e13065e 100644 --- a/server/test/test_seq_batch_plan.cpp +++ b/server/test/test_seq_batch_plan.cpp @@ -131,9 +131,22 @@ int main() { burst.decode[0].committed_tokens = {8, 9, 10}; burst.decode[0].ddtree_steps = 1; burst.decode[0].ddtree_accepted_tokens = 3; + burst.decode[0].ddtree_suspensions = 1; burst.decode[0].target_forwards = 1; CHECK(validate_step_result(work, burst, 2).empty()); + SeqEngine::StepResult orphan_suspension = good; + orphan_suspension.decode[0].ddtree_suspensions = 1; + CHECK(!validate_step_result(work, orphan_suspension, 2).empty()); + + SeqEngine::StepResult failed_suspension = good; + failed_suspension.decode[0].failed = true; + failed_suspension.decode[0].token = -1; + failed_suspension.decode[0].error = "decode failed"; + failed_suspension.decode[0].ddtree_steps = 1; + failed_suspension.decode[0].ddtree_suspensions = 1; + CHECK(!validate_step_result(work, failed_suspension, 2).empty()); + // A scheduler stop in the committed prefix must hide the remaining burst // and final pending token. The backend state is discarded at retirement. std::vector delivered; diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index 46a5287d7..7988733b2 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -7,10 +7,13 @@ #include "qwen35/concurrency/qwen35_slot_manager.h" #include "host_check.h" +#include "scoped_env.h" #include #include #include +#include +#include using namespace dflash::common; @@ -489,6 +492,117 @@ int main() { CHECK(!mgr.has_prefill_prompt_at_least(768)); } + // Adaptive DDTree judges aggregate cohort yield rather than allowing one + // unlucky request to suppress higher-yield peers. Equality at the + // six-token average continues; a low average suspends all participants. + { + const luce_test::ScopedEnvVar adaptive("DFLASH_DDTREE_ADAPTIVE", nullptr); + CHECK(Qwen35SlotManager::ddtree_cohort_should_suspend(5, 1)); + CHECK(!Qwen35SlotManager::ddtree_cohort_should_suspend(6, 1)); + // Mixed synthetic cohort: one child-less row plus an eleven-token row + // exactly meets the continuation floor; one fewer token does not. + CHECK(!Qwen35SlotManager::ddtree_cohort_should_suspend(12, 2)); + CHECK(Qwen35SlotManager::ddtree_cohort_should_suspend(11, 2)); + + PagedKvPool pool(8, 2, /*block_size=*/16); + Qwen35SlotManager mgr(pool, 64); + auto a = admit(mgr, 101, prompt_tokens(4), greedy_sampler()); + auto b = admit(mgr, 102, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(is_admitted(b)); + CHECK(mgr.append_prefill(a.slot, 4).ok); + CHECK(mgr.append_prefill(b.slot, 4).ok); + mgr.commit_prefill(a.slot); + mgr.commit_prefill(b.slot); + CHECK(mgr.slot(a.slot).request_id == 101); + CHECK(mgr.slot(b.slot).request_id == 102); + CHECK(mgr.ddtree_speculation_allowed(a.slot)); + CHECK(mgr.ddtree_speculation_allowed(b.slot)); + + // The keep decision records a real sample without latching either + // participant, including the low-yield member. + CHECK(!mgr.record_ddtree_sample(a.slot, false)); + CHECK(!mgr.record_ddtree_sample(b.slot, false)); + CHECK(mgr.ddtree_speculation_allowed(a.slot)); + CHECK(mgr.ddtree_speculation_allowed(b.slot)); + CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 1); + CHECK(mgr.slot(b.slot).ddtree_sampled_steps == 1); + + // A low cohort decision is applied identically and atomically to both. + CHECK(mgr.record_ddtree_sample(a.slot, true)); + CHECK(mgr.record_ddtree_sample(b.slot, true)); + CHECK(!mgr.ddtree_speculation_allowed(a.slot)); + CHECK(!mgr.ddtree_speculation_allowed(b.slot)); + CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 2); + CHECK(mgr.slot(b.slot).ddtree_sampled_steps == 2); + + // Suspended requests cannot pay for another bad probe. + CHECK(!mgr.record_ddtree_sample(a.slot, true)); + CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 2); + + mgr.retire(a.slot); + auto reused = admit(mgr, 103, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(reused) && reused.slot == a.slot); + CHECK(mgr.slot(reused.slot).request_id == 103); + CHECK(mgr.ddtree_speculation_allowed(reused.slot)); + CHECK(mgr.slot(reused.slot).ddtree_sampled_steps == 0); + } + + // A failed residency barrier quarantines retirement ownership, and a + // later admission retries it before considering the slot reusable. + { + PagedKvPool pool(/*physical_block_count=*/1, + /*max_sequences=*/1, /*block_size=*/4); + bool fail_sync = false; + PagedKvResidencyTransferOps transfers{ + [](size_t bytes) -> void * { return new uint8_t[bytes]; }, + [](void * ptr) { delete[] static_cast(ptr); }, + [](PagedKvSequenceHandle, uint32_t, uint32_t, + void *, size_t) { return true; }, + [](PagedKvSequenceHandle, uint32_t, uint32_t, + const void *, size_t) { return true; }, + [&fail_sync] { return !fail_sync; }, + }; + PagedKvResidencyConfig config; + config.block_bytes = 16; + config.resident_budget_blocks = 1; + config.sink_blocks = 0; + config.tail_blocks = 0; + PagedKvResidencyManager residency(pool, config, std::move(transfers)); + Qwen35SlotManager mgr(pool, /*max_ctx=*/16, + /*speculative_headroom=*/1, &residency); + + auto first = admit(mgr, 301, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(first)); + CHECK(mgr.append_prefill(first.slot, 4).ok); + CHECK(residency.commit_pending_writes(mgr.slot(first.slot).handle) == + PagedKvResidencyStatus::Ok); + mgr.commit_prefill(first.slot); + + fail_sync = true; + CHECK(residency.evict_block( + mgr.slot(first.slot).handle, 0, + /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + mgr.retire(first.slot); + CHECK(!mgr.is_active(first.slot)); + CHECK(mgr.slot(first.slot).retiring()); + CHECK(pool.active_sequence_count() == 1); + + auto blocked = admit(mgr, 302, prompt_tokens(4), greedy_sampler()); + CHECK(!is_admitted(blocked)); + CHECK(is_busy(blocked)); + CHECK(mgr.slot(first.slot).retiring()); + CHECK(pool.active_sequence_count() == 1); + + fail_sync = false; + auto recovered = admit(mgr, 303, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(recovered)); + CHECK(recovered.slot == first.slot); + CHECK(mgr.is_prefilling(recovered.slot)); + CHECK(pool.active_sequence_count() == 1); + } + std::printf("OK test_seq_slot_manager (%d checks)\n", g_checks); return 0; } From ea5b737545ae310d11b871a53de976b8f51520fc Mon Sep 17 00:00:00 2001 From: Graffioh Date: Fri, 14 Aug 2026 11:16:14 +0000 Subject: [PATCH 05/42] perf(qwen35): adapt packed prefill budgeting --- .../concurrency/generate_ragged_prompts.py | 1 + .../concurrency/run_qwen36_concurrency.sh | 38 ++++++++++++++----- .../concurrency/qwen35_slot_manager.cpp | 9 +++++ server/src/qwen35/graph_builders.cpp | 6 +++ server/src/qwen35/qwen35_backend.cpp | 2 +- server/test/test_seq_slot_manager.cpp | 17 +++++++++ 6 files changed, 62 insertions(+), 11 deletions(-) diff --git a/harness/benchmarks/concurrency/generate_ragged_prompts.py b/harness/benchmarks/concurrency/generate_ragged_prompts.py index ca0b180a7..c04d4c434 100755 --- a/harness/benchmarks/concurrency/generate_ragged_prompts.py +++ b/harness/benchmarks/concurrency/generate_ragged_prompts.py @@ -10,6 +10,7 @@ PROFILES = { + "tiny": (64, 96, 128, 160), "short": (250, 350, 450, 550), "medium": (650, 850, 1150, 1350), "long": (2000, 2600, 3400, 4000), diff --git a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh index 62c1e513e..5583bacf4 100755 --- a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh +++ b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh @@ -40,7 +40,9 @@ if [[ $# -ne 0 ]]; then usage >&2; exit 2; fi for cmd in python3 curl sha256sum; do command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; }; done [[ -r "$MODEL" ]] || { echo "set MODEL to a readable GGUF" >&2; exit 2; } [[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } -[[ -x "$LLAMA_SERVER_BIN" ]] || { echo "missing llama.cpp server: $LLAMA_SERVER_BIN" >&2; exit 2; } +if [[ ",$VARIANTS," == *,llama,* ]]; then + [[ -x "$LLAMA_SERVER_BIN" ]] || { echo "missing llama.cpp server: $LLAMA_SERVER_BIN" >&2; exit 2; } +fi [[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } [[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ @@ -74,7 +76,9 @@ for c in "${client_list[@]}"; do [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } done for v in "${variant_list[@]}"; do - [[ "$v" == luce-k8 || "$v" == luce-k1 || "$v" == llama ]] || { echo "unknown variant $v" >&2; exit 2; } + [[ "$v" == luce-k8 || "$v" == luce-k1 || "$v" == luce-k16-b2 || + "$v" == luce-k16-b4 || "$v" == luce-k16-dyn || + "$v" == llama ]] || { echo "unknown variant $v" >&2; exit 2; } done mkdir -p "$OUT/prompts" @@ -95,7 +99,8 @@ stop_server() { fi server_pid="" } -trap stop_server EXIT INT TERM +trap stop_server EXIT +trap 'exit 130' INT TERM served_model_matches() { python3 - "$PORT" "$1" <<'PY' @@ -147,9 +152,9 @@ PY } write_metadata() { - local path="$1" variant="$2" workload="$3" clients="$4" repeat="$5" binary="$6" max_prefills="$7" command_file="$8" + local path="$1" variant="$2" workload="$3" clients="$4" repeat="$5" binary="$6" max_prefills="$7" mixed_budget="$8" idle_budget="$9" quantum="${10}" command_file="${11}" python3 -c 'import hashlib,json,pathlib,subprocess,sys -p,variant,workload,clients,repeat,binary,max_prefills,cmd_file,model_sha,prompts,repo=sys.argv[1:] +p,variant,workload,clients,repeat,binary,max_prefills,mixed,idle,quantum,cmd_file,model_sha,prompts,repo=sys.argv[1:] digest=lambda x: hashlib.sha256(pathlib.Path(x).read_bytes()).hexdigest() libs={} for line in subprocess.run(["ldd",binary],text=True,capture_output=True).stdout.splitlines(): @@ -165,18 +170,20 @@ if variant == "llama": raise RuntimeError(f"cannot identify llama.cpp source version from {binary} --version") obj={"variant":variant,"workload":workload,"clients":int(clients),"repeat":int(repeat), "max_concurrent_prefills":int(max_prefills),"server_binary":str(pathlib.Path(binary).resolve()), + "mixed_prefill_tokens":int(mixed),"idle_prefill_tokens":int(idle), + "prefill_allocation_quantum":int(quantum), "server_binary_sha256":digest(binary),"model_sha256":model_sha, "prompt_file_sha256":digest(prompts),"server_command":pathlib.Path(cmd_file).read_text().strip(), "resolved_shared_library_sha256":libs, "lucebox_git_head":lucebox_git_head if variant != "llama" else None, "server_version":server_version} pathlib.Path(p).write_text(json.dumps(obj,indent=2,sort_keys=True)+"\n")' \ - "$path" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$command_file" "$MODEL_SHA256" "$OUT/prompts/$workload.jsonl" "$REPO" + "$path" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$mixed_budget" "$idle_budget" "$quantum" "$command_file" "$MODEL_SHA256" "$OUT/prompts/$workload.jsonl" "$REPO" } run_case() { local repeat="$1" workload="$2" clients="$3" variant="$4" - local max_ctx timeout capacity max_prefills binary model_id + local max_ctx timeout capacity max_prefills mixed_budget idle_budget quantum binary model_id if [[ "$workload" == long ]]; then max_ctx=8192; timeout=1800 else @@ -188,20 +195,30 @@ run_case() { local -a command launch_env if [[ "$variant" == llama ]]; then binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 + mixed_budget=0; idle_budget=0; quantum=0 command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") launch_env=() else binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce - [[ "$variant" == luce-k8 ]] && max_prefills=8 || max_prefills=1 + case "$variant" in + luce-k8) max_prefills=8; mixed_budget=2048; idle_budget=4096; quantum=512 ;; + luce-k16-b2) max_prefills=16; mixed_budget=2048; idle_budget=2048; quantum=128 ;; + luce-k16-b4) max_prefills=16; mixed_budget=4096; idle_budget=4096; quantum=256 ;; + luce-k16-dyn) max_prefills=16; mixed_budget=2048; idle_budget=4096; quantum=256 ;; + *) max_prefills=1; mixed_budget=2048; idle_budget=4096; quantum=512 ;; + esac command=("$binary" "$MODEL" --target-device hip:0 --paged-attention --max-concurrency "$SLOTS" --kv-pool-tokens "$capacity" --max-ctx "$max_ctx" --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 --host 127.0.0.1 --port "$PORT" --model-name "$model_id") launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" - "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills") + "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills" + "DFLASH_MIXED_PREFILL_TOKENS=$mixed_budget" + "DFLASH_IDLE_PREFILL_TOKENS=$idle_budget" + "DFLASH_PREFILL_ALLOCATION_QUANTUM=$quantum") fi if ((${#launch_env[@]})); then printf 'env ' > "$case_dir/server-command.txt" @@ -210,7 +227,7 @@ run_case() { printf '%q ' "${command[@]}" > "$case_dir/server-command.txt" fi printf '\n' >> "$case_dir/server-command.txt" - write_metadata "$case_dir/server-metadata.json" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$case_dir/server-command.txt" + write_metadata "$case_dir/server-metadata.json" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$mixed_budget" "$idle_budget" "$quantum" "$case_dir/server-command.txt" echo "[run] $workload C=$clients repeat=$repeat variant=$variant" port_is_available || return 1 @@ -228,6 +245,7 @@ run_case() { --require-distinct-prompts --max-tokens "$WARMUP_TOKENS" --temperature 0 \ --ignore-eos --timeout "$timeout" --cooldown 0 --out "$case_dir/warmup.json" \ --label "$variant $workload C=$clients warmup" > "$case_dir/warmup.txt" + sleep 1 python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" \ --require-distinct-prompts --max-tokens "$MAX_TOKENS" --temperature 0 \ diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index e0e64170d..e6ef4ac39 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -88,6 +88,15 @@ bool Qwen35SlotManager::is_active(int slot) const { bool Qwen35SlotManager::is_prefilling(int slot) const { return is_active(slot) && slots_[(size_t)slot].prefilling(); } + +bool Qwen35SlotManager::has_prefill_prompt_at_least(int tokens) const { + if (tokens <= 0) return true; + return std::any_of(slots_.begin(), slots_.end(), + [tokens](const Qwen35Slot & slot) { + return slot.prefilling() && slot.prompt_len >= tokens; + }); +} + void Qwen35SlotManager::accumulate_residency_delta( Qwen35Slot & slot, const PagedKvResidencyStats & before) { if (!residency_) return; diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index f96349af3..1becd302c 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -417,6 +417,12 @@ bool build_target_step( n_prefill_segments, graph_capacity)) { return false; } + // Experimental adaptive prefill: node count depends on aggregate width, + // ragged layout, and fused decode rows. Use the already-supported maximum + // capacity so the benchmark does not rely on an incomplete shape proxy. + if (n_prefill_tokens > 0) { + graph_capacity = std::max(graph_capacity, 131072); + } // Persistent thread_local arena: rebuilt step graphs land at identical // addresses, keeping the ggml-cuda CUDA-graph cache key (nodes[0]) and diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index cb3cf7d0a..efc7e617b 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -390,7 +390,7 @@ bool Qwen35Backend::init() { const int max_concurrent_prefills = n_slots > 1 ? std::clamp( env_int_or_default("DFLASH_MAX_CONCURRENT_PREFILLS", 8), - 1, std::min(n_slots, 8)) + 1, n_slots) : 1; const int mixed_prefill_tokens = std::max( 1, env_int_or_default("DFLASH_MIXED_PREFILL_TOKENS", 2048)); diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index 7988733b2..7a2d7a26c 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -603,6 +603,23 @@ int main() { CHECK(pool.active_sequence_count() == 1); } + // Long-prefill policy follows active request length and clears on retire. + { + PagedKvPool pool(128, 2, /*block_size=*/16); + Qwen35SlotManager mgr(pool, 2048); + CHECK(!mgr.has_prefill_prompt_at_least(768)); + auto short_req = + admit(mgr, 201, prompt_tokens(512), greedy_sampler()); + CHECK(is_admitted(short_req)); + CHECK(!mgr.has_prefill_prompt_at_least(768)); + auto long_req = + admit(mgr, 202, prompt_tokens(800), greedy_sampler()); + CHECK(is_admitted(long_req)); + CHECK(mgr.has_prefill_prompt_at_least(768)); + mgr.retire(long_req.slot); + CHECK(!mgr.has_prefill_prompt_at_least(768)); + } + std::printf("OK test_seq_slot_manager (%d checks)\n", g_checks); return 0; } From e3cee3222d14ea04ba8ebf002c7853f7ddde1a33 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 06:20:09 +0000 Subject: [PATCH 06/42] 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 cc75a80f1763e49cd3c99ab9566f4808dfeb418f Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:24:56 +0200 Subject: [PATCH 07/42] qwen35: DSpark speculative decoding support Wire the DSpark drafter heads (low-rank Markov bigram correction + confidence head) into the qwen35 spec-decode loop, so Qwen3.8-27B DSpark drafters (e.g. RadixArk/Qwen3.8-27B-DSpark) run with full head support: - spec loop: markov-corrected greedy chain (fused single-graph variant with non-fused fallback) replaces plain argmax projection when the drafter ships DSpark heads; DDTree candidate top-k gets the markov bias too. Env-gated: DFLASH_QWEN35_DSPARK, DFLASH_QWEN35_FUSED_DSPARK, DFLASH_QWEN35_DSPARK_TREE (all default on). - target capture layers now follow the drafter GGUF's dflash.target_layer_ids instead of the evenly-spaced derivation; the Qwen3.8 drafter is trained on layers 4/16/28/40/52, not 1/16/31/46/61. - draft loader: dflash.mask_token_id from the drafter GGUF wins over the family default (Qwen3.8 drafter uses 248077, default was 248070), and optional YaRN rope scaling keys are parsed into DraftWeights. - draft graph: rope calls honor the drafter's YaRN config (previously hardcoded plain NEOX rope). - Qwen35DFlashTarget exposes lm_head for the fused head path. - convert_dflash_to_gguf.py: handle single-file DSpark releases (markov/ confidence heads inline in model.safetensors), transformers>=5 nested rope_parameters and dflash_config.mask_token_id, and emit YaRN scaling metadata. The confidence-gate adaptive block length is not wired yet (q_len sizes the per-request step buffers); the chain runs with the gate off. --- server/scripts/convert_dflash_to_gguf.py | 64 ++++++++++--- server/src/draft/draft_gguf_loader.cpp | 26 +++++ server/src/draft/draft_graph.cpp | 39 ++++---- server/src/qwen35/qwen35_backend.cpp | 115 +++++++++++++++++++++-- server/src/qwen35/qwen35_dflash_target.h | 1 + 5 files changed, 205 insertions(+), 40 deletions(-) diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index b904d7ea5..c4482f5f8 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -110,6 +110,20 @@ def pick(*keys): or c.get("aux_hidden_state_layer_ids")) if _tli: a["capture_layer_ids"] = [int(x) for x in _tli] + # Newer HF configs (transformers >= 5.x, e.g. the Qwen3.8 DSpark + # drafter) nest rope_theta / YaRN under rope_parameters and + # mask_token_id under dflash_config instead of top-level. + rp = c.get("rope_parameters") or c.get("rope_scaling") or {} + if isinstance(rp, dict): + if rp.get("rope_theta") is not None: + a["rope_theta"] = float(rp["rope_theta"]) + if str(rp.get("rope_type", "")).lower() == "yarn": + a["yarn_factor"] = float(rp.get("factor", 0.0)) + a["yarn_orig_ctx"] = int(rp.get("original_max_position_embeddings", 0)) + a["yarn_beta_fast"] = float(rp.get("beta_fast", 32.0)) + a["yarn_beta_slow"] = float(rp.get("beta_slow", 1.0)) + if dfc.get("mask_token_id") is not None: + a["mask_token_id"] = int(dfc["mask_token_id"]) print(f"[info] read arch from {cfg_path}") else: print(f"[warn] no config.json next to safetensors; using 27B defaults") @@ -248,18 +262,30 @@ def bytes_to_np(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray: } +# Alias sets per head tensor: SpecForge sidecar names, DS4 MTP-shard names, +# and single-file releases (e.g. RadixArk Qwen3.8-27B-DSpark) that carry the +# heads inline in the main model.safetensors. +DSPARK_MARKOV_W1_KEYS = ("dspark_markov_head.markov_w1.weight", + "mtp.2.markov_head.markov_w1.weight", + "markov_head.markov_w1.weight") +DSPARK_MARKOV_W2_KEYS = ("dspark_markov_head.markov_w2.weight", + "mtp.2.markov_head.markov_w2.weight", + "markov_head.markov_w2.weight") +DSPARK_CONF_W_KEYS = ("dspark_confidence_head.weight", + "mtp.2.confidence_head.proj.weight", + "confidence_head.proj.weight") +DSPARK_CONF_B_KEYS = ("dspark_confidence_head.bias", + "mtp.2.confidence_head.proj.bias", + "confidence_head.proj.bias") + DSPARK_TENSOR_MAP = { - ("dspark_markov_head.markov_w1.weight", - "mtp.2.markov_head.markov_w1.weight"): ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), - ("dspark_markov_head.markov_w2.weight", - "mtp.2.markov_head.markov_w2.weight"): ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), + DSPARK_MARKOV_W1_KEYS: ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), + DSPARK_MARKOV_W2_KEYS: ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), } DSPARK_CONFIDENCE_TENSOR_MAP = { - ("dspark_confidence_head.weight", - "mtp.2.confidence_head.proj.weight"): ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), - ("dspark_confidence_head.bias", - "mtp.2.confidence_head.proj.bias"): ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), + DSPARK_CONF_W_KEYS: ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), + DSPARK_CONF_B_KEYS: ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), } @@ -372,8 +398,8 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): return print(f"[info] reading DSpark aux heads from {aux_path}") - w1 = resolved[("dspark_markov_head.markov_w1.weight", "mtp.2.markov_head.markov_w1.weight")][1] - w2 = resolved[("dspark_markov_head.markov_w2.weight", "mtp.2.markov_head.markov_w2.weight")][1] + w1 = resolved[DSPARK_MARKOV_W1_KEYS][1] + w2 = resolved[DSPARK_MARKOV_W2_KEYS][1] vocab = int(w1.shape[0]) rank = int(w1.shape[1]) if tuple(w2.shape) != (vocab, rank): @@ -397,8 +423,8 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): conf_missing.append(names) continue conf_resolved[names] = (found_name, tensor, spec) - weight_names = ("dspark_confidence_head.weight", "mtp.2.confidence_head.proj.weight") - bias_names = ("dspark_confidence_head.bias", "mtp.2.confidence_head.proj.bias") + weight_names = DSPARK_CONF_W_KEYS + bias_names = DSPARK_CONF_B_KEYS if weight_names not in conf_resolved: if conf_missing: print("[warn] incomplete DSpark confidence head; Markov head will still load") @@ -470,6 +496,12 @@ def main(): writer.add_uint32(f"{ARCH}.vocab_size", a["vocab"]) writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", a["rms_eps"]) writer.add_float32(f"{ARCH}.rope.freq_base", a["rope_theta"]) + if a.get("yarn_factor", 0.0) > 1.0: + writer.add_string(f"{ARCH}.rope.scaling.type", "yarn") + writer.add_float32(f"{ARCH}.rope.scaling.factor", a["yarn_factor"]) + writer.add_uint32(f"{ARCH}.rope.scaling.original_context_length", a["yarn_orig_ctx"]) + writer.add_float32(f"{ARCH}.rope.scaling.beta_fast", a["yarn_beta_fast"]) + writer.add_float32(f"{ARCH}.rope.scaling.beta_slow", a["yarn_beta_slow"]) # DFlash-specific hyperparameters writer.add_uint32(f"{ARCH}.dflash.n_target_layers", a["n_target_layers"]) @@ -534,7 +566,13 @@ def sort_key(t): if not args.no_aux_heads: aux_path = args.aux_heads if args.aux_heads is not None else args.safetensors.parent / "dflash_aux_heads.pt" add_domino_aux_heads(writer, ARCH, aux_path) - add_dspark_aux_heads(writer, ARCH, aux_path) + # DSpark heads may live in a sidecar (.pt / .safetensors) or inline in + # the main safetensors (single-file releases like RadixArk + # Qwen3.8-27B-DSpark). Fall back to the main file when no sidecar exists. + dspark_aux = aux_path + if dspark_aux is not None and not dspark_aux.exists(): + dspark_aux = args.safetensors + add_dspark_aux_heads(writer, ARCH, dspark_aux) print(f"[info] writing {args.out_gguf}") writer.write_header_to_file() diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index e5a04721c..c882adfd9 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -209,6 +209,14 @@ bool load_draft_gguf(const std::string & path, if (target) { out.mask_token_id = target->mask_token_id; } + // The drafter's own MASK id wins over the family default: newer drafters + // (e.g. the Qwen3.8 DSpark release) are trained with a different mask + // token than the target-side default, and drafting with the wrong mask + // embedding silently destroys acceptance. + { + const uint32_t mask_meta = read_u32("dflash.mask_token_id", 0); + if (mask_meta != 0) out.mask_token_id = (int32_t)mask_meta; + } // Upper bounds on hparams. Guards against malformed/hostile GGUFs that // would otherwise trigger huge allocations or signed-int overflow when @@ -245,6 +253,24 @@ bool load_draft_gguf(const std::string & path, if (out.rope_theta == 0.0f) { fprintf(stderr, "[draft-gguf] WARNING: rope.freq_base not found in GGUF, draft RoPE will be wrong\n"); } + // YaRN rope scaling (optional). Drafters trained with YaRN (e.g. Qwen3.8 + // DSpark: factor 32, orig ctx 8192) apply it at every position; plain + // RoPE at inference silently degrades acceptance. + { + const float yarn_factor = read_f32("rope.scaling.factor", 0.0f); + if (yarn_factor > 1.0f) { + out.rope_freq_scale = 1.0f / yarn_factor; + out.rope_ext_factor = 1.0f; + out.rope_attn_factor = read_f32("rope.scaling.attn_factor", 1.0f); + out.rope_beta_fast = read_f32("rope.scaling.beta_fast", 32.0f); + out.rope_beta_slow = read_f32("rope.scaling.beta_slow", 1.0f); + out.rope_n_ctx_orig = (int)read_u32("rope.scaling.original_context_length", 0); + fprintf(stderr, + "[draft-gguf] YaRN rope: factor=%.1f orig_ctx=%d beta=%.1f/%.1f\n", + yarn_factor, out.rope_n_ctx_orig, + out.rope_beta_fast, out.rope_beta_slow); + } + } out.layers.assign((size_t)n_layer, DraftLayer{}); auto g = [&](const char * name) -> ggml_tensor * { diff --git a/server/src/draft/draft_graph.cpp b/server/src/draft/draft_graph.cpp index 472c214c9..5886177dc 100644 --- a/server/src/draft/draft_graph.cpp +++ b/server/src/draft/draft_graph.cpp @@ -40,6 +40,19 @@ namespace dflash::common { +// RoPE with the drafter's scaling config. YaRN-trained drafters (e.g. the +// Qwen3.8 DSpark release: factor 32, orig ctx 8192) apply the scaled rotary +// at every position, so plain-RoPE inference silently degrades acceptance. +static ggml_tensor * draft_rope(ggml_context * ctx, ggml_tensor * t, + ggml_tensor * positions, + const DraftWeights & w) { + return ggml_rope_ext(ctx, t, positions, /*freq_factors=*/nullptr, + w.head_dim, GGML_ROPE_TYPE_NEOX, w.rope_n_ctx_orig, + w.rope_theta, w.rope_freq_scale, + w.rope_ext_factor, w.rope_attn_factor, + w.rope_beta_fast, w.rope_beta_slow); +} + // Feature fusion shared by the legacy one-shot graph and the cached-KV // builders: optional per-capture RMSNorm slices, fc projection, hidden_norm. // Row-independent, so it is bit-identical whether run over the full window @@ -83,7 +96,6 @@ DraftGraphOutputs build_draft_graph( const int n_kv = w.n_head_kv; const int head_dim = w.head_dim; const float eps = DFLASH27B_RMS_EPS; - const float rope_base = w.rope_theta; // ── 1. Feature fusion: target_feat = rms_norm(fc @ target_hidden_cat, hidden_norm) // fc: [5*hidden, hidden] (ggml: ne[0]=5*hidden, ne[1]=hidden) @@ -185,14 +197,8 @@ DraftGraphOutputs build_draft_graph( pk = ggml_view_1d(ctx, in.positions_k, eff_total_k, ctx_offset * ggml_element_size(in.positions_k)); } - Q = ggml_rope_ext(ctx, Q, in.positions_q, /*freq_factors=*/nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, /*n_ctx_orig=*/0, - rope_base, /*freq_scale=*/1.0f, - /*ext_factor=*/0.0f, /*attn_factor=*/1.0f, - /*beta_fast=*/0.0f, /*beta_slow=*/0.0f); - K = ggml_rope_ext(ctx, K, pk, nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, 0, - rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Q = draft_rope(ctx, Q, in.positions_q, w); + K = draft_rope(ctx, K, pk, w); // ── 2e. Permute into the layout flash_attn_ext wants // q: [n_embd_k=head_dim, n_batch=q_len, n_head, ne3] @@ -309,11 +315,7 @@ static void draft_ctx_kv_rows( K = ggml_reshape_3d(ctx, K, w.head_dim, w.n_head_kv, n); K = ggml_rms_norm(ctx, K, eps); K = ggml_mul (ctx, K, L.k_norm); - K = ggml_rope_ext(ctx, K, positions, /*freq_factors=*/nullptr, - w.head_dim, GGML_ROPE_TYPE_NEOX, /*n_ctx_orig=*/0, - w.rope_theta, /*freq_scale=*/1.0f, - /*ext_factor=*/0.0f, /*attn_factor=*/1.0f, - /*beta_fast=*/0.0f, /*beta_slow=*/0.0f); + K = draft_rope(ctx, K, positions, w); // rope output is contiguous [head_dim, n_kv, n] → head-major rows view *k_rows_out = ggml_view_2d(ctx, K, (int64_t)w.head_dim * w.n_head_kv, n, K->nb[2], 0); @@ -356,7 +358,6 @@ DraftGraphOutputs build_draft_kv_step( const int n_kv = w.n_head_kv; const int head_dim = w.head_dim; const float eps = DFLASH27B_RMS_EPS; - const float rope_base = w.rope_theta; const int kv_total = cache.kv_total; static const bool disable_attn_gate = @@ -380,18 +381,14 @@ DraftGraphOutputs build_draft_kv_step( Q = ggml_reshape_3d(ctx, Q, head_dim, n_head, q_len); Q = ggml_rms_norm(ctx, Q, eps); Q = ggml_mul (ctx, Q, L.q_norm); - Q = ggml_rope_ext(ctx, Q, in.positions_q, nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, 0, - rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Q = draft_rope(ctx, Q, in.positions_q, w); // ── noise K/V into the scratch cache slots ggml_tensor * Kn = ggml_mul_mat(ctx, L.wk, hn); Kn = ggml_reshape_3d(ctx, Kn, head_dim, n_kv, q_len); Kn = ggml_rms_norm(ctx, Kn, eps); Kn = ggml_mul (ctx, Kn, L.k_norm); - Kn = ggml_rope_ext(ctx, Kn, in.positions_q, nullptr, - head_dim, GGML_ROPE_TYPE_NEOX, 0, - rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kn = draft_rope(ctx, Kn, in.positions_q, w); ggml_tensor * Kn_rows = ggml_view_2d(ctx, Kn, (int64_t)head_dim * n_kv, q_len, Kn->nb[2], 0); ggml_tensor * Vn_rows = ggml_mul_mat(ctx, L.wv, hn); // [kv_dim, q_len] diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index f15901fb0..6ca81085f 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -13,6 +13,7 @@ #include "common/geometric_sampler_cuda.h" #include #endif +#include "common/dspark_head.h" #include "common/io_utils.h" #include "common/restore_delta.h" #include "qwen35_tensor_parallel.h" @@ -25,6 +26,7 @@ #include "flashprefill.h" #include +#include #include #include #include @@ -149,6 +151,33 @@ static bool qwen35_empty_visible_output(const std::vector & tokens, return true; } +// Drafters trained on explicit target layers (GGUF dflash.target_layer_ids) +// override the evenly-spaced derivation: capturing different layers than the +// drafter was trained on silently destroys acceptance. +static void apply_drafter_capture_layer_ids(const DraftWeights & dw, TargetWeights & w) { + if (dw.capture_layer_ids.empty()) return; + const int n = (int)dw.capture_layer_ids.size(); + bool ok = (n == w.n_capture_layers); + for (int k = 0; ok && k < n; k++) + ok = dw.capture_layer_ids[k] >= 0 && dw.capture_layer_ids[k] < w.n_layer; + if (!ok) { + std::fprintf(stderr, + "[draft] drafter target_layer_ids invalid (n=%d, slots=%d); " + "keeping derived capture layers\n", n, w.n_capture_layers); + return; + } + bool changed = false; + for (int k = 0; k < n; k++) { + changed |= w.capture_layer_ids[k] != dw.capture_layer_ids[k]; + w.capture_layer_ids[k] = dw.capture_layer_ids[k]; + } + if (changed) { + std::printf("[draft] target capture layers from drafter GGUF:"); + for (int k = 0; k < n; k++) std::printf(" %d", w.capture_layer_ids[k]); + std::printf("\n"); + } +} + // ── Construction / destruction ────────────────────────────────────────── Qwen35Backend::Qwen35Backend(const Qwen35Config & cfg) : cfg_(cfg) {} @@ -253,6 +282,7 @@ bool Qwen35Backend::init() { return false; } std::printf("[draft] loaded\n"); + apply_drafter_capture_layer_ids(dw_, w_); if (cfg_.draft_swa_window > 0) { dw_.swa_window = cfg_.draft_swa_window; @@ -607,6 +637,7 @@ bool Qwen35Backend::unpark(ParkTarget target) { std::fprintf(stderr, "[unpark] draft: %s\n", dflash27b_last_error()); return false; } + apply_drafter_capture_layer_ids(dw_, w_); // Re-apply rope overrides after reload. if (dw_.rope_theta != w_.rope_theta && w_.rope_theta > 0.0f) dw_.rope_theta = w_.rope_theta; @@ -2139,6 +2170,14 @@ bool Qwen35Backend::sync_local_draft_features(int start_pos, int n_tokens) { // ── DFlash speculative decode loop ───────────────────────────────────── +static bool qwen35_dspark_enabled() { + static const bool kEnabled = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK"); + return e == nullptr || std::string(e) != "0"; + }(); + return kEnabled; +} + bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io, @@ -2504,12 +2543,58 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // DDTree consumes top-K rows directly. Avoid projecting the same // hidden block once for argmax and again for top-K on every step. if (!use_tree_verify) { - if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { - std::fprintf(stderr, "spec-decode: projection failed\n"); - step_graph_destroy(draft_sg); - return false; + // DSpark heads (markov bigram correction + optional confidence + // gate) when the drafter ships them; mirrors the laguna hook. + bool used_dspark = false; + if (qwen35_dspark_enabled() && dw_.dspark.enabled && + q_len > 1 && !sampled_verify && !use_remote_draft) { + static std::atomic s_dspark_logged{false}; + if (!s_dspark_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DSpark Markov head active for greedy chain decode " + "(rank=%d vocab=%d confidence_dim=%d)\n", + dw_.dspark.markov_rank, dw_.dspark.vocab_size, + dw_.dspark.confidence_dim); + } + static const bool fused_dspark = []() { + const char * e = std::getenv("DFLASH_QWEN35_FUSED_DSPARK"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool ds_ok = false; + if (fused_dspark) { + ds_ok = dspark_markov_correct_greedy_chain_fused( + dw_, draft_backend_, target->lm_head_tensor(), + local_hidden.data(), q_len, last_tok, draft_tok); + } + if (!ds_ok) { + // threshold 0 = confidence gate off: q_len sizes the + // step buffers for the whole request, so the truncated + // chain the gate produces cannot be verified here yet. + ds_ok = dspark_markov_correct_greedy_chain(dw_, draft_backend_, *target, + local_hidden.data(), q_len, + last_tok, + /*confidence_threshold=*/0.0f, + draft_tok); + } + if (ds_ok) { + used_dspark = true; + } else { + static std::atomic s_dspark_warned{false}; + if (!s_dspark_warned.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DSpark Markov head failed; falling back to " + "base DFlash projection\n"); + } + } + } + if (!used_dspark) { + if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { + std::fprintf(stderr, "spec-decode: projection failed\n"); + step_graph_destroy(draft_sg); + return false; + } + draft_tok[0] = last_tok; } - draft_tok[0] = last_tok; } if (use_tree_verify) { @@ -2518,7 +2603,25 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector top_lp; std::vector top_ids; const auto profile_project_start = profile_start(); - if (!target->project_hidden_to_topk(local_hidden.data(), q_len, K, + static const bool dspark_tree = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK_TREE"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool topk_ok = false; + if (dspark_tree && qwen35_dspark_enabled() && dw_.dspark.enabled) { + static std::atomic s_dstree_logged{false}; + if (!s_dstree_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DSpark Markov head active for DDTree candidates\n"); + } + topk_ok = dspark_markov_project_topk(dw_, draft_backend_, + target->lm_head_tensor(), + local_hidden.data(), q_len, K, + cfg_.ddtree_temp, last_tok, + top_lp, top_ids); + } + if (!topk_ok && + !target->project_hidden_to_topk(local_hidden.data(), q_len, K, cfg_.ddtree_temp, top_lp, top_ids)) { std::fprintf(stderr, "spec-decode: ddtree topk projection failed\n"); step_graph_destroy(draft_sg); diff --git a/server/src/qwen35/qwen35_dflash_target.h b/server/src/qwen35/qwen35_dflash_target.h index 3c8864b6b..cc8a37c3d 100644 --- a/server/src/qwen35/qwen35_dflash_target.h +++ b/server/src/qwen35/qwen35_dflash_target.h @@ -75,6 +75,7 @@ class Qwen35DFlashTarget : public DFlashTarget { int hidden_size() const override { return w_.n_embd; } int mask_token_id() const override; + ggml_tensor * lm_head_tensor() override { return w_.output; } const std::vector & capture_layer_ids() const override; // kvflash mode: verify writes are slot-mapped via the pager and the From 49a3d61a3653dc3cdda9ff78ce8b127295a5361a Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 16:55:07 +0000 Subject: [PATCH 08/42] test(feature-gate): repair stacked baseline assertions --- server/test/test_feature_gate.cpp | 40 ++++++++++++++++--------------- 1 file changed, 21 insertions(+), 19 deletions(-) diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index fb1a5571a..7b3838874 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -386,47 +386,47 @@ void test_feature_gate_paged_attention_requires_plain_ar_decode() { 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)); + CHECK(gate_result( + concurrent_ddtree, "qwen35", PlacementBackend::Cuda).empty()); + CHECK(gate_result( + concurrent_ddtree, "qwen35", PlacementBackend::Hip).empty()); BackendArgs tensor_ddtree = concurrent_ddtree; - TEST_ASSERT(parse_placement_device_list( + CHECK(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)); + CHECK(!gate_result( + tensor_ddtree, "qwen35", PlacementBackend::Cuda).empty()); 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)); + CHECK(gate_result(concurrent_ddtree, "qwen35", + PlacementBackend::Hip, concurrent_pflash).empty()); BackendArgs concurrent_plain = base; concurrent_plain.max_concurrency = 16; - TEST_ASSERT(gate_accepts(concurrent_plain, "qwen35", - PlacementBackend::Hip, concurrent_pflash)); + CHECK(gate_result(concurrent_plain, "qwen35", + PlacementBackend::Hip, concurrent_pflash).empty()); BackendFeatureConfig concurrent_kvflash; concurrent_kvflash.kvflash_enabled = true; - TEST_ASSERT(gate_accepts(concurrent_plain, "qwen35", - PlacementBackend::Hip, concurrent_kvflash)); + CHECK(gate_result(concurrent_plain, "qwen35", + PlacementBackend::Hip, concurrent_kvflash).empty()); 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)); + CHECK(!gate_result( + bad_budget, "qwen35", PlacementBackend::Hip).empty()); } 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)); + CHECK(!gate_result( + remote_ddtree, "qwen35", PlacementBackend::Hip).empty()); BackendArgs windowed = base; windowed.fa_window = 4096; @@ -540,9 +540,11 @@ void test_feature_gate_parallel_and_kv_pool_rules() { ((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)); + CHECK(gate_result( + tree_pool, "qwen35", PlacementBackend::Cuda).empty()); tree_pool.kv_pool_tokens = max_tree_pool_tokens + PAGED_BLOCK_SIZE; - TEST_ASSERT(!gate_accepts(tree_pool, "qwen35", PlacementBackend::Cuda)); + CHECK(!gate_result( + tree_pool, "qwen35", PlacementBackend::Cuda).empty()); // The automatic pool is memory-derived, so a logical slot/context product // larger than the physical tensor address space is legal. From 2e9be33d5227a74880be52601aa4cfb054290646 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 17:04:57 +0000 Subject: [PATCH 09/42] feat(server): add per-request decode policy --- server/src/common/backend_args.h | 2 + server/src/common/backend_factory.cpp | 2 + server/src/common/concurrency/seq_engine.h | 4 ++ server/src/common/feature_gate.cpp | 24 +++++++--- server/src/common/speculation_policy.h | 48 +++++++++++++++++++ .../qwen35/concurrency/qwen35_seq_engine.cpp | 7 ++- server/src/qwen35/qwen35_backend.h | 2 + server/src/server/http_server.cpp | 17 +++++++ server/src/server/http_server.h | 5 ++ server/src/server/scheduler.cpp | 3 ++ server/src/server/server_main.cpp | 19 +++++++- server/test/test_feature_gate.cpp | 28 +++++++++++ server/test/test_server_unit.cpp | 23 +++++++++ 13 files changed, 175 insertions(+), 9 deletions(-) create mode 100644 server/src/common/speculation_policy.h diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h index ee7b55d33..66b6feecd 100644 --- a/server/src/common/backend_args.h +++ b/server/src/common/backend_args.h @@ -10,6 +10,7 @@ #include "placement/remote_target_shard_config.h" #include "prefill_attention_mode.h" +#include "speculation_policy.h" namespace dflash::common { // Server-owned features that participate in backend admission even though @@ -81,6 +82,7 @@ struct BackendArgs { bool ddtree_chain_seed = true; int verify_width = 0; // chain spec verify width; 0 = adaptive bool use_feature_mirror = false; + SpeculationPolicy speculation_policy = SpeculationPolicy::Adaptive; }; } // namespace dflash::common diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 0d1ccc61b..d49d2584a 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -284,6 +284,7 @@ std::unique_ptr create_backend( cfg.ddtree_chain_seed = args.ddtree_chain_seed; cfg.use_feature_mirror = args.use_feature_mirror; + cfg.speculation_policy = args.speculation_policy; auto backend = std::make_unique(cfg); if (!backend->init()) { std::fprintf(stderr, "[backend_factory] Qwen35Backend init failed\n"); @@ -310,6 +311,7 @@ std::unique_ptr create_backend( cfg.ddtree_chain_seed = args.ddtree_chain_seed; cfg.use_feature_mirror = args.use_feature_mirror; + cfg.speculation_policy = args.speculation_policy; auto backend = std::make_unique(cfg); if (!backend->init()) { std::fprintf(stderr, "[backend_factory] Qwen35MoeBackend init failed\n"); diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 6371e2552..812bc54dd 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -48,6 +48,7 @@ // engine through it before wiring it up. #pragma once +#include "common/speculation_policy.h" #include #include @@ -185,6 +186,9 @@ class SeqEngine { // False when scheduler-side policy may replace the sampled token // before it is committed (currently the thinking-budget close hook). bool allow_speculation = true; + // Effective server default + per-request override. The hard safety + // check above always wins over this policy. + SpeculationPolicy speculation_policy = SpeculationPolicy::Adaptive; }; struct DecodeOutput { int slot = -1; diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index 63099832b..dfbd1e604 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -177,6 +177,10 @@ std::string check_feature_compatibility( const bool concurrent_local_ddtree = concurrent_local_paged_qwen && args.draft_path != nullptr && args.ddtree_mode; + const bool concurrent_local_chain = + concurrent_local_paged_qwen && args.draft_path != nullptr && + !args.ddtree_mode && + args.speculation_policy != SpeculationPolicy::Never; if (args.ddtree_mode && (args.ddtree_budget < 1 || args.ddtree_budget > 255)) { @@ -200,13 +204,19 @@ std::string check_feature_compatibility( args.remote_target_shard.enabled()) { return "--paged-attention requires one local target device"; } - if (args.remote_draft.enabled()) { + if (args.remote_draft.enabled() && + args.speculation_policy != SpeculationPolicy::Never) { 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.ddtree_mode && !concurrent_local_ddtree) { + return "paged DDTree requires concurrent local target/draft execution " + "on one device"; + } + if (args.draft_path != nullptr && !args.ddtree_mode && + args.speculation_policy != SpeculationPolicy::Never && + !concurrent_local_chain) { + return "paged chain speculation requires concurrent local target/draft " + "execution on one device"; } if (args.fa_window != 0) { return "--paged-attention requires full attention (--fa-window 0)"; @@ -261,7 +271,9 @@ std::string check_feature_compatibility( const int64_t tree_scratch = concurrent_local_ddtree ? (int64_t)args.max_concurrency * paged_token_capacity(args.ddtree_budget + 1) - : 0; + : concurrent_local_chain + ? (int64_t)args.max_concurrency * paged_token_capacity(16) + : 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 * diff --git a/server/src/common/speculation_policy.h b/server/src/common/speculation_policy.h new file mode 100644 index 000000000..8758d8ced --- /dev/null +++ b/server/src/common/speculation_policy.h @@ -0,0 +1,48 @@ +// Decode policy shared by the HTTP layer, scheduler, and sequence engines. + +#pragma once + +#include +#include + +namespace dflash::common { + +enum class SpeculationPolicy { + Adaptive, + Always, + Never, +}; + +inline const char * speculation_policy_name(SpeculationPolicy policy) { + switch (policy) { + case SpeculationPolicy::Adaptive: return "adaptive"; + case SpeculationPolicy::Always: return "speculation"; + case SpeculationPolicy::Never: return "ar"; + } + return "adaptive"; +} + +inline bool parse_speculation_policy( + std::string_view value, SpeculationPolicy & policy) { + if (value == "adaptive") { + policy = SpeculationPolicy::Adaptive; + return true; + } + if (value == "speculation") { + policy = SpeculationPolicy::Always; + return true; + } + if (value == "ar") { + policy = SpeculationPolicy::Never; + return true; + } + return false; +} + +inline SpeculationPolicy resolve_speculation_policy( + SpeculationPolicy server_default, + const std::optional & request_override) { + return request_override.value_or(server_default); +} + +} // namespace dflash::common diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 52949acfc..349abd3da 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -151,10 +151,13 @@ bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { return value ? std::max(0, std::atoi(value)) : 0; }(); for (const StepInput & in : plan.decode) { - if (!in.allow_speculation || in.slot < 0 || + if (!in.allow_speculation || + in.speculation_policy == SpeculationPolicy::Never || + in.slot < 0 || in.slot >= slots_.slot_count() || !slots_.slot(in.slot).decoding() || - !slots_.ddtree_speculation_allowed(in.slot) || + (in.speculation_policy != SpeculationPolicy::Always && + !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()) { diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index 62fee8524..e7170c5af 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -18,6 +18,7 @@ #include "placement/remote_draft_config.h" #include "step_graph.h" #include "ddtree.h" +#include "common/speculation_policy.h" #include "dflash_feature_ring.h" #include "common/dflash_draft_kv.h" #include "common/concurrency/paged_kv_pool.h" @@ -86,6 +87,7 @@ struct Qwen35Config { float ddtree_temp = 1.0f; bool ddtree_chain_seed = true; bool use_feature_mirror = false; + SpeculationPolicy speculation_policy = SpeculationPolicy::Adaptive; }; // ── Backend class ─────────────────────────────────────────────────────── diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 52a4a845c..a66370003 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", speculation_policy_name(config.decode_mode)}, {"server", server}, {"model", { {"arch", config.arch}, @@ -777,6 +778,7 @@ json build_props_body(const ServerConfig & config, {"kv_cache_v", config.kv_cache_v}, {"lazy_draft", config.lazy_draft}, {"draft_residency", draft_residency_policy_name(config.draft_residency)}, + {"decode_mode", speculation_policy_name(config.decode_mode)}, {"target_sharding", config.target_sharding}, // Prefill chunk size (bargs.chunk). Surfaced so snapshot // tooling captures the full config — bench consumers @@ -1661,6 +1663,21 @@ 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; + if (body.contains("decode_mode")) { + if (!body["decode_mode"].is_string()) { + send_error(fd, 400, + "decode_mode must be ar, speculation, or adaptive"); + return false; + } + SpeculationPolicy policy; + if (!parse_speculation_policy( + body["decode_mode"].get(), policy)) { + send_error(fd, 400, + "decode_mode must be ar, speculation, or adaptive"); + return false; + } + req.decode_mode = policy; + } // Accept the output-token names used by each supported API dialect. // Default when the client omits all three: --default-max-tokens, so diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 0fa0ef718..3644dea01 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -19,6 +19,7 @@ #include "tokenizer.h" #include "chat_template.h" #include "tool_memory.h" +#include "common/speculation_policy.h" #include "prefix_cache.h" #include "disk_prefix_cache.h" #include "freeze_history.h" @@ -33,6 +34,7 @@ #include #include +#include #include #include #include @@ -212,6 +214,8 @@ struct ServerConfig { bool lazy_draft = false; // legacy alias for request-scoped draft residency DraftResidencyPolicy draft_residency = DraftResidencyPolicy::Auto; + // Default speculative-decode policy; individual requests may override it. + SpeculationPolicy decode_mode = SpeculationPolicy::Adaptive; // Disk prefix cache std::string disk_cache_dir; // empty = disabled size_t disk_cache_budget_mb = 4096; // max disk usage in MB @@ -335,6 +339,7 @@ struct ParsedRequest { DiskPrefixCachePolicy disk_cache_policy; // PPP: stable pin cut for tool-heavy requests (0 = use default boundary). int pin_end_token = 0; + std::optional decode_mode; }; // Parse request sampler fields, applying model-card defaults where present. diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 045786cb2..c4adf8419 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -704,6 +704,9 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { input.token = slots[(size_t)i].pending_tok; input.allow_speculation = slots[(size_t)i].hook.close_token_ids.empty(); + input.speculation_policy = resolve_speculation_policy( + config_.decode_mode, + slots[(size_t)i].job->req.decode_mode); 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..67f8058fd 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -122,6 +122,8 @@ 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 Speculative decode policy: ar, speculation,\n" + " or adaptive (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 +414,15 @@ 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) { + SpeculationPolicy policy; + if (!parse_speculation_policy(argv[++i], policy)) { + std::fprintf(stderr, + "[server] --decode-mode expects ar, speculation, or adaptive\n"); + return 2; + } + bargs.speculation_policy = policy; + 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] != '-') { @@ -1115,6 +1126,8 @@ int main(int argc, char ** argv) { } std::fprintf(stderr, "[server] │ ddtree_budget = %d\n", bargs.ddtree_budget); std::fprintf(stderr, "[server] │ prefix_cache = %d slots\n", sconfig.prefix_cache_cap); + std::fprintf(stderr, "[server] │ decode_mode = %s\n", + speculation_policy_name(bargs.speculation_policy)); std::fprintf(stderr, "[server] │ prefill_cache = %d slots\n", sconfig.prefill_cache_cap); std::fprintf(stderr, "[server] │ cors = %s\n", sconfig.enable_cors ? "ON" : "off"); std::fprintf(stderr, "[server] │ cache_type_k = %s\n", @@ -1158,7 +1171,11 @@ int main(int argc, char ** argv) { 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.speculative_enabled = + bargs.speculation_policy != SpeculationPolicy::Never && + (bargs.ddtree_mode || + (bargs.paged_attention && bargs.max_concurrency > 1 && + bargs.draft_path != nullptr)); 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_feature_gate.cpp b/server/test/test_feature_gate.cpp index 7b3838874..64cbc7f3b 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -376,6 +376,20 @@ void test_feature_gate_paged_attention_requires_plain_ar_decode() { BackendArgs draft = base; draft.draft_path = "/nonexistent/draft.gguf"; CHECK(!gate_result(draft, "qwen35", PlacementBackend::Cuda).empty()); + // A local DSpark chain cluster is admitted under concurrent paged + // serving on either GPU backend. Forced AR may carry an unused drafter + // even without concurrent slots. + BackendArgs concurrent_chain = draft; + concurrent_chain.max_concurrency = 16; + CHECK(gate_result( + concurrent_chain, "qwen35", PlacementBackend::Cuda).empty()); + CHECK(gate_result( + concurrent_chain, "qwen35", PlacementBackend::Hip).empty()); + + BackendArgs forced_ar = draft; + forced_ar.speculation_policy = SpeculationPolicy::Never; + CHECK(gate_result( + forced_ar, "qwen35", PlacementBackend::Cuda).empty()); BackendArgs ddtree = base; ddtree.ddtree_mode = true; @@ -545,6 +559,20 @@ void test_feature_gate_parallel_and_kv_pool_rules() { tree_pool.kv_pool_tokens = max_tree_pool_tokens + PAGED_BLOCK_SIZE; CHECK(!gate_result( tree_pool, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs chain_pool = paged; + chain_pool.max_concurrency = 16; + chain_pool.draft_path = "/nonexistent/draft.gguf"; + const long long chain_scratch = + (long long)chain_pool.max_concurrency * paged_token_capacity(16); + const long long max_chain_pool_tokens = + ((long long)INT_MAX - PAGED_BLOCK_SIZE - chain_scratch) / + PAGED_BLOCK_SIZE * PAGED_BLOCK_SIZE; + chain_pool.kv_pool_tokens = max_chain_pool_tokens; + CHECK(gate_result( + chain_pool, "qwen35", PlacementBackend::Hip).empty()); + chain_pool.kv_pool_tokens = max_chain_pool_tokens + PAGED_BLOCK_SIZE; + CHECK(!gate_result( + chain_pool, "qwen35", PlacementBackend::Hip).empty()); // The automatic pool is memory-derived, so a logical slot/context product // larger than the physical tensor address space is legal. diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 2ff471a3d..fe790b6e9 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -2580,6 +2580,26 @@ TEST_CASE(ServerUnitFixture, test_parse_request_sampler_applies_defaults_and_ove TEST_ASSERT(std::fabs(sampler.pres_pen - 0.3f) < 0.001f); TEST_ASSERT(std::fabs(sampler.rep_pen - 1.1f) < 0.001f); } +TEST_CASE(ServerUnitFixture, test_speculation_policy_parse_name_and_fold) { + SpeculationPolicy policy = SpeculationPolicy::Adaptive; + TEST_ASSERT(parse_speculation_policy("ar", policy)); + TEST_ASSERT(policy == SpeculationPolicy::Never); + TEST_ASSERT(std::string(speculation_policy_name(policy)) == "ar"); + + TEST_ASSERT(parse_speculation_policy("speculation", policy)); + TEST_ASSERT(policy == SpeculationPolicy::Always); + TEST_ASSERT(std::string(speculation_policy_name(policy)) == "speculation"); + + TEST_ASSERT(parse_speculation_policy("adaptive", policy)); + TEST_ASSERT(policy == SpeculationPolicy::Adaptive); + TEST_ASSERT(!parse_speculation_policy("always", policy)); + + TEST_ASSERT(resolve_speculation_policy( + SpeculationPolicy::Always, std::nullopt) == SpeculationPolicy::Always); + TEST_ASSERT(resolve_speculation_policy( + SpeculationPolicy::Always, SpeculationPolicy::Never) == + SpeculationPolicy::Never); +} TEST_CASE(ServerUnitFixture, test_require_messages_array_rejects_invalid) { const json valid = {{"messages", json::array({ @@ -4928,6 +4948,7 @@ TEST_CASE(ServerUnitFixture, test_props_runtime_shape) { cfg.chunk = 512; cfg.target_device = "auto:0"; cfg.draft_device = "auto:0"; + cfg.decode_mode = SpeculationPolicy::Always; TEST_ASSERT(cfg.admission_coalesce_ms == 20); Tokenizer tok; @@ -4943,6 +4964,8 @@ TEST_CASE(ServerUnitFixture, test_props_runtime_shape) { TEST_ASSERT(rt["kv_cache_v"].get() == "tq3_0"); TEST_ASSERT(rt["lazy_draft"].get() == false); TEST_ASSERT(rt["draft_residency"].get() == "persistent"); + TEST_ASSERT(rt["decode_mode"].get() == "speculation"); + TEST_ASSERT(body["decode_mode"].get() == "speculation"); TEST_ASSERT(rt["target_sharding"].get() == false); TEST_ASSERT(rt["chunk"].get() == 512); TEST_ASSERT(rt["target_device"].get() == "auto:0"); From 37423f2f5eb8dd064062f2e3e9f818bdd3f376c6 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 17:08:33 +0000 Subject: [PATCH 10/42] feat(draft): export calibrated pre-norm hidden state --- server/src/common/dflash_draft_kv.cpp | 7 +++-- server/src/common/dflash_draft_kv.h | 1 + server/src/draft/draft_graph.cpp | 7 ++++- server/src/draft/draft_graph.h | 1 + server/test/smoke_draft_graph.cpp | 37 ++++++++++++++++++++++++++- 5 files changed, 49 insertions(+), 4 deletions(-) diff --git a/server/src/common/dflash_draft_kv.cpp b/server/src/common/dflash_draft_kv.cpp index 05ec25aca..1442a495a 100644 --- a/server/src/common/dflash_draft_kv.cpp +++ b/server/src/common/dflash_draft_kv.cpp @@ -96,9 +96,12 @@ bool draft_kv_init(DraftKvState & st, si.mask_swa = st.mask_swa; si.lm_head = lm_head; DraftGraphOutputs go = build_draft_kv_step(st.g_ctx, st.gf, dw, st.cache, si); - if (!go.hidden_states) return false; + if (!go.hidden_prenorm || !go.hidden_states) return false; + st.hidden_prenorm = go.hidden_prenorm; st.hidden_states = go.hidden_states; st.logits = go.logits; + ggml_set_output(st.hidden_prenorm); + ggml_build_forward_expand(st.gf, st.hidden_prenorm); ggml_set_output(st.hidden_states); ggml_build_forward_expand(st.gf, st.hidden_states); if (st.logits) { @@ -142,7 +145,7 @@ void draft_kv_free(DraftKvState & st) { if (st.mem_ctx) { ggml_free(st.mem_ctx); st.mem_ctx = nullptr; } st.meta_arena.clear(); st.meta_arena.shrink_to_fit(); - st.hidden_states = st.logits = nullptr; + st.hidden_prenorm = st.hidden_states = st.logits = nullptr; st.cache.k.clear(); st.cache.v.clear(); st.slot_pos.clear(); diff --git a/server/src/common/dflash_draft_kv.h b/server/src/common/dflash_draft_kv.h index 888f9b46b..9f80044da 100644 --- a/server/src/common/dflash_draft_kv.h +++ b/server/src/common/dflash_draft_kv.h @@ -64,6 +64,7 @@ struct DraftKvState { ggml_context * g_ctx = nullptr; ggml_cgraph * gf = nullptr; ggml_gallocr_t galloc = nullptr; + ggml_tensor * hidden_prenorm = nullptr; ggml_tensor * hidden_states = nullptr; ggml_tensor * logits = nullptr; // iff lm_head passed at init diff --git a/server/src/draft/draft_graph.cpp b/server/src/draft/draft_graph.cpp index 5886177dc..6ad18e8ec 100644 --- a/server/src/draft/draft_graph.cpp +++ b/server/src/draft/draft_graph.cpp @@ -268,12 +268,15 @@ DraftGraphOutputs build_draft_graph( } } - // ── 3. Final norm + // ── 3. Final norm. DSpark's confidence head is calibrated on h before + // this normalization, so retain both tensors as distinct graph outputs. + ggml_set_name(h, "draft_hidden_prenorm"); ggml_tensor * out = ggml_rms_norm(ctx, h, eps); out = ggml_mul(ctx, out, w.out_norm); ggml_set_name(out, "draft_hidden_out"); DraftGraphOutputs og{}; + og.hidden_prenorm = h; og.hidden_states = out; og.logits = nullptr; @@ -447,10 +450,12 @@ DraftGraphOutputs build_draft_kv_step( } ggml_tensor * out = ggml_rms_norm(ctx, h, eps); + ggml_set_name(h, "draft_kv_hidden_prenorm"); out = ggml_mul(ctx, out, w.out_norm); ggml_set_name(out, "draft_kv_hidden_out"); DraftGraphOutputs og{}; + og.hidden_prenorm = h; og.hidden_states = out; og.logits = nullptr; if (in.lm_head) { diff --git a/server/src/draft/draft_graph.h b/server/src/draft/draft_graph.h index b89429dac..1be162963 100644 --- a/server/src/draft/draft_graph.h +++ b/server/src/draft/draft_graph.h @@ -30,6 +30,7 @@ struct DraftGraphInputs { }; struct DraftGraphOutputs { + ggml_tensor * hidden_prenorm; // [hidden, q_len, 1] before final RMSNorm ggml_tensor * hidden_states; // [hidden, q_len, 1] (always set) ggml_tensor * logits; // [vocab, q_len, 1] (non-null iff lm_head was provided) }; diff --git a/server/test/smoke_draft_graph.cpp b/server/test/smoke_draft_graph.cpp index 544f8b51f..57e9b69dd 100644 --- a/server/test/smoke_draft_graph.cpp +++ b/server/test/smoke_draft_graph.cpp @@ -20,6 +20,7 @@ #include "ggml-alloc.h" #include "ggml-backend.h" #include "ggml-cuda.h" +#include #include #include @@ -103,11 +104,22 @@ int main(int argc, char ** argv) { gi.positions_k = pos_k; DraftGraphOutputs go = build_draft_graph(gctx, w, gi); - if (!go.hidden_states) { std::fprintf(stderr, "build_draft_graph returned null\n"); return 1; } + if (!go.hidden_prenorm || !go.hidden_states) { + std::fprintf(stderr, "build_draft_graph returned null output\n"); + return 1; + } + ggml_tensor * rebuilt_hidden = ggml_rms_norm( + gctx, go.hidden_prenorm, DFLASH27B_RMS_EPS); + rebuilt_hidden = ggml_mul(gctx, rebuilt_hidden, w.out_norm); + ggml_set_name(rebuilt_hidden, "rebuilt_draft_hidden_out"); + ggml_set_output(go.hidden_prenorm); ggml_set_output(go.hidden_states); + ggml_set_output(rebuilt_hidden); ggml_cgraph * gf = ggml_new_graph(gctx); + ggml_build_forward_expand(gf, go.hidden_prenorm); ggml_build_forward_expand(gf, go.hidden_states); + ggml_build_forward_expand(gf, rebuilt_hidden); std::printf("graph built: n_nodes=%d\n", ggml_graph_n_nodes(gf)); // ── 5. Allocate graph + all input tensors on the backend @@ -158,6 +170,29 @@ int main(int argc, char ** argv) { } std::vector out(n_out_elems); ggml_backend_tensor_get(go.hidden_states, out.data(), 0, sizeof(float) * out.size()); + std::vector prenorm(n_out_elems); + std::vector rebuilt(n_out_elems); + ggml_backend_tensor_get(go.hidden_prenorm, prenorm.data(), 0, + sizeof(float) * prenorm.size()); + ggml_backend_tensor_get(rebuilt_hidden, rebuilt.data(), 0, + sizeof(float) * rebuilt.size()); + + double max_prenorm_delta = 0.0; + double max_rebuild_error = 0.0; + for (size_t i = 0; i < out.size(); ++i) { + max_prenorm_delta = std::max( + max_prenorm_delta, std::fabs((double)prenorm[i] - out[i])); + max_rebuild_error = std::max( + max_rebuild_error, std::fabs((double)rebuilt[i] - out[i])); + } + if (max_prenorm_delta < 1e-5 || max_rebuild_error > 1e-5) { + std::fprintf(stderr, + "FAIL: pre-norm calibration output delta=%.8g rebuild_error=%.8g\n", + max_prenorm_delta, max_rebuild_error); + return 1; + } + std::printf("pre-norm export OK: delta=%.6g rebuild_error=%.6g\n", + max_prenorm_delta, max_rebuild_error); int n_nan = 0, n_inf = 0; double sum = 0.0, sumsq = 0.0; From 2fdfa648bd0491001e8a0064e86f370c7681c943 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 17:29:33 +0000 Subject: [PATCH 11/42] feat(qwen35): add concurrent DSpark chain execution --- server/CMakeLists.txt | 10 + .../common/concurrency/chain_spec_shapes.h | 86 +++ server/src/common/concurrency/seq_engine.h | 6 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 610 +++++++++++++++++- .../qwen35/concurrency/qwen35_seq_engine.h | 14 +- server/src/qwen35/qwen35_backend.cpp | 40 +- server/src/server/scheduler.cpp | 6 + server/test/test_chain_spec_shapes.cpp | 54 ++ server/test/test_seq_batch_plan.cpp | 16 + 9 files changed, 833 insertions(+), 9 deletions(-) create mode 100644 server/src/common/concurrency/chain_spec_shapes.h create mode 100644 server/test/test_chain_spec_shapes.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index ad6ef874f..6b763b1e8 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1442,6 +1442,16 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/src) list(APPEND _raw_unit_test_targets test_ddtree_path) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_chain_spec_shapes.cpp") + # Pure host-side DSpark chain topology and mixed-launch arithmetic. + add_executable(test_chain_spec_shapes + test/test_chain_spec_shapes.cpp + src/common/ddtree.cpp) + target_include_directories(test_chain_spec_shapes PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_chain_spec_shapes) + 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) diff --git a/server/src/common/concurrency/chain_spec_shapes.h b/server/src/common/concurrency/chain_spec_shapes.h new file mode 100644 index 000000000..12907bef7 --- /dev/null +++ b/server/src/common/concurrency/chain_spec_shapes.h @@ -0,0 +1,86 @@ +// Pure host-side shape helpers for path-shaped DSpark verification. + +#pragma once + +#include "common/ddtree.h" + +#include +#include +#include + +namespace dflash::common { + +inline int chain_decode_bucket_width(int lanes) { + static constexpr int buckets[] = { + 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, + }; + if (lanes <= 0) return 0; + for (int bucket : buckets) { + if (bucket >= lanes) return bucket; + } + return 64; +} + +// draft_tokens[0] is the already-pending root; positions 1.. form the +// proposal. DDTree's flat indices then coincide with chain depth. +inline DDTree make_dspark_chain_tree( + const std::vector & draft_tokens) { + DDTree tree; + if (draft_tokens.size() <= 1) return tree; + + tree.n_nodes = static_cast(draft_tokens.size()) - 1; + tree.token_ids.assign(draft_tokens.begin() + 1, draft_tokens.end()); + tree.depths.resize(static_cast(tree.n_nodes)); + tree.parents.resize(static_cast(tree.n_nodes) + 1); + tree.child_maps.resize(static_cast(tree.n_nodes) + 1); + tree.parents[0] = -1; + for (int node = 1; node <= tree.n_nodes; ++node) { + tree.depths[static_cast(node) - 1] = node; + tree.parents[static_cast(node)] = node - 1; + tree.child_maps[static_cast(node) - 1] + [tree.token_ids[static_cast(node) - 1]] = node; + } + + const int width = tree.n_nodes + 1; + tree.visibility.assign(static_cast(width) * width, 0); + for (int row = 0; row < width; ++row) { + for (int col = 0; col <= row; ++col) { + tree.visibility[static_cast(row) * width + col] = 1; + } + } + return tree; +} + +struct ChainLaunchShape { + int spec_lanes = 0; + int tree_bucket = 0; + int tree_rows = 0; + int ar_lanes = 0; + int ar_bucket = 0; + int accepted_rows = 0; + int commit_rows = 0; +}; + +inline ChainLaunchShape chain_launch_shape( + const std::vector & admitted, + const std::vector & accepted_lengths, + int tree_width) { + ChainLaunchShape shape; + const size_t count = admitted.size(); + for (size_t i = 0; i < count; ++i) { + if (admitted[i]) { + ++shape.spec_lanes; + if (i < accepted_lengths.size()) { + shape.accepted_rows += std::max(0, accepted_lengths[i]); + } + } + } + shape.ar_lanes = static_cast(count) - shape.spec_lanes; + shape.tree_bucket = chain_decode_bucket_width(shape.spec_lanes); + shape.tree_rows = shape.tree_bucket * std::max(0, tree_width); + shape.ar_bucket = chain_decode_bucket_width(shape.ar_lanes); + shape.commit_rows = shape.accepted_rows + shape.ar_bucket; + return shape; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 812bc54dd..c67a2357d 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -207,6 +207,8 @@ class SeqEngine { uint64_t ddtree_steps = 0; uint64_t ddtree_accepted_tokens = 0; uint64_t ddtree_suspensions = 0; + uint64_t spec_steps = 0; + uint64_t spec_accepted_tokens = 0; uint64_t target_forwards = 0; uint64_t kvflash_page_ins = 0; uint64_t kvflash_page_outs = 0; @@ -336,6 +338,8 @@ inline std::string validate_step_result( return "failed decode exposes token payload"; if (output.ddtree_suspensions != 0) return "failed decode carries DDTree suspension telemetry"; + if (output.spec_steps != 0 || output.spec_accepted_tokens != 0) + return "failed decode carries chain speculation telemetry"; } else { if (output.token < 0) return "successful decode has no pending token"; @@ -351,6 +355,8 @@ inline std::string validate_step_result( return "decode output burst contains an invalid token"; if (output.ddtree_suspensions > output.ddtree_steps) return "DDTree suspension has no successful DDTree step"; + if (output.spec_accepted_tokens != 0 && output.spec_steps == 0) + return "chain acceptance has no successful speculation 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 349abd3da..3edd65d49 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -14,6 +14,8 @@ #include "common/sampler.h" #include "common/ddtree.h" #include "common/geometric_draft_topk_cuda.h" +#include "common/dspark_head.h" +#include "common/concurrency/chain_spec_shapes.h" #include "internal.h" #include @@ -40,7 +42,7 @@ int decode_bucket_width(int live_count) { 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 tree_scratch_stride, SpecMode spec_mode, int max_prefills, int mixed_prefill_tokens, int long_mixed_prefill_tokens, int long_prefill_threshold, int idle_prefill_tokens, int prefill_quantum) @@ -54,7 +56,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( 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) { + tree_scratch_stride_(tree_scratch_stride), spec_mode_(spec_mode) { const int n_slots = slots_.slot_count(); slot_draft_kv_.resize((size_t)n_slots); @@ -141,7 +143,8 @@ DraftKvState * Qwen35SeqEngine::ensure_slot_draft_kv(int slot) { } bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { - if (tree_width_ <= 1 || !capture_features_ || !plan.prefills.empty() || + if (spec_mode_ != SpecMode::ddtree || 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; @@ -169,6 +172,590 @@ bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { } return true; } +bool Qwen35SeqEngine::chain_spec_input_eligible( + const StepInput & in) const { + if (spec_mode_ != SpecMode::dspark_chain || !capture_features_ || + tree_width_ <= 1 || tree_width_ > 16 || + b_.dw_.block_size != tree_width_ || !b_.dw_.dspark.enabled || + !in.allow_speculation || + in.speculation_policy == SpeculationPolicy::Never || + in.slot < 0 || in.slot >= slots_.slot_count() || + !slots_.slot(in.slot).decoding() || + slots_.slot(in.slot).sampler.needs_logit_processing() || + slots_.slot(in.slot).cur_pos < 1 || + slots_.slot(in.slot).cur_pos >= slots_.max_context()) { + return false; + } + const char * floor_value = std::getenv("DFLASH_MIN_TOKENS"); + const int floor = floor_value + ? std::max(0, std::atoi(floor_value)) : 0; + return slots_.slot(in.slot).generated_tokens() >= floor; +} +std::optional Qwen35SeqEngine::step_chain_spec( + const StepPlan & plan, const std::vector & admitted) { + StepResult result; + const std::vector & inputs = plan.decode; + if (admitted.size() != inputs.size() || !plan.prefills.empty()) { + result.error = "invalid DSpark chain admission plan"; + return result; + } + + const int T = tree_width_; + const int hidden = b_.w_.n_embd; + const int n_head_kv = b_.w_.n_head_kv; + const int n_slots = slots_.slot_count(); + int spec_count = 0; + for (uint8_t value : admitted) spec_count += value != 0; + if (spec_count == 0) return std::nullopt; + const int tree_bucket = chain_decode_bucket_width(spec_count); + + struct Proposal { + size_t input_index = 0; + int slot = -1; + int32_t root = -1; + DDTree tree; + std::vector flat; + std::vector accepted; + std::vector path; + std::vector confidence; + int32_t verify_bonus = -1; + int32_t pending = -1; + }; + struct ArLane { + size_t input_index = 0; + int slot = -1; + int32_t token = -1; + int position = -1; + int64_t physical_row = -1; + int32_t pending = -1; + }; + + std::vector proposals; + proposals.reserve(static_cast(spec_count)); + std::vector proposal_for_input(inputs.size(), -1); + std::vector drafted_slots; + drafted_slots.reserve(inputs.size()); + + auto proposal_fallback = [&]() -> std::optional { + for (int slot : drafted_slots) { + if (slot >= 0 && slot < static_cast(slot_draft_kv_.size()) && + slot_draft_kv_[static_cast(slot)]) { + draft_kv_reset(*slot_draft_kv_[static_cast(slot)]); + } + } + return std::nullopt; + }; + + const char * draft_always_env = std::getenv("DFLASH_SPEC_DRAFT_ALWAYS"); + const bool draft_always = + draft_always_env && std::atoi(draft_always_env) != 0; + std::vector noise(static_cast(T), b_.w_.mask_token_id); + std::vector noise_embed(static_cast(hidden) * T); + std::vector local_hidden(static_cast(hidden) * T); + std::vector prenorm_hidden(static_cast(hidden) * T); + + // Serial C3 baseline: each slot owns a context-KV ring. C6 replaces this + // loop with the packed multi-lane graph while retaining it as fallback. + for (size_t i = 0; i < inputs.size(); ++i) { + const StepInput & in = inputs[i]; + const bool hard_eligible = chain_spec_input_eligible(in); + if (admitted[i] && !hard_eligible) return proposal_fallback(); + if (!hard_eligible || (!admitted[i] && !draft_always)) continue; + + DraftKvState * draft = ensure_slot_draft_kv(in.slot); + DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); + if (!draft || !mirror) return proposal_fallback(); + drafted_slots.push_back(in.slot); + + noise[0] = in.token; + std::fill(noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed(noise.data(), T, 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_get_async( + b_.draft_backend_, draft->hidden_states, local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + ggml_backend_tensor_get_async( + b_.draft_backend_, draft->hidden_prenorm, prenorm_hidden.data(), 0, + sizeof(float) * prenorm_hidden.size()); + ggml_backend_synchronize(b_.draft_backend_); + + std::vector draft_tokens; + std::vector confidence; + if (!dspark_markov_correct_greedy_chain_fused( + b_.dw_, b_.draft_backend_, b_.w_.output, + local_hidden.data(), T, in.token, draft_tokens, + &confidence, prenorm_hidden.data()) || + static_cast(draft_tokens.size()) != T) { + return proposal_fallback(); + } + + if (!admitted[i]) continue; + Proposal proposal; + proposal.input_index = i; + proposal.slot = in.slot; + proposal.root = in.token; + proposal.flat = std::move(draft_tokens); + proposal.tree = make_dspark_chain_tree(proposal.flat); + proposal.confidence = std::move(confidence); + if (proposal.tree.n_nodes + 1 != T) return proposal_fallback(); + proposal_for_input[i] = static_cast(proposals.size()); + proposals.push_back(std::move(proposal)); + } + if (static_cast(proposals.size()) != spec_count) { + return proposal_fallback(); + } + + // Launch 1: scratch-only packed path-tree verification. + StepGraph & tree_sg = b_.sg_; + int max_prefix = 1; + for (const Proposal & proposal : proposals) { + max_prefix = std::max( + max_prefix, slots_.slot(proposal.slot).cur_pos); + } + if (!build_target_step_paged_tree( + tree_sg, b_.w_, b_.cache_, b_.target_backend_, + T, tree_bucket, max_prefix, + tree_scratch_base_, tree_scratch_stride_, + b_.cfg_.kq_stride_pad)) { + result.error = "packed DSpark chain verify graph build failed"; + return result; + } + + const int total_tree = T * tree_bucket; + std::vector flat_tokens(static_cast(total_tree), 0); + std::vector parents(static_cast(total_tree), -1); + std::vector sizes(static_cast(tree_bucket), 0); + std::vector tree_slots(static_cast(tree_bucket), -1); + std::vector tree_state_slots( + static_cast(tree_bucket), 0); + std::vector query_slots(static_cast(total_tree), -1); + std::vector tree_rows( + static_cast(total_tree) * n_head_kv, scratch_row_); + std::vector tree_positions( + static_cast(4) * total_tree, 0); + std::vector tree_embed( + static_cast(hidden) * total_tree, 0.0f); + seq_lens_.assign(static_cast(n_slots), 0); + + for (int lane = 0; lane < spec_count; ++lane) { + const Proposal & proposal = proposals[static_cast(lane)]; + const int base = lane * T; + sizes[static_cast(lane)] = T; + tree_slots[static_cast(lane)] = proposal.slot; + tree_state_slots[static_cast(lane)] = proposal.slot; + seq_lens_[static_cast(proposal.slot)] = + slots_.slot(proposal.slot).cur_pos; + for (int node = 0; node < T; ++node) { + const int row = base + node; + flat_tokens[static_cast(row)] = + proposal.flat[static_cast(node)]; + parents[static_cast(row)] = node == 0 + ? -1 : proposal.tree.parents[static_cast(node)]; + query_slots[static_cast(row)] = proposal.slot; + const int depth = node == 0 + ? 0 : proposal.tree.depths[static_cast(node) - 1]; + const int position = + slots_.slot(proposal.slot).cur_pos + depth; + tree_positions[static_cast(0) * total_tree + row] = + position; + tree_positions[static_cast(1) * total_tree + row] = + position; + tree_positions[static_cast(2) * total_tree + row] = + position; + for (int head = 0; head < n_head_kv; ++head) { + tree_rows[static_cast(head) * total_tree + row] = + static_cast(tree_scratch_base_) + + static_cast(proposal.slot) * + tree_scratch_stride_ + node; + } + } + } + + if (!b_.w_.embedder.embed( + flat_tokens.data(), total_tree, tree_embed.data())) { + result.error = "packed DSpark chain 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_positions.data(), 0, + sizeof(int32_t) * tree_positions.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()); + 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, + 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 DSpark chain verify compute failed"; + return result; + } + + std::vector posterior(static_cast(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 lane = 0; lane < spec_count; ++lane) { + Proposal & proposal = proposals[static_cast(lane)]; + const int32_t * lane_posterior = + posterior.data() + static_cast(lane) * T; + proposal.accepted = follow_verified_tree( + proposal.tree, lane_posterior, proposal.verify_bonus); + const int room = + slots_.max_context() - slots_.slot(proposal.slot).cur_pos; + truncate_verified_path( + proposal.accepted, static_cast(std::max(0, room)), + lane_posterior, proposal.verify_bonus); + if (proposal.accepted.empty()) { + result.error = "DSpark accepted path has no context headroom"; + return result; + } + proposal.path.reserve(proposal.accepted.size()); + for (int flat_index : proposal.accepted) { + proposal.path.push_back(flat_index == 0 + ? proposal.root + : proposal.tree.token_ids[ + static_cast(flat_index) - 1]); + } + replay_total += static_cast(proposal.path.size()); + } + + // Stage accepted path segments and all non-admitted AR peers. Nothing is + // published to slot history until the combined durable graph succeeds. + std::vector replay_segments; + std::vector replay_tokens; + std::vector replay_slots; + std::vector replay_positions; + std::vector replay_physical; + replay_segments.reserve(static_cast(spec_count)); + replay_tokens.reserve(static_cast(replay_total)); + replay_slots.reserve(static_cast(replay_total)); + replay_positions.reserve(static_cast(replay_total)); + replay_physical.reserve(static_cast(replay_total)); + seq_lens_.assign(static_cast(n_slots), 0); + int max_kv_len = 1; + int replay_offset = 0; + + for (Proposal & proposal : proposals) { + const Qwen35SlotManager::StepAppend app = slots_.append_tokens( + proposal.slot, proposal.path.data(), + static_cast(proposal.path.size())); + const bool table_ok = slots_.residency_active() || + upload_block_table_delta( + proposal.slot, app.first_new_block, + app.new_blocks.data(), app.new_blocks.size()); + if (!app.ok || app.physical_rows.size() != proposal.path.size() || + !table_ok) { + result.error = app.busy + ? "paged KV pool exhausted during DSpark chain commit" + : "DSpark accepted-path K/V append failed"; + return result; + } + replay_segments.push_back({ + replay_offset, static_cast(proposal.path.size()), + proposal.slot, + }); + for (size_t row = 0; row < proposal.path.size(); ++row) { + replay_tokens.push_back(proposal.path[row]); + replay_slots.push_back(proposal.slot); + replay_positions.push_back(app.position + static_cast(row)); + replay_physical.push_back(app.physical_rows[row]); + } + replay_offset += static_cast(proposal.path.size()); + const int seq_len = + app.position + static_cast(proposal.path.size()); + seq_lens_[static_cast(proposal.slot)] = seq_len; + max_kv_len = std::max(max_kv_len, seq_len); + } + + std::vector ar_lanes; + ar_lanes.reserve(inputs.size() - static_cast(spec_count)); + std::vector ar_for_input(inputs.size(), -1); + for (size_t i = 0; i < inputs.size(); ++i) { + if (admitted[i]) continue; + const StepInput & in = inputs[i]; + const Qwen35SlotManager::StepAppend app = + slots_.append_token(in.slot, in.token); + if (!app.ok) { + result.error = app.busy + ? "paged KV pool exhausted during mixed AR commit" + : "mixed AR K/V append failed"; + return result; + } + 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) { + result.error = "mixed AR block-table update failed"; + return result; + } + ArLane lane; + lane.input_index = i; + lane.slot = in.slot; + lane.token = in.token; + lane.position = app.position; + lane.physical_row = app.physical_row; + ar_for_input[i] = static_cast(ar_lanes.size()); + ar_lanes.push_back(lane); + seq_lens_[static_cast(in.slot)] = app.position + 1; + max_kv_len = std::max(max_kv_len, app.position + 1); + } + if (!upload_all_active_block_tables()) { + result.error = "DSpark mixed-step block-table refresh failed"; + return result; + } + + // Launch 2: accepted path segments + compact AR rows in the same builder + // combination already used by mixed prefill/decode. + const int ar_count = static_cast(ar_lanes.size()); + const int ar_bucket = chain_decode_bucket_width(ar_count); + const int n_total = replay_total + ar_bucket; + const int gather_rows = spec_count + ar_bucket; + StepGraph & durable_sg = b_.sg_; + if (!build_target_step( + durable_sg, b_.w_, b_.cache_, b_.target_backend_, + 0, n_total, false, true, false, 0, 0, + b_.cfg_.kq_stride_pad, false, false, false, true, + ar_bucket > 0 ? ar_bucket : 1, 0, max_kv_len, + replay_total, replay_segments.data(), + static_cast(replay_segments.size()), + gather_rows, ar_bucket > 0) || + !durable_sg.kv_write_rows || !durable_sg.target_feat_rows || + !durable_sg.paged_query_seq_ids || + !durable_sg.paged_query_positions || + !durable_sg.logits_row_indices || !durable_sg.argmax_tokens) { + result.error = "DSpark mixed commit/AR graph build failed"; + return result; + } + + std::vector durable_tokens(static_cast(n_total), 0); + std::copy(replay_tokens.begin(), replay_tokens.end(), + durable_tokens.begin()); + for (int lane = 0; lane < ar_count; ++lane) { + durable_tokens[static_cast(replay_total + lane)] = + ar_lanes[static_cast(lane)].token; + } + embed_buf_.resize(static_cast(hidden) * n_total); + if (!b_.w_.embedder.embed( + durable_tokens.data(), n_total, embed_buf_.data())) { + result.error = "DSpark mixed commit/AR embedding failed"; + return result; + } + ggml_backend_tensor_set( + durable_sg.inp_embed, embed_buf_.data(), 0, + sizeof(float) * embed_buf_.size()); + + pos_buf_.assign(static_cast(4) * n_total, 0); + for (int row = 0; row < replay_total; ++row) { + const int position = replay_positions[static_cast(row)]; + pos_buf_[static_cast(0) * n_total + row] = position; + pos_buf_[static_cast(1) * n_total + row] = position; + pos_buf_[static_cast(2) * n_total + row] = position; + } + for (int lane = 0; lane < ar_count; ++lane) { + const int row = replay_total + lane; + const int position = ar_lanes[static_cast(lane)].position; + pos_buf_[static_cast(0) * n_total + row] = position; + pos_buf_[static_cast(1) * n_total + row] = position; + pos_buf_[static_cast(2) * n_total + row] = position; + } + ggml_backend_tensor_set(durable_sg.positions, pos_buf_.data(), 0, + sizeof(int32_t) * pos_buf_.size()); + + rows_buf_.assign( + static_cast(n_total) * n_head_kv, scratch_row_); + for (int head = 0; head < n_head_kv; ++head) { + for (int row = 0; row < replay_total; ++row) { + rows_buf_[static_cast(head) * n_total + row] = + replay_physical[static_cast(row)]; + } + for (int lane = 0; lane < ar_count; ++lane) { + rows_buf_[static_cast(head) * n_total + + replay_total + lane] = + ar_lanes[static_cast(lane)].physical_row; + } + } + ggml_backend_tensor_set( + durable_sg.kv_write_rows, rows_buf_.data(), 0, + sizeof(int64_t) * rows_buf_.size()); + + query_slot_ids_.assign(static_cast(n_total), -1); + query_positions_.assign(static_cast(n_total), -1); + for (int row = 0; row < replay_total; ++row) { + query_slot_ids_[static_cast(row)] = + replay_slots[static_cast(row)]; + query_positions_[static_cast(row)] = + replay_positions[static_cast(row)]; + } + for (int lane = 0; lane < ar_count; ++lane) { + const int row = replay_total + lane; + query_slot_ids_[static_cast(row)] = + ar_lanes[static_cast(lane)].slot; + query_positions_[static_cast(row)] = + ar_lanes[static_cast(lane)].position; + } + logits_rows_.clear(); + logits_rows_.reserve(static_cast(gather_rows)); + int path_end = 0; + for (const Proposal & proposal : proposals) { + path_end += static_cast(proposal.path.size()); + logits_rows_.push_back(path_end - 1); + } + for (int lane = 0; lane < ar_bucket; ++lane) { + logits_rows_.push_back(replay_total + lane); + } + ggml_backend_tensor_set( + durable_sg.paged_query_seq_ids, query_slot_ids_.data(), 0, + sizeof(int32_t) * query_slot_ids_.size()); + ggml_backend_tensor_set( + durable_sg.paged_query_positions, query_positions_.data(), 0, + sizeof(int32_t) * query_positions_.size()); + ggml_backend_tensor_set( + durable_sg.logits_row_indices, logits_rows_.data(), 0, + sizeof(int32_t) * logits_rows_.size()); + + active_slot_ids_.assign(static_cast(ar_bucket), -1); + state_slot_ids_.assign(static_cast(ar_bucket), 0); + for (int lane = 0; lane < ar_count; ++lane) { + active_slot_ids_[static_cast(lane)] = + ar_lanes[static_cast(lane)].slot; + state_slot_ids_[static_cast(lane)] = + ar_lanes[static_cast(lane)].slot; + } + if (ar_bucket > 0) { + ggml_backend_tensor_set( + durable_sg.active_slot_ids, active_slot_ids_.data(), 0, + sizeof(int32_t) * active_slot_ids_.size()); + ggml_backend_tensor_set( + durable_sg.state_slot_ids, state_slot_ids_.data(), 0, + sizeof(int32_t) * state_slot_ids_.size()); + } + + const int feature_cap = b_.cache_.target_feat_cap; + const int dead_feature_row = feature_cap * n_slots; + feature_rows_.assign( + static_cast(n_total), dead_feature_row); + for (int row = 0; row < replay_total; ++row) { + feature_rows_[static_cast(row)] = + replay_slots[static_cast(row)] * feature_cap + + replay_positions[static_cast(row)] % feature_cap; + } + for (int lane = 0; lane < ar_count; ++lane) { + const ArLane & ar = ar_lanes[static_cast(lane)]; + feature_rows_[static_cast(replay_total + lane)] = + ar.slot * feature_cap + ar.position % feature_cap; + } + ggml_backend_tensor_set( + durable_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_, durable_sg.gf) != + GGML_STATUS_SUCCESS) { + result.error = "DSpark mixed commit/AR compute failed"; + return result; + } + + argmax_buf_.assign(static_cast(gather_rows), -1); + ggml_backend_tensor_get_async( + b_.target_backend_, durable_sg.argmax_tokens, + argmax_buf_.data(), 0, + sizeof(int32_t) * argmax_buf_.size()); + ggml_backend_synchronize(b_.target_backend_); + for (int lane = 0; lane < spec_count; ++lane) { + Proposal & proposal = proposals[static_cast(lane)]; + proposal.pending = argmax_buf_[static_cast(lane)]; + if (proposal.pending < 0) { + result.error = "DSpark durable replay produced invalid token"; + return result; + } + } + for (int lane = 0; lane < ar_count; ++lane) { + ArLane & ar = ar_lanes[static_cast(lane)]; + ar.pending = + argmax_buf_[static_cast(spec_count + lane)]; + if (ar.pending < 0) { + result.error = "mixed AR durable step produced invalid token"; + return result; + } + } + + std::vector write_slots; + write_slots.reserve(inputs.size()); + for (const StepInput & in : inputs) write_slots.push_back(in.slot); + if (!commit_residency_writes(write_slots)) { + result.error = "DSpark mixed-step KV write commit failed"; + return result; + } + for (const StepInput & in : inputs) { + slots_.commit_step(in.slot); + std::string reselect_error; + if (!maybe_reselect_residency(in.slot, reselect_error)) { + result.error = reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error; + return result; + } + } + + result.decode.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + DecodeOutput out; + out.slot = inputs[i].slot; + if (admitted[i]) { + Proposal & proposal = + proposals[static_cast(proposal_for_input[i])]; + out.token = proposal.pending; + out.spec_steps = 1; + out.spec_accepted_tokens = + proposal.path.size() > 1 + ? static_cast(proposal.path.size() - 1) + : 0; + out.target_forwards = 2; + out.committed_tokens.assign( + proposal.path.begin() + 1, proposal.path.end()); + } else { + ArLane & ar = + ar_lanes[static_cast(ar_for_input[i])]; + out.token = ar.pending; + out.target_forwards = 1; + } + attach_residency_telemetry(out); + result.decode.push_back(std::move(out)); + } + return result; +} + std::optional Qwen35SeqEngine::step_ddtree( const StepPlan & plan) { @@ -845,6 +1432,23 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // by taking the existing packed AR path for this iteration. } + if (spec_mode_ == SpecMode::dspark_chain && plan.prefills.empty()) { + std::vector admitted(inputs.size(), 0); + bool any_admitted = false; + for (size_t i = 0; i < inputs.size(); ++i) { + admitted[i] = chain_spec_input_eligible(inputs[i]) && + inputs[i].speculation_policy == + SpeculationPolicy::Always; + any_admitted = any_admitted || admitted[i] != 0; + } + if (any_admitted) { + std::optional speculative = + step_chain_spec(plan, admitted); + if (speculative) return std::move(*speculative); + // Proposal setup failed before target/cache mutation. Preserve + // service through the ordinary packed AR path this iteration. + } + } const TargetWeights & w = b_.w_; StepGraph & sg = b_.sg_; const int hidden = w.n_embd; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 75275df39..c0d7f5fcc 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -39,6 +39,12 @@ class Qwen35Backend; class Qwen35SeqEngine final : public SeqEngine { public: + enum class SpecMode { + none, + ddtree, + dspark_chain, + }; + // `pool` and `backend` must outlive the engine. `scratch_row` is the // first row of the block appended past the pool's index space, used as // the K/V write destination of graph-bucket padding rows. @@ -46,7 +52,9 @@ class Qwen35SeqEngine final : public SeqEngine { Qwen35SeqEngine(Qwen35Backend & backend, PagedKvPool & pool, int max_ctx, int64_t scratch_row, int tree_width = 0, int tree_scratch_base = 0, - int tree_scratch_stride = 0, int max_prefills = 8, + int tree_scratch_stride = 0, + SpecMode spec_mode = SpecMode::none, + int max_prefills = 8, int mixed_prefill_tokens = 2048, int long_mixed_prefill_tokens = 4096, int long_prefill_threshold = 768, @@ -119,9 +127,12 @@ class Qwen35SeqEngine final : public SeqEngine { DraftFeatureMirror * slot_feature_mirror(int slot); DraftKvState * ensure_slot_draft_kv(int slot); bool ddtree_eligible(const StepPlan & plan) const; + bool chain_spec_input_eligible(const StepInput & input) 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); + std::optional step_chain_spec( + const StepPlan & plan, const std::vector & admitted); Qwen35Backend & b_; Qwen35SlotManager slots_; @@ -130,6 +141,7 @@ class Qwen35SeqEngine final : public SeqEngine { int tree_scratch_base_ = 0; int tree_scratch_stride_ = 0; bool capture_features_ = false; + SpecMode spec_mode_ = SpecMode::none; ggml_context * feature_view_ctx_ = nullptr; std::vector slot_feature_mirrors_; std::vector> slot_draft_kv_; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 5adac6e28..eb6d160da 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -127,6 +127,7 @@ static int dflash_min_tokens_floor() { return value; } + static FILE * open_dflash_floor_log() { #if defined(_WIN32) // Simple append-mode log on Windows (no file size check). @@ -205,6 +206,8 @@ static int64_t concurrent_fixed_cache_bytes( } } // namespace +static bool qwen35_dspark_enabled(); + #define IS_EOS_TOK(tok, w) \ ( ((w).eos_chat_id >= 0 && (tok) == (w).eos_chat_id) \ || ((w).eos_id >= 0 && (tok) == (w).eos_id ) ) @@ -436,13 +439,27 @@ 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 && + const bool concurrent_local_draft = + n_slots > 1 && cfg_.draft_path && !use_remote_draft && !tensor_parallel && !split_gpus_ && target_backend_ == draft_backend_; + const bool concurrent_local_ddtree = + concurrent_local_draft && cfg_.ddtree_mode; + const bool concurrent_local_chain = + concurrent_local_draft && !cfg_.ddtree_mode && dw_.dspark.enabled && + qwen35_dspark_enabled() && + cfg_.speculation_policy != SpeculationPolicy::Never; + const bool concurrent_spec_tree = + concurrent_local_ddtree || concurrent_local_chain; + const Qwen35SeqEngine::SpecMode spec_mode = concurrent_local_ddtree + ? Qwen35SeqEngine::SpecMode::ddtree + : concurrent_local_chain + ? Qwen35SeqEngine::SpecMode::dspark_chain + : Qwen35SeqEngine::SpecMode::none; const int tree_width = concurrent_local_ddtree - ? cfg_.ddtree_budget + 1 : 0; - const int tree_stride = concurrent_local_ddtree + ? cfg_.ddtree_budget + 1 + : concurrent_local_chain ? dw_.block_size : 0; + const int tree_stride = concurrent_spec_tree ? paged_token_capacity(tree_width) : 0; const int64_t concurrent_scratch_tokens = (int64_t)n_slots * tree_stride + PAGED_BLOCK_SIZE; @@ -510,7 +527,7 @@ 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_spec_tree)) { std::fprintf(stderr, "cache: %s\n", dflash27b_last_error()); return false; } @@ -575,6 +592,7 @@ bool Qwen35Backend::init() { seq_engine_ = std::make_unique( *this, *paged_kv_pool_, cfg_.device.max_ctx, dead_scratch_row, tree_width, tree_scratch_base, tree_stride, + spec_mode, max_concurrent_prefills, mixed_prefill_tokens, long_mixed_prefill_tokens, long_prefill_threshold, idle_prefill_tokens, prefill_quantum); @@ -585,6 +603,18 @@ bool Qwen35Backend::init() { cfg_.ddtree_budget, tree_width, adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); } + if (concurrent_local_chain) { + std::fprintf(stderr, + "[parallel-dspark] enabled width=%d mode=packed-chain-verify " + "decode_mode=%s draft=q4-mix-compatible\n", + tree_width, + speculation_policy_name(cfg_.speculation_policy)); + } else if (concurrent_local_draft && !cfg_.ddtree_mode && + cfg_.speculation_policy != SpeculationPolicy::Never) { + std::fprintf(stderr, + "[parallel-dspark] disabled: drafter lacks usable DSpark " + "Markov/confidence heads; using packed AR\n"); + } 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 c4adf8419..420c785dd 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -42,6 +42,8 @@ struct SchedSlot { uint64_t ddtree_steps = 0; uint64_t ddtree_accepted_tokens = 0; uint64_t ddtree_suspensions = 0; + uint64_t spec_steps = 0; + uint64_t spec_accepted_tokens = 0; uint64_t target_forwards = 0; uint64_t kvflash_page_ins = 0; uint64_t kvflash_page_outs = 0; @@ -328,6 +330,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { {"ddtree_steps", s.ddtree_steps}, {"ddtree_accepted_tokens", s.ddtree_accepted_tokens}, {"ddtree_suspensions", s.ddtree_suspensions}, + {"spec_steps", s.spec_steps}, + {"spec_accepted_tokens", s.spec_accepted_tokens}, {"target_forwards", s.target_forwards}, {"kvflash_page_ins", s.kvflash_page_ins}, {"kvflash_page_outs", s.kvflash_page_outs}, @@ -758,6 +762,8 @@ 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.spec_steps += out.spec_steps; + s.spec_accepted_tokens += out.spec_accepted_tokens; 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_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp new file mode 100644 index 000000000..42b4a20b2 --- /dev/null +++ b/server/test/test_chain_spec_shapes.cpp @@ -0,0 +1,54 @@ +#include "common/concurrency/chain_spec_shapes.h" +#include "host_check.h" + +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +int main() { + const std::vector draft = {10, 11, 12, 13}; + const DDTree tree = make_dspark_chain_tree(draft); + CHECK(tree.n_nodes == 3); + CHECK((tree.token_ids == std::vector{11, 12, 13})); + CHECK((tree.depths == std::vector{1, 2, 3})); + CHECK((tree.parents == std::vector{-1, 0, 1, 2})); + + int pending = -1; + const int32_t full_posterior[] = {11, 12, 13, 14}; + std::vector accepted = + follow_verified_tree(tree, full_posterior, pending); + CHECK((accepted == std::vector{0, 1, 2, 3})); + CHECK(pending == 14); + + const int32_t rejected_posterior[] = {11, 99, 13, 14}; + accepted = follow_verified_tree(tree, rejected_posterior, pending); + CHECK((accepted == std::vector{0, 1})); + CHECK(pending == 99); + + CHECK(truncate_verified_path( + accepted, 1, rejected_posterior, pending)); + CHECK((accepted == std::vector{0})); + CHECK(pending == 11); + + const ChainLaunchShape mixed = chain_launch_shape( + {1, 0, 1, 0, 0, 0}, {4, 0, 2, 0, 0, 0}, 16); + CHECK(mixed.spec_lanes == 2); + CHECK(mixed.tree_bucket == 2); + CHECK(mixed.tree_rows == 32); + CHECK(mixed.ar_lanes == 4); + CHECK(mixed.ar_bucket == 4); + CHECK(mixed.accepted_rows == 6); + CHECK(mixed.commit_rows == 10); + + const ChainLaunchShape all_spec = chain_launch_shape( + {1, 1, 1}, {1, 2, 3}, 16); + CHECK(all_spec.tree_bucket == 3); + CHECK(all_spec.ar_bucket == 0); + CHECK(all_spec.commit_rows == 6); + + std::printf("chain spec shape tests passed: %d checks\n", g_checks); + return 0; +} diff --git a/server/test/test_seq_batch_plan.cpp b/server/test/test_seq_batch_plan.cpp index e1e13065e..2860cc695 100644 --- a/server/test/test_seq_batch_plan.cpp +++ b/server/test/test_seq_batch_plan.cpp @@ -135,6 +135,22 @@ int main() { burst.decode[0].target_forwards = 1; CHECK(validate_step_result(work, burst, 2).empty()); + SeqEngine::StepResult chain_burst = good; + chain_burst.decode[0].committed_tokens = {8, 9}; + chain_burst.decode[0].spec_steps = 1; + chain_burst.decode[0].spec_accepted_tokens = 2; + chain_burst.decode[0].target_forwards = 2; + CHECK(validate_step_result(work, chain_burst, 2).empty()); + + SeqEngine::StepResult orphan_chain_acceptance = good; + orphan_chain_acceptance.decode[0].spec_accepted_tokens = 1; + CHECK(!validate_step_result(work, orphan_chain_acceptance, 2).empty()); + + SeqEngine::StepResult failed_chain = good; + failed_chain.decode[0] = {0, -1, true, "decode failed"}; + failed_chain.decode[0].spec_steps = 1; + CHECK(!validate_step_result(work, failed_chain, 2).empty()); + SeqEngine::StepResult orphan_suspension = good; orphan_suspension.decode[0].ddtree_suspensions = 1; CHECK(!validate_step_result(work, orphan_suspension, 2).empty()); From e13ee33407cca4c1fb83745f7da3d268eb740f65 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 17:38:49 +0000 Subject: [PATCH 12/42] feat(concurrency): add adaptive speculation gate --- server/CMakeLists.txt | 16 + .../common/concurrency/spec_cost_profile.h | 107 +++++ .../src/common/concurrency/speculation_gate.h | 382 ++++++++++++++++++ server/test/test_spec_cost_profile.cpp | 80 ++++ server/test/test_speculation_gate.cpp | 192 +++++++++ 5 files changed, 777 insertions(+) create mode 100644 server/src/common/concurrency/spec_cost_profile.h create mode 100644 server/src/common/concurrency/speculation_gate.h create mode 100644 server/test/test_spec_cost_profile.cpp create mode 100644 server/test/test_speculation_gate.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 6b763b1e8..28c7e5f55 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1452,6 +1452,22 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/test) list(APPEND _raw_unit_test_targets test_chain_spec_shapes) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_speculation_gate.cpp") + 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_spec_cost_profile.cpp") + add_executable(test_spec_cost_profile + test/test_spec_cost_profile.cpp) + target_include_directories(test_spec_cost_profile PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_spec_cost_profile) + 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) diff --git a/server/src/common/concurrency/spec_cost_profile.h b/server/src/common/concurrency/spec_cost_profile.h new file mode 100644 index 000000000..19f9d7b1e --- /dev/null +++ b/server/src/common/concurrency/spec_cost_profile.h @@ -0,0 +1,107 @@ +// Pure startup profiling protocol for monotone speculation cost tables. + +#pragma once + +#include "common/concurrency/speculation_gate.h" + +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct SpecProfileGrid { + std::vector tree_rows; + std::vector step_rows; + std::vector draft_lanes; +}; + +inline void sort_unique_positive(std::vector & values) { + values.erase(std::remove_if(values.begin(), values.end(), + [](int value) { return value <= 0; }), + values.end()); + std::sort(values.begin(), values.end()); + values.erase(std::unique(values.begin(), values.end()), values.end()); +} + +inline SpecProfileGrid build_spec_profile_grid( + int max_concurrency, int tree_width, int max_accept, + const std::function & bucket) { + SpecProfileGrid grid; + if (max_concurrency <= 0 || tree_width <= 0 || max_accept <= 0) + return grid; + auto bucketed = [&](int lanes) { + if (lanes <= 0) return 0; + return std::max(lanes, bucket ? bucket(lanes) : lanes); + }; + for (int lanes = 1; lanes <= max_concurrency; ++lanes) { + grid.tree_rows.push_back(bucketed(lanes) * tree_width); + grid.draft_lanes.push_back(lanes); + } + for (int concurrency = 1; concurrency <= max_concurrency; ++concurrency) { + grid.step_rows.push_back(bucketed(concurrency)); + for (int spec_lanes = 1; spec_lanes <= concurrency; ++spec_lanes) { + const int ar_rows = bucketed(concurrency - spec_lanes); + for (int accepted = spec_lanes; + accepted <= spec_lanes * max_accept; ++accepted) { + grid.step_rows.push_back(accepted + ar_rows); + } + } + } + sort_unique_positive(grid.tree_rows); + sort_unique_positive(grid.step_rows); + sort_unique_positive(grid.draft_lanes); + return grid; +} + +struct SpecProfileResult { + SpecCostSeries table; + std::string error; + bool ok() const { return error.empty() && table.valid(); } +}; + +template +SpecProfileResult profile_monotonic_costs( + std::vector indices, Runner && runner, int reps = 5) { + SpecProfileResult result; + if (reps <= 0) { + result.error = "profiling repetitions must be positive"; + return result; + } + sort_unique_positive(indices); + if (indices.empty()) { + result.error = "profiling grid is empty"; + return result; + } + result.table.indices = indices; + result.table.costs.reserve(indices.size()); + for (int index : indices) { + (void)runner(index); // graph capture / allocator warmup + std::vector samples; + samples.reserve(static_cast(reps)); + for (int rep = 0; rep < reps; ++rep) { + const double sample = runner(index); + if (!std::isfinite(sample) || sample <= 0.0) { + result.error = "profiling runner returned an invalid cost"; + result.table = {}; + return result; + } + samples.push_back(sample); + } + std::sort(samples.begin(), samples.end()); + double median = samples[(size_t)reps / 2]; + if (reps % 2 == 0) { + median = 0.5 * (samples[(size_t)reps / 2 - 1] + median); + } + if (!result.table.costs.empty()) { + median = std::max(median, result.table.costs.back()); + } + result.table.costs.push_back(median); + } + return result; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h new file mode 100644 index 000000000..a7b443f91 --- /dev/null +++ b/server/src/common/concurrency/speculation_gate.h @@ -0,0 +1,382 @@ +// Per-request adaptive speculation policy over startup-profiled costs. +// Pure host code: no graph, backend, or scheduler types belong here. + +#pragma once + +#include "common/speculation_policy.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct SpecGateConfig { + int stale_after_tokens = 64; +}; + +struct SpecCostLookup { + double cost = std::numeric_limits::infinity(); + int requested_index = 0; + int profiled_index = 0; + bool clamped = false; + bool rounded_up = false; +}; + +struct SpecCostSeries { + std::vector indices; + std::vector costs; + + bool valid() const { + if (indices.empty() || indices.size() != costs.size()) return false; + for (size_t i = 0; i < indices.size(); ++i) { + if (indices[i] < 0 || !std::isfinite(costs[i]) || costs[i] <= 0.0) + return false; + if (i > 0 && (indices[i] <= indices[i - 1] || + costs[i] < costs[i - 1])) + return false; + } + return true; + } + + SpecCostLookup lookup(int index) const { + SpecCostLookup result; + result.requested_index = index; + if (!valid()) return result; + auto it = std::lower_bound(indices.begin(), indices.end(), index); + if (it == indices.end()) { + result.profiled_index = indices.back(); + result.cost = costs.back(); + result.clamped = true; + return result; + } + const size_t pos = static_cast(it - indices.begin()); + result.profiled_index = *it; + result.cost = costs[pos]; + result.clamped = index < indices.front() || index > indices.back(); + result.rounded_up = *it != index && !result.clamped; + return result; + } +}; + +struct SpecCostTables { + SpecCostSeries tree_cost; + SpecCostSeries step_cost; + SpecCostSeries draft_cost; + + bool valid() const { + return tree_cost.valid() && step_cost.valid() && draft_cost.valid(); + } +}; + +struct SpecCandidate { + uint64_t request_id = 0; + int slot = -1; + SpeculationPolicy policy = SpeculationPolicy::Adaptive; + bool eligible = false; + int generated_tokens = 0; + // NaN means no calibrated confidence is available. Otherwise this is + // already the survival-product expected yield, including the root. + double confidence_yield = std::numeric_limits::quiet_NaN(); +}; + +struct SpecStepGeometry { + int tree_width = 1; + std::function bucket = [](int lanes) { + return std::max(0, lanes); + }; + + int bucketed_lanes(int lanes) const { + return lanes <= 0 ? 0 : std::max(lanes, bucket ? bucket(lanes) : lanes); + } + int tree_rows(int spec_lanes) const { + return bucketed_lanes(spec_lanes) * std::max(1, tree_width); + } + int step_rows(int concurrency, int spec_lanes, + double expected_spec_tokens) const { + const int accepted_rows = std::max( + spec_lanes, static_cast(std::lround(expected_spec_tokens))); + return accepted_rows + bucketed_lanes(concurrency - spec_lanes); + } +}; + +struct SpecPlanScore { + uint64_t request_id = 0; + int slot = -1; + double expected_yield = 1.0; + bool forced = false; + bool admitted = false; +}; + +struct SpecPlan { + bool valid = true; + std::string error; + int concurrency = 0; + int admitted_count = 0; + int tree_rows = 0; + int step_rows = 0; + int draft_lanes = 0; + double expected_tokens = 0.0; + double predicted_cost = 0.0; + double goodput = 0.0; + double ar_goodput = 0.0; + bool cost_lookup_clamped = false; + std::vector ordered; + std::vector admitted_request_ids; + std::vector admitted_slots; +}; + +inline double confidence_survival_yield( + const std::vector & confidences, int max_accept) { + if (max_accept <= 1) return 1.0; + double expected = 1.0; + double survival = 1.0; + const int depth = std::min( + static_cast(confidences.size()), max_accept - 1); + for (int i = 0; i < depth; ++i) { + const double c = std::clamp(confidences[(size_t)i], 0.0, 1.0); + survival *= c; + expected += survival; + } + return std::clamp(expected, 1.0, static_cast(max_accept)); +} + +class SpeculationGate { +public: + using ClampLogger = std::function; + + SpeculationGate(SpecGateConfig config, SpecCostTables costs, + SpecStepGeometry geometry, int max_accept, + ClampLogger clamp_logger = {}) + : config_(config), costs_(std::move(costs)), + geometry_(std::move(geometry)), + max_accept_(std::max(1, max_accept)), + prior_yield_(std::max(1, max_accept)), + clamp_logger_(std::move(clamp_logger)) {} + + bool valid() const { + return config_.stale_after_tokens >= 1 && costs_.valid() && + geometry_.tree_width >= 1 && max_accept_ >= 1; + } + + // draft_lanes_override prices always-drafting. -1 means admitted-only. + SpecPlan plan(int concurrency, + const std::vector & candidates, + int k_cap, int draft_lanes_override = -1) { + SpecPlan out; + out.concurrency = concurrency; + if (!valid() || concurrency < 0 || + static_cast(concurrency) != candidates.size() || + k_cap < 0) { + out.valid = false; + out.error = "invalid speculation gate inputs"; + return out; + } + + struct Ranked { + const SpecCandidate * candidate = nullptr; + double score = 1.0; + bool forced = false; + }; + std::vector forced; + std::vector adaptive; + forced.reserve(candidates.size()); + adaptive.reserve(candidates.size()); + + for (const SpecCandidate & candidate : candidates) { + if (!candidate.eligible || + candidate.policy == SpeculationPolicy::Never) { + continue; + } + const double score = score_candidate(candidate); + Ranked ranked{&candidate, score, + candidate.policy == SpeculationPolicy::Always}; + (ranked.forced ? forced : adaptive).push_back(ranked); + } + auto request_order = [](const Ranked & a, const Ranked & b) { + return a.candidate->request_id < b.candidate->request_id; + }; + std::sort(forced.begin(), forced.end(), request_order); + std::sort(adaptive.begin(), adaptive.end(), + [](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; + }); + + if (static_cast(forced.size()) > k_cap) { + out.valid = false; + out.error = "forced speculation exceeds executor capacity"; + return out; + } + + std::vector ranked; + ranked.reserve(forced.size() + adaptive.size()); + ranked.insert(ranked.end(), forced.begin(), forced.end()); + ranked.insert(ranked.end(), adaptive.begin(), adaptive.end()); + for (const Ranked & item : ranked) { + out.ordered.push_back({item.candidate->request_id, + item.candidate->slot, + item.score, item.forced, false}); + } + + const int forced_count = static_cast(forced.size()); + const int max_k = std::min(k_cap, ranked.size()); + double expected_sum = 0.0; + for (int i = 0; i < forced_count; ++i) expected_sum += ranked[i].score; + + const SpecCostLookup ar_lookup = costs_.step_cost.lookup( + geometry_.bucketed_lanes(concurrency)); + report_clamp("step", ar_lookup, out); + out.ar_goodput = concurrency == 0 ? 0.0 + : static_cast(concurrency) / ar_lookup.cost; + + int best_k = forced_count; + double best_goodput = -1.0; + double best_cost = 0.0; + double best_expected = 0.0; + int best_tree_rows = 0; + int best_step_rows = geometry_.bucketed_lanes(concurrency); + int best_draft_lanes = 0; + + for (int k = forced_count; k <= max_k; ++k) { + if (k > forced_count) expected_sum += ranked[k - 1].score; + const double expected_tokens = + static_cast(concurrency - k) + expected_sum; + double cost = 0.0; + int tree_rows = 0; + int step_rows = geometry_.bucketed_lanes(concurrency); + const int draft_lanes = draft_lanes_override >= 0 + ? draft_lanes_override : k; + + if (k == 0 && draft_lanes == 0) { + cost = ar_lookup.cost; + } else { + if (k > 0) { + tree_rows = geometry_.tree_rows(k); + const SpecCostLookup tree = costs_.tree_cost.lookup(tree_rows); + report_clamp("tree", tree, out); + cost += tree.cost; + step_rows = geometry_.step_rows( + concurrency, k, expected_sum); + } + const SpecCostLookup step = costs_.step_cost.lookup(step_rows); + report_clamp("step", step, out); + cost += step.cost; + if (draft_lanes > 0) { + const SpecCostLookup draft = + costs_.draft_cost.lookup(draft_lanes); + report_clamp("draft", draft, out); + cost += draft.cost; + } + } + const double goodput = expected_tokens / cost; + if (goodput > best_goodput) { + best_goodput = goodput; + best_k = k; + best_cost = cost; + best_expected = expected_tokens; + best_tree_rows = tree_rows; + best_step_rows = step_rows; + best_draft_lanes = draft_lanes; + } + } + + out.admitted_count = best_k; + out.goodput = std::max(0.0, best_goodput); + out.predicted_cost = best_cost; + out.expected_tokens = best_expected; + out.tree_rows = best_tree_rows; + out.step_rows = best_step_rows; + out.draft_lanes = best_draft_lanes; + for (int i = 0; i < best_k; ++i) { + out.ordered[(size_t)i].admitted = true; + out.admitted_request_ids.push_back(ranked[(size_t)i].candidate->request_id); + out.admitted_slots.push_back(ranked[(size_t)i].candidate->slot); + } + return out; + } + + void observe(uint64_t request_id, double emitted_tokens, + int generated_tokens) { + if (!std::isfinite(emitted_tokens) || emitted_tokens < 1.0 || + emitted_tokens > static_cast(max_accept_) || + generated_tokens < 0) { + return; + } + RequestState & state = states_[request_id]; + ++state.rounds; + state.mean_yield += + (emitted_tokens - state.mean_yield) / state.rounds; + state.tokens_at_last_spec = generated_tokens; + ++global_rounds_; + prior_yield_ += (emitted_tokens - prior_yield_) / global_rounds_; + } + + void forget(uint64_t request_id) { states_.erase(request_id); } + bool has_state(uint64_t request_id) const { + return states_.find(request_id) != states_.end(); + } + int rounds(uint64_t request_id) const { + auto it = states_.find(request_id); + return it == states_.end() ? 0 : it->second.rounds; + } + double mean_yield(uint64_t request_id) const { + auto it = states_.find(request_id); + return it == states_.end() ? 0.0 : it->second.mean_yield; + } + double prior_yield() const { return prior_yield_; } + const SpecCostTables & costs() const { return costs_; } + +private: + struct RequestState { + double mean_yield = 0.0; + int rounds = 0; + int tokens_at_last_spec = 0; + }; + + double score_candidate(const SpecCandidate & candidate) { + if (std::isfinite(candidate.confidence_yield)) { + return std::clamp(candidate.confidence_yield, 1.0, + static_cast(max_accept_)); + } + auto it = states_.find(candidate.request_id); + if (it != states_.end() && it->second.rounds > 0) { + const int idle_tokens = + candidate.generated_tokens - it->second.tokens_at_last_spec; + if (idle_tokens >= config_.stale_after_tokens) { + it->second = RequestState{}; + } else { + return std::clamp(it->second.mean_yield, 1.0, + static_cast(max_accept_)); + } + } + return std::clamp(prior_yield_, 1.0, + static_cast(max_accept_)); + } + + void report_clamp(const char * name, const SpecCostLookup & lookup, + SpecPlan & plan) const { + if (!lookup.clamped) return; + plan.cost_lookup_clamped = true; + if (clamp_logger_) + clamp_logger_(name, lookup.requested_index, lookup.profiled_index); + } + + SpecGateConfig config_; + SpecCostTables costs_; + SpecStepGeometry geometry_; + int max_accept_ = 1; + double prior_yield_ = 1.0; + uint64_t global_rounds_ = 0; + std::unordered_map states_; + ClampLogger clamp_logger_; +}; + +} // namespace dflash::common diff --git a/server/test/test_spec_cost_profile.cpp b/server/test/test_spec_cost_profile.cpp new file mode 100644 index 000000000..6a2e4e87d --- /dev/null +++ b/server/test/test_spec_cost_profile.cpp @@ -0,0 +1,80 @@ +#include "common/concurrency/spec_cost_profile.h" +#include "host_check.h" + +#include +#include +#include +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +int main() { + const SpecProfileGrid grid = build_spec_profile_grid( + 3, 16, 4, [](int lanes) { return lanes; }); + CHECK((grid.tree_rows == std::vector{16, 32, 48})); + CHECK((grid.draft_lanes == std::vector{1, 2, 3})); + CHECK(std::binary_search(grid.step_rows.begin(), grid.step_rows.end(), 1)); + CHECK(std::binary_search(grid.step_rows.begin(), grid.step_rows.end(), 12)); + CHECK(std::is_sorted(grid.step_rows.begin(), grid.step_rows.end())); + CHECK(std::adjacent_find(grid.step_rows.begin(), grid.step_rows.end()) == + grid.step_rows.end()); + + const SpecProfileGrid bucketed = build_spec_profile_grid( + 5, 7, 7, [](int lanes) { + if (lanes <= 1) return 1; + if (lanes <= 2) return 2; + if (lanes <= 4) return 4; + return 6; + }); + CHECK((bucketed.tree_rows == std::vector{7, 14, 28, 42})); + CHECK(bucketed.step_rows.back() == 35); + CHECK(build_spec_profile_grid(0, 16, 4, {}).tree_rows.empty()); + + std::unordered_map calls; + const std::unordered_map intended{{1, 10.0}, {2, 5.0}}; + const double noise[] = {-2.0, 1.0, 0.0, 2.0, -1.0}; + SpecProfileResult profiled = profile_monotonic_costs( + {2, 1, 2}, [&](int index) { + const int call = calls[index]++; + if (call == 0) return 10000.0; // discarded warmup + return intended.at(index) + noise[(call - 1) % 5]; + }); + CHECK(profiled.ok()); + CHECK((profiled.table.indices == std::vector{1, 2})); + CHECK(profiled.table.costs.size() == 2); + CHECK(profiled.table.costs[0] == 10.0); + CHECK(profiled.table.costs[1] == 10.0); // monotone clamp + CHECK(calls[1] == 6 && calls[2] == 6); + + int even_calls = 0; + profiled = profile_monotonic_costs({4}, [&](int) { + ++even_calls; + static const double samples[] = {999.0, 1.0, 2.0, 3.0, 4.0}; + return samples[(even_calls - 1) % 5]; + }, 4); + CHECK(profiled.ok()); + CHECK(profiled.table.costs[0] == 2.5); + CHECK(even_calls == 5); + + SpecProfileResult bad = profile_monotonic_costs( + {}, [](int) { return 1.0; }); + CHECK(!bad.ok() && !bad.error.empty()); + bad = profile_monotonic_costs( + {1}, [](int) { return 1.0; }, 0); + CHECK(!bad.ok()); + int invalid_calls = 0; + bad = profile_monotonic_costs({1}, [&](int) { + ++invalid_calls; + return invalid_calls == 2 + ? std::numeric_limits::quiet_NaN() : 1.0; + }); + CHECK(!bad.ok()); + CHECK(bad.table.indices.empty()); + + std::printf("spec cost profile tests passed: %d checks\n", g_checks); + return 0; +} diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp new file mode 100644 index 000000000..1276ba208 --- /dev/null +++ b/server/test/test_speculation_gate.cpp @@ -0,0 +1,192 @@ +#include "common/concurrency/speculation_gate.h" +#include "host_check.h" + +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +static SpecCostSeries series(int max_index, double cost) { + SpecCostSeries out; + for (int i = 1; i <= max_index; ++i) { + out.indices.push_back(i); + out.costs.push_back(cost); + } + return out; +} + +static SpecCostTables constant_costs(double tree, double step, double draft) { + return {series(128, tree), series(128, step), series(16, draft)}; +} + +static SpecStepGeometry geometry() { + SpecStepGeometry out; + out.tree_width = 4; + out.bucket = [](int lanes) { return lanes; }; + return out; +} + +static SpecCandidate candidate( + uint64_t id, int slot, double confidence, + SpeculationPolicy policy = SpeculationPolicy::Adaptive, + bool eligible = true, int generated = 0) { + return {id, slot, policy, eligible, generated, confidence}; +} + +int main() { + CHECK(std::abs(confidence_survival_yield({0.5f, 0.5f}, 4) - 1.75) < 1e-9); + CHECK(confidence_survival_yield({2.0f, -1.0f}, 4) == 2.0); + CHECK(confidence_survival_yield({}, 4) == 1.0); + + SpeculationGate costly({}, constant_costs(100.0, 10.0, 100.0), + geometry(), 4); + CHECK(costly.valid()); + SpecPlan plan = costly.plan(2, { + candidate(1, 0, 4.0), candidate(2, 1, 4.0)}, 2); + CHECK(plan.valid); + CHECK(plan.admitted_count == 0); + CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); + + // A high-yield request pays for one speculative lane. Adding the + // confidence-1 freeloader cannot improve the numerator and loses the + // tie because the argmax is scanned from the smaller prefix. + SpeculationGate prefix({}, constant_costs(1.0, 10.0, 1.0), + geometry(), 4); + plan = prefix.plan(3, { + candidate(10, 0, 4.0), candidate(11, 1, 1.0), + candidate(12, 2, 4.0, SpeculationPolicy::Never)}, 3); + CHECK(plan.admitted_count == 1); + CHECK((plan.admitted_request_ids == std::vector{10})); + CHECK(plan.ordered.size() == 2); + CHECK(plan.ordered[0].admitted); + CHECK(!plan.ordered[1].admitted); + + // Identical optimistic cold priors are worth trying at C=1 but not at + // C=8 when the profiled tree launch crosses an occupancy boundary. + SpecCostTables crossover = constant_costs(1.0, 4.0, 1.0); + for (size_t i = 0; i < crossover.tree_cost.indices.size(); ++i) { + if (crossover.tree_cost.indices[i] > 4) + crossover.tree_cost.costs[i] = 20.0; + } + SpeculationGate cold_low({}, crossover, geometry(), 4); + plan = cold_low.plan(1, {candidate(20, 0, NAN)}, 1); + CHECK(plan.admitted_count == 1); + SpeculationGate cold_high({}, crossover, geometry(), 4); + std::vector eight; + for (int i = 0; i < 8; ++i) eight.push_back(candidate(30 + i, i, NAN)); + plan = cold_high.plan(8, eight, 8); + CHECK(plan.admitted_count == 0); + + // Always and Never partition before adaptive ordering. + plan = costly.plan(3, { + candidate(1, 0, 1.0, SpeculationPolicy::Never), + candidate(2, 1, 1.0, SpeculationPolicy::Always), + candidate(3, 2, 4.0)}, 2); + CHECK(plan.valid); + CHECK(plan.admitted_count >= 1); + CHECK(plan.admitted_request_ids.front() == 2); + CHECK(plan.ordered.front().forced); + plan = costly.plan(2, { + candidate(1, 0, 4.0, SpeculationPolicy::Always), + candidate(2, 1, 4.0, SpeculationPolicy::Always)}, 1); + CHECK(!plan.valid); + CHECK(!plan.error.empty()); + + SpeculationGate stateful({64}, constant_costs(1.0, 10.0, 1.0), + geometry(), 4); + stateful.observe(100, 2.0, 5); + stateful.observe(100, 4.0, 9); + CHECK(stateful.rounds(100) == 2); + CHECK(std::abs(stateful.mean_yield(100) - 3.0) < 1e-12); + stateful.observe(100, 0.0, 10); + stateful.observe(100, 5.0, 10); + CHECK(stateful.rounds(100) == 2); + + // A stale non-speculating request resets from its low request-local mean + // to the now-higher deployment prior and re-enters through the argmax. + SpeculationGate stale({64}, constant_costs(1.0, 10.0, 1.0), + geometry(), 4); + stale.observe(200, 1.0, 5); + stale.observe(999, 4.0, 5); + CHECK(std::abs(stale.prior_yield() - 2.5) < 1e-12); + plan = stale.plan(1, {candidate(200, 0, NAN, + SpeculationPolicy::Adaptive, true, 6)}, 1); + CHECK(plan.ordered[0].expected_yield == 1.0); + plan = stale.plan(1, {candidate(200, 0, NAN, + SpeculationPolicy::Adaptive, true, 69)}, 1); + CHECK(std::abs(plan.ordered[0].expected_yield - 2.5) < 1e-12); + CHECK(stale.rounds(200) == 0); + + // State follows request IDs, survives a temporary eligibility loss, and + // is explicitly forgotten at retirement/slot reuse. + stale.observe(300, 3.0, 4); + plan = stale.plan(1, {candidate(300, 0, NAN, + SpeculationPolicy::Adaptive, false, 5)}, 1); + CHECK(plan.ordered.empty()); + CHECK(stale.has_state(300)); + stale.forget(300); + CHECK(!stale.has_state(300)); + plan = stale.plan(1, {candidate(301, 0, NAN)}, 1); + CHECK(plan.ordered[0].request_id == 301); + + // All-Never is a pure AR plan and does not allocate request state. + SpeculationGate never({}, constant_costs(1.0, 2.0, 1.0), geometry(), 4); + plan = never.plan(1, {candidate(400, 0, NAN, + SpeculationPolicy::Never)}, 1); + CHECK(plan.admitted_count == 0); + CHECK(plan.ordered.empty()); + CHECK(!never.has_state(400)); + + // C=1, capacity zero, malformed shapes, and always-draft pricing. + plan = prefix.plan(1, {candidate(500, 0, 4.0)}, 1); + CHECK(plan.admitted_count == 1); + plan = prefix.plan(1, {candidate(500, 0, 4.0)}, 0); + CHECK(plan.admitted_count == 0); + plan = prefix.plan(2, {candidate(500, 0, 4.0)}, 1); + CHECK(!plan.valid); + plan = prefix.plan(1, {candidate(500, 0, 4.0)}, 1, 4); + CHECK(plan.draft_lanes == 4); + + SpecCostSeries sparse{{2, 4}, {1.0, 2.0}}; + CHECK(sparse.valid()); + SpecCostLookup lookup = sparse.lookup(1); + CHECK(lookup.clamped && lookup.profiled_index == 2); + lookup = sparse.lookup(3); + CHECK(!lookup.clamped && lookup.rounded_up && lookup.profiled_index == 4); + lookup = sparse.lookup(9); + CHECK(lookup.clamped && lookup.profiled_index == 4); + SpecCostSeries invalid{{2, 1}, {1.0, 2.0}}; + CHECK(!invalid.valid()); + + int clamp_logs = 0; + SpecCostTables tiny{series(1, 1.0), series(1, 2.0), series(1, 1.0)}; + SpeculationGate clamped({}, tiny, geometry(), 4, + [&](const char *, int, int) { ++clamp_logs; }); + plan = clamped.plan(2, { + candidate(1, 0, 4.0), candidate(2, 1, 4.0)}, 2); + CHECK(plan.cost_lookup_clamped); + CHECK(clamp_logs > 0); + + // M1 convergence endpoints. + SpeculationGate pays({}, constant_costs(1.0, 10.0, 1.0), geometry(), 4); + for (int step = 0; step < 3; ++step) { + plan = pays.plan(2, { + candidate(600, 0, NAN, SpeculationPolicy::Adaptive, true, step), + candidate(601, 1, NAN, SpeculationPolicy::Adaptive, true, step)}, 2); + CHECK(plan.admitted_count == 2); + pays.observe(600, 4.0, step + 1); + pays.observe(601, 4.0, step + 1); + } + SpeculationGate cannot({}, constant_costs(100.0, 10.0, 100.0), + geometry(), 4); + for (int step = 0; step < 3; ++step) { + plan = cannot.plan(8, eight, 8); + CHECK(plan.admitted_count == 0); + } + + std::printf("speculation gate tests passed: %d checks\n", g_checks); + return 0; +} From df0195a3046fe0b3a1b55847181d5e33d01495b5 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 18:03:36 +0000 Subject: [PATCH 13/42] feat(qwen35): wire adaptive DSpark speculation --- .../qwen35/concurrency/qwen35_seq_engine.cpp | 510 +++++++++++++++++- .../qwen35/concurrency/qwen35_seq_engine.h | 9 + server/src/qwen35/qwen35_backend.cpp | 13 + 3 files changed, 526 insertions(+), 6 deletions(-) diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 3edd65d49..1a83b7878 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -16,9 +16,12 @@ #include "common/geometric_draft_topk_cuda.h" #include "common/dspark_head.h" #include "common/concurrency/chain_spec_shapes.h" +#include "common/concurrency/spec_cost_profile.h" #include "internal.h" #include +#include +#include #include #include #include @@ -59,6 +62,9 @@ Qwen35SeqEngine::Qwen35SeqEngine( tree_scratch_stride_(tree_scratch_stride), spec_mode_(spec_mode) { const int n_slots = slots_.slot_count(); slot_draft_kv_.resize((size_t)n_slots); + last_survival_score_.assign( + (size_t)n_slots, std::numeric_limits::quiet_NaN()); + last_survival_generated_.assign((size_t)n_slots, -1); // The concurrent DDTree stack is gated to a local same-device drafter. // Build metadata-only BF16 views over each slot's disjoint target feature @@ -113,6 +119,387 @@ Qwen35SeqEngine::~Qwen35SeqEngine() { } } +bool Qwen35SeqEngine::spec_gate_debug_enabled() const { + const char * value = std::getenv("DFLASH_SPEC_GATE_LOG"); + return value && std::atoi(value) != 0; +} + +bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { + speculation_gate_.reset(); + if (spec_mode_ != SpecMode::dspark_chain || !capture_features_ || + tree_width_ <= 1 || tree_width_ > 16 || slots_.residency_active()) { + std::fprintf(stderr, + "[spec-profile] disabled: chain/features unavailable or " + "concurrent KVFlash residency active; adaptive requests use AR\n"); + return false; + } + + const int n_slots = slots_.slot_count(); + const int T = tree_width_; + const int hidden = b_.w_.n_embd; + const int n_head_kv = b_.w_.n_head_kv; + const int max_profile_ctx = std::min( + slots_.max_context() - T, b_.cache_.target_feat_cap); + if (n_slots < 1 || max_profile_ctx < 1) return false; + const int ctx_tokens = std::clamp(context_tokens, 1, max_profile_ctx); + + std::string profile_error; + std::vector synthetic_slots; + synthetic_slots.reserve((size_t)n_slots); + auto cleanup = [&]() { + for (int slot : synthetic_slots) { + if (slot >= 0 && slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)slot]); + } + if (slots_.is_active(slot)) slots_.retire(slot); + if (slot >= 0 && slot < (int)last_survival_score_.size()) { + last_survival_score_[(size_t)slot] = + std::numeric_limits::quiet_NaN(); + last_survival_generated_[(size_t)slot] = -1; + } + } + }; + + // Fabricate page tables and steady-state sequence lengths without a + // target prefill. Zero K/V and captured target features once, then every + // timed graph sees a deterministic context at the requested length. + const int32_t profile_token = b_.w_.mask_token_id >= 0 + ? b_.w_.mask_token_id : 0; + std::vector prompt((size_t)ctx_tokens, profile_token); + const SamplerCfg greedy{}; + for (int lane = 0; lane < n_slots; ++lane) { + AdmitResult admitted = admit( + std::numeric_limits::max() - (uint64_t)lane, + prompt, greedy); + if (admitted.status != AdmitResult::Status::admitted) { + profile_error = admitted.error.empty() + ? "synthetic slot admission failed" : admitted.error; + cleanup(); + std::fprintf(stderr, "[spec-profile] %s\n", profile_error.c_str()); + return false; + } + synthetic_slots.push_back(admitted.slot); + Qwen35SlotManager::PrefillChunk chunk = + slots_.append_prefill(admitted.slot, ctx_tokens); + if (!chunk.ok || + !upload_block_table_delta(admitted.slot, chunk.first_new_block, + chunk.new_blocks.data(), + chunk.new_blocks.size())) { + profile_error = "synthetic paged context allocation failed"; + cleanup(); + std::fprintf(stderr, "[spec-profile] %s\n", profile_error.c_str()); + return false; + } + slots_.commit_prefill(admitted.slot); + } + for (ggml_tensor * tensor : b_.cache_.attn_k) { + if (tensor) ggml_backend_tensor_memset(tensor, 0, 0, ggml_nbytes(tensor)); + } + for (ggml_tensor * tensor : b_.cache_.attn_v) { + if (tensor) ggml_backend_tensor_memset(tensor, 0, 0, ggml_nbytes(tensor)); + } + if (b_.cache_.target_feat) { + ggml_backend_tensor_memset( + b_.cache_.target_feat, 0, 0, ggml_nbytes(b_.cache_.target_feat)); + } + seq_lens_.assign((size_t)n_slots, ctx_tokens); + ggml_backend_tensor_set( + b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + ggml_backend_synchronize(b_.target_backend_); + + const SpecProfileGrid grid = build_spec_profile_grid( + n_slots, T, T, [](int lanes) { + return chain_decode_bucket_width(lanes); + }); + + int prepared_tree_rows = -1; + auto tree_runner = [&](int total_rows) -> double { + if (!profile_error.empty()) + return std::numeric_limits::infinity(); + if (prepared_tree_rows != total_rows) { + if (total_rows <= 0 || total_rows % T != 0) { + profile_error = "invalid tree profiling shape"; + return std::numeric_limits::infinity(); + } + const int bucket = total_rows / T; + const int live = std::min(bucket, n_slots); + StepGraph & sg = b_.sg_; + if (!build_target_step_paged_tree( + sg, b_.w_, b_.cache_, b_.target_backend_, + T, bucket, ctx_tokens, + tree_scratch_base_, tree_scratch_stride_, + b_.cfg_.kq_stride_pad)) { + profile_error = "tree profiling graph build failed"; + return std::numeric_limits::infinity(); + } + + std::vector tokens((size_t)total_rows, profile_token); + std::vector embeddings((size_t)hidden * total_rows); + std::vector parents((size_t)total_rows, -1); + std::vector sizes((size_t)bucket, 0); + std::vector active((size_t)bucket, -1); + std::vector state((size_t)bucket, 0); + std::vector queries((size_t)total_rows, -1); + std::vector positions((size_t)4 * total_rows, 0); + std::vector rows( + (size_t)n_head_kv * total_rows, scratch_row_); + if (!b_.w_.embedder.embed( + tokens.data(), total_rows, embeddings.data())) { + profile_error = "tree profiling embedding failed"; + return std::numeric_limits::infinity(); + } + for (int lane = 0; lane < live; ++lane) { + const int slot = synthetic_slots[(size_t)lane]; + sizes[(size_t)lane] = T; + active[(size_t)lane] = slot; + state[(size_t)lane] = slot; + for (int node = 0; node < T; ++node) { + const int row = lane * T + node; + parents[(size_t)row] = node == 0 ? -1 : node - 1; + queries[(size_t)row] = slot; + for (int axis = 0; axis < 3; ++axis) { + positions[(size_t)axis * total_rows + row] = + ctx_tokens + node; + } + for (int head = 0; head < n_head_kv; ++head) { + rows[(size_t)head * total_rows + row] = + (int64_t)tree_scratch_base_ + + (int64_t)slot * tree_scratch_stride_ + node; + } + } + } + ggml_backend_tensor_set(sg.inp_embed, embeddings.data(), 0, + sizeof(float) * embeddings.size()); + ggml_backend_tensor_set(sg.positions, positions.data(), 0, + sizeof(int32_t) * positions.size()); + ggml_backend_tensor_set(sg.parent_ids, parents.data(), 0, + sizeof(int32_t) * parents.size()); + ggml_backend_tensor_set(sg.tree_sizes, sizes.data(), 0, + sizeof(int32_t) * sizes.size()); + if (detail::target_paged_tree_active_slots_need_upload(sg)) { + ggml_backend_tensor_set(sg.active_slot_ids, active.data(), 0, + sizeof(int32_t) * active.size()); + } + ggml_backend_tensor_set(sg.state_slot_ids, state.data(), 0, + sizeof(int32_t) * state.size()); + ggml_backend_tensor_set(sg.paged_query_seq_ids, queries.data(), 0, + sizeof(int32_t) * queries.size()); + ggml_backend_tensor_set(sg.kv_write_rows, rows.data(), 0, + sizeof(int64_t) * rows.size()); + prepared_tree_rows = total_rows; + } + const auto start = std::chrono::steady_clock::now(); + if (ggml_backend_graph_compute(b_.target_backend_, b_.sg_.gf) != + GGML_STATUS_SUCCESS) { + profile_error = "tree profiling compute failed"; + return std::numeric_limits::infinity(); + } + ggml_backend_synchronize(b_.target_backend_); + // The serving path rebuilds this graph for every decode iteration. + // Rebuild between samples too: replaying the same captured target + // graph is not a supported lifecycle on the HIP graph backend. + prepared_tree_rows = -1; + return std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + }; + + int prepared_step_rows = -1; + auto step_runner = [&](int total_rows) -> double { + if (!profile_error.empty()) + return std::numeric_limits::infinity(); + if (prepared_step_rows != total_rows) { + if (total_rows <= 0 || total_rows > n_slots * T) { + profile_error = "invalid durable-step profiling shape"; + return std::numeric_limits::infinity(); + } + std::vector segments; + int offset = 0; + for (int lane = 0; offset < total_rows; ++lane) { + const int length = std::min(T, total_rows - offset); + segments.push_back({ + offset, length, synthetic_slots[(size_t)lane]}); + offset += length; + } + StepGraph & sg = b_.sg_; + if (!build_target_step( + sg, b_.w_, b_.cache_, b_.target_backend_, + 0, total_rows, false, true, false, 0, 0, + b_.cfg_.kq_stride_pad, false, false, false, true, + 1, 0, ctx_tokens + T, + total_rows, segments.data(), (int)segments.size(), + (int)segments.size(), false) || + !sg.kv_write_rows || !sg.target_feat_rows || + !sg.paged_query_seq_ids || !sg.paged_query_positions || + !sg.logits_row_indices) { + profile_error = "durable-step profiling graph build failed"; + return std::numeric_limits::infinity(); + } + + std::vector tokens((size_t)total_rows, profile_token); + std::vector embeddings((size_t)hidden * total_rows); + std::vector positions((size_t)4 * total_rows, 0); + std::vector rows( + (size_t)n_head_kv * total_rows, scratch_row_); + std::vector queries((size_t)total_rows, -1); + std::vector feature_rows( + (size_t)total_rows, + b_.cache_.target_feat_cap * n_slots); + std::vector logits_rows; + logits_rows.reserve(segments.size()); + if (!b_.w_.embedder.embed( + tokens.data(), total_rows, embeddings.data())) { + profile_error = "durable-step profiling embedding failed"; + return std::numeric_limits::infinity(); + } + for (const QwenPrefillSegment & segment : segments) { + for (int j = 0; j < segment.n_tokens; ++j) { + const int row = segment.token_offset + j; + queries[(size_t)row] = segment.seq_slot; + for (int axis = 0; axis < 3; ++axis) { + positions[(size_t)axis * total_rows + row] = + ctx_tokens + j; + } + for (int head = 0; head < n_head_kv; ++head) { + rows[(size_t)head * total_rows + row] = + (int64_t)tree_scratch_base_ + + (int64_t)segment.seq_slot * tree_scratch_stride_ + j; + } + } + logits_rows.push_back(segment.token_offset + segment.n_tokens - 1); + } + ggml_backend_tensor_set(sg.inp_embed, embeddings.data(), 0, + sizeof(float) * embeddings.size()); + ggml_backend_tensor_set(sg.positions, positions.data(), 0, + sizeof(int32_t) * positions.size()); + ggml_backend_tensor_set(sg.kv_write_rows, rows.data(), 0, + sizeof(int64_t) * rows.size()); + ggml_backend_tensor_set(sg.paged_query_seq_ids, queries.data(), 0, + sizeof(int32_t) * queries.size()); + ggml_backend_tensor_set(sg.paged_query_positions, + positions.data(), 0, + sizeof(int32_t) * total_rows); + ggml_backend_tensor_set(sg.target_feat_rows, feature_rows.data(), 0, + sizeof(int32_t) * feature_rows.size()); + ggml_backend_tensor_set(sg.logits_row_indices, logits_rows.data(), 0, + sizeof(int32_t) * logits_rows.size()); + prepared_step_rows = total_rows; + } + const auto start = std::chrono::steady_clock::now(); + if (ggml_backend_graph_compute(b_.target_backend_, b_.sg_.gf) != + GGML_STATUS_SUCCESS) { + profile_error = "durable-step profiling compute failed"; + return std::numeric_limits::infinity(); + } + ggml_backend_synchronize(b_.target_backend_); + prepared_step_rows = -1; + return std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + }; + + std::vector noise((size_t)T, b_.w_.mask_token_id); + noise[0] = profile_token; + std::vector noise_embed((size_t)hidden * T); + std::vector local_hidden((size_t)hidden * T); + std::vector prenorm_hidden((size_t)hidden * T); + if (!b_.w_.embedder.embed(noise.data(), T, noise_embed.data())) { + cleanup(); + return false; + } + auto draft_runner = [&](int lanes) -> double { + if (!profile_error.empty()) + return std::numeric_limits::infinity(); + const auto start = std::chrono::steady_clock::now(); + for (int lane = 0; lane < lanes; ++lane) { + const int slot = synthetic_slots[(size_t)lane]; + DraftKvState * draft = ensure_slot_draft_kv(slot); + DraftFeatureMirror * mirror = slot_feature_mirror(slot); + if (!draft || !mirror || + !draft_kv_begin_step(*draft, b_.dw_, b_.draft_backend_, + *mirror, ctx_tokens)) { + profile_error = "draft profiling setup failed"; + return std::numeric_limits::infinity(); + } + 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) { + profile_error = "draft profiling backbone failed"; + return std::numeric_limits::infinity(); + } + ggml_backend_tensor_get_async( + b_.draft_backend_, draft->hidden_states, + local_hidden.data(), 0, sizeof(float) * local_hidden.size()); + ggml_backend_tensor_get_async( + b_.draft_backend_, draft->hidden_prenorm, + prenorm_hidden.data(), 0, + sizeof(float) * prenorm_hidden.size()); + ggml_backend_synchronize(b_.draft_backend_); + std::vector draft_tokens; + std::vector confidence; + if (!dspark_markov_correct_greedy_chain_fused( + b_.dw_, b_.draft_backend_, b_.w_.output, + local_hidden.data(), T, profile_token, draft_tokens, + &confidence, prenorm_hidden.data()) || + (int)draft_tokens.size() != T) { + profile_error = "draft profiling chain failed"; + return std::numeric_limits::infinity(); + } + } + return std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + }; + + SpecProfileResult tree = profile_monotonic_costs( + grid.tree_rows, tree_runner); + SpecProfileResult step = profile_monotonic_costs( + grid.step_rows, step_runner); + SpecProfileResult draft = profile_monotonic_costs( + grid.draft_lanes, draft_runner); + cleanup(); + if (!tree.ok() || !step.ok() || !draft.ok() || !profile_error.empty()) { + std::fprintf(stderr, "[spec-profile] failed: %s%s%s%s\n", + profile_error.c_str(), tree.error.c_str(), step.error.c_str(), + draft.error.c_str()); + return false; + } + + SpecCostTables tables{ + std::move(tree.table), std::move(step.table), std::move(draft.table)}; + SpecStepGeometry geometry; + geometry.tree_width = T; + geometry.bucket = [](int lanes) { + return chain_decode_bucket_width(lanes); + }; + speculation_gate_ = std::make_unique( + SpecGateConfig{}, tables, geometry, T, + [](const char * table, int requested, int profiled) { + std::fprintf(stderr, + "[spec-gate] %s_cost index %d outside profile; clamped to %d\n", + table, requested, profiled); + }); + if (!speculation_gate_->valid()) { + speculation_gate_.reset(); + return false; + } + + auto print_table = [](const char * name, const SpecCostSeries & series) { + std::fprintf(stderr, "[spec-profile] %s", name); + for (size_t i = 0; i < series.indices.size(); ++i) { + std::fprintf(stderr, "%s%d:%.1fus", + i == 0 ? " " : ",", series.indices[i], series.costs[i]); + } + std::fprintf(stderr, "\n"); + }; + std::fprintf(stderr, + "[spec-profile] context=%d reps=5 mode=serial-draft\n", ctx_tokens); + print_table("tree_cost", tables.tree_cost); + print_table("step_cost", tables.step_cost); + print_table("draft_cost", tables.draft_cost); + return true; +} + DraftFeatureMirror * Qwen35SeqEngine::slot_feature_mirror(int slot) { if (!capture_features_ || slot < 0 || slot >= (int)slot_feature_mirrors_.size()) { @@ -744,6 +1131,20 @@ std::optional Qwen35SeqEngine::step_chain_spec( out.target_forwards = 2; out.committed_tokens.assign( proposal.path.begin() + 1, proposal.path.end()); + const double survival = confidence_survival_yield( + proposal.confidence, tree_width_); + if (proposal.slot >= 0 && + proposal.slot < (int)last_survival_score_.size()) { + last_survival_score_[(size_t)proposal.slot] = survival; + last_survival_generated_[(size_t)proposal.slot] = + slots_.slot(proposal.slot).generated_tokens(); + } + if (speculation_gate_) { + speculation_gate_->observe( + slots_.slot(proposal.slot).request_id, + (double)proposal.path.size(), + slots_.slot(proposal.slot).generated_tokens()); + } } else { ArLane & ar = ar_lanes[static_cast(ar_for_input[i])]; @@ -1192,6 +1593,12 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( slot_draft_kv_[(size_t)result.slot]) { draft_kv_reset(*slot_draft_kv_[(size_t)result.slot]); } + if (result.slot >= 0 && + result.slot < (int)last_survival_score_.size()) { + last_survival_score_[(size_t)result.slot] = + std::numeric_limits::quiet_NaN(); + last_survival_generated_[(size_t)result.slot] = -1; + } } return result; } @@ -1434,16 +1841,100 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (spec_mode_ == SpecMode::dspark_chain && plan.prefills.empty()) { std::vector admitted(inputs.size(), 0); - bool any_admitted = false; - for (size_t i = 0; i < inputs.size(); ++i) { - admitted[i] = chain_spec_input_eligible(inputs[i]) && - inputs[i].speculation_policy == - SpeculationPolicy::Always; - any_admitted = any_admitted || admitted[i] != 0; + SpecPlan gate_plan; + bool have_gate_plan = false; + + const char * force_value = std::getenv("DFLASH_SPEC_GATE_FORCE"); + const std::string force = force_value ? force_value : ""; + const char * draft_always_value = + std::getenv("DFLASH_SPEC_DRAFT_ALWAYS"); + const bool draft_always = + draft_always_value && std::atoi(draft_always_value) != 0; + + if (speculation_gate_) { + std::vector candidates; + candidates.reserve(inputs.size()); + int drafting_lanes = 0; + for (const StepInput & in : inputs) { + const Qwen35Slot & seq = slots_.slot(in.slot); + SpeculationPolicy policy = in.speculation_policy; + if (policy == SpeculationPolicy::Adaptive) { + if (force == "all") policy = SpeculationPolicy::Always; + if (force == "none") policy = SpeculationPolicy::Never; + } + const bool eligible = chain_spec_input_eligible(in); + drafting_lanes += eligible ? 1 : 0; + double confidence = + std::numeric_limits::quiet_NaN(); + if (eligible && in.slot >= 0 && + in.slot < (int)last_survival_score_.size() && + last_survival_generated_[(size_t)in.slot] == + seq.generated_tokens()) { + confidence = last_survival_score_[(size_t)in.slot]; + } + candidates.push_back({ + seq.request_id, in.slot, policy, eligible, + seq.generated_tokens(), confidence, + }); + } + gate_plan = speculation_gate_->plan( + (int)inputs.size(), candidates, (int)inputs.size(), + draft_always ? drafting_lanes : -1); + have_gate_plan = true; + if (!gate_plan.valid) { + return fail_step(gate_plan.error.empty() + ? "adaptive speculation gate failed" : gate_plan.error); + } + for (int slot : gate_plan.admitted_slots) { + for (size_t i = 0; i < inputs.size(); ++i) { + if (inputs[i].slot == slot) admitted[i] = 1; + } + } + } else { + // A missing/failed profile degrades Adaptive to AR. Explicit + // speculation remains a reliable forced-mode oracle. + for (size_t i = 0; i < inputs.size(); ++i) { + SpeculationPolicy policy = inputs[i].speculation_policy; + if (policy == SpeculationPolicy::Adaptive && + force == "all") { + policy = SpeculationPolicy::Always; + } else if (policy == SpeculationPolicy::Adaptive && + force == "none") { + policy = SpeculationPolicy::Never; + } + admitted[i] = chain_spec_input_eligible(inputs[i]) && + policy == SpeculationPolicy::Always; + } } + + const bool any_admitted = std::any_of( + admitted.begin(), admitted.end(), + [](uint8_t value) { return value != 0; }); if (any_admitted) { + const auto started = std::chrono::steady_clock::now(); std::optional speculative = step_chain_spec(plan, admitted); + const double measured_us = + std::chrono::duration( + std::chrono::steady_clock::now() - started).count(); + if (have_gate_plan && spec_gate_debug_enabled()) { + std::fprintf(stderr, + "[spec-gate] C=%zu k=%d scores=[", + inputs.size(), gate_plan.admitted_count); + for (size_t i = 0; i < gate_plan.ordered.size(); ++i) { + const SpecPlanScore & score = gate_plan.ordered[i]; + std::fprintf(stderr, "%s%llu:%.3f%s", + i == 0 ? "" : ",", + (unsigned long long)score.request_id, + score.expected_yield, + score.admitted ? "*" : ""); + } + std::fprintf(stderr, + "] G(k)=%.6f G(0)=%.6f predicted=%.1fus " + "measured=%.1fus\n", + gate_plan.goodput, gate_plan.ar_goodput, + gate_plan.predicted_cost, measured_us); + } if (speculative) return std::move(*speculative); // Proposal setup failed before target/cache mutation. Preserve // service through the ordinary packed AR path this iteration. @@ -1829,6 +2320,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { void Qwen35SeqEngine::retire(int slot) { if (!slots_.is_active(slot)) return; + const uint64_t request_id = slots_.slot(slot).request_id; + if (speculation_gate_) speculation_gate_->forget(request_id); + if (slot >= 0 && slot < (int)last_survival_score_.size()) { + last_survival_score_[(size_t)slot] = + std::numeric_limits::quiet_NaN(); + last_survival_generated_[(size_t)slot] = -1; + } slots_.retire(slot); } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index c0d7f5fcc..127a8757a 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" @@ -29,6 +30,7 @@ #include #include +#include #include #include #include @@ -70,6 +72,9 @@ class Qwen35SeqEngine final : public SeqEngine { const SamplerCfg & sampler) override; StepResult step(const StepPlan & plan) override; + // Fabricate a steady-state paged context and profile the three launch + // families used by the DSpark gate. Called once from backend init. + bool profile_spec_costs(int context_tokens); StepPlanLimits step_plan_limits(int decode_rows) const override { const bool mixed = decode_rows > 0; const int per_sequence = mixed ? 512 : 2048; @@ -128,6 +133,7 @@ class Qwen35SeqEngine final : public SeqEngine { DraftKvState * ensure_slot_draft_kv(int slot); bool ddtree_eligible(const StepPlan & plan) const; bool chain_spec_input_eligible(const StepInput & input) const; + bool spec_gate_debug_enabled() 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); @@ -146,6 +152,9 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector slot_feature_mirrors_; std::vector> slot_draft_kv_; + std::unique_ptr speculation_gate_; + std::vector last_survival_score_; + std::vector last_survival_generated_; // Hoisted per-step buffers (reused across step() calls). std::vector output_rows_; std::vector live_tokens_; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index eb6d160da..c195cdb40 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -596,6 +596,19 @@ bool Qwen35Backend::init() { max_concurrent_prefills, mixed_prefill_tokens, long_mixed_prefill_tokens, long_prefill_threshold, idle_prefill_tokens, prefill_quantum); + if (concurrent_local_chain && + cfg_.speculation_policy == SpeculationPolicy::Adaptive) { + int profile_ctx = 4096; + if (const char * value = + std::getenv("DFLASH_SPEC_PROFILE_CONTEXT")) { + profile_ctx = std::max(1, std::atoi(value)); + } + if (!seq_engine_->profile_spec_costs(profile_ctx)) { + std::fprintf(stderr, + "[parallel-dspark] adaptive profile unavailable; " + "Adaptive requests degrade to packed AR\n"); + } + } if (concurrent_local_ddtree) { const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); std::fprintf(stderr, From 4bf139d48ca3646f8372e2d812b0f7ad79159fd9 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 18:36:34 +0000 Subject: [PATCH 14/42] feat(draft): batch concurrent DSpark drafting --- server/CMakeLists.txt | 11 + server/src/common/dflash_draft_kv.cpp | 209 ++++++++ server/src/common/dflash_draft_kv.h | 36 ++ server/src/common/dspark_head.cpp | 114 +++++ server/src/common/dspark_head.h | 25 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 445 ++++++++++++++---- .../qwen35/concurrency/qwen35_seq_engine.h | 15 + server/test/test_chain_spec_shapes.cpp | 10 + server/test/test_dspark_batched_head.cpp | 224 +++++++++ 9 files changed, 1008 insertions(+), 81 deletions(-) create mode 100644 server/test/test_dspark_batched_head.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 28c7e5f55..d9ba62836 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1452,6 +1452,17 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/test) list(APPEND _raw_unit_test_targets test_chain_spec_shapes) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dspark_batched_head.cpp") + add_executable(test_dspark_batched_head + test/test_dspark_batched_head.cpp) + target_include_directories(test_dspark_batched_head PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + target_link_libraries(test_dspark_batched_head PRIVATE + dflash_common ggml ggml-cpu + ${DFLASH27B_GGML_BACKEND_TARGET} ggml-base) + list(APPEND _raw_unit_test_targets test_dspark_batched_head) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_speculation_gate.cpp") add_executable(test_speculation_gate test/test_speculation_gate.cpp) diff --git a/server/src/common/dflash_draft_kv.cpp b/server/src/common/dflash_draft_kv.cpp index 1442a495a..9077da6a8 100644 --- a/server/src/common/dflash_draft_kv.cpp +++ b/server/src/common/dflash_draft_kv.cpp @@ -1,4 +1,5 @@ #include "dflash_draft_kv.h" +#include "dspark_head.h" #include #include @@ -321,4 +322,212 @@ bool draft_kv_begin_step(DraftKvState & st, return true; } +void draft_kv_batch_free(DraftKvBatchGraph & batch) { + if (batch.galloc) { + ggml_gallocr_free(batch.galloc); + batch.galloc = nullptr; + } + if (batch.g_ctx) { + ggml_free(batch.g_ctx); + batch.g_ctx = nullptr; + } + batch.gf = nullptr; + batch.seed_tokens = nullptr; + batch.token_depths.clear(); + batch.confidence_depths.clear(); + batch.lane_states.clear(); + batch.meta_arena.clear(); + batch.n_lanes = 0; + batch.q_len = 0; + batch.has_confidence = false; + batch.built_for = nullptr; + batch.built_lm_head = nullptr; +} + +static bool draft_kv_batch_build( + DraftKvBatchGraph & batch, + const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const std::vector & lane_states) { + if (!backend || !lm_head || lane_states.empty() || + dw.block_size <= 1) { + return false; + } + for (DraftKvState * state : lane_states) { + if (!state || !state->mem_buf || state->q_len != dw.block_size || + state->built_for != (const void *)&dw) { + return false; + } + } + + draft_kv_batch_free(batch); + const int n_lanes = static_cast(lane_states.size()); + const size_t arena_size = + (32u + 16u * (size_t)n_lanes) * 1024u * 1024u; + batch.meta_arena.resize(arena_size); + ggml_init_params gp{}; + gp.mem_size = batch.meta_arena.size(); + gp.mem_buffer = batch.meta_arena.data(); + gp.no_alloc = true; + batch.g_ctx = ggml_init(gp); + if (!batch.g_ctx) { + draft_kv_batch_free(batch); + return false; + } + batch.gf = ggml_new_graph_custom( + batch.g_ctx, 4096 * n_lanes + 2048, false); + batch.seed_tokens = + ggml_new_tensor_1d(batch.g_ctx, GGML_TYPE_I32, n_lanes); + ggml_set_input(batch.seed_tokens); + + std::vector hidden; + std::vector prenorm; + hidden.reserve((size_t)n_lanes); + prenorm.reserve((size_t)n_lanes); + for (DraftKvState * state : lane_states) { + DraftKvAppendInputs append{}; + append.n_rows = state->a_step; + append.feat = state->ap_feat; + append.positions = state->ap_pos; + append.rows = state->ap_rows; + if (!build_draft_kv_append( + batch.g_ctx, batch.gf, dw, state->cache, append)) { + draft_kv_batch_free(batch); + return false; + } + + DraftKvStepInputs step{}; + step.noise_embed = state->inp_embed; + step.positions_q = state->pos_q; + step.noise_rows = state->noise_rows; + step.mask_full = state->mask_full; + step.mask_swa = state->mask_swa; + DraftGraphOutputs outputs = build_draft_kv_step( + batch.g_ctx, batch.gf, dw, state->cache, step); + if (!outputs.hidden_prenorm || !outputs.hidden_states) { + draft_kv_batch_free(batch); + return false; + } + hidden.push_back(outputs.hidden_states); + prenorm.push_back(outputs.hidden_prenorm); + } + + DSparkBatchedChainOutputs chain; + if (!build_dspark_markov_batched_chain( + batch.g_ctx, batch.gf, dw, lm_head, hidden, prenorm, + batch.seed_tokens, dw.block_size, true, chain) || + chain.n_lanes != n_lanes || + static_cast(chain.tokens.size()) != dw.block_size - 1) { + draft_kv_batch_free(batch); + return false; + } + + batch.galloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!batch.galloc || + !ggml_gallocr_alloc_graph(batch.galloc, batch.gf)) { + std::fprintf(stderr, "[draft-kv-batch] graph alloc failed lanes=%d\n", + n_lanes); + draft_kv_batch_free(batch); + return false; + } + + batch.n_lanes = n_lanes; + batch.q_len = dw.block_size; + batch.has_confidence = + !chain.confidence.empty() && chain.confidence[0] != nullptr; + batch.built_for = &dw; + batch.built_lm_head = lm_head; + batch.lane_states = lane_states; + batch.token_depths = std::move(chain.tokens); + batch.confidence_depths = std::move(chain.confidence); + std::fprintf(stderr, + "[draft-kv-batch] packed graph ready lanes=%d q_len=%d " + "confidence=%s\n", + n_lanes, dw.block_size, + batch.has_confidence ? "on" : "off"); + return true; +} + +bool draft_kv_batch_compute( + DraftKvBatchGraph & batch, + const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const std::vector & lane_states, + const std::vector & seed_tokens, + std::vector> & draft_tokens, + std::vector> & confidences) { + draft_tokens.clear(); + confidences.clear(); + if (lane_states.empty() || + seed_tokens.size() != lane_states.size()) { + return false; + } + const bool reusable = + batch.gf && batch.built_for == (const void *)&dw && + batch.built_lm_head == lm_head && + batch.lane_states == lane_states; + if (!reusable && + !draft_kv_batch_build( + batch, dw, backend, lm_head, lane_states)) { + return false; + } + + ggml_backend_tensor_set( + batch.seed_tokens, seed_tokens.data(), 0, + sizeof(int32_t) * seed_tokens.size()); + if (ggml_backend_graph_compute(backend, batch.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[draft-kv-batch] graph compute failed lanes=%d\n", + batch.n_lanes); + return false; + } + + const int depths = batch.q_len - 1; + std::vector depth_tokens( + (size_t)depths * batch.n_lanes); + std::vector depth_confidence( + (size_t)depths * batch.n_lanes); + for (int depth = 0; depth < depths; ++depth) { + ggml_backend_tensor_get_async( + backend, batch.token_depths[(size_t)depth], + depth_tokens.data() + (size_t)depth * batch.n_lanes, + 0, sizeof(int32_t) * (size_t)batch.n_lanes); + if (batch.has_confidence && + batch.confidence_depths[(size_t)depth]) { + ggml_backend_tensor_get_async( + backend, batch.confidence_depths[(size_t)depth], + depth_confidence.data() + + (size_t)depth * batch.n_lanes, + 0, sizeof(float) * (size_t)batch.n_lanes); + } + } + ggml_backend_synchronize(backend); + + draft_tokens.assign( + (size_t)batch.n_lanes, + std::vector((size_t)batch.q_len)); + confidences.assign((size_t)batch.n_lanes, {}); + for (int lane = 0; lane < batch.n_lanes; ++lane) { + draft_tokens[(size_t)lane][0] = + seed_tokens[(size_t)lane]; + if (batch.has_confidence) { + confidences[(size_t)lane].resize((size_t)depths); + } + for (int depth = 0; depth < depths; ++depth) { + draft_tokens[(size_t)lane][(size_t)depth + 1] = + depth_tokens[(size_t)depth * batch.n_lanes + lane]; + if (batch.has_confidence) { + confidences[(size_t)lane][(size_t)depth] = + depth_confidence[ + (size_t)depth * batch.n_lanes + lane]; + } + } + } + return true; +} + } // namespace dflash::common diff --git a/server/src/common/dflash_draft_kv.h b/server/src/common/dflash_draft_kv.h index 9f80044da..04842b6e6 100644 --- a/server/src/common/dflash_draft_kv.h +++ b/server/src/common/dflash_draft_kv.h @@ -100,4 +100,40 @@ bool draft_kv_begin_step(DraftKvState & st, const DraftFeatureMirror & ring, int committed); +struct DraftKvBatchGraph { + DraftKvBatchGraph() = default; + DraftKvBatchGraph(const DraftKvBatchGraph &) = delete; + DraftKvBatchGraph & operator=(const DraftKvBatchGraph &) = delete; + + int n_lanes = 0; + int q_len = 0; + bool has_confidence = false; + const void * built_for = nullptr; + ggml_tensor * built_lm_head = nullptr; + std::vector lane_states; + + std::vector meta_arena; + ggml_context * g_ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_tensor * seed_tokens = nullptr; + std::vector token_depths; + std::vector confidence_depths; +}; + +void draft_kv_batch_free(DraftKvBatchGraph & batch); + +// All lane states must already have draft_kv_begin_step() inputs and +// inp_embed uploaded. The graph is rebuilt only when the ordered state cohort +// changes, then replays as one backend compute and one synchronization. +bool draft_kv_batch_compute( + DraftKvBatchGraph & batch, + const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const std::vector & lane_states, + const std::vector & seed_tokens, + std::vector> & draft_tokens, + std::vector> & confidences); + } // namespace dflash::common diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp index f0df52c10..4146ceacc 100644 --- a/server/src/common/dspark_head.cpp +++ b/server/src/common/dspark_head.cpp @@ -307,6 +307,120 @@ bool build_markov_chain_graph(const DraftWeights & dw, } // namespace +bool build_dspark_markov_batched_chain( + ggml_context * ctx, + ggml_cgraph * gf, + const DraftWeights & dw, + ggml_tensor * lm_head, + const std::vector & hidden_by_lane, + const std::vector & prenorm_by_lane, + ggml_tensor * seed_tokens, + int q_len, + bool want_confidence, + DSparkBatchedChainOutputs & out) { + out = {}; + const int n_lanes = static_cast(hidden_by_lane.size()); + const int hdim = dw.n_embd; + if (!ctx || !gf || !lm_head || !seed_tokens || n_lanes <= 0 || + q_len <= 1 || hdim <= 0 || !dw.dspark.enabled || + !dw.dspark.markov_w1 || !dw.dspark.markov_w2 || + seed_tokens->ne[0] != n_lanes || + prenorm_by_lane.size() != hidden_by_lane.size()) { + return false; + } + const int vocab = static_cast(lm_head->ne[1]); + if (vocab <= 0 || + (dw.dspark.vocab_size > 0 && vocab != dw.dspark.vocab_size)) { + return false; + } + for (int lane = 0; lane < n_lanes; ++lane) { + ggml_tensor * hidden = hidden_by_lane[(size_t)lane]; + ggml_tensor * prenorm = prenorm_by_lane[(size_t)lane]; + if (!hidden || hidden->ne[0] != hdim || hidden->ne[1] < q_len || + !prenorm || prenorm->ne[0] != hdim || prenorm->ne[1] < q_len) { + return false; + } + } + + const bool have_confidence = want_confidence && + dw.dspark.confidence_w && dw.dspark.confidence_b && + (dw.dspark.confidence_dim == hdim || + dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank); + const int n_depths = q_len - 1; + std::vector hidden_by_depth((size_t)n_depths); + std::vector confidence_by_depth((size_t)n_depths); + + auto concat_lane_column = [&](const std::vector & sources, + int column) -> ggml_tensor * { + ggml_tensor * packed = nullptr; + for (ggml_tensor * source : sources) { + ggml_tensor * lane = ggml_view_2d( + ctx, source, hdim, 1, source->nb[1], + (size_t)column * source->nb[1]); + packed = packed ? ggml_concat(ctx, packed, lane, 1) : lane; + } + return packed; + }; + + ggml_tensor * all_hidden = nullptr; + for (int depth = 0; depth < n_depths; ++depth) { + hidden_by_depth[(size_t)depth] = + concat_lane_column(hidden_by_lane, depth + 1); + confidence_by_depth[(size_t)depth] = + concat_lane_column(prenorm_by_lane, depth + 1); + if (!hidden_by_depth[(size_t)depth] || + !confidence_by_depth[(size_t)depth]) { + return false; + } + all_hidden = all_hidden + ? ggml_concat(ctx, all_hidden, hidden_by_depth[(size_t)depth], 1) + : hidden_by_depth[(size_t)depth]; + } + + // Depth-major layout makes each Markov step a contiguous [vocab, lanes] + // view while retaining one lm_head projection for the whole cohort. + ggml_tensor * base = ggml_mul_mat(ctx, lm_head, all_hidden); + ggml_tensor * prev_ids = seed_tokens; + out.n_lanes = n_lanes; + out.q_len = q_len; + out.tokens.assign((size_t)n_depths, nullptr); + out.confidence.assign((size_t)n_depths, nullptr); + + for (int depth = 0; depth < n_depths; ++depth) { + ggml_tensor * prev_emb = + ggml_get_rows(ctx, dw.dspark.markov_w1, prev_ids); + ggml_tensor * bias = + ggml_mul_mat(ctx, dw.dspark.markov_w2, prev_emb); + ggml_tensor * base_depth = ggml_view_2d( + ctx, base, vocab, n_lanes, base->nb[1], + (size_t)depth * (size_t)n_lanes * base->nb[1]); + ggml_tensor * corrected = ggml_add(ctx, base_depth, bias); + ggml_tensor * tok = ggml_argmax(ctx, corrected); + ggml_set_output(tok); + ggml_build_forward_expand(gf, tok); + out.tokens[(size_t)depth] = tok; + + if (have_confidence) { + ggml_tensor * conf_in = confidence_by_depth[(size_t)depth]; + if (dw.dspark.confidence_dim == + hdim + dw.dspark.markov_rank) { + conf_in = ggml_concat(ctx, conf_in, prev_emb, 0); + } + ggml_tensor * conf = + ggml_mul_mat(ctx, dw.dspark.confidence_w, conf_in); + conf = ggml_add( + ctx, conf, + ggml_reshape_2d(ctx, dw.dspark.confidence_b, 1, 1)); + conf = ggml_sigmoid(ctx, conf); + ggml_set_output(conf); + ggml_build_forward_expand(gf, conf); + out.confidence[(size_t)depth] = conf; + } + prev_ids = tok; + } + return true; +} + bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, ggml_backend_t backend, ggml_tensor * lm_head, diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h index 9b97b261d..6259482e0 100644 --- a/server/src/common/dspark_head.h +++ b/server/src/common/dspark_head.h @@ -36,6 +36,31 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, std::vector * confidence_out = nullptr, const float * confidence_hidden = nullptr); +// Outputs embedded in a caller-owned graph for a lane-batched DSpark chain. +// Each depth tensor is shaped [n_lanes], with confidence in [1, n_lanes]. +struct DSparkBatchedChainOutputs { + int n_lanes = 0; + int q_len = 0; + std::vector tokens; + std::vector confidence; +}; + +// Append a depth-major, multi-lane Markov chain to an existing draft graph. +// The lane backbones remain independent; their hidden tensors stay on-device. +// One lm_head matmul covers every (depth, lane), then each depth performs a +// batched Markov lookup, correction, argmax, and calibrated confidence head. +bool build_dspark_markov_batched_chain( + ggml_context * ctx, + ggml_cgraph * gf, + const DraftWeights & dw, + ggml_tensor * lm_head, + const std::vector & hidden_by_lane, + const std::vector & prenorm_by_lane, + ggml_tensor * seed_tokens, + int q_len, + bool want_confidence, + DSparkBatchedChainOutputs & out); + // DDTree candidate generation with the Markov correction: base logits for // all n_tokens positions in ONE lm_head matmul; rows 1..n-1 get the low-rank // previous-token bias chained along the main (argmax) path; top-K extracted diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 1a83b7878..146e66c64 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -62,6 +62,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( tree_scratch_stride_(tree_scratch_stride), spec_mode_(spec_mode) { const int n_slots = slots_.slot_count(); slot_draft_kv_.resize((size_t)n_slots); + prepared_chain_drafts_.resize((size_t)n_slots); last_survival_score_.assign( (size_t)n_slots, std::numeric_limits::quiet_NaN()); last_survival_generated_.assign((size_t)n_slots, -1); @@ -105,6 +106,11 @@ Qwen35SeqEngine::Qwen35SeqEngine( } Qwen35SeqEngine::~Qwen35SeqEngine() { + draft_kv_batch_free(batch_draft_graph_); + for (std::unique_ptr & state : dummy_draft_kv_) { + if (state) draft_kv_free(*state); + } + dummy_draft_kv_.clear(); for (std::unique_ptr & state : slot_draft_kv_) { if (state) draft_kv_free(*state); } @@ -403,6 +409,7 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { std::vector noise_embed((size_t)hidden * T); std::vector local_hidden((size_t)hidden * T); std::vector prenorm_hidden((size_t)hidden * T); + const bool profile_batched = batched_drafting_enabled(); if (!b_.w_.embedder.embed(noise.data(), T, noise_embed.data())) { cleanup(); return false; @@ -411,41 +418,109 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { if (!profile_error.empty()) return std::numeric_limits::infinity(); const auto start = std::chrono::steady_clock::now(); + std::vector states; + std::vector seeds; + states.reserve((size_t)chain_decode_bucket_width(lanes)); + seeds.reserve(states.capacity()); for (int lane = 0; lane < lanes; ++lane) { const int slot = synthetic_slots[(size_t)lane]; DraftKvState * draft = ensure_slot_draft_kv(slot); DraftFeatureMirror * mirror = slot_feature_mirror(slot); if (!draft || !mirror || - !draft_kv_begin_step(*draft, b_.dw_, b_.draft_backend_, - *mirror, ctx_tokens)) { + !draft_kv_begin_step( + *draft, b_.dw_, b_.draft_backend_, + *mirror, ctx_tokens)) { profile_error = "draft profiling setup failed"; return std::numeric_limits::infinity(); } - 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) { - profile_error = "draft profiling backbone failed"; + ggml_backend_tensor_set( + draft->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + states.push_back(draft); + seeds.push_back(profile_token); + } + + if (profile_batched) { + const int bucket = chain_decode_bucket_width(lanes); + const int dummy_count = bucket - lanes; + const int cap = std::min( + slot_feature_mirrors_[0].cap, + std::max(1, b_.cfg_.draft_ctx_max)); + while ((int)dummy_draft_kv_.size() < dummy_count) { + auto dummy = std::make_unique(); + if (!draft_kv_init( + *dummy, b_.dw_, b_.draft_backend_, + cap, nullptr)) { + draft_kv_free(*dummy); + break; + } + dummy_draft_kv_.push_back(std::move(dummy)); + } + if ((int)dummy_draft_kv_.size() < dummy_count) { + profile_error = "draft profiling dummy allocation failed"; return std::numeric_limits::infinity(); } - ggml_backend_tensor_get_async( - b_.draft_backend_, draft->hidden_states, - local_hidden.data(), 0, sizeof(float) * local_hidden.size()); - ggml_backend_tensor_get_async( - b_.draft_backend_, draft->hidden_prenorm, - prenorm_hidden.data(), 0, - sizeof(float) * prenorm_hidden.size()); - ggml_backend_synchronize(b_.draft_backend_); - std::vector draft_tokens; - std::vector confidence; - if (!dspark_markov_correct_greedy_chain_fused( - b_.dw_, b_.draft_backend_, b_.w_.output, - local_hidden.data(), T, profile_token, draft_tokens, - &confidence, prenorm_hidden.data()) || - (int)draft_tokens.size() != T) { - profile_error = "draft profiling chain failed"; + for (int i = 0; i < dummy_count; ++i) { + DraftKvState * dummy = + dummy_draft_kv_[(size_t)i].get(); + if (!draft_kv_begin_step( + *dummy, b_.dw_, b_.draft_backend_, + slot_feature_mirrors_[0], 1)) { + profile_error = "draft profiling dummy setup failed"; + return std::numeric_limits::infinity(); + } + ggml_backend_tensor_set( + dummy->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + states.push_back(dummy); + seeds.push_back(profile_token); + } + std::vector> draft_tokens; + std::vector> confidence; + if (!draft_kv_batch_compute( + batch_draft_graph_, b_.dw_, b_.draft_backend_, + b_.w_.output, states, seeds, + draft_tokens, confidence) || + draft_tokens.size() < (size_t)lanes) { + profile_error = "batched draft profiling compute failed"; return std::numeric_limits::infinity(); } + for (int lane = 0; lane < lanes; ++lane) { + if ((int)draft_tokens[(size_t)lane].size() != T) { + profile_error = "batched draft profiling shape failed"; + return std::numeric_limits::infinity(); + } + } + } else { + for (int lane = 0; lane < lanes; ++lane) { + DraftKvState * draft = states[(size_t)lane]; + if (ggml_backend_graph_compute( + b_.draft_backend_, draft->gf) != + GGML_STATUS_SUCCESS) { + profile_error = "draft profiling backbone failed"; + return std::numeric_limits::infinity(); + } + ggml_backend_tensor_get_async( + b_.draft_backend_, draft->hidden_states, + local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + ggml_backend_tensor_get_async( + b_.draft_backend_, draft->hidden_prenorm, + prenorm_hidden.data(), 0, + sizeof(float) * prenorm_hidden.size()); + ggml_backend_synchronize(b_.draft_backend_); + std::vector draft_tokens; + std::vector confidence; + if (!dspark_markov_correct_greedy_chain_fused( + b_.dw_, b_.draft_backend_, b_.w_.output, + local_hidden.data(), T, profile_token, + draft_tokens, &confidence, + prenorm_hidden.data()) || + (int)draft_tokens.size() != T) { + profile_error = "draft profiling chain failed"; + return std::numeric_limits::infinity(); + } + } } return std::chrono::duration( std::chrono::steady_clock::now() - start).count(); @@ -493,7 +568,9 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { std::fprintf(stderr, "\n"); }; std::fprintf(stderr, - "[spec-profile] context=%d reps=5 mode=serial-draft\n", ctx_tokens); + "[spec-profile] context=%d reps=5 mode=%s-draft\n", + ctx_tokens, + profile_batched ? "batched" : "serial"); print_table("tree_cost", tables.tree_cost); print_table("step_cost", tables.step_cost); print_table("draft_cost", tables.draft_cost); @@ -528,6 +605,195 @@ DraftKvState * Qwen35SeqEngine::ensure_slot_draft_kv(int slot) { } return state.get(); } +bool Qwen35SeqEngine::batched_drafting_enabled() const { + const char * value = std::getenv("DFLASH_SPEC_BATCHED_DRAFT"); + return !value || std::atoi(value) != 0; +} + +bool Qwen35SeqEngine::draft_always_enabled() const { + const char * value = std::getenv("DFLASH_SPEC_DRAFT_ALWAYS"); + return !value || std::atoi(value) != 0; +} + +bool Qwen35SeqEngine::prepare_chain_drafts( + const std::vector & inputs, + const std::vector & selected) { + if (selected.size() != inputs.size()) return false; + + const int T = tree_width_; + const int hidden = b_.w_.n_embd; + struct Lane { + size_t input_index = 0; + int slot = -1; + int32_t seed = -1; + DraftKvState * state = nullptr; + DraftFeatureMirror * mirror = nullptr; + }; + std::vector lanes; + lanes.reserve(inputs.size()); + std::vector noise((size_t)T, b_.w_.mask_token_id); + std::vector noise_embed((size_t)hidden * T); + + for (size_t i = 0; i < inputs.size(); ++i) { + if (!selected[i]) continue; + const StepInput & in = inputs[i]; + if (in.slot >= 0 && + in.slot < (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)in.slot].valid = false; + } + if (!chain_spec_input_eligible(in)) return false; + DraftKvState * state = ensure_slot_draft_kv(in.slot); + DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); + if (!state || !mirror || + !draft_kv_begin_step( + *state, b_.dw_, b_.draft_backend_, *mirror, + slots_.slot(in.slot).cur_pos)) { + return false; + } + noise[0] = in.token; + std::fill( + noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), T, noise_embed.data())) { + return false; + } + ggml_backend_tensor_set( + state->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + lanes.push_back({i, in.slot, in.token, state, mirror}); + } + if (lanes.empty()) return true; + + std::vector> drafts; + std::vector> confidences; + bool used_batch = false; + if (batched_drafting_enabled()) { + const int bucket = + chain_decode_bucket_width((int)lanes.size()); + std::vector batch_states; + std::vector seeds; + batch_states.reserve((size_t)bucket); + seeds.reserve((size_t)bucket); + for (const Lane & lane : lanes) { + batch_states.push_back(lane.state); + seeds.push_back(lane.seed); + } + + const int dummy_count = bucket - (int)lanes.size(); + const int cap = std::min( + lanes[0].mirror->cap, + std::max(1, b_.cfg_.draft_ctx_max)); + while ((int)dummy_draft_kv_.size() < dummy_count) { + auto dummy = std::make_unique(); + if (!draft_kv_init( + *dummy, b_.dw_, b_.draft_backend_, cap, nullptr)) { + draft_kv_free(*dummy); + break; + } + dummy_draft_kv_.push_back(std::move(dummy)); + } + if ((int)dummy_draft_kv_.size() >= dummy_count) { + noise[0] = lanes[0].seed; + std::fill( + noise.begin() + 1, noise.end(), + b_.w_.mask_token_id); + bool dummy_ok = b_.w_.embedder.embed( + noise.data(), T, noise_embed.data()); + for (int i = 0; dummy_ok && i < dummy_count; ++i) { + DraftKvState * dummy = + dummy_draft_kv_[(size_t)i].get(); + dummy_ok = draft_kv_begin_step( + *dummy, b_.dw_, b_.draft_backend_, + *lanes[0].mirror, 1); + if (dummy_ok) { + ggml_backend_tensor_set( + dummy->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + batch_states.push_back(dummy); + seeds.push_back(lanes[0].seed); + } + } + if (dummy_ok) { + used_batch = draft_kv_batch_compute( + batch_draft_graph_, b_.dw_, + b_.draft_backend_, b_.w_.output, + batch_states, seeds, drafts, confidences); + } + } + if (!used_batch) { + static bool warned = false; + if (!warned) { + warned = true; + std::fprintf(stderr, + "[draft-kv-batch] unavailable; using serial fallback\n"); + } + } + } + + if (!used_batch) { + drafts.resize(lanes.size()); + confidences.resize(lanes.size()); + std::vector local_hidden((size_t)hidden * T); + std::vector prenorm_hidden((size_t)hidden * T); + for (size_t lane = 0; lane < lanes.size(); ++lane) { + DraftKvState * state = lanes[lane].state; + if (ggml_backend_graph_compute( + b_.draft_backend_, state->gf) != + GGML_STATUS_SUCCESS) { + return false; + } + ggml_backend_tensor_get_async( + b_.draft_backend_, state->hidden_states, + local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + ggml_backend_tensor_get_async( + b_.draft_backend_, state->hidden_prenorm, + prenorm_hidden.data(), 0, + sizeof(float) * prenorm_hidden.size()); + ggml_backend_synchronize(b_.draft_backend_); + if (!dspark_markov_correct_greedy_chain_fused( + b_.dw_, b_.draft_backend_, b_.w_.output, + local_hidden.data(), T, lanes[lane].seed, + drafts[lane], &confidences[lane], + prenorm_hidden.data())) { + return false; + } + } + } + + if (drafts.size() < lanes.size() || + confidences.size() < lanes.size()) { + return false; + } + static bool missing_confidence_warned = false; + for (size_t lane = 0; lane < lanes.size(); ++lane) { + if ((int)drafts[lane].size() != T) return false; + const Lane & info = lanes[lane]; + PreparedChainDraft & prepared = + prepared_chain_drafts_[(size_t)info.slot]; + prepared.valid = true; + prepared.generated = + slots_.slot(info.slot).generated_tokens(); + prepared.root = info.seed; + prepared.tokens = std::move(drafts[lane]); + prepared.confidence = std::move(confidences[lane]); + + double score = std::numeric_limits::quiet_NaN(); + if (!prepared.confidence.empty()) { + score = confidence_survival_yield( + prepared.confidence, tree_width_); + } else if (!missing_confidence_warned) { + missing_confidence_warned = true; + std::fprintf(stderr, + "[spec-gate] calibrated confidence unavailable; " + "using measured yield/prior\n"); + } + last_survival_score_[(size_t)info.slot] = score; + last_survival_generated_[(size_t)info.slot] = + prepared.generated; + } + return true; +} bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { if (spec_mode_ != SpecMode::ddtree || tree_width_ <= 1 || @@ -633,70 +899,55 @@ std::optional Qwen35SeqEngine::step_chain_spec( return std::nullopt; }; - const char * draft_always_env = std::getenv("DFLASH_SPEC_DRAFT_ALWAYS"); - const bool draft_always = - draft_always_env && std::atoi(draft_always_env) != 0; - std::vector noise(static_cast(T), b_.w_.mask_token_id); - std::vector noise_embed(static_cast(hidden) * T); - std::vector local_hidden(static_cast(hidden) * T); - std::vector prenorm_hidden(static_cast(hidden) * T); - - // Serial C3 baseline: each slot owns a context-KV ring. C6 replaces this - // loop with the packed multi-lane graph while retaining it as fallback. + std::vector need_prepare(inputs.size(), 0); for (size_t i = 0; i < inputs.size(); ++i) { const StepInput & in = inputs[i]; const bool hard_eligible = chain_spec_input_eligible(in); if (admitted[i] && !hard_eligible) return proposal_fallback(); - if (!hard_eligible || (!admitted[i] && !draft_always)) continue; - - DraftKvState * draft = ensure_slot_draft_kv(in.slot); - DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); - if (!draft || !mirror) return proposal_fallback(); + if (!admitted[i]) continue; drafted_slots.push_back(in.slot); + const PreparedChainDraft & prepared = + prepared_chain_drafts_[(size_t)in.slot]; + const Qwen35Slot & seq = slots_.slot(in.slot); + need_prepare[i] = + !prepared.valid || + prepared.generated != seq.generated_tokens() || + prepared.root != in.token || + (int)prepared.tokens.size() != T; + } + if (std::any_of( + need_prepare.begin(), need_prepare.end(), + [](uint8_t value) { return value != 0; }) && + !prepare_chain_drafts(inputs, need_prepare)) { + 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(), T, 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_get_async( - b_.draft_backend_, draft->hidden_states, local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); - ggml_backend_tensor_get_async( - b_.draft_backend_, draft->hidden_prenorm, prenorm_hidden.data(), 0, - sizeof(float) * prenorm_hidden.size()); - ggml_backend_synchronize(b_.draft_backend_); - - std::vector draft_tokens; - std::vector confidence; - if (!dspark_markov_correct_greedy_chain_fused( - b_.dw_, b_.draft_backend_, b_.w_.output, - local_hidden.data(), T, in.token, draft_tokens, - &confidence, prenorm_hidden.data()) || - static_cast(draft_tokens.size()) != T) { + for (size_t i = 0; i < inputs.size(); ++i) { + if (!admitted[i]) continue; + const StepInput & in = inputs[i]; + PreparedChainDraft & prepared = + prepared_chain_drafts_[(size_t)in.slot]; + if (!prepared.valid || + prepared.generated != + slots_.slot(in.slot).generated_tokens() || + prepared.root != in.token || + (int)prepared.tokens.size() != T) { return proposal_fallback(); } - if (!admitted[i]) continue; Proposal proposal; proposal.input_index = i; proposal.slot = in.slot; proposal.root = in.token; - proposal.flat = std::move(draft_tokens); + proposal.flat = std::move(prepared.tokens); + proposal.confidence = std::move(prepared.confidence); + prepared.valid = false; proposal.tree = make_dspark_chain_tree(proposal.flat); - proposal.confidence = std::move(confidence); - if (proposal.tree.n_nodes + 1 != T) return proposal_fallback(); - proposal_for_input[i] = static_cast(proposals.size()); + if (proposal.tree.n_nodes + 1 != T) { + return proposal_fallback(); + } + proposal_for_input[i] = + static_cast(proposals.size()); proposals.push_back(std::move(proposal)); } if (static_cast(proposals.size()) != spec_count) { @@ -1598,6 +1849,10 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( last_survival_score_[(size_t)result.slot] = std::numeric_limits::quiet_NaN(); last_survival_generated_[(size_t)result.slot] = -1; + if (result.slot < + (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)result.slot].valid = false; + } } } return result; @@ -1840,16 +2095,41 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (spec_mode_ == SpecMode::dspark_chain && plan.prefills.empty()) { + const auto chain_started = std::chrono::steady_clock::now(); std::vector admitted(inputs.size(), 0); SpecPlan gate_plan; bool have_gate_plan = false; const char * force_value = std::getenv("DFLASH_SPEC_GATE_FORCE"); const std::string force = force_value ? force_value : ""; - const char * draft_always_value = - std::getenv("DFLASH_SPEC_DRAFT_ALWAYS"); - const bool draft_always = - draft_always_value && std::atoi(draft_always_value) != 0; + bool draft_always = draft_always_enabled(); + if (draft_always && speculation_gate_) { + std::vector draft_selected(inputs.size(), 0); + for (size_t i = 0; i < inputs.size(); ++i) { + SpeculationPolicy policy = + inputs[i].speculation_policy; + if (policy == SpeculationPolicy::Adaptive) { + if (force == "all") { + policy = SpeculationPolicy::Always; + } else if (force == "none") { + policy = SpeculationPolicy::Never; + } + } + draft_selected[i] = + policy != SpeculationPolicy::Never && + chain_spec_input_eligible(inputs[i]); + } + const bool any_draft = std::any_of( + draft_selected.begin(), draft_selected.end(), + [](uint8_t value) { return value != 0; }); + if (any_draft && + !prepare_chain_drafts(inputs, draft_selected)) { + draft_always = false; + std::fprintf(stderr, + "[draft-kv-batch] current-step drafting failed; " + "using admitted-only fallback\n"); + } + } if (speculation_gate_) { std::vector candidates; @@ -1863,7 +2143,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (force == "none") policy = SpeculationPolicy::Never; } const bool eligible = chain_spec_input_eligible(in); - drafting_lanes += eligible ? 1 : 0; + drafting_lanes += + eligible && policy != SpeculationPolicy::Never ? 1 : 0; double confidence = std::numeric_limits::quiet_NaN(); if (eligible && in.slot >= 0 && @@ -1911,12 +2192,11 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { admitted.begin(), admitted.end(), [](uint8_t value) { return value != 0; }); if (any_admitted) { - const auto started = std::chrono::steady_clock::now(); std::optional speculative = step_chain_spec(plan, admitted); const double measured_us = std::chrono::duration( - std::chrono::steady_clock::now() - started).count(); + std::chrono::steady_clock::now() - chain_started).count(); if (have_gate_plan && spec_gate_debug_enabled()) { std::fprintf(stderr, "[spec-gate] C=%zu k=%d scores=[", @@ -2326,6 +2606,9 @@ void Qwen35SeqEngine::retire(int slot) { last_survival_score_[(size_t)slot] = std::numeric_limits::quiet_NaN(); last_survival_generated_[(size_t)slot] = -1; + if (slot < (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)slot].valid = false; + } } slots_.retire(slot); } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 127a8757a..73d971cd2 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -134,6 +134,18 @@ class Qwen35SeqEngine final : public SeqEngine { bool ddtree_eligible(const StepPlan & plan) const; bool chain_spec_input_eligible(const StepInput & input) const; bool spec_gate_debug_enabled() const; + struct PreparedChainDraft { + bool valid = false; + int generated = -1; + int32_t root = -1; + std::vector tokens; + std::vector confidence; + }; + bool prepare_chain_drafts( + const std::vector & inputs, + const std::vector & selected); + bool batched_drafting_enabled() const; + bool draft_always_enabled() 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); @@ -152,6 +164,9 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector slot_feature_mirrors_; std::vector> slot_draft_kv_; + DraftKvBatchGraph batch_draft_graph_; + std::vector> dummy_draft_kv_; + std::vector prepared_chain_drafts_; std::unique_ptr speculation_gate_; std::vector last_survival_score_; std::vector last_survival_generated_; diff --git a/server/test/test_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp index 42b4a20b2..33c011f1c 100644 --- a/server/test/test_chain_spec_shapes.cpp +++ b/server/test/test_chain_spec_shapes.cpp @@ -15,6 +15,16 @@ int main() { CHECK((tree.token_ids == std::vector{11, 12, 13})); CHECK((tree.depths == std::vector{1, 2, 3})); CHECK((tree.parents == std::vector{-1, 0, 1, 2})); + const std::vector bucket_inputs = { + 1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 13, 16, 17, + }; + const std::vector bucket_expected = { + 1, 2, 3, 4, 6, 6, 8, 8, 12, 12, 16, 16, 24, + }; + for (size_t i = 0; i < bucket_inputs.size(); ++i) { + CHECK(chain_decode_bucket_width(bucket_inputs[i]) == + bucket_expected[i]); + } int pending = -1; const int32_t full_posterior[] = {11, 12, 13, 14}; diff --git a/server/test/test_dspark_batched_head.cpp b/server/test/test_dspark_batched_head.cpp new file mode 100644 index 000000000..64ee5634f --- /dev/null +++ b/server/test/test_dspark_batched_head.cpp @@ -0,0 +1,224 @@ +#include "common/dspark_head.h" +#include "host_check.h" + +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +int main() { + constexpr int hidden = 3; + constexpr int vocab = 5; + constexpr int rank = 2; + constexpr int q_len = 4; + constexpr int lanes = 2; + constexpr int confidence_dim = hidden + rank; + + ggml_backend_t backend = ggml_backend_cpu_init(); + CHECK(backend != nullptr); + if (!backend) return 1; + + ggml_init_params weights_params{}; + weights_params.mem_size = ggml_tensor_overhead() * 8; + weights_params.no_alloc = true; + ggml_context * weights_ctx = ggml_init(weights_params); + CHECK(weights_ctx != nullptr); + if (!weights_ctx) { + ggml_backend_free(backend); + return 1; + } + + ggml_tensor * lm_head = ggml_new_tensor_2d( + weights_ctx, GGML_TYPE_F32, hidden, vocab); + ggml_tensor * markov_w1 = ggml_new_tensor_2d( + weights_ctx, GGML_TYPE_F32, rank, vocab); + ggml_tensor * markov_w2 = ggml_new_tensor_2d( + weights_ctx, GGML_TYPE_F32, rank, vocab); + ggml_tensor * confidence_w = ggml_new_tensor_2d( + weights_ctx, GGML_TYPE_F32, confidence_dim, 1); + ggml_tensor * confidence_b = ggml_new_tensor_1d( + weights_ctx, GGML_TYPE_F32, 1); + ggml_backend_buffer_t weights_buf = + ggml_backend_alloc_ctx_tensors(weights_ctx, backend); + CHECK(weights_buf != nullptr); + if (!weights_buf) { + ggml_free(weights_ctx); + ggml_backend_free(backend); + return 1; + } + + std::vector lm((size_t)hidden * vocab); + std::vector w1((size_t)rank * vocab); + std::vector w2((size_t)rank * vocab); + std::vector cw((size_t)confidence_dim); + for (int token = 0; token < vocab; ++token) { + for (int h = 0; h < hidden; ++h) { + lm[(size_t)token * hidden + h] = + 0.031f * (float)(token + 1) * (float)(h + 1); + } + for (int r = 0; r < rank; ++r) { + w1[(size_t)token * rank + r] = + 0.017f * (float)(token + 1 + r); + w2[(size_t)token * rank + r] = + 0.013f * (float)(token + 1) * (float)(r + 1); + } + } + for (int i = 0; i < confidence_dim; ++i) { + cw[(size_t)i] = 0.021f * (float)(i + 1); + } + const float cb = -0.11f; + ggml_backend_tensor_set(lm_head, lm.data(), 0, sizeof(float) * lm.size()); + ggml_backend_tensor_set( + markov_w1, w1.data(), 0, sizeof(float) * w1.size()); + ggml_backend_tensor_set( + markov_w2, w2.data(), 0, sizeof(float) * w2.size()); + ggml_backend_tensor_set( + confidence_w, cw.data(), 0, sizeof(float) * cw.size()); + ggml_backend_tensor_set(confidence_b, &cb, 0, sizeof(cb)); + + DraftWeights dw; + dw.n_embd = hidden; + dw.block_size = q_len; + dw.dspark.enabled = true; + dw.dspark.markov_rank = rank; + dw.dspark.vocab_size = vocab; + dw.dspark.confidence_dim = confidence_dim; + dw.dspark.markov_w1 = markov_w1; + dw.dspark.markov_w2 = markov_w2; + dw.dspark.confidence_w = confidence_w; + dw.dspark.confidence_b = confidence_b; + + std::vector> hidden_host( + lanes, std::vector((size_t)hidden * q_len)); + std::vector> prenorm_host( + lanes, std::vector((size_t)hidden * q_len)); + const int32_t seeds[lanes] = {1, 3}; + for (int lane = 0; lane < lanes; ++lane) { + for (int position = 0; position < q_len; ++position) { + for (int h = 0; h < hidden; ++h) { + const size_t index = + (size_t)position * hidden + h; + hidden_host[(size_t)lane][index] = + 0.07f * (float)(1 + lane + 2 * position + h); + prenorm_host[(size_t)lane][index] = + hidden_host[(size_t)lane][index] + + 0.019f * (float)(h + 1); + } + } + } + + std::vector> serial_tokens(lanes); + std::vector> serial_confidence(lanes); + for (int lane = 0; lane < lanes; ++lane) { + CHECK(dspark_markov_correct_greedy_chain_fused( + dw, backend, lm_head, hidden_host[(size_t)lane].data(), + q_len, seeds[lane], serial_tokens[(size_t)lane], + &serial_confidence[(size_t)lane], + prenorm_host[(size_t)lane].data())); + } + + std::vector arena(4u * 1024u * 1024u); + ggml_init_params graph_params{}; + graph_params.mem_size = arena.size(); + graph_params.mem_buffer = arena.data(); + graph_params.no_alloc = true; + ggml_context * graph_ctx = ggml_init(graph_params); + CHECK(graph_ctx != nullptr); + ggml_cgraph * graph = + ggml_new_graph_custom(graph_ctx, 2048, false); + + std::vector hidden_inputs(lanes); + std::vector prenorm_inputs(lanes); + for (int lane = 0; lane < lanes; ++lane) { + hidden_inputs[(size_t)lane] = ggml_new_tensor_2d( + graph_ctx, GGML_TYPE_F32, hidden, q_len); + prenorm_inputs[(size_t)lane] = ggml_new_tensor_2d( + graph_ctx, GGML_TYPE_F32, hidden, q_len); + ggml_set_input(hidden_inputs[(size_t)lane]); + ggml_set_input(prenorm_inputs[(size_t)lane]); + } + ggml_tensor * seed_input = + ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, lanes); + ggml_set_input(seed_input); + + DSparkBatchedChainOutputs outputs; + CHECK(build_dspark_markov_batched_chain( + graph_ctx, graph, dw, lm_head, + hidden_inputs, prenorm_inputs, seed_input, + q_len, true, outputs)); + CHECK(outputs.n_lanes == lanes); + CHECK(outputs.q_len == q_len); + CHECK(outputs.tokens.size() == (size_t)(q_len - 1)); + CHECK(outputs.confidence.size() == (size_t)(q_len - 1)); + + ggml_gallocr_t allocator = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + CHECK(allocator != nullptr); + const bool graph_allocated = + allocator && ggml_gallocr_alloc_graph(allocator, graph); + CHECK(graph_allocated); + if (graph_allocated) { + for (int lane = 0; lane < lanes; ++lane) { + ggml_backend_tensor_set( + hidden_inputs[(size_t)lane], + hidden_host[(size_t)lane].data(), 0, + sizeof(float) * hidden_host[(size_t)lane].size()); + ggml_backend_tensor_set( + prenorm_inputs[(size_t)lane], + prenorm_host[(size_t)lane].data(), 0, + sizeof(float) * prenorm_host[(size_t)lane].size()); + } + ggml_backend_tensor_set( + seed_input, seeds, 0, sizeof(seeds)); + CHECK(ggml_backend_graph_compute(backend, graph) == + GGML_STATUS_SUCCESS); + + std::vector depth_tokens( + (size_t)(q_len - 1) * lanes); + std::vector depth_confidence( + (size_t)(q_len - 1) * lanes); + for (int depth = 0; depth < q_len - 1; ++depth) { + ggml_backend_tensor_get_async( + backend, outputs.tokens[(size_t)depth], + depth_tokens.data() + (size_t)depth * lanes, + 0, sizeof(int32_t) * lanes); + ggml_backend_tensor_get_async( + backend, outputs.confidence[(size_t)depth], + depth_confidence.data() + (size_t)depth * lanes, + 0, sizeof(float) * lanes); + } + ggml_backend_synchronize(backend); + + for (int lane = 0; lane < lanes; ++lane) { + CHECK(serial_tokens[(size_t)lane].size() == (size_t)q_len); + CHECK(serial_confidence[(size_t)lane].size() == + (size_t)(q_len - 1)); + CHECK(serial_tokens[(size_t)lane][0] == seeds[lane]); + for (int depth = 0; depth < q_len - 1; ++depth) { + CHECK(serial_tokens[(size_t)lane][(size_t)depth + 1] == + depth_tokens[(size_t)depth * lanes + lane]); + CHECK(std::fabs( + serial_confidence[(size_t)lane][(size_t)depth] - + depth_confidence[ + (size_t)depth * lanes + lane]) < 1e-6f); + } + } + } + if (allocator) ggml_gallocr_free(allocator); + + ggml_free(graph_ctx); + ggml_backend_buffer_free(weights_buf); + ggml_free(weights_ctx); + ggml_backend_free(backend); + std::printf( + "DSpark batched-head parity tests passed: %d checks\n", + g_checks); + return 0; +} From ab61d32d63a018f83db8512621353574eb3d7f97 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 18:56:03 +0000 Subject: [PATCH 15/42] feat(harness): add Qwen3.8 DSpark oracle matrix --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 47 +++ .../concurrency/generate_dspark_prompts.py | 115 +++++++ .../concurrency/run_qwen38_dspark_matrix.sh | 325 ++++++++++++++++++ .../concurrency/summarize_feature_matrix.py | 258 +++++++++++++- .../concurrency/test_feature_tools.py | 162 +++++++++ .../concurrency/verify_feature_metrics.py | 82 ++++- .../concurrency/write_feature_metadata.py | 8 +- 7 files changed, 989 insertions(+), 8 deletions(-) create mode 100755 harness/benchmarks/concurrency/generate_dspark_prompts.py create mode 100755 harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 62fce6a46..9d19f396a 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -3,6 +3,53 @@ The bounded Strix Halo measurements collected for the draft implementation are recorded in [`STRIX_HALO_RESULTS.md`](STRIX_HALO_RESULTS.md). +## Qwen3.8 DSpark adaptive matrix + +`run_qwen38_dspark_matrix.sh` is the C7 acceptance-gated speculation matrix. +It uses only the [RadixArk Qwen3.8-27B-DSpark](https://huggingface.co/RadixArk/Qwen3.8-27B-DSpark) +source. `DRAFT_MODEL` must point to the **q4-mix requantized drafter** produced +from that repository. + +```bash +MODEL=/opt/models/Qwen3.8-27B-Q4_K_M.gguf \ +DRAFT_MODEL=/opt/models/Qwen3.8-27B-DSpark-RadixArk-q4-mix.gguf \ +REPEATS=5 \ +harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh +``` + +The default fresh-process matrix is: + +- `ar`, `speculation`, `adaptive-on`, and `adaptive-off`, where the adaptive + suffix records `DFLASH_SPEC_DRAFT_ALWAYS`. +- Live concurrency `C ∈ {1,2,3,4,6,8}` over the checked-in HumanEval and + GSM8K cohorts plus deterministic prose prompts. +- A fixed C=6 north-star cohort with two code and four chat requests. +- Batched drafting enabled for every row; adaptive startup profiling uses a + 4096-token synthetic context. + +Every process records the target, server, shared-library, and drafter hashes, +the literal command and launch environment, startup pool dimensions, request +IDs, and terminal concurrency counters. The proof rejects forced-speculation +rows unless every measured request has positive `spec_steps`. Adaptive rows +may legitimately choose k=0, but must show both the packed DSpark startup +marker and a completed startup cost profile. Chain rows must keep all +`ddtree_*` counters at zero, preserving the DDTree proof semantics below. + +For every workload/concurrency pair, the summarizer forms a paired oracle: + +```text +goodput oracle = max(ar, speculation) +TTFT oracle = min(ar, speculation) +``` + +Both adaptive policies must reach at least 0.995 of that oracle for mean and +median output goodput and inverse TTFT. The summary fails the run if any of +those four ratios misses the gate; p95 alone is never used as acceptance +evidence. Set `WORKLOADS`, `CLIENTS`, `DECODE_MODES`, or +`ADAPTIVE_DRAFT_ALWAYS` to select a smaller diagnostic subset. + +## Qwen3.6 DDTree/PFlash/KVFlash matrix + `run_qwen36_feature_matrix.sh` extends the PR #596 protocol with feature ablations for the complete Strix Halo configuration: diff --git a/harness/benchmarks/concurrency/generate_dspark_prompts.py b/harness/benchmarks/concurrency/generate_dspark_prompts.py new file mode 100755 index 000000000..2e7f22357 --- /dev/null +++ b/harness/benchmarks/concurrency/generate_dspark_prompts.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +"""Build deterministic Qwen3.8 DSpark concurrency workload manifests.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from generate_ragged_prompts import write_records + + +HERE = Path(__file__).resolve().parent +PROMPT_ROOT = HERE.parent / "prompts" +SOURCE_FILES = { + "humaneval": PROMPT_ROOT / "bench_he.jsonl", + "gsm8k": PROMPT_ROOT / "bench_gsm.jsonl", +} +PROSE_TOPICS = ( + "why reproducible benchmarks need immutable inputs and explicit hardware metadata", + "how admission control improves the reliability of a concurrent inference service", + "the tradeoff between latency, throughput, and fairness in token scheduling", + "why calibrated confidence is useful when deciding whether to speculate", + "how bounded caches turn memory pressure into a predictable operating policy", + "the value of fail-closed telemetry when validating an optimization", + "why deterministic greedy output is a strong regression oracle", + "how batching amortizes launch overhead without sharing request state", + "why startup profiling should use monotone cost tables", + "how paired fresh-process repeats reduce benchmark order bias", +) + + +def _message_prompt(row: dict[str, object], source: Path, line_no: int) -> str: + messages = row.get("messages") + if not isinstance(messages, list) or not messages: + raise ValueError(f"{source}:{line_no}: missing messages") + pieces: list[str] = [] + for message in messages: + if not isinstance(message, dict): + raise ValueError(f"{source}:{line_no}: message must be an object") + role = message.get("role") + content = message.get("content") + if not isinstance(role, str) or not isinstance(content, str) or not content: + raise ValueError(f"{source}:{line_no}: invalid message") + pieces.append(content if len(messages) == 1 else f"{role}: {content}") + return "\n\n".join(pieces) + + +def _source_records(profile: str) -> list[dict[str, object]]: + source = SOURCE_FILES[profile] + records: list[dict[str, object]] = [] + for line_no, raw in enumerate(source.read_text(encoding="utf-8").splitlines(), 1): + if not raw.strip(): + continue + row = json.loads(raw) + records.append({ + "id": str(row.get("id") or f"{profile}-{line_no:02d}"), + "suite": profile, + "prompt": _message_prompt(row, source, line_no), + }) + if len(records) < 8: + raise ValueError(f"{source}: need at least eight distinct prompts") + return records + + +def _prose_records() -> list[dict[str, object]]: + return [ + { + "id": f"prose-{index:02d}", + "suite": "prose", + "prompt": ( + "Write a clear, self-contained technical essay of about 500 words on " + f"{topic}. Include one concrete example and end with a concise conclusion." + ), + } + for index, topic in enumerate(PROSE_TOPICS, 1) + ] + + +def build_records(profile: str) -> list[dict[str, object]]: + if profile in SOURCE_FILES: + return _source_records(profile) + prose = _prose_records() + if profile == "prose": + return prose + if profile == "north-star": + code = _source_records("humaneval")[:2] + chat = [dict(row) for row in prose[:4]] + for row in code: + row["cohort"] = "2code+4chat" + for row in chat: + row["cohort"] = "2code+4chat" + return code + chat + raise ValueError(f"unknown DSpark workload {profile!r}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--profile", choices=("humaneval", "gsm8k", "prose", "north-star"), + 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)} {args.profile} prompts to {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh new file mode 100755 index 000000000..5ab4df50a --- /dev/null +++ b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh @@ -0,0 +1,325 @@ +#!/usr/bin/env bash +# Fresh-process Qwen3.8 DSpark mode matrix with fail-closed chain 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_dspark_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:-}" +LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" +OUT="${OUT:-$REPO/.harness-runs/qwen38-dspark-matrix-$(date -u +%Y%m%dT%H%M%SZ)}" +REPEATS="${REPEATS:-1}" +WORKLOADS="${WORKLOADS:-humaneval,gsm8k,prose,north-star}" +DECODE_MODES="${DECODE_MODES:-ar,speculation,adaptive}" +ADAPTIVE_DRAFT_ALWAYS="${ADAPTIVE_DRAFT_ALWAYS:-on,off}" +CLIENTS="${CLIENTS:-1,2,3,4,6,8}" +SLOTS="${SLOTS:-8}" +MAX_CTX="${MAX_CTX:-8192}" +MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" +MAX_TOKENS="${MAX_TOKENS:-256}" +WARMUP_TOKENS="${WARMUP_TOKENS:-16}" +PROFILE_CONTEXT="${PROFILE_CONTEXT:-4096}" +PORT="${PORT:-18138}" +COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-3}" +HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-900}" +REQUEST_TIMEOUT_SECONDS="${REQUEST_TIMEOUT_SECONDS:-1800}" +TARGET_DEVICE="${TARGET_DEVICE:-hip:0}" +DRAFT_DEVICE="${DRAFT_DEVICE:-hip:0}" +VISIBLE_DEVICES="${VISIBLE_DEVICES:-1}" + +usage() { + cat <<'EOF' +Usage: + MODEL=/path/Qwen3.8-27B-Q4_K_M.gguf \ + DRAFT_MODEL=/path/Qwen3.8-27B-DSpark-RadixArk-q4-mix.gguf \ + harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh + +The only supported drafter source for this matrix is: + https://huggingface.co/RadixArk/Qwen3.8-27B-DSpark + +DRAFT_MODEL is the q4-mix requantized drafter produced from that repository. +The default fresh-process matrix runs ar, forced speculation, adaptive with +always-drafting on, and adaptive with always-drafting off at live concurrency +1,2,3,4,6,8 over HumanEval, GSM8K, and prose. The 2-code+4-chat north-star row +runs only at C=6. The summary fails if either adaptive policy is below 0.995 of +the paired ar/speculation oracle in mean or median goodput or TTFT. + +Defaults select the second visible host GPU and address it as hip:0 inside the +process (VISIBLE_DEVICES=1, TARGET_DEVICE=hip:0, DRAFT_DEVICE=hip:0). Override +VISIBLE_DEVICES on a single-GPU Strix Halo host. 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 Qwen3.8 target GGUF" >&2; exit 2; } +[[ -r "$DRAFT_MODEL" ]] || { echo "set DRAFT_MODEL to the readable RadixArk q4-mix GGUF" >&2; exit 2; } +[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +for value_name in REPEATS SLOTS MAX_CTX MAX_CONCURRENT_PREFILLS MAX_TOKENS WARMUP_TOKENS PROFILE_CONTEXT HEALTH_TIMEOUT_SECONDS REQUEST_TIMEOUT_SECONDS; do + value="${!value_name}" + [[ "$value" =~ ^[1-9][0-9]*$ ]] || { echo "$value_name must be positive" >&2; exit 2; } +done +[[ "$PORT" =~ ^[1-9][0-9]*$ ]] || { echo "PORT must be positive" >&2; exit 2; } +[[ "$COOLDOWN_SECONDS" =~ ^[0-9]+$ ]] || { echo "COOLDOWN_SECONDS must be non-negative" >&2; exit 2; } +(( PROFILE_CONTEXT < MAX_CTX )) || { echo "PROFILE_CONTEXT must be below MAX_CTX" >&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 mode_list <<< "$DECODE_MODES" +IFS=, read -r -a adaptive_axis <<< "$ADAPTIVE_DRAFT_ALWAYS" +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 WORKLOADS "${workload_list[@]}" || exit 2 +reject_duplicates DECODE_MODES "${mode_list[@]}" || exit 2 +reject_duplicates ADAPTIVE_DRAFT_ALWAYS "${adaptive_axis[@]}" || exit 2 +reject_duplicates CLIENTS "${client_list[@]}" || exit 2 + +for workload in "${workload_list[@]}"; do + case "$workload" in + humaneval|gsm8k|prose|north-star) ;; + *) echo "unknown workload $workload" >&2; exit 2 ;; + esac +done +for mode in "${mode_list[@]}"; do + case "$mode" in + ar|speculation|adaptive) ;; + *) echo "unknown decode mode $mode" >&2; exit 2 ;; + esac +done +for axis in "${adaptive_axis[@]}"; do + [[ "$axis" == on || "$axis" == off ]] || { + echo "ADAPTIVE_DRAFT_ALWAYS entries must be on or off" >&2 + exit 2 + } +done +for clients in "${client_list[@]}"; do + [[ "$clients" =~ ^[1-9][0-9]*$ ]] || { echo "CLIENTS entries must be positive" >&2; exit 2; } + (( clients <= SLOTS )) || { echo "CLIENTS=$clients exceeds SLOTS=$SLOTS" >&2; exit 2; } +done + +variants=() +for mode in "${mode_list[@]}"; do + if [[ "$mode" == adaptive ]]; then + for axis in "${adaptive_axis[@]}"; do variants+=("adaptive-$axis"); done + else + variants+=("$mode") + fi +done +(( ${#variants[@]} > 0 )) || { echo "no decode variants selected" >&2; exit 2; } + +MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" +DRAFT_MODEL_SHA256="$(sha256sum "$DRAFT_MODEL" | 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 +} + +case_applicable() { + local workload="$1" clients="$2" + [[ "$workload" != north-star || "$clients" == 6 ]] +} + +run_case() { + local repeat="$1" workload="$2" clients="$3" variant="$4" + local decode_mode="$variant" draft_always="" + if [[ "$variant" == adaptive-* ]]; then + decode_mode=adaptive + draft_always="${variant#adaptive-}" + fi + + local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" + mkdir -p "$case_dir" + local capacity=$((SLOTS * MAX_CTX)) + local model_id=qwen38-dspark + local -a command=( + "$LUCE_SERVER_BIN" "$MODEL" + --draft "$DRAFT_MODEL" + --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_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 20 --draft-residency persistent + --decode-mode "$decode_mode" + --host 127.0.0.1 --port "$PORT" --model-name "$model_id" + ) + local -a launch_env=( + "HIP_VISIBLE_DEVICES=$VISIBLE_DEVICES" + "DFLASH_MAX_CONCURRENT_PREFILLS=$MAX_CONCURRENT_PREFILLS" + "DFLASH_SPEC_BATCHED_DRAFT=1" + ) + if [[ "$decode_mode" == adaptive ]]; then + launch_env+=( + "DFLASH_SPEC_GATE_LOG=1" + "DFLASH_SPEC_PROFILE_CONTEXT=$PROFILE_CONTEXT" + ) + if [[ "$draft_always" == on ]]; then + launch_env+=("DFLASH_SPEC_DRAFT_ALWAYS=1") + else + launch_env+=("DFLASH_SPEC_DRAFT_ALWAYS=0") + fi + fi + + printf 'env ' > "$case_dir/server-command.txt" + printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$case_dir/server-command.txt" + printf '\n' >> "$case_dir/server-command.txt" + + local -a metadata=( + python3 "$METADATA_TOOL" + --out "$case_dir/server-metadata.json" + --variant "$variant" --workload "$workload" + --clients "$clients" --repeat "$repeat" + --binary "$LUCE_SERVER_BIN" --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_CONCURRENT_PREFILLS" + --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_DEVICE" + --draft-model "$DRAFT_MODEL" --draft-model-sha256 "$DRAFT_MODEL_SHA256" + --decode-mode "$decode_mode" + ) + [[ -n "$draft_always" ]] && metadata+=(--draft-always "$draft_always") + local item + for item in "${launch_env[@]}"; do metadata+=(--launch-env "$item"); done + "${metadata[@]}" + + echo "[run] $workload C=$clients repeat=$repeat variant=$variant" + port_is_available || return 1 + env "${launch_env[@]}" "${command[@]}" > "$case_dir/server.log" 2>&1 & + server_pid=$! + if ! wait_health; then tail -n 160 "$case_dir/server.log" >&2 || true; return 1; fi + python3 "$RUNTIME_METADATA_TOOL" \ + --metadata "$case_dir/server-metadata.json" \ + --server-log "$case_dir/server.log" + + local 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 0 + --require-distinct-prompts --temperature 0 --ignore-eos + --require-effective-prompt-telemetry + --timeout "$REQUEST_TIMEOUT_SECONDS" --cooldown 0 + ) + local -a warmup_cmd=( + python3 "$CLIENT" "${common_client[@]}" + --max-tokens "$WARMUP_TOKENS" + --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" + --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 + + local -a proof_cmd=( + python3 "$PROOF_TOOL" + --bench "$case_dir/bench.json" + --server-log "$case_dir/server.log" + --out "$case_dir/feature-proof.json" + ) + [[ "$decode_mode" != ar ]] && proof_cmd+=(--expect chain) + "${proof_cmd[@]}" + 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]}" + if ! case_applicable "$workload" "$clients"; then + echo "[skip] north-star is fixed at C=6; C=$clients" + continue + fi + shift_by=$(((repeat + c_index) % ${#variants[@]})) + for ((i=0; i<${#variants[@]}; i++)); do + variant="${variants[$(((i + shift_by) % ${#variants[@]}))]}" + 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_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py index b789a5fbd..f6c1789dd 100755 --- a/harness/benchmarks/concurrency/summarize_feature_matrix.py +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Summarize Qwen3.6 concurrent feature ablations and activation proof.""" +"""Summarize Qwen3.6 feature ablations or Qwen3.8 DSpark oracle runs.""" from __future__ import annotations @@ -40,6 +40,9 @@ def load_reports(root: Path) -> list[dict]: expected_by_variant = { "ar": [], "ddtree": ["ddtree"], "pflash": ["pflash"], "kvflash": ["kvflash"], "full": ["ddtree", "kvflash", "pflash"], + "speculation": ["chain"], + "adaptive-on": ["chain"], + "adaptive-off": ["chain"], } if variant not in expected_by_variant: raise ValueError(f"{path}: unknown Lucebox variant {variant!r}") @@ -47,6 +50,20 @@ def load_reports(root: Path) -> list[dict]: raise ValueError( f"{proof_path}: expected_features does not match variant {variant}" ) + mode_by_variant = { + "ar": "ar", "speculation": "speculation", + "adaptive-on": "adaptive", "adaptive-off": "adaptive", + } + if variant in mode_by_variant: + mode = (meta.get("feature_config") or {}).get("decode_mode") + # `ar` is shared with the older Qwen3.6 matrix, whose metadata + # intentionally has no decode_mode field. + if ((variant != "ar" or mode is not None) + and (mode != mode_by_variant[variant] + or proof.get("decode_mode") != mode)): + raise ValueError( + f"{path}: decode_mode proof does not match variant {variant}" + ) reports.append({ "path": path, "report": report, "level": level, "meta": meta, "proof": proof, @@ -102,7 +119,7 @@ def output_stability(items: list[dict]) -> str: ) -def summarize(reports: list[dict]) -> str: +def summarize_qwen36(reports: list[dict]) -> str: grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) for item in reports: meta, level = item["meta"], item["level"] @@ -215,6 +232,243 @@ def summarize(reports: list[dict]) -> str: return "\n".join(lines) +DSPARK_VARIANTS = {"ar", "speculation", "adaptive-on", "adaptive-off"} +ORACLE_THRESHOLD = 0.995 + + +def _is_dspark_item(item: dict) -> bool: + config = item["meta"].get("feature_config") or {} + return config.get("decode_mode") in ("ar", "speculation", "adaptive") + + +def _positive_metric(item: dict, key: str) -> float: + value = item["level"].get(key) + if type(value) not in (int, float) or value <= 0: + raise ValueError(f"{item['path']}: missing positive {key}") + return float(value) + + +def _dspark_signature(item: dict) -> tuple[object, ...]: + base = run_signature(item) + meta = item["meta"] + config = meta.get("feature_config") or {} + draft_sha256 = config.get("draft_model_sha256") + binary_sha256 = meta.get("server_binary_sha256") + if ( + not isinstance(draft_sha256, str) or not draft_sha256 + or not isinstance(binary_sha256, str) or not binary_sha256 + ): + raise ValueError(f"{item['path']}: incomplete DSpark provenance") + return (*base, draft_sha256, binary_sha256) + + +def _mean(values: list[float]) -> float: + return statistics.fmean(values) + + +def _pair(values: list[float]) -> tuple[float, float]: + return _mean(values), median(values) + + +def _fmt_pair(values: list[float], digits: int = 2) -> str: + mean_value, median_value = _pair(values) + return f"{mean_value:.{digits}f}/{median_value:.{digits}f}" + + +def summarize_dspark(reports: list[dict]) -> str: + grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) + for item in reports: + meta, level = item["meta"], item["level"] + variant = str(meta.get("variant", "")) + if variant not in DSPARK_VARIANTS: + raise ValueError(f"{item['path']}: mixed or unknown DSpark variant {variant!r}") + config = meta.get("feature_config") or {} + expected_axis = ( + variant.removeprefix("adaptive-") + if variant.startswith("adaptive-") else None + ) + if config.get("draft_always") != expected_axis: + raise ValueError( + f"{item['path']}: draft_always does not match variant {variant}" + ) + key = (str(meta["workload"]), int(level["clients"]), 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({_dspark_signature(item) for item in items}) != 1: + raise ValueError(f"{key}: incompatible DSpark run metadata") + prompt_hashes = { + item["level"].get("selected_prompt_set_sha256") for item in items + } + if len(prompt_hashes) != 1 or not all( + isinstance(value, str) and value for value in prompt_hashes + ): + raise ValueError(f"{key}: missing or inconsistent selected prompts") + + lines = [ + "# Qwen3.8 DSpark adaptive concurrency matrix", "", + "Every row requires request-correlated telemetry. Forced speculation must " + "record positive chain steps for every request; adaptive rows must prove " + "the packed DSpark backend and startup profile were active.", + "", + "Mean/median are computed across paired fresh-process repeats. The oracle " + "uses max(AR, speculation) for goodput and min(AR, speculation) for TTFT " + "within each repeat.", + "", + "| Workload | C | Mode | N | Goodput mean/median | Oracle goodput " + "mean/median | Adaptive/oracle goodput mean/median | TTFT mean/median s | " + "Oracle/adaptive TTFT mean/median | Spec accepted/step | Spec steps | " + "Target forwards | Stable output |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " + "---: | ---: | ---: | :---: |", + ] + regressions: list[str] = [] + variant_order = {"ar": 0, "speculation": 1, "adaptive-on": 2, "adaptive-off": 3} + for workload, clients, variant in sorted( + grouped, key=lambda key: (key[0], key[1], variant_order[key[2]]) + ): + items = grouped[(workload, clients, variant)] + by_repeat = {int(item["meta"]["repeat"]): item for item in items} + goodputs = [_positive_metric(item, "aggregate_tok_s") for item in items] + ttfts = [_positive_metric(item, "ttft_median_s") for item in items] + controls = { + name: grouped.get((workload, clients, name), []) + for name in ("ar", "speculation") + } + oracle_goodputs: list[float] = [] + oracle_ttfts: list[float] = [] + if controls["ar"] and controls["speculation"]: + control_maps = { + name: {int(item["meta"]["repeat"]): item for item in values} + for name, values in controls.items() + } + if control_maps["ar"].keys() != control_maps["speculation"].keys(): + raise ValueError(f"{workload} C={clients}: oracle repeat sets differ") + for repeat in sorted(control_maps["ar"]): + ar = control_maps["ar"][repeat] + speculation = control_maps["speculation"][repeat] + if _dspark_signature(ar) != _dspark_signature(speculation): + raise ValueError( + f"{workload} C={clients} repeat={repeat}: oracle metadata differs" + ) + if ( + ar["level"].get("selected_prompt_set_sha256") + != speculation["level"].get("selected_prompt_set_sha256") + ): + raise ValueError( + f"{workload} C={clients} repeat={repeat}: oracle prompts differ" + ) + ar_output = ar["level"].get("selected_output_set_sha256") + speculation_output = speculation["level"].get( + "selected_output_set_sha256" + ) + if not isinstance(ar_output, str) or ar_output != speculation_output: + raise ValueError( + f"{workload} C={clients} repeat={repeat}: " + "AR/speculation outputs differ" + ) + oracle_goodputs.append(max( + _positive_metric(ar, "aggregate_tok_s"), + _positive_metric(speculation, "aggregate_tok_s"), + )) + oracle_ttfts.append(min( + _positive_metric(ar, "ttft_median_s"), + _positive_metric(speculation, "ttft_median_s"), + )) + + goodput_ratio = "—" + ttft_ratio = "—" + if variant.startswith("adaptive-"): + if not oracle_goodputs: + raise ValueError(f"{workload} C={clients}: adaptive row lacks AR/spec oracle") + oracle_repeats = { + int(item["meta"]["repeat"]) for item in controls["ar"] + } + if by_repeat.keys() != oracle_repeats: + raise ValueError( + f"{workload} C={clients} {variant}: repeat set differs from oracle" + ) + reference = controls["ar"][0] + if any(_dspark_signature(item) != _dspark_signature(reference) for item in items): + raise ValueError( + f"{workload} C={clients} {variant}: metadata differs from oracle" + ) + if any( + item["level"].get("selected_prompt_set_sha256") + != reference["level"].get("selected_prompt_set_sha256") + for item in items + ): + raise ValueError( + f"{workload} C={clients} {variant}: prompts differ from oracle" + ) + ar_by_repeat = { + int(item["meta"]["repeat"]): item for item in controls["ar"] + } + for repeat, item in by_repeat.items(): + adaptive_output = item["level"].get("selected_output_set_sha256") + ar_output = ar_by_repeat[repeat]["level"].get( + "selected_output_set_sha256" + ) + if not isinstance(adaptive_output, str) or adaptive_output != ar_output: + raise ValueError( + f"{workload} C={clients} {variant} repeat={repeat}: " + "adaptive/AR outputs differ" + ) + gp_mean, gp_median = _pair(goodputs) + oracle_gp_mean, oracle_gp_median = _pair(oracle_goodputs) + ttft_mean, ttft_median = _pair(ttfts) + oracle_ttft_mean, oracle_ttft_median = _pair(oracle_ttfts) + gp_ratios = (gp_mean / oracle_gp_mean, gp_median / oracle_gp_median) + ttft_ratios = ( + oracle_ttft_mean / ttft_mean, + oracle_ttft_median / ttft_median, + ) + goodput_ratio = f"{gp_ratios[0]:.3f}/{gp_ratios[1]:.3f}" + ttft_ratio = f"{ttft_ratios[0]:.3f}/{ttft_ratios[1]:.3f}" + if min(*gp_ratios, *ttft_ratios) < ORACLE_THRESHOLD: + regressions.append( + f"{workload} C={clients} {variant}: goodput={goodput_ratio} " + f"ttft={ttft_ratio}" + ) + + aggregates = [item["proof"]["aggregate"] for item in items] + spec_steps = sum(value["spec_steps"] for value in aggregates) + accepted = sum(value["spec_accepted_tokens"] for value in aggregates) + accepted_per_step = accepted / spec_steps if spec_steps else 0.0 + target_forwards = median([value["target_forwards"] for value in aggregates]) + stable = output_stability(items) + lines.append( + f"| {workload} | {clients} | {variant} | {len(items)} | " + f"{_fmt_pair(goodputs)} | " + f"{_fmt_pair(oracle_goodputs) if oracle_goodputs else 'n/a'} | " + f"{goodput_ratio} | {_fmt_pair(ttfts, 3)} | {ttft_ratio} | " + f"{accepted_per_step:.2f} | {spec_steps} | {target_forwards:.0f} | " + f"{stable} |" + ) + lines += [ + "", + f"Oracle-relative gate: every adaptive mean/median goodput and inverse " + f"TTFT ratio must be >= {ORACLE_THRESHOLD:.3f}.", + ] + if regressions: + raise ValueError( + "oracle-relative criterion failed: " + "; ".join(regressions) + ) + return "\n".join(lines) + + +def summarize(reports: list[dict]) -> str: + dspark = [_is_dspark_item(item) for item in reports] + if any(dspark): + if not all(dspark): + raise ValueError("cannot mix Qwen3.6 feature and Qwen3.8 DSpark rows") + return summarize_dspark(reports) + return summarize_qwen36(reports) + + def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("root", type=Path) diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index 5e5645f7d..38d54e42b 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -27,6 +27,7 @@ def load(name: str): generator = load("generate_feature_prompts") +dspark_generator = load("generate_dspark_prompts") proof = load("verify_feature_metrics") summary = load("summarize_feature_matrix") @@ -73,6 +74,8 @@ def metric(request_id: str, effective: int, page_outs: int = 1) -> dict: "ddtree_suspensions": 0, # Zero acceptance is legitimate and must not invalidate execution proof. "ddtree_accepted_tokens": 0, + "spec_steps": 0, + "spec_accepted_tokens": 0, "target_forwards": 3, "kvflash_page_ins": 0, "kvflash_page_outs": page_outs, @@ -101,6 +104,18 @@ def test_activation_profiles_are_disjoint_and_above_thresholds(self) -> None: self.assertGreaterEqual(min(row["target_words"] for row in pressure), 12000) self.assertTrue(all(row["activation_target"] == "pflash-auto" for row in compression)) + def test_dspark_profiles_reuse_public_fixtures_and_build_north_star(self) -> None: + humaneval = dspark_generator.build_records("humaneval") + gsm8k = dspark_generator.build_records("gsm8k") + prose = dspark_generator.build_records("prose") + north_star = dspark_generator.build_records("north-star") + self.assertGreaterEqual(len(humaneval), 8) + self.assertGreaterEqual(len(gsm8k), 8) + self.assertGreaterEqual(len(prose), 8) + self.assertEqual(len(north_star), 6) + self.assertEqual([row["suite"] for row in north_star[:2]], ["humaneval"] * 2) + self.assertEqual([row["suite"] for row in north_star[2:]], ["prose"] * 4) + self.assertEqual(len({row["prompt"] for row in north_star}), 6) class FeatureRunnerShellTests(unittest.TestCase): def run_invalid_matrix( @@ -453,6 +468,78 @@ def test_measured_request_requires_explicit_error_status(self) -> None: with self.assertRaisesRegex(ValueError, "explicit error status"): proof.measured_requests(input_report) + def test_forced_chain_requires_steps_and_matching_startup_proof(self) -> None: + input_report = report(variant="speculation") + input_report["server_metadata"]["feature_config"]["decode_mode"] = ( + "speculation" + ) + rows = [metric("r1", 2000), metric("r2", 2100)] + for row in rows: + row["ddtree_steps"] = 0 + row["ddtree_accepted_tokens"] = 0 + row["spec_steps"] = 3 + row["spec_accepted_tokens"] = 2 + startup = ( + "[parallel-dspark] enabled width=7 mode=packed-chain-verify " + "decode_mode=speculation draft=q4-mix-compatible" + ) + result = proof.verify(input_report, rows, {"chain"}, startup) + self.assertTrue(result["valid"], result["errors"]) + self.assertEqual(result["decode_mode"], "speculation") + self.assertEqual(result["aggregate"]["spec_steps"], 6) + self.assertEqual(result["aggregate"]["spec_accepted_tokens"], 4) + + rows[0]["spec_steps"] = 0 + rows[0]["spec_accepted_tokens"] = 0 + result = proof.verify(input_report, rows, {"chain"}, startup) + self.assertFalse(result["valid"]) + self.assertIn("spec_steps is zero", "\n".join(result["errors"])) + + result = proof.verify(input_report, rows[1:], {"chain"}, "") + self.assertFalse(result["valid"]) + self.assertIn("startup proof is missing", "\n".join(result["errors"])) + + def test_adaptive_chain_allows_ar_argmax_but_requires_profile(self) -> None: + input_report = report(variant="adaptive-on") + input_report["server_metadata"]["feature_config"]["decode_mode"] = ( + "adaptive" + ) + rows = [metric("r1", 2000), metric("r2", 2100)] + for row in rows: + row["ddtree_steps"] = 0 + row["ddtree_accepted_tokens"] = 0 + startup = "\n".join(( + "[spec-profile] context=4096 reps=5 mode=batched-draft", + "[parallel-dspark] enabled width=7 mode=packed-chain-verify " + "decode_mode=adaptive draft=q4-mix-compatible", + )) + result = proof.verify(input_report, rows, {"chain"}, startup) + self.assertTrue(result["valid"], result["errors"]) + + no_profile = startup.splitlines()[1] + result = proof.verify(input_report, rows, {"chain"}, no_profile) + self.assertFalse(result["valid"]) + self.assertIn("cost-profile proof is missing", "\n".join(result["errors"])) + + def test_ar_decode_mode_rejects_speculation_activity(self) -> None: + input_report = report(variant="ar") + input_report["server_metadata"]["feature_config"]["decode_mode"] = "ar" + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[0]["spec_steps"] = 1 + result = proof.verify(input_report, rows, set()) + self.assertFalse(result["valid"]) + self.assertIn( + "AR decode_mode emitted chain speculation", "\n".join(result["errors"]) + ) + + def test_spec_acceptance_without_step_is_rejected(self) -> None: + row = metric("r1", 2000) + row["spec_accepted_tokens"] = 1 + with self.assertRaisesRegex( + ValueError, "spec_accepted_tokens requires positive spec_steps", + ): + proof.aggregate_rows([row]) + class FeatureSummaryTests(unittest.TestCase): @staticmethod @@ -490,6 +577,53 @@ def item( }, }, } + @staticmethod + def dspark_item( + variant: str, goodput: float, ttft: float, *, repeat: int = 1, + ) -> dict: + mode = ( + "adaptive" if variant.startswith("adaptive-") else variant + ) + draft_always = ( + variant.removeprefix("adaptive-") + if variant.startswith("adaptive-") else None + ) + spec_steps = 0 if variant == "ar" else 8 + return { + "path": Path(f"/tmp/{variant}-r{repeat}/bench.json"), + "report": { + "max_tokens": 256, "ignore_eos": True, + "temperature": 0.0, "seed": 1, + }, + "meta": { + "workload": "humaneval", "variant": variant, + "repeat": repeat, "model_sha256": "a" * 64, + "server_binary_sha256": "b" * 64, + "feature_config": { + "decode_mode": mode, + "draft_always": draft_always, + "draft_model_sha256": "c" * 64, + }, + }, + "level": { + "clients": 6, + "aggregate_tok_s": goodput, + "ttft_median_s": ttft, + "selected_prompt_set_sha256": "same-prompts", + "selected_output_set_sha256": "same-output", + }, + "proof": { + "aggregate": { + "ddtree_steps": 0, + "ddtree_suspensions": 0, + "ddtree_accepted_tokens": 0, + "spec_steps": spec_steps, + "spec_accepted_tokens": spec_steps * 2, + "target_forwards": 16, + }, + }, + } + def test_summary_compares_feature_row_to_ar(self) -> None: text = summary.summarize([self.item("ar", 10.0), self.item("full", 12.0)]) @@ -585,6 +719,34 @@ def test_unstable_ar_control_suppresses_feature_delta(self) -> None: row = next(line for line in text.splitlines() if "| full |" in line) self.assertEqual(row.split("|")[7].strip(), "n/a") + def test_dspark_summary_enforces_mean_median_oracle_gate(self) -> None: + reports = [ + self.dspark_item("ar", 100.0, 1.0, repeat=1), + self.dspark_item("ar", 102.0, 1.1, repeat=2), + self.dspark_item("speculation", 98.0, 0.9, repeat=1), + self.dspark_item("speculation", 103.0, 1.0, repeat=2), + self.dspark_item("adaptive-on", 100.0, 0.9, repeat=1), + self.dspark_item("adaptive-on", 103.0, 1.0, repeat=2), + self.dspark_item("adaptive-off", 100.0, 0.9, repeat=1), + self.dspark_item("adaptive-off", 102.6, 1.0, repeat=2), + ] + text = summary.summarize(reports) + self.assertIn("Qwen3.8 DSpark adaptive concurrency matrix", text) + self.assertIn("Oracle-relative gate", text) + self.assertIn("| adaptive-on |", text) + self.assertIn("| adaptive-off |", text) + + def test_dspark_summary_rejects_oracle_regression(self) -> None: + reports = [ + self.dspark_item("ar", 100.0, 1.0), + self.dspark_item("speculation", 98.0, 0.9), + self.dspark_item("adaptive-on", 90.0, 1.5), + ] + with self.assertRaisesRegex( + ValueError, "oracle-relative criterion failed.*adaptive-on", + ): + summary.summarize(reports) + if __name__ == "__main__": unittest.main() diff --git a/harness/benchmarks/concurrency/verify_feature_metrics.py b/harness/benchmarks/concurrency/verify_feature_metrics.py index b8faf043b..b184b49e1 100644 --- a/harness/benchmarks/concurrency/verify_feature_metrics.py +++ b/harness/benchmarks/concurrency/verify_feature_metrics.py @@ -14,9 +14,13 @@ PREFIX = "[concurrency-metrics] " COUNTERS = ( "ddtree_steps", "ddtree_suspensions", "ddtree_accepted_tokens", + "spec_steps", "spec_accepted_tokens", "target_forwards", "kvflash_page_ins", "kvflash_page_outs", "kvflash_reselects", ) +DECODE_MODES = ("ar", "speculation", "adaptive") +DSPARK_STARTUP_PREFIX = "[parallel-dspark] enabled" +SPEC_PROFILE_PREFIX = "[spec-profile] context=" REQUIRED_KEYS = ( "request_id", "effective_prompt_tokens", *COUNTERS, "kvflash_resident_blocks", "pflash_applied", "pflash_input_tokens", @@ -108,11 +112,16 @@ def aggregate_rows(rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]: raise ValueError( f"{request_id}: ddtree_suspensions must be 0 or 1 per request" ) + if row["spec_accepted_tokens"] > 0 and row["spec_steps"] == 0: + raise ValueError( + f"{request_id}: spec_accepted_tokens requires positive spec_steps" + ) return aggregate def verify( report: dict[str, Any], markers: list[dict[str, Any]], expected: set[str], + server_log_text: str | None = None, ) -> dict[str, Any]: measured = measured_requests(report) all_rows = aggregate_rows(markers) @@ -122,6 +131,12 @@ def verify( if missing: errors.append(f"missing concurrency telemetry for {len(missing)} measured request(s): {missing}") + metadata = report.get("server_metadata") or {} + feature_config = metadata.get("feature_config") or {} + decode_mode = feature_config.get("decode_mode") + if decode_mode is not None and decode_mode not in DECODE_MODES: + errors.append(f"invalid recorded decode_mode {decode_mode!r}") + for request_id, measured_row in measured.items(): metric = rows.get(request_id) if metric is None: @@ -149,6 +164,24 @@ def verify( 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 "chain" in expected: + if metric["target_forwards"] <= 0: + errors.append( + f"{request_id}: chain decode requested but target_forwards is zero" + ) + if metric["ddtree_steps"] != 0 or metric["ddtree_accepted_tokens"] != 0: + errors.append( + f"{request_id}: chain run must keep DDTree counters at zero" + ) + if metric["ddtree_suspensions"] != 0: + errors.append( + f"{request_id}: chain run must keep ddtree_suspensions at zero" + ) + if decode_mode == "speculation" and metric["spec_steps"] <= 0: + errors.append( + f"{request_id}: forced chain speculation requested but " + "spec_steps 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") @@ -158,10 +191,42 @@ def verify( f"({metric['pflash_input_tokens']} -> {metric['pflash_output_tokens']})" ) + if "chain" in expected: + if decode_mode not in ("speculation", "adaptive"): + errors.append( + "chain requested but metadata decode_mode is not speculation or adaptive" + ) + log_text = server_log_text or "" + startup_mode = ( + isinstance(decode_mode, str) + and re.search( + rf"^.*{re.escape(DSPARK_STARTUP_PREFIX)}.*" + rf"decode_mode={re.escape(decode_mode)}\b.*" + r"draft=q4-mix-compatible.*$", + log_text, + flags=re.MULTILINE, + ) + ) + if not startup_mode: + errors.append( + "chain requested but matching packed DSpark startup proof is missing" + ) + if decode_mode == "adaptive" and SPEC_PROFILE_PREFIX not in log_text: + errors.append( + "adaptive chain requested but startup cost-profile proof is missing" + ) + elif decode_mode == "ar": + active_spec = sorted( + request_id for request_id, row in rows.items() + if row["spec_steps"] != 0 or row["spec_accepted_tokens"] != 0 + ) + if active_spec: + errors.append( + f"AR decode_mode emitted chain speculation for request(s): {active_spec}" + ) + 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") @@ -266,7 +331,8 @@ def verify( ) return { - "schema_version": 3, + "schema_version": 4, + "decode_mode": decode_mode, "expected_features": sorted(expected), "valid": not errors, "errors": errors, @@ -297,12 +363,18 @@ 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( + "--expect", action="append", + choices=("ddtree", "chain", "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)) + log_text = args.server_log.read_text(encoding="utf-8", errors="replace") + result = verify( + report, parse_markers(args.server_log), set(args.expect), log_text, + ) except Exception as exc: print(f"[proof] error: {exc}", file=sys.stderr) return 2 diff --git a/harness/benchmarks/concurrency/write_feature_metadata.py b/harness/benchmarks/concurrency/write_feature_metadata.py index 01e40146f..515cbcea1 100644 --- a/harness/benchmarks/concurrency/write_feature_metadata.py +++ b/harness/benchmarks/concurrency/write_feature_metadata.py @@ -90,6 +90,8 @@ def main() -> int: parser.add_argument("--draft-device", default=None) parser.add_argument("--draft-model", type=pathlib.Path) parser.add_argument("--draft-model-sha256") + parser.add_argument("--decode-mode", choices=("ar", "speculation", "adaptive")) + parser.add_argument("--draft-always", choices=("on", "off")) parser.add_argument("--ddtree", action="store_true") parser.add_argument("--fast-rollback", action="store_true") parser.add_argument("--ddtree-budget", type=int) @@ -138,6 +140,8 @@ def main() -> int: literal_flags += ["--target-device", args.target_device] if args.draft_device: literal_flags += ["--draft-device", args.draft_device] + if args.decode_mode: + literal_flags += ["--decode-mode", args.decode_mode] if args.ddtree: literal_flags += ["--ddtree"] if args.ddtree_budget is not None: @@ -154,7 +158,7 @@ def main() -> int: literal_flags += ["--kvflash", args.kvflash] obj = { - "schema_version": 3, + "schema_version": 4, "variant": args.variant, "workload": args.workload, "clients": args.clients, @@ -177,6 +181,8 @@ def main() -> int: "draft_device": args.draft_device, "draft_model": str(args.draft_model.resolve()) if args.draft_model else None, "draft_model_sha256": draft_model_sha256, + "decode_mode": args.decode_mode, + "draft_always": args.draft_always, "ddtree": args.ddtree, "fast_rollback": args.fast_rollback, "ddtree_budget": args.ddtree_budget, From 0fdc5f9889967de904ecc38a53ebc1a4b8e76b43 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 08:20:35 +0000 Subject: [PATCH 16/42] feat(concurrency): calibrate adaptive speculation confidence --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 19 ++- .../concurrency/run_qwen38_dspark_matrix.sh | 30 ++++- .../concurrency/summarize_feature_matrix.py | 109 ++++++++++++++-- .../concurrency/test_feature_tools.py | 28 +++- .../concurrency/write_feature_metadata.py | 2 + .../src/common/concurrency/speculation_gate.h | 120 ++++++++++++++++-- .../qwen35/concurrency/qwen35_seq_engine.cpp | 68 +++++++--- .../qwen35/concurrency/qwen35_seq_engine.h | 1 + server/test/test_speculation_gate.cpp | 60 +++++++++ 9 files changed, 379 insertions(+), 58 deletions(-) diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 9d19f396a..9f819fa67 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -19,8 +19,10 @@ harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh The default fresh-process matrix is: -- `ar`, `speculation`, `adaptive-on`, and `adaptive-off`, where the adaptive - suffix records `DFLASH_SPEC_DRAFT_ALWAYS`. +- `ar`, `speculation`, `adaptive-on`, `adaptive-off`, and + `adaptive-confidence-off`. The last arm keeps always-drafting enabled but + sets `DFLASH_SPEC_CONFIDENCE=0`, forcing the gate onto its measured/prior + fallback without changing the cost-aware top-k selector. - Live concurrency `C ∈ {1,2,3,4,6,8}` over the checked-in HumanEval and GSM8K cohorts plus deterministic prose prompts. - A fixed C=6 north-star cohort with two code and four chat requests. @@ -42,11 +44,14 @@ goodput oracle = max(ar, speculation) TTFT oracle = min(ar, speculation) ``` -Both adaptive policies must reach at least 0.995 of that oracle for mean and -median output goodput and inverse TTFT. The summary fails the run if any of -those four ratios misses the gate; p95 alone is never used as acceptance -evidence. Set `WORKLOADS`, `CLIENTS`, `DECODE_MODES`, or -`ADAPTIVE_DRAFT_ALWAYS` to select a smaller diagnostic subset. +Both confidence-enabled adaptive policies must reach at least 0.995 of that +oracle for mean and median output goodput and inverse TTFT. The summary fails +the run if any of those four ratios misses the gate; p95 alone is never used as +acceptance evidence. A second table reports paired goodput and inverse-TTFT +deltas between `adaptive-on` and `adaptive-confidence-off` without gating the +ablation. Set `WORKLOADS`, `CLIENTS`, `DECODE_MODES`, +`ADAPTIVE_DRAFT_ALWAYS`, or `CONFIDENCE_ABLATION=0` to select a smaller +diagnostic subset. ## Qwen3.6 DDTree/PFlash/KVFlash matrix diff --git a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh index 5ab4df50a..aaefe67e5 100755 --- a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh @@ -19,6 +19,7 @@ REPEATS="${REPEATS:-1}" WORKLOADS="${WORKLOADS:-humaneval,gsm8k,prose,north-star}" DECODE_MODES="${DECODE_MODES:-ar,speculation,adaptive}" ADAPTIVE_DRAFT_ALWAYS="${ADAPTIVE_DRAFT_ALWAYS:-on,off}" +CONFIDENCE_ABLATION="${CONFIDENCE_ABLATION:-1}" CLIENTS="${CLIENTS:-1,2,3,4,6,8}" SLOTS="${SLOTS:-8}" MAX_CTX="${MAX_CTX:-8192}" @@ -46,10 +47,12 @@ The only supported drafter source for this matrix is: DRAFT_MODEL is the q4-mix requantized drafter produced from that repository. The default fresh-process matrix runs ar, forced speculation, adaptive with -always-drafting on, and adaptive with always-drafting off at live concurrency -1,2,3,4,6,8 over HumanEval, GSM8K, and prose. The 2-code+4-chat north-star row -runs only at C=6. The summary fails if either adaptive policy is below 0.995 of -the paired ar/speculation oracle in mean or median goodput or TTFT. +always-drafting on, adaptive with always-drafting off, and an always-drafting +confidence-off ablation at live concurrency 1,2,3,4,6,8 over HumanEval, GSM8K, +and prose. The 2-code+4-chat north-star row runs only at C=6. The summary fails +if either confidence-enabled adaptive policy is below 0.995 of the paired +ar/speculation oracle in mean or median goodput or TTFT. The confidence +ablation delta is reported but not gated. Defaults select the second visible host GPU and address it as hip:0 inside the process (VISIBLE_DEVICES=1, TARGET_DEVICE=hip:0, DRAFT_DEVICE=hip:0). Override @@ -71,6 +74,7 @@ for value_name in REPEATS SLOTS MAX_CTX MAX_CONCURRENT_PREFILLS MAX_TOKENS WARMU done [[ "$PORT" =~ ^[1-9][0-9]*$ ]] || { echo "PORT must be positive" >&2; exit 2; } [[ "$COOLDOWN_SECONDS" =~ ^[0-9]+$ ]] || { echo "COOLDOWN_SECONDS must be non-negative" >&2; exit 2; } +[[ "$CONFIDENCE_ABLATION" == 0 || "$CONFIDENCE_ABLATION" == 1 ]] || { echo "CONFIDENCE_ABLATION must be 0 or 1" >&2; exit 2; } (( PROFILE_CONTEXT < MAX_CTX )) || { echo "PROFILE_CONTEXT must be below MAX_CTX" >&2; exit 2; } [[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } @@ -130,6 +134,9 @@ variants=() for mode in "${mode_list[@]}"; do if [[ "$mode" == adaptive ]]; then for axis in "${adaptive_axis[@]}"; do variants+=("adaptive-$axis"); done + if [[ "$CONFIDENCE_ABLATION" == 1 ]]; then + variants+=("adaptive-confidence-off") + fi else variants+=("$mode") fi @@ -193,10 +200,15 @@ case_applicable() { run_case() { local repeat="$1" workload="$2" clients="$3" variant="$4" - local decode_mode="$variant" draft_always="" - if [[ "$variant" == adaptive-* ]]; then + local decode_mode="$variant" draft_always="" confidence="" + if [[ "$variant" == adaptive-confidence-off ]]; then + decode_mode=adaptive + draft_always=on + confidence=off + elif [[ "$variant" == adaptive-* ]]; then decode_mode=adaptive draft_always="${variant#adaptive-}" + confidence=on fi local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" @@ -230,6 +242,11 @@ run_case() { else launch_env+=("DFLASH_SPEC_DRAFT_ALWAYS=0") fi + if [[ "$confidence" == on ]]; then + launch_env+=("DFLASH_SPEC_CONFIDENCE=1") + else + launch_env+=("DFLASH_SPEC_CONFIDENCE=0") + fi fi printf 'env ' > "$case_dir/server-command.txt" @@ -251,6 +268,7 @@ run_case() { --decode-mode "$decode_mode" ) [[ -n "$draft_always" ]] && metadata+=(--draft-always "$draft_always") + [[ -n "$confidence" ]] && metadata+=(--confidence "$confidence") local item for item in "${launch_env[@]}"; do metadata+=(--launch-env "$item"); done "${metadata[@]}" diff --git a/harness/benchmarks/concurrency/summarize_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py index f6c1789dd..8afb29bec 100755 --- a/harness/benchmarks/concurrency/summarize_feature_matrix.py +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -43,6 +43,7 @@ def load_reports(root: Path) -> list[dict]: "speculation": ["chain"], "adaptive-on": ["chain"], "adaptive-off": ["chain"], + "adaptive-confidence-off": ["chain"], } if variant not in expected_by_variant: raise ValueError(f"{path}: unknown Lucebox variant {variant!r}") @@ -53,6 +54,7 @@ def load_reports(root: Path) -> list[dict]: mode_by_variant = { "ar": "ar", "speculation": "speculation", "adaptive-on": "adaptive", "adaptive-off": "adaptive", + "adaptive-confidence-off": "adaptive", } if variant in mode_by_variant: mode = (meta.get("feature_config") or {}).get("decode_mode") @@ -232,7 +234,13 @@ def summarize_qwen36(reports: list[dict]) -> str: return "\n".join(lines) -DSPARK_VARIANTS = {"ar", "speculation", "adaptive-on", "adaptive-off"} +DSPARK_VARIANTS = { + "ar", + "speculation", + "adaptive-on", + "adaptive-off", + "adaptive-confidence-off", +} ORACLE_THRESHOLD = 0.995 @@ -275,6 +283,11 @@ def _fmt_pair(values: list[float], digits: int = 2) -> str: return f"{mean_value:.{digits}f}/{median_value:.{digits}f}" +def _fmt_percent_pair(values: list[float]) -> str: + mean_value, median_value = _pair(values) + return f"{mean_value * 100:+.1f}%/{median_value * 100:+.1f}%" + + def summarize_dspark(reports: list[dict]) -> str: grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) for item in reports: @@ -283,13 +296,20 @@ def summarize_dspark(reports: list[dict]) -> str: if variant not in DSPARK_VARIANTS: raise ValueError(f"{item['path']}: mixed or unknown DSpark variant {variant!r}") config = meta.get("feature_config") or {} - expected_axis = ( - variant.removeprefix("adaptive-") - if variant.startswith("adaptive-") else None - ) - if config.get("draft_always") != expected_axis: + expected_axes = { + "ar": (None, None), + "speculation": (None, None), + "adaptive-on": ("on", "on"), + "adaptive-off": ("off", "on"), + "adaptive-confidence-off": ("on", "off"), + } + expected_draft, expected_confidence = expected_axes[variant] + if ( + config.get("draft_always") != expected_draft + or config.get("confidence") != expected_confidence + ): raise ValueError( - f"{item['path']}: draft_always does not match variant {variant}" + f"{item['path']}: adaptive axes do not match variant {variant}" ) key = (str(meta["workload"]), int(level["clients"]), variant) grouped[key].append(item) @@ -326,7 +346,13 @@ def summarize_dspark(reports: list[dict]) -> str: "---: | ---: | ---: | :---: |", ] regressions: list[str] = [] - variant_order = {"ar": 0, "speculation": 1, "adaptive-on": 2, "adaptive-off": 3} + variant_order = { + "ar": 0, + "speculation": 1, + "adaptive-on": 2, + "adaptive-off": 3, + "adaptive-confidence-off": 4, + } for workload, clients, variant in sorted( grouped, key=lambda key: (key[0], key[1], variant_order[key[2]]) ): @@ -428,7 +454,10 @@ def summarize_dspark(reports: list[dict]) -> str: ) goodput_ratio = f"{gp_ratios[0]:.3f}/{gp_ratios[1]:.3f}" ttft_ratio = f"{ttft_ratios[0]:.3f}/{ttft_ratios[1]:.3f}" - if min(*gp_ratios, *ttft_ratios) < ORACLE_THRESHOLD: + if ( + variant != "adaptive-confidence-off" + and min(*gp_ratios, *ttft_ratios) < ORACLE_THRESHOLD + ): regressions.append( f"{workload} C={clients} {variant}: goodput={goodput_ratio} " f"ttft={ttft_ratio}" @@ -448,10 +477,68 @@ def summarize_dspark(reports: list[dict]) -> str: f"{accepted_per_step:.2f} | {spec_steps} | {target_forwards:.0f} | " f"{stable} |" ) + + ablation_keys = sorted({ + (workload, clients) + for workload, clients, variant in grouped + if ( + variant == "adaptive-confidence-off" + and (workload, clients, "adaptive-on") in grouped + ) + }) + if ablation_keys: + lines += [ + "", + "## Confidence ablation", + "", + "Paired per-repeat deltas compare adaptive always-drafting with " + "confidence enabled against the same policy with confidence hidden. " + "Positive values favor fresh confidence; these deltas are reported " + "but are not acceptance-gated.", + "", + "| Workload | C | Goodput delta mean/median | " + "Inverse-TTFT delta mean/median |", + "| :--- | ---: | ---: | ---: |", + ] + for workload, clients in ablation_keys: + confidence_items = grouped[(workload, clients, "adaptive-on")] + ablated_items = grouped[ + (workload, clients, "adaptive-confidence-off") + ] + confidence_by_repeat = { + int(item["meta"]["repeat"]): item for item in confidence_items + } + ablated_by_repeat = { + int(item["meta"]["repeat"]): item for item in ablated_items + } + if confidence_by_repeat.keys() != ablated_by_repeat.keys(): + raise ValueError( + f"{workload} C={clients}: confidence ablation repeats differ" + ) + goodput_deltas: list[float] = [] + inverse_ttft_deltas: list[float] = [] + for repeat in sorted(confidence_by_repeat): + enabled = confidence_by_repeat[repeat] + disabled = ablated_by_repeat[repeat] + goodput_deltas.append( + _positive_metric(enabled, "aggregate_tok_s") + / _positive_metric(disabled, "aggregate_tok_s") + - 1.0 + ) + inverse_ttft_deltas.append( + _positive_metric(disabled, "ttft_median_s") + / _positive_metric(enabled, "ttft_median_s") + - 1.0 + ) + lines.append( + f"| {workload} | {clients} | " + f"{_fmt_percent_pair(goodput_deltas)} | " + f"{_fmt_percent_pair(inverse_ttft_deltas)} |" + ) lines += [ "", - f"Oracle-relative gate: every adaptive mean/median goodput and inverse " - f"TTFT ratio must be >= {ORACLE_THRESHOLD:.3f}.", + f"Oracle-relative gate: every confidence-enabled adaptive mean/median " + f"goodput and inverse TTFT ratio must be >= {ORACLE_THRESHOLD:.3f}.", ] if regressions: raise ValueError( diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index 38d54e42b..297631256 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -180,6 +180,15 @@ def test_signals_exit_and_launch_environment_is_recorded(self) -> None: self.assertIn("'env ' > \"$case_dir/server-command.txt\"", runner) self.assertIn('"${launch_env[@]}" "${command[@]}"', runner) + def test_dspark_runner_has_explicit_confidence_ablation(self) -> None: + runner = (HERE / "run_qwen38_dspark_matrix.sh").read_text( + encoding="utf-8", + ) + self.assertIn('variants+=("adaptive-confidence-off")', runner) + self.assertIn('"DFLASH_SPEC_CONFIDENCE=1"', runner) + self.assertIn('"DFLASH_SPEC_CONFIDENCE=0"', runner) + self.assertIn('metadata+=(--confidence "$confidence")', 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") @@ -584,10 +593,15 @@ def dspark_item( mode = ( "adaptive" if variant.startswith("adaptive-") else variant ) - draft_always = ( - variant.removeprefix("adaptive-") - if variant.startswith("adaptive-") else None - ) + if variant == "adaptive-confidence-off": + draft_always = "on" + confidence = "off" + elif variant.startswith("adaptive-"): + draft_always = variant.removeprefix("adaptive-") + confidence = "on" + else: + draft_always = None + confidence = None spec_steps = 0 if variant == "ar" else 8 return { "path": Path(f"/tmp/{variant}-r{repeat}/bench.json"), @@ -602,6 +616,7 @@ def dspark_item( "feature_config": { "decode_mode": mode, "draft_always": draft_always, + "confidence": confidence, "draft_model_sha256": "c" * 64, }, }, @@ -729,12 +744,17 @@ def test_dspark_summary_enforces_mean_median_oracle_gate(self) -> None: self.dspark_item("adaptive-on", 103.0, 1.0, repeat=2), self.dspark_item("adaptive-off", 100.0, 0.9, repeat=1), self.dspark_item("adaptive-off", 102.6, 1.0, repeat=2), + self.dspark_item("adaptive-confidence-off", 75.0, 1.8, repeat=1), + self.dspark_item("adaptive-confidence-off", 80.0, 1.7, repeat=2), ] text = summary.summarize(reports) self.assertIn("Qwen3.8 DSpark adaptive concurrency matrix", text) self.assertIn("Oracle-relative gate", text) self.assertIn("| adaptive-on |", text) self.assertIn("| adaptive-off |", text) + self.assertIn("| adaptive-confidence-off |", text) + self.assertIn("## Confidence ablation", text) + self.assertIn("Positive values favor fresh confidence", text) def test_dspark_summary_rejects_oracle_regression(self) -> None: reports = [ diff --git a/harness/benchmarks/concurrency/write_feature_metadata.py b/harness/benchmarks/concurrency/write_feature_metadata.py index 515cbcea1..b2005ea25 100644 --- a/harness/benchmarks/concurrency/write_feature_metadata.py +++ b/harness/benchmarks/concurrency/write_feature_metadata.py @@ -92,6 +92,7 @@ def main() -> int: parser.add_argument("--draft-model-sha256") parser.add_argument("--decode-mode", choices=("ar", "speculation", "adaptive")) parser.add_argument("--draft-always", choices=("on", "off")) + parser.add_argument("--confidence", choices=("on", "off")) parser.add_argument("--ddtree", action="store_true") parser.add_argument("--fast-rollback", action="store_true") parser.add_argument("--ddtree-budget", type=int) @@ -183,6 +184,7 @@ def main() -> int: "draft_model_sha256": draft_model_sha256, "decode_mode": args.decode_mode, "draft_always": args.draft_always, + "confidence": args.confidence, "ddtree": args.ddtree, "fast_rollback": args.fast_rollback, "ddtree_budget": args.ddtree_budget, diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h index a7b443f91..fbf5caf46 100644 --- a/server/src/common/concurrency/speculation_gate.h +++ b/server/src/common/concurrency/speculation_gate.h @@ -81,8 +81,10 @@ struct SpecCandidate { SpeculationPolicy policy = SpeculationPolicy::Adaptive; bool eligible = false; int generated_tokens = 0; - // NaN means no calibrated confidence is available. Otherwise this is - // already the survival-product expected yield, including the root. + // NaN means no current-block confidence is available. Otherwise this is + // the survival-product expected yield, including the root. Speculator + // adapters provide per-position probability-like scores to + // confidence_survival_yield(); the gate owns calibration and clamping. double confidence_yield = std::numeric_limits::quiet_NaN(); }; @@ -106,10 +108,26 @@ struct SpecStepGeometry { } }; +enum class SpecScoreSource : uint8_t { + Confidence, + Measured, + Prior, +}; + +inline const char * spec_score_source_name(SpecScoreSource source) { + switch (source) { + case SpecScoreSource::Confidence: return "confidence"; + case SpecScoreSource::Measured: return "measured"; + case SpecScoreSource::Prior: return "prior"; + } + return "unknown"; +} + struct SpecPlanScore { uint64_t request_id = 0; int slot = -1; double expected_yield = 1.0; + SpecScoreSource source = SpecScoreSource::Prior; bool forced = false; bool admitted = false; }; @@ -132,6 +150,12 @@ struct SpecPlan { std::vector admitted_slots; }; +// Generic confidence contract for every chain speculator: `confidences[i]` +// is a probability-like, monotone-in-acceptance score for accepting position +// i of the current block conditioned on the preceding positions. The gate +// converts that adapter-owned vector into an expected emitted-token yield. It +// does not depend on how a producer obtains the scores (trained head, selector +// softmax, or a future speculator-specific readout). inline double confidence_survival_yield( const std::vector & confidences, int max_accept) { if (max_accept <= 1) return 1.0; @@ -148,6 +172,14 @@ inline double confidence_survival_yield( } class SpeculationGate { +private: + struct CandidateScore { + double expected_yield = 1.0; + double uncalibrated_confidence = + std::numeric_limits::quiet_NaN(); + SpecScoreSource source = SpecScoreSource::Prior; + }; + public: using ClampLogger = std::function; @@ -180,9 +212,18 @@ class SpeculationGate { return out; } + // A plan is consumed synchronously by the engine. Drop any abandoned + // prediction from a failed prior execution before recording this one. + for (const SpecCandidate & candidate : candidates) { + pending_confidence_.erase(candidate.request_id); + } + struct Ranked { const SpecCandidate * candidate = nullptr; double score = 1.0; + double uncalibrated_confidence = + std::numeric_limits::quiet_NaN(); + SpecScoreSource source = SpecScoreSource::Prior; bool forced = false; }; std::vector forced; @@ -195,8 +236,9 @@ class SpeculationGate { candidate.policy == SpeculationPolicy::Never) { continue; } - const double score = score_candidate(candidate); - Ranked ranked{&candidate, score, + const CandidateScore score = score_candidate(candidate); + Ranked ranked{&candidate, score.expected_yield, + score.uncalibrated_confidence, score.source, candidate.policy == SpeculationPolicy::Always}; (ranked.forced ? forced : adaptive).push_back(ranked); } @@ -223,7 +265,8 @@ class SpeculationGate { for (const Ranked & item : ranked) { out.ordered.push_back({item.candidate->request_id, item.candidate->slot, - item.score, item.forced, false}); + item.score, item.source, + item.forced, false}); } const int forced_count = static_cast(forced.size()); @@ -299,15 +342,22 @@ class SpeculationGate { out.ordered[(size_t)i].admitted = true; out.admitted_request_ids.push_back(ranked[(size_t)i].candidate->request_id); out.admitted_slots.push_back(ranked[(size_t)i].candidate->slot); + if (ranked[(size_t)i].source == SpecScoreSource::Confidence) { + pending_confidence_[ranked[(size_t)i].candidate->request_id] = + ranked[(size_t)i].uncalibrated_confidence; + } } return out; } void observe(uint64_t request_id, double emitted_tokens, int generated_tokens) { + auto pending = pending_confidence_.find(request_id); if (!std::isfinite(emitted_tokens) || emitted_tokens < 1.0 || emitted_tokens > static_cast(max_accept_) || generated_tokens < 0) { + if (pending != pending_confidence_.end()) + pending_confidence_.erase(pending); return; } RequestState & state = states_[request_id]; @@ -317,9 +367,18 @@ class SpeculationGate { state.tokens_at_last_spec = generated_tokens; ++global_rounds_; prior_yield_ += (emitted_tokens - prior_yield_) / global_rounds_; + if (pending != pending_confidence_.end()) { + calibration_predicted_ += pending->second; + calibration_realized_ += emitted_tokens; + ++calibration_observations_; + pending_confidence_.erase(pending); + } } - void forget(uint64_t request_id) { states_.erase(request_id); } + void forget(uint64_t request_id) { + states_.erase(request_id); + pending_confidence_.erase(request_id); + } bool has_state(uint64_t request_id) const { return states_.find(request_id) != states_.end(); } @@ -332,6 +391,20 @@ class SpeculationGate { return it == states_.end() ? 0.0 : it->second.mean_yield; } double prior_yield() const { return prior_yield_; } + double calibration_scale() const { + // Keep the identity scale until observations carry at least one + // block's worth of predicted mass; the block length supplies the + // threshold, so calibration adds no policy tuning constant. + if (calibration_observations_ == 0 || + calibration_predicted_ < static_cast(max_accept_)) + return 1.0; + return std::clamp( + calibration_realized_ / calibration_predicted_, + kCalibrationScaleMin, kCalibrationScaleMax); + } + uint64_t calibration_observations() const { + return calibration_observations_; + } const SpecCostTables & costs() const { return costs_; } private: @@ -341,10 +414,17 @@ class SpeculationGate { int tokens_at_last_spec = 0; }; - double score_candidate(const SpecCandidate & candidate) { + CandidateScore score_candidate(const SpecCandidate & candidate) { if (std::isfinite(candidate.confidence_yield)) { - return std::clamp(candidate.confidence_yield, 1.0, - static_cast(max_accept_)); + const double raw = std::clamp( + candidate.confidence_yield, 1.0, + static_cast(max_accept_)); + return { + std::clamp(calibration_scale() * raw, 1.0, + static_cast(max_accept_)), + raw, + SpecScoreSource::Confidence, + }; } auto it = states_.find(candidate.request_id); if (it != states_.end() && it->second.rounds > 0) { @@ -353,12 +433,20 @@ class SpeculationGate { if (idle_tokens >= config_.stale_after_tokens) { it->second = RequestState{}; } else { - return std::clamp(it->second.mean_yield, 1.0, - static_cast(max_accept_)); + return { + std::clamp(it->second.mean_yield, 1.0, + static_cast(max_accept_)), + std::numeric_limits::quiet_NaN(), + SpecScoreSource::Measured, + }; } } - return std::clamp(prior_yield_, 1.0, - static_cast(max_accept_)); + return { + std::clamp(prior_yield_, 1.0, + static_cast(max_accept_)), + std::numeric_limits::quiet_NaN(), + SpecScoreSource::Prior, + }; } void report_clamp(const char * name, const SpecCostLookup & lookup, @@ -375,7 +463,13 @@ class SpeculationGate { int max_accept_ = 1; double prior_yield_ = 1.0; uint64_t global_rounds_ = 0; + static constexpr double kCalibrationScaleMin = 0.25; + static constexpr double kCalibrationScaleMax = 4.0; + double calibration_predicted_ = 0.0; + double calibration_realized_ = 0.0; + uint64_t calibration_observations_ = 0; std::unordered_map states_; + std::unordered_map pending_confidence_; ClampLogger clamp_logger_; }; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 146e66c64..35c0df47e 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -40,6 +40,42 @@ int decode_bucket_width(int live_count) { return 64; } +void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, + double measured_us) { + int confidence = 0; + int measured = 0; + int prior = 0; + for (const SpecPlanScore & score : plan.ordered) { + switch (score.source) { + case SpecScoreSource::Confidence: ++confidence; break; + case SpecScoreSource::Measured: ++measured; break; + case SpecScoreSource::Prior: ++prior; break; + } + } + + std::fprintf(stderr, "[spec-gate] C=%d k=%d scores=[", + plan.concurrency, plan.admitted_count); + for (size_t i = 0; i < plan.ordered.size(); ++i) { + const SpecPlanScore & score = plan.ordered[i]; + std::fprintf(stderr, "%s%llu:%.3f/%s%s", + i == 0 ? "" : ",", + (unsigned long long)score.request_id, + score.expected_yield, + spec_score_source_name(score.source), + score.admitted ? "*" : ""); + } + std::fprintf(stderr, + "] sources=confidence:%d,measured:%d,prior:%d calibration=%.3f " + "G(k)=%.6f G(0)=%.6f predicted=%.1fus", + confidence, measured, prior, calibration_scale, + plan.goodput, plan.ar_goodput, plan.predicted_cost); + if (std::isfinite(measured_us)) { + std::fprintf(stderr, " measured=%.1fus\n", measured_us); + } else { + std::fprintf(stderr, " measured=ar-path\n"); + } +} + } // namespace Qwen35SeqEngine::Qwen35SeqEngine( @@ -615,6 +651,11 @@ bool Qwen35SeqEngine::draft_always_enabled() const { return !value || std::atoi(value) != 0; } +bool Qwen35SeqEngine::confidence_scoring_enabled() const { + const char * value = std::getenv("DFLASH_SPEC_CONFIDENCE"); + return !value || std::atoi(value) != 0; +} + bool Qwen35SeqEngine::prepare_chain_drafts( const std::vector & inputs, const std::vector & selected) { @@ -2134,6 +2175,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (speculation_gate_) { std::vector candidates; candidates.reserve(inputs.size()); + const bool use_confidence = confidence_scoring_enabled(); int drafting_lanes = 0; for (const StepInput & in : inputs) { const Qwen35Slot & seq = slots_.slot(in.slot); @@ -2147,7 +2189,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { eligible && policy != SpeculationPolicy::Never ? 1 : 0; double confidence = std::numeric_limits::quiet_NaN(); - if (eligible && in.slot >= 0 && + if (use_confidence && eligible && in.slot >= 0 && in.slot < (int)last_survival_score_.size() && last_survival_generated_[(size_t)in.slot] == seq.generated_tokens()) { @@ -2198,27 +2240,19 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { std::chrono::duration( std::chrono::steady_clock::now() - chain_started).count(); if (have_gate_plan && spec_gate_debug_enabled()) { - std::fprintf(stderr, - "[spec-gate] C=%zu k=%d scores=[", - inputs.size(), gate_plan.admitted_count); - for (size_t i = 0; i < gate_plan.ordered.size(); ++i) { - const SpecPlanScore & score = gate_plan.ordered[i]; - std::fprintf(stderr, "%s%llu:%.3f%s", - i == 0 ? "" : ",", - (unsigned long long)score.request_id, - score.expected_yield, - score.admitted ? "*" : ""); - } - std::fprintf(stderr, - "] G(k)=%.6f G(0)=%.6f predicted=%.1fus " - "measured=%.1fus\n", - gate_plan.goodput, gate_plan.ar_goodput, - gate_plan.predicted_cost, measured_us); + log_spec_gate_plan( + gate_plan, speculation_gate_->calibration_scale(), + measured_us); } if (speculative) return std::move(*speculative); // Proposal setup failed before target/cache mutation. Preserve // service through the ordinary packed AR path this iteration. } + if (!any_admitted && have_gate_plan && spec_gate_debug_enabled()) { + log_spec_gate_plan( + gate_plan, speculation_gate_->calibration_scale(), + std::numeric_limits::quiet_NaN()); + } } const TargetWeights & w = b_.w_; StepGraph & sg = b_.sg_; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 73d971cd2..ec6082a16 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -146,6 +146,7 @@ class Qwen35SeqEngine final : public SeqEngine { const std::vector & selected); bool batched_drafting_enabled() const; bool draft_always_enabled() const; + bool confidence_scoring_enabled() 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); diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index 1276ba208..6de2b19b4 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -62,6 +62,8 @@ int main() { CHECK((plan.admitted_request_ids == std::vector{10})); CHECK(plan.ordered.size() == 2); CHECK(plan.ordered[0].admitted); + CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); + CHECK(std::string(spec_score_source_name(plan.ordered[0].source)) == "confidence"); CHECK(!plan.ordered[1].admitted); // Identical optimistic cold priors are worth trying at C=1 but not at @@ -115,9 +117,11 @@ int main() { plan = stale.plan(1, {candidate(200, 0, NAN, SpeculationPolicy::Adaptive, true, 6)}, 1); CHECK(plan.ordered[0].expected_yield == 1.0); + CHECK(plan.ordered[0].source == SpecScoreSource::Measured); plan = stale.plan(1, {candidate(200, 0, NAN, SpeculationPolicy::Adaptive, true, 69)}, 1); CHECK(std::abs(plan.ordered[0].expected_yield - 2.5) < 1e-12); + CHECK(plan.ordered[0].source == SpecScoreSource::Prior); CHECK(stale.rounds(200) == 0); // State follows request IDs, survives a temporary eligibility loss, and @@ -170,6 +174,62 @@ int main() { CHECK(plan.cost_lookup_clamped); CHECK(clamp_logs > 0); + // Confidence is calibrated globally from admitted confidence-scored + // rounds. A 2x-overconfident signal converges to scale 0.5 and produces + // the same cost-aware cut as the true-yield oracle. + SpecCostTables calibration_costs = constant_costs(4.0, 10.0, 10.0); + SpeculationGate calibrated({}, calibration_costs, geometry(), 4); + CHECK(calibrated.calibration_scale() == 1.0); + CHECK(calibrated.calibration_observations() == 0); + plan = calibrated.plan(1, {candidate(700, 0, 4.0)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); + calibrated.observe(700, 2.0, 1); + CHECK(calibrated.calibration_observations() == 1); + CHECK(std::abs(calibrated.calibration_scale() - 0.5) < 1e-12); + plan = calibrated.plan(1, {candidate(700, 0, 4.0)}, 1); + SpeculationGate true_yield({}, calibration_costs, geometry(), 4); + SpecPlan oracle = true_yield.plan(1, {candidate(700, 0, 2.0)}, 1); + CHECK(plan.admitted_count == oracle.admitted_count); + calibrated.forget(700); + CHECK(std::abs(calibrated.calibration_scale() - 0.5) < 1e-12); + + // Prior- and measured-scored observations update only the degraded-mode + // history. They never contaminate confidence calibration. + SpeculationGate isolated({}, constant_costs(1.0, 10.0, 1.0), + geometry(), 4); + isolated.observe(800, 3.0, 1); + CHECK(isolated.calibration_observations() == 0); + plan = isolated.plan(1, { + candidate(800, 0, NAN, SpeculationPolicy::Adaptive, true, 2)}, 1); + CHECK(plan.ordered[0].source == SpecScoreSource::Measured); + isolated.observe(800, 2.0, 3); + CHECK(isolated.calibration_observations() == 0); + CHECK(isolated.calibration_scale() == 1.0); + plan = isolated.plan(1, {candidate(801, 0, NAN)}, 1); + CHECK(plan.ordered[0].source == SpecScoreSource::Prior); + + // The global ratio is bounded and survives request churn. + SpeculationGate lower_bound({}, constant_costs(1.0, 1.0, 1.0), + geometry(), 16); + plan = lower_bound.plan(1, { + candidate(900, 0, 16.0, SpeculationPolicy::Always)}, 1); + CHECK(plan.admitted_count == 1); + lower_bound.observe(900, 1.0, 1); + CHECK(lower_bound.calibration_scale() == 0.25); + lower_bound.forget(900); + CHECK(lower_bound.calibration_scale() == 0.25); + SpeculationGate upper_bound({}, constant_costs(1.0, 1.0, 1.0), + geometry(), 16); + for (uint64_t request_id = 901; request_id < 905; ++request_id) { + plan = upper_bound.plan(1, { + candidate(request_id, 0, 4.0, SpeculationPolicy::Always)}, 1); + upper_bound.observe(request_id, 16.0, 1); + if (request_id == 901) + CHECK(upper_bound.calibration_scale() == 1.0); + } + CHECK(upper_bound.calibration_scale() == 4.0); + // M1 convergence endpoints. SpeculationGate pays({}, constant_costs(1.0, 10.0, 1.0), geometry(), 4); for (int step = 0; step < 3; ++step) { From 9afca0b8ccf4abb446e7e34defffa3b038f149a2 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 08:45:46 +0000 Subject: [PATCH 17/42] fix(concurrency): activate calibration after 32 rounds --- .../src/common/concurrency/speculation_gate.h | 15 ++++--- .../qwen35/concurrency/qwen35_seq_engine.cpp | 42 +++++++++++++++++-- server/test/test_speculation_gate.cpp | 33 +++++++++------ 3 files changed, 69 insertions(+), 21 deletions(-) diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h index fbf5caf46..583821536 100644 --- a/server/src/common/concurrency/speculation_gate.h +++ b/server/src/common/concurrency/speculation_gate.h @@ -142,6 +142,9 @@ struct SpecPlan { int draft_lanes = 0; double expected_tokens = 0.0; double predicted_cost = 0.0; + // Pre-scale confidence yield for the admitted confidence-scored lanes. + // This is the denominator contribution shown in per-step telemetry. + double calibration_predicted_tokens = 0.0; double goodput = 0.0; double ar_goodput = 0.0; bool cost_lookup_clamped = false; @@ -343,6 +346,8 @@ class SpeculationGate { out.admitted_request_ids.push_back(ranked[(size_t)i].candidate->request_id); out.admitted_slots.push_back(ranked[(size_t)i].candidate->slot); if (ranked[(size_t)i].source == SpecScoreSource::Confidence) { + out.calibration_predicted_tokens += + ranked[(size_t)i].uncalibrated_confidence; pending_confidence_[ranked[(size_t)i].candidate->request_id] = ranked[(size_t)i].uncalibrated_confidence; } @@ -392,11 +397,10 @@ class SpeculationGate { } double prior_yield() const { return prior_yield_; } double calibration_scale() const { - // Keep the identity scale until observations carry at least one - // block's worth of predicted mass; the block length supplies the - // threshold, so calibration adds no policy tuning constant. - if (calibration_observations_ == 0 || - calibration_predicted_ < static_cast(max_accept_)) + // Calibration activation and bounds are fixed protocol constants, + // not workload policy tunables. + if (calibration_observations_ < kCalibrationMinObservations || + calibration_predicted_ <= 0.0) return 1.0; return std::clamp( calibration_realized_ / calibration_predicted_, @@ -463,6 +467,7 @@ class SpeculationGate { int max_accept_ = 1; double prior_yield_ = 1.0; uint64_t global_rounds_ = 0; + static constexpr uint64_t kCalibrationMinObservations = 32; static constexpr double kCalibrationScaleMin = 0.25; static constexpr double kCalibrationScaleMax = 4.0; double calibration_predicted_ = 0.0; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 35c0df47e..93e592192 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -40,7 +40,29 @@ int decode_bucket_width(int live_count) { return 64; } +double confidence_realized_tokens( + const SpecPlan & plan, const SeqEngine::StepResult & result) { + double realized = 0.0; + for (const SpecPlanScore & score : plan.ordered) { + if (!score.admitted || + score.source != SpecScoreSource::Confidence) { + continue; + } + const auto output = std::find_if( + result.decode.begin(), result.decode.end(), + [&](const SeqEngine::DecodeOutput & item) { + return item.slot == score.slot; + }); + if (output == result.decode.end() || output->failed) + return std::numeric_limits::quiet_NaN(); + realized += 1.0 + static_cast(output->spec_accepted_tokens); + } + return realized; +} + void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, + uint64_t calibration_rounds, + double calibration_realized_tokens, double measured_us) { int confidence = 0; int measured = 0; @@ -65,9 +87,18 @@ void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, score.admitted ? "*" : ""); } std::fprintf(stderr, - "] sources=confidence:%d,measured:%d,prior:%d calibration=%.3f " - "G(k)=%.6f G(0)=%.6f predicted=%.1fus", + "] sources=confidence:%d,measured:%d,prior:%d " + "calibration=%.3f rounds=%llu calib_tokens=%.3f/", confidence, measured, prior, calibration_scale, + (unsigned long long)calibration_rounds, + plan.calibration_predicted_tokens); + if (std::isfinite(calibration_realized_tokens)) { + std::fprintf(stderr, "%.3f", calibration_realized_tokens); + } else { + std::fprintf(stderr, "n/a"); + } + std::fprintf(stderr, + " G(k)=%.6f G(0)=%.6f predicted_cost=%.1fus", plan.goodput, plan.ar_goodput, plan.predicted_cost); if (std::isfinite(measured_us)) { std::fprintf(stderr, " measured=%.1fus\n", measured_us); @@ -2240,9 +2271,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { std::chrono::duration( std::chrono::steady_clock::now() - chain_started).count(); if (have_gate_plan && spec_gate_debug_enabled()) { + const double realized_tokens = speculative + ? confidence_realized_tokens(gate_plan, *speculative) + : std::numeric_limits::quiet_NaN(); log_spec_gate_plan( gate_plan, speculation_gate_->calibration_scale(), - measured_us); + speculation_gate_->calibration_observations(), + realized_tokens, measured_us); } if (speculative) return std::move(*speculative); // Proposal setup failed before target/cache mutation. Preserve @@ -2251,6 +2286,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (!any_admitted && have_gate_plan && spec_gate_debug_enabled()) { log_spec_gate_plan( gate_plan, speculation_gate_->calibration_scale(), + speculation_gate_->calibration_observations(), 0.0, std::numeric_limits::quiet_NaN()); } } diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index 6de2b19b4..84574c1e1 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -181,11 +181,15 @@ int main() { SpeculationGate calibrated({}, calibration_costs, geometry(), 4); CHECK(calibrated.calibration_scale() == 1.0); CHECK(calibrated.calibration_observations() == 0); - plan = calibrated.plan(1, {candidate(700, 0, 4.0)}, 1); - CHECK(plan.admitted_count == 1); - CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); - calibrated.observe(700, 2.0, 1); - CHECK(calibrated.calibration_observations() == 1); + for (int round = 0; round < 32; ++round) { + plan = calibrated.plan(1, {candidate(700, 0, 4.0)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); + CHECK(plan.calibration_predicted_tokens == 4.0); + calibrated.observe(700, 2.0, round + 1); + if (round < 31) CHECK(calibrated.calibration_scale() == 1.0); + } + CHECK(calibrated.calibration_observations() == 32); CHECK(std::abs(calibrated.calibration_scale() - 0.5) < 1e-12); plan = calibrated.plan(1, {candidate(700, 0, 4.0)}, 1); SpeculationGate true_yield({}, calibration_costs, geometry(), 4); @@ -212,21 +216,24 @@ int main() { // The global ratio is bounded and survives request churn. SpeculationGate lower_bound({}, constant_costs(1.0, 1.0, 1.0), geometry(), 16); - plan = lower_bound.plan(1, { - candidate(900, 0, 16.0, SpeculationPolicy::Always)}, 1); - CHECK(plan.admitted_count == 1); - lower_bound.observe(900, 1.0, 1); + for (int round = 0; round < 32; ++round) { + plan = lower_bound.plan(1, { + candidate(900, 0, 16.0, SpeculationPolicy::Always)}, 1); + CHECK(plan.admitted_count == 1); + lower_bound.observe(900, 1.0, round + 1); + if (round < 31) CHECK(lower_bound.calibration_scale() == 1.0); + } CHECK(lower_bound.calibration_scale() == 0.25); lower_bound.forget(900); CHECK(lower_bound.calibration_scale() == 0.25); SpeculationGate upper_bound({}, constant_costs(1.0, 1.0, 1.0), geometry(), 16); - for (uint64_t request_id = 901; request_id < 905; ++request_id) { + for (int round = 0; round < 32; ++round) { + const uint64_t request_id = 901 + static_cast(round); plan = upper_bound.plan(1, { candidate(request_id, 0, 4.0, SpeculationPolicy::Always)}, 1); - upper_bound.observe(request_id, 16.0, 1); - if (request_id == 901) - CHECK(upper_bound.calibration_scale() == 1.0); + upper_bound.observe(request_id, 16.0, round + 1); + if (round < 31) CHECK(upper_bound.calibration_scale() == 1.0); } CHECK(upper_bound.calibration_scale() == 4.0); From 50d4d5140b13b7e091b6c5201c152067bba10142 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 09:06:06 +0000 Subject: [PATCH 18/42] refactor(concurrency): make adaptive gate confidence-only --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 18 ++- .../concurrency/run_qwen38_dspark_matrix.sh | 14 +- .../concurrency/summarize_feature_matrix.py | 4 +- .../concurrency/test_feature_tools.py | 9 +- .../src/common/concurrency/speculation_gate.h | 122 +++++---------- .../qwen35/concurrency/qwen35_seq_engine.cpp | 38 ++--- server/test/test_speculation_gate.cpp | 140 ++++++------------ 7 files changed, 129 insertions(+), 216 deletions(-) diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 9f819fa67..e33e719c6 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -19,10 +19,12 @@ harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh The default fresh-process matrix is: -- `ar`, `speculation`, `adaptive-on`, `adaptive-off`, and - `adaptive-confidence-off`. The last arm keeps always-drafting enabled but - sets `DFLASH_SPEC_CONFIDENCE=0`, forcing the gate onto its measured/prior - fallback without changing the cost-aware top-k selector. +- `ar`, `speculation`, `adaptive-on`, and `adaptive-confidence-off`. The + last arm keeps always-drafting enabled but sets + `DFLASH_SPEC_CONFIDENCE=0`; with no historical fallback, adaptive requests + remain on AR while still paying drafting cost. This directly prices the + confidence signal. `ADAPTIVE_DRAFT_ALWAYS=on,off` adds the old admitted-only + drafting arm as an optional diagnostic, not a production acceptance path. - Live concurrency `C ∈ {1,2,3,4,6,8}` over the checked-in HumanEval and GSM8K cohorts plus deterministic prose prompts. - A fixed C=6 north-star cohort with two code and four chat requests. @@ -44,10 +46,10 @@ goodput oracle = max(ar, speculation) TTFT oracle = min(ar, speculation) ``` -Both confidence-enabled adaptive policies must reach at least 0.995 of that -oracle for mean and median output goodput and inverse TTFT. The summary fails -the run if any of those four ratios misses the gate; p95 alone is never used as -acceptance evidence. A second table reports paired goodput and inverse-TTFT +`adaptive-on` must reach at least 0.995 of that oracle for mean and median +output goodput and inverse TTFT. The summary fails the run if any of those four +ratios misses the gate; p95 alone is never used as acceptance evidence. A +second table reports paired goodput and inverse-TTFT deltas between `adaptive-on` and `adaptive-confidence-off` without gating the ablation. Set `WORKLOADS`, `CLIENTS`, `DECODE_MODES`, `ADAPTIVE_DRAFT_ALWAYS`, or `CONFIDENCE_ABLATION=0` to select a smaller diff --git a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh index aaefe67e5..745ecd36d 100755 --- a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh @@ -18,7 +18,7 @@ OUT="${OUT:-$REPO/.harness-runs/qwen38-dspark-matrix-$(date -u +%Y%m%dT%H%M%SZ)} REPEATS="${REPEATS:-1}" WORKLOADS="${WORKLOADS:-humaneval,gsm8k,prose,north-star}" DECODE_MODES="${DECODE_MODES:-ar,speculation,adaptive}" -ADAPTIVE_DRAFT_ALWAYS="${ADAPTIVE_DRAFT_ALWAYS:-on,off}" +ADAPTIVE_DRAFT_ALWAYS="${ADAPTIVE_DRAFT_ALWAYS:-on}" CONFIDENCE_ABLATION="${CONFIDENCE_ABLATION:-1}" CLIENTS="${CLIENTS:-1,2,3,4,6,8}" SLOTS="${SLOTS:-8}" @@ -47,12 +47,12 @@ The only supported drafter source for this matrix is: DRAFT_MODEL is the q4-mix requantized drafter produced from that repository. The default fresh-process matrix runs ar, forced speculation, adaptive with -always-drafting on, adaptive with always-drafting off, and an always-drafting -confidence-off ablation at live concurrency 1,2,3,4,6,8 over HumanEval, GSM8K, -and prose. The 2-code+4-chat north-star row runs only at C=6. The summary fails -if either confidence-enabled adaptive policy is below 0.995 of the paired -ar/speculation oracle in mean or median goodput or TTFT. The confidence -ablation delta is reported but not gated. +always-drafting on, and an always-drafting confidence-off ablation at live +concurrency 1,2,3,4,6,8 over HumanEval, GSM8K, and prose. The 2-code+4-chat +north-star row runs only at C=6. The summary fails if adaptive-on is below +0.995 of the paired ar/speculation oracle in mean or median goodput or TTFT. +The confidence ablation delta is reported but not gated. Set +ADAPTIVE_DRAFT_ALWAYS=on,off only for an additional diagnostic arm. Defaults select the second visible host GPU and address it as hip:0 inside the process (VISIBLE_DEVICES=1, TARGET_DEVICE=hip:0, DRAFT_DEVICE=hip:0). Override diff --git a/harness/benchmarks/concurrency/summarize_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py index 8afb29bec..0c520ee24 100755 --- a/harness/benchmarks/concurrency/summarize_feature_matrix.py +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -455,7 +455,7 @@ def summarize_dspark(reports: list[dict]) -> str: goodput_ratio = f"{gp_ratios[0]:.3f}/{gp_ratios[1]:.3f}" ttft_ratio = f"{ttft_ratios[0]:.3f}/{ttft_ratios[1]:.3f}" if ( - variant != "adaptive-confidence-off" + variant == "adaptive-on" and min(*gp_ratios, *ttft_ratios) < ORACLE_THRESHOLD ): regressions.append( @@ -537,7 +537,7 @@ def summarize_dspark(reports: list[dict]) -> str: ) lines += [ "", - f"Oracle-relative gate: every confidence-enabled adaptive mean/median " + f"Oracle-relative gate: adaptive-on mean/median " f"goodput and inverse TTFT ratio must be >= {ORACLE_THRESHOLD:.3f}.", ] if regressions: diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index 297631256..a99562fa5 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -185,6 +185,9 @@ def test_dspark_runner_has_explicit_confidence_ablation(self) -> None: encoding="utf-8", ) self.assertIn('variants+=("adaptive-confidence-off")', runner) + self.assertIn( + 'ADAPTIVE_DRAFT_ALWAYS="${ADAPTIVE_DRAFT_ALWAYS:-on}"', runner, + ) self.assertIn('"DFLASH_SPEC_CONFIDENCE=1"', runner) self.assertIn('"DFLASH_SPEC_CONFIDENCE=0"', runner) self.assertIn('metadata+=(--confidence "$confidence")', runner) @@ -742,8 +745,10 @@ def test_dspark_summary_enforces_mean_median_oracle_gate(self) -> None: self.dspark_item("speculation", 103.0, 1.0, repeat=2), self.dspark_item("adaptive-on", 100.0, 0.9, repeat=1), self.dspark_item("adaptive-on", 103.0, 1.0, repeat=2), - self.dspark_item("adaptive-off", 100.0, 0.9, repeat=1), - self.dspark_item("adaptive-off", 102.6, 1.0, repeat=2), + # Admitted-only drafting has no confidence bootstrap now and is + # diagnostic rather than acceptance-gated. + self.dspark_item("adaptive-off", 70.0, 2.0, repeat=1), + self.dspark_item("adaptive-off", 72.0, 1.9, repeat=2), self.dspark_item("adaptive-confidence-off", 75.0, 1.8, repeat=1), self.dspark_item("adaptive-confidence-off", 80.0, 1.7, repeat=2), ] diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h index 583821536..e2ac6b6df 100644 --- a/server/src/common/concurrency/speculation_gate.h +++ b/server/src/common/concurrency/speculation_gate.h @@ -17,10 +17,6 @@ namespace dflash::common { -struct SpecGateConfig { - int stale_after_tokens = 64; -}; - struct SpecCostLookup { double cost = std::numeric_limits::infinity(); int requested_index = 0; @@ -80,11 +76,11 @@ struct SpecCandidate { int slot = -1; SpeculationPolicy policy = SpeculationPolicy::Adaptive; bool eligible = false; - int generated_tokens = 0; - // NaN means no current-block confidence is available. Otherwise this is - // the survival-product expected yield, including the root. Speculator - // adapters provide per-position probability-like scores to - // confidence_survival_yield(); the gate owns calibration and clamping. + // NaN means no current-block confidence is available; an adaptive + // candidate then remains on AR. Otherwise this is the survival-product + // expected yield, including the root. Speculator adapters provide + // per-position probability-like scores to confidence_survival_yield(); + // the gate owns calibration and clamping. double confidence_yield = std::numeric_limits::quiet_NaN(); }; @@ -110,15 +106,13 @@ struct SpecStepGeometry { enum class SpecScoreSource : uint8_t { Confidence, - Measured, - Prior, + Unavailable, }; inline const char * spec_score_source_name(SpecScoreSource source) { switch (source) { case SpecScoreSource::Confidence: return "confidence"; - case SpecScoreSource::Measured: return "measured"; - case SpecScoreSource::Prior: return "prior"; + case SpecScoreSource::Unavailable: return "unavailable"; } return "unknown"; } @@ -127,7 +121,7 @@ struct SpecPlanScore { uint64_t request_id = 0; int slot = -1; double expected_yield = 1.0; - SpecScoreSource source = SpecScoreSource::Prior; + SpecScoreSource source = SpecScoreSource::Unavailable; bool forced = false; bool admitted = false; }; @@ -147,6 +141,7 @@ struct SpecPlan { double calibration_predicted_tokens = 0.0; double goodput = 0.0; double ar_goodput = 0.0; + int unavailable_count = 0; bool cost_lookup_clamped = false; std::vector ordered; std::vector admitted_request_ids; @@ -180,25 +175,21 @@ class SpeculationGate { double expected_yield = 1.0; double uncalibrated_confidence = std::numeric_limits::quiet_NaN(); - SpecScoreSource source = SpecScoreSource::Prior; + SpecScoreSource source = SpecScoreSource::Unavailable; }; public: using ClampLogger = std::function; - SpeculationGate(SpecGateConfig config, SpecCostTables costs, - SpecStepGeometry geometry, int max_accept, - ClampLogger clamp_logger = {}) - : config_(config), costs_(std::move(costs)), - geometry_(std::move(geometry)), + SpeculationGate(SpecCostTables costs, SpecStepGeometry geometry, + int max_accept, ClampLogger clamp_logger = {}) + : costs_(std::move(costs)), geometry_(std::move(geometry)), max_accept_(std::max(1, max_accept)), - prior_yield_(std::max(1, max_accept)), clamp_logger_(std::move(clamp_logger)) {} bool valid() const { - return config_.stale_after_tokens >= 1 && costs_.valid() && - geometry_.tree_width >= 1 && max_accept_ >= 1; + return costs_.valid() && geometry_.tree_width >= 1 && max_accept_ >= 1; } // draft_lanes_override prices always-drafting. -1 means admitted-only. @@ -216,7 +207,7 @@ class SpeculationGate { } // A plan is consumed synchronously by the engine. Drop any abandoned - // prediction from a failed prior execution before recording this one. + // prediction from a failed previous execution before recording this one. for (const SpecCandidate & candidate : candidates) { pending_confidence_.erase(candidate.request_id); } @@ -226,7 +217,7 @@ class SpeculationGate { double score = 1.0; double uncalibrated_confidence = std::numeric_limits::quiet_NaN(); - SpecScoreSource source = SpecScoreSource::Prior; + SpecScoreSource source = SpecScoreSource::Unavailable; bool forced = false; }; std::vector forced; @@ -240,6 +231,11 @@ class SpeculationGate { continue; } const CandidateScore score = score_candidate(candidate); + if (score.source == SpecScoreSource::Unavailable && + candidate.policy != SpeculationPolicy::Always) { + ++out.unavailable_count; + continue; + } Ranked ranked{&candidate, score.expected_yield, score.uncalibrated_confidence, score.source, candidate.policy == SpeculationPolicy::Always}; @@ -355,47 +351,21 @@ class SpeculationGate { return out; } - void observe(uint64_t request_id, double emitted_tokens, - int generated_tokens) { + void observe(uint64_t request_id, double emitted_tokens) { auto pending = pending_confidence_.find(request_id); + if (pending == pending_confidence_.end()) return; if (!std::isfinite(emitted_tokens) || emitted_tokens < 1.0 || - emitted_tokens > static_cast(max_accept_) || - generated_tokens < 0) { - if (pending != pending_confidence_.end()) - pending_confidence_.erase(pending); - return; - } - RequestState & state = states_[request_id]; - ++state.rounds; - state.mean_yield += - (emitted_tokens - state.mean_yield) / state.rounds; - state.tokens_at_last_spec = generated_tokens; - ++global_rounds_; - prior_yield_ += (emitted_tokens - prior_yield_) / global_rounds_; - if (pending != pending_confidence_.end()) { - calibration_predicted_ += pending->second; - calibration_realized_ += emitted_tokens; - ++calibration_observations_; + emitted_tokens > static_cast(max_accept_)) { pending_confidence_.erase(pending); + return; } + calibration_predicted_ += pending->second; + calibration_realized_ += emitted_tokens; + ++calibration_observations_; + pending_confidence_.erase(pending); } - void forget(uint64_t request_id) { - states_.erase(request_id); - pending_confidence_.erase(request_id); - } - bool has_state(uint64_t request_id) const { - return states_.find(request_id) != states_.end(); - } - int rounds(uint64_t request_id) const { - auto it = states_.find(request_id); - return it == states_.end() ? 0 : it->second.rounds; - } - double mean_yield(uint64_t request_id) const { - auto it = states_.find(request_id); - return it == states_.end() ? 0.0 : it->second.mean_yield; - } - double prior_yield() const { return prior_yield_; } + void forget(uint64_t request_id) { pending_confidence_.erase(request_id); } double calibration_scale() const { // Calibration activation and bounds are fixed protocol constants, // not workload policy tunables. @@ -412,13 +382,7 @@ class SpeculationGate { const SpecCostTables & costs() const { return costs_; } private: - struct RequestState { - double mean_yield = 0.0; - int rounds = 0; - int tokens_at_last_spec = 0; - }; - - CandidateScore score_candidate(const SpecCandidate & candidate) { + CandidateScore score_candidate(const SpecCandidate & candidate) const { if (std::isfinite(candidate.confidence_yield)) { const double raw = std::clamp( candidate.confidence_yield, 1.0, @@ -430,26 +394,10 @@ class SpeculationGate { SpecScoreSource::Confidence, }; } - auto it = states_.find(candidate.request_id); - if (it != states_.end() && it->second.rounds > 0) { - const int idle_tokens = - candidate.generated_tokens - it->second.tokens_at_last_spec; - if (idle_tokens >= config_.stale_after_tokens) { - it->second = RequestState{}; - } else { - return { - std::clamp(it->second.mean_yield, 1.0, - static_cast(max_accept_)), - std::numeric_limits::quiet_NaN(), - SpecScoreSource::Measured, - }; - } - } return { - std::clamp(prior_yield_, 1.0, - static_cast(max_accept_)), + 1.0, std::numeric_limits::quiet_NaN(), - SpecScoreSource::Prior, + SpecScoreSource::Unavailable, }; } @@ -461,19 +409,15 @@ class SpeculationGate { clamp_logger_(name, lookup.requested_index, lookup.profiled_index); } - SpecGateConfig config_; SpecCostTables costs_; SpecStepGeometry geometry_; int max_accept_ = 1; - double prior_yield_ = 1.0; - uint64_t global_rounds_ = 0; static constexpr uint64_t kCalibrationMinObservations = 32; static constexpr double kCalibrationScaleMin = 0.25; static constexpr double kCalibrationScaleMax = 4.0; double calibration_predicted_ = 0.0; double calibration_realized_ = 0.0; uint64_t calibration_observations_ = 0; - std::unordered_map states_; std::unordered_map pending_confidence_; ClampLogger clamp_logger_; }; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 93e592192..6570a803a 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -65,13 +65,11 @@ void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, double calibration_realized_tokens, double measured_us) { int confidence = 0; - int measured = 0; - int prior = 0; + int unavailable = plan.unavailable_count; for (const SpecPlanScore & score : plan.ordered) { switch (score.source) { case SpecScoreSource::Confidence: ++confidence; break; - case SpecScoreSource::Measured: ++measured; break; - case SpecScoreSource::Prior: ++prior; break; + case SpecScoreSource::Unavailable: ++unavailable; break; } } @@ -87,9 +85,9 @@ void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, score.admitted ? "*" : ""); } std::fprintf(stderr, - "] sources=confidence:%d,measured:%d,prior:%d " + "] sources=confidence:%d,unavailable:%d " "calibration=%.3f rounds=%llu calib_tokens=%.3f/", - confidence, measured, prior, calibration_scale, + confidence, unavailable, calibration_scale, (unsigned long long)calibration_rounds, plan.calibration_predicted_tokens); if (std::isfinite(calibration_realized_tokens)) { @@ -615,7 +613,7 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { return chain_decode_bucket_width(lanes); }; speculation_gate_ = std::make_unique( - SpecGateConfig{}, tables, geometry, T, + tables, geometry, T, [](const char * table, int requested, int profiled) { std::fprintf(stderr, "[spec-gate] %s_cost index %d outside profile; clamped to %d\n", @@ -706,13 +704,21 @@ bool Qwen35SeqEngine::prepare_chain_drafts( std::vector noise((size_t)T, b_.w_.mask_token_id); std::vector noise_embed((size_t)hidden * T); + // A failed current draft must not leave the previous block's confidence + // looking current. Successful lanes publish a fresh score below. + for (size_t i = 0; i < inputs.size(); ++i) { + if (!selected[i]) continue; + const int slot = inputs[i].slot; + if (slot < 0 || slot >= (int)prepared_chain_drafts_.size()) continue; + prepared_chain_drafts_[(size_t)slot].valid = false; + last_survival_score_[(size_t)slot] = + std::numeric_limits::quiet_NaN(); + last_survival_generated_[(size_t)slot] = -1; + } + for (size_t i = 0; i < inputs.size(); ++i) { if (!selected[i]) continue; const StepInput & in = inputs[i]; - if (in.slot >= 0 && - in.slot < (int)prepared_chain_drafts_.size()) { - prepared_chain_drafts_[(size_t)in.slot].valid = false; - } if (!chain_spec_input_eligible(in)) return false; DraftKvState * state = ensure_slot_draft_kv(in.slot); DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); @@ -857,8 +863,8 @@ bool Qwen35SeqEngine::prepare_chain_drafts( } else if (!missing_confidence_warned) { missing_confidence_warned = true; std::fprintf(stderr, - "[spec-gate] calibrated confidence unavailable; " - "using measured yield/prior\n"); + "[spec-gate] current confidence unavailable; " + "adaptive request remains AR\n"); } last_survival_score_[(size_t)info.slot] = score; last_survival_generated_[(size_t)info.slot] = @@ -1465,8 +1471,7 @@ std::optional Qwen35SeqEngine::step_chain_spec( if (speculation_gate_) { speculation_gate_->observe( slots_.slot(proposal.slot).request_id, - (double)proposal.path.size(), - slots_.slot(proposal.slot).generated_tokens()); + (double)proposal.path.size()); } } else { ArLane & ar = @@ -2227,8 +2232,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { confidence = last_survival_score_[(size_t)in.slot]; } candidates.push_back({ - seq.request_id, in.slot, policy, eligible, - seq.generated_tokens(), confidence, + seq.request_id, in.slot, policy, eligible, confidence, }); } gate_plan = speculation_gate_->plan( diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index 84574c1e1..0739cd6ca 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -32,8 +32,8 @@ static SpecStepGeometry geometry() { static SpecCandidate candidate( uint64_t id, int slot, double confidence, SpeculationPolicy policy = SpeculationPolicy::Adaptive, - bool eligible = true, int generated = 0) { - return {id, slot, policy, eligible, generated, confidence}; + bool eligible = true) { + return {id, slot, policy, eligible, confidence}; } int main() { @@ -41,7 +41,7 @@ int main() { CHECK(confidence_survival_yield({2.0f, -1.0f}, 4) == 2.0); CHECK(confidence_survival_yield({}, 4) == 1.0); - SpeculationGate costly({}, constant_costs(100.0, 10.0, 100.0), + SpeculationGate costly(constant_costs(100.0, 10.0, 100.0), geometry(), 4); CHECK(costly.valid()); SpecPlan plan = costly.plan(2, { @@ -53,7 +53,7 @@ int main() { // A high-yield request pays for one speculative lane. Adding the // confidence-1 freeloader cannot improve the numerator and loses the // tie because the argmax is scanned from the smaller prefix. - SpeculationGate prefix({}, constant_costs(1.0, 10.0, 1.0), + SpeculationGate prefix(constant_costs(1.0, 10.0, 1.0), geometry(), 4); plan = prefix.plan(3, { candidate(10, 0, 4.0), candidate(11, 1, 1.0), @@ -66,21 +66,19 @@ int main() { CHECK(std::string(spec_score_source_name(plan.ordered[0].source)) == "confidence"); CHECK(!plan.ordered[1].admitted); - // Identical optimistic cold priors are worth trying at C=1 but not at - // C=8 when the profiled tree launch crosses an occupancy boundary. + // Adaptive admission is confidence-only. Missing confidence stays on AR + // and is counted explicitly instead of consulting historical state. SpecCostTables crossover = constant_costs(1.0, 4.0, 1.0); - for (size_t i = 0; i < crossover.tree_cost.indices.size(); ++i) { - if (crossover.tree_cost.indices[i] > 4) - crossover.tree_cost.costs[i] = 20.0; - } - SpeculationGate cold_low({}, crossover, geometry(), 4); - plan = cold_low.plan(1, {candidate(20, 0, NAN)}, 1); - CHECK(plan.admitted_count == 1); - SpeculationGate cold_high({}, crossover, geometry(), 4); - std::vector eight; - for (int i = 0; i < 8; ++i) eight.push_back(candidate(30 + i, i, NAN)); - plan = cold_high.plan(8, eight, 8); + SpeculationGate confidence_only(crossover, geometry(), 4); + plan = confidence_only.plan(1, {candidate(20, 0, NAN)}, 1); CHECK(plan.admitted_count == 0); + CHECK(plan.ordered.empty()); + CHECK(plan.unavailable_count == 1); + plan = confidence_only.plan(2, { + candidate(20, 0, NAN), candidate(21, 1, 4.0)}, 2); + CHECK(plan.admitted_count == 1); + CHECK((plan.admitted_request_ids == std::vector{21})); + CHECK(plan.unavailable_count == 1); // Always and Never partition before adaptive ordering. plan = costly.plan(3, { @@ -97,52 +95,12 @@ int main() { CHECK(!plan.valid); CHECK(!plan.error.empty()); - SpeculationGate stateful({64}, constant_costs(1.0, 10.0, 1.0), - geometry(), 4); - stateful.observe(100, 2.0, 5); - stateful.observe(100, 4.0, 9); - CHECK(stateful.rounds(100) == 2); - CHECK(std::abs(stateful.mean_yield(100) - 3.0) < 1e-12); - stateful.observe(100, 0.0, 10); - stateful.observe(100, 5.0, 10); - CHECK(stateful.rounds(100) == 2); - - // A stale non-speculating request resets from its low request-local mean - // to the now-higher deployment prior and re-enters through the argmax. - SpeculationGate stale({64}, constant_costs(1.0, 10.0, 1.0), - geometry(), 4); - stale.observe(200, 1.0, 5); - stale.observe(999, 4.0, 5); - CHECK(std::abs(stale.prior_yield() - 2.5) < 1e-12); - plan = stale.plan(1, {candidate(200, 0, NAN, - SpeculationPolicy::Adaptive, true, 6)}, 1); - CHECK(plan.ordered[0].expected_yield == 1.0); - CHECK(plan.ordered[0].source == SpecScoreSource::Measured); - plan = stale.plan(1, {candidate(200, 0, NAN, - SpeculationPolicy::Adaptive, true, 69)}, 1); - CHECK(std::abs(plan.ordered[0].expected_yield - 2.5) < 1e-12); - CHECK(plan.ordered[0].source == SpecScoreSource::Prior); - CHECK(stale.rounds(200) == 0); - - // State follows request IDs, survives a temporary eligibility loss, and - // is explicitly forgotten at retirement/slot reuse. - stale.observe(300, 3.0, 4); - plan = stale.plan(1, {candidate(300, 0, NAN, - SpeculationPolicy::Adaptive, false, 5)}, 1); - CHECK(plan.ordered.empty()); - CHECK(stale.has_state(300)); - stale.forget(300); - CHECK(!stale.has_state(300)); - plan = stale.plan(1, {candidate(301, 0, NAN)}, 1); - CHECK(plan.ordered[0].request_id == 301); - - // All-Never is a pure AR plan and does not allocate request state. - SpeculationGate never({}, constant_costs(1.0, 2.0, 1.0), geometry(), 4); + // All-Never is a pure AR plan. + SpeculationGate never(constant_costs(1.0, 2.0, 1.0), geometry(), 4); plan = never.plan(1, {candidate(400, 0, NAN, SpeculationPolicy::Never)}, 1); CHECK(plan.admitted_count == 0); CHECK(plan.ordered.empty()); - CHECK(!never.has_state(400)); // C=1, capacity zero, malformed shapes, and always-draft pricing. plan = prefix.plan(1, {candidate(500, 0, 4.0)}, 1); @@ -167,7 +125,7 @@ int main() { int clamp_logs = 0; SpecCostTables tiny{series(1, 1.0), series(1, 2.0), series(1, 1.0)}; - SpeculationGate clamped({}, tiny, geometry(), 4, + SpeculationGate clamped(tiny, geometry(), 4, [&](const char *, int, int) { ++clamp_logs; }); plan = clamped.plan(2, { candidate(1, 0, 4.0), candidate(2, 1, 4.0)}, 2); @@ -178,7 +136,7 @@ int main() { // rounds. A 2x-overconfident signal converges to scale 0.5 and produces // the same cost-aware cut as the true-yield oracle. SpecCostTables calibration_costs = constant_costs(4.0, 10.0, 10.0); - SpeculationGate calibrated({}, calibration_costs, geometry(), 4); + SpeculationGate calibrated(calibration_costs, geometry(), 4); CHECK(calibrated.calibration_scale() == 1.0); CHECK(calibrated.calibration_observations() == 0); for (int round = 0; round < 32; ++round) { @@ -186,73 +144,73 @@ int main() { CHECK(plan.admitted_count == 1); CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); CHECK(plan.calibration_predicted_tokens == 4.0); - calibrated.observe(700, 2.0, round + 1); + calibrated.observe(700, 2.0); if (round < 31) CHECK(calibrated.calibration_scale() == 1.0); } CHECK(calibrated.calibration_observations() == 32); CHECK(std::abs(calibrated.calibration_scale() - 0.5) < 1e-12); plan = calibrated.plan(1, {candidate(700, 0, 4.0)}, 1); - SpeculationGate true_yield({}, calibration_costs, geometry(), 4); + SpeculationGate true_yield(calibration_costs, geometry(), 4); SpecPlan oracle = true_yield.plan(1, {candidate(700, 0, 2.0)}, 1); CHECK(plan.admitted_count == oracle.admitted_count); calibrated.forget(700); CHECK(std::abs(calibrated.calibration_scale() - 0.5) < 1e-12); - // Prior- and measured-scored observations update only the degraded-mode - // history. They never contaminate confidence calibration. - SpeculationGate isolated({}, constant_costs(1.0, 10.0, 1.0), + // Observations without a stashed confidence prediction do nothing. + // Explicit forced speculation remains available but does not calibrate + // when its confidence is missing. + SpeculationGate isolated(constant_costs(1.0, 10.0, 1.0), geometry(), 4); - isolated.observe(800, 3.0, 1); + isolated.observe(800, 3.0); CHECK(isolated.calibration_observations() == 0); + plan = isolated.plan(1, {candidate(800, 0, NAN)}, 1); + CHECK(plan.admitted_count == 0); + CHECK(plan.ordered.empty()); + CHECK(plan.unavailable_count == 1); plan = isolated.plan(1, { - candidate(800, 0, NAN, SpeculationPolicy::Adaptive, true, 2)}, 1); - CHECK(plan.ordered[0].source == SpecScoreSource::Measured); - isolated.observe(800, 2.0, 3); + candidate(801, 0, NAN, SpeculationPolicy::Always)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(plan.ordered[0].source == SpecScoreSource::Unavailable); + isolated.observe(801, 2.0); CHECK(isolated.calibration_observations() == 0); CHECK(isolated.calibration_scale() == 1.0); - plan = isolated.plan(1, {candidate(801, 0, NAN)}, 1); - CHECK(plan.ordered[0].source == SpecScoreSource::Prior); // The global ratio is bounded and survives request churn. - SpeculationGate lower_bound({}, constant_costs(1.0, 1.0, 1.0), + SpeculationGate lower_bound(constant_costs(1.0, 1.0, 1.0), geometry(), 16); for (int round = 0; round < 32; ++round) { plan = lower_bound.plan(1, { candidate(900, 0, 16.0, SpeculationPolicy::Always)}, 1); CHECK(plan.admitted_count == 1); - lower_bound.observe(900, 1.0, round + 1); + lower_bound.observe(900, 1.0); if (round < 31) CHECK(lower_bound.calibration_scale() == 1.0); } CHECK(lower_bound.calibration_scale() == 0.25); lower_bound.forget(900); CHECK(lower_bound.calibration_scale() == 0.25); - SpeculationGate upper_bound({}, constant_costs(1.0, 1.0, 1.0), + SpeculationGate upper_bound(constant_costs(1.0, 1.0, 1.0), geometry(), 16); for (int round = 0; round < 32; ++round) { const uint64_t request_id = 901 + static_cast(round); plan = upper_bound.plan(1, { candidate(request_id, 0, 4.0, SpeculationPolicy::Always)}, 1); - upper_bound.observe(request_id, 16.0, round + 1); + upper_bound.observe(request_id, 16.0); if (round < 31) CHECK(upper_bound.calibration_scale() == 1.0); } CHECK(upper_bound.calibration_scale() == 4.0); - // M1 convergence endpoints. - SpeculationGate pays({}, constant_costs(1.0, 10.0, 1.0), geometry(), 4); - for (int step = 0; step < 3; ++step) { - plan = pays.plan(2, { - candidate(600, 0, NAN, SpeculationPolicy::Adaptive, true, step), - candidate(601, 1, NAN, SpeculationPolicy::Adaptive, true, step)}, 2); - CHECK(plan.admitted_count == 2); - pays.observe(600, 4.0, step + 1); - pays.observe(601, 4.0, step + 1); - } - SpeculationGate cannot({}, constant_costs(100.0, 10.0, 100.0), + // Cost-aware endpoints remain direct functions of current confidence. + SpeculationGate pays(constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = pays.plan(2, { + candidate(600, 0, 4.0), candidate(601, 1, 4.0)}, 2); + CHECK(plan.admitted_count == 2); + SpeculationGate cannot(constant_costs(100.0, 10.0, 100.0), geometry(), 4); - for (int step = 0; step < 3; ++step) { - plan = cannot.plan(8, eight, 8); - CHECK(plan.admitted_count == 0); - } + std::vector eight; + for (int i = 0; i < 8; ++i) + eight.push_back(candidate(30 + i, i, 4.0)); + plan = cannot.plan(8, eight, 8); + CHECK(plan.admitted_count == 0); std::printf("speculation gate tests passed: %d checks\n", g_checks); return 0; From 167682e7c5a7fcc86ff7e4ee436d8f4443c34e4c Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 11:00:50 +0000 Subject: [PATCH 19/42] bench(qwen38): codify PR 625 baseline --- harness/benchmarks/QWEN38_PR625_BASELINE.md | 79 +++++++++++ harness/benchmarks/prompts/qwen38_pr625.jsonl | 2 + server/scripts/convert_dflash_to_gguf.py | 6 + server/scripts/prepare_qwen38_pr625_models.sh | 54 ++++++++ .../scripts/validate_qwen38_pr625_models.py | 126 ++++++++++++++++++ 5 files changed, 267 insertions(+) create mode 100644 harness/benchmarks/QWEN38_PR625_BASELINE.md create mode 100644 harness/benchmarks/prompts/qwen38_pr625.jsonl create mode 100755 server/scripts/prepare_qwen38_pr625_models.sh create mode 100755 server/scripts/validate_qwen38_pr625_models.py diff --git a/harness/benchmarks/QWEN38_PR625_BASELINE.md b/harness/benchmarks/QWEN38_PR625_BASELINE.md new file mode 100644 index 000000000..521ff9be8 --- /dev/null +++ b/harness/benchmarks/QWEN38_PR625_BASELINE.md @@ -0,0 +1,79 @@ +# Qwen3.8-27B PR #625 baseline + +Establish this dense, single-request baseline before measuring PR #626's +concurrent paged path. The two paths intentionally do not share cache or +attention settings. + +## Models + +Sources: + +- target: `bartowski/Qwen3.8-27B-GGUF`, `Qwen3.8-27B-IQ4_XS.gguf` +- drafter: `RadixArk/Qwen3.8-27B-DSpark`, `model.safetensors` + +Prepare the permanent local pair with: + +```bash +TARGET_SOURCE=/path/Qwen3.8-27B-IQ4_XS.gguf \ +DRAFT_SOURCE=/path/RadixArk-Qwen3.8-27B-DSpark/model.safetensors \ +LLAMA_QUANTIZE=/path/llama-quantize \ +server/scripts/prepare_qwen38_pr625_models.sh +``` + +This produces and validates: + +- target: pure IQ4_XS body, Q5_K `output.weight`, Q6_K `attn_v` and + `ssm_out` +- drafter: no YaRN, Q8_0 by default, capture layers `4,16,28,40,52`, mask + token `248077`. PR #625 does not publish the DSpark precision; set + `DRAFT_SCHEME=q4-mix` only for an explicit ablation. + +## Build on Radeon AI PRO R9700 + +Use ROCm 7.2, `gfx1201`, Release, and HIP graphs: + +```bash +cmake -S server -B server/build-pr625-r9700 -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_HIP_COMPILER=/opt/rocm/llvm/bin/clang++ \ + -DDFLASH27B_GPU_BACKEND=hip \ + -DDFLASH27B_HIP_ARCHITECTURES=gfx1201 \ + -DDFLASH27B_HIP_SM80_EQUIV=ON \ + -DGGML_HIP_GRAPHS=ON \ + -DDFLASH27B_FA_ALL_QUANTS=ON \ + -DDFLASH27B_SERVER=ON -DDFLASH27B_TESTS=OFF +cmake --build server/build-pr625-r9700 -j +``` + +## Dense single-request launch + +Use `HIP_VISIBLE_DEVICES=0` when the R9700 is the first physical GPU. Inside +the process it remains `hip:0`. + +```bash +HIP_VISIBLE_DEVICES=0 \ +DFLASH_SINGLE_CHAIN_CHECKPOINT_F32=1 \ +DFLASH_FAST_ROLLBACK_THRESHOLD=1 \ +LUCE_Q8_MEMO=1 \ +DFLASH_KV_ROTATE=0 \ +server/build-pr625-r9700/dflash_server \ + models/.lucebox/qwen38-pr625/Qwen3.8-27B-PR625-IQ4_XS.gguf \ + --draft models/.lucebox/qwen38-pr625/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ + --target-device hip:0 --draft-device hip:0 \ + --fa-window 2048 --cache-type-k q8_0 --cache-type-v q8_0 \ + --max-ctx 8192 --prefix-cache-slots 0 --prefill-cache-slots 0 \ + --decode-mode speculation --host 127.0.0.1 --port 18140 +``` + +Do not pass `--paged-attention` or `--max-concurrency` for this baseline. For +the AR control, start a fresh process **without `--draft`** and use +`--decode-mode ar`; the dense backend otherwise has a loaded drafter and can +enter its original speculative loop. All target, cache, and attention settings +stay identical. Use greedy 300-token generations with +`harness/benchmarks/prompts/qwen38_pr625.jsonl`, and reject a measurement that +ends before the 300-token cap. + +PR #625 reported R9700 decode throughput of 34.3/34.4 tok/s for AR and +45.6/32.4 tok/s for DSpark on its code/prose prompts. Exact numeric parity +requires the original unpublished prompts; the structural check is that code +benefits while prose can remain below AR. diff --git a/harness/benchmarks/prompts/qwen38_pr625.jsonl b/harness/benchmarks/prompts/qwen38_pr625.jsonl new file mode 100644 index 000000000..635ec2eec --- /dev/null +++ b/harness/benchmarks/prompts/qwen38_pr625.jsonl @@ -0,0 +1,2 @@ +{"id":"code-long","suite":"code","prompt":"Write a single self-contained, production-quality Python module implementing an asynchronous bounded worker pool. Include type annotations, docstrings, graceful cancellation, backpressure, per-job timeouts, structured result objects, clean shutdown, an executable usage example, and comprehensive unittest tests. Return only Python code and make the module at least 250 lines long."} +{"id":"prose-long","suite":"prose","prompt":"Write a clear, self-contained technical essay of about 500 words on why reproducible benchmarks need immutable inputs and explicit hardware metadata. Include one concrete example and end with a concise conclusion."} diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index c4482f5f8..5ebc5cdef 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -466,6 +466,9 @@ def main(): help="optional Domino/DSpark aux-head .pt or DS4 MTP .safetensors file; defaults to dflash_aux_heads.pt next to the safetensors") ap.add_argument("--no-aux-heads", action="store_true", help="do not auto-embed Domino/DSpark aux-head tensors") + ap.add_argument("--no-yarn", action="store_true", + help="omit YaRN scaling metadata while retaining rope_theta; " + "this matches the PR #625 short-context Qwen3.8 DSpark artifact") args = ap.parse_args() if not args.safetensors.exists(): @@ -478,6 +481,9 @@ def main(): print(f"[info] {n_entries} tensor entries") a = load_arch(args.safetensors, header) + if args.no_yarn: + for key in ("yarn_factor", "yarn_orig_ctx", "yarn_beta_fast", "yarn_beta_slow"): + a.pop(key, None) writer = gguf.GGUFWriter(args.out_gguf, ARCH) diff --git a/server/scripts/prepare_qwen38_pr625_models.sh b/server/scripts/prepare_qwen38_pr625_models.sh new file mode 100755 index 000000000..839fe6028 --- /dev/null +++ b/server/scripts/prepare_qwen38_pr625_models.sh @@ -0,0 +1,54 @@ +#!/usr/bin/env bash +# Build and validate the exact Qwen3.8-27B model pair used by PR #625. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +REPO="$(cd -- "$SCRIPT_DIR/../.." && pwd -P)" + +TARGET_SOURCE="${TARGET_SOURCE:?set TARGET_SOURCE to bartowski Qwen3.8-27B IQ4_XS GGUF}" +DRAFT_SOURCE="${DRAFT_SOURCE:?set DRAFT_SOURCE to RadixArk Qwen3.8-27B-DSpark model.safetensors}" +LLAMA_QUANTIZE="${LLAMA_QUANTIZE:?set LLAMA_QUANTIZE to a llama-quantize binary}" +PYTHON="${PYTHON:-python3}" +OUT_DIR="${OUT_DIR:-$REPO/models/.lucebox/qwen38-pr625}" +DRAFT_SCHEME="${DRAFT_SCHEME:-q8_0}" + +[[ -r "$TARGET_SOURCE" ]] || { echo "unreadable TARGET_SOURCE: $TARGET_SOURCE" >&2; exit 2; } +[[ -r "$DRAFT_SOURCE" ]] || { echo "unreadable DRAFT_SOURCE: $DRAFT_SOURCE" >&2; exit 2; } +[[ -x "$LLAMA_QUANTIZE" ]] || { echo "LLAMA_QUANTIZE is not executable: $LLAMA_QUANTIZE" >&2; exit 2; } +command -v "$PYTHON" >/dev/null || { echo "PYTHON is unavailable: $PYTHON" >&2; exit 2; } +[[ "$DRAFT_SCHEME" == q8_0 || "$DRAFT_SCHEME" == q4-mix ]] || { + echo "DRAFT_SCHEME must be q8_0 or q4-mix" >&2 + exit 2 +} + +mkdir -p "$OUT_DIR" +work_dir="$(mktemp -d "$OUT_DIR/.prepare.XXXXXX")" +cleanup() { rm -rf -- "$work_dir"; } +trap cleanup EXIT + +draft_f16="$work_dir/Qwen3.8-27B-DSpark-RadixArk-no-yarn-f16.gguf" +draft_final="$work_dir/Qwen3.8-27B-DSpark-RadixArk-no-yarn-$DRAFT_SCHEME.gguf" +target_final="$work_dir/Qwen3.8-27B-PR625-IQ4_XS.gguf" + +"$PYTHON" "$SCRIPT_DIR/convert_dflash_to_gguf.py" \ + "$DRAFT_SOURCE" "$draft_f16" --no-yarn +"$PYTHON" "$SCRIPT_DIR/quantize_dflash_draft.py" \ + "$draft_f16" "$draft_final" --scheme "$DRAFT_SCHEME" + +# PR #625 target: pure IQ4_XS body, Q5_K output, Q6_K attn_v/ssm_out. +# The validator below deliberately catches quantizers that let --pure suppress +# explicit --tensor-type overrides. +"$LLAMA_QUANTIZE" \ + --allow-requantize --pure \ + --output-tensor-type q5_k \ + --tensor-type ssm_out=q6_k \ + --tensor-type attn_v=q6_k \ + "$TARGET_SOURCE" "$target_final" iq4_xs + +"$PYTHON" "$SCRIPT_DIR/validate_qwen38_pr625_models.py" \ + --target "$target_final" --draft "$draft_final" --draft-scheme "$DRAFT_SCHEME" + +mv -- "$target_final" "$OUT_DIR/Qwen3.8-27B-PR625-IQ4_XS.gguf" +mv -- "$draft_final" "$OUT_DIR/Qwen3.8-27B-DSpark-RadixArk-no-yarn-$DRAFT_SCHEME.gguf" + +echo "PR #625 model pair ready in $OUT_DIR" diff --git a/server/scripts/validate_qwen38_pr625_models.py b/server/scripts/validate_qwen38_pr625_models.py new file mode 100755 index 000000000..d3215f48d --- /dev/null +++ b/server/scripts/validate_qwen38_pr625_models.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +"""Fail closed unless target and drafter match the PR #625 Qwen3.8 recipe.""" + +import argparse +import sys +from collections import Counter +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "deps" / "llama.cpp" / "gguf-py")) + +from gguf import GGUFReader # noqa: E402 + + +DRAFT_ARCH = "qwen35-dflash-draft" + + +def require(condition: bool, message: str, errors: list[str]) -> None: + if not condition: + errors.append(message) + + +def field(reader: GGUFReader, name: str): + value = reader.fields.get(name) + return None if value is None else value.contents() + + +def validate_target(path: Path, errors: list[str]) -> None: + reader = GGUFReader(path) + tensors = {tensor.name: tensor for tensor in reader.tensors} + counts = Counter(tensor.tensor_type.name for tensor in reader.tensors) + + require(field(reader, "general.architecture") == "qwen35", + "target architecture must be qwen35", errors) + require(counts == Counter({"IQ4_XS": 440, "F32": 360, "Q6_K": 65, "Q5_K": 1}), + f"target tensor-type counts differ from PR #625: {dict(counts)}", errors) + require(tensors.get("output.weight") is not None and + tensors["output.weight"].tensor_type.name == "Q5_K", + "target output.weight must be Q5_K", errors) + + q6_names = { + tensor.name for tensor in reader.tensors if tensor.tensor_type.name == "Q6_K" + } + invalid_q6 = sorted( + name for name in q6_names + if not (name.endswith("ssm_out.weight") or name.endswith("attn_v.weight")) + ) + require(not invalid_q6, + f"target has unexpected Q6_K tensors: {invalid_q6}", errors) + require(all( + tensor.tensor_type.name == "Q6_K" + for name, tensor in tensors.items() + if name.endswith("ssm_out.weight") or name.endswith("attn_v.weight") + ), "every target ssm_out/attn_v tensor must be Q6_K", errors) + + +def validate_draft(path: Path, scheme: str, errors: list[str]) -> None: + reader = GGUFReader(path) + prefix = DRAFT_ARCH + "." + counts = Counter(tensor.tensor_type.name for tensor in reader.tensors) + + require(field(reader, "general.architecture") == DRAFT_ARCH, + f"drafter architecture must be {DRAFT_ARCH}", errors) + expected_counts = { + "f16": Counter({"F16": 39, "F32": 23}), + "q8_0": Counter({"Q8_0": 39, "F32": 23}), + "q4-mix": Counter({"Q4_0": 35, "F32": 23, "Q8_0": 4}), + }[scheme] + require(counts == expected_counts, + f"drafter tensor-type counts differ from {scheme}: {dict(counts)}", errors) + require(field(reader, prefix + "rope.freq_base") == 10_000_000.0, + "drafter rope.freq_base must be 10000000", errors) + require(not any("rope.scaling" in name for name in reader.fields), + "drafter must not contain YaRN/rope.scaling metadata", errors) + + expected = { + "dflash.n_target_layers": 5, + "dflash.block_size": 7, + "dflash.mask_token_id": 248077, + "dflash.target_layer_ids": [4, 16, 28, 40, 52], + "dflash.dspark.enabled": 1, + "dflash.dspark.markov_rank": 256, + "dflash.dspark.vocab_size": 248320, + "dflash.dspark.confidence_dim": 5376, + "dflash.dspark.confidence.enabled": 1, + } + for key, wanted in expected.items(): + actual = field(reader, prefix + key) + require(actual == wanted, + f"drafter {key} must be {wanted!r}, got {actual!r}", errors) + + if scheme == "q4-mix": + invalid_q8 = sorted( + tensor.name for tensor in reader.tensors + if tensor.tensor_type.name == "Q8_0" and not tensor.name.startswith("dflash.") + ) + require(not invalid_q8, + f"q4-mix has non-head Q8_0 tensors: {invalid_q8}", errors) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--target", required=True, type=Path) + parser.add_argument("--draft", required=True, type=Path) + parser.add_argument("--draft-scheme", choices=("f16", "q8_0", "q4-mix"), + default="q8_0") + args = parser.parse_args() + + errors: list[str] = [] + for label, path in (("target", args.target), ("draft", args.draft)): + if not path.is_file(): + errors.append(f"{label} is not a readable file: {path}") + if not errors: + validate_target(args.target, errors) + validate_draft(args.draft, args.draft_scheme, errors) + + if errors: + for error in errors: + print(f"error: {error}", file=sys.stderr) + return 1 + print(f"PR #625 target OK: {args.target}") + print(f"PR #625 no-YaRN {args.draft_scheme} drafter OK: {args.draft}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From c0fbe36ddae81e3f36d3d79bd81b5ef4bdec398d Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:25:50 +0200 Subject: [PATCH 20/42] qwen35: per-step verify length for DSpark confidence gate Verify/accept now run over v_len (the drafted chain's actual length) instead of the buffer-sizing q_len, so the DSpark confidence gate's adaptive block truncation is structurally supported. The gate itself stays off by default (DFLASH_QWEN35_DSPARK_CONFIDENCE_THRESHOLD=0): with the RadixArk Qwen3.8 drafter, any threshold in 0.1-0.5 truncates to the same short chain regardless of value, so the confidence scores coming out of the shared head path look mis-scaled and need a separate investigation before the gate can help. threshold=0 is bench-verified regression-free. --- server/src/qwen35/qwen35_backend.cpp | 44 ++++++++++++++++++++-------- 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index c195cdb40..a4df94f81 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -2472,6 +2472,22 @@ static bool qwen35_dspark_enabled() { return kEnabled; } +// Confidence-gate threshold for adaptive block length (0 = gate off, verify +// the full drafted block). The drafter's AcceptRatePredictor scores each +// draft position; the chain is truncated at the first position below the +// threshold and only the confident prefix is verified. +static float qwen35_dspark_confidence_threshold() { + static const float kThreshold = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK_CONFIDENCE_THRESHOLD"); + if (!e) return 0.0f; + float threshold = (float)std::atof(e); + if (threshold < 0.0f) threshold = 0.0f; + if (threshold > 1.0f) threshold = 1.0f; + return threshold; + }(); + return kThreshold; +} + bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io, @@ -2834,6 +2850,10 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, cfg_.ddtree_mode && target->supports_tree_verify() && kvflash_tree_ok && !use_remote_draft && q_len > 1 && tree_special_inactive; + // Chain-verify length for this step. The DSpark confidence gate may + // truncate the drafted block (adaptive block length); q_len stays the + // buffer-sizing upper bound. + int v_len = q_len; // DDTree consumes top-K rows directly. Avoid projecting the same // hidden block once for argmax and again for top-K on every step. if (!use_tree_verify) { @@ -2855,23 +2875,23 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return !(e && e[0] == '0' && e[1] == '\0'); }(); bool ds_ok = false; - if (fused_dspark) { + if (fused_dspark && qwen35_dspark_confidence_threshold() <= 0.0f) { ds_ok = dspark_markov_correct_greedy_chain_fused( dw_, draft_backend_, target->lm_head_tensor(), local_hidden.data(), q_len, last_tok, draft_tok); } if (!ds_ok) { - // threshold 0 = confidence gate off: q_len sizes the - // step buffers for the whole request, so the truncated - // chain the gate produces cannot be verified here yet. ds_ok = dspark_markov_correct_greedy_chain(dw_, draft_backend_, *target, local_hidden.data(), q_len, last_tok, - /*confidence_threshold=*/0.0f, + qwen35_dspark_confidence_threshold(), draft_tok); } if (ds_ok) { used_dspark = true; + // Confidence gate truncates the drafted chain: verify + // only the confident prefix this step. + v_len = std::max(1, (int)draft_tok.size()); } else { static std::atomic s_dspark_warned{false}; if (!s_dspark_warned.exchange(true)) { @@ -3179,7 +3199,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int hint_fill = 0; if (hint_tokens && n_generated < (int)hint_tokens->size()) { const int hint_avail = (int)hint_tokens->size() - n_generated; - hint_fill = std::min(hint_avail, q_len - 1); + hint_fill = std::min(hint_avail, v_len - 1); for (int i = 0; i < hint_fill; i++) { draft_tok[1 + i] = (*hint_tokens)[n_generated + i]; } @@ -3214,13 +3234,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int accept_n = 1; int bonus_tok = -1; if (sampled_verify) { - if (!target->read_verify_logits(q_len, verify_logits)) { + if (!target->read_verify_logits(v_len, verify_logits)) { std::fprintf(stderr, "spec-decode: verify logits read failed\n"); target->restore_kv(); step_graph_destroy(draft_sg); return false; } - const int vocab_v = (int)(verify_logits.size() / (size_t)q_len); + const int vocab_v = (int)(verify_logits.size() / (size_t)v_len); static const bool kSvDebug = []() { const char * e = std::getenv("DFLASH_SV_DEBUG"); return e != nullptr && std::string(e) == "1"; @@ -3229,7 +3249,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // Row-alignment check: CPU argmax over each bulk-read row must // equal the GPU argmax (target_tok). Divergence = misaligned // or stale bulk read. - for (int i = 0; i < q_len; i++) { + for (int i = 0; i < v_len; i++) { const float * row = verify_logits.data() + (size_t)i * vocab_v; int am = 0; float best = row[0]; for (int v = 1; v < vocab_v; v++) @@ -3254,7 +3274,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, verify_history = out_tokens; verify_history.push_back(draft_tok[0]); bool mismatched = false; - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < v_len - 1; i++) { const int s = sample_logits( verify_logits.data() + (size_t)i * vocab_v, vocab_v, sampler_, verify_history, sampler_rng_); @@ -3275,11 +3295,11 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } (void)mismatched; } else { - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < v_len - 1; i++) { if (draft_tok[i + 1] == target_tok[i]) accept_n++; else break; } - bonus_tok = (accept_n < q_len) ? target_tok[accept_n - 1] : -1; + bonus_tok = (accept_n < v_len) ? target_tok[accept_n - 1] : -1; } // Track hint acceptance telemetry. if (hint_fill > 0) { From 248b5017f9ae6b7be57f04e4b978b36844413226 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:10:08 +0200 Subject: [PATCH 21/42] ggml: fused DeltaNet decode kernels for HIP - ggml_ssm_conv_step: one kernel for the causal-conv decode/verify step (history window + silu(conv) + in-place history write-back + optional rollback window copy) replacing transpose/concat/ssm_conv/silu/cpy. - ggml_gated_delta_net_set_raw_gates: the GDN kernel applies sigmoid(beta) and softplus(alpha + dt_bias) * A itself. - ADD + RMS_NORM + MUL fusion (residual add materialized alongside the normalized output) in the CUDA/HIP graph evaluator. - legacy pool MAX_BUFFERS 256 -> 1024: LUCE_Q8_MEMO holds ~300 pooled buffers per evaluation; a full pool freed in-flight buffers with cudaFree and produced illegal memory accesses on long prefills. --- server/deps/llama.cpp/ggml/include/ggml.h | 32 +++++- .../deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp | 2 + .../ggml/src/ggml-cuda/gated_delta_net.cu | 86 +++++++++----- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 59 +++++++++- .../deps/llama.cpp/ggml/src/ggml-cuda/norm.cu | 81 +++++++++++++ .../llama.cpp/ggml/src/ggml-cuda/norm.cuh | 3 + .../llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu | 106 ++++++++++++++++++ server/deps/llama.cpp/ggml/src/ggml.c | 70 ++++++++++++ 8 files changed, 410 insertions(+), 29 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index beb6fc817..567e29094 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -221,7 +221,7 @@ #define GGML_MAX_DIMS 4 #define GGML_MAX_PARAMS 2048 -#define GGML_MAX_SRC 10 +#define GGML_MAX_SRC 12 #define GGML_MAX_N_THREADS 512 #define GGML_MAX_OP_PARAMS 64 @@ -2758,6 +2758,25 @@ extern "C" { struct ggml_tensor * c, struct ggml_tensor * parent_ids); + // dflash extension: fused causal-conv step for recurrent decode/verify. + // Replaces transpose + concat(state, x) + ssm_conv + silu + state + // write-back with one kernel. + // x: [C, T, S] f32, rows contiguous (token stride may be + // larger than C, e.g. a row-slice of a stacked GEMV) + // c: [K, C] f32 depthwise conv weights + // conv_state: [K-1, C, S] f32 history; READ, then OVERWRITTEN in + // place with the last K-1 conv inputs + // conv_input_out: optional [>= K-1+T, C, S] f32; receives the full + // conv window (history rows then x rows) per channel, + // for speculative-decode rollback. May be a view. + // Returns silu(conv(x)) as [C, T, S]. CUDA/HIP only. + GGML_API struct ggml_tensor * ggml_ssm_conv_step( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * c, + struct ggml_tensor * conv_state, + struct ggml_tensor * conv_input_out); + GGML_API struct ggml_tensor * ggml_ssm_scan( struct ggml_context * ctx, struct ggml_tensor * s, @@ -2904,6 +2923,17 @@ extern "C" { struct ggml_tensor * tensor, bool skip_intermediate); + // dflash extension: let the kernel derive the gates from the raw + // projections instead of graph-side sigmoid/softplus ops: + // beta_val = sigmoid(beta_raw) + // g_val = exp(softplus(alpha_raw + dt_bias[h]) * A[h]) + // `g` then carries alpha_raw and `beta` carries beta_raw (both [1,H,T,S]); + // dt_bias and A are [H] f32. Only for the non-tree, non-KDA CUDA/HIP path. + GGML_API void ggml_gated_delta_net_set_raw_gates( + struct ggml_tensor * tensor, + struct ggml_tensor * dt_bias, + struct ggml_tensor * A); + // dflash extension: tree-mode gated delta net for DDTree-style // speculative decoding verify. `parent_ids` is an int32 tensor of shape // [n_tokens, n_seqs] where entry [t, s] is the index within sequence s of diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp index b10e8c75d..2c05b85f2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp @@ -9332,6 +9332,8 @@ void ggml_compute_forward_flash_attn_back( static void ggml_compute_forward_ssm_conv_f32( const ggml_compute_params * params, ggml_tensor * dst) { + // dflash: the fused step mode (ggml_ssm_conv_step) is CUDA/HIP only + GGML_ASSERT(ggml_get_op_params_i32(dst, 0) == 0 && "ggml_ssm_conv_step is not supported on CPU"); const ggml_tensor * src0 = dst->src[0]; // conv_x const ggml_tensor * src1 = dst->src[1]; // conv1d.weight diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu index 76f2de9da..e9f416dd8 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu @@ -97,7 +97,9 @@ gated_delta_net_cuda(const float * q, int64_t sb3, const uint3 neqk1_magic, const uint3 rq3_magic, - float scale) { + float scale, + const float * gate_bias, // raw-gate mode: dt_bias[H], else nullptr + const float * gate_A) { // raw-gate mode: A[H] const uint32_t h_idx = blockIdx.x; const uint32_t sequence = blockIdx.y; // each warp owns one column, using warp-level primitives to reduce across rows @@ -196,7 +198,9 @@ gated_delta_net_cuda(const float * q, const float * beta_t = beta + gb_offset; const float * g_t = g + gb_offset * (KDA ? S_v : 1); - const float beta_val = *beta_t; + // raw-gate mode: beta = sigmoid(beta_raw); g = softplus(alpha_raw + bias) * A + const bool raw_gates = gate_bias != nullptr; + const float beta_val = raw_gates ? 1.0f / (1.0f + expf(-(*beta_t))) : *beta_t; // Cache k and q in registers float k_reg[rows_per_lane]; @@ -209,7 +213,12 @@ gated_delta_net_cuda(const float * q, } if constexpr (!KDA) { - const float g_val = expf(*g_t); + float g_log = *g_t; + if (raw_gates) { + const float a = g_log + gate_bias[h_idx]; + g_log = ((a > 20.0f) ? a : logf(1.0f + expf(a))) * gate_A[h_idx]; + } + const float g_val = expf(g_log); // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] float kv_shard = 0.0f; @@ -318,7 +327,9 @@ gated_delta_net_cuda_grouped_cols(const float * q, int64_t sb3, const uint3 neqk1_magic, const uint3 rq3_magic, - float scale) { + float scale, + const float * gate_bias, // raw-gate mode: dt_bias[H], else nullptr + const float * gate_A) { // raw-gate mode: A[H] static_assert(S_v == 128, "grouped GDN kernel is specialized for S_v=128"); static_assert(WIDTH == 16, "grouped GDN kernel expects 16-lane subgroups"); static_assert(COLS == 4, "grouped GDN kernel expects 4 columns per subgroup"); @@ -387,8 +398,16 @@ gated_delta_net_cuda_grouped_cols(const float * q, float g_val = 0.0f; float beta_val = 0.0f; if (threadIdx.x == 0) { - g_val = expf(g[gb_offset]); - beta_val = beta[gb_offset]; + if (gate_bias != nullptr) { + // raw-gate mode: g = exp(softplus(alpha_raw + bias) * A), beta = sigmoid(beta_raw) + const float a = g[gb_offset] + gate_bias[h_idx]; + const float sp = (a > 20.0f) ? a : logf(1.0f + expf(a)); + g_val = expf(sp * gate_A[h_idx]); + beta_val = 1.0f / (1.0f + expf(-beta[gb_offset])); + } else { + g_val = expf(g[gb_offset]); + beta_val = beta[gb_offset]; + } } g_val = __shfl_sync(0xffffffffU, g_val, 0); beta_val = __shfl_sync(0xffffffffU, beta_val, 0); @@ -497,7 +516,8 @@ static void launch_gated_delta_net( int64_t sv1, int64_t sv2, int64_t sv3, int64_t sb1, int64_t sb2, int64_t sb3, int64_t neqk1, int64_t rq3, - float scale, cudaStream_t stream) { + float scale, cudaStream_t stream, + const float * gate_bias = nullptr, const float * gate_A = nullptr) { //TODO: Add chunked kernel for even faster pre-fill const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; const int num_warps = 4; @@ -521,19 +541,19 @@ static void launch_gated_delta_net( gated_delta_net_cuda<16, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; case 32: gated_delta_net_cuda<32, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; case 64: { gated_delta_net_cuda<64, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; } case 128: { @@ -552,7 +572,7 @@ static void launch_gated_delta_net( gated_delta_net_cuda_grouped_cols<128, cols, width, 32, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } else if (warp_size == 64) { constexpr int groups_per_warp = 64 / width; dim3 grouped_grid_dims(H, n_seqs, (groups + column_groups_per_block * groups_per_warp - 1) / (column_groups_per_block * groups_per_warp)); @@ -560,24 +580,24 @@ static void launch_gated_delta_net( gated_delta_net_cuda_grouped_cols<128, cols, width, 64, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, - sb1, sb2, sb3, neqk1_magic, rq3_magic, scale); + sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } break; } @@ -693,6 +713,18 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * const bool tree_mode = (parent_ids_d != nullptr); const bool skip_intermediate = ggml_get_op_params_i32(dst, 0) != 0; + // dflash raw-gate mode: src[9] = dt_bias[H], src[10] = A[H]; src[8] + // remains available for the optional active-slot map. The kernel + // applies sigmoid / softplus+bias / A itself (see ggml_gated_delta_net_set_raw_gates). + const bool raw_gates = ggml_get_op_params_i32(dst, 2) != 0; + const float * gate_bias_d = nullptr; + const float * gate_A_d = nullptr; + if (raw_gates) { + GGML_ASSERT(dst->src[9] && dst->src[10]); + GGML_ASSERT(!kda && !tree_mode); + gate_bias_d = (const float *) dst->src[9]->data; + gate_A_d = (const float *) dst->src[10]->data; + } const bool write_intermediate = tree_mode || !skip_intermediate || persist_inter_d != nullptr; // Macro to expand KDA × TREE_MODE × WRITE_INTER for a given InterT. @@ -704,35 +736,35 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * if (tree_mode) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } else if (write_intermediate) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } else { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } \ } else { \ if (tree_mode) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } else if (write_intermediate) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } else { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ - S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream); \ + S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ } \ } \ } while (0) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 1c55d51d1..543417701 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -473,7 +473,11 @@ const ggml_cuda_device_info & ggml_cuda_info() { // buffer pool for cuda (legacy) struct ggml_cuda_pool_leg : public ggml_cuda_pool { - static const int MAX_BUFFERS = 256; + // 1024 (upstream 256): LUCE_Q8_MEMO keeps one pooled q8_1 activation + // buffer per quantized matmul alive across a whole graph evaluation + // (~300 on a 64-layer hybrid), and a full pool falls back to freeing + // in-flight buffers with cudaFree. + static const int MAX_BUFFERS = 1024; int device; struct ggml_cuda_buffer { @@ -4311,6 +4315,49 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } } + // dflash: residual ADD + RMS_NORM + MUL. The add output stays live (it is + // the next residual), so this is a subgraph fusion with two outputs. + if (ops.size() == 3 && ops.begin()[0] == GGML_OP_ADD && ops.begin()[1] == GGML_OP_RMS_NORM && + ops.begin()[2] == GGML_OP_MUL) { + if (!ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx, node_idx + 2 })) { + return false; + } + const ggml_tensor * add = cgraph->nodes[node_idx]; + const ggml_tensor * rms = cgraph->nodes[node_idx + 1]; + const ggml_tensor * mul = cgraph->nodes[node_idx + 2]; + if (rms->src[0] != add) { + return false; + } + const ggml_tensor * w = nullptr; + if (mul->src[0] == rms) { + w = mul->src[1]; + } else if (mul->src[1] == rms) { + w = mul->src[0]; + } else { + return false; + } + const ggml_tensor * a = add->src[0]; + const ggml_tensor * b = add->src[1]; + if (a->type != GGML_TYPE_F32 || b->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32 || + add->type != GGML_TYPE_F32 || rms->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32) { + return false; + } + if (!ggml_is_contiguous(a) || !ggml_is_contiguous(b) || !ggml_is_contiguous(w) || + !ggml_is_contiguous(add) || !ggml_is_contiguous(mul)) { + return false; + } + if (!ggml_are_same_shape(a, b) || !ggml_are_same_shape(a, add) || !ggml_are_same_shape(a, mul)) { + return false; + } + if (w->ne[0] != a->ne[0] || ggml_nelements(w) != a->ne[0]) { + return false; + } + if (ggml_backend_buft_is_cuda_split(a->buffer->buft) || ggml_backend_buft_is_cuda_split(b->buffer->buft)) { + return false; + } + return true; + } + if (!ggml_can_fuse(cgraph, node_idx, ops)) { return false; } @@ -4910,6 +4957,12 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) { + ggml_cuda_op_add_rms_norm_mul_fused(*cuda_ctx, node, cgraph->nodes[i+1], cgraph->nodes[i+2]); + i += 2; + continue; + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD}, {})) { ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i+1], cgraph->nodes[i+2]); i += 2; @@ -6157,6 +6210,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g } } case GGML_OP_SSM_CONV: { + // dflash fused step mode handles any channel count + if (ggml_get_op_params_i32(op, 0) == 1) { + return true; + } // assumes d_inner % threads == 0 return op->src[0]->ne[1] % 128 == 0; } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu index ef98f675a..696a6f441 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu @@ -150,6 +150,57 @@ static __global__ void rms_norm_f32(const float * x, } } +// dflash: residual add fused into the following rms_norm * weight. +// sum = a + b (written to sum_out; it is the next residual) +// dst = rms_norm(sum) * w +// All of a, b, sum_out, dst are contiguous [ncols, R]; w is [ncols]. +template +static __global__ void add_rms_norm_mul_f32(const float * __restrict__ a, + const float * __restrict__ b, + float * __restrict__ sum_out, + float * __restrict__ dst, + const float * __restrict__ w, + const int ncols, + const float eps) { + const int64_t row = blockIdx.x; + const int tid = threadIdx.x; + + a += row * ncols; + b += row * ncols; + sum_out += row * ncols; + dst += row * ncols; + + float tmp = 0.0f; + for (int col = tid; col < ncols; col += block_size) { + const float s = a[col] + b[col]; + sum_out[col] = s; + tmp += s * s; + } + + extern __shared__ float s_sum[]; + tmp = block_reduce(tmp, s_sum); + + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); + + for (int col = tid; col < ncols; col += block_size) { + dst[col] = scale * sum_out[col] * w[col]; + } +} + +static void add_rms_norm_mul_f32_cuda(const float * a, const float * b, float * sum_out, float * dst, + const float * w, const int ncols, const int64_t nrows, + const float eps, cudaStream_t stream) { + const dim3 blocks_num(nrows, 1, 1); + if (ncols < 1024) { + const dim3 block_dims(256, 1, 1); + add_rms_norm_mul_f32<256><<>>(a, b, sum_out, dst, w, ncols, eps); + } else { + const dim3 block_dims(1024, 1, 1); + add_rms_norm_mul_f32<1024><<>>(a, b, sum_out, dst, w, ncols, eps); + } +} + template static __global__ void rms_norm_back_f32( const float * grad, const float * xf, float * dst, const int ncols, const float eps) { @@ -533,6 +584,36 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * eps, stream); } +// dflash: ADD (residual) + RMS_NORM + MUL in one launch. `add_tensor` is the +// residual add node (its output is materialized), `rms_tensor` is elided, +// `mul_tensor` receives the normalized * weight result. +void ggml_cuda_op_add_rms_norm_mul_fused(ggml_backend_cuda_context & ctx, + ggml_tensor * add_tensor, + ggml_tensor * rms_tensor, + ggml_tensor * mul_tensor) { + const ggml_tensor * a = add_tensor->src[0]; + const ggml_tensor * b = add_tensor->src[1]; + const ggml_tensor * w = (mul_tensor->src[0] == rms_tensor) ? mul_tensor->src[1] : mul_tensor->src[0]; + + float eps = 0.0f; + memcpy(&eps, rms_tensor->op_params, sizeof(float)); + GGML_ASSERT(eps >= 0.0f); + + GGML_ASSERT(a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32 && w->type == GGML_TYPE_F32); + GGML_ASSERT(add_tensor->type == GGML_TYPE_F32 && mul_tensor->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(a) && ggml_is_contiguous(b) && ggml_is_contiguous(w)); + GGML_ASSERT(ggml_is_contiguous(add_tensor) && ggml_is_contiguous(mul_tensor)); + GGML_ASSERT(ggml_are_same_shape(a, b) && ggml_are_same_shape(a, add_tensor) && ggml_are_same_shape(a, mul_tensor)); + GGML_ASSERT(w->ne[0] == a->ne[0] && ggml_nelements(w) == a->ne[0]); + + const int ncols = (int) a->ne[0]; + const int64_t nrows = ggml_nrows(a); + + add_rms_norm_mul_f32_cuda((const float *) a->data, (const float *) b->data, + (float *) add_tensor->data, (float *) mul_tensor->data, + (const float *) w->data, ncols, nrows, eps, ctx.stream()); +} + void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh index a74f63767..6313a98ce 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh @@ -16,3 +16,6 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// dflash: residual ADD + RMS_NORM + MUL fusion (see norm.cu) +void ggml_cuda_op_add_rms_norm_mul_fused(ggml_backend_cuda_context & ctx, ggml_tensor * add_tensor, ggml_tensor * rms_tensor, ggml_tensor * mul_tensor); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu index e6ce26f72..8c82cb2ac 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ssm-conv.cu @@ -244,7 +244,113 @@ static void ssm_conv_f32_cuda(const float * src0, const float * src1, const int } } +// dflash: fused conv step (see ggml_ssm_conv_step). One thread per channel +// walks the token loop with the K-1 history in registers, writes silu(conv), +// the optional rollback window and the new history in place. +template +static __global__ void ssm_conv_step_f32(const float * __restrict__ x, const int x_nb1, const int x_nb2, + const float * __restrict__ w, const int w_nb1, + float * state, const int st_nb1, const int st_nb2, + float * __restrict__ y, const int y_nb1, const int y_nb2, + float * ci, const int ci_nb1, const int ci_nb2, + const int C, const int T) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + const int s = blockIdx.y; + if (c >= C) return; + + const float * xs = (const float *) ((const char *) x + (size_t) s * x_nb2) + c; + float * st = (float *) ((char *) state + (size_t) s * st_nb2 + (size_t) c * st_nb1); + float * ys = (float *) ((char *) y + (size_t) s * y_nb2) + c; + float * cs = ci ? (float *) ((char *) ci + (size_t) s * ci_nb2 + (size_t) c * ci_nb1) : nullptr; + const float * wc = (const float *) ((const char *) w + (size_t) c * w_nb1); + + const int xs_stride = x_nb1 / sizeof(float); + const int ys_stride = y_nb1 / sizeof(float); + + float wt[K]; + float win[K]; // oldest first; win[K-1] is the current input +#pragma unroll + for (int k = 0; k < K; k++) { + wt[k] = wc[k]; + } +#pragma unroll + for (int j = 0; j < K - 1; j++) { + win[j] = st[j]; + if (cs) cs[j] = win[j]; + } + for (int t = 0; t < T; t++) { + const float xt = xs[(size_t) t * xs_stride]; + win[K - 1] = xt; + float acc = 0.0f; +#pragma unroll + for (int k = 0; k < K; k++) { + acc += win[k] * wt[k]; + } + ys[(size_t) t * ys_stride] = ggml_cuda_op_silu_single(acc); + if (cs) cs[K - 1 + t] = xt; +#pragma unroll + for (int j = 0; j < K - 1; j++) { + win[j] = win[j + 1]; + } + } +#pragma unroll + for (int j = 0; j < K - 1; j++) { + st[j] = win[j]; + } +} + +static void ggml_cuda_op_ssm_conv_step(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * w = dst->src[1]; + ggml_tensor * st = dst->src[2]; + ggml_tensor * ci = dst->src[3]; + + const int K = (int) w->ne[0]; + const int C = (int) w->ne[1]; + const int T = (int) dst->ne[1]; + const int S = (int) dst->ne[2]; + + GGML_ASSERT(x->type == GGML_TYPE_F32 && w->type == GGML_TYPE_F32 && st->type == GGML_TYPE_F32); + GGML_ASSERT(x->nb[0] == sizeof(float)); + GGML_ASSERT(w->nb[0] == sizeof(float)); + GGML_ASSERT(st->nb[0] == sizeof(float) && st->nb[1] == (size_t) (K - 1) * sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + if (ci) { + GGML_ASSERT(ci->type == GGML_TYPE_F32 && ci->nb[0] == sizeof(float)); + GGML_ASSERT(ci->ne[0] >= K - 1 + T); + } + + const int threads = 256; + const dim3 blocks((C + threads - 1) / threads, S, 1); + cudaStream_t stream = ctx.stream(); + + auto launch = [&](auto KK) { + constexpr int kK = decltype(KK)::value; + ssm_conv_step_f32<<>>( + (const float *) x->data, (int) x->nb[1], (int) x->nb[2], + (const float *) w->data, (int) w->nb[1], + (float *) st->data, (int) st->nb[1], (int) st->nb[2], + (float *) dst->data, (int) dst->nb[1], (int) dst->nb[2], + ci ? (float *) ci->data : nullptr, ci ? (int) ci->nb[1] : 0, ci ? (int) ci->nb[2] : 0, + C, T); + }; + switch (K) { + case 3: launch(std::integral_constant{}); break; + case 4: launch(std::integral_constant{}); break; + case 5: launch(std::integral_constant{}); break; + case 9: launch(std::integral_constant{}); break; + default: GGML_ABORT("ssm_conv_step only supports kernel sizes 3, 4, 5, 9."); + } +} + void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * silu_dst) { + // dflash: fused step mode (silu already applied by the kernel) + if (ggml_get_op_params_i32(dst, 0) == 1) { + GGML_ASSERT(silu_dst == nullptr); + ggml_cuda_op_ssm_conv_step(ctx, dst); + return; + } + const struct ggml_tensor * src0 = dst->src[0]; // conv_x const struct ggml_tensor * src1 = dst->src[1]; // conv1d.weight // dflash27b_ggml: optional src[2] = parent_ids (i32) enables tree mode diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 77a590cca..b46a6df2a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5941,6 +5941,53 @@ struct ggml_tensor * ggml_ssm_conv_tree( return result; } +// dflash: fused conv step. Same op id as ggml_ssm_conv; op_params[0] = 1 +// marks step mode, srcs are (x, c, conv_state, conv_input_out). +struct ggml_tensor * ggml_ssm_conv_step( + struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * c, + struct ggml_tensor * conv_state, + struct ggml_tensor * conv_input_out) { + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(c->type == GGML_TYPE_F32); + GGML_ASSERT(conv_state->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_matrix(c)); + GGML_ASSERT(ggml_is_contiguous(c)); + GGML_ASSERT(x->nb[0] == sizeof(float)); + GGML_ASSERT(x->ne[3] == 1); + + const int64_t d_conv = c->ne[0]; + const int64_t d_inner = c->ne[1]; + const int64_t n_t = x->ne[1]; + const int64_t n_s = x->ne[2]; + + GGML_ASSERT(x->ne[0] == d_inner); + GGML_ASSERT(conv_state->ne[0] == d_conv - 1); + GGML_ASSERT(conv_state->ne[1] == d_inner); + GGML_ASSERT(conv_state->ne[2] == n_s); + GGML_ASSERT(conv_state->nb[0] == sizeof(float)); + GGML_ASSERT(conv_state->nb[1] == (size_t)(d_conv - 1) * sizeof(float)); + if (conv_input_out) { + GGML_ASSERT(conv_input_out->type == GGML_TYPE_F32); + GGML_ASSERT(conv_input_out->ne[0] >= d_conv - 1 + n_t); + GGML_ASSERT(conv_input_out->ne[1] == d_inner); + GGML_ASSERT(conv_input_out->ne[2] == n_s); + GGML_ASSERT(conv_input_out->nb[0] == sizeof(float)); + } + + struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_inner, n_t, n_s); + ggml_set_op_params_i32(result, 0, 1); // step mode + + result->op = GGML_OP_SSM_CONV; + result->src[0] = x; + result->src[1] = c; + result->src[2] = conv_state; + result->src[3] = conv_input_out; + + return result; +} + // ggml_ssm_scan struct ggml_tensor * ggml_ssm_scan( @@ -6754,6 +6801,29 @@ void ggml_gated_delta_net_set_skip_intermediate( tensor->nb[3] = tensor->nb[2]*tensor->ne[2]; } +// dflash: raw-gate mode (see ggml.h). src[8] is reserved for the optional +// active-slot map; dt_bias -> src[9], A -> src[10], +// op_params[2] = 1. +void ggml_gated_delta_net_set_raw_gates( + struct ggml_tensor * tensor, + struct ggml_tensor * dt_bias, + struct ggml_tensor * A) { + GGML_ASSERT(tensor != NULL); + GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); + GGML_ASSERT(dt_bias != NULL && A != NULL); + GGML_ASSERT(dt_bias->type == GGML_TYPE_F32 && A->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(dt_bias) && ggml_is_contiguous(A)); + const struct ggml_tensor * v = tensor->src[2]; + GGML_ASSERT(ggml_nelements(dt_bias) == v->ne[1]); + GGML_ASSERT(ggml_nelements(A) == v->ne[1]); + // scalar gate only (no KDA), no tree mode + GGML_ASSERT(tensor->src[3]->ne[0] == 1); + GGML_ASSERT(tensor->src[6] == NULL); + tensor->src[9] = dt_bias; + tensor->src[10] = A; + ggml_set_op_params_i32(tensor, 2, 1); +} + // dflash: tree-mode variant. Same op, with parent_ids plumbed into // src[6] so the CUDA kernel can branch-reload state at DFS transitions. struct ggml_tensor * ggml_gated_delta_net_tree( From 7a1c77ea0aa1c312b8d97e3c89f44091dc579c5f Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:10:26 +0200 Subject: [PATCH 22/42] ggml: 64x64 MMQ tiles for dense verify widths on RDNA Rename the RDNA small-tile macro to GGML_CUDA_MMQ_SMALL_TILE and apply it to IQ4_XS/Q4_K/Q5_K/Q6_K/Q8_0 in addition to the ROCmFPX formats. At spec-decode verify widths (N<=16) the 128-row tile leaves a 5120-row projection with only 40 blocks on a 64-CU gfx1201; 64x64/4-warp tiles measured +12-23% on those shapes (verify step 43.8 -> 39.7 ms on Qwen3.8-27B) at ~8% prefill cost. --- .../deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh | 25 +++++++++++-------- .../template-instances/generate_cu_files.py | 12 ++++++--- .../template-instances/mmq-instance-iq4_xs.cu | 1 + .../mmq-instance-q2_0_rocmfp2.cu | 2 +- .../mmq-instance-q2_1_rocmfp2_mix.cu | 2 +- .../mmq-instance-q3_0_rocmfpx.cu | 2 +- .../mmq-instance-q3_1_rocmfp3_mix.cu | 2 +- .../mmq-instance-q4_0_rocmfp4_fast.cu | 2 +- .../template-instances/mmq-instance-q4_k.cu | 2 +- .../template-instances/mmq-instance-q5_k.cu | 1 + .../template-instances/mmq-instance-q6_k.cu | 1 + .../template-instances/mmq-instance-q8_0.cu | 1 + 12 files changed, 34 insertions(+), 19 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index 876a8ba45..266bf6006 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -107,9 +107,14 @@ struct tile_x_sizes { int sc; }; -// RDNA uses 128x128, eight-warp MMQ tiles by default. Q4_K narrows the row -// dimension to 128x64, while ROCmFPX uses 64x64 four-warp tiles. Their -// unpacking pressure makes the smaller tiles faster on gfx1151. +// RDNA uses 128x128, eight-warp MMQ tiles by default. Template instances +// compiled with GGML_CUDA_MMQ_SMALL_TILE use 64x64, four-warp tiles: +// - ROCmFPX formats: their unpacking pressure makes the smaller tile faster +// on gfx1151; +// - IQ4_XS / Q6_K / Q8_0 (dense hybrid targets): at spec-decode verify +// widths (N<=16) the 128-row tile leaves a 5120-row projection with only +// 40 blocks on a 64-CU gfx1201; the small tile measured +12-23% there +// (mmq_probe) at the cost of ~8% prefill throughput. #ifndef LUCEBOX_RDNA_MMQ_TILE_OVERRIDE #define LUCEBOX_RDNA_MMQ_TILE_OVERRIDE 1 #endif @@ -122,7 +127,7 @@ struct tile_x_sizes { static int get_mmq_x_max_host(const int cc) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #else return 128; @@ -139,7 +144,7 @@ static int get_mmq_x_max_host(const int cc) { static constexpr __device__ int get_mmq_x_max_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #else return 128; @@ -169,7 +174,7 @@ static constexpr __device__ int get_mmq_x_max_device() { static int get_mmq_y_host(const int cc) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #elif defined(LUCEBOX_RDNA_MMQ_Y) return LUCEBOX_RDNA_MMQ_Y; @@ -191,7 +196,7 @@ static constexpr __device__ int get_iter_k([[maybe_unused]] const ggml_type type static constexpr __device__ int get_mmq_y_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 64; #elif defined(LUCEBOX_RDNA_MMQ_Y) return LUCEBOX_RDNA_MMQ_Y; @@ -346,7 +351,7 @@ static constexpr __device__ int mmq_get_granularity_device(const int /*mmq_x*/) #if defined(GGML_USE_HIP) static int mmq_get_nwarps_host(const int cc, const int warp_size) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 4; #elif defined(LUCEBOX_RDNA_MMQ_Y) return 4; @@ -364,7 +369,7 @@ static int mmq_get_nwarps_host(const int /*cc*/, const int warp_size) { static constexpr __device__ int mmq_get_nwarps_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(GGML_CUDA_MMQ_SMALL_TILE) return 4; #elif defined(LUCEBOX_RDNA_MMQ_Y) return 4; @@ -4223,7 +4228,7 @@ template #if defined(GGML_USE_HIP) // RDNA4 is compute-bound on MMQ (WMMA path); allow compiler to use more VGPRs // (minBlocks=1 matches NVIDIA Volta+ behavior and reduces register spilling). -#if defined(RDNA4) && !defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(RDNA4) && !defined(GGML_CUDA_MMQ_SMALL_TILE) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 1) #elif defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py index f87396f5a..5e3a1f00d 100755 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -102,10 +102,16 @@ def get_short_name(long_quant_name): "GGML_TYPE_Q2_1_ROCMFP2_MIX", "GGML_TYPE_Q3_0_ROCMFPX", "GGML_TYPE_Q3_1_ROCMFP3_MIX", + # Dense hybrid (Qwen3.5/3.8) verify widths N<=16 on gfx1201: the + # 128-row tile leaves a 5120-row projection with only 40 blocks; + # 64x64/4-warp tiles measured +12-23% on those shapes (mmq_probe). + "GGML_TYPE_IQ4_XS", + "GGML_TYPE_Q4_K", + "GGML_TYPE_Q5_K", + "GGML_TYPE_Q6_K", + "GGML_TYPE_Q8_0", }: - guard = "#define GGML_CUDA_ROCMFPX_MMQ_TILE 1\n" - if type == "GGML_TYPE_Q4_K": - guard = "#define LUCEBOX_RDNA_MMQ_Y 64\n" + guard = "#define GGML_CUDA_MMQ_SMALL_TILE 1\n" f.write(SOURCE_MMQ.format(type=type, guard=guard)) for type in range(1, 17): diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu index 1eb3b7430..5e2a1127a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu index 8221e1d1e..b00cd9a0c 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q2_0_ROCMFP2); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu index 647b4572f..f73033e33 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_1_rocmfp2_mix.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q2_1_ROCMFP2_MIX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu index 2380af75c..486782982 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q3_0_ROCMFPX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu index 1873e073f..92197f871 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_1_rocmfp3_mix.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q3_1_ROCMFP3_MIX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu index 94a2bb0f5..92cb4653d 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu index dcf47b2c2..f9a206d20 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#define LUCEBOX_RDNA_MMQ_Y 64 +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q4_K); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu index a2e90ffd5..7cf43f75e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q5_K); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu index 470938fef..8bc6b7434 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q6_K); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu index 974477bbb..fb8fcf911 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu @@ -1,5 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. +#define GGML_CUDA_MMQ_SMALL_TILE 1 #include "../mmq.cuh" DECL_MMQ_CASE(GGML_TYPE_Q8_0); From d50f928203bfc672d8d57703f463f375e1a6963d Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:10:27 +0200 Subject: [PATCH 23/42] qwen35: stacked projections and fused DeltaNet decode graph - loader places attn_gate|attn_qkv and ssm_beta|ssm_alpha back to back and exposes zero-copy stacked aliases (L.wqkv_z, L.ssm_ba): one GEMV each instead of two (DFLASH_QWEN35_NO_STACK=1 disables). - FFN uses ggml_swiglu_split so the backend fuses gate/up/GLU into one vector kernel at decode. - DeltaNet block: single l2_norm over the q|k slab, ggml_ssm_conv_step, raw-gate gated_delta_net (in place, no state copy), no q/k head repeat (the kernel broadcasts). DFLASH_QWEN35_NO_FUSED_KERNELS=1 keeps the op-by-op graph for A/B. - DFLASH_KV_ROTATE=0 skips the FWHT K/Q rotation (precision-neutral with q8_0/f16 caches, two fewer launches per attention layer). Qwen3.8-27B IQ4_XS on R9700: plain decode 30.4 -> 33.8 tok/s with identical greedy output. --- server/src/internal.h | 8 + server/src/qwen35/gguf_target_loader.cpp | 103 +++++++- server/src/qwen35/qwen35_target_graph.cpp | 303 +++++++++++++++------- 3 files changed, 311 insertions(+), 103 deletions(-) diff --git a/server/src/internal.h b/server/src/internal.h index 9cb2a03b6..1aabee02f 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -76,6 +76,13 @@ struct TargetLayer { ggml_tensor * ssm_dt_bias = nullptr; // [dt_rank] per-head alpha bias ggml_tensor * ssm_norm = nullptr; // [head_v_dim] ggml_tensor * ssm_out = nullptr; // output projection after delta-net + // Zero-copy stacked projections (set by the loader when the two source + // tensors share a type and were placed back to back in the weight buffer): + // wqkv_z: rows [0, n_z) = wqkv_gate (z), rows [n_z, ...) = wqkv + // ssm_ba: rows [0, dt_rank) = ssm_beta, rows [dt_rank, ...) = ssm_alpha + // One GEMV each instead of two; nullptr when stacking was not possible. + ggml_tensor * wqkv_z = nullptr; + ggml_tensor * ssm_ba = nullptr; // MoE FFN (qwen35moe only; nullptr on dense qwen35) ggml_tensor * ffn_gate_inp = nullptr; // [hidden, n_expert] router @@ -147,6 +154,7 @@ struct CpuEmbedder { struct TargetWeights { ggml_context * ctx = nullptr; + ggml_context * stack_ctx = nullptr; // owns the stacked alias tensors ggml_backend_t backend = nullptr; ggml_backend_buffer_t buf = nullptr; diff --git a/server/src/qwen35/gguf_target_loader.cpp b/server/src/qwen35/gguf_target_loader.cpp index 1f917ee69..4c41acb75 100644 --- a/server/src/qwen35/gguf_target_loader.cpp +++ b/server/src/qwen35/gguf_target_loader.cpp @@ -680,16 +680,71 @@ bool load_target_gguf_partial(const std::string & path, if (!t || !should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, plan.load_output, plan.skip_expert_tensors)) { continue; } - alloc_total = align_up_size(alloc_total, alignment); TargetTensorAlloc a; a.tensor = t; a.file_offset = gguf_get_data_offset(gctx) + gguf_get_tensor_offset(gctx, tid); a.file_size = gguf_get_tensor_size(gctx, tid); - a.buffer_offset = alloc_total; - alloc_total += ggml_backend_buft_get_alloc_size(buft, t); allocs.push_back(a); } + // Stacked projections: place each (first, second) pair back to back in the + // weight buffer so one alias tensor spanning both rows serves a single + // GEMV. Only for the plain single-buffer path (the TP meta allocator + // places tensors itself) and only when the pair shares type/ne0 and the + // first tensor's byte size keeps the second one aligned. + const bool can_stack = !plan.metadata_only && !ggml_backend_buft_is_meta(buft) && + std::getenv("DFLASH_QWEN35_NO_STACK") == nullptr; + if (can_stack) { + auto find_alloc = [&](const std::string & name) -> int { + for (size_t i = 0; i < allocs.size(); i++) { + if (name == allocs[i].tensor->name) return (int)i; + } + return -1; + }; + // (first, second) suffix pairs; the alias tensor stacks first's rows + // then second's, so they are emitted in that order whichever member + // the file lists first. + static const char * const kPairs[][2] = { + { ".attn_gate.weight", ".attn_qkv.weight" }, + { ".ssm_beta.weight", ".ssm_alpha.weight" }, + }; + std::vector ordered; + ordered.reserve(allocs.size()); + std::vector taken(allocs.size(), false); + for (size_t i = 0; i < allocs.size(); i++) { + if (taken[i]) continue; + const std::string name = allocs[i].tensor->name; + int first = -1, second = -1; + if (name.rfind("blk.", 0) == 0) { + for (const auto & pr : kPairs) { + for (int m = 0; m < 2; m++) { + const size_t pos = name.find(pr[m]); + if (pos == std::string::npos) continue; + const std::string prefix = name.substr(0, pos); + first = find_alloc(prefix + pr[0]); + second = find_alloc(prefix + pr[1]); + break; + } + if (first >= 0 || second >= 0) break; + } + } + if (first >= 0 && second >= 0 && !taken[(size_t)first] && !taken[(size_t)second]) { + taken[(size_t)first] = taken[(size_t)second] = true; + ordered.push_back(allocs[(size_t)first]); + ordered.push_back(allocs[(size_t)second]); + continue; + } + taken[i] = true; + ordered.push_back(allocs[i]); + } + allocs.swap(ordered); + } + for (TargetTensorAlloc & a : allocs) { + alloc_total = align_up_size(alloc_total, alignment); + a.buffer_offset = alloc_total; + alloc_total += ggml_backend_buft_get_alloc_size(buft, a.tensor); + } + // The generic meta buffer allocator must see all tensors together so it // can allocate each device from its actual slices. The legacy loader's // monolithic backing buffer would reserve alloc_total on every rank. @@ -793,6 +848,47 @@ bool load_target_gguf_partial(const std::string & path, return false; } } + if (can_stack) { + // Alias tensors over adjacent pairs. They read the same bytes as + // the two source tensors (no copy, no extra VRAM). + ggml_init_params sip{}; + sip.mem_size = (2 * n_layer + 8) * ggml_tensor_overhead(); + sip.mem_buffer = nullptr; + sip.no_alloc = true; + out.stack_ctx = ggml_init(sip); + int n_stacked = 0; + auto make_stack = [&](ggml_tensor * first, ggml_tensor * second, + const char * name) -> ggml_tensor * { + if (!first || !second || !out.stack_ctx) return nullptr; + if (first->type != second->type || first->ne[0] != second->ne[0]) return nullptr; + if (!ggml_is_contiguous(first) || !ggml_is_contiguous(second)) return nullptr; + const char * f = (const char *)first->data; + const char * sd = (const char *)second->data; + if (!f || !sd || sd != f + ggml_nbytes(first)) return nullptr; + ggml_tensor * st = ggml_new_tensor_2d(out.stack_ctx, first->type, + first->ne[0], first->ne[1] + second->ne[1]); + // The alias must not need padding the backend would want to + // clear past its end (that would scribble on the next tensor). + if (ggml_backend_buft_get_alloc_size(buft, st) != ggml_nbytes(st)) return nullptr; + ggml_set_name(st, name); + if (ggml_backend_tensor_alloc(out.buf, st, first->data) != GGML_STATUS_SUCCESS) { + return nullptr; + } + n_stacked++; + return st; + }; + for (int il = 0; il < (int)n_layer; il++) { + TargetLayer & L = out.layers[il]; + char nm[96]; + std::snprintf(nm, sizeof(nm), "blk.%d.attn_gate_qkv.stacked", il); + L.wqkv_z = make_stack(L.wqkv_gate, L.wqkv, nm); + std::snprintf(nm, sizeof(nm), "blk.%d.ssm_beta_alpha.stacked", il); + L.ssm_ba = make_stack(L.ssm_beta, L.ssm_alpha, nm); + } + if (n_stacked > 0) { + std::fprintf(stderr, "[loader] stacked %d projection pairs (zero-copy aliases)\n", n_stacked); + } + } } const size_t data_start = gguf_get_data_offset(gctx); @@ -958,6 +1054,7 @@ bool load_target_gguf_partial(const std::string & path, void free_target_weights(TargetWeights & w) { if (w.buf) { ggml_backend_buffer_free(w.buf); w.buf = nullptr; } if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; } + if (w.stack_ctx) { ggml_free(w.stack_ctx); w.stack_ctx = nullptr; } // CpuEmbedder destructor handles the mmap automatically. w.moe_hybrid.reset(); w.layers.clear(); diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 5f10dcdc9..ccea1c550 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -156,7 +156,13 @@ bool create_target_cache_partial(const TargetWeights & w, // Graph-level FWHT K-rotation (TurboQuant-style outlier spreading with // standard quant types that keep fast FA kernel paths on all arches). // Skip for TQ3_0 K cache — that type already applies WHT during quantization. - out.kv_k_rotated = (kv_k_type != GGML_TYPE_TQ3_0); + // DFLASH_KV_ROTATE=0 turns it off (two fewer launches per attention layer; + // with q8_0/f16 caches the rotation is precision-neutral). + static const bool kv_rotate_env = []() { + const char * e = std::getenv("DFLASH_KV_ROTATE"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + out.kv_k_rotated = (kv_k_type != GGML_TYPE_TQ3_0) && kv_rotate_env; const bool needs_256_stride = kv_k_type == GGML_TYPE_TQ3_0 || kv_v_type == GGML_TYPE_TQ3_0; @@ -689,10 +695,19 @@ bool ensure_ssm_snapshot(TargetCache & c, ggml_backend_t backend) { static ggml_tensor * build_swiglu_ffn(ggml_context * ctx, ggml_tensor * cur, const TargetLayer & L) { - ggml_tensor * gate = apply_scale2(ctx, ggml_mul_mat(ctx, L.w_gate, cur), L.w_gate_s); // [inter, n_tokens] - gate = ggml_silu(ctx, gate); - ggml_tensor * up = apply_scale2(ctx, ggml_mul_mat(ctx, L.w_up, cur), L.w_up_s); - ggml_tensor * gu = ggml_mul(ctx, gate, up); + ggml_tensor * gate = ggml_mul_mat(ctx, L.w_gate, cur); // [inter, n_tokens] + ggml_tensor * up = ggml_mul_mat(ctx, L.w_up, cur); + ggml_tensor * gu; + if (L.w_gate_s == 1.0f && L.w_up_s == 1.0f) { + // GLU node right after the two matmuls: the CUDA/HIP backend fuses + // mul_mat(gate) + mul_mat(up) + swiglu into a single vector kernel + // for single-token decode. + gu = ggml_swiglu_split(ctx, gate, up); + } else { + gate = ggml_silu(ctx, apply_scale2(ctx, gate, L.w_gate_s)); + up = apply_scale2(ctx, up, L.w_up_s); + gu = ggml_mul(ctx, gate, up); + } return apply_scale2(ctx, ggml_mul_mat(ctx, L.w_down, gu), L.w_down_s); // [hidden, n_tokens] } @@ -1108,23 +1123,53 @@ static ggml_tensor * build_delta_net_block( GGML_ASSERT(!ragged || (!cap && !parent_ids)); const bool can_skip_gdn_intermediate = skip_gdn_intermediate && !parent_ids && !cap; - // ── Whole-batch projections ───────────────────────────────────── - // qkv_mixed = wqkv @ cur [10240, n_tokens] - ggml_tensor * qkv_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv, cur), L.wqkv_s); - - // z = wqkv_gate @ cur [inner, n_tokens] - ggml_tensor * z = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv_gate, cur), L.wqkv_gate_s); + // Row slices of stacked projections are strided for multi-token inputs. + // Materialize only the small beta/alpha slices; qkv keeps its explicit + // column stride and z is made contiguous at the final per-segment gate. + auto contig = [&](ggml_tensor * t) { + return ggml_is_contiguous(t) ? t : ggml_cont(ctx, t); + }; - // beta = sigmoid(ssm_beta @ cur) [dt_rank, n_tokens] - ggml_tensor * beta_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_beta, cur), L.ssm_beta_s); - beta_2d = ggml_sigmoid(ctx, beta_2d); + // ── Whole-batch projections ───────────────────────────────────── + // One GEMM over the zero-copy stacked (z | qkv) alias when possible. + ggml_tensor * qkv_2d = nullptr; + ggml_tensor * z = nullptr; + const bool stacked_qkv_z = + L.wqkv_z && L.wqkv_s == 1.0f && L.wqkv_gate_s == 1.0f; + if (stacked_qkv_z) { + const int64_t n_z = L.wqkv_gate->ne[1]; + ggml_tensor * qkvz = ggml_mul_mat(ctx, L.wqkv_z, cur); + const size_t e = ggml_element_size(qkvz); + z = ggml_view_2d(ctx, qkvz, n_z, n_tokens, qkvz->nb[1], 0); + qkv_2d = ggml_view_2d(ctx, qkvz, conv_channels, n_tokens, + qkvz->nb[1], (size_t)n_z * e); + } else { + qkv_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv, cur), L.wqkv_s); + z = apply_scale2(ctx, ggml_mul_mat(ctx, L.wqkv_gate, cur), L.wqkv_gate_s); + } + + // One GEMM over the zero-copy stacked (beta | alpha) alias when possible. + ggml_tensor * beta_2d = nullptr; + ggml_tensor * alpha_2d = nullptr; + const bool stacked_ba = + L.ssm_ba && L.ssm_beta_s == 1.0f && L.ssm_alpha_s == 1.0f; + if (stacked_ba) { + ggml_tensor * ba = ggml_mul_mat(ctx, L.ssm_ba, cur); + const size_t e = ggml_element_size(ba); + beta_2d = contig(ggml_view_2d( + ctx, ba, num_v_heads, n_tokens, ba->nb[1], 0)); + alpha_2d = contig(ggml_view_2d( + ctx, ba, num_v_heads, n_tokens, ba->nb[1], + (size_t)num_v_heads * e)); + } else { + beta_2d = apply_scale2( + ctx, ggml_mul_mat(ctx, L.ssm_beta, cur), L.ssm_beta_s); + alpha_2d = apply_scale2( + ctx, ggml_mul_mat(ctx, L.ssm_alpha, cur), L.ssm_alpha_s); + } - // alpha = ssm_alpha @ cur [dt_rank, n_tokens] - // g = softplus(alpha + ssm_dt_bias) * ssm_a (-A_log.exp() * softplus) - ggml_tensor * alpha = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_alpha, cur), L.ssm_alpha_s); - alpha = ggml_add(ctx, alpha, L.ssm_dt_bias); - alpha = ggml_softplus(ctx, alpha); - ggml_tensor * g_2d = ggml_mul(ctx, alpha, L.ssm_a); + static const bool fused_kernels_env = + std::getenv("DFLASH_QWEN35_NO_FUSED_KERNELS") == nullptr; // ── Token-axis segments: prompt chunks first, then the decode batch ── struct DeltaSeg { @@ -1187,19 +1232,48 @@ static ggml_tensor * build_delta_net_block( // 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 && !seg_tree) || + const bool dense_chain = !ragged && !active_slot_ids && !seg_tree; + const bool inplace_state = dense_chain || (seg_active && !seg_tree) || (allow_inplace_state && can_skip_gdn_intermediate && !ragged && n_seq_tokens == 1); - ggml_tensor * qkv_mixed = ggml_reshape_3d(ctx, - seg_cols(qkv_2d, seg.off, seg_tokens), - conv_channels, n_seq_tokens, seg_seqs); + ggml_tensor * qkv_seg = seg_cols(qkv_2d, seg.off, seg_tokens); + ggml_tensor * qkv_mixed = stacked_qkv_z + ? ggml_view_3d(ctx, qkv_seg, + conv_channels, n_seq_tokens, seg_seqs, + qkv_seg->nb[1], qkv_seg->nb[1] * n_seq_tokens, 0) + : ggml_reshape_3d(ctx, qkv_seg, + conv_channels, n_seq_tokens, seg_seqs); + + // Chunked delta-net path is opt-in and chain-only. + bool use_chunked = false; + if (can_skip_gdn_intermediate && n_seq_tokens > 1) { + if (const char * s_env = std::getenv("DFLASH27B_CHUNKED")) { + use_chunked = (std::atoi(s_env) != 0); + } + } + const bool fused_conv = fused_kernels_env && !seg_active && !seg_tree; + const bool raw_gates = fused_kernels_env && !seg_tree && !use_chunked; + ggml_tensor * beta = ggml_reshape_4d(ctx, seg_cols(beta_2d, seg.off, seg_tokens), 1, num_v_heads, n_seq_tokens, seg_seqs); - ggml_tensor * g_tensor = ggml_reshape_4d(ctx, - seg_cols(g_2d, seg.off, seg_tokens), - 1, num_v_heads, n_seq_tokens, seg_seqs); + ggml_tensor * alpha = ggml_reshape_3d(ctx, + seg_cols(alpha_2d, seg.off, seg_tokens), + num_v_heads, n_seq_tokens, seg_seqs); + ggml_tensor * g_tensor = nullptr; + if (raw_gates) { + // The kernel applies sigmoid(beta) and softplus(alpha + dt_bias) * A. + g_tensor = ggml_reshape_4d( + ctx, alpha, 1, num_v_heads, n_seq_tokens, seg_seqs); + } else { + beta = ggml_sigmoid(ctx, beta); + alpha = ggml_add(ctx, alpha, L.ssm_dt_bias); + alpha = ggml_softplus(ctx, alpha); + g_tensor = ggml_mul(ctx, alpha, L.ssm_a); + g_tensor = ggml_reshape_4d( + ctx, g_tensor, 1, num_v_heads, n_seq_tokens, seg_seqs); + } // ── Fetch conv state [kernel-1, conv_channels] and prepend to qkv_mixed // along the token axis to form the convolution input. @@ -1218,66 +1292,83 @@ static ggml_tensor * build_delta_net_block( w.ssm_d_conv - 1, conv_channels, seg_seqs); } - // qkv_mixed currently is [conv_channels, n_tokens, n_seqs]; we need - // [n_tokens, conv_channels, n_seqs] to concat on dim 0. - ggml_tensor * qkv_T = ggml_transpose(ctx, qkv_mixed); - - ggml_tensor * conv_input = ggml_concat(ctx, conv_states_r, qkv_T, 0); - // I0 domain: [0,K_conv-2] are prefix-history rows; tree token flat slot t - // (root-inclusive, including synthetic root t=0) is stored at - // conv_input row (K_conv-1)+t. - // conv_input: [kernel-1 + n_tokens, conv_channels, n_seqs] - - // For spec-decode rollback: copy the full conv_input into the persistent - // cache buffer via an in-graph ggml_cpy. This avoids marking conv_input as - // 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) { - // 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; - } 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); + ggml_tensor * conv_out = nullptr; + if (fused_conv) { + // One kernel: window = [conv_state | x], silu(conv), history + // write-back, and (when capturing) the rollback window copy. + ggml_tensor * ci_dst = nullptr; + if (cap && cap->conv_input) { + const int64_t ci_len = (w.ssm_d_conv - 1) + n_tokens; + ci_dst = (ci_len == cap->conv_input->ne[0]) + ? cap->conv_input + : 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); + } + conv_out = ggml_ssm_conv_step(ctx, qkv_mixed, L.ssm_conv1d, conv_states_r, ci_dst); + } else { + // qkv_mixed currently is [conv_channels, n_tokens, n_seqs]; we need + // [n_tokens, conv_channels, n_seqs] to concat on dim 0. + ggml_tensor * qkv_T = ggml_transpose(ctx, qkv_mixed); + + ggml_tensor * conv_input = ggml_concat(ctx, conv_states_r, qkv_T, 0); + // I0 domain: [0,K_conv-2] are prefix-history rows; tree token flat slot t + // (root-inclusive, including synthetic root t=0) is stored at + // conv_input row (K_conv-1)+t. + // conv_input: [kernel-1 + n_tokens, conv_channels, n_seqs] + + // For spec-decode rollback: copy the full conv_input into the persistent + // cache buffer via an in-graph ggml_cpy. This avoids marking conv_input as + // 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) { + // 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; + } 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); + } + GGML_ASSERT(ggml_nelements(conv_input) == ggml_nelements(dst)); + ggml_build_forward_expand(gf, ggml_cpy(ctx, conv_input, dst)); } - GGML_ASSERT(ggml_nelements(conv_input) == ggml_nelements(dst)); - ggml_build_forward_expand(gf, ggml_cpy(ctx, conv_input, dst)); - } - // ── Save the last (kernel-1) steps back to the conv state - ggml_tensor * last_conv = ggml_view_3d(ctx, conv_input, - 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 && !seg_tree) { - const int64_t slab = - (int64_t)(w.ssm_d_conv - 1) * conv_channels; - ggml_tensor * compact_last = ggml_reshape_2d( - ctx, ggml_cont(ctx, last_conv), slab, seg_seqs); - ggml_tensor * all_conv = ggml_reshape_2d( - 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)); - } else if (!seg_tree) { - ggml_build_forward_expand(gf, ggml_cpy(ctx, last_conv, seg.conv_st)); - } - - // ── 1D conv + silu - // Tree mode: use the parent-chain-aware variant so sibling nodes gather - // 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_ssm_conv (ctx, conv_input, L.ssm_conv1d); - conv_out = ggml_silu(ctx, conv_out); + // ── Save the last (kernel-1) steps back to conv_state + ggml_tensor * last_conv = ggml_view_3d(ctx, conv_input, + 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 && !seg_tree) { + const int64_t slab = + (int64_t)(w.ssm_d_conv - 1) * conv_channels; + ggml_tensor * compact_last = ggml_reshape_2d( + ctx, ggml_cont(ctx, last_conv), slab, seg_seqs); + ggml_tensor * all_conv = ggml_reshape_2d( + 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)); + } else if (!seg_tree) { + ggml_build_forward_expand( + gf, ggml_cpy(ctx, last_conv, seg.conv_st)); + } + + // ── 1D conv + silu + // Tree mode: use the parent-chain-aware variant so sibling nodes gather + // 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). + conv_out = parent_ids + ? ggml_ssm_conv_tree(ctx, conv_input, L.ssm_conv1d, parent_ids) + : ggml_ssm_conv (ctx, conv_input, L.ssm_conv1d); + conv_out = ggml_silu(ctx, conv_out); + } // conv_out: [conv_channels, n_tokens, n_seqs] const int64_t q_offset = 0; @@ -1306,13 +1397,29 @@ static ggml_tensor * build_delta_net_block( row_size * n_seq_tokens, v_offset * elt); - // L2 norm on Q and K - q_c = ggml_l2_norm(ctx, q_c, w.rms_eps); - k_c = ggml_l2_norm(ctx, k_c, w.rms_eps); - - // Repeat Q and K from num_k_heads to num_v_heads so they match V's layout - // (only needed if not using the fused op's broadcast support). - if (num_k_heads != num_v_heads) { + // L2 norm on Q and K: q and k heads are adjacent in conv_out, so one + // launch over the [head_k_dim, 2*num_k_heads] slab normalizes both. + { + ggml_tensor * qk_c = ggml_view_4d(ctx, conv_out, + head_k_dim, 2 * num_k_heads, n_seq_tokens, seg_seqs, + head_k_dim * elt, + row_size, + row_size * n_seq_tokens, + q_offset * elt); + ggml_tensor * qk_n = ggml_l2_norm(ctx, qk_c, w.rms_eps); // contiguous [hd, 2*Hk, T, S] + const size_t ne_ = ggml_element_size(qk_n); + q_c = ggml_view_4d(ctx, qk_n, head_k_dim, num_k_heads, n_seq_tokens, seg_seqs, + qk_n->nb[1], qk_n->nb[2], qk_n->nb[3], 0); + k_c = ggml_view_4d(ctx, qk_n, head_k_dim, num_k_heads, n_seq_tokens, seg_seqs, + qk_n->nb[1], qk_n->nb[2], qk_n->nb[3], + (size_t)num_k_heads * head_k_dim * ne_); + } + + // Repeat Q and K from num_k_heads to num_v_heads so they match V's layout. + // The fused gated_delta_net kernels broadcast heads themselves (v head h + // reads q/k head h % num_k_heads, the same tiling ggml_repeat produces), + // so only the chunked path needs the materialized copies. + if (num_k_heads != num_v_heads && use_chunked) { q_c = ggml_repeat_4d(ctx, q_c, head_k_dim, num_v_heads, n_seq_tokens, seg_seqs); k_c = ggml_repeat_4d(ctx, k_c, head_k_dim, num_v_heads, n_seq_tokens, seg_seqs); } @@ -1371,13 +1478,6 @@ static ggml_tensor * build_delta_net_block( // default — port produces correct shape but slightly wrong final state, // 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 (const char * s_env = std::getenv("DFLASH27B_CHUNKED")) { - use_chunked = (std::atoi(s_env) != 0); - } - } - ggml_tensor * output = nullptr; if (use_chunked) { @@ -1410,6 +1510,9 @@ static ggml_tensor * build_delta_net_block( result->src[7] = persist_inter; } } + if (raw_gates) { + ggml_gated_delta_net_set_raw_gates(result, L.ssm_dt_bias, L.ssm_a); + } if (can_skip_gdn_intermediate) { ggml_gated_delta_net_set_skip_intermediate(result, true); } @@ -1466,7 +1569,7 @@ static ggml_tensor * build_delta_net_block( // ── Gated output norm: rms_norm(output) * silu(z_4d) ggml_tensor * z_4d = ggml_reshape_4d(ctx, - seg_cols(z, seg.off, seg_tokens), + contig(seg_cols(z, seg.off, seg_tokens)), head_v_dim, num_v_heads, n_seq_tokens, seg_seqs); ggml_tensor * output_n = ggml_rms_norm(ctx, rms_norm_input_f32(ctx, output), w.rms_eps); output_n = ggml_mul(ctx, output_n, L.ssm_norm); From bcd0742d478313121ec80480a72dc62da45c79d0 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:10:27 +0200 Subject: [PATCH 24/42] qwen35: adaptive speculation policy and chain-path profiling - Qwen35AdaptiveSpecPolicy: EMA of accepted draft tokens per step; below 0.8*(spec_step_ratio-1) the loop runs a burst of plain-decode steps (seed-only verify, no drafter/heads/snapshot/rollback, features still captured) and probes again afterwards. Env DFLASH_QWEN35_SPEC_STEP_RATIO (default 1.7, 0 disables) and DFLASH_QWEN35_AR_BURST (default 40). Low-acceptance prose 28.1 -> 32.4 tok/s, code/mixed unchanged. - Confidence gate now uses the fused Markov graph and truncates on the host; DFLASH_QWEN35_DSPARK_CONF_DEBUG=1 prints per-position scores. - spec-profile hooks for the chain path (project/snapshot/verify/ rollback/feature). --- server/src/qwen35/qwen35_backend.cpp | 317 ++++++++++++++++++--------- 1 file changed, 217 insertions(+), 100 deletions(-) diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index a4df94f81..345afeeb1 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -2488,6 +2488,37 @@ static float qwen35_dspark_confidence_threshold() { return kThreshold; } +// Adaptive speculation policy: a spec step (draft + heads + width-q verify) +// costs about DFLASH_QWEN35_SPEC_STEP_RATIO plain-decode steps, so it only +// pays off while the drafter gets more than (ratio - 1) of its tokens +// accepted per step. Below that (low-acceptance prose) the loop runs +// DFLASH_QWEN35_AR_BURST plain-decode steps inside the spec loop (target +// forward on the seed token only, features still captured for the drafter), +// then probes with one spec step. Set DFLASH_QWEN35_SPEC_STEP_RATIO=0 to +// disable the policy. +struct Qwen35AdaptiveSpecPolicy { + float step_ratio = 1.7f; // spec step cost / plain step cost (measured, gfx1201 IQ4_XS w8) + int burst = 40; // plain-decode steps per burst (each burst ends with one spec probe step) + float ema_alpha = 0.1f; // slow EMA: ~10-step memory so bursty acceptance does not flap + bool enabled() const { return step_ratio > 1.0f && burst > 0; } + // Enter a burst only clearly below break-even (hysteresis against noise). + float accept_threshold() const { return 0.8f * (step_ratio - 1.0f); } +}; + +static Qwen35AdaptiveSpecPolicy qwen35_adaptive_spec_policy() { + static const Qwen35AdaptiveSpecPolicy kPolicy = []() { + Qwen35AdaptiveSpecPolicy p; + if (const char * e = std::getenv("DFLASH_QWEN35_SPEC_STEP_RATIO")) { + p.step_ratio = (float)std::atof(e); + } + if (const char * e = std::getenv("DFLASH_QWEN35_AR_BURST")) { + p.burst = std::atoi(e); + } + return p; + }(); + return kPolicy; +} + bool Qwen35Backend::do_spec_decode(int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io, @@ -2678,8 +2709,25 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, auto t_dec0 = std::chrono::steady_clock::now(); + // Adaptive speculation state (see Qwen35AdaptiveSpecPolicy). + const Qwen35AdaptiveSpecPolicy adaptive = qwen35_adaptive_spec_policy(); + // Start well above the burst threshold so an unlucky opening does not + // park a predictable stream in plain decode; low-acceptance text still + // settles into bursts within a couple of dozen steps. + float accepted_ema = 2.0f * adaptive.accept_threshold(); + int ar_burst_left = 0; + int n_ar_burst_steps = 0; + while (n_generated < n_gen) { const int need_commit_budget = n_gen - n_generated; + // Plain-decode step inside the spec loop: no drafter forward, verify + // the seed token only. Features are still captured, so the drafter + // resumes cleanly on the next probe step. + const bool ar_step = adaptive.enabled() && ar_burst_left > 0; + if (ar_step) { + ar_burst_left--; + n_ar_burst_steps++; + } // Budget hook: no tail-off here. The close-token injection fires // during the emit phase (step 8) after acceptance+replay, mirroring @@ -2718,105 +2766,107 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return false; } - // 2. Draft compute + // 2. Draft compute (skipped on plain-decode burst steps) constexpr int DRAFT_CTX_MAX_DEFAULT = 2048; - const int ring_cap = use_remote_draft ? remote_draft_.ring_cap() : feature_mirror_.cap; - const int draft_ctx = std::min(committed, - std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max))); - const int draft_start = committed - draft_ctx; - int mirror_slot0 = 0; - const bool use_mirror_view = - !use_remote_draft && - draft_feature_mirror_can_view(feature_mirror_, committed, draft_ctx, mirror_slot0); - - const auto profile_draft_start = profile_start(); - if (use_remote_draft) { - local_hidden.clear(); - if (!remote_draft_.propose(committed, draft_ctx, noise_embed, local_hidden)) { - std::fprintf(stderr, "spec-decode: remote draft propose failed\n"); - step_graph_destroy(draft_sg); - return false; - } - } else { - // [TAG_DRAFT_KV] ring-cached drafter context KV: append newly - // committed rows instead of re-encoding the whole feature window. - static const bool draft_kv_on = []() { - const char * e = std::getenv("DFLASH_DRAFT_KV"); - return !(e && e[0] == '0' && e[1] == '\0'); - }(); - bool use_draft_kv = draft_kv_on && feature_mirror_.target_feat != nullptr; - if (use_draft_kv && draft_kv_.gf && - draft_kv_.built_for != (const void *)&dw_) { - draft_kv_free(draft_kv_); - } - if (use_draft_kv && !draft_kv_.gf) { - const int kv_cap = std::min(ring_cap, - std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)); - if (!draft_kv_init(draft_kv_, dw_, draft_backend_, kv_cap, nullptr)) { - draft_kv_free(draft_kv_); - use_draft_kv = false; - std::fprintf(stderr, - "spec-decode: draft-kv init failed; using legacy draft path\n"); - } - } - if (use_draft_kv) { - if (!draft_kv_begin_step(draft_kv_, dw_, draft_backend_, - feature_mirror_, committed)) { - std::fprintf(stderr, "spec-decode: draft-kv step prep failed\n"); - step_graph_destroy(draft_sg); - return false; - } - ggml_backend_tensor_set(draft_kv_.inp_embed, noise_embed.data(), 0, - sizeof(float) * noise_embed.size()); - if (ggml_backend_graph_compute(draft_backend_, draft_kv_.gf) != - GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "spec-decode: draft-kv compute failed\n"); + if (!ar_step) { + const int ring_cap = use_remote_draft ? remote_draft_.ring_cap() : feature_mirror_.cap; + const int draft_ctx = std::min(committed, + std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max))); + const int draft_start = committed - draft_ctx; + int mirror_slot0 = 0; + const bool use_mirror_view = + !use_remote_draft && + draft_feature_mirror_can_view(feature_mirror_, committed, draft_ctx, mirror_slot0); + + const auto profile_draft_start = profile_start(); + if (use_remote_draft) { + local_hidden.clear(); + if (!remote_draft_.propose(committed, draft_ctx, noise_embed, local_hidden)) { + std::fprintf(stderr, "spec-decode: remote draft propose failed\n"); step_graph_destroy(draft_sg); return false; } - local_hidden.resize((size_t)hidden * q_len); - ggml_backend_tensor_get(draft_kv_.hidden_states, local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); } else { - if (!build_draft_step(draft_sg, dw_, /*lm_head=*/nullptr, draft_backend_, - draft_ctx, use_mirror_view ? &feature_mirror_ : nullptr, - committed, - /*ctx_len_max=*/std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)))) { - std::fprintf(stderr, "spec-decode: draft build failed\n"); - step_graph_destroy(draft_sg); - return false; - } - if (!use_mirror_view && - !copy_feature_ring_range_to_tensor(feature_mirror_, draft_sg.target_hidden_cat, - draft_start, draft_ctx)) { - std::fprintf(stderr, "spec-decode: feature copy failed\n"); - step_graph_destroy(draft_sg); - return false; + // [TAG_DRAFT_KV] ring-cached drafter context KV: append newly + // committed rows instead of re-encoding the whole feature window. + static const bool draft_kv_on = []() { + const char * e = std::getenv("DFLASH_DRAFT_KV"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool use_draft_kv = draft_kv_on && feature_mirror_.target_feat != nullptr; + if (use_draft_kv && draft_kv_.gf && + draft_kv_.built_for != (const void *)&dw_) { + draft_kv_free(draft_kv_); } - ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0, - sizeof(float) * noise_embed.size()); - pos_k.resize((size_t)draft_ctx + q_len); - for (int i = 0; i < q_len; i++) pos_q[i] = draft_ctx + i; - for (int i = 0; i < draft_ctx + q_len; i++) pos_k[i] = i; - ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0, - sizeof(int32_t) * pos_q.size()); - ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0, - sizeof(int32_t) * pos_k.size()); - - auto st = ggml_backend_graph_compute(draft_backend_, draft_sg.gf); - if (st != GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "spec-decode: draft compute failed\n"); - step_graph_destroy(draft_sg); - return false; + if (use_draft_kv && !draft_kv_.gf) { + const int kv_cap = std::min(ring_cap, + std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)); + if (!draft_kv_init(draft_kv_, dw_, draft_backend_, kv_cap, nullptr)) { + draft_kv_free(draft_kv_); + use_draft_kv = false; + std::fprintf(stderr, + "spec-decode: draft-kv init failed; using legacy draft path\n"); + } } + if (use_draft_kv) { + if (!draft_kv_begin_step(draft_kv_, dw_, draft_backend_, + feature_mirror_, committed)) { + std::fprintf(stderr, "spec-decode: draft-kv step prep failed\n"); + step_graph_destroy(draft_sg); + return false; + } + ggml_backend_tensor_set(draft_kv_.inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + if (ggml_backend_graph_compute(draft_backend_, draft_kv_.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "spec-decode: draft-kv compute failed\n"); + step_graph_destroy(draft_sg); + return false; + } + local_hidden.resize((size_t)hidden * q_len); + ggml_backend_tensor_get(draft_kv_.hidden_states, local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + } else { + if (!build_draft_step(draft_sg, dw_, /*lm_head=*/nullptr, draft_backend_, + draft_ctx, use_mirror_view ? &feature_mirror_ : nullptr, + committed, + /*ctx_len_max=*/std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, cfg_.draft_ctx_max)))) { + std::fprintf(stderr, "spec-decode: draft build failed\n"); + step_graph_destroy(draft_sg); + return false; + } + if (!use_mirror_view && + !copy_feature_ring_range_to_tensor(feature_mirror_, draft_sg.target_hidden_cat, + draft_start, draft_ctx)) { + std::fprintf(stderr, "spec-decode: feature copy failed\n"); + step_graph_destroy(draft_sg); + return false; + } + ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + pos_k.resize((size_t)draft_ctx + q_len); + for (int i = 0; i < q_len; i++) pos_q[i] = draft_ctx + i; + for (int i = 0; i < draft_ctx + q_len; i++) pos_k[i] = i; + ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0, + sizeof(int32_t) * pos_q.size()); + ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0, + sizeof(int32_t) * pos_k.size()); + + auto st = ggml_backend_graph_compute(draft_backend_, draft_sg.gf); + if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "spec-decode: draft compute failed\n"); + step_graph_destroy(draft_sg); + return false; + } - // Read draft hidden states to host for LM-head projection. - local_hidden.resize((size_t)hidden * q_len); - ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); + // Read draft hidden states to host for LM-head projection. + local_hidden.resize((size_t)hidden * q_len); + ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + } } - } - profile_add(profile_draft_s, profile_draft_start); + profile_add(profile_draft_s, profile_draft_start); + } // !ar_step // ── DDTree tree-structured verify ──────────────────────────────── // When --ddtree is on and the target supports tree verify, build a @@ -2848,7 +2898,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, kvflash_pager_.identity_prefix_covers(committed)); const bool use_tree_verify = cfg_.ddtree_mode && target->supports_tree_verify() && kvflash_tree_ok && - !use_remote_draft && q_len > 1 && tree_special_inactive; + !use_remote_draft && q_len > 1 && tree_special_inactive && !ar_step; // Chain-verify length for this step. The DSpark confidence gate may // truncate the drafted block (adaptive block length); q_len stays the @@ -2856,7 +2906,11 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int v_len = q_len; // DDTree consumes top-K rows directly. Avoid projecting the same // hidden block once for argmax and again for top-K on every step. - if (!use_tree_verify) { + if (ar_step) { + draft_tok.assign(1, last_tok); + v_len = 1; + } else if (!use_tree_verify) { + const auto profile_project_start = profile_start(); // DSpark heads (markov bigram correction + optional confidence // gate) when the drafter ships them; mirrors the laguna hook. bool used_dspark = false; @@ -2875,16 +2929,41 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return !(e && e[0] == '0' && e[1] == '\0'); }(); bool ds_ok = false; - if (fused_dspark && qwen35_dspark_confidence_threshold() <= 0.0f) { + const float conf_threshold = qwen35_dspark_confidence_threshold(); + if (fused_dspark) { + // One graph for every candidate: markov-corrected tokens + // plus (when gated) the confidence score per position. + std::vector conf_scores; ds_ok = dspark_markov_correct_greedy_chain_fused( dw_, draft_backend_, target->lm_head_tensor(), - local_hidden.data(), q_len, last_tok, draft_tok); + local_hidden.data(), q_len, last_tok, draft_tok, + conf_threshold > 0.0f ? &conf_scores : nullptr); + if (ds_ok && conf_threshold > 0.0f) { + // Truncate the chain at the first low-confidence + // position: draft_tok[0] is the seed, candidate i + // scores conf_scores[i-1]. + size_t keep = 1; + while (keep < draft_tok.size() && + keep - 1 < conf_scores.size() && + conf_scores[keep - 1] >= conf_threshold) { + ++keep; + } + static const bool conf_debug = []() { + const char * e = std::getenv("DFLASH_QWEN35_DSPARK_CONF_DEBUG"); + return e && e[0] == '1'; + }(); + if (conf_debug) { + std::fprintf(stderr, "[dspark-conf] keep=%zu/%zu:", keep, draft_tok.size()); + for (float c : conf_scores) std::fprintf(stderr, " %.3f", c); + std::fprintf(stderr, "\n"); + } + draft_tok.resize(keep); + } } if (!ds_ok) { ds_ok = dspark_markov_correct_greedy_chain(dw_, draft_backend_, *target, local_hidden.data(), q_len, - last_tok, - qwen35_dspark_confidence_threshold(), + last_tok, conf_threshold, draft_tok); } if (ds_ok) { @@ -2909,6 +2988,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } draft_tok[0] = last_tok; } + profile_add(profile_project_s, profile_project_start); } if (use_tree_verify) { @@ -3210,13 +3290,18 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, io.observer("draft", draft_tok); } - // 4. Verify: snapshot KV, run target forward over draft tokens - if (!target->snapshot_kv()) { + // 4. Verify: snapshot KV, run target forward over draft tokens. + // A plain-decode step verifies only the (always accepted) seed, so + // it never rolls back: skip the snapshot copy. + const auto profile_snapshot_start = profile_start(); + if (!ar_step && !target->snapshot_kv()) { step_graph_destroy(draft_sg); return false; } + profile_add(profile_snapshot_s, profile_snapshot_start); int verify_last_tok = -1; + const auto profile_verify_start = profile_start(); if (!target->verify_batch(draft_tok, committed, verify_last_tok, &target_tok, /*capture_ssm_intermediates=*/true)) { std::fprintf(stderr, "spec-decode: verify failed\n"); @@ -3224,6 +3309,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, step_graph_destroy(draft_sg); return false; } + profile_add(profile_verify_s, profile_verify_start); target_forwards++; // 5. Acceptance. Greedy: longest matching prefix between draft and @@ -3324,7 +3410,14 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int replay_last_tok = -1; bool fast_rolled_back = false; - if (use_fast_rollback) { + if (ar_step) { + // Seed-only verify: the recurrent state already sits after the + // one committed token; nothing to restore. + bonus_tok = -1; + commit_n = std::min(accept_n, need_commit_budget); + replay_last_tok = target_tok[commit_n - 1]; + fast_rolled_back = true; + } else if (use_fast_rollback) { // Fast rollback: restore SSM from captured intermediates, skip replay. // Implicit bonus: target_tok[commit_n-1] seeds next draft as draft_tok[0], // always accepted on next step. @@ -3333,7 +3426,10 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // budget (need_commit_budget), so committing accept_n would emit // more tokens than requested. commit_n was already clamped above. commit_n = std::min(accept_n, need_commit_budget); - if (target->rollback_to(committed, commit_n)) { + const auto profile_rollback_start = profile_start(); + const bool rolled = target->rollback_to(committed, commit_n); + profile_add(profile_rollback_s, profile_rollback_start); + if (rolled) { replay_last_tok = target_tok[commit_n - 1]; fast_rolled_back = true; rollback_diag.record_fast_rollback(accept_n); @@ -3359,11 +3455,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, for (int i = 0; i < commit_n; i++) { replay_batch[i] = (i < accept_n) ? draft_tok[i] : bonus_tok; } + const auto profile_replay_start = profile_start(); if (!target->verify_batch(replay_batch, committed, replay_last_tok, nullptr)) { std::fprintf(stderr, "spec-decode: replay failed\n"); step_graph_destroy(draft_sg); return false; } + profile_add(profile_replay_s, profile_replay_start); target_forwards++; } @@ -3380,10 +3478,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, return false; } } else if (feature_mirror_.target_feat && cache_.target_feat) { + const auto profile_feature_start = profile_start(); if (!sync_local_draft_features(committed, commit_n)) { step_graph_destroy(draft_sg); return false; } + profile_add(profile_feature_s, profile_feature_start); } // 8. Emit committed tokens (stop at EOS) @@ -3527,6 +3627,18 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, n_accept_sum += std::min(accept_n, emitted); n_draft_steps++; + // Adaptive policy update on real spec steps: EMA of accepted draft + // tokens (the seed is always accepted); a low EMA schedules a burst + // of plain-decode steps, the step after the burst is a spec probe. + if (adaptive.enabled() && !ar_step) { + const float accepted_drafts = (float)std::max(0, accept_n - 1); + accepted_ema = (1.0f - adaptive.ema_alpha) * accepted_ema + + adaptive.ema_alpha * accepted_drafts; + if (accepted_ema < adaptive.accept_threshold()) { + ar_burst_left = adaptive.burst; + } + } + // Notify observer with accepted tokens for this step. if (io.observer) { io.observer("verify", replay_tok); @@ -3613,6 +3725,11 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, n_generated > 0 ? n_generated / decode_s : 0.0, n_draft_steps, n_accept_sum, total_draft_pos, accept_pct, n_draft_steps > 0 ? (double)n_generated / (double)n_draft_steps : 0.0); + if (n_ar_burst_steps > 0) { + std::fprintf(stderr, "[spec-decode] adaptive: %d of %d steps ran as plain decode " + "(accept threshold %.2f drafts/step, burst %d)\n", + n_ar_burst_steps, n_draft_steps, adaptive.accept_threshold(), adaptive.burst); + } if (tp_profile) { std::fprintf(stderr, "[spec-profile] draft=%.3fs project=%.3fs snapshot=%.3fs " From 56b2e0572ab167be57224e5c75f7f2b71e7fc2db Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 18 Aug 2026 23:44:24 +0200 Subject: [PATCH 25/42] ggml: FA vec kernel splits short KV spans across two blocks launch_fattn was told the vec kernel consumes D keys per step; it walks nthreads (128) per step, so a 256-key window at head_dim 256 ran as one block per head. Passing nthreads lets it use two blocks per head plus the combine pass: Qwen3.8-27B plain decode 34.3 -> 34.6 tok/s on R9700, identical output. --- server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh index bcf1dd804..85a2af718 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn-vec.cuh @@ -534,7 +534,10 @@ void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggm const bool need_f16_K = type_K == GGML_TYPE_F16; const bool need_f16_V = type_V == GGML_TYPE_F16; constexpr size_t nbytes_shared = 0; - launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false); + // The kernel walks the KV sequence in steps of nthreads (not D); telling + // launch_fattn so lets it split a short KV span (e.g. a 256-token window + // at head_dim 256) across two blocks per head instead of one. + launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, nthreads, need_f16_K, need_f16_V, false); } template From 6136339d643388c55674da7a3217a56b04d83bfe Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 19 Aug 2026 00:53:23 +0200 Subject: [PATCH 26/42] qwen35: adaptive policy probe step reacts fast The first spec step after a plain-decode burst updates the acceptance EMA with alpha 0.5 so a stream that became predictable leaves plain decode immediately; step ratio and start value keep the measured best balance (45.7 / 31.8 / 40.4 tok/s code / prose / mixed). --- server/src/qwen35/qwen35_backend.cpp | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 345afeeb1..fd0ed4e6b 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -2497,7 +2497,9 @@ static float qwen35_dspark_confidence_threshold() { // then probes with one spec step. Set DFLASH_QWEN35_SPEC_STEP_RATIO=0 to // disable the policy. struct Qwen35AdaptiveSpecPolicy { - float step_ratio = 1.7f; // spec step cost / plain step cost (measured, gfx1201 IQ4_XS w8) + float step_ratio = 1.7f; // spec/plain step cost; the measured 1.9 (54 vs 28.6 ms) is deliberately + // under-stated: a higher threshold costs more on bursty code/mixed streams than + // it saves on prose (measured 45.6/40.4/32.4 vs 42.5/36.4/32.3 tok/s) int burst = 40; // plain-decode steps per burst (each burst ends with one spec probe step) float ema_alpha = 0.1f; // slow EMA: ~10-step memory so bursty acceptance does not flap bool enabled() const { return step_ratio > 1.0f && burst > 0; } @@ -2712,11 +2714,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // Adaptive speculation state (see Qwen35AdaptiveSpecPolicy). const Qwen35AdaptiveSpecPolicy adaptive = qwen35_adaptive_spec_policy(); // Start well above the burst threshold so an unlucky opening does not - // park a predictable stream in plain decode; low-acceptance text still - // settles into bursts within a couple of dozen steps. + // park a predictable stream in plain decode. The probe step that ends a + // burst updates the EMA with a fast alpha (see below) so a stream that + // turned predictable leaves plain decode quickly. float accepted_ema = 2.0f * adaptive.accept_threshold(); int ar_burst_left = 0; int n_ar_burst_steps = 0; + bool probe_step = false; // first spec step after a burst while (n_generated < n_gen) { const int need_commit_budget = n_gen - n_generated; @@ -2727,6 +2731,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (ar_step) { ar_burst_left--; n_ar_burst_steps++; + probe_step = (ar_burst_left == 0); } // Budget hook: no tail-off here. The close-token injection fires @@ -3632,8 +3637,11 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // of plain-decode steps, the step after the burst is a spec probe. if (adaptive.enabled() && !ar_step) { const float accepted_drafts = (float)std::max(0, accept_n - 1); - accepted_ema = (1.0f - adaptive.ema_alpha) * accepted_ema + - adaptive.ema_alpha * accepted_drafts; + // A probe (first spec step after a burst) weighs its result + // heavily: it is the only evidence about the current text. + const float alpha = probe_step ? 0.5f : adaptive.ema_alpha; + accepted_ema = (1.0f - alpha) * accepted_ema + alpha * accepted_drafts; + probe_step = false; if (accepted_ema < adaptive.accept_threshold()) { ar_burst_left = adaptive.burst; } From b0957a75a71f59b14f1d957d950ce456596d93c5 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 19 Aug 2026 01:36:54 +0200 Subject: [PATCH 27/42] qwen35: adaptive policy uses the measured spec/plain step-time ratio The break-even acceptance now follows live EMAs of the spec-step and plain-step wall times (default 1.9 until both are measured), so it is right for any drafter block size (width-8 DSpark and width-16 DFlash measure ~1.8 on gfx1201). --- server/src/qwen35/qwen35_backend.cpp | 30 +++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index fd0ed4e6b..a17bd4d4a 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -2497,14 +2497,15 @@ static float qwen35_dspark_confidence_threshold() { // then probes with one spec step. Set DFLASH_QWEN35_SPEC_STEP_RATIO=0 to // disable the policy. struct Qwen35AdaptiveSpecPolicy { - float step_ratio = 1.7f; // spec/plain step cost; the measured 1.9 (54 vs 28.6 ms) is deliberately - // under-stated: a higher threshold costs more on bursty code/mixed streams than - // it saves on prose (measured 45.6/40.4/32.4 vs 42.5/36.4/32.3 tok/s) + float step_ratio = 1.9f; // spec/plain step cost used until both step kinds have been timed + // (measured 54-55 vs 28.6 ms on gfx1201 for width-8 and width-16 verify) int burst = 40; // plain-decode steps per burst (each burst ends with one spec probe step) float ema_alpha = 0.1f; // slow EMA: ~10-step memory so bursty acceptance does not flap bool enabled() const { return step_ratio > 1.0f && burst > 0; } // Enter a burst only clearly below break-even (hysteresis against noise). - float accept_threshold() const { return 0.8f * (step_ratio - 1.0f); } + // `ratio` is the live spec/plain step-time ratio once measured. + float accept_threshold(float ratio) const { return 0.8f * (ratio - 1.0f); } + float accept_threshold() const { return accept_threshold(step_ratio); } }; static Qwen35AdaptiveSpecPolicy qwen35_adaptive_spec_policy() { @@ -2721,6 +2722,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int ar_burst_left = 0; int n_ar_burst_steps = 0; bool probe_step = false; // first spec step after a burst + // Live step-time EMAs (seconds) for the break-even ratio; 0 = not yet measured. + double t_spec_step_ema = 0.0; + double t_ar_step_ema = 0.0; + auto live_step_ratio = [&]() { + return (t_spec_step_ema > 0.0 && t_ar_step_ema > 0.0) + ? (float)(t_spec_step_ema / t_ar_step_ema) : adaptive.step_ratio; + }; while (n_generated < n_gen) { const int need_commit_budget = n_gen - n_generated; @@ -2733,6 +2741,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, n_ar_burst_steps++; probe_step = (ar_burst_left == 0); } + const auto t_step_start = std::chrono::steady_clock::now(); // Budget hook: no tail-off here. The close-token injection fires // during the emit phase (step 8) after acceptance+replay, mirroring @@ -3635,6 +3644,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // Adaptive policy update on real spec steps: EMA of accepted draft // tokens (the seed is always accepted); a low EMA schedules a burst // of plain-decode steps, the step after the burst is a spec probe. + if (adaptive.enabled()) { + const double t_step = std::chrono::duration( + std::chrono::steady_clock::now() - t_step_start).count(); + double & t_ema = ar_step ? t_ar_step_ema : t_spec_step_ema; + t_ema = (t_ema > 0.0) ? 0.9 * t_ema + 0.1 * t_step : t_step; + } if (adaptive.enabled() && !ar_step) { const float accepted_drafts = (float)std::max(0, accept_n - 1); // A probe (first spec step after a burst) weighs its result @@ -3642,7 +3657,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, const float alpha = probe_step ? 0.5f : adaptive.ema_alpha; accepted_ema = (1.0f - alpha) * accepted_ema + alpha * accepted_drafts; probe_step = false; - if (accepted_ema < adaptive.accept_threshold()) { + if (accepted_ema < adaptive.accept_threshold(live_step_ratio())) { ar_burst_left = adaptive.burst; } } @@ -3735,8 +3750,9 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, n_draft_steps > 0 ? (double)n_generated / (double)n_draft_steps : 0.0); if (n_ar_burst_steps > 0) { std::fprintf(stderr, "[spec-decode] adaptive: %d of %d steps ran as plain decode " - "(accept threshold %.2f drafts/step, burst %d)\n", - n_ar_burst_steps, n_draft_steps, adaptive.accept_threshold(), adaptive.burst); + "(step ratio %.2f, accept threshold %.2f drafts/step, burst %d)\n", + n_ar_burst_steps, n_draft_steps, live_step_ratio(), + adaptive.accept_threshold(live_step_ratio()), adaptive.burst); } if (tp_profile) { std::fprintf(stderr, From fb63013a8a4608d0d811d2a0decb30baa6f373a0 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 19 Aug 2026 02:22:19 +0200 Subject: [PATCH 28/42] qwen35: DFlash 2 drafter support (dynamic convs + candidate selector) DFlash 2 (z-lab/inco, e.g. z-lab/Qwen3.8-27B-DFlash2) is the DFlash backbone plus a grouped dynamic causal conv around attention and MLP in every layer and a candidate selector head (top-k lm_head candidates per block position, one path scored by a low-rank bigram form). - converter: maps attention_conv/mlp_conv (base kernels F32, kernel projections) and candidate_selector tensors, emits dflash2.* metadata, reads block_size from dflash_config, emits SWA pattern for drafters with causal sliding layers. - loader: DraftConvWeights per layer, DraftSelectorWeights, shape checks. - draft graph: conv prepare/finish (two taps over the block, per-element base + per-group dynamic coefficient) in both the stateless and the cached-KV builders. - selector chain: top-k via the target's GPU top-k (kMaxK 8 -> 16), one cached graph for hproj + codebook row gathers, host path search. - spec loop uses the selector before the DSpark/argmax paths. Qwen3.8-27B IQ4_XS on R9700, q8_0 drafter, greedy: 109.9 code / 50.7 prose / 111.8 mixed tok/s (DSpark drafter: 45.6 / 32.4 / 38.6); avg 5.9-6.0 accepted tokens per 8-token block on code, ~2.7 on prose. --- server/CMakeLists.txt | 1 + server/scripts/convert_dflash_to_gguf.py | 38 ++++- server/src/common/dflash2_head.cpp | 156 ++++++++++++++++++ server/src/common/dflash2_head.h | 29 ++++ .../src/common/geometric_draft_topk_cuda.cu | 3 +- server/src/draft/draft_gguf_loader.cpp | 82 ++++++++- server/src/draft/draft_graph.cpp | 103 +++++++++++- server/src/internal.h | 32 ++++ server/src/qwen35/qwen35_backend.cpp | 27 ++- 9 files changed, 463 insertions(+), 8 deletions(-) create mode 100644 server/src/common/dflash2_head.cpp create mode 100644 server/src/common/dflash2_head.h diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index d9ba62836..57f2ff462 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -439,6 +439,7 @@ add_library(dflash_common STATIC src/common/dynamic_backend.cpp src/common/domino_head.cpp src/common/dspark_head.cpp + src/common/dflash2_head.cpp src/common/target_shard_ipc.cpp src/common/target_shard_ipc_daemon.cpp src/common/dflash_feature_ring.cpp diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index 5ebc5cdef..74f71bce2 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -124,6 +124,23 @@ def pick(*keys): a["yarn_beta_slow"] = float(rp.get("beta_slow", 1.0)) if dfc.get("mask_token_id") is not None: a["mask_token_id"] = int(dfc["mask_token_id"]) + if dfc.get("block_size") is not None: + a["block_size"] = int(dfc["block_size"]) + # DFlash 2 (z-lab/inco): grouped dynamic convs + candidate selector. + if dfc.get("conv_kernel_size") is not None: + a["conv_kernel_size"] = int(dfc["conv_kernel_size"]) + a["conv_group_size"] = int(dfc.get("conv_group_size", 16)) + if dfc.get("selector_rank") is not None: + a["selector_rank"] = int(dfc["selector_rank"]) + a["selector_top_k"] = int(dfc.get("selector_top_k", 16)) + # Per-layer sliding-window / causal attention (Qwen3.6-style drafters + # and DFlash 2). HF: layer_types + sliding_window; a top-level + # is_causal=false (DFlash 2) makes every layer bidirectional, which is + # our default (no SWA pattern emitted). + lt = c.get("layer_types") + if lt and c.get("sliding_window") and c.get("is_causal", None) is not False: + a["swa_window"] = int(c["sliding_window"]) + a["swa_pattern"] = [str(x) == "sliding_attention" for x in lt] print(f"[info] read arch from {cfg_path}") else: print(f"[warn] no config.json next to safetensors; using 27B defaults") @@ -196,8 +213,17 @@ def map_name(name: str) -> str | None: "mlp.gate_proj.weight": f"blk.{i}.ffn_gate.weight", "mlp.up_proj.weight": f"blk.{i}.ffn_up.weight", "mlp.down_proj.weight": f"blk.{i}.ffn_down.weight", + # DFlash 2 grouped dynamic convs + "attention_conv.base_kernel": f"blk.{i}.attn_conv.base", + "attention_conv.kernel_projection.weight": f"blk.{i}.attn_conv.proj.weight", + "mlp_conv.base_kernel": f"blk.{i}.ffn_conv.base", + "mlp_conv.kernel_projection.weight": f"blk.{i}.ffn_conv.proj.weight", } return layer_map.get(rest) + # DFlash 2 candidate selector + if name == "candidate_selector.hidden_projection.weight": return "dflash.selector.hproj.weight" + if name == "candidate_selector.predecessor_codebook": return "dflash.selector.pred_cb" + if name == "candidate_selector.successor_codebook": return "dflash.selector.succ_cb" return None @@ -522,6 +548,15 @@ def main(): elif _cap_ids: print(f"[warn] capture_layer_ids len {len(_cap_ids)} != n_target_layers " f"{a['n_target_layers']}; not embedding ids", file=sys.stderr) + if a.get("swa_pattern"): + writer.add_uint32(f"{ARCH}.attention.sliding_window", a["swa_window"]) + writer.add_array(f"{ARCH}.attention.sliding_window_pattern", [bool(x) for x in a["swa_pattern"]]) + if a.get("conv_kernel_size"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_kernel_size", a["conv_kernel_size"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_group_size", a["conv_group_size"]) + if a.get("selector_rank"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_rank", a["selector_rank"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_top_k", a["selector_top_k"]) # Walk + add tensors. Sort: dflash.* singletons first, then output_*, # then per-layer in numeric order — keeps the on-disk layout stable. @@ -560,7 +595,8 @@ def sort_key(t): is_norm = ( gguf_name.endswith("_norm.weight") or gguf_name == "output_norm.weight" or - gguf_name == "dflash.hidden_norm.weight" + gguf_name == "dflash.hidden_norm.weight" or + gguf_name.endswith("_conv.base") # DFlash 2 conv base kernels [2, K, hidden] ) if is_norm: arr = arr.astype(" +#include +#include +#include + +namespace dflash::common { + +namespace { + +// Selector projection graph, built once per (drafter, backend, n_cand, K). +struct SelectorGraph { + const DraftWeights * dw = nullptr; + ggml_backend_t backend = nullptr; + int n_cand = 0; + int K = 0; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_tensor * inp_hidden = nullptr; + ggml_tensor * inp_succ = nullptr; + ggml_tensor * inp_pred = nullptr; + ggml_tensor * hproj = nullptr; + ggml_tensor * succ = nullptr; + ggml_tensor * pred = nullptr; +}; + +SelectorGraph & selector_graph() { + static thread_local SelectorGraph g; + return g; +} + +void selector_graph_free(SelectorGraph & g) { + if (g.galloc) { ggml_gallocr_free(g.galloc); g.galloc = nullptr; } + if (g.ctx) { ggml_free(g.ctx); g.ctx = nullptr; } + g.gf = nullptr; + g.dw = nullptr; + g.n_cand = 0; + g.K = 0; +} + +} // namespace + +bool dflash2_select_chain(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok) { + const DraftSelectorWeights & sel = dw.selector; + if (!sel.enabled || !sel.hproj || !sel.pred_cb || !sel.succ_cb) return false; + if (q_len <= 1 || !local_hidden || !backend) return false; + const int hdim = dw.n_embd; + const int rank = sel.rank; + const int K = sel.top_k; + const int n_cand = q_len - 1; + if (hdim <= 0 || rank <= 0 || K <= 0) return false; + + // 1. Top-k candidates (log-probs) per block position through the target + // lm_head. Position 0 of local_hidden is the seed slot; candidates are + // rows 1 .. q_len-1. + std::vector cand_lp; + std::vector cand_ids; + if (!target.project_hidden_to_topk(local_hidden + (size_t)hdim, n_cand, K, /*temperature=*/1.0f, + cand_lp, cand_ids)) { + return false; + } + if (cand_lp.size() != (size_t)n_cand * K || cand_ids.size() != (size_t)n_cand * K) return false; + + // 2. One graph on the draft backend: hproj(h) for every candidate position, + // successor rows for every candidate, predecessor rows for the seed and + // every candidate (the path picks its predecessor among them). The + // graph shape only depends on (n_cand, K), so it is built once and + // reused across steps. + const int n_rows_pred = 1 + n_cand * K; + SelectorGraph & g = selector_graph(); + if (!g.ctx || g.dw != &dw || g.backend != backend || g.n_cand != n_cand || g.K != K) { + selector_graph_free(g); + const size_t arena_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead() + 4096; + g.arena.assign(arena_size, 0); + ggml_init_params ip{}; + ip.mem_size = g.arena.size(); + ip.mem_buffer = g.arena.data(); + ip.no_alloc = true; + g.ctx = ggml_init(ip); + if (!g.ctx) return false; + g.gf = ggml_new_graph(g.ctx); + g.inp_hidden = ggml_new_tensor_2d(g.ctx, GGML_TYPE_F32, hdim, n_cand); + g.inp_succ = ggml_new_tensor_1d(g.ctx, GGML_TYPE_I32, n_cand * K); + g.inp_pred = ggml_new_tensor_1d(g.ctx, GGML_TYPE_I32, n_rows_pred); + ggml_set_input(g.inp_hidden); + ggml_set_input(g.inp_succ); + ggml_set_input(g.inp_pred); + g.hproj = ggml_mul_mat(g.ctx, sel.hproj, g.inp_hidden); // [rank, n_cand] + g.succ = ggml_get_rows(g.ctx, sel.succ_cb, g.inp_succ); // [rank, n_cand*K] f32 + g.pred = ggml_get_rows(g.ctx, sel.pred_cb, g.inp_pred); // [rank, 1+n_cand*K] f32 + ggml_set_output(g.hproj); + ggml_set_output(g.succ); + ggml_set_output(g.pred); + ggml_build_forward_expand(g.gf, g.hproj); + ggml_build_forward_expand(g.gf, g.succ); + ggml_build_forward_expand(g.gf, g.pred); + g.galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!g.galloc || !ggml_gallocr_alloc_graph(g.galloc, g.gf)) { + std::fprintf(stderr, "dflash2_select_chain: gallocr_alloc_graph failed\n"); + selector_graph_free(g); + return false; + } + g.dw = &dw; g.backend = backend; g.n_cand = n_cand; g.K = K; + } + + std::vector pred_ids((size_t)n_rows_pred); + pred_ids[0] = last_tok; + std::memcpy(pred_ids.data() + 1, cand_ids.data(), sizeof(int32_t) * (size_t)n_cand * K); + ggml_backend_tensor_set(g.inp_hidden, local_hidden + (size_t)hdim, 0, sizeof(float) * (size_t)hdim * n_cand); + ggml_backend_tensor_set(g.inp_succ, cand_ids.data(), 0, sizeof(int32_t) * (size_t)n_cand * K); + ggml_backend_tensor_set(g.inp_pred, pred_ids.data(), 0, sizeof(int32_t) * (size_t)n_rows_pred); + if (ggml_backend_graph_compute(backend, g.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "dflash2_select_chain: graph_compute failed\n"); + return false; + } + std::vector h_hproj((size_t)rank * n_cand); + std::vector h_succ((size_t)rank * n_cand * K); + std::vector h_pred((size_t)rank * n_rows_pred); + ggml_backend_tensor_get_async(backend, g.hproj, h_hproj.data(), 0, sizeof(float) * h_hproj.size()); + ggml_backend_tensor_get_async(backend, g.succ, h_succ.data(), 0, sizeof(float) * h_succ.size()); + ggml_backend_tensor_get_async(backend, g.pred, h_pred.data(), 0, sizeof(float) * h_pred.size()); + ggml_backend_synchronize(backend); + + // 3. Path search: greedy over the candidates, conditioned on the previous pick. + draft_tok.assign((size_t)q_len, last_tok); + int prev_row = 0; // row in h_pred: 0 = seed, 1 + i*K + k = candidate k of position i + for (int i = 0; i < n_cand; ++i) { + const float * pr = h_pred.data() + (size_t)prev_row * rank; + const float * hp = h_hproj.data() + (size_t)i * rank; + float best = -INFINITY; + int best_k = 0; + for (int k = 0; k < K; ++k) { + const float * sc = h_succ.data() + ((size_t)i * K + k) * rank; + float dot = 0.0f; + for (int r = 0; r < rank; ++r) dot += pr[r] * hp[r] * sc[r]; + const float score = cand_lp[(size_t)i * K + k] + dot; + if (score > best) { best = score; best_k = k; } + } + draft_tok[(size_t)i + 1] = cand_ids[(size_t)i * K + best_k]; + prev_row = 1 + i * K + best_k; + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash2_head.h b/server/src/common/dflash2_head.h new file mode 100644 index 000000000..446a8646c --- /dev/null +++ b/server/src/common/dflash2_head.h @@ -0,0 +1,29 @@ +#pragma once + +#include "dflash_target.h" +#include "internal.h" + +#include +#include + +namespace dflash::common { + +// DFlash 2 candidate selector for greedy chain drafting. +// +// For every drafted block position the target lm_head logits are reduced to +// the selector's top-k candidates (log-probs, so per-position constants do +// not matter for the argmax), then one path is traced through them: +// score(c) = logp(c) + < pred_cb[prev] * hproj(h_pos), succ_cb[c] > +// prev = argmax_c score(c) +// starting from the block seed `last_tok`. Runs the projections (hproj GEMV +// and codebook row gathers) in one small graph on `backend`, the k-way path +// search on the host. Fills draft_tok = [last_tok, tok_1 .. tok_{q_len-1}]. +bool dflash2_select_chain(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok); + +} // namespace dflash::common diff --git a/server/src/common/geometric_draft_topk_cuda.cu b/server/src/common/geometric_draft_topk_cuda.cu index 71086c98a..ba287656d 100644 --- a/server/src/common/geometric_draft_topk_cuda.cu +++ b/server/src/common/geometric_draft_topk_cuda.cu @@ -13,7 +13,7 @@ namespace dflash::common { namespace { -constexpr int kMaxK = 8; // ddtree_K is 8 in practice; K>kMaxK → CPU fallback +constexpr int kMaxK = 16; // ddtree_K is 8, the DFlash 2 selector uses 16; K>kMaxK → CPU fallback constexpr int kBlock = 256; // threads per block (power of two for the reduction) constexpr int kMaxSplit = 128; // max vocab splits per position (combine-block cap) @@ -380,6 +380,7 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, switch (K) { DFLASH_TOPK_CASE(1) DFLASH_TOPK_CASE(2) DFLASH_TOPK_CASE(3) DFLASH_TOPK_CASE(4) DFLASH_TOPK_CASE(5) DFLASH_TOPK_CASE(6) DFLASH_TOPK_CASE(7) DFLASH_TOPK_CASE(8) + DFLASH_TOPK_CASE(12) DFLASH_TOPK_CASE(16) default: break; } #undef DFLASH_TOPK_CASE diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index c882adfd9..58c203195 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -75,7 +75,7 @@ int count_attn_gate_layers(const DraftWeights & w) { bool check_shape_1d(const ggml_tensor * t, int64_t ne0, const char * name, char * buf, size_t buf_sz) { if (!t || t->ne[0] != ne0) { - std::snprintf(buf, buf_sz, "draft GGUF: Domino tensor %s shape mismatch: got [%lld], expected [%lld]", + std::snprintf(buf, buf_sz, "draft GGUF: tensor %s shape mismatch: got [%lld], expected [%lld]", name, t ? (long long)t->ne[0] : -1LL, (long long)ne0); return false; } @@ -86,7 +86,7 @@ bool check_shape_2d(const ggml_tensor * t, int64_t ne0, int64_t ne1, const char * name, char * buf, size_t buf_sz) { if (!t || t->ne[0] != ne0 || t->ne[1] != ne1) { std::snprintf(buf, buf_sz, - "draft GGUF: Domino tensor %s shape mismatch: got [%lld,%lld], expected [%lld,%lld]", + "draft GGUF: tensor %s shape mismatch: got [%lld,%lld], expected [%lld,%lld]", name, t ? (long long)t->ne[0] : -1LL, t ? (long long)t->ne[1] : -1LL, @@ -96,6 +96,21 @@ bool check_shape_2d(const ggml_tensor * t, int64_t ne0, int64_t ne1, return true; } +bool check_shape_3d(const ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t ne2, + const char * name, char * buf, size_t buf_sz) { + if (!t || t->ne[0] != ne0 || t->ne[1] != ne1 || t->ne[2] != ne2) { + std::snprintf(buf, buf_sz, + "draft GGUF: tensor %s shape mismatch: got [%lld,%lld,%lld], expected [%lld,%lld,%lld]", + name, + t ? (long long)t->ne[0] : -1LL, + t ? (long long)t->ne[1] : -1LL, + t ? (long long)t->ne[2] : -1LL, + (long long)ne0, (long long)ne1, (long long)ne2); + return false; + } + return true; +} + } // namespace bool load_draft_gguf(const std::string & path, @@ -327,6 +342,11 @@ bool load_draft_gguf(const std::string & path, L.w_gate = fnd("ffn_gate.weight"); L.w_up = fnd("ffn_up.weight"); L.w_down = fnd("ffn_down.weight"); + // DFlash 2 grouped dynamic convs (optional) + L.attn_conv.base = fnd("attn_conv.base"); + L.attn_conv.proj = fnd("attn_conv.proj.weight"); + L.mlp_conv.base = fnd("ffn_conv.base"); + L.mlp_conv.proj = fnd("ffn_conv.proj.weight"); if (!L.attn_norm || !L.ffn_norm || !L.wq || !L.wk || !L.wv || !L.wo || !L.q_norm || !L.k_norm || !L.w_gate || !L.w_up || !L.w_down) { char b[128]; @@ -477,6 +497,64 @@ bool load_draft_gguf(const std::string & path, out.dspark.confidence_dim); } + // DFlash 2: dynamic convs in every layer + candidate selector head. + { + const int conv_k = (int)read_u32("dflash.dflash2.conv_kernel_size", 0); + int n_conv = 0; + for (const DraftLayer & L : out.layers) { + if (L.attn_conv.present() && L.mlp_conv.present()) n_conv++; + } + if (n_conv > 0 || conv_k > 0) { + if (n_conv != out.n_layer || conv_k <= 0) { + set_last_error("draft GGUF: DFlash 2 conv tensors/metadata incomplete " + "(need attn_conv/ffn_conv base+proj in every layer and " + "dflash.dflash2.conv_kernel_size)"); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + out.conv_kernel_size = conv_k; + out.conv_group_size = (int)read_u32("dflash.dflash2.conv_group_size", 16); + const DraftLayer & L0 = out.layers[0]; + const int64_t groups = out.n_embd / out.conv_group_size; + char shape_err[192]; + if (!check_shape_3d(L0.attn_conv.base, out.n_embd, conv_k, 2, "attn_conv.base", shape_err, sizeof(shape_err)) || + !check_shape_2d(L0.attn_conv.proj, out.n_embd, 2 * conv_k * groups, "attn_conv.proj", shape_err, sizeof(shape_err))) { + set_last_error(shape_err); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + std::fprintf(stderr, "[draft GGUF] DFlash 2 dynamic convs: kernel=%d group=%d\n", + out.conv_kernel_size, out.conv_group_size); + } + out.selector = DraftSelectorWeights{}; + out.selector.hproj = g("dflash.selector.hproj.weight"); + out.selector.pred_cb = g("dflash.selector.pred_cb"); + out.selector.succ_cb = g("dflash.selector.succ_cb"); + const uint32_t sel_rank = read_u32("dflash.dflash2.selector_rank", 0); + if (out.selector.hproj || out.selector.pred_cb || out.selector.succ_cb || sel_rank) { + if (!out.selector.hproj || !out.selector.pred_cb || !out.selector.succ_cb) { + set_last_error("draft GGUF: DFlash 2 selector tensors incomplete " + "(hproj.weight, pred_cb, succ_cb)"); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + out.selector.rank = sel_rank ? (int)sel_rank : (int)out.selector.hproj->ne[1]; + out.selector.top_k = (int)read_u32("dflash.dflash2.selector_top_k", 16); + char shape_err[192]; + const int64_t R = out.selector.rank; + if (!check_shape_2d(out.selector.hproj, out.n_embd, R, "selector.hproj", shape_err, sizeof(shape_err)) || + !check_shape_2d(out.selector.pred_cb, R, out.selector.pred_cb->ne[1], "selector.pred_cb", shape_err, sizeof(shape_err)) || + !check_shape_2d(out.selector.succ_cb, R, out.selector.pred_cb->ne[1], "selector.succ_cb", shape_err, sizeof(shape_err))) { + set_last_error(shape_err); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } + out.selector.enabled = true; + std::fprintf(stderr, "[draft GGUF] DFlash 2 selector enabled: rank=%d top_k=%d vocab=%lld\n", + out.selector.rank, out.selector.top_k, (long long)out.selector.pred_cb->ne[1]); + } + } + // GGUF Qwen3.6 drafters carry SWA metadata emitted by the converter: // dflash-draft.attention.sliding_window = 2048 // dflash-draft.attention.sliding_window_pattern = [true,true,true,true,false] diff --git a/server/src/draft/draft_graph.cpp b/server/src/draft/draft_graph.cpp index 6ad18e8ec..0178fd40b 100644 --- a/server/src/draft/draft_graph.cpp +++ b/server/src/draft/draft_graph.cpp @@ -85,6 +85,69 @@ static ggml_tensor * draft_fuse_features( return target_feat; } +// ── DFlash 2 grouped dynamic causal conv ──────────────────────────── +// +// Two taps over the draft block (positions within the block; the block's +// first slot has no predecessor). For each tap k the coefficient is a +// per-element base kernel plus a per-group dynamic kernel projected from +// the block's normalized hidden state: +// dyn = proj @ x_norm [2*K*groups, q_len] +// coef_s_k = base[s][k] (per element) + dyn[s][k] (per group, broadcast) +// out = sum_k coef_s_k * shift_k(x) +// s = 0 ("prepare", applied to the sub-block input) or 1 ("finish", applied +// to the sub-block output); both use the dyn computed from the input. +struct DraftDynConv { + ggml_tensor * dyn = nullptr; // [2*K*groups, q_len] +}; + +static DraftDynConv draft_dyn_conv_kernel(ggml_context * ctx, + const DraftConvWeights & cw, + ggml_tensor * x_norm) { + DraftDynConv dc; + dc.dyn = ggml_mul_mat(ctx, cw.proj, x_norm); // [2*K*groups, q_len] + return dc; +} + +static ggml_tensor * draft_dyn_conv_apply(ggml_context * ctx, + const DraftWeights & w, + const DraftConvWeights & cw, + const DraftDynConv & dc, + int s, // 0 = prepare, 1 = finish + ggml_tensor * x) { // [hidden, q_len] + const int64_t hidden = x->ne[0]; + const int64_t q_len = x->ne[1]; + const int K = w.conv_kernel_size; + const int64_t gs = w.conv_group_size; + const int64_t groups = hidden / gs; + const size_t e = ggml_element_size(dc.dyn); + + ggml_tensor * out = nullptr; + for (int k = 0; k < K; ++k) { + // shift_k(x): column l takes x[:, l-k], zero for l < k + ggml_tensor * xs = x; + if (k > 0) { + if (q_len <= k) break; + ggml_tensor * head = ggml_view_2d(ctx, x, hidden, q_len - k, x->nb[1], 0); + xs = ggml_pad_ext(ctx, head, 0, 0, k, 0, 0, 0, 0, 0); // [hidden, q_len] + } + // per-group dynamic coefficient for (s, k): rows [(s*K+k)*groups, +groups) + ggml_tensor * dyn_sk = ggml_view_3d(ctx, dc.dyn, 1, groups, q_len, + e, dc.dyn->nb[1], + (size_t)((s * K + k) * groups) * e); + ggml_tensor * xs3 = ggml_reshape_3d(ctx, xs, gs, groups, q_len); + ggml_tensor * dyn3 = ggml_repeat(ctx, dyn_sk, xs3); // [gs, groups, q_len] + // per-element base coefficient base[s][k]: [hidden] at offset (s*K+k)*hidden + ggml_tensor * base_sk = ggml_view_3d(ctx, cw.base, gs, groups, 1, + cw.base->nb[0] * gs, cw.base->nb[0] * hidden, + (size_t)(s * K + k) * cw.base->nb[1]); + ggml_tensor * coef = ggml_add(ctx, dyn3, base_sk); // broadcast over q_len + ggml_tensor * term = ggml_mul(ctx, xs3, coef); + term = ggml_reshape_2d(ctx, term, hidden, q_len); + out = out ? ggml_add(ctx, out, term) : term; + } + return out; +} + DraftGraphOutputs build_draft_graph( ggml_context * ctx, const DraftWeights & w, @@ -130,10 +193,18 @@ DraftGraphOutputs build_draft_graph( const int eff_total_k = eff_ctx + q_len; const int ctx_offset = use_swa ? (ctx_len - w.swa_window) : 0; + const bool dyn_conv = w.conv_kernel_size > 0 && L.attn_conv.present() && L.mlp_conv.present(); + if (!disable_attn) { // ── 2a. Attention pre-norm ggml_tensor * hn = ggml_rms_norm(ctx, h, eps); hn = ggml_mul(ctx, hn, L.attn_norm); + // DFlash 2: dynamic conv "prepare" on the attention input + DraftDynConv attn_dc; + if (dyn_conv) { + attn_dc = draft_dyn_conv_kernel(ctx, L.attn_conv, hn); + hn = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 0, hn); + } std::snprintf(probe_name, sizeof(probe_name), "draft_l%d_hn", il); ggml_set_name(hn, probe_name); @@ -241,6 +312,9 @@ DraftGraphOutputs build_draft_graph( // ── 2g. Output projection + residual // wo: [q_dim, hidden] (ne[0]=q_dim, ne[1]=hidden) ggml_tensor * attn_out = ggml_mul_mat(ctx, L.wo, attn); // [hidden, q_len] + if (dyn_conv) { + attn_out = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 1, attn_out); + } std::snprintf(probe_name, sizeof(probe_name), "draft_l%d_attn_out", il); ggml_set_name(attn_out, probe_name); h = ggml_add(ctx, h, attn_out); @@ -252,6 +326,11 @@ DraftGraphOutputs build_draft_graph( // ── 2h. FFN pre-norm ggml_tensor * hf = ggml_rms_norm(ctx, h, eps); hf = ggml_mul(ctx, hf, L.ffn_norm); + DraftDynConv mlp_dc; + if (dyn_conv) { + mlp_dc = draft_dyn_conv_kernel(ctx, L.mlp_conv, hf); + hf = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 0, hf); + } // ── 2i. SwiGLU: down(silu(gate(x)) * up(x)) // w_gate, w_up: [hidden, intermediate] @@ -261,6 +340,9 @@ DraftGraphOutputs build_draft_graph( ggml_tensor * u = ggml_mul_mat(ctx, L.w_up, hf); // [inter, q_len] ggml_tensor * gu = ggml_mul(ctx, g, u); ggml_tensor * ffn_out = ggml_mul_mat(ctx, L.w_down, gu); // [hidden, q_len] + if (dyn_conv) { + ffn_out = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 1, ffn_out); + } h = ggml_add(ctx, h, ffn_out); std::snprintf(probe_name, sizeof(probe_name), "draft_l%d_h_after_ffn", il); @@ -374,10 +456,16 @@ DraftGraphOutputs build_draft_kv_step( for (int il = 0; il < w.n_layer; il++) { const DraftLayer & L = w.layers[il]; const bool layer_is_swa = L.is_swa && !disable_swa; + const bool dyn_conv = w.conv_kernel_size > 0 && L.attn_conv.present() && L.mlp_conv.present(); - // ── attention pre-norm + // ── attention pre-norm (+ DFlash 2 dynamic conv "prepare") ggml_tensor * hn = ggml_rms_norm(ctx, h, eps); hn = ggml_mul(ctx, hn, L.attn_norm); + DraftDynConv attn_dc; + if (dyn_conv) { + attn_dc = draft_dyn_conv_kernel(ctx, L.attn_conv, hn); + hn = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 0, hn); + } // ── Q from noise, per-head RMSNorm, RoPE at absolute positions ggml_tensor * Q = ggml_mul_mat(ctx, L.wq, hn); @@ -436,16 +524,27 @@ DraftGraphOutputs build_draft_kv_step( attn = ggml_reshape_2d(ctx, attn, head_dim * n_head, q_len); ggml_tensor * attn_out = ggml_mul_mat(ctx, L.wo, attn); + if (dyn_conv) { + attn_out = draft_dyn_conv_apply(ctx, w, L.attn_conv, attn_dc, 1, attn_out); + } h = ggml_add(ctx, h, attn_out); - // ── FFN + // ── FFN (+ DFlash 2 dynamic conv prepare/finish) ggml_tensor * hf = ggml_rms_norm(ctx, h, eps); hf = ggml_mul(ctx, hf, L.ffn_norm); + DraftDynConv mlp_dc; + if (dyn_conv) { + mlp_dc = draft_dyn_conv_kernel(ctx, L.mlp_conv, hf); + hf = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 0, hf); + } ggml_tensor * g = ggml_mul_mat(ctx, L.w_gate, hf); g = ggml_silu(ctx, g); ggml_tensor * u = ggml_mul_mat(ctx, L.w_up, hf); ggml_tensor * gu = ggml_mul(ctx, g, u); ggml_tensor * ffn_out = ggml_mul_mat(ctx, L.w_down, gu); + if (dyn_conv) { + ffn_out = draft_dyn_conv_apply(ctx, w, L.mlp_conv, mlp_dc, 1, ffn_out); + } h = ggml_add(ctx, h, ffn_out); } diff --git a/server/src/internal.h b/server/src/internal.h index 1aabee02f..e4a6ea878 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -242,6 +242,18 @@ void free_target_weights(TargetWeights & w); // ─── Draft weights (z-lab DFlash, bf16) ─────────────────────────── +// DFlash 2 grouped dynamic causal conv (two taps over the draft block, one +// instance before/after attention and one before/after the MLP): +// dyn = proj @ x_norm [2*K*groups, q_len] +// prepare = sum_k (base[0][k] + dyn[0][k]) * shift_k(x_norm) +// finish = sum_k (base[1][k] + dyn[1][k]) * shift_k(sub_block_out) +// base is per element, dyn per group of conv_group_size elements. +struct DraftConvWeights { + ggml_tensor * base = nullptr; // [hidden, K, 2] f32 + ggml_tensor * proj = nullptr; // [hidden, 2*K*groups] + bool present() const { return base != nullptr && proj != nullptr; } +}; + struct DraftLayer { ggml_tensor * attn_norm; ggml_tensor * ffn_norm; @@ -255,6 +267,8 @@ struct DraftLayer { ggml_tensor * w_gate; ggml_tensor * w_up; ggml_tensor * w_down; + DraftConvWeights attn_conv; // optional DFlash 2 conv around attention + DraftConvWeights mlp_conv; // optional DFlash 2 conv around the MLP bool is_swa = false; // true for SWA layers (Qwen3.6 pattern) bool attn_gate_per_head = false; }; @@ -288,6 +302,18 @@ struct DraftDSparkWeights { ggml_tensor * confidence_b = nullptr; // [1] f32 }; +// DFlash 2 candidate selector: top-k candidates per block position from the +// target lm_head logits, then one path through them scored by a low-rank +// bigram form unary[c] + . +struct DraftSelectorWeights { + bool enabled = false; + int rank = 0; + int top_k = 0; + ggml_tensor * hproj = nullptr; // [hidden, rank] + ggml_tensor * pred_cb = nullptr; // [rank, vocab] predecessor codebook + ggml_tensor * succ_cb = nullptr; // [rank, vocab] successor codebook +}; + struct DraftWeights { ggml_context * ctx = nullptr; ggml_backend_t backend = nullptr; @@ -332,6 +358,12 @@ struct DraftWeights { // Optional DSpark/DeepSpec-style Markov correction head. When present, // greedy chain decode adds a low-rank previous-token bias before argmax. DraftDSparkWeights dspark; + + // Optional DFlash 2 pieces: dynamic convs live in the layers, the + // selector replaces argmax/markov projection for the drafted chain. + int conv_kernel_size = 0; // 0 = no dynamic convs + int conv_group_size = 0; + DraftSelectorWeights selector; }; bool load_draft_safetensors(const std::string & path, diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index a17bd4d4a..575b10439 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -16,6 +16,7 @@ #include #endif #include "common/dspark_head.h" +#include "common/dflash2_head.h" #include "common/io_utils.h" #include "common/restore_delta.h" #include "qwen35_tensor_parallel.h" @@ -2925,10 +2926,32 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, v_len = 1; } else if (!use_tree_verify) { const auto profile_project_start = profile_start(); + // DFlash 2 selector (top-k candidates + low-rank path score) when + // the drafter ships it. + bool used_dspark = false; + if (dw_.selector.enabled && q_len > 1 && !sampled_verify && !use_remote_draft) { + static std::atomic s_sel_logged{false}; + if (!s_sel_logged.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DFlash 2 selector active for greedy chain decode " + "(rank=%d top_k=%d)\n", dw_.selector.rank, dw_.selector.top_k); + } + if (dflash2_select_chain(dw_, draft_backend_, *target, + local_hidden.data(), q_len, last_tok, draft_tok)) { + used_dspark = true; + v_len = std::max(1, (int)draft_tok.size()); + } else { + static std::atomic s_sel_warned{false}; + if (!s_sel_warned.exchange(true)) { + std::fprintf(stderr, + "[qwen35-spec] DFlash 2 selector failed; falling back to " + "base DFlash projection\n"); + } + } + } // DSpark heads (markov bigram correction + optional confidence // gate) when the drafter ships them; mirrors the laguna hook. - bool used_dspark = false; - if (qwen35_dspark_enabled() && dw_.dspark.enabled && + if (!used_dspark && qwen35_dspark_enabled() && dw_.dspark.enabled && q_len > 1 && !sampled_verify && !use_remote_draft) { static std::atomic s_dspark_logged{false}; if (!s_dspark_logged.exchange(true)) { From 547f5c29b10f60152dcf4b7a7a0b2ac174f6bf2e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 19 Aug 2026 02:53:43 +0200 Subject: [PATCH 29/42] ggml: skip the pathological mmq_x=32 small tile on RDNA With the 64-row/4-warp tile the mmq_x=32 instantiation runs at 180 GB/s on gfx1201 (17408x5120 IQ4_XS) against 443 GB/s at mmq_x=16 and 315 at 48, so N=17..32 batches (DDTree budgets, prefill remainders) took 2.4x longer than N=16 or N=40. Choose the next tile instead. --- server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index 266bf6006..29a1d4db2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -4790,6 +4790,14 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda if (mmq_x % granularity != 0 || mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size, nwarps) > smpbo) { continue; } +#if defined(GGML_CUDA_MMQ_SMALL_TILE) + // The 64-row/4-warp tile is pathological at mmq_x == 32 on gfx1201 + // (17408x5120 IQ4_XS: N=16 443 GB/s, N=24..32 180 GB/s, N=48 315 GB/s + // in mmq_probe); a wider tile with more padding is still faster. + if (LUCEBOX_RDNA_TILE_HOST(cc) && mmq_x == 32) { + continue; + } +#endif const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; From 382173899e6ee14eac319745f2798144cc69f531 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 11:23:14 +0000 Subject: [PATCH 30/42] bench(qwen38): pin PR 625 q8 draft baseline --- harness/benchmarks/QWEN38_PR625_BASELINE.md | 10 +++++++--- server/scripts/prepare_qwen38_pr625_models.sh | 10 ++++++---- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/harness/benchmarks/QWEN38_PR625_BASELINE.md b/harness/benchmarks/QWEN38_PR625_BASELINE.md index 521ff9be8..207825bf4 100644 --- a/harness/benchmarks/QWEN38_PR625_BASELINE.md +++ b/harness/benchmarks/QWEN38_PR625_BASELINE.md @@ -24,9 +24,13 @@ This produces and validates: - target: pure IQ4_XS body, Q5_K `output.weight`, Q6_K `attn_v` and `ssm_out` -- drafter: no YaRN, Q8_0 by default, capture layers `4,16,28,40,52`, mask - token `248077`. PR #625 does not publish the DSpark precision; set - `DRAFT_SCHEME=q4-mix` only for an explicit ablation. +- drafter: no YaRN, Q8_0, capture layers `4,16,28,40,52`, mask token + `248077`. PR #625 does not spell out the DSpark quantization command, but + Q8_0 reproduces its reported compute regime: this setup measured a 1.78x + speculative/plain step-time ratio on the R9700, versus the PR's ~1.8x. + The unquantized F16 drafter measured 4.20x and is therefore not the + benchmark artifact. Set `DRAFT_SCHEME=f16` or `DRAFT_SCHEME=q4-mix` only + for an explicit ablation. ## Build on Radeon AI PRO R9700 diff --git a/server/scripts/prepare_qwen38_pr625_models.sh b/server/scripts/prepare_qwen38_pr625_models.sh index 839fe6028..16681efb2 100755 --- a/server/scripts/prepare_qwen38_pr625_models.sh +++ b/server/scripts/prepare_qwen38_pr625_models.sh @@ -16,8 +16,8 @@ DRAFT_SCHEME="${DRAFT_SCHEME:-q8_0}" [[ -r "$DRAFT_SOURCE" ]] || { echo "unreadable DRAFT_SOURCE: $DRAFT_SOURCE" >&2; exit 2; } [[ -x "$LLAMA_QUANTIZE" ]] || { echo "LLAMA_QUANTIZE is not executable: $LLAMA_QUANTIZE" >&2; exit 2; } command -v "$PYTHON" >/dev/null || { echo "PYTHON is unavailable: $PYTHON" >&2; exit 2; } -[[ "$DRAFT_SCHEME" == q8_0 || "$DRAFT_SCHEME" == q4-mix ]] || { - echo "DRAFT_SCHEME must be q8_0 or q4-mix" >&2 +[[ "$DRAFT_SCHEME" == f16 || "$DRAFT_SCHEME" == q8_0 || "$DRAFT_SCHEME" == q4-mix ]] || { + echo "DRAFT_SCHEME must be f16, q8_0, or q4-mix" >&2 exit 2 } @@ -32,8 +32,10 @@ target_final="$work_dir/Qwen3.8-27B-PR625-IQ4_XS.gguf" "$PYTHON" "$SCRIPT_DIR/convert_dflash_to_gguf.py" \ "$DRAFT_SOURCE" "$draft_f16" --no-yarn -"$PYTHON" "$SCRIPT_DIR/quantize_dflash_draft.py" \ - "$draft_f16" "$draft_final" --scheme "$DRAFT_SCHEME" +if [[ "$DRAFT_SCHEME" != f16 ]]; then + "$PYTHON" "$SCRIPT_DIR/quantize_dflash_draft.py" \ + "$draft_f16" "$draft_final" --scheme "$DRAFT_SCHEME" +fi # PR #625 target: pure IQ4_XS body, Q5_K output, Q6_K attn_v/ssm_out. # The validator below deliberately catches quantizers that let --pure suppress From b9c040561dbd3c958bef27a6356b3688d9bb885a Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 13:34:17 +0000 Subject: [PATCH 31/42] bench(concurrency): profile DSpark adaptive rounds --- .../QWEN38_DSPARK_ADAPTIVE_SELECTION.md | 143 +++ .../concurrency/analyze_gate_decisions.py | 872 ++++++++++++++++++ .../concurrency/generate_dspark_prompts.py | 55 +- .../concurrency/run_qwen38_dspark_matrix.sh | 57 +- .../concurrency/test_feature_tools.py | 188 ++++ .../concurrency/write_feature_metadata.py | 12 + .../qwen38_dspark_adaptive_selection.jsonl | 6 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 133 ++- .../qwen35/concurrency/qwen35_seq_engine.h | 12 +- 9 files changed, 1466 insertions(+), 12 deletions(-) create mode 100644 harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md create mode 100644 harness/benchmarks/concurrency/analyze_gate_decisions.py create mode 100644 harness/benchmarks/prompts/qwen38_dspark_adaptive_selection.jsonl diff --git a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md new file mode 100644 index 000000000..1e0d22818 --- /dev/null +++ b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md @@ -0,0 +1,143 @@ +# Qwen3.8 DSpark adaptive-selection prompts + +This six-request workload is a prompt-selection fixture for the concurrent +adaptive speculation gate. It is deliberately balanced rather than +representative: + +- two prompts where dense DSpark is a strong win; +- two prompts where dense DSpark is only a marginal win; +- two prompts where dense autoregressive decode wins. + +The machine-readable source is +`prompts/qwen38_dspark_adaptive_selection.jsonl`. Every row records its +selection class, expected dense oracle, and the R9700 screening measurement +that justified the label. +The exact completion lengths and ordered output hashes are retained locally in +`QWEN38_DSPARK_ADAPTIVE_SELECTION_R9700.json`. + +DSpark is the fixed proposal mechanism for this fixture. The benchmark is +about adaptive activation: which lanes the gate admits at each live +concurrency, whether the selected `k` beats matched pure AR, how much rejected +always-drafting costs, and which target phase dominates a bad decision. + + +## Dense screening baseline + +Measured 2026-08-19 on the Radeon AI PRO R9700 (`gfx1201`, ROCm 7.2) with: + +- target: Qwen3.8-27B PR #625 requantization, IQ4_XS body, Q5_K output, + Q6_K `attn_v` and `ssm_out`; +- drafter: RadixArk Qwen3.8-27B DSpark, no YaRN, Q8_0, block width 8; +- build: Release, `gfx1201`, HIP graphs; +- dense cache/attention: Q8_0 K/V, `--fa-window 2048`; +- greedy chat requests, up to 256 output tokens; +- fresh AR process without a drafter; +- forced DSpark with `DFLASH_QWEN35_SPEC_STEP_RATIO=0`, so the dense + decoder could not hide weak speculation behind its own AR bursts. + +The throughput columns are model-side decode rates. This was a one-repeat +selection screen, not a publication measurement. + +| Prompt | Class | Dense oracle | AR tok/s | Forced spec tok/s | Spec/AR | Accept | Commit/step | +| :--- | :--- | :--- | ---: | ---: | ---: | ---: | ---: | +| `he_09 sum_product` | strong win | speculation | 34.89 | 46.35 | 1.328 | 35.8% | 2.51 | +| `he_10 rolling_max` | strong win | speculation | 34.89 | 50.30 | 1.442 | 38.9% | 2.72 | +| `he_02 separate_paren_groups` | marginal win | speculation | 34.95 | 37.69 | 1.078 | 29.0% | 2.03 | +| `he_03 truncate_number` | marginal win | speculation | 35.15 | 37.44 | 1.065 | 28.9% | 2.02 | +| `he_08 filter_by_substring` | loss | AR | 35.04 | 33.95 | 0.969 | 26.2% | 1.84 | +| `prose-01 reproducibility` | loss | AR | 34.90 | 27.10 | 0.777 | 20.8% | 1.45 | + +All six selected rows produced identical ordered content hashes under dense AR +and forced dense speculation. Prompts with mismatched hashes were rejected +from the performance fixture: HumanEval 01, 05, 06, and 07, plus prose 02, +03, and 04. Keep those rejected prompts as correctness diagnostics; do not +interpret their throughput as a valid speculative win or loss. + +## Concurrent activation benchmark + +The smallest adversarial cohort is fixed at C=3: + +- two dense strong-win prompts (`sum_product` and `rolling_max`); +- the strongest dense AR win (`reproducibility`). + +This is intentionally difficult. A useful gate must learn that dense labels +are only priors: concurrency can move the break-even point enough that a +formerly strong speculative sample should run as AR. + +Run three paired fresh-process repeats: + +```bash +MODEL=/path/Qwen3.8-27B-PR625-IQ4_XS.gguf \ +DRAFT_MODEL=/path/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ +LUCE_SERVER_BIN=server/build-pr625-r9700/dflash_server \ +VISIBLE_DEVICES=0 \ +WORKLOADS=adaptive-selection-c3 CLIENTS=3 \ +DECODE_MODES=ar,speculation,adaptive REPEATS=3 \ +harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh +``` + +The paged concurrent executor requires full attention, represented by +`FA_WINDOW=0`. At this fixture's maximum context (at most 134 prompt tokens +plus 256 generated tokens), a 2048-token window would not discard context, so +full attention does not change the attended token set. It may still change +kernel overhead; the paired controls measure that runtime. The runner uses +Q8_0 K/V by default and records both cache types and the window in every case. + +For the wider selection boundary, run all six prompts at C=6: + +Run the selection workload only at C=6: + +```bash +MODEL=/path/Qwen3.8-27B-PR625-IQ4_XS.gguf \ +DRAFT_MODEL=/path/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ +LUCE_SERVER_BIN=server/build-pr625-r9700/dflash_server \ +VISIBLE_DEVICES=0 \ +WORKLOADS=adaptive-selection CLIENTS=6 \ +DECODE_MODES=ar,speculation,adaptive REPEATS=3 \ +harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh +``` + +The dense oracle labels are priors, not a hard assertion about the concurrent +executor. In particular, the two marginal dense wins may correctly become AR +choices after packed-tree, replay, padding, and synchronization costs are +included. The useful adaptive behavior is: + +1. rank the two strong wins above the two known losses; +2. keep the prose loss and normally the HumanEval loss in AR; +3. make the marginal pair expose the actual concurrent break-even boundary; +4. preserve the ordered greedy output hashes across AR, forced speculation, + and adaptive modes. + +## Profiling outputs + +`DFLASH_STEP_TIMING=1` is enabled by default. Every case retains: + +- `server.log`: startup, warmup, and measured request evidence; +- `benchmark-server.log`: only the measured window; +- `bench.json`: request and aggregate throughput; +- `feature-proof.json`: request-correlated execution/correctness proof. + +The matrix root adds: + +- `profiling.json`: complete machine-readable gate, request, shape, and phase + distributions; +- `profiling.md`: concise activation-regret and bottleneck tables; +- `summary.md`: oracle-gated aggregate summary, written only when output + stability and adaptive criteria pass. + +Read the report in this order: + +1. Require identical outputs between AR, forced speculation, and adaptive. + A mismatch is a correctness failure, not a throughput result. +2. Use paired concurrent AR/speculation measurements as the empirical oracle; + never use the dense prompt label as the final activation answer. +3. Inspect `Activation outcome against matched pure AR`. It compares each + `(live, k, path)` shape to pure AR at the same live concurrency. +4. Inspect `Gate prediction calibration` for predicted-versus-realized + goodput and realized-versus-AR regret. +5. Use phase attribution to choose the next optimization: draft tax on k=0, + verify, or the structural replay forward. + +The profiler's `total_us` uses one common decode-round origin for pure AR, +adaptive k=0, and speculative rounds. `draft_us` is a subset of that wall +time; do not add it a second time. diff --git a/harness/benchmarks/concurrency/analyze_gate_decisions.py b/harness/benchmarks/concurrency/analyze_gate_decisions.py new file mode 100644 index 000000000..da5cf2fb7 --- /dev/null +++ b/harness/benchmarks/concurrency/analyze_gate_decisions.py @@ -0,0 +1,872 @@ +#!/usr/bin/env python3 +"""Analyze adaptive gate decisions and per-phase decode timing. + +Inputs may be individual matrix case directories or a completed matrix root. +The analyzer joins benchmark request IDs to engine IDs, preserves prompt +selection labels, compares paired AR/speculation controls, and summarizes +machine-readable [step-timing] records. When benchmark-server.log exists it is +preferred so warmup rounds cannot contaminate the measured distributions. +""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import statistics +from collections import defaultdict +from pathlib import Path +from typing import Any + + +NUMBER = r"(?:[0-9]+(?:\.[0-9]+)?|\.[0-9]+)(?:[eE][+-]?[0-9]+)?" +GATE_RE = re.compile( + rf"\[spec-gate\] C=(?P\d+) k=(?P\d+) " + rf"scores=\[(?P[^\]]*)\].*?" + rf"G\(k\)=(?P{NUMBER}) G\(0\)=(?P{NUMBER}).*?" + rf"predicted_cost=(?P{NUMBER})us " + rf"measured=(?:(?P{NUMBER})us|ar-path)" +) +SCORE_RE = re.compile( + rf"(?P\d+):(?P{NUMBER}|nan)/" + r"(?P[a-z_-]+)(?P\*?)" +) +METRIC_RE = re.compile(r"\[concurrency-metrics\] (?P\{.*\})") +TIMING_RE = re.compile(r"\[step-timing\] (?P\{.*\})") +TIMING_COUNT_FIELDS = ( + "live", "k", "emitted_tokens", "accepted_tokens", "target_forwards", +) + + +def _json_object(match: re.Match[str], path: Path, line_no: int) -> dict[str, Any]: + try: + value = json.loads(match.group("json")) + except json.JSONDecodeError as exc: + raise ValueError(f"{path}:{line_no}: invalid profiling JSON: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"{path}:{line_no}: profiling record must be an object") + return value + + +def _validate_timing(row: dict[str, Any], path: Path, line_no: int) -> None: + if row.get("path") not in ("ar", "spec"): + raise ValueError(f"{path}:{line_no}: invalid step-timing path") + for key in TIMING_COUNT_FIELDS: + value = row.get(key) + if type(value) is not int or value < 0: + raise ValueError( + f"{path}:{line_no}: step-timing {key} must be non-negative int" + ) + total = row.get("total_us") + if type(total) not in (int, float) or not math.isfinite(total) or total <= 0: + raise ValueError(f"{path}:{line_no}: step-timing total_us must be positive") + for key, value in row.items(): + if key.endswith("_us") and ( + type(value) not in (int, float) + or not math.isfinite(value) + or value < 0 + ): + raise ValueError( + f"{path}:{line_no}: step-timing {key} must be non-negative" + ) + + +def parse_server_log( + path: Path, +) -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]], list[dict[str, Any]]]: + rounds: list[dict[str, Any]] = [] + metrics: dict[str, dict[str, Any]] = {} + timings: list[dict[str, Any]] = [] + for line_no, line in enumerate( + path.read_text(encoding="utf-8", errors="replace").splitlines(), 1, + ): + gate = GATE_RE.search(line) + if gate: + entries = [] + for score in SCORE_RE.finditer(gate.group("scores")): + raw_score = score.group("score") + entries.append({ + "request": int(score.group("rid")), + "score": ( + float(raw_score) if raw_score != "nan" + else math.nan + ), + "source": score.group("source"), + "admitted": score.group("admitted") == "*", + }) + rounds.append({ + "concurrency": int(gate.group("c")), + "k": int(gate.group("k")), + "predicted_goodput_tok_s": float(gate.group("gk")) * 1e6, + "ar_goodput_tok_s": float(gate.group("g0")) * 1e6, + "predicted_cost_us": float(gate.group("predicted")), + "measured_cost_us": ( + float(gate.group("measured")) + if gate.group("measured") is not None else None + ), + "entries": entries, + }) + continue + metric = METRIC_RE.search(line) + if metric: + row = _json_object(metric, path, line_no) + request_id = row.get("request_id") + if isinstance(request_id, str): + if request_id in metrics: + raise ValueError( + f"{path}:{line_no}: duplicate request metric {request_id}" + ) + metrics[request_id] = row + continue + timing = TIMING_RE.search(line) + if timing: + row = _json_object(timing, path, line_no) + _validate_timing(row, path, line_no) + timings.append(row) + return rounds, metrics, timings + + +def _prompt_records(prompt_file: Path | None) -> list[dict[str, Any]]: + if prompt_file is None or not prompt_file.exists(): + return [] + records = [] + for line_no, line in enumerate( + prompt_file.read_text(encoding="utf-8").splitlines(), 1, + ): + if not line.strip(): + continue + value = json.loads(line) + if not isinstance(value, dict): + raise ValueError(f"{prompt_file}:{line_no}: prompt must be an object") + records.append(value) + return records + + +def prompt_names(bench: dict[str, Any], prompt_file: Path | None) -> dict[str, dict]: + """Map wire request_id to prompt identity and selection metadata.""" + records = _prompt_records(prompt_file) + out: dict[str, dict] = {} + for level in bench.get("levels", []): + for detail in level.get("requests_detail", []): + request_id = detail.get("request_id") + index = detail.get("prompt_index") + if not isinstance(request_id, str) or type(index) is not int: + continue + source = records[index] if index < len(records) else {} + out[request_id] = { + "prompt_index": index, + "prompt_id": str(source.get("id", f"prompt-{index}")), + "selection_class": source.get("selection_class"), + "expected_dense_oracle": source.get("expected_dense_oracle"), + "dense_r9700_baseline": source.get("dense_r9700_baseline"), + "decode_tok_s": detail.get("request_decode_tok_s"), + "output_sha256": detail.get("content_sha256"), + } + return out + + +def _find_prompt_file( + case_dir: Path, workload: str, explicit: Path | None, +) -> Path | None: + if explicit is not None: + return explicit + for parent in (case_dir, *case_dir.parents): + candidate = parent / "prompts" / f"{workload}.jsonl" + if candidate.is_file(): + return candidate + return None + + +def _percentile(values: list[float], fraction: float) -> float | None: + if not values: + return None + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def _timing_summary(rows: list[dict[str, Any]]) -> dict[str, Any]: + emitted = sum(int(row["emitted_tokens"]) for row in rows) + accepted = sum(int(row["accepted_tokens"]) for row in rows) + forwards = sum(int(row["target_forwards"]) for row in rows) + total_us = sum(float(row["total_us"]) for row in rows) + spec_lane_steps = sum(int(row["k"]) for row in rows) + phase_keys = sorted({ + key for row in rows for key in row + if key.endswith("_us") and key != "total_us" + }) + means = { + key: statistics.fmean(float(row.get(key, 0.0)) for row in rows) + for key in phase_keys + } + medians = { + key: statistics.median(float(row.get(key, 0.0)) for row in rows) + for key in phase_keys + } + p95 = { + key: _percentile( + [float(row.get(key, 0.0)) for row in rows], 0.95, + ) + for key in phase_keys + } + return { + "rounds": len(rows), + "total_wall_us": total_us, + "emitted_tokens": emitted, + "accepted_tokens": accepted, + "target_forwards": forwards, + "emitted_tokens_per_target_forward": ( + emitted / forwards if forwards else None + ), + "accepted_tokens_per_spec_lane_step": ( + accepted / spec_lane_steps if spec_lane_steps else None + ), + "round_goodput_tok_s": ( + emitted * 1e6 / total_us if total_us else None + ), + "mean_total_us": statistics.fmean( + float(row["total_us"]) for row in rows + ) if rows else None, + "p50_total_us": statistics.median( + float(row["total_us"]) for row in rows + ) if rows else None, + "p95_total_us": _percentile( + [float(row["total_us"]) for row in rows], 0.95, + ), + "phase_mean_us": means, + "phase_p50_us": medians, + "phase_p95_us": p95, + } + + +def _gate_summary( + rounds: list[dict[str, Any]], timings: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], int]: + gate_by_k: dict[int, list[dict[str, Any]]] = defaultdict(list) + timing_by_k: dict[int, list[dict[str, Any]]] = defaultdict(list) + for row in rounds: + gate_by_k[int(row["k"])].append(row) + for row in timings: + timing_by_k[int(row["k"])].append(row) + + summaries = [] + mismatch = 0 + for k in sorted(gate_by_k): + gate_rows = gate_by_k[k] + timing_rows = timing_by_k.get(k, []) + paired = min(len(gate_rows), len(timing_rows)) + mismatch += abs(len(gate_rows) - len(timing_rows)) + paired_timing = timing_rows[:paired] + realized_tokens = sum( + int(row["emitted_tokens"]) for row in paired_timing + ) + realized_us = sum(float(row["total_us"]) for row in paired_timing) + predicted = statistics.fmean( + float(row["predicted_goodput_tok_s"]) for row in gate_rows + ) + realized = ( + realized_tokens * 1e6 / realized_us if realized_us else None + ) + measured_costs = [ + float(row["measured_cost_us"]) for row in gate_rows + if row["measured_cost_us"] is not None + ] + summaries.append({ + "k": k, + "gate_rounds": len(gate_rows), + "timing_rounds": len(timing_rows), + "paired_rounds": paired, + "predicted_goodput_mean_tok_s": predicted, + "predicted_ar_goodput_mean_tok_s": statistics.fmean( + float(row["ar_goodput_tok_s"]) for row in gate_rows + ), + "realized_goodput_tok_s": realized, + "realized_over_predicted": ( + realized / predicted if realized is not None and predicted else None + ), + "realized_over_predicted_ar": ( + realized / statistics.fmean( + float(row["ar_goodput_tok_s"]) for row in gate_rows + ) + if realized is not None else None + ), + "predicted_cost_mean_us": statistics.fmean( + float(row["predicted_cost_us"]) for row in gate_rows + ), + "gate_measured_cost_mean_us": ( + statistics.fmean(measured_costs) if measured_costs else None + ), + "timed_wall_mean_us": ( + statistics.fmean(float(row["total_us"]) for row in paired_timing) + if paired_timing else None + ), + }) + return summaries, mismatch + + +def _class_summary(requests: list[dict[str, Any]]) -> list[dict[str, Any]]: + grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) + for row in requests: + selection_class = row.get("selection_class") + if isinstance(selection_class, str): + grouped[selection_class].append(row) + out = [] + for selection_class, rows in sorted(grouped.items()): + gate_rounds = sum(int(row["gate_rounds"]) for row in rows) + admitted = sum(int(row["admitted_rounds"]) for row in rows) + spec_steps = sum(int(row["spec_steps"]) for row in rows) + accepted = sum(int(row["spec_accepted_tokens"]) for row in rows) + out.append({ + "selection_class": selection_class, + "requests": len(rows), + "gate_rounds": gate_rounds, + "admitted_fraction": admitted / gate_rounds if gate_rounds else 0.0, + "spec_steps": spec_steps, + "accepted_tokens": accepted, + "commit_per_spec_step": ( + (spec_steps + accepted) / spec_steps if spec_steps else None + ), + }) + return out + + +def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, Any]: + bench_path = case_dir / "bench.json" + bench = json.loads(bench_path.read_text(encoding="utf-8")) + metadata = bench.get("server_metadata") or {} + workload = str(metadata.get("workload") or "") + variant = str(metadata.get("variant") or case_dir.name) + repeat = metadata.get("repeat") + clients = metadata.get("clients") + if type(clients) is not int: + levels = bench.get("levels") or [] + clients = levels[0].get("clients") if levels else None + + log_path = case_dir / "benchmark-server.log" + if not log_path.is_file(): + log_path = case_dir / "server.log" + rounds, metrics, timings = parse_server_log(log_path) + resolved_prompt_file = _find_prompt_file( + case_dir, workload, prompt_file, + ) + by_wire = prompt_names(bench, resolved_prompt_file) + + engine_to_prompt: dict[int, dict[str, Any]] = {} + for wire_id, row in metrics.items(): + info = by_wire.get(wire_id) + engine_id = row.get("engine_request_id") + if info is None or type(engine_id) is not int: + continue + engine_to_prompt[engine_id] = { + **info, + "spec_accepted_tokens": row.get("spec_accepted_tokens", 0), + "spec_steps": row.get("spec_steps", 0), + "target_forwards": row.get("target_forwards", 0), + "output_tokens": row.get("output_tokens", 0), + } + + per_request: dict[int, dict[str, Any]] = defaultdict( + lambda: {"rounds": 0, "admitted": 0, "score_sum": 0.0, "scored": 0} + ) + k_histogram: dict[int, int] = defaultdict(int) + for entry in rounds: + k_histogram[int(entry["k"])] += 1 + for score in entry["entries"]: + stats = per_request[int(score["request"])] + stats["rounds"] += 1 + if score["admitted"]: + stats["admitted"] += 1 + if score["source"] == "confidence" and math.isfinite(score["score"]): + stats["score_sum"] += score["score"] + stats["scored"] += 1 + + requests = [] + for engine_id, prompt in sorted( + engine_to_prompt.items(), key=lambda item: item[1]["prompt_index"], + ): + stats = per_request[engine_id] + steps = prompt.get("spec_steps", 0) + accepted = prompt.get("spec_accepted_tokens", 0) + requests.append({ + "engine_request_id": engine_id, + **prompt, + "gate_rounds": stats["rounds"], + "admitted_rounds": stats["admitted"], + "admitted_fraction": ( + stats["admitted"] / stats["rounds"] if stats["rounds"] else 0.0 + ), + "mean_confidence_yield": ( + stats["score_sum"] / stats["scored"] if stats["scored"] else None + ), + "commit_per_spec_step": ( + (steps + accepted) / steps if steps else None + ), + }) + + levels = bench.get("levels") or [] + aggregate_tok_s = ( + levels[0].get("aggregate_tok_s") if len(levels) == 1 else None + ) + by_path = { + path: _timing_summary([ + row for row in timings if row["path"] == path + ]) + for path in ("ar", "spec") + if any(row["path"] == path for row in timings) + } + shape_groups: dict[tuple[str, int, int], list[dict[str, Any]]] = defaultdict(list) + for row in timings: + shape_groups[ + (str(row["path"]), int(row["live"]), int(row["k"])) + ].append(row) + by_shape = [ + { + "path": path, + "live": live, + "k": k, + **_timing_summary(rows), + } + for (path, live, k), rows in sorted(shape_groups.items()) + ] + ar_timings = [row for row in timings if row["path"] == "ar"] + drafted_ar = [row for row in ar_timings if float(row.get("draft_us", 0)) > 0] + ar_wall = sum(float(row["total_us"]) for row in ar_timings) + ar_draft = sum(float(row.get("draft_us", 0)) for row in ar_timings) + gate_by_k, gate_timing_mismatch = _gate_summary(rounds, timings) + + return { + "case": str(case_dir), + "log": str(log_path), + "workload": workload, + "clients": clients, + "repeat": repeat, + "variant": variant, + "aggregate_tok_s": aggregate_tok_s, + "gate_rounds": len(rounds), + "k_histogram": dict(sorted(k_histogram.items())), + "gate_by_k": gate_by_k, + "gate_timing_count_mismatch": gate_timing_mismatch, + "timing": { + "records": len(timings), + "by_path": by_path, + "by_shape": by_shape, + "draft_tax_on_ar": { + "ar_rounds": len(ar_timings), + "drafted_ar_rounds": len(drafted_ar), + "draft_us": ar_draft, + "ar_wall_us": ar_wall, + "fraction_of_ar_wall": ar_draft / ar_wall if ar_wall else None, + }, + }, + "selection_classes": _class_summary(requests), + "requests": requests, + } + + +def discover_case_dirs(paths: list[Path]) -> list[Path]: + case_dirs: set[Path] = set() + for path in paths: + if path.name == "bench.json" and path.is_file(): + case_dirs.add(path.parent) + elif (path / "bench.json").is_file(): + case_dirs.add(path) + elif path.is_dir(): + case_dirs.update(item.parent for item in path.rglob("bench.json")) + else: + raise ValueError(f"{path}: not a benchmark case or matrix root") + if not case_dirs: + raise ValueError("no bench.json files found") + return sorted(case_dirs) + + +def compare_prompts(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: + grouped: dict[tuple[Any, ...], dict[str, dict[str, Any]]] = defaultdict(dict) + metadata: dict[tuple[Any, ...], dict[str, Any]] = {} + for case in cases: + for request in case["requests"]: + key = ( + case["workload"], case["clients"], case["repeat"], + request["prompt_id"], + ) + grouped[key][case["variant"]] = request + metadata[key] = { + name: request.get(name) for name in ( + "selection_class", "expected_dense_oracle", + "dense_r9700_baseline", + ) + } + + comparisons = [] + for key in sorted(grouped, key=lambda value: tuple(str(item) for item in value)): + variants = grouped[key] + if "ar" not in variants or "speculation" not in variants: + continue + ar_rate = variants["ar"].get("decode_tok_s") + spec_rate = variants["speculation"].get("decode_tok_s") + ratio = ( + spec_rate / ar_rate + if type(ar_rate) in (int, float) and ar_rate > 0 + and type(spec_rate) in (int, float) else None + ) + empirical_oracle = ( + "speculation" if ratio is not None and ratio > 1.0 else "ar" + ) + row = { + "workload": key[0], + "clients": key[1], + "repeat": key[2], + "prompt_id": key[3], + **metadata[key], + "ar_decode_tok_s": ar_rate, + "speculation_decode_tok_s": spec_rate, + "speculation_over_ar": ratio, + "empirical_concurrent_oracle": empirical_oracle, + "matches_expected_dense_oracle": ( + empirical_oracle == metadata[key].get("expected_dense_oracle") + if metadata[key].get("expected_dense_oracle") is not None else None + ), + "output_stable": ( + variants["ar"].get("output_sha256") + == variants["speculation"].get("output_sha256") + and variants["ar"].get("output_sha256") is not None + ), + "adaptive": {}, + } + for name, request in sorted(variants.items()): + if name.startswith("adaptive-"): + row["adaptive"][name] = { + "decode_tok_s": request.get("decode_tok_s"), + "admitted_fraction": request.get("admitted_fraction"), + "mean_confidence_yield": request.get( + "mean_confidence_yield" + ), + "commit_per_spec_step": request.get( + "commit_per_spec_step" + ), + } + comparisons.append(row) + return comparisons + + +def compare_activation_shapes( + cases: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Compare every non-control round shape with pure AR at the same live C.""" + grouped: dict[tuple[Any, ...], dict[str, dict[str, Any]]] = defaultdict(dict) + for case in cases: + grouped[ + (case["workload"], case["clients"], case["repeat"]) + ][case["variant"]] = case + + comparisons = [] + for key, variants in sorted( + grouped.items(), key=lambda item: tuple(str(value) for value in item[0]), + ): + ar_case = variants.get("ar") + if ar_case is None: + continue + ar_shapes = { + int(shape["live"]): shape + for shape in ar_case["timing"]["by_shape"] + if shape["path"] == "ar" and int(shape["k"]) == 0 + } + ar_aggregate = ar_case.get("aggregate_tok_s") + for variant, case in sorted(variants.items()): + if variant == "ar": + continue + aggregate = case.get("aggregate_tok_s") + case["aggregate_over_ar"] = ( + aggregate / ar_aggregate + if type(aggregate) in (int, float) + and type(ar_aggregate) in (int, float) + and ar_aggregate > 0 else None + ) + for shape in case["timing"]["by_shape"]: + baseline = ar_shapes.get(int(shape["live"])) + if baseline is None: + continue + realized = shape.get("round_goodput_tok_s") + ar_rate = baseline.get("round_goodput_tok_s") + ratio = ( + realized / ar_rate + if type(realized) in (int, float) + and type(ar_rate) in (int, float) + and ar_rate > 0 else None + ) + path = str(shape["path"]) + comparisons.append({ + "workload": key[0], + "clients": key[1], + "repeat": key[2], + "variant": variant, + "path": path, + "live": shape["live"], + "k": shape["k"], + "rounds": shape["rounds"], + "realized_goodput_tok_s": realized, + "pure_ar_goodput_tok_s": ar_rate, + "realized_over_pure_ar": ratio, + "mean_wall_us": shape.get("mean_total_us"), + "mean_draft_us": ( + shape.get("phase_mean_us") or {} + ).get("draft_us"), + "activation_outcome": ( + "profitable" if path == "spec" and ratio is not None + and ratio > 1.0 + else "unprofitable" if path == "spec" + else "ar-with-draft-tax" if ( + (shape.get("phase_mean_us") or {}).get( + "draft_us", 0.0 + ) > 0 + ) else "ar" + ), + }) + for case in cases: + if case["variant"] == "ar": + case["aggregate_over_ar"] = 1.0 + else: + case.setdefault("aggregate_over_ar", None) + return comparisons + + +def _fmt(value: Any, digits: int = 2) -> str: + return f"{value:.{digits}f}" if type(value) in (int, float) else "n/a" + + +def render_markdown(report: dict[str, Any]) -> str: + lines = [ + "# Adaptive speculation profiling", + "", + "Round goodput is emitted decode tokens divided by the common measured " + "decode-round wall. Draft time is already inside that wall and is also " + "reported separately as attribution.", + "", + "## Case overview", + "", + "| Workload | C | Repeat | Variant | Benchmark tok/s | AR rounds | " + "Spec rounds | Timed tok/s | vs AR | AR draft tax |", + "| :--- | ---: | ---: | :--- | ---: | ---: | ---: | ---: | ---: | " + "---: |", + ] + for case in report["cases"]: + paths = case["timing"]["by_path"] + ar = paths.get("ar", {}) + spec = paths.get("spec", {}) + total_tokens = sum( + value.get("emitted_tokens", 0) for value in paths.values() + ) + total_wall = sum( + value.get("total_wall_us", 0.0) for value in paths.values() + ) + timed_rate = total_tokens * 1e6 / total_wall if total_wall else None + tax = case["timing"]["draft_tax_on_ar"]["fraction_of_ar_wall"] + tax_text = f"{tax * 100:.1f}%" if tax is not None else "n/a" + lines.append( + f"| {case['workload']} | {case['clients']} | {case['repeat']} | " + f"{case['variant']} | {_fmt(case['aggregate_tok_s'])} | " + f"{ar.get('rounds', 0)} | {spec.get('rounds', 0)} | " + f"{_fmt(timed_rate)} | {_fmt(case['aggregate_over_ar'], 3)} | " + f"{tax_text} |" + ) + + if report["activation_comparisons"]: + lines += [ + "", + "## Activation outcome against matched pure AR", + "", + "Each row compares the measured round shape with pure AR at the same " + "number of live requests. A speculative ratio below 1 is an " + "activation error for that concurrency shape.", + "", + "| Workload | Variant | Live | k | Path | Rounds | Actual tok/s | " + "Pure AR tok/s | Actual/AR | Draft mean us | Outcome |", + "| :--- | :--- | ---: | ---: | :--- | ---: | ---: | ---: | ---: | " + "---: | :--- |", + ] + for row in report["activation_comparisons"]: + lines.append( + f"| {row['workload']} | {row['variant']} | {row['live']} | " + f"{row['k']} | {row['path']} | {row['rounds']} | " + f"{_fmt(row['realized_goodput_tok_s'])} | " + f"{_fmt(row['pure_ar_goodput_tok_s'])} | " + f"{_fmt(row['realized_over_pure_ar'], 3)} | " + f"{_fmt(row['mean_draft_us'], 1)} | " + f"{row['activation_outcome']} |" + ) + + phase_rows = [ + (case, path, summary) + for case in report["cases"] + for path, summary in case["timing"]["by_path"].items() + ] + if phase_rows: + lines += [ + "", + "## Mean phase attribution", + "", + "Draft is a subset of pre-round time, so it must not be added to " + "the other columns a second time. Verify and replay are the two " + "target forwards on speculative rounds.", + "", + "| Variant | Path | Rounds | Wall us | Draft us | AR graph us | " + "Verify us | Replay us | Build us | Readback us |", + "| :--- | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " + "---: | ---: |", + ] + for case, path, summary in phase_rows: + phase = summary["phase_mean_us"] + build_us = sum( + phase.get(key, 0.0) for key in ( + "graph_build_us", "graph_prepare_us", + "verify_build_us", "replay_build_us", + ) + ) + read_us = sum( + phase.get(key, 0.0) for key in ( + "posterior_read_us", "sample_read_us", + ) + ) + lines.append( + f"| {case['variant']} | {path} | {summary['rounds']} | " + f"{_fmt(summary['mean_total_us'], 1)} | " + f"{_fmt(phase.get('draft_us'), 1)} | " + f"{_fmt(phase.get('graph_exec_us'), 1)} | " + f"{_fmt(phase.get('verify_exec_us'), 1)} | " + f"{_fmt(phase.get('replay_exec_us'), 1)} | " + f"{_fmt(build_us, 1)} | {_fmt(read_us, 1)} |" + ) + + if report["prompt_comparisons"]: + lines += [ + "", + "## Per-prompt concurrent oracle", + "", + "| Workload | C | Prompt | Class | AR tok/s | Spec tok/s | Spec/AR | " + "Concurrent oracle | Dense label agrees | Stable |", + "| :--- | ---: | :--- | :--- | ---: | ---: | ---: | :--- | " + ":---: | :---: |", + ] + for row in report["prompt_comparisons"]: + lines.append( + f"| {row['workload']} | {row['clients']} | {row['prompt_id']} | " + f"{row.get('selection_class') or 'n/a'} | " + f"{_fmt(row['ar_decode_tok_s'])} | " + f"{_fmt(row['speculation_decode_tok_s'])} | " + f"{_fmt(row['speculation_over_ar'], 3)} | " + f"{row['empirical_concurrent_oracle']} | " + f"{row['matches_expected_dense_oracle']} | " + f"{row['output_stable']} |" + ) + + adaptive_classes = [ + (case, item) + for case in report["cases"] if case["variant"].startswith("adaptive-") + for item in case["selection_classes"] + ] + if adaptive_classes: + lines += [ + "", + "## Adaptive decisions by prompt class", + "", + "| Workload | C | Variant | Class | Gate rounds | Admitted | " + "Commit/spec step |", + "| :--- | ---: | :--- | :--- | ---: | ---: | ---: |", + ] + for case, item in adaptive_classes: + lines.append( + f"| {case['workload']} | {case['clients']} | {case['variant']} | " + f"{item['selection_class']} | {item['gate_rounds']} | " + f"{_fmt(item['admitted_fraction'] * 100, 1)}% | " + f"{_fmt(item['commit_per_spec_step'])} |" + ) + + gate_rows = [ + (case, item) + for case in report["cases"] + for item in case["gate_by_k"] + ] + if gate_rows: + lines += [ + "", + "## Gate prediction calibration", + "", + "| Workload | C | Variant | k | Rounds | Pred tok/s | " + "Realized tok/s | Realized/pred | Realized/AR | Pred cost us | " + "Timed wall us |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | " + "---: | ---: | ---: |", + ] + for case, item in gate_rows: + lines.append( + f"| {case['workload']} | {case['clients']} | " + f"{case['variant']} | {item['k']} | {item['gate_rounds']} | " + f"{_fmt(item['predicted_goodput_mean_tok_s'])} | " + f"{_fmt(item['realized_goodput_tok_s'])} | " + f"{_fmt(item['realized_over_predicted'], 3)} | " + f"{_fmt(item['realized_over_predicted_ar'], 3)} | " + f"{_fmt(item['predicted_cost_mean_us'], 1)} | " + f"{_fmt(item['timed_wall_mean_us'], 1)} |" + ) + + lines.append("") + return "\n".join(lines) + + +def build_report( + paths: list[Path], prompt_file: Path | None = None, +) -> dict[str, Any]: + cases = [ + analyze_case(case_dir, prompt_file) + for case_dir in discover_case_dirs(paths) + ] + activation_comparisons = compare_activation_shapes(cases) + return { + "schema_version": 3, + "cases": cases, + "prompt_comparisons": compare_prompts(cases), + "activation_comparisons": activation_comparisons, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("paths", nargs="+", type=Path) + parser.add_argument( + "--prompt-file", type=Path, default=None, + help="override prompt metadata (normally discovered from matrix root)", + ) + parser.add_argument("--out", type=Path) + parser.add_argument("--markdown-out", type=Path) + args = parser.parse_args() + report = build_report(args.paths, args.prompt_file) + rendered = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.out: + args.out.write_text(rendered, encoding="utf-8") + if args.markdown_out: + args.markdown_out.write_text( + render_markdown(report), encoding="utf-8", + ) + if args.out or args.markdown_out: + destinations = ", ".join( + str(path) for path in (args.out, args.markdown_out) + if path is not None + ) + print( + f"[profile] cases={len(report['cases'])} " + f"prompt_comparisons={len(report['prompt_comparisons'])} " + f"activation_comparisons={len(report['activation_comparisons'])} " + f"wrote {destinations}" + ) + else: + print(rendered, end="") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/generate_dspark_prompts.py b/harness/benchmarks/concurrency/generate_dspark_prompts.py index 2e7f22357..2b3830e40 100755 --- a/harness/benchmarks/concurrency/generate_dspark_prompts.py +++ b/harness/benchmarks/concurrency/generate_dspark_prompts.py @@ -16,6 +16,17 @@ "humaneval": PROMPT_ROOT / "bench_he.jsonl", "gsm8k": PROMPT_ROOT / "bench_gsm.jsonl", } +ADAPTIVE_SELECTION_SOURCE = ( + PROMPT_ROOT / "qwen38_dspark_adaptive_selection.jsonl" +) +# Mixed C=3 cohort: two strong dense speculation wins plus the strongest +# dense speculation loss. The adaptive win condition is a 2-spec+1-AR split +# that beats both uniform modes. +ADAPTIVE_SELECTION_C3_IDS = ( + "adaptive-he-09-sum-product", + "adaptive-he-10-rolling-max", + "adaptive-prose-01-reproducibility", +) PROSE_TOPICS = ( "why reproducible benchmarks need immutable inputs and explicit hardware metadata", "how admission control improves the reliability of a concurrent inference service", @@ -77,10 +88,49 @@ def _prose_records() -> list[dict[str, object]]: ] +def _adaptive_selection_records() -> list[dict[str, object]]: + records: list[dict[str, object]] = [] + for line_no, raw in enumerate( + ADAPTIVE_SELECTION_SOURCE.read_text(encoding="utf-8").splitlines(), 1, + ): + if not raw.strip(): + continue + row = json.loads(raw) + if not isinstance(row.get("id"), str) or not isinstance(row.get("prompt"), str): + raise ValueError( + f"{ADAPTIVE_SELECTION_SOURCE}:{line_no}: invalid id or prompt" + ) + if row.get("expected_dense_oracle") not in ("ar", "speculation"): + raise ValueError( + f"{ADAPTIVE_SELECTION_SOURCE}:{line_no}: invalid dense oracle" + ) + baseline = row.get("dense_r9700_baseline") + if not isinstance(baseline, dict) or baseline.get("lossless") is not True: + raise ValueError( + f"{ADAPTIVE_SELECTION_SOURCE}:{line_no}: baseline must be lossless" + ) + records.append(row) + if len(records) != 6: + raise ValueError( + f"{ADAPTIVE_SELECTION_SOURCE}: expected exactly six prompts" + ) + return records + + def build_records(profile: str) -> list[dict[str, object]]: if profile in SOURCE_FILES: return _source_records(profile) prose = _prose_records() + if profile == "adaptive-selection": + return _adaptive_selection_records() + if profile == "adaptive-selection-c3": + by_id = {row["id"]: row for row in _adaptive_selection_records()} + missing = [pid for pid in ADAPTIVE_SELECTION_C3_IDS if pid not in by_id] + if missing: + raise ValueError( + f"{ADAPTIVE_SELECTION_SOURCE}: missing C=3 prompts {missing}" + ) + return [dict(by_id[pid]) for pid in ADAPTIVE_SELECTION_C3_IDS] if profile == "prose": return prose if profile == "north-star": @@ -97,7 +147,10 @@ def build_records(profile: str) -> list[dict[str, object]]: def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( - "--profile", choices=("humaneval", "gsm8k", "prose", "north-star"), + "--profile", choices=( + "humaneval", "gsm8k", "prose", "north-star", + "adaptive-selection", "adaptive-selection-c3", + ), required=True, ) parser.add_argument("--out", type=Path, required=True) diff --git a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh index 745ecd36d..608d291f5 100755 --- a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh @@ -10,6 +10,7 @@ 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}" +ANALYZER="${ANALYZER:-$SCRIPT_DIR/analyze_gate_decisions.py}" MODEL="${MODEL:-}" DRAFT_MODEL="${DRAFT_MODEL:-}" @@ -20,6 +21,7 @@ WORKLOADS="${WORKLOADS:-humaneval,gsm8k,prose,north-star}" DECODE_MODES="${DECODE_MODES:-ar,speculation,adaptive}" ADAPTIVE_DRAFT_ALWAYS="${ADAPTIVE_DRAFT_ALWAYS:-on}" CONFIDENCE_ABLATION="${CONFIDENCE_ABLATION:-1}" +STEP_TIMING="${STEP_TIMING:-1}" CLIENTS="${CLIENTS:-1,2,3,4,6,8}" SLOTS="${SLOTS:-8}" MAX_CTX="${MAX_CTX:-8192}" @@ -27,6 +29,9 @@ MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" MAX_TOKENS="${MAX_TOKENS:-256}" WARMUP_TOKENS="${WARMUP_TOKENS:-16}" PROFILE_CONTEXT="${PROFILE_CONTEXT:-4096}" +CACHE_TYPE_K="${CACHE_TYPE_K:-q8_0}" +CACHE_TYPE_V="${CACHE_TYPE_V:-q8_0}" +FA_WINDOW="${FA_WINDOW:-0}" PORT="${PORT:-18138}" COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-3}" HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-900}" @@ -38,14 +43,17 @@ VISIBLE_DEVICES="${VISIBLE_DEVICES:-1}" usage() { cat <<'EOF' Usage: - MODEL=/path/Qwen3.8-27B-Q4_K_M.gguf \ - DRAFT_MODEL=/path/Qwen3.8-27B-DSpark-RadixArk-q4-mix.gguf \ + MODEL=/path/Qwen3.8-27B-PR625-IQ4_XS.gguf \ + DRAFT_MODEL=/path/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh The only supported drafter source for this matrix is: https://huggingface.co/RadixArk/Qwen3.8-27B-DSpark -DRAFT_MODEL is the q4-mix requantized drafter produced from that repository. +DRAFT_MODEL is a GGUF produced from that repository. The labeled R9700 +adaptive-selection baseline uses the no-YaRN Q8_0 artifact; changing draft +quantization changes both acceptance and cost, so it invalidates those stored +oracle ratios. The default fresh-process matrix runs ar, forced speculation, adaptive with always-drafting on, and an always-drafting confidence-off ablation at live concurrency 1,2,3,4,6,8 over HumanEval, GSM8K, and prose. The 2-code+4-chat @@ -53,6 +61,20 @@ north-star row runs only at C=6. The summary fails if adaptive-on is below 0.995 of the paired ar/speculation oracle in mean or median goodput or TTFT. The confidence ablation delta is reported but not gated. Set ADAPTIVE_DRAFT_ALWAYS=on,off only for an additional diagnostic arm. +The labeled adaptive-selection workload is fixed at C=6. Run it with +WORKLOADS=adaptive-selection CLIENTS=6 to compare adaptive decisions against +two strong speculation wins, two marginal wins, and two AR wins. +The adaptive-selection-c3 workload is fixed at C=3: two strong speculation +wins plus the strongest AR win. Run it with WORKLOADS=adaptive-selection-c3 +CLIENTS=3; the target adaptive outcome is a 2-spec+1-AR split that beats both +uniform decode modes. +STEP_TIMING=1 (the default) records phase-attributed JSON for every measured +decode round and writes profiling.json plus profiling.md. Set STEP_TIMING=0 +only when measuring the small overhead of diagnostic logging itself. +The screened prompt labels used Q8 KV and FA_WINDOW=2048. The concurrent +paged executor requires full attention, so this runner pins Q8 KV but defaults +FA_WINDOW=0 and records all three values in every case. Dense labels are priors; +the paired concurrent AR/speculation controls are the runtime oracle. Defaults select the second visible host GPU and address it as hip:0 inside the process (VISIBLE_DEVICES=1, TARGET_DEVICE=hip:0, DRAFT_DEVICE=hip:0). Override @@ -73,8 +95,11 @@ for value_name in REPEATS SLOTS MAX_CTX MAX_CONCURRENT_PREFILLS MAX_TOKENS WARMU [[ "$value" =~ ^[1-9][0-9]*$ ]] || { echo "$value_name must be positive" >&2; exit 2; } done [[ "$PORT" =~ ^[1-9][0-9]*$ ]] || { echo "PORT must be positive" >&2; exit 2; } +[[ "$FA_WINDOW" == 0 ]] || { echo "paged concurrency requires FA_WINDOW=0" >&2; exit 2; } +[[ -n "$CACHE_TYPE_K" && -n "$CACHE_TYPE_V" ]] || { echo "cache types must be non-empty" >&2; exit 2; } [[ "$COOLDOWN_SECONDS" =~ ^[0-9]+$ ]] || { echo "COOLDOWN_SECONDS must be non-negative" >&2; exit 2; } [[ "$CONFIDENCE_ABLATION" == 0 || "$CONFIDENCE_ABLATION" == 1 ]] || { echo "CONFIDENCE_ABLATION must be 0 or 1" >&2; exit 2; } +[[ "$STEP_TIMING" == 0 || "$STEP_TIMING" == 1 ]] || { echo "STEP_TIMING must be 0 or 1" >&2; exit 2; } (( PROFILE_CONTEXT < MAX_CTX )) || { echo "PROFILE_CONTEXT must be below MAX_CTX" >&2; exit 2; } [[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } @@ -109,7 +134,7 @@ reject_duplicates CLIENTS "${client_list[@]}" || exit 2 for workload in "${workload_list[@]}"; do case "$workload" in - humaneval|gsm8k|prose|north-star) ;; + humaneval|gsm8k|prose|north-star|adaptive-selection|adaptive-selection-c3) ;; *) echo "unknown workload $workload" >&2; exit 2 ;; esac done @@ -195,7 +220,11 @@ PY case_applicable() { local workload="$1" clients="$2" - [[ "$workload" != north-star || "$clients" == 6 ]] + case "$workload" in + north-star|adaptive-selection) [[ "$clients" == 6 ]] ;; + adaptive-selection-c3) [[ "$clients" == 3 ]] ;; + *) return 0 ;; + esac } run_case() { @@ -221,7 +250,8 @@ run_case() { --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_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 + --cache-type-k "$CACHE_TYPE_K" --cache-type-v "$CACHE_TYPE_V" + --fa-window "$FA_WINDOW" --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 20 --draft-residency persistent --decode-mode "$decode_mode" @@ -231,6 +261,7 @@ run_case() { "HIP_VISIBLE_DEVICES=$VISIBLE_DEVICES" "DFLASH_MAX_CONCURRENT_PREFILLS=$MAX_CONCURRENT_PREFILLS" "DFLASH_SPEC_BATCHED_DRAFT=1" + "DFLASH_STEP_TIMING=$STEP_TIMING" ) if [[ "$decode_mode" == adaptive ]]; then launch_env+=( @@ -266,6 +297,8 @@ run_case() { --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_DEVICE" --draft-model "$DRAFT_MODEL" --draft-model-sha256 "$DRAFT_MODEL_SHA256" --decode-mode "$decode_mode" + --cache-type-k "$CACHE_TYPE_K" --cache-type-v "$CACHE_TYPE_V" + --fa-window "$FA_WINDOW" ) [[ -n "$draft_always" ]] && metadata+=(--draft-always "$draft_always") [[ -n "$confidence" ]] && metadata+=(--confidence "$confidence") @@ -286,7 +319,7 @@ run_case() { local -a common_client=( --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" --clients "$clients" --prompt-file "$prompts" --prompt-offset 0 - --require-distinct-prompts --temperature 0 --ignore-eos + --require-distinct-prompts --temperature 0 --require-effective-prompt-telemetry --timeout "$REQUEST_TIMEOUT_SECONDS" --cooldown 0 ) @@ -298,6 +331,10 @@ run_case() { ) "${warmup_cmd[@]}" > "$case_dir/warmup.txt" + # Keep a measured-only log alongside the full startup/warmup log. This + # prevents warmup rounds from contaminating gate and phase distributions. + local benchmark_log_offset + benchmark_log_offset="$(wc -c < "$case_dir/server.log")" local -a benchmark_cmd=( python3 "$CLIENT" "${common_client[@]}" --max-tokens "$MAX_TOKENS" @@ -307,6 +344,8 @@ run_case() { ) "${benchmark_cmd[@]}" | tee "$case_dir/bench.txt" stop_server + tail -c "+$((benchmark_log_offset + 1))" "$case_dir/server.log" \ + > "$case_dir/benchmark-server.log" local -a proof_cmd=( python3 "$PROOF_TOOL" @@ -325,7 +364,7 @@ for ((repeat=1; repeat<=REPEATS; repeat++)); do for c_index in "${!client_list[@]}"; do clients="${client_list[$c_index]}" if ! case_applicable "$workload" "$clients"; then - echo "[skip] north-star is fixed at C=6; C=$clients" + echo "[skip] $workload is pinned to one concurrency; C=$clients" continue fi shift_by=$(((repeat + c_index) % ${#variants[@]})) @@ -339,5 +378,7 @@ for ((repeat=1; repeat<=REPEATS; repeat++)); do done (( active_cases > 0 )) || { echo "no applicable benchmark cases" >&2; exit 2; } +python3 "$ANALYZER" "$OUT" \ + --out "$OUT/profiling.json" --markdown-out "$OUT/profiling.md" python3 "$SUMMARIZER" "$OUT" --out "$OUT/summary.md" echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index a99562fa5..6109469ce 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -30,6 +30,7 @@ def load(name: str): dspark_generator = load("generate_dspark_prompts") proof = load("verify_feature_metrics") summary = load("summarize_feature_matrix") +gate_analysis = load("analyze_gate_decisions") def report( @@ -109,6 +110,7 @@ def test_dspark_profiles_reuse_public_fixtures_and_build_north_star(self) -> Non gsm8k = dspark_generator.build_records("gsm8k") prose = dspark_generator.build_records("prose") north_star = dspark_generator.build_records("north-star") + selection = dspark_generator.build_records("adaptive-selection") self.assertGreaterEqual(len(humaneval), 8) self.assertGreaterEqual(len(gsm8k), 8) self.assertGreaterEqual(len(prose), 8) @@ -116,6 +118,39 @@ def test_dspark_profiles_reuse_public_fixtures_and_build_north_star(self) -> Non self.assertEqual([row["suite"] for row in north_star[:2]], ["humaneval"] * 2) self.assertEqual([row["suite"] for row in north_star[2:]], ["prose"] * 4) self.assertEqual(len({row["prompt"] for row in north_star}), 6) + self.assertEqual(len(selection), 6) + self.assertEqual( + [row["selection_class"] for row in selection], + [ + "speculation_strong_win", + "speculation_strong_win", + "speculation_marginal_win", + "speculation_marginal_win", + "speculation_loss", + "speculation_loss", + ], + ) + self.assertEqual( + [row["expected_dense_oracle"] for row in selection], + ["speculation"] * 4 + ["ar"] * 2, + ) + selection_c3 = dspark_generator.build_records("adaptive-selection-c3") + self.assertEqual( + [row["id"] for row in selection_c3], + list(dspark_generator.ADAPTIVE_SELECTION_C3_IDS), + ) + self.assertEqual( + [row["selection_class"] for row in selection_c3], + [ + "speculation_strong_win", + "speculation_strong_win", + "speculation_loss", + ], + ) + self.assertEqual( + [row["expected_dense_oracle"] for row in selection_c3], + ["speculation", "speculation", "ar"], + ) class FeatureRunnerShellTests(unittest.TestCase): def run_invalid_matrix( @@ -192,6 +227,18 @@ def test_dspark_runner_has_explicit_confidence_ablation(self) -> None: self.assertIn('"DFLASH_SPEC_CONFIDENCE=0"', runner) self.assertIn('metadata+=(--confidence "$confidence")', runner) + def test_dspark_runner_profiles_only_the_measured_window(self) -> None: + runner = (HERE / "run_qwen38_dspark_matrix.sh").read_text( + encoding="utf-8", + ) + self.assertIn('STEP_TIMING="${STEP_TIMING:-1}"', runner) + self.assertIn('"DFLASH_STEP_TIMING=$STEP_TIMING"', runner) + self.assertIn('CACHE_TYPE_K="${CACHE_TYPE_K:-q8_0}"', runner) + self.assertIn('CACHE_TYPE_V="${CACHE_TYPE_V:-q8_0}"', runner) + self.assertIn('FA_WINDOW="${FA_WINDOW:-0}"', runner) + self.assertIn('"$case_dir/benchmark-server.log"', runner) + self.assertIn('--markdown-out "$OUT/profiling.md"', 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") @@ -218,6 +265,147 @@ def test_llama_only_does_not_require_lucebox_binary(self) -> None: +class GateAnalysisTests(unittest.TestCase): + def test_measured_step_timing_is_joined_and_summarized(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + prompt_dir = root / "prompts" + prompt_dir.mkdir() + (prompt_dir / "selection.jsonl").write_text( + json.dumps({ + "id": "friendly", + "prompt": "hello", + "selection_class": "speculation_strong_win", + "expected_dense_oracle": "speculation", + }) + "\n", + encoding="utf-8", + ) + case = root / "selection" / "c1" / "r1" / "adaptive-on" + case.mkdir(parents=True) + (case / "bench.json").write_text(json.dumps({ + "server_metadata": { + "workload": "selection", + "variant": "adaptive-on", + "clients": 1, + "repeat": 1, + }, + "levels": [{ + "aggregate_tok_s": 9.0, + "clients": 1, + "requests_detail": [{ + "request_id": "wire-1", + "prompt_index": 0, + "request_decode_tok_s": 8.0, + "content_sha256": "abc", + }], + }], + }), encoding="utf-8") + timing = { + "path": "ar", "live": 1, "k": 0, "decode_bucket": 1, + "n_prefill": 0, "max_kv_len": 10, + "draft_us": 20.0, "draft_lanes": 1, + "pre_us": 30.0, "graph_build_us": 5.0, + "graph_prepare_us": 5.0, "graph_exec_us": 40.0, + "sample_read_us": 10.0, "finish_us": 10.0, + "total_us": 100.0, "accepted_tokens": 0, + "emitted_tokens": 1, "target_forwards": 1, + } + metric_row = { + "request_id": "wire-1", "engine_request_id": 7, + "spec_accepted_tokens": 0, "spec_steps": 0, + "target_forwards": 1, "output_tokens": 1, + } + (case / "benchmark-server.log").write_text( + "[spec-gate] C=1 k=0 scores=[7:1.000/confidence] " + "sources=confidence:1,unavailable:0 " + "G(k)=0.010000 G(0)=0.020000 " + "predicted_cost=50.0us measured=ar-path\n" + f"[step-timing] {json.dumps(timing)}\n" + f"[concurrency-metrics] {json.dumps(metric_row)}\n", + encoding="utf-8", + ) + + report = gate_analysis.analyze_case(case) + self.assertEqual(report["log"], str(case / "benchmark-server.log")) + self.assertEqual(report["requests"][0]["prompt_id"], "friendly") + self.assertEqual(report["timing"]["by_path"]["ar"]["rounds"], 1) + self.assertAlmostEqual( + report["timing"]["draft_tax_on_ar"]["fraction_of_ar_wall"], + 0.2, + ) + self.assertAlmostEqual( + report["gate_by_k"][0]["realized_goodput_tok_s"], 10000.0, + ) + self.assertEqual(report["gate_timing_count_mismatch"], 0) + + def test_paired_controls_produce_a_per_prompt_concurrent_oracle(self) -> None: + base_request = { + "prompt_id": "p1", + "selection_class": "speculation_strong_win", + "expected_dense_oracle": "speculation", + "dense_r9700_baseline": {"spec_over_ar": 1.4}, + "output_sha256": "same", + } + cases = [] + for variant, rate in ( + ("ar", 10.0), ("speculation", 15.0), ("adaptive-on", 14.0), + ): + cases.append({ + "workload": "selection", "clients": 3, "repeat": 1, + "variant": variant, + "requests": [{ + **base_request, + "decode_tok_s": rate, + "admitted_fraction": 0.75, + "mean_confidence_yield": 2.5, + "commit_per_spec_step": 2.0, + }], + }) + comparison = gate_analysis.compare_prompts(cases)[0] + self.assertEqual( + comparison["empirical_concurrent_oracle"], "speculation", + ) + self.assertAlmostEqual(comparison["speculation_over_ar"], 1.5) + self.assertTrue(comparison["matches_expected_dense_oracle"]) + self.assertTrue(comparison["output_stable"]) + self.assertEqual( + comparison["adaptive"]["adaptive-on"]["admitted_fraction"], 0.75, + ) + + def test_activation_shapes_are_compared_at_the_same_live_concurrency(self) -> None: + def shape(path: str, live: int, k: int, rate: float, draft: float) -> dict: + return { + "path": path, "live": live, "k": k, "rounds": 4, + "round_goodput_tok_s": rate, "mean_total_us": 100.0, + "phase_mean_us": {"draft_us": draft}, + } + + cases = [ + { + "workload": "selection", "clients": 3, "repeat": 1, + "variant": "ar", "aggregate_tok_s": 60.0, + "timing": {"by_shape": [ + shape("ar", 2, 0, 50.0, 0.0), + shape("ar", 3, 0, 70.0, 0.0), + ]}, + }, + { + "workload": "selection", "clients": 3, "repeat": 1, + "variant": "adaptive-on", "aggregate_tok_s": 42.0, + "timing": {"by_shape": [ + shape("ar", 2, 0, 40.0, 10.0), + shape("spec", 3, 2, 49.0, 15.0), + ]}, + }, + ] + rows = gate_analysis.compare_activation_shapes(cases) + self.assertEqual(len(rows), 2) + self.assertAlmostEqual(cases[1]["aggregate_over_ar"], 0.7) + self.assertEqual(rows[0]["activation_outcome"], "ar-with-draft-tax") + self.assertEqual(rows[1]["activation_outcome"], "unprofitable") + self.assertAlmostEqual(rows[1]["realized_over_pure_ar"], 0.7) + + class FeatureProofTests(unittest.TestCase): def test_full_below_pool_passes_without_page_traffic(self) -> None: rows = [ diff --git a/harness/benchmarks/concurrency/write_feature_metadata.py b/harness/benchmarks/concurrency/write_feature_metadata.py index b2005ea25..e6c48073b 100644 --- a/harness/benchmarks/concurrency/write_feature_metadata.py +++ b/harness/benchmarks/concurrency/write_feature_metadata.py @@ -91,6 +91,9 @@ def main() -> int: parser.add_argument("--draft-model", type=pathlib.Path) parser.add_argument("--draft-model-sha256") parser.add_argument("--decode-mode", choices=("ar", "speculation", "adaptive")) + parser.add_argument("--cache-type-k") + parser.add_argument("--cache-type-v") + parser.add_argument("--fa-window", type=int) parser.add_argument("--draft-always", choices=("on", "off")) parser.add_argument("--confidence", choices=("on", "off")) parser.add_argument("--ddtree", action="store_true") @@ -143,6 +146,12 @@ def main() -> int: literal_flags += ["--draft-device", args.draft_device] if args.decode_mode: literal_flags += ["--decode-mode", args.decode_mode] + if args.cache_type_k: + literal_flags += ["--cache-type-k", args.cache_type_k] + if args.cache_type_v: + literal_flags += ["--cache-type-v", args.cache_type_v] + if args.fa_window is not None: + literal_flags += ["--fa-window", str(args.fa_window)] if args.ddtree: literal_flags += ["--ddtree"] if args.ddtree_budget is not None: @@ -183,6 +192,9 @@ def main() -> int: "draft_model": str(args.draft_model.resolve()) if args.draft_model else None, "draft_model_sha256": draft_model_sha256, "decode_mode": args.decode_mode, + "cache_type_k": args.cache_type_k, + "cache_type_v": args.cache_type_v, + "fa_window": args.fa_window, "draft_always": args.draft_always, "confidence": args.confidence, "ddtree": args.ddtree, diff --git a/harness/benchmarks/prompts/qwen38_dspark_adaptive_selection.jsonl b/harness/benchmarks/prompts/qwen38_dspark_adaptive_selection.jsonl new file mode 100644 index 000000000..a2721c465 --- /dev/null +++ b/harness/benchmarks/prompts/qwen38_dspark_adaptive_selection.jsonl @@ -0,0 +1,6 @@ +{"dense_r9700_baseline":{"ar_decode_tok_s":34.89,"lossless":true,"spec_accept_pct":35.8,"spec_avg_commit":2.51,"spec_decode_tok_s":46.35,"spec_over_ar":1.328},"expected_dense_oracle":"speculation","id":"adaptive-he-09-sum-product","prompt":"Complete the following Python function.\n\nfrom typing import List, Tuple\n\ndef sum_product(numbers: List[int]) -> Tuple[int, int]:\n \"\"\" For a given list of integers, return a tuple consisting of a sum and a product of all the integers in a list.\n Empty sum should be equal to 0 and empty product should be equal to 1.\n >>> sum_product([])\n (0, 1)\n >>> sum_product([1, 2, 3, 4])\n (10, 24)\n \"\"\"\n","selection_class":"speculation_strong_win","source_id":"he_09","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":34.89,"lossless":true,"spec_accept_pct":38.9,"spec_avg_commit":2.72,"spec_decode_tok_s":50.3,"spec_over_ar":1.442},"expected_dense_oracle":"speculation","id":"adaptive-he-10-rolling-max","prompt":"Complete the following Python function.\n\nfrom typing import List\n\ndef rolling_max(numbers: List[int]) -> List[int]:\n \"\"\" From a given list of integers, generate a list of rolling maximum element found until given moment\n in the sequence.\n >>> rolling_max([1, 2, 3, 2, 3, 4, 2])\n [1, 2, 3, 3, 3, 4, 4]\n \"\"\"\n","selection_class":"speculation_strong_win","source_id":"he_10","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":34.95,"lossless":true,"spec_accept_pct":29.0,"spec_avg_commit":2.03,"spec_decode_tok_s":37.69,"spec_over_ar":1.078},"expected_dense_oracle":"speculation","id":"adaptive-he-02-separate-paren-groups","prompt":"Complete the following Python function.\n\nfrom typing import List\n\ndef separate_paren_groups(paren_string: str) -> List[str]:\n \"\"\" Input to this function is a string containing multiple groups of nested parentheses. Your goal is to\n separate those group into separate strings and return the list of those.\n Separate groups are balanced (each open brace is properly closed) and not nested within each other\n Ignore any spaces in the input string.\n >>> separate_paren_groups('( ) (( )) (( )( ))')\n ['()', '(())', '(()())']\n \"\"\"\n","selection_class":"speculation_marginal_win","source_id":"he_02","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":35.15,"lossless":true,"spec_accept_pct":28.9,"spec_avg_commit":2.02,"spec_decode_tok_s":37.44,"spec_over_ar":1.065},"expected_dense_oracle":"speculation","id":"adaptive-he-03-truncate-number","prompt":"Complete the following Python function.\n\ndef truncate_number(number: float) -> float:\n \"\"\" Given a positive floating point number, it can be decomposed into\n and integer part (largest integer smaller than given number) and decimals\n (leftover part always smaller than 1).\n\n Return the decimal part of the number.\n >>> truncate_number(3.5)\n 0.5\n \"\"\"\n","selection_class":"speculation_marginal_win","source_id":"he_03","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":35.04,"lossless":true,"spec_accept_pct":26.2,"spec_avg_commit":1.84,"spec_decode_tok_s":33.95,"spec_over_ar":0.969},"expected_dense_oracle":"ar","id":"adaptive-he-08-filter-by-substring","prompt":"Complete the following Python function.\n\nfrom typing import List\n\ndef filter_by_substring(strings: List[str], substring: str) -> List[str]:\n \"\"\" Filter an input list of strings only for ones that contain given substring\n >>> filter_by_substring([], 'a')\n []\n >>> filter_by_substring(['abc', 'bacd', 'cde', 'array'], 'a')\n ['abc', 'bacd', 'array']\n \"\"\"\n","selection_class":"speculation_loss","source_id":"he_08","suite":"humaneval"} +{"dense_r9700_baseline":{"ar_decode_tok_s":34.9,"lossless":true,"spec_accept_pct":20.8,"spec_avg_commit":1.45,"spec_decode_tok_s":27.1,"spec_over_ar":0.777},"expected_dense_oracle":"ar","id":"adaptive-prose-01-reproducibility","prompt":"Write a clear, self-contained technical essay of about 500 words on why reproducible benchmarks need immutable inputs and explicit hardware metadata. Include one concrete example and end with a concise conclusion.","selection_class":"speculation_loss","source_id":"prose-01","suite":"prose"} diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 6570a803a..c879a4eb2 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -195,6 +195,14 @@ bool Qwen35SeqEngine::spec_gate_debug_enabled() const { return value && std::atoi(value) != 0; } +bool Qwen35SeqEngine::step_timing_enabled() { + static const bool enabled = []() { + const char * value = std::getenv("DFLASH_STEP_TIMING"); + return value && std::atoi(value) != 0; + }(); + return enabled; +} + bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { speculation_gate_.reset(); if (spec_mode_ != SpecMode::dspark_chain || !capture_features_ || @@ -690,6 +698,28 @@ bool Qwen35SeqEngine::prepare_chain_drafts( const std::vector & selected) { if (selected.size() != inputs.size()) return false; + // Accumulate the full drafting wall (draft graph compute + fused + // selector + readbacks) into the round's [step-timing] attribution, + // including early-failure paths. + struct DraftTimer { + Qwen35SeqEngine * engine; + std::chrono::steady_clock::time_point start; + int lanes; + ~DraftTimer() { + if (!engine) return; + engine->round_draft_us_ += + std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + engine->round_draft_lanes_ += lanes; + } + }; + DraftTimer draft_timer{ + step_timing_enabled() ? this : nullptr, + std::chrono::steady_clock::now(), + (int)std::count_if( + selected.begin(), selected.end(), + [](uint8_t value) { return value != 0; })}; + const int T = tree_width_; const int hidden = b_.w_.n_embd; struct Lane { @@ -923,7 +953,8 @@ bool Qwen35SeqEngine::chain_spec_input_eligible( return slots_.slot(in.slot).generated_tokens() >= floor; } std::optional Qwen35SeqEngine::step_chain_spec( - const StepPlan & plan, const std::vector & admitted) { + const StepPlan & plan, const std::vector & admitted, + std::chrono::steady_clock::time_point round_started) { StepResult result; const std::vector & inputs = plan.decode; if (admitted.size() != inputs.size() || !plan.prefills.empty()) { @@ -940,6 +971,20 @@ std::optional Qwen35SeqEngine::step_chain_spec( if (spec_count == 0) return std::nullopt; const int tree_bucket = chain_decode_bucket_width(spec_count); + // Optional per-phase wall attribution ([step-timing]). Timestamps mark + // phase boundaries; graph computes are synchronous on this backend, so + // each *_exec span covers upload + compute up to its trailing sync. + const bool timing = step_timing_enabled(); + using timing_clock = std::chrono::steady_clock; + const auto t_round_start = round_started; + timing_clock::time_point t_verify_build_start, t_verify_build_end, + t_verify_exec_end, t_posterior_end, t_commit_end, + t_replay_build_end, t_replay_exec_end, t_sample_end; + auto span_us = [](timing_clock::time_point from, + timing_clock::time_point to) { + return std::chrono::duration(to - from).count(); + }; + struct Proposal { size_t input_index = 0; int slot = -1; @@ -1039,6 +1084,7 @@ std::optional Qwen35SeqEngine::step_chain_spec( max_prefix = std::max( max_prefix, slots_.slot(proposal.slot).cur_pos); } + t_verify_build_start = timing_clock::now(); if (!build_target_step_paged_tree( tree_sg, b_.w_, b_.cache_, b_.target_backend_, T, tree_bucket, max_prefix, @@ -1047,6 +1093,7 @@ std::optional Qwen35SeqEngine::step_chain_spec( result.error = "packed DSpark chain verify graph build failed"; return result; } + t_verify_build_end = timing_clock::now(); const int total_tree = T * tree_bucket; std::vector flat_tokens(static_cast(total_tree), 0); @@ -1132,11 +1179,13 @@ std::optional Qwen35SeqEngine::step_chain_spec( result.error = "packed DSpark chain verify compute failed"; return result; } + t_verify_exec_end = timing_clock::now(); std::vector posterior(static_cast(total_tree), -1); ggml_backend_tensor_get( tree_sg.argmax_tokens, posterior.data(), 0, sizeof(int32_t) * posterior.size()); + t_posterior_end = timing_clock::now(); int replay_total = 0; for (int lane = 0; lane < spec_count; ++lane) { @@ -1249,6 +1298,7 @@ std::optional Qwen35SeqEngine::step_chain_spec( result.error = "DSpark mixed-step block-table refresh failed"; return result; } + t_commit_end = timing_clock::now(); // Launch 2: accepted path segments + compact AR rows in the same builder // combination already used by mixed prefill/decode. @@ -1272,6 +1322,7 @@ std::optional Qwen35SeqEngine::step_chain_spec( result.error = "DSpark mixed commit/AR graph build failed"; return result; } + t_replay_build_end = timing_clock::now(); std::vector durable_tokens(static_cast(n_total), 0); std::copy(replay_tokens.begin(), replay_tokens.end(), @@ -1402,6 +1453,7 @@ std::optional Qwen35SeqEngine::step_chain_spec( result.error = "DSpark mixed commit/AR compute failed"; return result; } + t_replay_exec_end = timing_clock::now(); argmax_buf_.assign(static_cast(gather_rows), -1); ggml_backend_tensor_get_async( @@ -1409,6 +1461,7 @@ std::optional Qwen35SeqEngine::step_chain_spec( argmax_buf_.data(), 0, sizeof(int32_t) * argmax_buf_.size()); ggml_backend_synchronize(b_.target_backend_); + t_sample_end = timing_clock::now(); for (int lane = 0; lane < spec_count; ++lane) { Proposal & proposal = proposals[static_cast(lane)]; proposal.pending = argmax_buf_[static_cast(lane)]; @@ -1482,6 +1535,37 @@ std::optional Qwen35SeqEngine::step_chain_spec( attach_residency_telemetry(out); result.decode.push_back(std::move(out)); } + if (timing) { + const auto t_round_end = timing_clock::now(); + std::fprintf(stderr, + "[step-timing] {\"path\":\"spec\",\"live\":%d,\"k\":%d," + "\"tree_bucket\":%d,\"tree_rows\":%d,\"replay_rows\":%d," + "\"ar_lanes\":%d,\"ar_bucket\":%d,\"max_kv_len\":%d," + "\"draft_us\":%.1f,\"draft_lanes\":%d," + "\"pre_us\":%.1f,\"verify_build_us\":%.1f," + "\"verify_exec_us\":%.1f,\"posterior_read_us\":%.1f," + "\"commit_cpu_us\":%.1f,\"replay_build_us\":%.1f," + "\"replay_exec_us\":%.1f,\"sample_read_us\":%.1f," + "\"finish_us\":%.1f,\"total_us\":%.1f," + "\"accepted_tokens\":%d,\"emitted_tokens\":%d," + "\"target_forwards\":%d}\n", + (int)inputs.size(), spec_count, + tree_bucket, T * tree_bucket, replay_total, + ar_count, ar_bucket, max_kv_len, + round_draft_us_, round_draft_lanes_, + span_us(t_round_start, t_verify_build_start), + span_us(t_verify_build_start, t_verify_build_end), + span_us(t_verify_build_end, t_verify_exec_end), + span_us(t_verify_exec_end, t_posterior_end), + span_us(t_posterior_end, t_commit_end), + span_us(t_commit_end, t_replay_build_end), + span_us(t_replay_build_end, t_replay_exec_end), + span_us(t_replay_exec_end, t_sample_end), + span_us(t_sample_end, t_round_end), + span_us(t_round_start, t_round_end), + replay_total - spec_count, replay_total + ar_count, + 2 * spec_count + ar_count); + } return result; } @@ -2171,7 +2255,18 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // by taking the existing packed AR path for this iteration. } + // One common wall-clock origin makes pure AR, adaptive k=0, and admitted + // speculative rounds directly comparable. In adaptive always-draft mode + // round_draft_us_ is a subset of this wall, not an extra cost to add. + const bool timing = step_timing_enabled(); + using timing_clock = std::chrono::steady_clock; + const auto decode_round_started = timing + ? timing_clock::now() : timing_clock::time_point{}; + if (spec_mode_ == SpecMode::dspark_chain && plan.prefills.empty()) { + // New chain round: restart the [step-timing] draft attribution. + round_draft_us_ = 0.0; + round_draft_lanes_ = 0; const auto chain_started = std::chrono::steady_clock::now(); std::vector admitted(inputs.size(), 0); SpecPlan gate_plan; @@ -2270,7 +2365,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { [](uint8_t value) { return value != 0; }); if (any_admitted) { std::optional speculative = - step_chain_spec(plan, admitted); + step_chain_spec(plan, admitted, decode_round_started); const double measured_us = std::chrono::duration( std::chrono::steady_clock::now() - chain_started).count(); @@ -2298,6 +2393,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { StepGraph & sg = b_.sg_; const int hidden = w.n_embd; const int n_head_kv = w.n_head_kv; + timing_clock::time_point t_ar_build_start, t_ar_build_end, + t_ar_exec_start, t_ar_exec_end, t_ar_read_end; decode_outputs.reserve(inputs.size()); prefill_outputs.reserve(plan.prefills.size()); @@ -2415,6 +2512,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { : std::max(1, n_commits)) : 0; + if (timing) t_ar_build_start = timing_clock::now(); bool built = false; if (with_prefill) { built = build_target_step( @@ -2462,6 +2560,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { !sg.logits_row_indices))) { return fail_step("packed prefill/decode graph build failed"); } + if (timing) t_ar_build_end = timing_clock::now(); embed_buf_.resize((size_t)hidden * n_total); int token_offset = 0; @@ -2602,6 +2701,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { b_.target_backend_, b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, sizeof(int32_t) * seq_lens_.size()); + if (timing) t_ar_exec_start = timing_clock::now(); ggml_status st = GGML_STATUS_FAILED; { const Qwen35RoctxRange roctx_compute( @@ -2611,6 +2711,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (st != GGML_STATUS_SUCCESS) { return fail_step("packed prefill/decode compute failed"); } + if (timing) t_ar_exec_end = timing_clock::now(); const int decode_row0 = with_prefill ? n_commits : 0; const int argmax_rows = with_prefill ? gather_rows : decode_bucket; @@ -2623,6 +2724,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { "qwen35.argmax_readback", roctx_metadata); ggml_backend_synchronize(b_.target_backend_); } + if (timing) t_ar_read_end = timing_clock::now(); std::vector write_slots; write_slots.reserve(live_slot_ids_.size() + prefills.size()); write_slots.insert(write_slots.end(), @@ -2669,6 +2771,33 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } prefill_outputs.push_back(std::move(out)); } + if (timing && plan.prefills.empty() && live_count > 0) { + const auto t_ar_end = timing_clock::now(); + auto span_us = [](timing_clock::time_point from, + timing_clock::time_point to) { + return std::chrono::duration( + to - from).count(); + }; + std::fprintf(stderr, + "[step-timing] {\"path\":\"ar\",\"live\":%d,\"k\":0," + "\"decode_bucket\":%d,\"n_prefill\":0,\"max_kv_len\":%d," + "\"draft_us\":%.1f,\"draft_lanes\":%d," + "\"pre_us\":%.1f,\"graph_build_us\":%.1f," + "\"graph_prepare_us\":%.1f,\"graph_exec_us\":%.1f," + "\"sample_read_us\":%.1f,\"finish_us\":%.1f," + "\"total_us\":%.1f,\"accepted_tokens\":0," + "\"emitted_tokens\":%d,\"target_forwards\":%d}\n", + live_count, decode_bucket, max_kv_len, + round_draft_us_, round_draft_lanes_, + span_us(decode_round_started, t_ar_build_start), + span_us(t_ar_build_start, t_ar_build_end), + span_us(t_ar_build_end, t_ar_exec_start), + span_us(t_ar_exec_start, t_ar_exec_end), + span_us(t_ar_exec_end, t_ar_read_end), + span_us(t_ar_read_end, t_ar_end), + span_us(decode_round_started, t_ar_end), + live_count, live_count); + } return result; } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index ec6082a16..94bff3717 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -29,6 +29,7 @@ #include "qwen35_slot_manager.h" #include +#include #include #include #include @@ -147,11 +148,16 @@ class Qwen35SeqEngine final : public SeqEngine { bool batched_drafting_enabled() const; bool draft_always_enabled() const; bool confidence_scoring_enabled() const; + // DFLASH_STEP_TIMING=1 emits one [step-timing] JSON line per decode + // round attributing wall time to draft, verify, readback, CPU commit, + // replay, and packed-AR phases. Diagnostic only; off by default. + static bool step_timing_enabled(); // 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_chain_spec( - const StepPlan & plan, const std::vector & admitted); + const StepPlan & plan, const std::vector & admitted, + std::chrono::steady_clock::time_point round_started); Qwen35Backend & b_; Qwen35SlotManager slots_; @@ -171,6 +177,10 @@ class Qwen35SeqEngine final : public SeqEngine { std::unique_ptr speculation_gate_; std::vector last_survival_score_; std::vector last_survival_generated_; + // Per-round draft cost accumulator for [step-timing]; reset at the top + // of each dspark_chain round, accumulated by prepare_chain_drafts. + double round_draft_us_ = 0.0; + int round_draft_lanes_ = 0; // Hoisted per-step buffers (reused across step() calls). std::vector output_rows_; std::vector live_tokens_; From 73546172af18a2bbe2ab7e9f21ab062a286497b6 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 14:40:33 +0000 Subject: [PATCH 32/42] feat(concurrency): make DSpark activation request-sticky --- .../QWEN38_DSPARK_ADAPTIVE_SELECTION.md | 22 +- .../benchmarks/concurrency/FEATURE_MATRIX.md | 37 +- .../concurrency/analyze_gate_decisions.py | 160 +++++- .../concurrency/run_qwen38_dspark_matrix.sh | 44 +- .../concurrency/summarize_feature_matrix.py | 39 +- .../concurrency/test_feature_tools.py | 171 +++++- .../src/common/concurrency/speculation_gate.h | 532 +++++++++++++++--- server/src/common/model_backend.h | 8 + server/src/common/speculation_policy.h | 17 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 400 +++++++++---- .../qwen35/concurrency/qwen35_seq_engine.h | 4 +- server/src/qwen35/qwen35_backend.cpp | 53 +- server/src/qwen35/qwen35_backend.h | 5 + server/src/server/scheduler.cpp | 18 + server/test/test_server_unit.cpp | 20 + server/test/test_speculation_gate.cpp | 376 ++++++++++--- 16 files changed, 1516 insertions(+), 390 deletions(-) diff --git a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md index 1e0d22818..53d00e227 100644 --- a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md +++ b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md @@ -17,8 +17,8 @@ The exact completion lengths and ordered output hashes are retained locally in DSpark is the fixed proposal mechanism for this fixture. The benchmark is about adaptive activation: which lanes the gate admits at each live -concurrency, whether the selected `k` beats matched pure AR, how much rejected -always-drafting costs, and which target phase dominates a bad decision. +concurrency, whether the selected `k` beats matched pure AR, whether rejected +lanes avoid drafting, and which target phase dominates a bad decision. ## Dense screening baseline @@ -100,7 +100,19 @@ harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh The dense oracle labels are priors, not a hard assertion about the concurrent executor. In particular, the two marginal dense wins may correctly become AR choices after packed-tree, replay, padding, and synchronization costs are -included. The useful adaptive behavior is: +included. Adaptive mode makes one activation decision per request: it scores +the request at its first target-decode step, then keeps the chosen AR or +speculation mode through EOS/max tokens. + +The prefill logits produce the first sampled output token before a +target-decode step exists. Activation therefore happens before the first +target-decode execution, with no preliminary AR decode round; a request that +retires directly from prefill (for example, `max_tokens=1`) has no decode mode +to activate. Mixed-prefill and min-token-floor rounds may temporarily execute +through the safe AR graph after the sticky decision is made; they do not +reclassify the request. + +The useful adaptive behavior is: 1. rank the two strong wins above the two known losses; 2. keep the prose loss and normally the HumanEval loss in AR; @@ -135,8 +147,8 @@ Read the report in this order: `(live, k, path)` shape to pure AR at the same live concurrency. 4. Inspect `Gate prediction calibration` for predicted-versus-realized goodput and realized-versus-AR regret. -5. Use phase attribution to choose the next optimization: draft tax on k=0, - verify, or the structural replay forward. +5. Use phase attribution to choose the next optimization: unexpected draft + work on k=0, verify, or the structural replay forward. The profiler's `total_us` uses one common decode-round origin for pure AR, adaptive k=0, and speculative rounds. `draft_us` is a subset of that wall diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index e33e719c6..b4668e5e1 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -7,36 +7,45 @@ recorded in [`STRIX_HALO_RESULTS.md`](STRIX_HALO_RESULTS.md). `run_qwen38_dspark_matrix.sh` is the C7 acceptance-gated speculation matrix. It uses only the [RadixArk Qwen3.8-27B-DSpark](https://huggingface.co/RadixArk/Qwen3.8-27B-DSpark) -source. `DRAFT_MODEL` must point to the **q4-mix requantized drafter** produced -from that repository. +source. `DRAFT_MODEL` must point to a GGUF produced from that repository. The +labeled R9700 baseline uses the no-YaRN Q8_0 artifact. ```bash MODEL=/opt/models/Qwen3.8-27B-Q4_K_M.gguf \ -DRAFT_MODEL=/opt/models/Qwen3.8-27B-DSpark-RadixArk-q4-mix.gguf \ +DRAFT_MODEL=/opt/models/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ REPEATS=5 \ harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh ``` The default fresh-process matrix is: -- `ar`, `speculation`, `adaptive-on`, and `adaptive-confidence-off`. The - last arm keeps always-drafting enabled but sets - `DFLASH_SPEC_CONFIDENCE=0`; with no historical fallback, adaptive requests - remain on AR while still paying drafting cost. This directly prices the - confidence signal. `ADAPTIVE_DRAFT_ALWAYS=on,off` adds the old admitted-only - drafting arm as an optional diagnostic, not a production acceptance path. +- `ar`, `speculation`, `adaptive-on`, and `adaptive-confidence-off`. + `adaptive-on` batch-scores each new adaptive request once at its first + target-decode step, chooses AR or speculation, and keeps that mode through + retirement. In a decode-only round the one-time bootstrap draft is reused + immediately when the request chooses speculation; later AR rounds do not + draft. A target-decode step exists only after prefill has produced the first + sampled token, so one-token requests that retire directly from prefill never + enter adaptive activation. Mixed-prefill and min-token-floor rounds are + temporary execution constraints: the sticky decision is still committed, + but speculation starts on the first executor-safe round. The last arm sets + `DFLASH_SPEC_CONFIDENCE=0`, suppressing activation scoring so it remains a + draft-free AR ablation. There is no eager-draft matrix axis. - Live concurrency `C ∈ {1,2,3,4,6,8}` over the checked-in HumanEval and GSM8K cohorts plus deterministic prose prompts. - A fixed C=6 north-star cohort with two code and four chat requests. -- Batched drafting enabled for every row; adaptive startup profiling uses a - 4096-token synthetic context. +- Batched drafting enabled for DSpark rows; adaptive drafts the one-time cold + batch and thereafter only requests whose sticky mode is speculation. + Startup profiling uses a 4096-token synthetic context. Every process records the target, server, shared-library, and drafter hashes, the literal command and launch environment, startup pool dimensions, request IDs, and terminal concurrency counters. The proof rejects forced-speculation rows unless every measured request has positive `spec_steps`. Adaptive rows -may legitimately choose k=0, but must show both the packed DSpark startup -marker and a completed startup cost profile. Chain rows must keep all +may legitimately choose k=0, but must show exactly one finite +`[spec-activation]` AR/speculation decision for every measured adaptive +request, plus both the packed DSpark startup marker and a completed startup +cost profile. Chain rows must keep all `ddtree_*` counters at zero, preserving the DDTree proof semantics below. For every workload/concurrency pair, the summarizer forms a paired oracle: @@ -52,7 +61,7 @@ ratios misses the gate; p95 alone is never used as acceptance evidence. A second table reports paired goodput and inverse-TTFT deltas between `adaptive-on` and `adaptive-confidence-off` without gating the ablation. Set `WORKLOADS`, `CLIENTS`, `DECODE_MODES`, -`ADAPTIVE_DRAFT_ALWAYS`, or `CONFIDENCE_ABLATION=0` to select a smaller +or `CONFIDENCE_ABLATION=0` to select a smaller diagnostic subset. ## Qwen3.6 DDTree/PFlash/KVFlash matrix diff --git a/harness/benchmarks/concurrency/analyze_gate_decisions.py b/harness/benchmarks/concurrency/analyze_gate_decisions.py index da5cf2fb7..1ff37b3ae 100644 --- a/harness/benchmarks/concurrency/analyze_gate_decisions.py +++ b/harness/benchmarks/concurrency/analyze_gate_decisions.py @@ -34,6 +34,7 @@ ) METRIC_RE = re.compile(r"\[concurrency-metrics\] (?P\{.*\})") TIMING_RE = re.compile(r"\[step-timing\] (?P\{.*\})") +ACTIVATION_RE = re.compile(r"\[spec-activation\]\s+(?P.*)$") TIMING_COUNT_FIELDS = ( "live", "k", "emitted_tokens", "accepted_tokens", "target_forwards", ) @@ -72,12 +73,44 @@ def _validate_timing(row: dict[str, Any], path: Path, line_no: int) -> None: ) +def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: + for key in ("request_id", "slot"): + value = row.get(key) + if type(value) is not int or value < 0: + raise ValueError( + f"{path}:{line_no}: spec-activation {key} must be a " + "non-negative int" + ) + for key in ("initial_confidence", "calibrated_yield"): + value = row.get(key) + if ( + type(value) not in (int, float) + or not math.isfinite(value) + or value < 1.0 + ): + raise ValueError( + f"{path}:{line_no}: spec-activation {key} must be finite " + "and at least 1" + ) + if row.get("decision") not in ("ar", "speculation"): + raise ValueError( + f"{path}:{line_no}: spec-activation decision must be ar or " + "speculation" + ) + + def parse_server_log( path: Path, -) -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]], list[dict[str, Any]]]: +) -> tuple[ + list[dict[str, Any]], + dict[str, dict[str, Any]], + list[dict[str, Any]], + list[dict[str, Any]], +]: rounds: list[dict[str, Any]] = [] metrics: dict[str, dict[str, Any]] = {} timings: list[dict[str, Any]] = [] + activations: list[dict[str, Any]] = [] for line_no, line in enumerate( path.read_text(encoding="utf-8", errors="replace").splitlines(), 1, ): @@ -108,6 +141,12 @@ def parse_server_log( "entries": entries, }) continue + activation = ACTIVATION_RE.search(line) + if activation: + row = _json_object(activation, path, line_no) + _validate_activation(row, path, line_no) + activations.append(row) + continue metric = METRIC_RE.search(line) if metric: row = _json_object(metric, path, line_no) @@ -124,7 +163,7 @@ def parse_server_log( row = _json_object(timing, path, line_no) _validate_timing(row, path, line_no) timings.append(row) - return rounds, metrics, timings + return rounds, metrics, timings, activations def _prompt_records(prompt_file: Path | None) -> list[dict[str, Any]]: @@ -336,6 +375,95 @@ def _class_summary(requests: list[dict[str, Any]]) -> list[dict[str, Any]]: return out +def _activation_proof( + variant: str, + activations: list[dict[str, Any]], + metrics: dict[str, dict[str, Any]], + measured_wire_ids: set[str], + log_path: Path, +) -> tuple[dict[str, Any], dict[int, dict[str, Any]]]: + expected_ids: set[int] = set() + engine_to_wire: dict[int, str] = {} + mapping_errors: list[str] = [] + for wire_id, metric in metrics.items(): + engine_id = metric.get("engine_request_id") + if type(engine_id) is not int or engine_id < 0: + mapping_errors.append( + f"metric {wire_id} has no non-negative integer engine_request_id" + ) + continue + previous = engine_to_wire.get(engine_id) + if previous is not None and previous != wire_id: + mapping_errors.append( + f"engine request {engine_id} maps both {previous} and {wire_id}" + ) + continue + engine_to_wire[engine_id] = wire_id + expected_ids.add(engine_id) + + by_id: dict[int, list[dict[str, Any]]] = defaultdict(list) + for activation in activations: + by_id[int(activation["request_id"])].append(activation) + + required = variant == "adaptive-on" + if required: + missing_metric_ids = sorted(measured_wire_ids - set(metrics)) + unknown_metric_ids = sorted(set(metrics) - measured_wire_ids) + duplicate_ids = sorted( + request_id for request_id, rows in by_id.items() if len(rows) != 1 + ) + unknown_ids = sorted(set(by_id) - expected_ids) + missing_ids = sorted(expected_ids - set(by_id)) + errors = list(mapping_errors) + if missing_metric_ids: + errors.append( + "missing engine metrics for measured wire requests " + + ",".join(missing_metric_ids) + ) + if unknown_metric_ids: + errors.append( + "engine metrics reference unknown measured wire requests " + + ",".join(unknown_metric_ids) + ) + if duplicate_ids: + errors.append( + "duplicate activations for engine requests " + + ",".join(str(value) for value in duplicate_ids) + ) + if unknown_ids: + errors.append( + "activations reference unknown engine requests " + + ",".join(str(value) for value in unknown_ids) + ) + if missing_ids: + errors.append( + "missing activations for engine requests " + + ",".join(str(value) for value in missing_ids) + ) + if errors: + raise ValueError( + f"{log_path}: adaptive-on activation proof failed: " + + "; ".join(errors) + ) + + unique = {request_id: rows[0] for request_id, rows in by_id.items()} + decisions = { + decision: sum( + 1 for row in activations if row["decision"] == decision + ) + for decision in ("ar", "speculation") + } + return ({ + "required": required, + "validation": "passed" if required else "not-required", + "measured_engine_requests": len(expected_ids), + "records": len(activations), + "unique_requests": len(by_id), + "matched_requests": len(set(by_id) & expected_ids), + "decision_counts": decisions, + }, unique) + + def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, Any]: bench_path = case_dir / "bench.json" bench = json.loads(bench_path.read_text(encoding="utf-8")) @@ -351,11 +479,14 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A log_path = case_dir / "benchmark-server.log" if not log_path.is_file(): log_path = case_dir / "server.log" - rounds, metrics, timings = parse_server_log(log_path) + rounds, metrics, timings, activations = parse_server_log(log_path) resolved_prompt_file = _find_prompt_file( case_dir, workload, prompt_file, ) by_wire = prompt_names(bench, resolved_prompt_file) + activation_summary, activation_by_engine = _activation_proof( + variant, activations, metrics, set(by_wire), log_path, + ) engine_to_prompt: dict[int, dict[str, Any]] = {} for wire_id, row in metrics.items(): @@ -393,6 +524,7 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A stats = per_request[engine_id] steps = prompt.get("spec_steps", 0) accepted = prompt.get("spec_accepted_tokens", 0) + activation = activation_by_engine.get(engine_id) requests.append({ "engine_request_id": engine_id, **prompt, @@ -404,6 +536,20 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A "mean_confidence_yield": ( stats["score_sum"] / stats["scored"] if stats["scored"] else None ), + "activation_slot": ( + activation.get("slot") if activation is not None else None + ), + "initial_confidence": ( + activation.get("initial_confidence") + if activation is not None else None + ), + "calibrated_yield": ( + activation.get("calibrated_yield") + if activation is not None else None + ), + "activation_decision": ( + activation.get("decision") if activation is not None else None + ), "commit_per_spec_step": ( (steps + accepted) / steps if steps else None ), @@ -448,6 +594,7 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A "repeat": repeat, "variant": variant, "aggregate_tok_s": aggregate_tok_s, + "activation": activation_summary, "gate_rounds": len(rounds), "k_histogram": dict(sorted(k_histogram.items())), "gate_by_k": gate_by_k, @@ -543,6 +690,11 @@ def compare_prompts(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: row["adaptive"][name] = { "decode_tok_s": request.get("decode_tok_s"), "admitted_fraction": request.get("admitted_fraction"), + "activation_decision": request.get( + "activation_decision" + ), + "initial_confidence": request.get("initial_confidence"), + "calibrated_yield": request.get("calibrated_yield"), "mean_confidence_yield": request.get( "mean_confidence_yield" ), @@ -827,7 +979,7 @@ def build_report( ] activation_comparisons = compare_activation_shapes(cases) return { - "schema_version": 3, + "schema_version": 4, "cases": cases, "prompt_comparisons": compare_prompts(cases), "activation_comparisons": activation_comparisons, diff --git a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh index 608d291f5..cad371f53 100755 --- a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh @@ -19,7 +19,6 @@ OUT="${OUT:-$REPO/.harness-runs/qwen38-dspark-matrix-$(date -u +%Y%m%dT%H%M%SZ)} REPEATS="${REPEATS:-1}" WORKLOADS="${WORKLOADS:-humaneval,gsm8k,prose,north-star}" DECODE_MODES="${DECODE_MODES:-ar,speculation,adaptive}" -ADAPTIVE_DRAFT_ALWAYS="${ADAPTIVE_DRAFT_ALWAYS:-on}" CONFIDENCE_ABLATION="${CONFIDENCE_ABLATION:-1}" STEP_TIMING="${STEP_TIMING:-1}" CLIENTS="${CLIENTS:-1,2,3,4,6,8}" @@ -54,13 +53,14 @@ DRAFT_MODEL is a GGUF produced from that repository. The labeled R9700 adaptive-selection baseline uses the no-YaRN Q8_0 artifact; changing draft quantization changes both acceptance and cost, so it invalidates those stored oracle ratios. -The default fresh-process matrix runs ar, forced speculation, adaptive with -always-drafting on, and an always-drafting confidence-off ablation at live -concurrency 1,2,3,4,6,8 over HumanEval, GSM8K, and prose. The 2-code+4-chat -north-star row runs only at C=6. The summary fails if adaptive-on is below -0.995 of the paired ar/speculation oracle in mean or median goodput or TTFT. -The confidence ablation delta is reported but not gated. Set -ADAPTIVE_DRAFT_ALWAYS=on,off only for an additional diagnostic arm. +The default fresh-process matrix runs ar, forced speculation, one-shot +adaptive activation with confidence scored once for every request, and a +draft-free confidence-off AR ablation at live concurrency 1,2,3,4,6,8 over +HumanEval, GSM8K, and prose. An adaptive request keeps its selected +speculation/AR mode until retirement. The 2-code+4-chat north-star row runs +only at C=6. The summary fails if adaptive-on is below 0.995 of the paired +ar/speculation oracle in mean or median goodput or TTFT. The confidence +ablation delta is reported but not gated; there is no eager-draft matrix axis. The labeled adaptive-selection workload is fixed at C=6. Run it with WORKLOADS=adaptive-selection CLIENTS=6 to compare adaptive decisions against two strong speculation wins, two marginal wins, and two AR wins. @@ -88,7 +88,7 @@ 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 Qwen3.8 target GGUF" >&2; exit 2; } -[[ -r "$DRAFT_MODEL" ]] || { echo "set DRAFT_MODEL to the readable RadixArk q4-mix GGUF" >&2; exit 2; } +[[ -r "$DRAFT_MODEL" ]] || { echo "set DRAFT_MODEL to a readable RadixArk DSpark GGUF" >&2; exit 2; } [[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } for value_name in REPEATS SLOTS MAX_CTX MAX_CONCURRENT_PREFILLS MAX_TOKENS WARMUP_TOKENS PROFILE_CONTEXT HEALTH_TIMEOUT_SECONDS REQUEST_TIMEOUT_SECONDS; do value="${!value_name}" @@ -113,7 +113,6 @@ fi IFS=, read -r -a workload_list <<< "$WORKLOADS" IFS=, read -r -a mode_list <<< "$DECODE_MODES" -IFS=, read -r -a adaptive_axis <<< "$ADAPTIVE_DRAFT_ALWAYS" IFS=, read -r -a client_list <<< "$CLIENTS" reject_duplicates() { local list_name="$1" value @@ -129,7 +128,6 @@ reject_duplicates() { } reject_duplicates WORKLOADS "${workload_list[@]}" || exit 2 reject_duplicates DECODE_MODES "${mode_list[@]}" || exit 2 -reject_duplicates ADAPTIVE_DRAFT_ALWAYS "${adaptive_axis[@]}" || exit 2 reject_duplicates CLIENTS "${client_list[@]}" || exit 2 for workload in "${workload_list[@]}"; do @@ -144,12 +142,6 @@ for mode in "${mode_list[@]}"; do *) echo "unknown decode mode $mode" >&2; exit 2 ;; esac done -for axis in "${adaptive_axis[@]}"; do - [[ "$axis" == on || "$axis" == off ]] || { - echo "ADAPTIVE_DRAFT_ALWAYS entries must be on or off" >&2 - exit 2 - } -done for clients in "${client_list[@]}"; do [[ "$clients" =~ ^[1-9][0-9]*$ ]] || { echo "CLIENTS entries must be positive" >&2; exit 2; } (( clients <= SLOTS )) || { echo "CLIENTS=$clients exceeds SLOTS=$SLOTS" >&2; exit 2; } @@ -158,7 +150,7 @@ done variants=() for mode in "${mode_list[@]}"; do if [[ "$mode" == adaptive ]]; then - for axis in "${adaptive_axis[@]}"; do variants+=("adaptive-$axis"); done + variants+=("adaptive-on") if [[ "$CONFIDENCE_ABLATION" == 1 ]]; then variants+=("adaptive-confidence-off") fi @@ -229,14 +221,12 @@ case_applicable() { run_case() { local repeat="$1" workload="$2" clients="$3" variant="$4" - local decode_mode="$variant" draft_always="" confidence="" + local decode_mode="$variant" confidence="" if [[ "$variant" == adaptive-confidence-off ]]; then decode_mode=adaptive - draft_always=on confidence=off - elif [[ "$variant" == adaptive-* ]]; then + elif [[ "$variant" == adaptive-on ]]; then decode_mode=adaptive - draft_always="${variant#adaptive-}" confidence=on fi @@ -268,11 +258,6 @@ run_case() { "DFLASH_SPEC_GATE_LOG=1" "DFLASH_SPEC_PROFILE_CONTEXT=$PROFILE_CONTEXT" ) - if [[ "$draft_always" == on ]]; then - launch_env+=("DFLASH_SPEC_DRAFT_ALWAYS=1") - else - launch_env+=("DFLASH_SPEC_DRAFT_ALWAYS=0") - fi if [[ "$confidence" == on ]]; then launch_env+=("DFLASH_SPEC_CONFIDENCE=1") else @@ -300,8 +285,9 @@ run_case() { --cache-type-k "$CACHE_TYPE_K" --cache-type-v "$CACHE_TYPE_V" --fa-window "$FA_WINDOW" ) - [[ -n "$draft_always" ]] && metadata+=(--draft-always "$draft_always") - [[ -n "$confidence" ]] && metadata+=(--confidence "$confidence") + if [[ "$decode_mode" == adaptive ]]; then + metadata+=(--draft-always off --confidence "$confidence") + fi local item for item in "${launch_env[@]}"; do metadata+=(--launch-env "$item"); done "${metadata[@]}" diff --git a/harness/benchmarks/concurrency/summarize_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py index 0c520ee24..aba59d4d3 100755 --- a/harness/benchmarks/concurrency/summarize_feature_matrix.py +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -42,7 +42,6 @@ def load_reports(root: Path) -> list[dict]: "kvflash": ["kvflash"], "full": ["ddtree", "kvflash", "pflash"], "speculation": ["chain"], "adaptive-on": ["chain"], - "adaptive-off": ["chain"], "adaptive-confidence-off": ["chain"], } if variant not in expected_by_variant: @@ -53,7 +52,7 @@ def load_reports(root: Path) -> list[dict]: ) mode_by_variant = { "ar": "ar", "speculation": "speculation", - "adaptive-on": "adaptive", "adaptive-off": "adaptive", + "adaptive-on": "adaptive", "adaptive-confidence-off": "adaptive", } if variant in mode_by_variant: @@ -238,7 +237,6 @@ def summarize_qwen36(reports: list[dict]) -> str: "ar", "speculation", "adaptive-on", - "adaptive-off", "adaptive-confidence-off", } ORACLE_THRESHOLD = 0.995 @@ -296,20 +294,24 @@ def summarize_dspark(reports: list[dict]) -> str: if variant not in DSPARK_VARIANTS: raise ValueError(f"{item['path']}: mixed or unknown DSpark variant {variant!r}") config = meta.get("feature_config") or {} - expected_axes = { - "ar": (None, None), - "speculation": (None, None), - "adaptive-on": ("on", "on"), - "adaptive-off": ("off", "on"), - "adaptive-confidence-off": ("on", "off"), + expected_confidence = { + "ar": None, + "speculation": None, + "adaptive-on": "on", + "adaptive-confidence-off": "off", } - expected_draft, expected_confidence = expected_axes[variant] + draft_always = config.get("draft_always") + valid_draft_policy = ( + draft_always in (None, "off") + if variant.startswith("adaptive-") + else draft_always is None + ) if ( - config.get("draft_always") != expected_draft - or config.get("confidence") != expected_confidence + not valid_draft_policy + or config.get("confidence") != expected_confidence[variant] ): raise ValueError( - f"{item['path']}: adaptive axes do not match variant {variant}" + f"{item['path']}: adaptive metadata does not match variant {variant}" ) key = (str(meta["workload"]), int(level["clients"]), variant) grouped[key].append(item) @@ -350,8 +352,7 @@ def summarize_dspark(reports: list[dict]) -> str: "ar": 0, "speculation": 1, "adaptive-on": 2, - "adaptive-off": 3, - "adaptive-confidence-off": 4, + "adaptive-confidence-off": 3, } for workload, clients, variant in sorted( grouped, key=lambda key: (key[0], key[1], variant_order[key[2]]) @@ -491,10 +492,10 @@ def summarize_dspark(reports: list[dict]) -> str: "", "## Confidence ablation", "", - "Paired per-repeat deltas compare adaptive always-drafting with " - "confidence enabled against the same policy with confidence hidden. " - "Positive values favor fresh confidence; these deltas are reported " - "but are not acceptance-gated.", + "Paired per-repeat deltas compare one-shot adaptive activation " + "against the draft-free confidence-off AR fallback. Positive values " + "favor per-request activation; these deltas are reported but are not " + "acceptance-gated.", "", "| Workload | C | Goodput delta mean/median | " "Inverse-TTFT delta mean/median |", diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index 6109469ce..c12af11a7 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -5,6 +5,7 @@ import importlib.util import json +import math import os import subprocess import sys @@ -219,13 +220,17 @@ def test_dspark_runner_has_explicit_confidence_ablation(self) -> None: runner = (HERE / "run_qwen38_dspark_matrix.sh").read_text( encoding="utf-8", ) + self.assertIn('variants+=("adaptive-on")', runner) self.assertIn('variants+=("adaptive-confidence-off")', runner) - self.assertIn( - 'ADAPTIVE_DRAFT_ALWAYS="${ADAPTIVE_DRAFT_ALWAYS:-on}"', runner, - ) self.assertIn('"DFLASH_SPEC_CONFIDENCE=1"', runner) self.assertIn('"DFLASH_SPEC_CONFIDENCE=0"', runner) - self.assertIn('metadata+=(--confidence "$confidence")', runner) + self.assertIn( + 'metadata+=(--draft-always off --confidence "$confidence")', + runner, + ) + self.assertNotIn("ADAPTIVE_DRAFT_ALWAYS", runner) + self.assertNotIn("DFLASH_SPEC_DRAFT_ALWAYS", runner) + self.assertNotIn("RadixArk q4-mix GGUF", runner) def test_dspark_runner_profiles_only_the_measured_window(self) -> None: runner = (HERE / "run_qwen38_dspark_matrix.sh").read_text( @@ -266,6 +271,54 @@ def test_llama_only_does_not_require_lucebox_binary(self) -> None: class GateAnalysisTests(unittest.TestCase): + @staticmethod + def _write_activation_case( + root: Path, + variant: str, + activations: list[dict], + engine_ids: tuple[int, ...] = (7, 8), + ) -> Path: + case = root / "selection" / "c2" / "r1" / variant + case.mkdir(parents=True) + details = [ + { + "request_id": f"wire-{index}", + "prompt_index": index - 1, + "request_decode_tok_s": 8.0, + "content_sha256": f"hash-{index}", + } + for index in range(1, len(engine_ids) + 1) + ] + (case / "bench.json").write_text(json.dumps({ + "server_metadata": { + "workload": "selection", "variant": variant, + "clients": len(engine_ids), "repeat": 1, + }, + "levels": [{ + "aggregate_tok_s": 9.0, + "clients": len(engine_ids), + "requests_detail": details, + }], + }), encoding="utf-8") + metric_rows = [ + { + "request_id": f"wire-{index}", + "engine_request_id": engine_id, + "spec_accepted_tokens": 0, "spec_steps": 0, + "target_forwards": 1, "output_tokens": 1, + } + for index, engine_id in enumerate(engine_ids, 1) + ] + lines = [ + f"[spec-activation] {json.dumps(row)}" for row in activations + ] + [ + f"[concurrency-metrics] {json.dumps(row)}" for row in metric_rows + ] + (case / "benchmark-server.log").write_text( + "\n".join(lines) + "\n", encoding="utf-8", + ) + return case + def test_measured_step_timing_is_joined_and_summarized(self) -> None: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) @@ -315,11 +368,17 @@ def test_measured_step_timing_is_joined_and_summarized(self) -> None: "spec_accepted_tokens": 0, "spec_steps": 0, "target_forwards": 1, "output_tokens": 1, } + activation = { + "request_id": 7, "slot": 0, + "initial_confidence": 3.25, "calibrated_yield": 1.75, + "decision": "ar", + } (case / "benchmark-server.log").write_text( "[spec-gate] C=1 k=0 scores=[7:1.000/confidence] " "sources=confidence:1,unavailable:0 " "G(k)=0.010000 G(0)=0.020000 " "predicted_cost=50.0us measured=ar-path\n" + f"[spec-activation] {json.dumps(activation)}\n" f"[step-timing] {json.dumps(timing)}\n" f"[concurrency-metrics] {json.dumps(metric_row)}\n", encoding="utf-8", @@ -337,6 +396,80 @@ def test_measured_step_timing_is_joined_and_summarized(self) -> None: report["gate_by_k"][0]["realized_goodput_tok_s"], 10000.0, ) self.assertEqual(report["gate_timing_count_mismatch"], 0) + self.assertEqual(report["activation"]["validation"], "passed") + self.assertEqual(report["activation"]["records"], 1) + request = report["requests"][0] + self.assertEqual(request["activation_slot"], 0) + self.assertEqual(request["initial_confidence"], 3.25) + self.assertEqual(request["calibrated_yield"], 1.75) + self.assertEqual(request["activation_decision"], "ar") + + def test_adaptive_on_activation_proof_fails_closed(self) -> None: + activation_7 = { + "request_id": 7, "slot": 0, + "initial_confidence": 2.0, "calibrated_yield": 1.25, + "decision": "speculation", + } + activation_8 = { + "request_id": 8, "slot": 1, + "initial_confidence": 1.0, "calibrated_yield": 1.0, + "decision": "ar", + } + activation_9 = {**activation_8, "request_id": 9} + cases = ( + ("missing", [activation_7], "missing activations.*8"), + ( + "duplicate", [activation_7, activation_7, activation_8], + "duplicate activations.*7", + ), + ( + "unknown", [activation_7, activation_8, activation_9], + "unknown engine requests.*9", + ), + ) + for label, activations, message in cases: + with self.subTest(label=label), tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", activations, + ) + with self.assertRaisesRegex(ValueError, message): + gate_analysis.analyze_case(case) + + def test_confidence_off_is_exempt_from_activation_coverage(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-confidence-off", [], + ) + report = gate_analysis.analyze_case(case) + self.assertFalse(report["activation"]["required"]) + self.assertEqual(report["activation"]["validation"], "not-required") + self.assertEqual(report["activation"]["records"], 0) + self.assertTrue(all( + row["activation_decision"] is None + for row in report["requests"] + )) + + def test_activation_record_fields_are_strictly_validated(self) -> None: + valid = { + "request_id": 7, "slot": 0, + "initial_confidence": 1.0, "calibrated_yield": 1.0, + "decision": "ar", + } + cases = ( + ({**valid, "request_id": True}, "request_id"), + ({**valid, "slot": -1}, "slot"), + ({**valid, "initial_confidence": 0.99}, "initial_confidence"), + ({**valid, "calibrated_yield": math.nan}, "calibrated_yield"), + ({**valid, "decision": "undecided"}, "decision"), + ) + for row, message in cases: + with self.subTest(field=message), tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(row)}\n", encoding="utf-8", + ) + with self.assertRaisesRegex(ValueError, message): + gate_analysis.parse_server_log(path) def test_paired_controls_produce_a_per_prompt_concurrent_oracle(self) -> None: base_request = { @@ -784,12 +917,11 @@ def dspark_item( mode = ( "adaptive" if variant.startswith("adaptive-") else variant ) - if variant == "adaptive-confidence-off": - draft_always = "on" - confidence = "off" - elif variant.startswith("adaptive-"): - draft_always = variant.removeprefix("adaptive-") - confidence = "on" + if variant.startswith("adaptive-"): + draft_always = "off" + confidence = ( + "off" if variant == "adaptive-confidence-off" else "on" + ) else: draft_always = None confidence = None @@ -933,10 +1065,6 @@ def test_dspark_summary_enforces_mean_median_oracle_gate(self) -> None: self.dspark_item("speculation", 103.0, 1.0, repeat=2), self.dspark_item("adaptive-on", 100.0, 0.9, repeat=1), self.dspark_item("adaptive-on", 103.0, 1.0, repeat=2), - # Admitted-only drafting has no confidence bootstrap now and is - # diagnostic rather than acceptance-gated. - self.dspark_item("adaptive-off", 70.0, 2.0, repeat=1), - self.dspark_item("adaptive-off", 72.0, 1.9, repeat=2), self.dspark_item("adaptive-confidence-off", 75.0, 1.8, repeat=1), self.dspark_item("adaptive-confidence-off", 80.0, 1.7, repeat=2), ] @@ -944,10 +1072,21 @@ def test_dspark_summary_enforces_mean_median_oracle_gate(self) -> None: self.assertIn("Qwen3.8 DSpark adaptive concurrency matrix", text) self.assertIn("Oracle-relative gate", text) self.assertIn("| adaptive-on |", text) - self.assertIn("| adaptive-off |", text) self.assertIn("| adaptive-confidence-off |", text) self.assertIn("## Confidence ablation", text) - self.assertIn("Positive values favor fresh confidence", text) + self.assertIn("Positive values favor per-request activation", text) + + def test_dspark_summary_rejects_eager_draft_metadata(self) -> None: + adaptive = self.dspark_item("adaptive-on", 100.0, 0.9) + adaptive["meta"]["feature_config"]["draft_always"] = "on" + with self.assertRaisesRegex( + ValueError, "adaptive metadata does not match variant adaptive-on", + ): + summary.summarize([ + self.dspark_item("ar", 100.0, 1.0), + self.dspark_item("speculation", 98.0, 0.9), + adaptive, + ]) def test_dspark_summary_rejects_oracle_regression(self) -> None: reports = [ diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h index e2ac6b6df..a4c782d49 100644 --- a/server/src/common/concurrency/speculation_gate.h +++ b/server/src/common/concurrency/speculation_gate.h @@ -17,6 +17,17 @@ namespace dflash::common { +struct SpecGateConfig { + // Offline calibration for the first request cohort, before any admitted + // request has produced online yield evidence. Generic callers default to + // neutral; a concrete speculator may install its fitted prior. + double initial_yield_scale = 1.0; + double yield_ema_alpha = 0.20; + double calibration_ema_alpha = 0.20; + double cost_ema_alpha = 0.20; + double adaptive_gain_margin = 0.02; +}; + struct SpecCostLookup { double cost = std::numeric_limits::infinity(); int requested_index = 0; @@ -75,12 +86,19 @@ struct SpecCandidate { uint64_t request_id = 0; int slot = -1; SpeculationPolicy policy = SpeculationPolicy::Adaptive; - bool eligible = false; - // NaN means no current-block confidence is available; an adaptive - // candidate then remains on AR. Otherwise this is the survival-product - // expected yield, including the root. Speculator adapters provide - // per-position probability-like scores to confidence_survival_yield(); - // the gate owns calibration and clamping. + // scoreable is a request-lifetime capability: the engine can produce the + // mandatory one-time activation score from the committed feature mirror. + bool scoreable = false; + // can_speculate is also request-lifetime (for example, false for an + // unsupported sampler or thinking hook). Temporary blockers such as a + // min-token floor or a mixed-prefill round must not clear this flag; the + // engine executes AR for that round without changing the sticky decision. + bool can_speculate = false; + // NaN means the adapter did not supply a score in this plan. An undecided + // adaptive request with no retained score is returned for its one-time + // bootstrap. The first finite score is retained for the lifetime of the + // request and later scores are ignored. This is the survival-product + // expected yield, including the root; the gate owns calibration/clamping. double confidence_yield = std::numeric_limits::quiet_NaN(); }; @@ -106,24 +124,46 @@ struct SpecStepGeometry { enum class SpecScoreSource : uint8_t { Confidence, + InitialConfidence, Unavailable, }; inline const char * spec_score_source_name(SpecScoreSource source) { switch (source) { case SpecScoreSource::Confidence: return "confidence"; + case SpecScoreSource::InitialConfidence: return "initial"; case SpecScoreSource::Unavailable: return "unavailable"; } return "unknown"; } +enum class SpecDecision : uint8_t { + Undecided, + AR, + Speculation, +}; + +inline const char * spec_decision_name(SpecDecision decision) { + switch (decision) { + case SpecDecision::Undecided: return "undecided"; + case SpecDecision::AR: return "ar"; + case SpecDecision::Speculation: return "speculation"; + } + return "unknown"; +} + struct SpecPlanScore { uint64_t request_id = 0; int slot = -1; double expected_yield = 1.0; SpecScoreSource source = SpecScoreSource::Unavailable; + SpecDecision decision = SpecDecision::Undecided; bool forced = false; bool admitted = false; + // True only in the plan that commits this request's immutable adaptive + // AR/speculation decision. This supports one activation record per + // request without treating later sticky execution as a new decision. + bool newly_decided = false; }; struct SpecPlan { @@ -135,17 +175,30 @@ struct SpecPlan { int step_rows = 0; int draft_lanes = 0; double expected_tokens = 0.0; + // Startup-profiled cost before online correction, and the shape-local + // correction applied to it. predicted_cost is their product. + double profiled_cost = 0.0; + double cost_scale = 1.0; double predicted_cost = 0.0; - // Pre-scale confidence yield for the admitted confidence-scored lanes. - // This is the denominator contribution shown in per-step telemetry. + // Final calibrated expected yield for admitted confidence-scored lanes. + // This is directly comparable with realized emitted tokens in telemetry. double calibration_predicted_tokens = 0.0; double goodput = 0.0; double ar_goodput = 0.0; int unavailable_count = 0; + // False means at least one adaptive request still needs its one-time + // confidence bootstrap. No new adaptive decisions are committed in that + // plan. The immediate post-bootstrap replan commits all undecided lanes. + bool decisions_committed = false; bool cost_lookup_clamped = false; std::vector ordered; std::vector admitted_request_ids; std::vector admitted_slots; + // Cold requests are unavailable only until their one-time initialization. + // The engine drafts these together and immediately replans before choosing + // the request's sticky AR/speculation mode. + std::vector bootstrap_request_ids; + std::vector bootstrap_slots; }; // Generic confidence contract for every chain speculator: `confidences[i]` @@ -171,6 +224,59 @@ inline double confidence_survival_yield( class SpeculationGate { private: + struct RequestState { + double initial_confidence = + std::numeric_limits::quiet_NaN(); + SpecDecision decision = SpecDecision::Undecided; + double yield_ema = std::numeric_limits::quiet_NaN(); + uint64_t yield_observations = 0; + double calibration_ratio_ema = 1.0; + uint64_t calibration_observations = 0; + }; + + struct PendingConfidence { + double confidence_yield = + std::numeric_limits::quiet_NaN(); + }; + + struct CostShape { + int concurrency = 0; + int admitted_count = 0; + int tree_rows = 0; + int step_rows = 0; + int draft_lanes = 0; + + bool operator==(const CostShape & other) const { + return concurrency == other.concurrency && + admitted_count == other.admitted_count && + tree_rows == other.tree_rows && + step_rows == other.step_rows && + draft_lanes == other.draft_lanes; + } + }; + + struct CostShapeHash { + size_t operator()(const CostShape & shape) const { + size_t seed = 0; + auto mix = [&](int value) { + seed ^= std::hash{}(value) + + static_cast(0x9e3779b9U) + + (seed << 6) + (seed >> 2); + }; + mix(shape.concurrency); + mix(shape.admitted_count); + mix(shape.tree_rows); + mix(shape.step_rows); + mix(shape.draft_lanes); + return seed; + } + }; + + struct CostState { + double scale = 1.0; + uint64_t observations = 0; + }; + struct CandidateScore { double expected_yield = 1.0; double uncalibrated_confidence = @@ -183,13 +289,33 @@ class SpeculationGate { const char * table, int requested, int profiled)>; SpeculationGate(SpecCostTables costs, SpecStepGeometry geometry, - int max_accept, ClampLogger clamp_logger = {}) - : costs_(std::move(costs)), geometry_(std::move(geometry)), + int max_accept, ClampLogger clamp_logger = {}, + SpecGateConfig config = {}) + : config_(config), costs_(std::move(costs)), + geometry_(std::move(geometry)), max_accept_(std::max(1, max_accept)), clamp_logger_(std::move(clamp_logger)) {} + SpeculationGate(SpecGateConfig config, SpecCostTables costs, + SpecStepGeometry geometry, int max_accept, + ClampLogger clamp_logger = {}) + : SpeculationGate(std::move(costs), std::move(geometry), max_accept, + std::move(clamp_logger), config) {} + bool valid() const { - return costs_.valid() && geometry_.tree_width >= 1 && max_accept_ >= 1; + auto valid_alpha = [](double value) { + return std::isfinite(value) && value > 0.0 && value <= 1.0; + }; + return valid_alpha(config_.yield_ema_alpha) && + valid_alpha(config_.calibration_ema_alpha) && + valid_alpha(config_.cost_ema_alpha) && + std::isfinite(config_.initial_yield_scale) && + config_.initial_yield_scale >= kCalibrationScaleMin && + config_.initial_yield_scale <= kCalibrationScaleMax && + std::isfinite(config_.adaptive_gain_margin) && + config_.adaptive_gain_margin >= 0.0 && + costs_.valid() && geometry_.tree_width >= 1 && + max_accept_ >= 1; } // draft_lanes_override prices always-drafting. -1 means admitted-only. @@ -218,34 +344,92 @@ class SpeculationGate { double uncalibrated_confidence = std::numeric_limits::quiet_NaN(); SpecScoreSource source = SpecScoreSource::Unavailable; + SpecDecision decision = SpecDecision::Undecided; bool forced = false; + bool commit_candidate = false; }; std::vector forced; - std::vector adaptive; + std::vector undecided; + std::vector forced_ar; forced.reserve(candidates.size()); - adaptive.reserve(candidates.size()); + undecided.reserve(candidates.size()); + forced_ar.reserve(candidates.size()); for (const SpecCandidate & candidate : candidates) { - if (!candidate.eligible || - candidate.policy == SpeculationPolicy::Never) { + if (candidate.policy == SpeculationPolicy::Never) { + continue; + } + const SpecDecision prior_decision = decision(candidate.request_id); + if (candidate.policy == SpeculationPolicy::Adaptive && + prior_decision == SpecDecision::AR) { + // A one-shot AR decision removes the request from all future + // adaptive rankings until forget(). continue; } const CandidateScore score = score_candidate(candidate); - if (score.source == SpecScoreSource::Unavailable && - candidate.policy != SpeculationPolicy::Always) { + const bool adaptive_undecided = + candidate.policy == SpeculationPolicy::Adaptive && + prior_decision == SpecDecision::Undecided; + if (adaptive_undecided && + score.source == SpecScoreSource::Unavailable) { ++out.unavailable_count; + if (candidate.scoreable) { + out.bootstrap_request_ids.push_back(candidate.request_id); + out.bootstrap_slots.push_back(candidate.slot); + } else { + // Never silently turn a missing mandatory score into a + // decision. A deployment that cannot score an adaptive + // DSpark request must handle that contract failure + // explicitly rather than inventing an AR activation. + out.valid = false; + out.error = "adaptive request is not scoreable"; + return out; + } + continue; + } + if (adaptive_undecided && !candidate.can_speculate) { + // Still retain its mandatory initial score for telemetry, but + // permanently unsupported execution commits directly to AR. + forced_ar.push_back({ + &candidate, score.expected_yield, + score.uncalibrated_confidence, score.source, + prior_decision, false, true}); + continue; + } + if (candidate.policy == SpeculationPolicy::Always && + prior_decision != SpecDecision::Speculation && + !candidate.can_speculate) { continue; } Ranked ranked{&candidate, score.expected_yield, score.uncalibrated_confidence, score.source, - candidate.policy == SpeculationPolicy::Always}; - (ranked.forced ? forced : adaptive).push_back(ranked); + prior_decision, + candidate.policy == SpeculationPolicy::Always || + prior_decision == SpecDecision::Speculation, + adaptive_undecided}; + (ranked.forced ? forced : undecided).push_back(ranked); } + + // Selection is atomic across the newly arriving batch. If any cold + // adaptive request still needs its score, defer every undecided lane + // until the engine publishes the batched bootstrap and immediately + // replans. Previously decided speculation and explicit Always lanes + // can still execute in the bootstrap plan. + if (!out.bootstrap_slots.empty()) undecided.clear(); + out.decisions_committed = out.bootstrap_slots.empty(); + if (out.decisions_committed) { + for (const Ranked & item : forced_ar) { + request_states_[item.candidate->request_id].decision = + SpecDecision::AR; + } + } + auto request_order = [](const Ranked & a, const Ranked & b) { return a.candidate->request_id < b.candidate->request_id; }; std::sort(forced.begin(), forced.end(), request_order); - std::sort(adaptive.begin(), adaptive.end(), + std::sort(forced_ar.begin(), forced_ar.end(), request_order); + std::sort(undecided.begin(), undecided.end(), [](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; @@ -258,14 +442,16 @@ class SpeculationGate { } std::vector ranked; - ranked.reserve(forced.size() + adaptive.size()); + ranked.reserve(forced.size() + undecided.size()); ranked.insert(ranked.end(), forced.begin(), forced.end()); - ranked.insert(ranked.end(), adaptive.begin(), adaptive.end()); + ranked.insert(ranked.end(), undecided.begin(), undecided.end()); for (const Ranked & item : ranked) { out.ordered.push_back({item.candidate->request_id, item.candidate->slot, item.score, item.source, - item.forced, false}); + item.decision, + item.forced, false, + item.commit_candidate}); } const int forced_count = static_cast(forced.size()); @@ -276,122 +462,286 @@ class SpeculationGate { const SpecCostLookup ar_lookup = costs_.step_cost.lookup( geometry_.bucketed_lanes(concurrency)); report_clamp("step", ar_lookup, out); + const CostShape ar_shape{ + concurrency, 0, 0, geometry_.bucketed_lanes(concurrency), 0}; + const double ar_scale = cost_scale(ar_shape); out.ar_goodput = concurrency == 0 ? 0.0 - : static_cast(concurrency) / ar_lookup.cost; - - int best_k = forced_count; - double best_goodput = -1.0; - double best_cost = 0.0; - double best_expected = 0.0; - int best_tree_rows = 0; - int best_step_rows = geometry_.bucketed_lanes(concurrency); - int best_draft_lanes = 0; + : static_cast(concurrency) / + (ar_lookup.cost * ar_scale); + + struct PlanPoint { + int k = 0; + double goodput = -1.0; + double profiled_cost = 0.0; + double cost_scale = 1.0; + double expected_tokens = 0.0; + int tree_rows = 0; + int step_rows = 0; + int draft_lanes = 0; + }; + PlanPoint baseline; + PlanPoint best; for (int k = forced_count; k <= max_k; ++k) { if (k > forced_count) expected_sum += ranked[k - 1].score; const double expected_tokens = static_cast(concurrency - k) + expected_sum; - double cost = 0.0; + double profiled_cost = 0.0; int tree_rows = 0; int step_rows = geometry_.bucketed_lanes(concurrency); const int draft_lanes = draft_lanes_override >= 0 ? draft_lanes_override : k; if (k == 0 && draft_lanes == 0) { - cost = ar_lookup.cost; + profiled_cost = ar_lookup.cost; } else { if (k > 0) { tree_rows = geometry_.tree_rows(k); const SpecCostLookup tree = costs_.tree_cost.lookup(tree_rows); report_clamp("tree", tree, out); - cost += tree.cost; + profiled_cost += tree.cost; step_rows = geometry_.step_rows( concurrency, k, expected_sum); } const SpecCostLookup step = costs_.step_cost.lookup(step_rows); report_clamp("step", step, out); - cost += step.cost; + profiled_cost += step.cost; if (draft_lanes > 0) { const SpecCostLookup draft = costs_.draft_cost.lookup(draft_lanes); report_clamp("draft", draft, out); - cost += draft.cost; + profiled_cost += draft.cost; } } - const double goodput = expected_tokens / cost; - if (goodput > best_goodput) { - best_goodput = goodput; - best_k = k; - best_cost = cost; - best_expected = expected_tokens; - best_tree_rows = tree_rows; - best_step_rows = step_rows; - best_draft_lanes = draft_lanes; + const CostShape shape{ + concurrency, k, tree_rows, step_rows, draft_lanes}; + const double scale = cost_scale(shape); + const double goodput = + expected_tokens / (profiled_cost * scale); + const PlanPoint point{ + k, goodput, profiled_cost, scale, expected_tokens, + tree_rows, step_rows, draft_lanes}; + if (k == forced_count) baseline = point; + if (goodput > best.goodput) { + best = point; } } - out.admitted_count = best_k; - out.goodput = std::max(0.0, best_goodput); - out.predicted_cost = best_cost; - out.expected_tokens = best_expected; - out.tree_rows = best_tree_rows; - out.step_rows = best_step_rows; - out.draft_lanes = best_draft_lanes; - for (int i = 0; i < best_k; ++i) { + // Explicit Always and sticky-Speculation lanes establish the + // non-negotiable baseline. The safety margin applies only to the + // one-shot admission of additional undecided adaptive lanes. + if (best.k > forced_count && + best.goodput < baseline.goodput * + (1.0 + config_.adaptive_gain_margin)) { + best = baseline; + } + + out.admitted_count = best.k; + out.goodput = std::max(0.0, best.goodput); + out.profiled_cost = best.profiled_cost; + out.cost_scale = best.cost_scale; + out.predicted_cost = best.profiled_cost * best.cost_scale; + out.expected_tokens = best.expected_tokens; + out.tree_rows = best.tree_rows; + out.step_rows = best.step_rows; + out.draft_lanes = best.draft_lanes; + for (size_t i = 0; i < ranked.size(); ++i) { + if (!ranked[i].commit_candidate) continue; + const SpecDecision committed = static_cast(i) < best.k + ? SpecDecision::Speculation : SpecDecision::AR; + request_states_[ranked[i].candidate->request_id].decision = + committed; + out.ordered[i].decision = committed; + } + for (int i = 0; i < best.k; ++i) { out.ordered[(size_t)i].admitted = true; out.admitted_request_ids.push_back(ranked[(size_t)i].candidate->request_id); out.admitted_slots.push_back(ranked[(size_t)i].candidate->slot); - if (ranked[(size_t)i].source == SpecScoreSource::Confidence) { + if (std::isfinite( + ranked[(size_t)i].uncalibrated_confidence)) { out.calibration_predicted_tokens += - ranked[(size_t)i].uncalibrated_confidence; + ranked[(size_t)i].score; pending_confidence_[ranked[(size_t)i].candidate->request_id] = - ranked[(size_t)i].uncalibrated_confidence; + {ranked[(size_t)i].uncalibrated_confidence}; + } + } + if (out.decisions_committed) { + for (const Ranked & item : forced_ar) { + out.ordered.push_back({ + item.candidate->request_id, + item.candidate->slot, + item.score, + item.source, + SpecDecision::AR, + false, + false, + true, + }); } } return out; } - void observe(uint64_t request_id, double emitted_tokens) { + // `confidence_yield` may carry the freshly drafted current-block score. + // It calibrates this observation (including forced/probe executions) but + // never overwrites the one-shot activation score retained for the request. + // `generated_tokens` is kept for source compatibility with adapters that + // already pass it; one-shot decisions have no age or refresh semantics. + void observe(uint64_t request_id, double emitted_tokens, + int /* generated_tokens */, + double confidence_yield = + std::numeric_limits::quiet_NaN()) { auto pending = pending_confidence_.find(request_id); - if (pending == pending_confidence_.end()) return; if (!std::isfinite(emitted_tokens) || emitted_tokens < 1.0 || emitted_tokens > static_cast(max_accept_)) { + if (pending != pending_confidence_.end()) + pending_confidence_.erase(pending); + return; + } + + RequestState & state = request_states_[request_id]; + update_ema(state.yield_ema, state.yield_observations, + emitted_tokens, config_.yield_ema_alpha); + + double raw = std::numeric_limits::quiet_NaN(); + if (std::isfinite(confidence_yield)) { + raw = std::clamp(confidence_yield, 1.0, + static_cast(max_accept_)); + if (!std::isfinite(state.initial_confidence)) + state.initial_confidence = raw; + } else if (pending != pending_confidence_.end()) { + raw = pending->second.confidence_yield; + } + if (std::isfinite(raw)) { + const double ratio = std::clamp( + emitted_tokens / raw, + kCalibrationScaleMin, kCalibrationScaleMax); + update_ema(state.calibration_ratio_ema, + state.calibration_observations, ratio, + config_.calibration_ema_alpha); + update_ema(global_calibration_ratio_ema_, + global_calibration_observations_, ratio, + config_.calibration_ema_alpha); + } + if (pending != pending_confidence_.end()) pending_confidence_.erase(pending); + } + + // Compatibility for adapters that do not publish a current-block score. + void observe(uint64_t request_id, double emitted_tokens) { + observe(request_id, emitted_tokens, 0); + } + + void observe_cost(const SpecPlan & plan, double measured_us) { + if (!plan.valid || !std::isfinite(measured_us) || measured_us <= 0.0 || + !std::isfinite(plan.profiled_cost) || plan.profiled_cost <= 0.0) { return; } - calibration_predicted_ += pending->second; - calibration_realized_ += emitted_tokens; - ++calibration_observations_; - pending_confidence_.erase(pending); + const CostShape shape{ + plan.concurrency, plan.admitted_count, plan.tree_rows, + plan.step_rows, plan.draft_lanes}; + const double ratio = std::clamp( + measured_us / plan.profiled_cost, + kCostScaleMin, kCostScaleMax); + CostState & state = cost_states_[shape]; + update_ema(state.scale, state.observations, ratio, + config_.cost_ema_alpha); } - void forget(uint64_t request_id) { pending_confidence_.erase(request_id); } + void forget(uint64_t request_id) { + request_states_.erase(request_id); + pending_confidence_.erase(request_id); + } + + bool has_state(uint64_t request_id) const { + return request_states_.find(request_id) != request_states_.end(); + } + bool has_confidence(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state != request_states_.end() && + std::isfinite(state->second.initial_confidence); + } + double initial_confidence(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state == request_states_.end() + ? std::numeric_limits::quiet_NaN() + : state->second.initial_confidence; + } + SpecDecision decision(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state == request_states_.end() + ? SpecDecision::Undecided : state->second.decision; + } + double yield_ema(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state == request_states_.end() + ? std::numeric_limits::quiet_NaN() + : state->second.yield_ema; + } + + double calibration_scale(uint64_t request_id) const { + auto state = request_states_.find(request_id); + if (state != request_states_.end() && + state->second.calibration_observations > 0) { + double scale = state->second.calibration_ratio_ema; + // The emitted-yield EWMA is a second request-local calibration + // view. Blend it as a ratio against the initial raw score so it + // refines that score instead of becoming a score by itself. + if (state->second.yield_observations > 0 && + std::isfinite(state->second.yield_ema) && + std::isfinite(state->second.initial_confidence)) { + const double yield_scale = std::clamp( + state->second.yield_ema / + state->second.initial_confidence, + kCalibrationScaleMin, kCalibrationScaleMax); + scale = 0.5 * (scale + yield_scale); + } + return std::clamp( + scale, kCalibrationScaleMin, kCalibrationScaleMax); + } + return calibration_scale(); + } double calibration_scale() const { - // Calibration activation and bounds are fixed protocol constants, - // not workload policy tunables. - if (calibration_observations_ < kCalibrationMinObservations || - calibration_predicted_ <= 0.0) - return 1.0; - return std::clamp( - calibration_realized_ / calibration_predicted_, - kCalibrationScaleMin, kCalibrationScaleMax); + return global_calibration_observations_ == 0 + ? config_.initial_yield_scale : global_calibration_ratio_ema_; } uint64_t calibration_observations() const { - return calibration_observations_; + return global_calibration_observations_; } const SpecCostTables & costs() const { return costs_; } private: - CandidateScore score_candidate(const SpecCandidate & candidate) const { + CandidateScore score_candidate(const SpecCandidate & candidate) { + bool accepted_initial_score = false; if (std::isfinite(candidate.confidence_yield)) { const double raw = std::clamp( candidate.confidence_yield, 1.0, static_cast(max_accept_)); + RequestState & state = request_states_[candidate.request_id]; + if (!std::isfinite(state.initial_confidence)) { + state.initial_confidence = raw; + accepted_initial_score = true; + } + return { + std::clamp( + calibration_scale(candidate.request_id) * + state.initial_confidence, + 1.0, static_cast(max_accept_)), + state.initial_confidence, + accepted_initial_score ? SpecScoreSource::Confidence + : SpecScoreSource::InitialConfidence, + }; + } + auto state = request_states_.find(candidate.request_id); + if (state != request_states_.end() && + std::isfinite(state->second.initial_confidence)) { return { - std::clamp(calibration_scale() * raw, 1.0, - static_cast(max_accept_)), - raw, - SpecScoreSource::Confidence, + std::clamp( + calibration_scale(candidate.request_id) * + state->second.initial_confidence, + 1.0, static_cast(max_accept_)), + state->second.initial_confidence, + SpecScoreSource::InitialConfidence, }; } return { @@ -401,6 +751,19 @@ class SpeculationGate { }; } + static void update_ema(double & value, uint64_t & observations, + double sample, double alpha) { + value = observations == 0 + ? sample : (1.0 - alpha) * value + alpha * sample; + ++observations; + } + + double cost_scale(const CostShape & shape) const { + auto state = cost_states_.find(shape); + return state == cost_states_.end() || state->second.observations == 0 + ? 1.0 : state->second.scale; + } + void report_clamp(const char * name, const SpecCostLookup & lookup, SpecPlan & plan) const { if (!lookup.clamped) return; @@ -409,16 +772,19 @@ class SpeculationGate { clamp_logger_(name, lookup.requested_index, lookup.profiled_index); } + SpecGateConfig config_; SpecCostTables costs_; SpecStepGeometry geometry_; int max_accept_ = 1; - static constexpr uint64_t kCalibrationMinObservations = 32; static constexpr double kCalibrationScaleMin = 0.25; static constexpr double kCalibrationScaleMax = 4.0; - double calibration_predicted_ = 0.0; - double calibration_realized_ = 0.0; - uint64_t calibration_observations_ = 0; - std::unordered_map pending_confidence_; + static constexpr double kCostScaleMin = 0.25; + static constexpr double kCostScaleMax = 4.0; + double global_calibration_ratio_ema_ = 1.0; + uint64_t global_calibration_observations_ = 0; + std::unordered_map request_states_; + std::unordered_map pending_confidence_; + std::unordered_map cost_states_; ClampLogger clamp_logger_; }; diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 12445e6b2..422043ade 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -335,6 +335,14 @@ struct ModelBackend { // and stays valid until shutdown(). virtual SeqEngine * seq_engine() { return nullptr; } + // Request-level decode_mode is consumed by the concurrent scheduler. + // Report only capabilities whose resources were actually initialized; + // the scheduler rejects unsupported non-AR modes before claiming a slot. + virtual ConcurrentDecodeCapabilities concurrent_decode_capabilities() + const { + return {}; + } + // ── Snapshots ──────────────────────────────────────────────────── // With right-sized CPU-resident snapshots, each slot costs only // ~(cur_pos × 5 KB) of system RAM, so we can afford many slots. diff --git a/server/src/common/speculation_policy.h b/server/src/common/speculation_policy.h index 8758d8ced..aa6b25c1d 100644 --- a/server/src/common/speculation_policy.h +++ b/server/src/common/speculation_policy.h @@ -13,6 +13,23 @@ enum class SpeculationPolicy { Never, }; +// Runtime capabilities for the concurrent decode path. Forced speculation +// only needs an executable draft/verify chain; adaptive additionally needs +// the activation gate and its cost profile. AR is always supported. +struct ConcurrentDecodeCapabilities { + bool forced_speculation = false; + bool adaptive = false; + + constexpr bool supports(SpeculationPolicy policy) const { + switch (policy) { + case SpeculationPolicy::Always: return forced_speculation; + case SpeculationPolicy::Adaptive: return adaptive; + case SpeculationPolicy::Never: return true; + } + return false; + } +}; + inline const char * speculation_policy_name(SpeculationPolicy policy) { switch (policy) { case SpeculationPolicy::Adaptive: return "adaptive"; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index c879a4eb2..f606878fc 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -45,7 +45,7 @@ double confidence_realized_tokens( double realized = 0.0; for (const SpecPlanScore & score : plan.ordered) { if (!score.admitted || - score.source != SpecScoreSource::Confidence) { + score.source == SpecScoreSource::Unavailable) { continue; } const auto output = std::find_if( @@ -65,10 +65,12 @@ void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, double calibration_realized_tokens, double measured_us) { int confidence = 0; + int initial = 0; int unavailable = plan.unavailable_count; for (const SpecPlanScore & score : plan.ordered) { switch (score.source) { case SpecScoreSource::Confidence: ++confidence; break; + case SpecScoreSource::InitialConfidence: ++initial; break; case SpecScoreSource::Unavailable: ++unavailable; break; } } @@ -84,10 +86,18 @@ void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, spec_score_source_name(score.source), score.admitted ? "*" : ""); } + std::fprintf(stderr, "] decisions=["); + for (size_t i = 0; i < plan.ordered.size(); ++i) { + const SpecPlanScore & score = plan.ordered[i]; + std::fprintf(stderr, "%s%llu:%s", + i == 0 ? "" : ",", + (unsigned long long)score.request_id, + spec_decision_name(score.decision)); + } std::fprintf(stderr, - "] sources=confidence:%d,unavailable:%d " + "] sources=confidence:%d,initial:%d,unavailable:%d " "calibration=%.3f rounds=%llu calib_tokens=%.3f/", - confidence, unavailable, calibration_scale, + confidence, initial, unavailable, calibration_scale, (unsigned long long)calibration_rounds, plan.calibration_predicted_tokens); if (std::isfinite(calibration_realized_tokens)) { @@ -96,7 +106,9 @@ void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, std::fprintf(stderr, "n/a"); } std::fprintf(stderr, + " profiled_cost=%.1fus cost_scale=%.3f" " G(k)=%.6f G(0)=%.6f predicted_cost=%.1fus", + plan.profiled_cost, plan.cost_scale, plan.goodput, plan.ar_goodput, plan.predicted_cost); if (std::isfinite(measured_us)) { std::fprintf(stderr, " measured=%.1fus\n", measured_us); @@ -105,6 +117,21 @@ void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, } } +void log_spec_activations(const SpecPlan & plan, + const SpeculationGate & gate) { + for (const SpecPlanScore & score : plan.ordered) { + if (!score.newly_decided) continue; + const double initial = gate.initial_confidence(score.request_id); + std::fprintf(stderr, + "[spec-activation] {\"request_id\":%llu,\"slot\":%d," + "\"initial_confidence\":%.6f,\"calibrated_yield\":%.6f," + "\"decision\":\"%s\"}\n", + (unsigned long long)score.request_id, score.slot, + initial, score.expected_yield, + spec_decision_name(score.decision)); + } +} + } // namespace Qwen35SeqEngine::Qwen35SeqEngine( @@ -130,7 +157,6 @@ Qwen35SeqEngine::Qwen35SeqEngine( prepared_chain_drafts_.resize((size_t)n_slots); last_survival_score_.assign( (size_t)n_slots, std::numeric_limits::quiet_NaN()); - last_survival_generated_.assign((size_t)n_slots, -1); // The concurrent DDTree stack is gated to a local same-device drafter. // Build metadata-only BF16 views over each slot's disjoint target feature @@ -209,7 +235,8 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { tree_width_ <= 1 || tree_width_ > 16 || slots_.residency_active()) { std::fprintf(stderr, "[spec-profile] disabled: chain/features unavailable or " - "concurrent KVFlash residency active; adaptive requests use AR\n"); + "concurrent KVFlash residency active; adaptive capability " + "unavailable\n"); return false; } @@ -235,7 +262,6 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { if (slot >= 0 && slot < (int)last_survival_score_.size()) { last_survival_score_[(size_t)slot] = std::numeric_limits::quiet_NaN(); - last_survival_generated_[(size_t)slot] = -1; } } }; @@ -683,11 +709,6 @@ bool Qwen35SeqEngine::batched_drafting_enabled() const { return !value || std::atoi(value) != 0; } -bool Qwen35SeqEngine::draft_always_enabled() const { - const char * value = std::getenv("DFLASH_SPEC_DRAFT_ALWAYS"); - return !value || std::atoi(value) != 0; -} - bool Qwen35SeqEngine::confidence_scoring_enabled() const { const char * value = std::getenv("DFLASH_SPEC_CONFIDENCE"); return !value || std::atoi(value) != 0; @@ -714,7 +735,7 @@ bool Qwen35SeqEngine::prepare_chain_drafts( } }; DraftTimer draft_timer{ - step_timing_enabled() ? this : nullptr, + this, std::chrono::steady_clock::now(), (int)std::count_if( selected.begin(), selected.end(), @@ -734,22 +755,21 @@ bool Qwen35SeqEngine::prepare_chain_drafts( std::vector noise((size_t)T, b_.w_.mask_token_id); std::vector noise_embed((size_t)hidden * T); - // A failed current draft must not leave the previous block's confidence - // looking current. Successful lanes publish a fresh score below. + // Proposal validity is current-block-specific. Do not clear the last + // published confidence before the draft succeeds: bootstrap must either + // publish a finite score or fail, and a later draft failure must not erase + // the immutable activation score already owned by the gate. for (size_t i = 0; i < inputs.size(); ++i) { if (!selected[i]) continue; const int slot = inputs[i].slot; if (slot < 0 || slot >= (int)prepared_chain_drafts_.size()) continue; prepared_chain_drafts_[(size_t)slot].valid = false; - last_survival_score_[(size_t)slot] = - std::numeric_limits::quiet_NaN(); - last_survival_generated_[(size_t)slot] = -1; } for (size_t i = 0; i < inputs.size(); ++i) { if (!selected[i]) continue; const StepInput & in = inputs[i]; - if (!chain_spec_input_eligible(in)) return false; + if (!chain_confidence_input_scoreable(in)) return false; DraftKvState * state = ensure_slot_draft_kv(in.slot); DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); if (!state || !mirror || @@ -894,11 +914,9 @@ bool Qwen35SeqEngine::prepare_chain_drafts( missing_confidence_warned = true; std::fprintf(stderr, "[spec-gate] current confidence unavailable; " - "adaptive request remains AR\n"); + "cannot score adaptive activation\n"); } last_survival_score_[(size_t)info.slot] = score; - last_survival_generated_[(size_t)info.slot] = - prepared.generated; } return true; } @@ -933,18 +951,36 @@ bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { } return true; } +bool Qwen35SeqEngine::chain_confidence_input_scoreable( + const StepInput & in) const { + const int hidden = b_.dw_.n_embd; + const bool have_confidence = + b_.dw_.dspark.confidence_w && + b_.dw_.dspark.confidence_b && + (b_.dw_.dspark.confidence_dim == hidden || + b_.dw_.dspark.confidence_dim == + hidden + b_.dw_.dspark.markov_rank); + return spec_mode_ == SpecMode::dspark_chain && capture_features_ && + tree_width_ > 1 && tree_width_ <= 16 && + b_.dw_.block_size == tree_width_ && b_.dw_.dspark.enabled && + have_confidence && + in.slot >= 0 && in.slot < slots_.slot_count() && + slots_.slot(in.slot).decoding() && + slots_.slot(in.slot).cur_pos >= 1 && + slots_.slot(in.slot).cur_pos < slots_.max_context(); +} + +bool Qwen35SeqEngine::chain_spec_request_capable( + const StepInput & in) const { + return chain_confidence_input_scoreable(in) && + in.allow_speculation && + in.speculation_policy != SpeculationPolicy::Never && + !slots_.slot(in.slot).sampler.needs_logit_processing(); +} + bool Qwen35SeqEngine::chain_spec_input_eligible( const StepInput & in) const { - if (spec_mode_ != SpecMode::dspark_chain || !capture_features_ || - tree_width_ <= 1 || tree_width_ > 16 || - b_.dw_.block_size != tree_width_ || !b_.dw_.dspark.enabled || - !in.allow_speculation || - in.speculation_policy == SpeculationPolicy::Never || - in.slot < 0 || in.slot >= slots_.slot_count() || - !slots_.slot(in.slot).decoding() || - slots_.slot(in.slot).sampler.needs_logit_processing() || - slots_.slot(in.slot).cur_pos < 1 || - slots_.slot(in.slot).cur_pos >= slots_.max_context()) { + if (!chain_spec_request_capable(in)) { return false; } const char * floor_value = std::getenv("DFLASH_MIN_TOKENS"); @@ -1461,20 +1497,14 @@ std::optional Qwen35SeqEngine::step_chain_spec( argmax_buf_.data(), 0, sizeof(int32_t) * argmax_buf_.size()); ggml_backend_synchronize(b_.target_backend_); - t_sample_end = timing_clock::now(); for (int lane = 0; lane < spec_count; ++lane) { - Proposal & proposal = proposals[static_cast(lane)]; - proposal.pending = argmax_buf_[static_cast(lane)]; - if (proposal.pending < 0) { + if (argmax_buf_[static_cast(lane)] < 0) { result.error = "DSpark durable replay produced invalid token"; return result; } } for (int lane = 0; lane < ar_count; ++lane) { - ArLane & ar = ar_lanes[static_cast(lane)]; - ar.pending = - argmax_buf_[static_cast(spec_count + lane)]; - if (ar.pending < 0) { + if (argmax_buf_[static_cast(spec_count + lane)] < 0) { result.error = "mixed AR durable step produced invalid token"; return result; } @@ -1487,8 +1517,36 @@ std::optional Qwen35SeqEngine::step_chain_spec( result.error = "DSpark mixed-step KV write commit failed"; return result; } + + // Publish the fed root/path before sampling the next token, matching the + // ordinary AR path's penalty history, RNG, and min-token-floor semantics. for (const StepInput & in : inputs) { slots_.commit_step(in.slot); + } + for (int lane = 0; lane < spec_count; ++lane) { + Proposal & proposal = proposals[static_cast(lane)]; + proposal.pending = sample_graph_row( + proposal.slot, lane, + &argmax_buf_[static_cast(lane)], &logits_buf_); + if (proposal.pending < 0) { + result.error = "DSpark durable replay sampling failed"; + return result; + } + } + for (int lane = 0; lane < ar_count; ++lane) { + ArLane & ar = ar_lanes[static_cast(lane)]; + const int gathered_row = spec_count + lane; + ar.pending = sample_graph_row( + ar.slot, gathered_row, + &argmax_buf_[static_cast(gathered_row)], &logits_buf_); + if (ar.pending < 0) { + result.error = "mixed AR durable sampling failed"; + return result; + } + } + t_sample_end = timing_clock::now(); + + for (const StepInput & in : inputs) { std::string reselect_error; if (!maybe_reselect_residency(in.slot, reselect_error)) { result.error = reselect_error.empty() @@ -1513,18 +1571,17 @@ std::optional Qwen35SeqEngine::step_chain_spec( out.target_forwards = 2; out.committed_tokens.assign( proposal.path.begin() + 1, proposal.path.end()); - const double survival = confidence_survival_yield( - proposal.confidence, tree_width_); - if (proposal.slot >= 0 && - proposal.slot < (int)last_survival_score_.size()) { - last_survival_score_[(size_t)proposal.slot] = survival; - last_survival_generated_[(size_t)proposal.slot] = - slots_.slot(proposal.slot).generated_tokens(); - } if (speculation_gate_) { + const double block_confidence = + proposal.slot >= 0 && + proposal.slot < (int)last_survival_score_.size() + ? last_survival_score_[(size_t)proposal.slot] + : std::numeric_limits::quiet_NaN(); speculation_gate_->observe( slots_.slot(proposal.slot).request_id, - (double)proposal.path.size()); + (double)proposal.path.size(), + slots_.slot(proposal.slot).generated_tokens(), + block_confidence); } } else { ArLane & ar = @@ -2009,7 +2066,6 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( result.slot < (int)last_survival_score_.size()) { last_survival_score_[(size_t)result.slot] = std::numeric_limits::quiet_NaN(); - last_survival_generated_[(size_t)result.slot] = -1; if (result.slot < (int)prepared_chain_drafts_.size()) { prepared_chain_drafts_[(size_t)result.slot].valid = false; @@ -2256,58 +2312,34 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } // One common wall-clock origin makes pure AR, adaptive k=0, and admitted - // speculative rounds directly comparable. In adaptive always-draft mode - // round_draft_us_ is a subset of this wall, not an extra cost to add. + // speculative rounds directly comparable and feeds the online cost model + // even when diagnostic phase logging is disabled. const bool timing = step_timing_enabled(); using timing_clock = std::chrono::steady_clock; - const auto decode_round_started = timing + const bool gate_cost_timing = + spec_mode_ == SpecMode::dspark_chain && + speculation_gate_ != nullptr; + const auto decode_round_started = timing || gate_cost_timing ? timing_clock::now() : timing_clock::time_point{}; + std::optional pending_ar_gate_plan; - if (spec_mode_ == SpecMode::dspark_chain && plan.prefills.empty()) { + if (spec_mode_ == SpecMode::dspark_chain && !inputs.empty()) { // New chain round: restart the [step-timing] draft attribution. round_draft_us_ = 0.0; round_draft_lanes_ = 0; - const auto chain_started = std::chrono::steady_clock::now(); + const auto chain_started = decode_round_started; + const bool chain_execution_available = plan.prefills.empty(); std::vector admitted(inputs.size(), 0); SpecPlan gate_plan; bool have_gate_plan = false; const char * force_value = std::getenv("DFLASH_SPEC_GATE_FORCE"); const std::string force = force_value ? force_value : ""; - bool draft_always = draft_always_enabled(); - if (draft_always && speculation_gate_) { - std::vector draft_selected(inputs.size(), 0); - for (size_t i = 0; i < inputs.size(); ++i) { - SpeculationPolicy policy = - inputs[i].speculation_policy; - if (policy == SpeculationPolicy::Adaptive) { - if (force == "all") { - policy = SpeculationPolicy::Always; - } else if (force == "none") { - policy = SpeculationPolicy::Never; - } - } - draft_selected[i] = - policy != SpeculationPolicy::Never && - chain_spec_input_eligible(inputs[i]); - } - const bool any_draft = std::any_of( - draft_selected.begin(), draft_selected.end(), - [](uint8_t value) { return value != 0; }); - if (any_draft && - !prepare_chain_drafts(inputs, draft_selected)) { - draft_always = false; - std::fprintf(stderr, - "[draft-kv-batch] current-step drafting failed; " - "using admitted-only fallback\n"); - } - } if (speculation_gate_) { std::vector candidates; candidates.reserve(inputs.size()); const bool use_confidence = confidence_scoring_enabled(); - int drafting_lanes = 0; for (const StepInput & in : inputs) { const Qwen35Slot & seq = slots_.slot(in.slot); SpeculationPolicy policy = in.speculation_policy; @@ -2315,37 +2347,121 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (force == "all") policy = SpeculationPolicy::Always; if (force == "none") policy = SpeculationPolicy::Never; } - const bool eligible = chain_spec_input_eligible(in); - drafting_lanes += - eligible && policy != SpeculationPolicy::Never ? 1 : 0; + // The confidence-off arm is an explicit AR ablation. It must + // not pay the one-time activation draft. + if (!use_confidence && + policy == SpeculationPolicy::Adaptive) { + policy = SpeculationPolicy::Never; + } + const bool scoreable = + chain_confidence_input_scoreable(in); + const bool can_speculate = + chain_spec_request_capable(in); double confidence = std::numeric_limits::quiet_NaN(); - if (use_confidence && eligible && in.slot >= 0 && + if (use_confidence && scoreable && in.slot >= 0 && in.slot < (int)last_survival_score_.size() && - last_survival_generated_[(size_t)in.slot] == - seq.generated_tokens()) { + std::isfinite( + last_survival_score_[(size_t)in.slot])) { confidence = last_survival_score_[(size_t)in.slot]; } candidates.push_back({ - seq.request_id, in.slot, policy, eligible, confidence, + seq.request_id, in.slot, policy, + scoreable, can_speculate, confidence, }); } gate_plan = speculation_gate_->plan( (int)inputs.size(), candidates, (int)inputs.size(), - draft_always ? drafting_lanes : -1); + -1); have_gate_plan = true; if (!gate_plan.valid) { return fail_step(gate_plan.error.empty() ? "adaptive speculation gate failed" : gate_plan.error); } + auto replan_with_published_confidence = [&]() { + for (SpecCandidate & candidate : candidates) { + if (!candidate.scoreable || + candidate.policy == SpeculationPolicy::Never) { + continue; + } + const int slot = candidate.slot; + if (slot >= 0 && + slot < (int)last_survival_score_.size() && + std::isfinite( + last_survival_score_[(size_t)slot])) { + candidate.confidence_yield = + last_survival_score_[(size_t)slot]; + } + } + gate_plan = speculation_gate_->plan( + (int)inputs.size(), candidates, (int)inputs.size(), + -1); + return gate_plan.valid; + }; + + // Bootstrap every cold eligible request together. This is the + // request's one-time confidence initialization, not eager drafting + // on every later k=0 round. + if (use_confidence && !gate_plan.bootstrap_slots.empty()) { + std::vector bootstrap(inputs.size(), 0); + for (int slot : gate_plan.bootstrap_slots) { + for (size_t i = 0; i < inputs.size(); ++i) { + if (inputs[i].slot == slot) bootstrap[i] = 1; + } + } + if (!prepare_chain_drafts(inputs, bootstrap)) { + return fail_step( + "adaptive confidence bootstrap failed"); + } + for (int slot : gate_plan.bootstrap_slots) { + if (slot < 0 || + slot >= (int)last_survival_score_.size() || + !std::isfinite( + last_survival_score_[(size_t)slot])) { + return fail_step( + "adaptive confidence bootstrap produced no score"); + } + } + if (!replan_with_published_confidence()) { + return fail_step(gate_plan.error.empty() + ? "adaptive speculation gate failed" : gate_plan.error); + } + if (!gate_plan.bootstrap_slots.empty() || + !gate_plan.decisions_committed) { + return fail_step( + "adaptive confidence bootstrap did not commit modes"); + } + } + log_spec_activations(gate_plan, *speculation_gate_); + + // A temporary round constraint never changes the sticky request + // mode. Execute the exact selected cohort only when every selected + // lane is runnable; otherwise this mixed/floor-constrained round + // uses packed AR and is excluded from gate cost feedback. + bool execution_matches_plan = chain_execution_available; for (int slot : gate_plan.admitted_slots) { - for (size_t i = 0; i < inputs.size(); ++i) { - if (inputs[i].slot == slot) admitted[i] = 1; + auto input = std::find_if( + inputs.begin(), inputs.end(), + [&](const StepInput & item) { + return item.slot == slot; + }); + if (input == inputs.end() || + !chain_spec_input_eligible(*input)) { + execution_matches_plan = false; + break; + } + } + if (execution_matches_plan) { + for (int slot : gate_plan.admitted_slots) { + for (size_t i = 0; i < inputs.size(); ++i) { + if (inputs[i].slot == slot) admitted[i] = 1; + } } } } else { - // A missing/failed profile degrades Adaptive to AR. Explicit - // speculation remains a reliable forced-mode oracle. + // Explicit forced speculation remains usable without a cost + // profile. Adaptive mode fails explicitly: silently serving an + // unscored request as AR violates its activation contract. for (size_t i = 0; i < inputs.size(); ++i) { SpeculationPolicy policy = inputs[i].speculation_policy; if (policy == SpeculationPolicy::Adaptive && @@ -2354,8 +2470,12 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } else if (policy == SpeculationPolicy::Adaptive && force == "none") { policy = SpeculationPolicy::Never; + } else if (policy == SpeculationPolicy::Adaptive) { + return fail_step( + "adaptive speculation gate unavailable"); } - admitted[i] = chain_spec_input_eligible(inputs[i]) && + admitted[i] = chain_execution_available && + chain_spec_input_eligible(inputs[i]) && policy == SpeculationPolicy::Always; } } @@ -2369,24 +2489,35 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const double measured_us = std::chrono::duration( std::chrono::steady_clock::now() - chain_started).count(); - if (have_gate_plan && spec_gate_debug_enabled()) { - const double realized_tokens = speculative + const bool spec_completed = + speculative && speculative->error.empty(); + const bool cost_sample_valid = + have_gate_plan && spec_completed && + round_draft_lanes_ == gate_plan.draft_lanes; + if (cost_sample_valid) { + speculation_gate_->observe_cost(gate_plan, measured_us); + } + if (have_gate_plan && spec_gate_debug_enabled() && + spec_completed) { + const double realized_tokens = spec_completed ? confidence_realized_tokens(gate_plan, *speculative) : std::numeric_limits::quiet_NaN(); log_spec_gate_plan( gate_plan, speculation_gate_->calibration_scale(), speculation_gate_->calibration_observations(), - realized_tokens, measured_us); + realized_tokens, + cost_sample_valid + ? measured_us + : std::numeric_limits::quiet_NaN()); } if (speculative) return std::move(*speculative); // Proposal setup failed before target/cache mutation. Preserve // service through the ordinary packed AR path this iteration. } - if (!any_admitted && have_gate_plan && spec_gate_debug_enabled()) { - log_spec_gate_plan( - gate_plan, speculation_gate_->calibration_scale(), - speculation_gate_->calibration_observations(), 0.0, - std::numeric_limits::quiet_NaN()); + if (!any_admitted && have_gate_plan) { + if (gate_plan.admitted_count == 0 && plan.prefills.empty()) { + pending_ar_gate_plan = gate_plan; + } } } const TargetWeights & w = b_.w_; @@ -2771,32 +2902,52 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } prefill_outputs.push_back(std::move(out)); } - if (timing && plan.prefills.empty() && live_count > 0) { + if (plan.prefills.empty() && live_count > 0 && + (timing || pending_ar_gate_plan.has_value())) { const auto t_ar_end = timing_clock::now(); auto span_us = [](timing_clock::time_point from, timing_clock::time_point to) { return std::chrono::duration( to - from).count(); }; - std::fprintf(stderr, - "[step-timing] {\"path\":\"ar\",\"live\":%d,\"k\":0," - "\"decode_bucket\":%d,\"n_prefill\":0,\"max_kv_len\":%d," - "\"draft_us\":%.1f,\"draft_lanes\":%d," - "\"pre_us\":%.1f,\"graph_build_us\":%.1f," - "\"graph_prepare_us\":%.1f,\"graph_exec_us\":%.1f," - "\"sample_read_us\":%.1f,\"finish_us\":%.1f," - "\"total_us\":%.1f,\"accepted_tokens\":0," - "\"emitted_tokens\":%d,\"target_forwards\":%d}\n", - live_count, decode_bucket, max_kv_len, - round_draft_us_, round_draft_lanes_, - span_us(decode_round_started, t_ar_build_start), - span_us(t_ar_build_start, t_ar_build_end), - span_us(t_ar_build_end, t_ar_exec_start), - span_us(t_ar_exec_start, t_ar_exec_end), - span_us(t_ar_exec_end, t_ar_read_end), - span_us(t_ar_read_end, t_ar_end), - span_us(decode_round_started, t_ar_end), - live_count, live_count); + if (pending_ar_gate_plan) { + const double measured_us = + span_us(decode_round_started, t_ar_end); + const bool cost_sample_valid = round_draft_lanes_ == 0; + if (cost_sample_valid) { + speculation_gate_->observe_cost( + *pending_ar_gate_plan, measured_us); + } + if (spec_gate_debug_enabled()) { + log_spec_gate_plan( + *pending_ar_gate_plan, speculation_gate_->calibration_scale(), + speculation_gate_->calibration_observations(), 0.0, + cost_sample_valid + ? measured_us + : std::numeric_limits::quiet_NaN()); + } + } + if (timing) { + std::fprintf(stderr, + "[step-timing] {\"path\":\"ar\",\"live\":%d,\"k\":0," + "\"decode_bucket\":%d,\"n_prefill\":0,\"max_kv_len\":%d," + "\"draft_us\":%.1f,\"draft_lanes\":%d," + "\"pre_us\":%.1f,\"graph_build_us\":%.1f," + "\"graph_prepare_us\":%.1f,\"graph_exec_us\":%.1f," + "\"sample_read_us\":%.1f,\"finish_us\":%.1f," + "\"total_us\":%.1f,\"accepted_tokens\":0," + "\"emitted_tokens\":%d,\"target_forwards\":%d}\n", + live_count, decode_bucket, max_kv_len, + round_draft_us_, round_draft_lanes_, + span_us(decode_round_started, t_ar_build_start), + span_us(t_ar_build_start, t_ar_build_end), + span_us(t_ar_build_end, t_ar_exec_start), + span_us(t_ar_exec_start, t_ar_exec_end), + span_us(t_ar_exec_end, t_ar_read_end), + span_us(t_ar_read_end, t_ar_end), + span_us(decode_round_started, t_ar_end), + live_count, live_count); + } } return result; } @@ -2808,7 +2959,6 @@ void Qwen35SeqEngine::retire(int slot) { if (slot >= 0 && slot < (int)last_survival_score_.size()) { last_survival_score_[(size_t)slot] = std::numeric_limits::quiet_NaN(); - last_survival_generated_[(size_t)slot] = -1; if (slot < (int)prepared_chain_drafts_.size()) { prepared_chain_drafts_[(size_t)slot].valid = false; } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 94bff3717..9b12edf91 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -133,6 +133,8 @@ class Qwen35SeqEngine final : public SeqEngine { DraftFeatureMirror * slot_feature_mirror(int slot); DraftKvState * ensure_slot_draft_kv(int slot); bool ddtree_eligible(const StepPlan & plan) const; + bool chain_confidence_input_scoreable(const StepInput & input) const; + bool chain_spec_request_capable(const StepInput & input) const; bool chain_spec_input_eligible(const StepInput & input) const; bool spec_gate_debug_enabled() const; struct PreparedChainDraft { @@ -146,7 +148,6 @@ class Qwen35SeqEngine final : public SeqEngine { const std::vector & inputs, const std::vector & selected); bool batched_drafting_enabled() const; - bool draft_always_enabled() const; bool confidence_scoring_enabled() const; // DFLASH_STEP_TIMING=1 emits one [step-timing] JSON line per decode // round attributing wall time to draft, verify, readback, CPU commit, @@ -176,7 +177,6 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector prepared_chain_drafts_; std::unique_ptr speculation_gate_; std::vector last_survival_score_; - std::vector last_survival_generated_; // Per-round draft cost accumulator for [step-timing]; reset at the top // of each dspark_chain round, accumulated by prepare_chain_drafts. double round_draft_us_ = 0.0; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 575b10439..46e87575c 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -277,6 +277,7 @@ KvFlashAutoBudget Qwen35Backend::make_kvflash_budget(const TargetWeights & w, bool Qwen35Backend::init() { configure_concurrent_hipblaslt_default(cfg_); + concurrent_decode_capabilities_ = {}; const bool use_remote_draft = cfg_.remote_draft.enabled(); const bool tensor_parallel = cfg_.device.is_tensor_parallel(); @@ -450,6 +451,12 @@ bool Qwen35Backend::init() { concurrent_local_draft && !cfg_.ddtree_mode && dw_.dspark.enabled && qwen35_dspark_enabled() && cfg_.speculation_policy != SpeculationPolicy::Never; + const bool concurrent_confidence_head = + concurrent_local_chain && dw_.dspark.confidence_w && + dw_.dspark.confidence_b && + (dw_.dspark.confidence_dim == dw_.n_embd || + dw_.dspark.confidence_dim == + dw_.n_embd + dw_.dspark.markov_rank); const bool concurrent_spec_tree = concurrent_local_ddtree || concurrent_local_chain; const Qwen35SeqEngine::SpecMode spec_mode = concurrent_local_ddtree @@ -597,18 +604,35 @@ bool Qwen35Backend::init() { max_concurrent_prefills, mixed_prefill_tokens, long_mixed_prefill_tokens, long_prefill_threshold, idle_prefill_tokens, prefill_quantum); - if (concurrent_local_chain && - cfg_.speculation_policy == SpeculationPolicy::Adaptive) { + concurrent_decode_capabilities_.forced_speculation = + concurrent_local_ddtree || concurrent_confidence_head; + concurrent_decode_capabilities_.adaptive = + concurrent_local_ddtree; + // Per-request decode_mode may select Adaptive even when the + // server default is forced speculation. Build the gate whenever + // the concurrent DSpark chain exists so such requests cannot + // silently fall through to unscored AR. + if (concurrent_confidence_head) { int profile_ctx = 4096; if (const char * value = std::getenv("DFLASH_SPEC_PROFILE_CONTEXT")) { profile_ctx = std::max(1, std::atoi(value)); } - if (!seq_engine_->profile_spec_costs(profile_ctx)) { + const bool profile_ready = + seq_engine_->profile_spec_costs(profile_ctx); + concurrent_decode_capabilities_.adaptive = profile_ready; + if (!profile_ready) { std::fprintf(stderr, - "[parallel-dspark] adaptive profile unavailable; " - "Adaptive requests degrade to packed AR\n"); + "[parallel-dspark] adaptive capability unavailable: " + "cost profile failed; adaptive requests will be " + "rejected at admission (forced speculation remains " + "available)\n"); } + } else if (concurrent_local_chain) { + std::fprintf(stderr, + "[parallel-dspark] activation unavailable: drafter has " + "no compatible confidence head; speculation/adaptive " + "requests will be rejected at admission\n"); } if (concurrent_local_ddtree) { const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); @@ -617,17 +641,25 @@ bool Qwen35Backend::init() { cfg_.ddtree_budget, tree_width, adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); } - if (concurrent_local_chain) { + if (concurrent_confidence_head) { std::fprintf(stderr, "[parallel-dspark] enabled width=%d mode=packed-chain-verify " "decode_mode=%s draft=q4-mix-compatible\n", tree_width, speculation_policy_name(cfg_.speculation_policy)); - } else if (concurrent_local_draft && !cfg_.ddtree_mode && - cfg_.speculation_policy != SpeculationPolicy::Never) { + } else if (!cfg_.ddtree_mode && + cfg_.speculation_policy != SpeculationPolicy::Never && + !concurrent_local_chain) { + std::fprintf(stderr, + "[parallel-dspark] unavailable for this concurrent " + "configuration; speculation/adaptive requests will be " + "rejected at admission\n"); + } else if (!cfg_.ddtree_mode && concurrent_local_draft && + dw_.dspark.enabled && qwen35_dspark_enabled()) { std::fprintf(stderr, - "[parallel-dspark] disabled: drafter lacks usable DSpark " - "Markov/confidence heads; using packed AR\n"); + "[parallel-dspark] disabled by decode_mode=ar; " + "per-request speculation/adaptive overrides will be " + "rejected at admission\n"); } std::printf("[parallel] %d decode slots, up to %d packed prefills " "(mixed short/long %d/%d at >=%d tokens, " @@ -1327,6 +1359,7 @@ DFlashTarget * Qwen35Backend::dflash_target() { void Qwen35Backend::shutdown() { const bool use_remote_draft = cfg_.remote_draft.enabled(); + concurrent_decode_capabilities_ = {}; seq_engine_.reset(); end_paged_sequence(); paged_kv_residency_.reset(); diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index e7170c5af..6f9545b5d 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -149,6 +149,10 @@ class Qwen35Backend : public ModelBackend { // attention); null otherwise, which is what tells the server to serve // one request at a time through generate(). SeqEngine * seq_engine() override; + ConcurrentDecodeCapabilities concurrent_decode_capabilities() + const override { + return concurrent_decode_capabilities_; + } // EOS identity of the loaded weights. Model-level, so it stays on the // backend and is shared by the AR decode path and the engine. @@ -318,6 +322,7 @@ class Qwen35Backend : public ModelBackend { // hence the friendship — and owns everything else concurrent serving // needs (Qwen35SlotManager, slot prefill, the batched decode step). std::unique_ptr seq_engine_; + ConcurrentDecodeCapabilities concurrent_decode_capabilities_; friend class Qwen35SeqEngine; // DFLASH_MIN_TOKENS floor for the slot paths (mirrors do_ar_decode's diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 420c785dd..9c595ff17 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -430,6 +430,24 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { return AdmissionDisposition::Retired; } + const SpeculationPolicy decode_mode = resolve_speculation_policy( + config_.decode_mode, req.decode_mode); + const ConcurrentDecodeCapabilities decode_capabilities = + backend_.concurrent_decode_capabilities(); + if (!decode_capabilities.supports(decode_mode)) { + const std::string message = + std::string("decode_mode=") + + speculation_policy_name(decode_mode) + + " is unavailable for this server's concurrent decode " + "configuration"; + std::fprintf(stderr, + "[server] concurrent admission rejected %s: %s\n", + req.response_id.c_str(), message.c_str()); + send_error(job->fd, 409, message); + finish_job(job); + return AdmissionDisposition::Retired; + } + if (!job->announced) { job->announced = true; std::fprintf(stderr, diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index fe790b6e9..829bdbaae 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -2599,6 +2599,26 @@ TEST_CASE(ServerUnitFixture, test_speculation_policy_parse_name_and_fold) { TEST_ASSERT(resolve_speculation_policy( SpeculationPolicy::Always, SpeculationPolicy::Never) == SpeculationPolicy::Never); + + const ConcurrentDecodeCapabilities ar_only{}; + TEST_ASSERT(ar_only.supports(SpeculationPolicy::Never)); + TEST_ASSERT(!ar_only.supports(SpeculationPolicy::Always)); + TEST_ASSERT(!ar_only.supports(SpeculationPolicy::Adaptive)); + + const ConcurrentDecodeCapabilities forced_only{ + /*forced_speculation=*/true, + /*adaptive=*/false, + }; + TEST_ASSERT(forced_only.supports(SpeculationPolicy::Never)); + TEST_ASSERT(forced_only.supports(SpeculationPolicy::Always)); + TEST_ASSERT(!forced_only.supports(SpeculationPolicy::Adaptive)); + + const ConcurrentDecodeCapabilities adaptive{ + /*forced_speculation=*/true, + /*adaptive=*/true, + }; + TEST_ASSERT(adaptive.supports(SpeculationPolicy::Always)); + TEST_ASSERT(adaptive.supports(SpeculationPolicy::Adaptive)); } TEST_CASE(ServerUnitFixture, test_require_messages_array_rejects_invalid) { diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index 0739cd6ca..411e5fcd7 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -3,6 +3,7 @@ #include #include +#include #include using namespace dflash::common; @@ -32,27 +33,49 @@ static SpecStepGeometry geometry() { static SpecCandidate candidate( uint64_t id, int slot, double confidence, SpeculationPolicy policy = SpeculationPolicy::Adaptive, - bool eligible = true) { - return {id, slot, policy, eligible, confidence}; + bool scoreable = true, bool can_speculate = true) { + return {id, slot, policy, scoreable, can_speculate, confidence}; } int main() { CHECK(std::abs(confidence_survival_yield({0.5f, 0.5f}, 4) - 1.75) < 1e-9); CHECK(confidence_survival_yield({2.0f, -1.0f}, 4) == 2.0); CHECK(confidence_survival_yield({}, 4) == 1.0); + CHECK(std::string(spec_decision_name(SpecDecision::Undecided)) == + "undecided"); + CHECK(std::string(spec_decision_name(SpecDecision::AR)) == "ar"); + CHECK(std::string(spec_decision_name(SpecDecision::Speculation)) == + "speculation"); + // A first finite score is a complete one-shot evaluation. Expensive + // speculation commits both requests to sticky AR and prices k=0 as the + // exact pure-AR candidate (no draft tax). SpeculationGate costly(constant_costs(100.0, 10.0, 100.0), geometry(), 4); CHECK(costly.valid()); SpecPlan plan = costly.plan(2, { candidate(1, 0, 4.0), candidate(2, 1, 4.0)}, 2); CHECK(plan.valid); + CHECK(plan.decisions_committed); + CHECK(plan.admitted_count == 0); + CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); + CHECK(plan.draft_lanes == 0); + CHECK(plan.profiled_cost == 10.0); + CHECK(plan.cost_scale == 1.0); + CHECK(plan.predicted_cost == plan.profiled_cost); + CHECK(costly.decision(1) == SpecDecision::AR); + CHECK(costly.decision(2) == SpecDecision::AR); + CHECK(costly.initial_confidence(1) == 4.0); + plan = costly.plan(2, { + candidate(1, 0, NAN), candidate(2, 1, NAN)}, 2); + CHECK(plan.bootstrap_slots.empty()); + CHECK(plan.ordered.empty()); CHECK(plan.admitted_count == 0); CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); - // A high-yield request pays for one speculative lane. Adding the - // confidence-1 freeloader cannot improve the numerator and loses the - // tie because the argmax is scanned from the smaller prefix. + // The initial ranking is a prefix argmax. The selected request remains + // speculative and the rejected request remains AR even if later inputs + // present reversed scores; the first score is immutable. SpeculationGate prefix(constant_costs(1.0, 10.0, 1.0), geometry(), 4); plan = prefix.plan(3, { @@ -62,54 +85,130 @@ int main() { CHECK((plan.admitted_request_ids == std::vector{10})); CHECK(plan.ordered.size() == 2); CHECK(plan.ordered[0].admitted); + CHECK(plan.ordered[0].decision == SpecDecision::Speculation); + CHECK(plan.ordered[1].decision == SpecDecision::AR); CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); - CHECK(std::string(spec_score_source_name(plan.ordered[0].source)) == "confidence"); - CHECK(!plan.ordered[1].admitted); + CHECK(std::string(spec_score_source_name(plan.ordered[0].source)) == + "confidence"); + CHECK(prefix.decision(10) == SpecDecision::Speculation); + CHECK(prefix.decision(11) == SpecDecision::AR); + plan = prefix.plan(3, { + candidate(10, 0, 1.0), candidate(11, 1, 4.0), + candidate(12, 2, 4.0, SpeculationPolicy::Never)}, 3); + CHECK((plan.admitted_request_ids == std::vector{10})); + CHECK(plan.ordered.size() == 1); + CHECK(plan.ordered[0].forced); + CHECK(plan.ordered[0].source == SpecScoreSource::InitialConfidence); + CHECK(prefix.initial_confidence(10) == 4.0); + CHECK(prefix.initial_confidence(11) == 1.0); - // Adaptive admission is confidence-only. Missing confidence stays on AR - // and is counted explicitly instead of consulting historical state. + // A cold batch is atomic: every scoreable undecided lane without a score + // is returned for bootstrap, and no scored-but-undecided peer commits until + // the immediate replan. The replan commits every lane exactly once. SpecCostTables crossover = constant_costs(1.0, 4.0, 1.0); - SpeculationGate confidence_only(crossover, geometry(), 4); - plan = confidence_only.plan(1, {candidate(20, 0, NAN)}, 1); + SpeculationGate bootstrap(crossover, geometry(), 4); + plan = bootstrap.plan(2, { + candidate(20, 0, NAN), candidate(21, 1, 4.0)}, 2); + CHECK(plan.valid); + CHECK(!plan.decisions_committed); CHECK(plan.admitted_count == 0); CHECK(plan.ordered.empty()); CHECK(plan.unavailable_count == 1); - plan = confidence_only.plan(2, { - candidate(20, 0, NAN), candidate(21, 1, 4.0)}, 2); + CHECK((plan.bootstrap_request_ids == std::vector{20})); + CHECK((plan.bootstrap_slots == std::vector{0})); + CHECK(bootstrap.decision(20) == SpecDecision::Undecided); + CHECK(bootstrap.decision(21) == SpecDecision::Undecided); + CHECK(bootstrap.initial_confidence(21) == 4.0); + + plan = bootstrap.plan(2, { + candidate(20, 0, 1.0), candidate(21, 1, NAN)}, 2); + CHECK(plan.valid); + CHECK(plan.decisions_committed); + CHECK(plan.bootstrap_slots.empty()); CHECK(plan.admitted_count == 1); CHECK((plan.admitted_request_ids == std::vector{21})); - CHECK(plan.unavailable_count == 1); + CHECK(bootstrap.decision(20) == SpecDecision::AR); + CHECK(bootstrap.decision(21) == SpecDecision::Speculation); + CHECK(bootstrap.initial_confidence(20) == 1.0); + CHECK(bootstrap.initial_confidence(21) == 4.0); + + plan = bootstrap.plan(2, { + candidate(20, 0, 4.0), candidate(21, 1, 1.0)}, 2); + CHECK(plan.decisions_committed); + CHECK(plan.bootstrap_slots.empty()); + CHECK(plan.admitted_count == 1); + CHECK((plan.admitted_request_ids == std::vector{21})); + CHECK(plan.ordered.size() == 1); + CHECK(plan.ordered[0].forced); + CHECK(plan.ordered[0].decision == SpecDecision::Speculation); + CHECK(bootstrap.initial_confidence(20) == 1.0); + CHECK(bootstrap.initial_confidence(21) == 4.0); - // Always and Never partition before adaptive ordering. - plan = costly.plan(3, { - candidate(1, 0, 1.0, SpeculationPolicy::Never), - candidate(2, 1, 1.0, SpeculationPolicy::Always), - candidate(3, 2, 4.0)}, 2); + // Request-lifetime scoreability is separate from permanent executor + // support. Unsupported requests still bootstrap and retain a confidence, + // then commit directly to AR. Missing mandatory scoring is an explicit + // contract error and never silently becomes an AR decision. + SpeculationGate support(crossover, geometry(), 4); + plan = support.plan(1, { + candidate(30, 0, NAN, SpeculationPolicy::Adaptive, true, false)}, 1); + CHECK(plan.valid); + CHECK(!plan.decisions_committed); + CHECK((plan.bootstrap_slots == std::vector{0})); + CHECK(support.decision(30) == SpecDecision::Undecided); + plan = support.plan(1, { + candidate(30, 0, 4.0, SpeculationPolicy::Adaptive, true, false)}, 1); + CHECK(plan.valid); + CHECK(plan.decisions_committed); + CHECK(plan.admitted_count == 0); + CHECK(plan.ordered.size() == 1); + CHECK(plan.ordered[0].decision == SpecDecision::AR); + CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); + CHECK(support.decision(30) == SpecDecision::AR); + CHECK(support.initial_confidence(30) == 4.0); + plan = support.plan(1, { + candidate(31, 0, NAN, SpeculationPolicy::Adaptive, false, false)}, 1); + CHECK(!plan.valid); + CHECK(!plan.error.empty()); + CHECK(support.decision(31) == SpecDecision::Undecided); + + // Explicit Always/Never are configured execution policies, not adaptive + // activation decisions. Always remains a non-negotiable baseline. + SpeculationGate policies(constant_costs(100.0, 10.0, 100.0), + geometry(), 4); + plan = policies.plan(3, { + candidate(40, 0, 1.0, SpeculationPolicy::Never), + candidate(41, 1, NAN, SpeculationPolicy::Always), + candidate(42, 2, 4.0)}, 2); CHECK(plan.valid); CHECK(plan.admitted_count >= 1); - CHECK(plan.admitted_request_ids.front() == 2); + CHECK(plan.admitted_request_ids.front() == 41); CHECK(plan.ordered.front().forced); - plan = costly.plan(2, { - candidate(1, 0, 4.0, SpeculationPolicy::Always), - candidate(2, 1, 4.0, SpeculationPolicy::Always)}, 1); + CHECK(plan.ordered.front().source == SpecScoreSource::Unavailable); + CHECK(policies.decision(41) == SpecDecision::Undecided); + CHECK(policies.decision(42) == SpecDecision::Speculation); + plan = policies.plan(2, { + candidate(43, 0, 4.0, SpeculationPolicy::Always), + candidate(44, 1, 4.0, SpeculationPolicy::Always)}, 1); CHECK(!plan.valid); CHECK(!plan.error.empty()); - // All-Never is a pure AR plan. SpeculationGate never(constant_costs(1.0, 2.0, 1.0), geometry(), 4); - plan = never.plan(1, {candidate(400, 0, NAN, - SpeculationPolicy::Never)}, 1); + plan = never.plan(1, { + candidate(45, 0, NAN, SpeculationPolicy::Never)}, 1); CHECK(plan.admitted_count == 0); CHECK(plan.ordered.empty()); + CHECK(never.decision(45) == SpecDecision::Undecided); - // C=1, capacity zero, malformed shapes, and always-draft pricing. - plan = prefix.plan(1, {candidate(500, 0, 4.0)}, 1); - CHECK(plan.admitted_count == 1); - plan = prefix.plan(1, {candidate(500, 0, 4.0)}, 0); + // Capacity, malformed shapes, always-draft pricing, and lookup clamps. + SpeculationGate capacity(constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = capacity.plan(1, {candidate(50, 0, 4.0)}, 0); CHECK(plan.admitted_count == 0); - plan = prefix.plan(2, {candidate(500, 0, 4.0)}, 1); + CHECK(capacity.decision(50) == SpecDecision::AR); + plan = capacity.plan(2, {candidate(51, 0, 4.0)}, 1); CHECK(!plan.valid); - plan = prefix.plan(1, {candidate(500, 0, 4.0)}, 1, 4); + SpeculationGate always_draft( + constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = always_draft.plan(1, {candidate(52, 0, 4.0)}, 1, 4); CHECK(plan.draft_lanes == 4); SpecCostSeries sparse{{2, 4}, {1.0, 2.0}}; @@ -128,89 +227,200 @@ int main() { SpeculationGate clamped(tiny, geometry(), 4, [&](const char *, int, int) { ++clamp_logs; }); plan = clamped.plan(2, { - candidate(1, 0, 4.0), candidate(2, 1, 4.0)}, 2); + candidate(60, 0, 4.0), candidate(61, 1, 4.0)}, 2); CHECK(plan.cost_lookup_clamped); CHECK(clamp_logs > 0); - // Confidence is calibrated globally from admitted confidence-scored - // rounds. A 2x-overconfident signal converges to scale 0.5 and produces - // the same cost-aware cut as the true-yield oracle. - SpecCostTables calibration_costs = constant_costs(4.0, 10.0, 10.0); - SpeculationGate calibrated(calibration_costs, geometry(), 4); - CHECK(calibrated.calibration_scale() == 1.0); - CHECK(calibrated.calibration_observations() == 0); - for (int round = 0; round < 32; ++round) { - plan = calibrated.plan(1, {candidate(700, 0, 4.0)}, 1); - CHECK(plan.admitted_count == 1); - CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); - CHECK(plan.calibration_predicted_tokens == 4.0); - calibrated.observe(700, 2.0); - if (round < 31) CHECK(calibrated.calibration_scale() == 1.0); - } - CHECK(calibrated.calibration_observations() == 32); - CHECK(std::abs(calibrated.calibration_scale() - 0.5) < 1e-12); + // Calibration is request-local and active after one observation. An + // explicit current-block confidence calibrates that observation without + // replacing the immutable initial score; the global ratio helps only new + // undecided requests. Telemetry sums the final calibrated admitted score. + SpecGateConfig fast; + fast.yield_ema_alpha = 0.5; + fast.calibration_ema_alpha = 0.5; + SpeculationGate calibrated(fast, crossover, geometry(), 4); plan = calibrated.plan(1, {candidate(700, 0, 4.0)}, 1); - SpeculationGate true_yield(calibration_costs, geometry(), 4); - SpecPlan oracle = true_yield.plan(1, {candidate(700, 0, 2.0)}, 1); - CHECK(plan.admitted_count == oracle.admitted_count); - calibrated.forget(700); + CHECK(calibrated.decision(700) == SpecDecision::Speculation); + CHECK(calibrated.initial_confidence(700) == 4.0); + calibrated.observe(700, 2.0, 1, 4.0); + CHECK(calibrated.calibration_observations() == 1); CHECK(std::abs(calibrated.calibration_scale() - 0.5) < 1e-12); + CHECK(std::abs(calibrated.calibration_scale(700) - 0.5) < 1e-12); + CHECK(std::abs(calibrated.yield_ema(700) - 2.0) < 1e-12); - // Observations without a stashed confidence prediction do nothing. - // Explicit forced speculation remains available but does not calibrate - // when its confidence is missing. + plan = calibrated.plan(1, {candidate(700, 0, 1.0)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(plan.ordered[0].forced); + CHECK(plan.ordered[0].source == SpecScoreSource::InitialConfidence); + CHECK(plan.ordered[0].expected_yield == 2.0); + CHECK(plan.calibration_predicted_tokens == 2.0); + CHECK(calibrated.initial_confidence(700) == 4.0); + + calibrated.observe(700, 4.0, 2, 2.0); + CHECK(std::abs(calibrated.calibration_scale() - 1.25) < 1e-12); + CHECK(std::abs(calibrated.calibration_scale(700) - 1.0) < 1e-12); + CHECK(std::abs(calibrated.yield_ema(700) - 3.0) < 1e-12); + CHECK(calibrated.initial_confidence(700) == 4.0); + plan = calibrated.plan(1, {candidate(700, 0, NAN)}, 1); + CHECK(plan.ordered[0].expected_yield == 4.0); + CHECK(plan.calibration_predicted_tokens == 4.0); + + plan = calibrated.plan(1, {candidate(701, 1, 2.0)}, 1); + CHECK(std::abs(plan.ordered[0].expected_yield - 2.5) < 1e-12); + CHECK(std::abs(calibrated.calibration_scale(701) - 1.25) < 1e-12); + calibrated.observe(701, 1.0, 1, 2.0); + CHECK(std::abs(calibrated.calibration_scale(701) - 0.5) < 1e-12); + CHECK(std::abs(calibrated.calibration_scale(700) - 1.0) < 1e-12); + + // A deployment may install a request-level offline fit for the cold + // cohort. It calibrates the first irreversible choice immediately, while + // the immutable raw confidence remains available for telemetry. + SpecGateConfig fitted_config; + fitted_config.initial_yield_scale = 0.5; + SpeculationGate fitted( + fitted_config, constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = fitted.plan(1, {candidate(702, 0, 4.0)}, 1); + CHECK(fitted.calibration_scale() == 0.5); + CHECK(fitted.initial_confidence(702) == 4.0); + CHECK(plan.ordered[0].expected_yield == 2.0); + CHECK(plan.ordered[0].newly_decided); + + // Yield evidence alone does not invent an activation score. Forced/probe + // execution can publish a score and train calibration immediately while + // leaving the configured policy outside the adaptive decision state. SpeculationGate isolated(constant_costs(1.0, 10.0, 1.0), geometry(), 4); - isolated.observe(800, 3.0); + isolated.observe(800, 3.0, 1); CHECK(isolated.calibration_observations() == 0); + CHECK(!isolated.has_confidence(800)); plan = isolated.plan(1, {candidate(800, 0, NAN)}, 1); - CHECK(plan.admitted_count == 0); - CHECK(plan.ordered.empty()); - CHECK(plan.unavailable_count == 1); + CHECK((plan.bootstrap_slots == std::vector{0})); + CHECK(isolated.decision(800) == SpecDecision::Undecided); + plan = isolated.plan(1, { candidate(801, 0, NAN, SpeculationPolicy::Always)}, 1); CHECK(plan.admitted_count == 1); - CHECK(plan.ordered[0].source == SpecScoreSource::Unavailable); - isolated.observe(801, 2.0); - CHECK(isolated.calibration_observations() == 0); - CHECK(isolated.calibration_scale() == 1.0); + isolated.observe(801, 2.0, 1, 4.0); + CHECK(isolated.has_confidence(801)); + CHECK(isolated.initial_confidence(801) == 4.0); + CHECK(isolated.calibration_observations() == 1); + CHECK(std::abs(isolated.calibration_scale(801) - 0.5) < 1e-12); + CHECK(isolated.decision(801) == SpecDecision::Undecided); + plan = isolated.plan(1, {candidate(801, 0, NAN)}, 1); + CHECK(plan.bootstrap_slots.empty()); + CHECK(isolated.decision(801) == SpecDecision::Speculation); + + // Forget is the only adaptive state reset. It drops the sticky decision, + // immutable score, yield, and local calibration while preserving the + // deployment-global calibration prior for future requests. + const double global_before_forget = calibrated.calibration_scale(); + calibrated.forget(700); + CHECK(!calibrated.has_state(700)); + CHECK(!calibrated.has_confidence(700)); + CHECK(std::isnan(calibrated.initial_confidence(700))); + CHECK(calibrated.decision(700) == SpecDecision::Undecided); + CHECK(calibrated.calibration_scale() == global_before_forget); + plan = calibrated.plan(1, {candidate(700, 0, NAN)}, 1); + CHECK(!plan.decisions_committed); + CHECK((plan.bootstrap_slots == std::vector{0})); - // The global ratio is bounded and survives request churn. + // Calibration clamps are active on the first forced/probe observation. SpeculationGate lower_bound(constant_costs(1.0, 1.0, 1.0), geometry(), 16); - for (int round = 0; round < 32; ++round) { - plan = lower_bound.plan(1, { - candidate(900, 0, 16.0, SpeculationPolicy::Always)}, 1); - CHECK(plan.admitted_count == 1); - lower_bound.observe(900, 1.0); - if (round < 31) CHECK(lower_bound.calibration_scale() == 1.0); - } + plan = lower_bound.plan(1, { + candidate(900, 0, 16.0, SpeculationPolicy::Always)}, 1); + lower_bound.observe(900, 1.0, 1); + CHECK(lower_bound.calibration_observations() == 1); CHECK(lower_bound.calibration_scale() == 0.25); lower_bound.forget(900); CHECK(lower_bound.calibration_scale() == 0.25); SpeculationGate upper_bound(constant_costs(1.0, 1.0, 1.0), geometry(), 16); - for (int round = 0; round < 32; ++round) { - const uint64_t request_id = 901 + static_cast(round); - plan = upper_bound.plan(1, { - candidate(request_id, 0, 4.0, SpeculationPolicy::Always)}, 1); - upper_bound.observe(request_id, 16.0); - if (round < 31) CHECK(upper_bound.calibration_scale() == 1.0); - } + plan = upper_bound.plan(1, { + candidate(901, 0, 4.0, SpeculationPolicy::Always)}, 1); + upper_bound.observe(901, 16.0, 1); + CHECK(upper_bound.calibration_observations() == 1); CHECK(upper_bound.calibration_scale() == 4.0); - // Cost-aware endpoints remain direct functions of current confidence. + // Shape-local total-cost feedback changes only future undecided choices. + // It cannot flip a request whose one-shot decision is already sticky. + SpeculationGate cost_feedback( + constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = cost_feedback.plan(1, {candidate(950, 0, 4.0)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(plan.profiled_cost == 12.0); + CHECK(plan.cost_scale == 1.0); + CHECK(cost_feedback.decision(950) == SpecDecision::Speculation); + cost_feedback.observe_cost(plan, 48.0); + plan = cost_feedback.plan(1, {candidate(950, 0, NAN)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(plan.ordered[0].forced); + CHECK(plan.profiled_cost == 12.0); + CHECK(plan.cost_scale == 4.0); + CHECK(plan.predicted_cost == 48.0); + CHECK(cost_feedback.decision(950) == SpecDecision::Speculation); + + plan = cost_feedback.plan(1, {candidate(951, 0, 4.0)}, 1); + CHECK(plan.admitted_count == 0); + CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); + CHECK(cost_feedback.decision(951) == SpecDecision::AR); + plan = cost_feedback.plan(2, { + candidate(952, 0, 4.0), candidate(953, 1, 4.0)}, 2); + CHECK(plan.admitted_count == 2); + CHECK(plan.cost_scale == 1.0); + + SpeculationGate ar_feedback( + constant_costs(100.0, 10.0, 100.0), geometry(), 4); + SpecPlan ar_plan = ar_feedback.plan(1, {candidate(960, 0, 4.0)}, 1); + CHECK(ar_plan.admitted_count == 0); + CHECK(ar_feedback.decision(960) == SpecDecision::AR); + ar_feedback.observe_cost(ar_plan, 20.0); + ar_plan = ar_feedback.plan(1, {candidate(960, 0, NAN)}, 1); + CHECK(ar_plan.admitted_count == 0); + CHECK(ar_plan.profiled_cost == 10.0); + CHECK(ar_plan.cost_scale == 2.0); + CHECK(ar_plan.predicted_cost == 20.0); + CHECK(std::abs(ar_plan.goodput - ar_plan.ar_goodput) < 1e-12); + + // Adaptive gains below the default 2% safety margin commit AR. A zero + // margin admits the same new request, while explicit Always is unchanged. + const SpecCostTables near_break_even = + constant_costs(1.0, 100.0, 1.0); + SpeculationGate margin_gate(near_break_even, geometry(), 4); + plan = margin_gate.plan(1, {candidate(970, 0, 1.04)}, 1); + CHECK(plan.admitted_count == 0); + CHECK(margin_gate.decision(970) == SpecDecision::AR); + plan = margin_gate.plan(1, {candidate(970, 0, 4.0)}, 1); + CHECK(plan.admitted_count == 0); + CHECK(plan.ordered.empty()); + CHECK(margin_gate.initial_confidence(970) == 1.04); + + SpecGateConfig zero_margin; + zero_margin.adaptive_gain_margin = 0.0; + SpeculationGate no_margin( + zero_margin, near_break_even, geometry(), 4); + plan = no_margin.plan(1, {candidate(971, 0, 1.04)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(no_margin.decision(971) == SpecDecision::Speculation); + plan = margin_gate.plan(1, { + candidate(972, 0, 1.0, SpeculationPolicy::Always)}, 1); + CHECK(plan.admitted_count == 1); + CHECK(plan.ordered[0].forced); + + // Cost-aware endpoints remain deterministic for each new request. SpeculationGate pays(constant_costs(1.0, 10.0, 1.0), geometry(), 4); plan = pays.plan(2, { - candidate(600, 0, 4.0), candidate(601, 1, 4.0)}, 2); + candidate(980, 0, 4.0), candidate(981, 1, 4.0)}, 2); CHECK(plan.admitted_count == 2); SpeculationGate cannot(constant_costs(100.0, 10.0, 100.0), geometry(), 4); std::vector eight; for (int i = 0; i < 8; ++i) - eight.push_back(candidate(30 + i, i, 4.0)); + eight.push_back(candidate(990 + i, i, 4.0)); plan = cannot.plan(8, eight, 8); CHECK(plan.admitted_count == 0); + for (int i = 0; i < 8; ++i) + CHECK(cannot.decision(990 + i) == SpecDecision::AR); std::printf("speculation gate tests passed: %d checks\n", g_checks); return 0; From fccb39a561cbe518d782f6d047b7801f67f9d396 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 15:28:51 +0000 Subject: [PATCH 33/42] fix(concurrency): enforce one-shot DSpark activation --- .../QWEN38_DSPARK_ADAPTIVE_SELECTION.md | 33 +- .../benchmarks/concurrency/FEATURE_MATRIX.md | 39 +- .../concurrency/analyze_gate_decisions.py | 117 ++++- .../concurrency/run_qwen38_dspark_matrix.sh | 5 +- .../concurrency/test_feature_tools.py | 121 ++++- .../common/concurrency/chain_spec_shapes.h | 42 ++ server/src/common/concurrency/seq_engine.h | 44 +- .../src/common/concurrency/speculation_gate.h | 112 +++-- .../qwen35/concurrency/qwen35_seq_engine.cpp | 475 +++++++++++++----- .../qwen35/concurrency/qwen35_seq_engine.h | 11 +- server/src/server/scheduler.cpp | 14 +- server/test/seq_engine_contract.h | 2 +- server/test/test_chain_spec_shapes.cpp | 38 ++ server/test/test_seq_batch_plan.cpp | 19 + server/test/test_seq_engine_contract.cpp | 21 + server/test/test_speculation_gate.cpp | 92 +++- 16 files changed, 959 insertions(+), 226 deletions(-) diff --git a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md index 53d00e227..ffe6cb453 100644 --- a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md +++ b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md @@ -100,17 +100,32 @@ harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh The dense oracle labels are priors, not a hard assertion about the concurrent executor. In particular, the two marginal dense wins may correctly become AR choices after packed-tree, replay, padding, and synchronization costs are -included. Adaptive mode makes one activation decision per request: it scores -the request at its first target-decode step, then keeps the chosen AR or -speculation mode through EOS/max tokens. +included. Adaptive mode performs one evaluation per request at its first +target-decode boundary. A successful evaluation produces the request's +confidence score and commits either AR or speculation through EOS/max tokens; +a failed evaluation commits AR without inventing a score. The request is never +evaluated again. Later execution observations calibrate activations +for future requests, not the mode of a request already in flight. The prefill logits produce the first sampled output token before a -target-decode step exists. Activation therefore happens before the first -target-decode execution, with no preliminary AR decode round; a request that -retires directly from prefill (for example, `max_tokens=1`) has no decode mode -to activate. Mixed-prefill and min-token-floor rounds may temporarily execute -through the safe AR graph after the sticky decision is made; they do not -reclassify the request. +target-decode step exists. The one-time draft and activation therefore finish +before the first target-decode execution. A request that chooses speculation +reuses that proposal immediately, so it has no preliminary AR decode round; a +request that retires directly from prefill (for example, `max_tokens=1`) has no +target-decode mode to activate. + +`DFLASH_MIN_TOKENS` does not delay speculation or switch the request to AR. +Below the floor, verification stops before an accepted EOS and durable replay +samples the replacement non-EOS token from the kept prefix tip. At or above the +floor, the EOS is kept and deeper tokens are discarded. + +Spec decode and prompt prefill use incompatible target graphs. When at least +one request in the live decode cohort chose speculation, the executor advances +that decode cohort and reports every selected prefill slice as deferred; the +scheduler retries those unchanged prompts after the Spec wave drains. This +preserves the literal request mode but can increase TTFT for requests queued +behind a long Spec wave. An all-AR cohort keeps the fused mixed +prefill/decode path and its continuous-batching TTFT behavior. The useful adaptive behavior is: diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index b4668e5e1..6db4a0d64 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -20,17 +20,27 @@ harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh The default fresh-process matrix is: - `ar`, `speculation`, `adaptive-on`, and `adaptive-confidence-off`. - `adaptive-on` batch-scores each new adaptive request once at its first - target-decode step, chooses AR or speculation, and keeps that mode through - retirement. In a decode-only round the one-time bootstrap draft is reused - immediately when the request chooses speculation; later AR rounds do not - draft. A target-decode step exists only after prefill has produced the first - sampled token, so one-token requests that retire directly from prefill never - enter adaptive activation. Mixed-prefill and min-token-floor rounds are - temporary execution constraints: the sticky decision is still committed, - but speculation starts on the first executor-safe round. The last arm sets + `adaptive-on` batch-evaluates each new adaptive request exactly once at its + first target-decode boundary. A successful score chooses AR or speculation; + a failed evaluation chooses AR without a synthetic score. That mode remains + fixed through retirement. If speculation is chosen, the one-time bootstrap + proposal is reused immediately; a request that chooses AR never drafts again. + A target-decode step exists only after prefill has produced the first sampled + token, so one-token requests that retire directly from prefill never enter + adaptive activation. Later execution evidence calibrates activations for new + requests and never re-evaluates an active request. The last arm sets `DFLASH_SPEC_CONFIDENCE=0`, suppressing activation scoring so it remains a draft-free AR ablation. There is no eager-draft matrix axis. +- A Spec request runs the speculative executor immediately even below + `DFLASH_MIN_TOKENS`. An accepted EOS below the floor is excluded from the + committed path and durable replay samples the non-EOS replacement; an allowed + EOS ends the path. The floor never creates an AR warm-up phase. +- Spec decode is graph-exclusive with prompt prefill. If a scheduler plan + contains any sticky-Spec lane, decode advances and selected prefill slices are + reported as deferred with no prompt progress. They are retried after the Spec + wave drains. This can raise TTFT for queued requests, so the inverse-TTFT gate + measures an intentional mode-fidelity tradeoff. All-AR cohorts continue + using the fused mixed prefill/decode path. - Live concurrency `C ∈ {1,2,3,4,6,8}` over the checked-in HumanEval and GSM8K cohorts plus deterministic prose prompts. - A fixed C=6 north-star cohort with two code and four chat requests. @@ -42,10 +52,13 @@ Every process records the target, server, shared-library, and drafter hashes, the literal command and launch environment, startup pool dimensions, request IDs, and terminal concurrency counters. The proof rejects forced-speculation rows unless every measured request has positive `spec_steps`. Adaptive rows -may legitimately choose k=0, but must show exactly one finite -`[spec-activation]` AR/speculation decision for every measured adaptive -request, plus both the packed DSpark startup marker and a completed startup -cost profile. Chain rows must keep all +may legitimately choose k=0, but must show exactly one `[spec-activation]` +AR/speculation decision for every measured adaptive request. A scored +activation carries finite confidence/yield and no fallback reason; a failed +one carries null scores, sticky AR, and +`fallback_reason=confidence_evaluation_failed`. The proof also requires both +the packed DSpark startup marker and a completed startup cost profile. Chain +rows must keep all `ddtree_*` counters at zero, preserving the DDTree proof semantics below. For every workload/concurrency pair, the summarizer forms a paired oracle: diff --git a/harness/benchmarks/concurrency/analyze_gate_decisions.py b/harness/benchmarks/concurrency/analyze_gate_decisions.py index 1ff37b3ae..7995ac79d 100644 --- a/harness/benchmarks/concurrency/analyze_gate_decisions.py +++ b/harness/benchmarks/concurrency/analyze_gate_decisions.py @@ -81,21 +81,59 @@ def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: f"{path}:{line_no}: spec-activation {key} must be a " "non-negative int" ) + decision = row.get("decision") + if decision not in ("ar", "speculation"): + raise ValueError( + f"{path}:{line_no}: spec-activation decision must be ar or " + "speculation" + ) + evaluation = row.get("evaluation") + if evaluation not in ("scored", "failed"): + raise ValueError( + f"{path}:{line_no}: spec-activation evaluation must be scored " + "or failed" + ) + if "fallback_reason" not in row: + raise ValueError( + f"{path}:{line_no}: spec-activation fallback_reason is required" + ) for key in ("initial_confidence", "calibrated_yield"): - value = row.get(key) - if ( - type(value) not in (int, float) - or not math.isfinite(value) - or value < 1.0 - ): + if key not in row: raise ValueError( - f"{path}:{line_no}: spec-activation {key} must be finite " - "and at least 1" + f"{path}:{line_no}: spec-activation {key} is required" ) - if row.get("decision") not in ("ar", "speculation"): + if evaluation == "scored": + for key in ("initial_confidence", "calibrated_yield"): + value = row.get(key) + if ( + type(value) not in (int, float) + or not math.isfinite(value) + or value < 1.0 + ): + raise ValueError( + f"{path}:{line_no}: spec-activation {key} must be finite " + "and at least 1 for a scored evaluation" + ) + if row["fallback_reason"] is not None: + raise ValueError( + f"{path}:{line_no}: scored spec-activation fallback_reason " + "must be null" + ) + return + if row.get("initial_confidence") is not None or ( + row.get("calibrated_yield") is not None + ): raise ValueError( - f"{path}:{line_no}: spec-activation decision must be ar or " - "speculation" + f"{path}:{line_no}: failed spec-activation scores must be null" + ) + if decision != "ar": + raise ValueError( + f"{path}:{line_no}: failed spec-activation decision must be ar" + ) + if row["fallback_reason"] != "confidence_evaluation_failed": + raise ValueError( + f"{path}:{line_no}: failed spec-activation fallback_reason must " + "be confidence_evaluation_failed" ) @@ -440,6 +478,39 @@ def _activation_proof( "missing activations for engine requests " + ",".join(str(value) for value in missing_ids) ) + for request_id in sorted(expected_ids & set(by_id)): + rows = by_id[request_id] + if len(rows) != 1: + continue + metric = metrics[engine_to_wire[request_id]] + spec_steps = metric.get("spec_steps") + target_forwards = metric.get("target_forwards") + if ( + type(spec_steps) is not int + or spec_steps < 0 + or type(target_forwards) is not int + or target_forwards < 0 + ): + errors.append( + f"engine request {request_id} has invalid execution " + "counters" + ) + continue + decision = rows[0]["decision"] + if decision == "ar" and spec_steps != 0: + errors.append( + f"AR activation for engine request {request_id} " + f"executed {spec_steps} speculation steps" + ) + if decision == "speculation" and ( + spec_steps == 0 or target_forwards != 2 * spec_steps + ): + errors.append( + f"Spec activation for engine request {request_id} " + "contains a non-speculative target step " + f"(spec_steps={spec_steps}, " + f"target_forwards={target_forwards})" + ) if errors: raise ValueError( f"{log_path}: adaptive-on activation proof failed: " @@ -453,14 +524,22 @@ def _activation_proof( ) for decision in ("ar", "speculation") } + evaluations = { + evaluation: sum( + 1 for row in activations if row["evaluation"] == evaluation + ) + for evaluation in ("scored", "failed") + } return ({ "required": required, "validation": "passed" if required else "not-required", + "execution_validation": "passed" if required else "not-required", "measured_engine_requests": len(expected_ids), "records": len(activations), "unique_requests": len(by_id), "matched_requests": len(set(by_id) & expected_ids), "decision_counts": decisions, + "evaluation_counts": evaluations, }, unique) @@ -550,6 +629,14 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A "activation_decision": ( activation.get("decision") if activation is not None else None ), + "activation_evaluation": ( + activation.get("evaluation") + if activation is not None else None + ), + "activation_fallback_reason": ( + activation.get("fallback_reason") + if activation is not None else None + ), "commit_per_spec_step": ( (steps + accepted) / steps if steps else None ), @@ -693,6 +780,12 @@ def compare_prompts(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: "activation_decision": request.get( "activation_decision" ), + "activation_evaluation": request.get( + "activation_evaluation" + ), + "activation_fallback_reason": request.get( + "activation_fallback_reason" + ), "initial_confidence": request.get("initial_confidence"), "calibrated_yield": request.get("calibrated_yield"), "mean_confidence_yield": request.get( @@ -979,7 +1072,7 @@ def build_report( ] activation_comparisons = compare_activation_shapes(cases) return { - "schema_version": 4, + "schema_version": 5, "cases": cases, "prompt_comparisons": compare_prompts(cases), "activation_comparisons": activation_comparisons, diff --git a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh index cad371f53..3444781dc 100755 --- a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh @@ -54,8 +54,9 @@ adaptive-selection baseline uses the no-YaRN Q8_0 artifact; changing draft quantization changes both acceptance and cost, so it invalidates those stored oracle ratios. The default fresh-process matrix runs ar, forced speculation, one-shot -adaptive activation with confidence scored once for every request, and a -draft-free confidence-off AR ablation at live concurrency 1,2,3,4,6,8 over +adaptive activation with one confidence evaluation per request (a failed +evaluation falls back to sticky AR), and a draft-free confidence-off AR +ablation at live concurrency 1,2,3,4,6,8 over HumanEval, GSM8K, and prose. An adaptive request keeps its selected speculation/AR mode until retirement. The 2-code+4-chat north-star row runs only at C=6. The summary fails if adaptive-on is below 0.995 of the paired diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index c12af11a7..29414cf29 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -277,6 +277,7 @@ def _write_activation_case( variant: str, activations: list[dict], engine_ids: tuple[int, ...] = (7, 8), + execution: dict[int, tuple[int, int]] | None = None, ) -> Path: case = root / "selection" / "c2" / "r1" / variant case.mkdir(parents=True) @@ -304,8 +305,12 @@ def _write_activation_case( { "request_id": f"wire-{index}", "engine_request_id": engine_id, - "spec_accepted_tokens": 0, "spec_steps": 0, - "target_forwards": 1, "output_tokens": 1, + "spec_accepted_tokens": 0, + "spec_steps": (execution or {}).get(engine_id, (0, 1))[0], + "target_forwards": (execution or {}).get( + engine_id, (0, 1) + )[1], + "output_tokens": 1, } for index, engine_id in enumerate(engine_ids, 1) ] @@ -371,6 +376,7 @@ def test_measured_step_timing_is_joined_and_summarized(self) -> None: activation = { "request_id": 7, "slot": 0, "initial_confidence": 3.25, "calibrated_yield": 1.75, + "evaluation": "scored", "fallback_reason": None, "decision": "ar", } (case / "benchmark-server.log").write_text( @@ -403,16 +409,20 @@ def test_measured_step_timing_is_joined_and_summarized(self) -> None: self.assertEqual(request["initial_confidence"], 3.25) self.assertEqual(request["calibrated_yield"], 1.75) self.assertEqual(request["activation_decision"], "ar") + self.assertEqual(request["activation_evaluation"], "scored") + self.assertIsNone(request["activation_fallback_reason"]) def test_adaptive_on_activation_proof_fails_closed(self) -> None: activation_7 = { "request_id": 7, "slot": 0, "initial_confidence": 2.0, "calibrated_yield": 1.25, + "evaluation": "scored", "fallback_reason": None, "decision": "speculation", } activation_8 = { "request_id": 8, "slot": 1, "initial_confidence": 1.0, "calibrated_yield": 1.0, + "evaluation": "scored", "fallback_reason": None, "decision": "ar", } activation_9 = {**activation_8, "request_id": 9} @@ -449,18 +459,118 @@ def test_confidence_off_is_exempt_from_activation_coverage(self) -> None: for row in report["requests"] )) + def test_adaptive_on_activation_proof_enforces_sticky_execution(self) -> None: + activation = { + "request_id": 7, "slot": 0, + "initial_confidence": 2.0, "calibrated_yield": 1.5, + "evaluation": "scored", "fallback_reason": None, + "decision": "speculation", + } + with tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", [activation], + engine_ids=(7,), execution={7: (2, 4)}, + ) + report = gate_analysis.analyze_case(case) + self.assertEqual( + report["activation"]["execution_validation"], "passed", + ) + bad_cases = ( + ( + {**activation, "decision": "ar"}, {7: (1, 2)}, + "AR activation.*executed 1 speculation steps", + ), + ( + activation, {7: (2, 5)}, + "Spec activation.*contains a non-speculative target step", + ), + ( + activation, {7: (0, 1)}, + "Spec activation.*contains a non-speculative target step", + ), + ) + for row, execution, message in bad_cases: + with self.subTest(message=message), tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", [row], + engine_ids=(7,), execution=execution, + ) + with self.assertRaisesRegex(ValueError, message): + gate_analysis.analyze_case(case) + + def test_failed_evaluation_activation_is_request_local_and_explicit( + self, + ) -> None: + failed = { + "request_id": 7, "slot": 0, + "initial_confidence": None, "calibrated_yield": None, + "evaluation": "failed", + "fallback_reason": "confidence_evaluation_failed", + "decision": "ar", + } + scored = { + "request_id": 8, "slot": 1, + "initial_confidence": 2.0, "calibrated_yield": 1.5, + "evaluation": "scored", "fallback_reason": None, + "decision": "speculation", + } + with tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", [failed, scored], + execution={7: (0, 1), 8: (2, 4)}, + ) + report = gate_analysis.analyze_case(case) + self.assertEqual( + report["activation"]["evaluation_counts"], + {"scored": 1, "failed": 1}, + ) + by_id = { + row["engine_request_id"]: row for row in report["requests"] + } + self.assertEqual(by_id[7]["activation_decision"], "ar") + self.assertEqual(by_id[7]["activation_evaluation"], "failed") + self.assertIsNone(by_id[7]["initial_confidence"]) + self.assertEqual( + by_id[7]["activation_fallback_reason"], + "confidence_evaluation_failed", + ) + self.assertEqual( + by_id[8]["activation_evaluation"], "scored", + ) + def test_activation_record_fields_are_strictly_validated(self) -> None: valid = { "request_id": 7, "slot": 0, "initial_confidence": 1.0, "calibrated_yield": 1.0, + "evaluation": "scored", "fallback_reason": None, "decision": "ar", } + failed = { + "request_id": 7, "slot": 0, + "initial_confidence": None, "calibrated_yield": None, + "evaluation": "failed", + "fallback_reason": "confidence_evaluation_failed", + "decision": "ar", + } + missing_evaluation = dict(valid) + missing_evaluation.pop("evaluation") + missing_fallback_reason = dict(valid) + missing_fallback_reason.pop("fallback_reason") + missing_failed_score = dict(failed) + missing_failed_score.pop("initial_confidence") cases = ( ({**valid, "request_id": True}, "request_id"), ({**valid, "slot": -1}, "slot"), ({**valid, "initial_confidence": 0.99}, "initial_confidence"), ({**valid, "calibrated_yield": math.nan}, "calibrated_yield"), ({**valid, "decision": "undecided"}, "decision"), + (missing_evaluation, "evaluation"), + (missing_fallback_reason, "fallback_reason"), + (missing_failed_score, "initial_confidence"), + ({**valid, "fallback_reason": "failure"}, "fallback_reason"), + ({**failed, "initial_confidence": 1.0}, "scores must be null"), + ({**failed, "decision": "speculation"}, "decision must be ar"), + ({**failed, "fallback_reason": "draft_failed"}, "fallback_reason"), ) for row, message in cases: with self.subTest(field=message), tempfile.TemporaryDirectory() as tmp: @@ -470,6 +580,13 @@ def test_activation_record_fields_are_strictly_validated(self) -> None: ) with self.assertRaisesRegex(ValueError, message): gate_analysis.parse_server_log(path) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(failed)}\n", encoding="utf-8", + ) + _, _, _, activations = gate_analysis.parse_server_log(path) + self.assertEqual(activations, [failed]) def test_paired_controls_produce_a_per_prompt_concurrent_oracle(self) -> None: base_request = { diff --git a/server/src/common/concurrency/chain_spec_shapes.h b/server/src/common/concurrency/chain_spec_shapes.h index 12907bef7..0b392d32d 100644 --- a/server/src/common/concurrency/chain_spec_shapes.h +++ b/server/src/common/concurrency/chain_spec_shapes.h @@ -83,4 +83,46 @@ inline ChainLaunchShape chain_launch_shape( return shape; } +// Proposal preparation is independent per request. A requested speculative +// lane whose proposal failed must be removed from both executor cohorts: it is +// a lane-local failure, never an AR fallback for that decode step. +enum class ChainLaneDisposition : uint8_t { + AR, + Speculation, + Failed, +}; + +inline ChainLaneDisposition chain_lane_disposition( + bool requested_speculation, bool proposal_failed) { + if (proposal_failed) return ChainLaneDisposition::Failed; + return requested_speculation + ? ChainLaneDisposition::Speculation + : ChainLaneDisposition::AR; +} + +inline bool chain_lane_executes(ChainLaneDisposition disposition) { + return disposition != ChainLaneDisposition::Failed; +} + +// The pending root at path[0] was sampled by the preceding target step, so +// the ordinary sampler has already applied the min-token EOS floor to it. +// Accepted children would bypass that sampler. Stop before an EOS that is +// still below the floor so replay samples a replacement from the kept +// tip's exact logits. Once the floor is met, keep the EOS itself but discard +// deeper accepted tokens that the scheduler would hide after retirement. +template +inline size_t chain_min_tokens_safe_prefix( + const std::vector & path, + int generated_tokens_before_root, + int min_tokens, + IsEos is_eos) { + const int generated = std::max(0, generated_tokens_before_root); + for (size_t child = 1; child < path.size(); ++child) { + if (!is_eos(path[child])) continue; + return generated + static_cast(child) < min_tokens + ? child : child + 1; + } + return path.size(); +} + } // namespace dflash::common diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index c67a2357d..62147caa8 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -6,12 +6,15 @@ // paged KV cache and execute a batched decode step. Any additional // per-sequence model state is owned by the concrete engine, not by this // interface. admit() claims a slot and queues its prompt without compute. -// Each step() then advances a scheduler-selected cohort of prompt slices -// alongside the complete live decode batch. Once a prefill completes, the -// scheduler 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. +// Each step() advances the complete live decode batch and may also advance a +// scheduler-selected cohort of prompt slices. A decode graph that cannot mix +// with prefill reports those slices as deferred and leaves their prompt state +// unchanged; the scheduler selects them again after the exclusive decode wave. +// Once a prefill completes, the 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 preserve authority to substitute a +// thinking-budget close token. // // The split of duties is deliberate and is the reason this interface exists // apart from ModelBackend: @@ -220,6 +223,11 @@ class SeqEngine { enum class Status { advanced, completed, + // Selected work intentionally made no progress because the same + // engine call executed an incompatible decode graph. No prompt + // state changed; the scheduler keeps it pending and selects it + // again after the exclusive decode wave. + deferred, failed, }; @@ -262,8 +270,9 @@ class SeqEngine { virtual StepPlanLimits step_plan_limits(int decode_rows) const = 0; // A successful result returns one decode output for every decode input and - // one explicit advanced/completed/failed result for every selected - // prefill. Invalid plans return a fatal error without advancing state. + // one explicit advanced/completed/deferred/failed result for every + // selected prefill. Invalid plans return a fatal error without advancing + // state. // Runtime failures are terminal for the live cohort and may follow partial // backend mutation, but expose no consumable payload. virtual StepResult step(const StepPlan & plan) = 0; @@ -275,6 +284,20 @@ class SeqEngine { virtual bool token_is_eos(int32_t token) const = 0; }; +// Scheduler fairness advances only when selected prompt work actually moved. +// A deferred or failed slice remains non-progress even though both are valid, +// explicit answers for the selected row. +inline bool prefill_result_made_progress( + const SeqEngine::StepResult & result) { + using Status = SeqEngine::PrefillOutput::Status; + return std::any_of( + result.prefills.begin(), result.prefills.end(), + [](const SeqEngine::PrefillOutput & output) { + return output.status == Status::advanced || + output.status == Status::completed; + }); +} + // 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. @@ -371,6 +394,7 @@ inline std::string validate_step_result( return "step returned duplicate prefill outputs"; if (output.status != PrefillStatus::advanced && output.status != PrefillStatus::completed && + output.status != PrefillStatus::deferred && output.status != PrefillStatus::failed) return "prefill output has an unknown status"; if (output.status == PrefillStatus::advanced && @@ -379,6 +403,10 @@ inline std::string validate_step_result( if (output.status == PrefillStatus::completed && (output.token < 0 || !output.error.empty())) return "completed prefill has invalid payload"; + if (output.status == PrefillStatus::deferred && + (plan.decode.empty() || output.token >= 0 || + !output.error.empty())) + return "deferred prefill has invalid payload or no decode peer"; if (output.status == PrefillStatus::failed && (output.token >= 0 || output.error.empty())) return "failed prefill has invalid payload"; diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h index a4c782d49..4256b487b 100644 --- a/server/src/common/concurrency/speculation_gate.h +++ b/server/src/common/concurrency/speculation_gate.h @@ -1,4 +1,7 @@ -// Per-request adaptive speculation policy over startup-profiled costs. +// Per-request adaptive speculation policy over startup-profiled costs. Every +// adaptive request receives one AR or speculation decision and keeps that +// decision until forget(). A failed confidence evaluation is represented +// explicitly and commits sticky AR without inventing a score. // Pure host code: no graph, backend, or scheduler types belong here. #pragma once @@ -87,18 +90,19 @@ struct SpecCandidate { int slot = -1; SpeculationPolicy policy = SpeculationPolicy::Adaptive; // scoreable is a request-lifetime capability: the engine can produce the - // mandatory one-time activation score from the committed feature mirror. + // one-time activation score from the committed feature mirror. bool scoreable = false; // can_speculate is also request-lifetime (for example, false for an - // unsupported sampler or thinking hook). Temporary blockers such as a - // min-token floor or a mixed-prefill round must not clear this flag; the - // engine executes AR for that round without changing the sticky decision. + // unsupported sampler or thinking hook). The min-token EOS policy is + // enforced inside the speculative path, and an incompatible prefill graph + // is deferred, so neither changes this capability or the chosen mode. bool can_speculate = false; - // NaN means the adapter did not supply a score in this plan. An undecided - // adaptive request with no retained score is returned for its one-time - // bootstrap. The first finite score is retained for the lifetime of the - // request and later scores are ignored. This is the survival-product - // expected yield, including the root; the gate owns calibration/clamping. + // While an adaptive request is Undecided, NaN requests its one-time + // bootstrap and a finite value is the preferred activation measurement. + // Evaluation failure explicitly falls back to sticky AR without inventing + // a score. Otherwise the gate commits exactly one mode from this + // survival-product expected yield, including the root; the gate owns + // calibration and clamping. double confidence_yield = std::numeric_limits::quiet_NaN(); }; @@ -143,6 +147,17 @@ enum class SpecDecision : uint8_t { Speculation, }; +enum class SpecEvaluationAction : uint8_t { + Score, + FallbackAR, +}; + +struct SpecPendingEvaluation { + uint64_t request_id = 0; + int slot = -1; + SpecEvaluationAction action = SpecEvaluationAction::Score; +}; + inline const char * spec_decision_name(SpecDecision decision) { switch (decision) { case SpecDecision::Undecided: return "undecided"; @@ -187,18 +202,19 @@ struct SpecPlan { double ar_goodput = 0.0; int unavailable_count = 0; // False means at least one adaptive request still needs its one-time - // confidence bootstrap. No new adaptive decisions are committed in that - // plan. The immediate post-bootstrap replan commits all undecided lanes. + // evaluation action. No new adaptive decisions are committed in that + // plan. The engine resolves every tagged action request-locally, then one + // immediate replan commits all remaining undecided lanes atomically. bool decisions_committed = false; bool cost_lookup_clamped = false; std::vector ordered; std::vector admitted_request_ids; std::vector admitted_slots; - // Cold requests are unavailable only until their one-time initialization. - // The engine drafts these together and immediately replans before choosing - // the request's sticky AR/speculation mode. - std::vector bootstrap_request_ids; - std::vector bootstrap_slots; + // Score actions are batched for one-time confidence initialization. + // FallbackAR actions cannot attempt scoring and must instead commit sticky + // AR with an explicit failed-evaluation activation. One tagged record keeps + // request identity and slot inseparable on all failure paths. + std::vector pending_evaluations; }; // Generic confidence contract for every chain speculator: `confidences[i]` @@ -228,6 +244,7 @@ class SpeculationGate { double initial_confidence = std::numeric_limits::quiet_NaN(); SpecDecision decision = SpecDecision::Undecided; + bool confidence_evaluation_failed = false; double yield_ema = std::numeric_limits::quiet_NaN(); uint64_t yield_observations = 0; double calibration_ratio_ema = 1.0; @@ -373,34 +390,23 @@ class SpeculationGate { if (adaptive_undecided && score.source == SpecScoreSource::Unavailable) { ++out.unavailable_count; - if (candidate.scoreable) { - out.bootstrap_request_ids.push_back(candidate.request_id); - out.bootstrap_slots.push_back(candidate.slot); - } else { - // Never silently turn a missing mandatory score into a - // decision. A deployment that cannot score an adaptive - // DSpark request must handle that contract failure - // explicitly rather than inventing an AR activation. - out.valid = false; - out.error = "adaptive request is not scoreable"; - return out; - } + out.pending_evaluations.push_back({ + candidate.request_id, candidate.slot, + candidate.scoreable + ? SpecEvaluationAction::Score + : SpecEvaluationAction::FallbackAR, + }); continue; } if (adaptive_undecided && !candidate.can_speculate) { - // Still retain its mandatory initial score for telemetry, but - // permanently unsupported execution commits directly to AR. + // Record the mandatory initial score in activation telemetry, + // but permanently unsupported execution commits directly to AR. forced_ar.push_back({ &candidate, score.expected_yield, score.uncalibrated_confidence, score.source, prior_decision, false, true}); continue; } - if (candidate.policy == SpeculationPolicy::Always && - prior_decision != SpecDecision::Speculation && - !candidate.can_speculate) { - continue; - } Ranked ranked{&candidate, score.expected_yield, score.uncalibrated_confidence, score.source, prior_decision, @@ -415,8 +421,10 @@ class SpeculationGate { // until the engine publishes the batched bootstrap and immediately // replans. Previously decided speculation and explicit Always lanes // can still execute in the bootstrap plan. - if (!out.bootstrap_slots.empty()) undecided.clear(); - out.decisions_committed = out.bootstrap_slots.empty(); + if (!out.pending_evaluations.empty()) { + undecided.clear(); + } + out.decisions_committed = out.pending_evaluations.empty(); if (out.decisions_committed) { for (const Ranked & item : forced_ar) { request_states_[item.candidate->request_id].decision = @@ -582,11 +590,10 @@ class SpeculationGate { return out; } - // `confidence_yield` may carry the freshly drafted current-block score. - // It calibrates this observation (including forced/probe executions) but - // never overwrites the one-shot activation score retained for the request. - // `generated_tokens` is kept for source compatibility with adapters that - // already pass it; one-shot decisions have no age or refresh semantics. + // `confidence_yield` may describe the completed speculative block. It + // calibrates later request activations (including forced executions) but + // cannot change this request's activation score or mode. `generated_tokens` + // remains only for source compatibility with existing adapters. void observe(uint64_t request_id, double emitted_tokens, int /* generated_tokens */, double confidence_yield = @@ -648,6 +655,20 @@ class SpeculationGate { config_.cost_ema_alpha); } + // Commit the explicit cold-evaluation failure policy. This is a real + // sticky AR decision, but intentionally has no initial confidence value. + // False means the request had already committed a mode. + bool commit_evaluation_fallback_ar(uint64_t request_id) { + RequestState & state = request_states_[request_id]; + if (state.decision != SpecDecision::Undecided) return false; + state.initial_confidence = + std::numeric_limits::quiet_NaN(); + state.confidence_evaluation_failed = true; + state.decision = SpecDecision::AR; + pending_confidence_.erase(request_id); + return true; + } + void forget(uint64_t request_id) { request_states_.erase(request_id); pending_confidence_.erase(request_id); @@ -661,6 +682,11 @@ class SpeculationGate { return state != request_states_.end() && std::isfinite(state->second.initial_confidence); } + bool confidence_evaluation_failed(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state != request_states_.end() && + state->second.confidence_evaluation_failed; + } double initial_confidence(uint64_t request_id) const { auto state = request_states_.find(request_id); return state == request_states_.end() diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index f606878fc..f9aa51473 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -125,6 +125,7 @@ void log_spec_activations(const SpecPlan & plan, std::fprintf(stderr, "[spec-activation] {\"request_id\":%llu,\"slot\":%d," "\"initial_confidence\":%.6f,\"calibrated_yield\":%.6f," + "\"evaluation\":\"scored\",\"fallback_reason\":null," "\"decision\":\"%s\"}\n", (unsigned long long)score.request_id, score.slot, initial, score.expected_yield, @@ -132,6 +133,16 @@ void log_spec_activations(const SpecPlan & plan, } } +void log_spec_evaluation_fallback(uint64_t request_id, int slot) { + std::fprintf(stderr, + "[spec-activation] {\"request_id\":%llu,\"slot\":%d," + "\"initial_confidence\":null,\"calibrated_yield\":null," + "\"evaluation\":\"failed\"," + "\"fallback_reason\":\"confidence_evaluation_failed\"," + "\"decision\":\"ar\"}\n", + (unsigned long long)request_id, slot); +} + } // namespace Qwen35SeqEngine::Qwen35SeqEngine( @@ -716,7 +727,9 @@ bool Qwen35SeqEngine::confidence_scoring_enabled() const { bool Qwen35SeqEngine::prepare_chain_drafts( const std::vector & inputs, - const std::vector & selected) { + const std::vector & selected, + bool force_serial, + bool fail_fast_batch) { if (selected.size() != inputs.size()) return false; // Accumulate the full drafting wall (draft graph compute + fused @@ -795,7 +808,8 @@ bool Qwen35SeqEngine::prepare_chain_drafts( std::vector> drafts; std::vector> confidences; bool used_batch = false; - if (batched_drafting_enabled()) { + const bool try_batched = !force_serial && batched_drafting_enabled(); + if (try_batched) { const int bucket = chain_decode_bucket_width((int)lanes.size()); std::vector batch_states; @@ -858,6 +872,32 @@ bool Qwen35SeqEngine::prepare_chain_drafts( } } + if (!used_batch && try_batched && fail_fast_batch) return false; + + if (!used_batch && try_batched) { + // A failed backend compute can leave a subset of the packed draft + // graph's cache writes visible. Rebuild every real lane from its + // captured target-feature ring before entering the serial fallback. + for (const Lane & lane : lanes) { + draft_kv_reset(*lane.state); + if (!draft_kv_begin_step( + *lane.state, b_.dw_, b_.draft_backend_, *lane.mirror, + slots_.slot(lane.slot).cur_pos)) { + return false; + } + noise[0] = lane.seed; + std::fill( + noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), T, noise_embed.data())) { + return false; + } + ggml_backend_tensor_set( + lane.state->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + } + } + if (!used_batch) { drafts.resize(lanes.size()); confidences.resize(lanes.size()); @@ -980,15 +1020,9 @@ bool Qwen35SeqEngine::chain_spec_request_capable( bool Qwen35SeqEngine::chain_spec_input_eligible( const StepInput & in) const { - if (!chain_spec_request_capable(in)) { - return false; - } - const char * floor_value = std::getenv("DFLASH_MIN_TOKENS"); - const int floor = floor_value - ? std::max(0, std::atoi(floor_value)) : 0; - return slots_.slot(in.slot).generated_tokens() >= floor; + return chain_spec_request_capable(in); } -std::optional Qwen35SeqEngine::step_chain_spec( +SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( const StepPlan & plan, const std::vector & admitted, std::chrono::steady_clock::time_point round_started) { StepResult result; @@ -1002,10 +1036,17 @@ std::optional Qwen35SeqEngine::step_chain_spec( const int hidden = b_.w_.n_embd; const int n_head_kv = b_.w_.n_head_kv; const int n_slots = slots_.slot_count(); - int spec_count = 0; - for (uint8_t value : admitted) spec_count += value != 0; - if (spec_count == 0) return std::nullopt; - const int tree_bucket = chain_decode_bucket_width(spec_count); + const int min_tokens = []() { + const char * value = std::getenv("DFLASH_MIN_TOKENS"); + return value ? std::max(0, std::atoi(value)) : 0; + }(); + const int requested_spec_count = static_cast(std::count_if( + admitted.begin(), admitted.end(), + [](uint8_t value) { return value != 0; })); + if (requested_spec_count == 0) { + result.error = "empty DSpark chain admission plan"; + return result; + } // Optional per-phase wall attribution ([step-timing]). Timestamps mark // phase boundaries; graph computes are synchronous on this backend, so @@ -1043,28 +1084,55 @@ std::optional Qwen35SeqEngine::step_chain_spec( }; std::vector proposals; - proposals.reserve(static_cast(spec_count)); + proposals.reserve(static_cast(requested_spec_count)); std::vector proposal_for_input(inputs.size(), -1); - std::vector drafted_slots; - drafted_slots.reserve(inputs.size()); + std::vector active_admitted = admitted; + std::vector retried(inputs.size(), 0); + std::vector proposal_errors(inputs.size()); - auto proposal_fallback = [&]() -> std::optional { - for (int slot : drafted_slots) { - if (slot >= 0 && slot < static_cast(slot_draft_kv_.size()) && - slot_draft_kv_[static_cast(slot)]) { - draft_kv_reset(*slot_draft_kv_[static_cast(slot)]); - } + auto clean_proposal_lane = [&](size_t i) { + const int slot = inputs[i].slot; + if (slot >= 0 && + slot < static_cast(prepared_chain_drafts_.size())) { + prepared_chain_drafts_[static_cast(slot)].valid = false; } - return std::nullopt; + if (slot >= 0 && slot < static_cast(slot_draft_kv_.size()) && + slot_draft_kv_[static_cast(slot)]) { + draft_kv_reset(*slot_draft_kv_[static_cast(slot)]); + } + }; + auto fail_proposal_lane = [&](size_t i, const char * error) { + clean_proposal_lane(i); + active_admitted[i] = 0; + proposal_errors[i] = error; + std::fprintf(stderr, + "[spec-proposal-failure] request_id=%llu slot=%d error=%s\n", + (unsigned long long)slots_.slot(inputs[i].slot).request_id, + inputs[i].slot, error); + }; + auto retry_proposal_lane = [&](size_t i) { + retried[i] = 1; + clean_proposal_lane(i); + std::vector selected(inputs.size(), 0); + selected[i] = 1; + if (prepare_chain_drafts(inputs, selected, /*force_serial=*/true)) { + return true; + } + fail_proposal_lane( + i, "DSpark proposal preparation failed after clean retry"); + return false; }; std::vector need_prepare(inputs.size(), 0); for (size_t i = 0; i < inputs.size(); ++i) { const StepInput & in = inputs[i]; const bool hard_eligible = chain_spec_input_eligible(in); - if (admitted[i] && !hard_eligible) return proposal_fallback(); if (!admitted[i]) continue; - drafted_slots.push_back(in.slot); + if (!hard_eligible) { + fail_proposal_lane( + i, "sticky speculation request became ineligible"); + continue; + } const PreparedChainDraft & prepared = prepared_chain_drafts_[(size_t)in.slot]; const Qwen35Slot & seq = slots_.slot(in.slot); @@ -1077,42 +1145,75 @@ std::optional Qwen35SeqEngine::step_chain_spec( if (std::any_of( need_prepare.begin(), need_prepare.end(), [](uint8_t value) { return value != 0; }) && - !prepare_chain_drafts(inputs, need_prepare)) { - return proposal_fallback(); + !prepare_chain_drafts( + inputs, need_prepare, /*force_serial=*/false, + /*fail_fast_batch=*/true)) { + // The packed prepare has no target-side effects. Reset its drafter + // state, then retry each affected request once through the serial path + // so one broken lane cannot fail or demote healthy peers. + for (size_t i = 0; i < inputs.size(); ++i) { + if (need_prepare[i]) clean_proposal_lane(i); + } + for (size_t i = 0; i < inputs.size(); ++i) { + if (need_prepare[i] && active_admitted[i]) { + retry_proposal_lane(i); + } + } } - for (size_t i = 0; i < inputs.size(); ++i) { - if (!admitted[i]) continue; + auto take_prepared_proposal = [&](size_t i, Proposal & proposal) { const StepInput & in = inputs[i]; PreparedChainDraft & prepared = prepared_chain_drafts_[(size_t)in.slot]; if (!prepared.valid || - prepared.generated != - slots_.slot(in.slot).generated_tokens() || + prepared.generated != slots_.slot(in.slot).generated_tokens() || prepared.root != in.token || (int)prepared.tokens.size() != T) { - return proposal_fallback(); + return false; } - Proposal proposal; - proposal.input_index = i; - proposal.slot = in.slot; - proposal.root = in.token; - proposal.flat = std::move(prepared.tokens); - proposal.confidence = std::move(prepared.confidence); + Proposal next; + next.input_index = i; + next.slot = in.slot; + next.root = in.token; + next.flat = std::move(prepared.tokens); + next.confidence = std::move(prepared.confidence); prepared.valid = false; - proposal.tree = make_dspark_chain_tree(proposal.flat); - if (proposal.tree.n_nodes + 1 != T) { - return proposal_fallback(); + next.tree = make_dspark_chain_tree(next.flat); + if (next.tree.n_nodes + 1 != T) return false; + proposal = std::move(next); + return true; + }; + + for (size_t i = 0; i < inputs.size(); ++i) { + if (!active_admitted[i]) continue; + Proposal proposal; + if (!take_prepared_proposal(i, proposal)) { + if (!retried[i] && retry_proposal_lane(i) && + take_prepared_proposal(i, proposal)) { + // The clean retry repaired a stale or malformed proposal. + } else { + if (active_admitted[i]) { + fail_proposal_lane( + i, "DSpark proposal remained invalid after clean retry"); + } + continue; + } } proposal_for_input[i] = static_cast(proposals.size()); proposals.push_back(std::move(proposal)); } - if (static_cast(proposals.size()) != spec_count) { - return proposal_fallback(); - } + const int spec_count = static_cast(proposals.size()); + const int tree_bucket = spec_count > 0 + ? chain_decode_bucket_width(spec_count) : 0; + auto lane_disposition = [&](size_t i) { + return chain_lane_disposition( + active_admitted[i] != 0, !proposal_errors[i].empty()); + }; + int replay_total = 0; + if (spec_count > 0) { // Launch 1: scratch-only packed path-tree verification. StepGraph & tree_sg = b_.sg_; int max_prefix = 1; @@ -1223,7 +1324,6 @@ std::optional Qwen35SeqEngine::step_chain_spec( sizeof(int32_t) * posterior.size()); t_posterior_end = timing_clock::now(); - int replay_total = 0; for (int lane = 0; lane < spec_count; ++lane) { Proposal & proposal = proposals[static_cast(lane)]; const int32_t * lane_posterior = @@ -1246,8 +1346,22 @@ std::optional Qwen35SeqEngine::step_chain_spec( : proposal.tree.token_ids[ static_cast(flat_index) - 1]); } + const size_t safe_prefix = chain_min_tokens_safe_prefix( + proposal.path, + slots_.slot(proposal.slot).generated_tokens(), + min_tokens, + [&](int32_t token) { return token_is_eos(token); }); + proposal.path.resize(safe_prefix); + proposal.accepted.resize(safe_prefix); replay_total += static_cast(proposal.path.size()); } + } else { + const auto no_verify = timing_clock::now(); + t_verify_build_start = no_verify; + t_verify_build_end = no_verify; + t_verify_exec_end = no_verify; + t_posterior_end = no_verify; + } // Stage accepted path segments and all non-admitted AR peers. Nothing is // published to slot history until the combined durable graph succeeds. @@ -1301,7 +1415,7 @@ std::optional Qwen35SeqEngine::step_chain_spec( ar_lanes.reserve(inputs.size() - static_cast(spec_count)); std::vector ar_for_input(inputs.size(), -1); for (size_t i = 0; i < inputs.size(); ++i) { - if (admitted[i]) continue; + if (lane_disposition(i) != ChainLaneDisposition::AR) continue; const StepInput & in = inputs[i]; const Qwen35SlotManager::StepAppend app = slots_.append_token(in.slot, in.token); @@ -1330,6 +1444,20 @@ std::optional Qwen35SeqEngine::step_chain_spec( seq_lens_[static_cast(in.slot)] = app.position + 1; max_kv_len = std::max(max_kv_len, app.position + 1); } + const int ar_count = static_cast(ar_lanes.size()); + if (spec_count == 0 && ar_count == 0) { + result.decode.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + DecodeOutput out; + out.slot = inputs[i].slot; + out.failed = true; + out.error = proposal_errors[i].empty() + ? "DSpark proposal lane made no progress" + : proposal_errors[i]; + result.decode.push_back(std::move(out)); + } + return result; + } if (!upload_all_active_block_tables()) { result.error = "DSpark mixed-step block-table refresh failed"; return result; @@ -1338,10 +1466,10 @@ std::optional Qwen35SeqEngine::step_chain_spec( // Launch 2: accepted path segments + compact AR rows in the same builder // combination already used by mixed prefill/decode. - const int ar_count = static_cast(ar_lanes.size()); const int ar_bucket = chain_decode_bucket_width(ar_count); const int n_total = replay_total + ar_bucket; const int gather_rows = spec_count + ar_bucket; + const bool has_replay = replay_total > 0; StepGraph & durable_sg = b_.sg_; if (!build_target_step( durable_sg, b_.w_, b_.cache_, b_.target_backend_, @@ -1352,8 +1480,11 @@ std::optional Qwen35SeqEngine::step_chain_spec( static_cast(replay_segments.size()), gather_rows, ar_bucket > 0) || !durable_sg.kv_write_rows || !durable_sg.target_feat_rows || - !durable_sg.paged_query_seq_ids || - !durable_sg.paged_query_positions || + (has_replay && + (!durable_sg.paged_query_seq_ids || + !durable_sg.paged_query_positions)) || + (ar_bucket > 0 && + (!durable_sg.active_slot_ids || !durable_sg.state_slot_ids)) || !durable_sg.logits_row_indices || !durable_sg.argmax_tokens) { result.error = "DSpark mixed commit/AR graph build failed"; return result; @@ -1436,12 +1567,14 @@ std::optional Qwen35SeqEngine::step_chain_spec( for (int lane = 0; lane < ar_bucket; ++lane) { logits_rows_.push_back(replay_total + lane); } - ggml_backend_tensor_set( - durable_sg.paged_query_seq_ids, query_slot_ids_.data(), 0, - sizeof(int32_t) * query_slot_ids_.size()); - ggml_backend_tensor_set( - durable_sg.paged_query_positions, query_positions_.data(), 0, - sizeof(int32_t) * query_positions_.size()); + if (has_replay) { + ggml_backend_tensor_set( + durable_sg.paged_query_seq_ids, query_slot_ids_.data(), 0, + sizeof(int32_t) * query_slot_ids_.size()); + ggml_backend_tensor_set( + durable_sg.paged_query_positions, query_positions_.data(), 0, + sizeof(int32_t) * query_positions_.size()); + } ggml_backend_tensor_set( durable_sg.logits_row_indices, logits_rows_.data(), 0, sizeof(int32_t) * logits_rows_.size()); @@ -1512,7 +1645,11 @@ std::optional Qwen35SeqEngine::step_chain_spec( std::vector write_slots; write_slots.reserve(inputs.size()); - for (const StepInput & in : inputs) write_slots.push_back(in.slot); + for (size_t i = 0; i < inputs.size(); ++i) { + if (chain_lane_executes(lane_disposition(i))) { + write_slots.push_back(inputs[i].slot); + } + } if (!commit_residency_writes(write_slots)) { result.error = "DSpark mixed-step KV write commit failed"; return result; @@ -1520,8 +1657,10 @@ std::optional Qwen35SeqEngine::step_chain_spec( // Publish the fed root/path before sampling the next token, matching the // ordinary AR path's penalty history, RNG, and min-token-floor semantics. - for (const StepInput & in : inputs) { - slots_.commit_step(in.slot); + for (size_t i = 0; i < inputs.size(); ++i) { + if (chain_lane_executes(lane_disposition(i))) { + slots_.commit_step(inputs[i].slot); + } } for (int lane = 0; lane < spec_count; ++lane) { Proposal & proposal = proposals[static_cast(lane)]; @@ -1546,9 +1685,10 @@ std::optional Qwen35SeqEngine::step_chain_spec( } t_sample_end = timing_clock::now(); - for (const StepInput & in : inputs) { + for (size_t i = 0; i < inputs.size(); ++i) { + if (!chain_lane_executes(lane_disposition(i))) continue; std::string reselect_error; - if (!maybe_reselect_residency(in.slot, reselect_error)) { + if (!maybe_reselect_residency(inputs[i].slot, reselect_error)) { result.error = reselect_error.empty() ? "KVFlash reselect failed" : reselect_error; return result; @@ -1559,7 +1699,14 @@ std::optional Qwen35SeqEngine::step_chain_spec( for (size_t i = 0; i < inputs.size(); ++i) { DecodeOutput out; out.slot = inputs[i].slot; - if (admitted[i]) { + const ChainLaneDisposition disposition = lane_disposition(i); + if (disposition == ChainLaneDisposition::Failed) { + out.failed = true; + out.error = proposal_errors[i]; + result.decode.push_back(std::move(out)); + continue; + } + if (disposition == ChainLaneDisposition::Speculation) { Proposal & proposal = proposals[static_cast(proposal_for_input[i])]; out.token = proposal.pending; @@ -1606,7 +1753,7 @@ std::optional Qwen35SeqEngine::step_chain_spec( "\"finish_us\":%.1f,\"total_us\":%.1f," "\"accepted_tokens\":%d,\"emitted_tokens\":%d," "\"target_forwards\":%d}\n", - (int)inputs.size(), spec_count, + spec_count + ar_count, spec_count, tree_bucket, T * tree_bucket, replay_total, ar_count, ar_bucket, max_kv_len, round_draft_us_, round_draft_lanes_, @@ -2328,7 +2475,6 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { round_draft_us_ = 0.0; round_draft_lanes_ = 0; const auto chain_started = decode_round_started; - const bool chain_execution_available = plan.prefills.empty(); std::vector admitted(inputs.size(), 0); SpecPlan gate_plan; bool have_gate_plan = false; @@ -2399,64 +2545,142 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { return gate_plan.valid; }; - // Bootstrap every cold eligible request together. This is the - // request's one-time confidence initialization, not eager drafting - // on every later k=0 round. - if (use_confidence && !gate_plan.bootstrap_slots.empty()) { + auto reset_evaluation_lane = [&](int slot) { + if (slot >= 0 && + slot < (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)slot] = {}; + } + if (slot >= 0 && slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)slot]); + } + if (slot >= 0 && + slot < (int)last_survival_score_.size()) { + last_survival_score_[(size_t)slot] = + std::numeric_limits::quiet_NaN(); + } + }; + auto commit_evaluation_fallback = + [&](const SpecPendingEvaluation & evaluation) { + reset_evaluation_lane(evaluation.slot); + if (speculation_gate_->commit_evaluation_fallback_ar( + evaluation.request_id)) { + log_spec_evaluation_fallback( + evaluation.request_id, evaluation.slot); + } + }; + + // Resolve every one-time evaluation action before one immediate + // replan. A packed bootstrap failure is retried once per lane so + // one broken request falls back to sticky AR without poisoning a + // healthy scored peer or the cohort. + if (!gate_plan.pending_evaluations.empty()) { + std::vector score_evaluations; std::vector bootstrap(inputs.size(), 0); - for (int slot : gate_plan.bootstrap_slots) { + for (const SpecPendingEvaluation & evaluation : + gate_plan.pending_evaluations) { + if (evaluation.action == + SpecEvaluationAction::FallbackAR) { + commit_evaluation_fallback(evaluation); + continue; + } + score_evaluations.push_back(evaluation); for (size_t i = 0; i < inputs.size(); ++i) { - if (inputs[i].slot == slot) bootstrap[i] = 1; + if (inputs[i].slot == evaluation.slot) { + bootstrap[i] = 1; + } } } - if (!prepare_chain_drafts(inputs, bootstrap)) { - return fail_step( - "adaptive confidence bootstrap failed"); - } - for (int slot : gate_plan.bootstrap_slots) { - if (slot < 0 || - slot >= (int)last_survival_score_.size() || - !std::isfinite( - last_survival_score_[(size_t)slot])) { - return fail_step( - "adaptive confidence bootstrap produced no score"); + + if (use_confidence && !score_evaluations.empty()) { + const bool batch_scored = prepare_chain_drafts( + inputs, bootstrap, /*force_serial=*/false, + /*fail_fast_batch=*/true); + if (!batch_scored) { + for (const SpecPendingEvaluation & evaluation : + score_evaluations) { + reset_evaluation_lane(evaluation.slot); + } + for (const SpecPendingEvaluation & evaluation : + score_evaluations) { + std::vector one(inputs.size(), 0); + for (size_t i = 0; i < inputs.size(); ++i) { + if (inputs[i].slot == evaluation.slot) { + one[i] = 1; + } + } + const bool lane_scored = prepare_chain_drafts( + inputs, one, /*force_serial=*/true); + if (!lane_scored || evaluation.slot < 0 || + evaluation.slot >= + (int)last_survival_score_.size() || + !std::isfinite(last_survival_score_[ + (size_t)evaluation.slot])) { + commit_evaluation_fallback(evaluation); + } + } + } else { + for (const SpecPendingEvaluation & evaluation : + score_evaluations) { + if (evaluation.slot < 0 || + evaluation.slot >= + (int)last_survival_score_.size() || + !std::isfinite(last_survival_score_[ + (size_t)evaluation.slot])) { + commit_evaluation_fallback(evaluation); + } + } + } + } else { + for (const SpecPendingEvaluation & evaluation : + score_evaluations) { + commit_evaluation_fallback(evaluation); } } + if (!replan_with_published_confidence()) { return fail_step(gate_plan.error.empty() ? "adaptive speculation gate failed" : gate_plan.error); } - if (!gate_plan.bootstrap_slots.empty() || + // This guard converts any unexpectedly unpublished score into + // the same request-local fallback rather than repeating the + // cold evaluation forever or failing a mixed cohort. + if (!gate_plan.pending_evaluations.empty()) { + const std::vector unresolved = + gate_plan.pending_evaluations; + for (const SpecPendingEvaluation & evaluation : unresolved) { + commit_evaluation_fallback(evaluation); + } + if (!replan_with_published_confidence()) { + return fail_step(gate_plan.error.empty() + ? "adaptive speculation gate failed" + : gate_plan.error); + } + } + if (!gate_plan.pending_evaluations.empty() || !gate_plan.decisions_committed) { return fail_step( - "adaptive confidence bootstrap did not commit modes"); + "adaptive activation state did not commit modes"); } } log_spec_activations(gate_plan, *speculation_gate_); - // A temporary round constraint never changes the sticky request - // mode. Execute the exact selected cohort only when every selected - // lane is runnable; otherwise this mixed/floor-constrained round - // uses packed AR and is excluded from gate cost feedback. - bool execution_matches_plan = chain_execution_available; + // Preserve every sticky Spec admission. Prefills are deferred + // below, the min-token floor is enforced inside the speculative + // path, and a later capability invariant failure becomes a + // lane-local error in step_chain_spec rather than an AR step. for (int slot : gate_plan.admitted_slots) { - auto input = std::find_if( - inputs.begin(), inputs.end(), - [&](const StepInput & item) { - return item.slot == slot; - }); - if (input == inputs.end() || - !chain_spec_input_eligible(*input)) { - execution_matches_plan = false; - break; - } - } - if (execution_matches_plan) { - for (int slot : gate_plan.admitted_slots) { - for (size_t i = 0; i < inputs.size(); ++i) { - if (inputs[i].slot == slot) admitted[i] = 1; + bool found = false; + for (size_t i = 0; i < inputs.size(); ++i) { + if (inputs[i].slot == slot) { + admitted[i] = 1; + found = true; } } + if (!found) { + return fail_step( + "speculation gate admitted a missing decode lane"); + } } } else { // Explicit forced speculation remains usable without a cost @@ -2474,9 +2698,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { return fail_step( "adaptive speculation gate unavailable"); } - admitted[i] = chain_execution_available && - chain_spec_input_eligible(inputs[i]) && - policy == SpeculationPolicy::Always; + admitted[i] = policy == SpeculationPolicy::Always; } } @@ -2484,23 +2706,46 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { admitted.begin(), admitted.end(), [](uint8_t value) { return value != 0; }); if (any_admitted) { - std::optional speculative = - step_chain_spec(plan, admitted, decode_round_started); + StepPlan chain_plan = plan; + chain_plan.prefills.clear(); + StepResult speculative = + step_chain_spec(chain_plan, admitted, decode_round_started); const double measured_us = std::chrono::duration( std::chrono::steady_clock::now() - chain_started).count(); - const bool spec_completed = - speculative && speculative->error.empty(); + const bool spec_completed = speculative.error.empty(); + bool proposal_failed = false; + if (spec_completed) { + for (size_t i = 0; i < inputs.size(); ++i) { + if (!admitted[i]) continue; + const auto output = std::find_if( + speculative.decode.begin(), speculative.decode.end(), + [&](const DecodeOutput & item) { + return item.slot == inputs[i].slot; + }); + if (output == speculative.decode.end() || output->failed) { + proposal_failed = true; + break; + } + } + } + if (spec_completed && !plan.prefills.empty()) { + speculative.prefills.reserve(plan.prefills.size()); + for (const PrefillSlice & slice : plan.prefills) { + speculative.prefills.push_back({ + slice.slot, PrefillOutput::Status::deferred, -1, {}}); + } + } const bool cost_sample_valid = - have_gate_plan && spec_completed && + have_gate_plan && spec_completed && !proposal_failed && round_draft_lanes_ == gate_plan.draft_lanes; if (cost_sample_valid) { speculation_gate_->observe_cost(gate_plan, measured_us); } if (have_gate_plan && spec_gate_debug_enabled() && - spec_completed) { + spec_completed && !proposal_failed) { const double realized_tokens = spec_completed - ? confidence_realized_tokens(gate_plan, *speculative) + ? confidence_realized_tokens(gate_plan, speculative) : std::numeric_limits::quiet_NaN(); log_spec_gate_plan( gate_plan, speculation_gate_->calibration_scale(), @@ -2510,9 +2755,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { ? measured_us : std::numeric_limits::quiet_NaN()); } - if (speculative) return std::move(*speculative); - // Proposal setup failed before target/cache mutation. Preserve - // service through the ordinary packed AR path this iteration. + return speculative; } if (!any_admitted && have_gate_plan) { if (gate_plan.admitted_count == 0 && plan.prefills.empty()) { diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 9b12edf91..b0d5f8751 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -146,17 +146,20 @@ class Qwen35SeqEngine final : public SeqEngine { }; bool prepare_chain_drafts( const std::vector & inputs, - const std::vector & selected); + const std::vector & selected, + bool force_serial = false, + bool fail_fast_batch = false); bool batched_drafting_enabled() const; bool confidence_scoring_enabled() const; // DFLASH_STEP_TIMING=1 emits one [step-timing] JSON line per decode // round attributing wall time to draft, verify, readback, CPU commit, // replay, and packed-AR phases. Diagnostic only; off by default. static bool step_timing_enabled(); - // nullopt means proposal setup failed before target/cache mutation and the - // caller may safely use the ordinary packed AR path for this iteration. + // DDTree preserves its legacy best-effort AR fallback. DSpark chain + // proposal failures are instead returned as lane-local DecodeOutput + // failures so a sticky speculation decision can never execute as AR. std::optional step_ddtree(const StepPlan & plan); - std::optional step_chain_spec( + StepResult step_chain_spec( const StepPlan & plan, const std::vector & admitted, std::chrono::steady_clock::time_point round_started); diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 9c595ff17..a11098fae 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -739,9 +739,6 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { engine.step_plan_limits((int)step_plan.decode.size()); step_plan.prefills = plan_prefill_slices( prefill_candidates, step_limits, prefill_round_robin_start); - if (!prefill_candidates.empty()) { - ++prefill_round_robin_start; - } SeqEngine::StepResult step_result = engine.step(step_plan); const std::string protocol_error = @@ -794,6 +791,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { }); } using PrefillStatus = SeqEngine::PrefillOutput::Status; + const bool prefill_progressed = + prefill_result_made_progress(step_result); for (const auto & out : step_result.prefills) { if (out.slot < 0 || out.slot >= n_slots) continue; SchedSlot & s = slots[(size_t)out.slot]; @@ -804,6 +803,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.finished = true; continue; } + if (out.status == PrefillStatus::deferred) { + // The exclusive decode graph made no prompt progress. Leave + // the slot untouched so normal FIFO planning retries it after + // the current decode wave drains. + continue; + } if (out.status == PrefillStatus::completed) { s.prefilling = false; publish_live_count(); @@ -817,6 +822,9 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { continue; } } + // Deferred means the engine intentionally left the selected slices + // untouched. Keep allocation fairness stable across such rounds. + if (prefill_progressed) ++prefill_round_robin_start; // Phase 4 — Non-blocking flush of every live slot's chunks. Progress // resets the stall clock; a reader that makes no progress for 30 s // or lets the buffer hit the cap is dropped (its slot retires). diff --git a/server/test/seq_engine_contract.h b/server/test/seq_engine_contract.h index 0361a893c..17073646c 100644 --- a/server/test/seq_engine_contract.h +++ b/server/test/seq_engine_contract.h @@ -151,7 +151,7 @@ inline std::vector check_seq_engine_contract(SeqEngine & engine) { require(remaining[(size_t)output.slot] > 1, "prefill reported advanced for its final token"); --remaining[(size_t)output.slot]; - } else { + } else if (output.status == PrefillStatus::completed) { require(remaining[(size_t)output.slot] == 1, "prefill reported completion before its final token"); remaining[(size_t)output.slot] = 0; diff --git a/server/test/test_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp index 33c011f1c..88bd6ccf7 100644 --- a/server/test/test_chain_spec_shapes.cpp +++ b/server/test/test_chain_spec_shapes.cpp @@ -59,6 +59,44 @@ int main() { CHECK(all_spec.ar_bucket == 0); CHECK(all_spec.commit_rows == 6); + const ChainLaunchShape ar_after_spec_failures = chain_launch_shape( + {0, 0}, {0, 0}, 16); + CHECK(ar_after_spec_failures.spec_lanes == 0); + CHECK(ar_after_spec_failures.tree_bucket == 0); + CHECK(ar_after_spec_failures.tree_rows == 0); + CHECK(ar_after_spec_failures.ar_lanes == 2); + CHECK(ar_after_spec_failures.ar_bucket == 2); + CHECK(ar_after_spec_failures.commit_rows == 2); + + CHECK(chain_lane_disposition(false, false) == + ChainLaneDisposition::AR); + CHECK(chain_lane_disposition(true, false) == + ChainLaneDisposition::Speculation); + CHECK(chain_lane_disposition(true, true) == + ChainLaneDisposition::Failed); + CHECK(chain_lane_disposition(false, true) == + ChainLaneDisposition::Failed); + CHECK(chain_lane_executes(ChainLaneDisposition::AR)); + CHECK(chain_lane_executes(ChainLaneDisposition::Speculation)); + CHECK(!chain_lane_executes(ChainLaneDisposition::Failed)); + + const auto eos = [](int32_t token) { return token == 2; }; + const std::vector eos_first_child = {10, 2, 11}; + CHECK(chain_min_tokens_safe_prefix( + eos_first_child, 0, 3, eos) == 1); + CHECK(chain_min_tokens_safe_prefix( + eos_first_child, 2, 3, eos) == 2); + const std::vector eos_second_child = {10, 11, 2, 12}; + CHECK(chain_min_tokens_safe_prefix( + eos_second_child, 0, 3, eos) == 2); + CHECK(chain_min_tokens_safe_prefix( + eos_second_child, 1, 3, eos) == 3); + const std::vector eos_root = {2, 11, 12}; + CHECK(chain_min_tokens_safe_prefix( + eos_root, 0, 3, eos) == eos_root.size()); + CHECK(chain_min_tokens_safe_prefix( + eos_first_child, 0, 0, eos) == 2); + std::printf("chain spec shape tests passed: %d checks\n", g_checks); return 0; } diff --git a/server/test/test_seq_batch_plan.cpp b/server/test/test_seq_batch_plan.cpp index 2860cc695..dff2122fa 100644 --- a/server/test/test_seq_batch_plan.cpp +++ b/server/test/test_seq_batch_plan.cpp @@ -126,6 +126,7 @@ int main() { good.prefills.push_back({ 1, SeqEngine::PrefillOutput::Status::advanced, -1, {}}); CHECK(validate_step_result(work, good, 2).empty()); + CHECK(prefill_result_made_progress(good)); SeqEngine::StepResult burst = good; burst.decode[0].committed_tokens = {8, 9, 10}; @@ -189,6 +190,23 @@ int main() { complete.prefills[0] = { 1, SeqEngine::PrefillOutput::Status::completed, 12, {}}; CHECK(validate_step_result(work, complete, 2).empty()); + CHECK(prefill_result_made_progress(complete)); + + SeqEngine::StepResult deferred = good; + deferred.prefills[0] = { + 1, SeqEngine::PrefillOutput::Status::deferred, -1, {}}; + CHECK(validate_step_result(work, deferred, 2).empty()); + CHECK(!prefill_result_made_progress(deferred)); + deferred.prefills[0].token = 12; + CHECK(!validate_step_result(work, deferred, 2).empty()); + deferred.prefills[0].token = -1; + deferred.prefills[0].error = "not an error"; + CHECK(!validate_step_result(work, deferred, 2).empty()); + SeqEngine::StepPlan idle_prefill = work; + idle_prefill.decode.clear(); + deferred.prefills[0].error.clear(); + deferred.decode.clear(); + CHECK(!validate_step_result(idle_prefill, deferred, 2).empty()); SeqEngine::StepResult missing_decode = good; missing_decode.decode.clear(); @@ -208,6 +226,7 @@ int main() { prefill_failure.prefills[0] = { 1, SeqEngine::PrefillOutput::Status::failed, -1, "prefill failed"}; CHECK(validate_step_result(work, prefill_failure, 2).empty()); + CHECK(!prefill_result_made_progress(prefill_failure)); SeqEngine::StepResult bad_row_failure = prefill_failure; bad_row_failure.prefills.back().error.clear(); diff --git a/server/test/test_seq_engine_contract.cpp b/server/test/test_seq_engine_contract.cpp index 63a62d18a..373aac441 100644 --- a/server/test/test_seq_engine_contract.cpp +++ b/server/test/test_seq_engine_contract.cpp @@ -22,6 +22,7 @@ struct Faults { bool lose_other_pending = false; bool overconsume_prefill = false; bool drop_second_completion = false; + bool defer_first_mixed_prefill = false; bool retire_leaks = false; bool burst_when_speculation_disabled = false; }; @@ -127,6 +128,14 @@ class FakeSeqEngine final : public SeqEngine { if (!slot.active || !slot.prefilling || slot.remaining <= 0) { continue; } + if (faults_.defer_first_mixed_prefill && + !plan.decode.empty() && !deferred_mixed_prefill_) { + deferred_mixed_prefill_ = true; + result.prefills.push_back({ + slice.slot, PrefillOutput::Status::deferred, -1, {}, + }); + continue; + } int consumed = std::min(slice.max_tokens, slot.remaining); if (faults_.overconsume_prefill) consumed = slice.max_tokens + 1; slot.remaining -= consumed; @@ -236,6 +245,7 @@ class FakeSeqEngine final : public SeqEngine { std::vector slots_; Faults faults_; FakeCapabilities capabilities_; + bool deferred_mixed_prefill_ = false; }; static void print_violations(const char * label, @@ -284,6 +294,17 @@ int main() { CHECK(violations.empty()); } + { + Faults faults; + faults.defer_first_mixed_prefill = true; + FakeSeqEngine engine(2, faults); + const auto violations = check_seq_engine_contract(engine); + if (!violations.empty()) { + print_violations("conforming-deferred-prefill", violations); + } + CHECK(violations.empty()); + } + struct Case { const char * label; bool Faults::*fault; diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index 411e5fcd7..efa9795ab 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -68,7 +68,7 @@ int main() { CHECK(costly.initial_confidence(1) == 4.0); plan = costly.plan(2, { candidate(1, 0, NAN), candidate(2, 1, NAN)}, 2); - CHECK(plan.bootstrap_slots.empty()); + CHECK(plan.pending_evaluations.empty()); CHECK(plan.ordered.empty()); CHECK(plan.admitted_count == 0); CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); @@ -114,8 +114,11 @@ int main() { CHECK(plan.admitted_count == 0); CHECK(plan.ordered.empty()); CHECK(plan.unavailable_count == 1); - CHECK((plan.bootstrap_request_ids == std::vector{20})); - CHECK((plan.bootstrap_slots == std::vector{0})); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].request_id == 20); + CHECK(plan.pending_evaluations[0].slot == 0); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::Score); CHECK(bootstrap.decision(20) == SpecDecision::Undecided); CHECK(bootstrap.decision(21) == SpecDecision::Undecided); CHECK(bootstrap.initial_confidence(21) == 4.0); @@ -124,7 +127,7 @@ int main() { candidate(20, 0, 1.0), candidate(21, 1, NAN)}, 2); CHECK(plan.valid); CHECK(plan.decisions_committed); - CHECK(plan.bootstrap_slots.empty()); + CHECK(plan.pending_evaluations.empty()); CHECK(plan.admitted_count == 1); CHECK((plan.admitted_request_ids == std::vector{21})); CHECK(bootstrap.decision(20) == SpecDecision::AR); @@ -135,7 +138,7 @@ int main() { plan = bootstrap.plan(2, { candidate(20, 0, 4.0), candidate(21, 1, 1.0)}, 2); CHECK(plan.decisions_committed); - CHECK(plan.bootstrap_slots.empty()); + CHECK(plan.pending_evaluations.empty()); CHECK(plan.admitted_count == 1); CHECK((plan.admitted_request_ids == std::vector{21})); CHECK(plan.ordered.size() == 1); @@ -146,14 +149,19 @@ int main() { // Request-lifetime scoreability is separate from permanent executor // support. Unsupported requests still bootstrap and retain a confidence, - // then commit directly to AR. Missing mandatory scoring is an explicit - // contract error and never silently becomes an AR decision. + // then commit directly to AR. A request that cannot evaluate confidence + // receives an explicit failed-evaluation action and sticky AR with no + // synthetic score. SpeculationGate support(crossover, geometry(), 4); plan = support.plan(1, { candidate(30, 0, NAN, SpeculationPolicy::Adaptive, true, false)}, 1); CHECK(plan.valid); CHECK(!plan.decisions_committed); - CHECK((plan.bootstrap_slots == std::vector{0})); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].request_id == 30); + CHECK(plan.pending_evaluations[0].slot == 0); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::Score); CHECK(support.decision(30) == SpecDecision::Undecided); plan = support.plan(1, { candidate(30, 0, 4.0, SpeculationPolicy::Adaptive, true, false)}, 1); @@ -167,9 +175,54 @@ int main() { CHECK(support.initial_confidence(30) == 4.0); plan = support.plan(1, { candidate(31, 0, NAN, SpeculationPolicy::Adaptive, false, false)}, 1); - CHECK(!plan.valid); - CHECK(!plan.error.empty()); + CHECK(plan.valid); + CHECK(!plan.decisions_committed); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].request_id == 31); + CHECK(plan.pending_evaluations[0].slot == 0); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::FallbackAR); CHECK(support.decision(31) == SpecDecision::Undecided); + CHECK(support.commit_evaluation_fallback_ar(31)); + CHECK(!support.commit_evaluation_fallback_ar(31)); + CHECK(support.decision(31) == SpecDecision::AR); + CHECK(support.confidence_evaluation_failed(31)); + CHECK(!support.has_confidence(31)); + CHECK(std::isnan(support.initial_confidence(31))); + plan = support.plan(1, {candidate(31, 0, 4.0)}, 1); + CHECK(plan.valid); + CHECK(plan.decisions_committed); + CHECK(plan.pending_evaluations.empty()); + CHECK(plan.ordered.empty()); + CHECK(support.decision(31) == SpecDecision::AR); + + // A failed lane and an already-scored healthy lane activate atomically: + // the former becomes sticky AR while the latter still receives its + // score-based mode on the immediate replan. + SpeculationGate mixed_activation(crossover, geometry(), 4); + plan = mixed_activation.plan(2, { + candidate(32, 0, NAN, SpeculationPolicy::Adaptive, false, false), + candidate(33, 1, 4.0)}, 2); + CHECK(plan.valid); + CHECK(!plan.decisions_committed); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].request_id == 32); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::FallbackAR); + CHECK(mixed_activation.decision(32) == SpecDecision::Undecided); + CHECK(mixed_activation.decision(33) == SpecDecision::Undecided); + CHECK(mixed_activation.commit_evaluation_fallback_ar(32)); + plan = mixed_activation.plan(2, { + candidate(32, 0, NAN, SpeculationPolicy::Adaptive, false, false), + candidate(33, 1, NAN)}, 2); + CHECK(plan.valid); + CHECK(plan.decisions_committed); + CHECK(plan.pending_evaluations.empty()); + CHECK(mixed_activation.decision(32) == SpecDecision::AR); + CHECK(mixed_activation.confidence_evaluation_failed(32)); + CHECK(mixed_activation.decision(33) == SpecDecision::Speculation); + CHECK(mixed_activation.initial_confidence(33) == 4.0); + CHECK((plan.admitted_request_ids == std::vector{33})); // Explicit Always/Never are configured execution policies, not adaptive // activation decisions. Always remains a non-negotiable baseline. @@ -191,6 +244,13 @@ int main() { candidate(44, 1, 4.0, SpeculationPolicy::Always)}, 1); CHECK(!plan.valid); CHECK(!plan.error.empty()); + plan = policies.plan(1, { + candidate(46, 0, NAN, SpeculationPolicy::Always, + /*scoreable=*/false, /*can_speculate=*/false)}, 1); + CHECK(plan.valid); + CHECK(plan.admitted_count == 1); + CHECK(plan.ordered.front().forced); + CHECK(plan.admitted_request_ids.front() == 46); SpeculationGate never(constant_costs(1.0, 2.0, 1.0), geometry(), 4); plan = never.plan(1, { @@ -294,7 +354,10 @@ int main() { CHECK(isolated.calibration_observations() == 0); CHECK(!isolated.has_confidence(800)); plan = isolated.plan(1, {candidate(800, 0, NAN)}, 1); - CHECK((plan.bootstrap_slots == std::vector{0})); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].slot == 0); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::Score); CHECK(isolated.decision(800) == SpecDecision::Undecided); plan = isolated.plan(1, { @@ -307,7 +370,7 @@ int main() { CHECK(std::abs(isolated.calibration_scale(801) - 0.5) < 1e-12); CHECK(isolated.decision(801) == SpecDecision::Undecided); plan = isolated.plan(1, {candidate(801, 0, NAN)}, 1); - CHECK(plan.bootstrap_slots.empty()); + CHECK(plan.pending_evaluations.empty()); CHECK(isolated.decision(801) == SpecDecision::Speculation); // Forget is the only adaptive state reset. It drops the sticky decision, @@ -322,7 +385,10 @@ int main() { CHECK(calibrated.calibration_scale() == global_before_forget); plan = calibrated.plan(1, {candidate(700, 0, NAN)}, 1); CHECK(!plan.decisions_committed); - CHECK((plan.bootstrap_slots == std::vector{0})); + CHECK(plan.pending_evaluations.size() == 1); + CHECK(plan.pending_evaluations[0].slot == 0); + CHECK(plan.pending_evaluations[0].action == + SpecEvaluationAction::Score); // Calibration clamps are active on the first forced/probe observation. SpeculationGate lower_bound(constant_costs(1.0, 1.0, 1.0), From 1626eeb02d7f8002ddd07ae6f47c4131a027d986 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 16:22:45 +0000 Subject: [PATCH 34/42] refactor(concurrency): isolate adaptive request scoring --- .../QWEN38_DSPARK_ADAPTIVE_SELECTION.md | 9 +- .../benchmarks/concurrency/FEATURE_MATRIX.md | 6 +- .../concurrency/analyze_gate_decisions.py | 16 +- .../concurrency/test_feature_tools.py | 20 +-- .../src/common/concurrency/speculation_gate.h | 170 +++--------------- .../qwen35/concurrency/qwen35_seq_engine.cpp | 44 ++--- server/test/test_speculation_gate.cpp | 132 +++----------- 7 files changed, 92 insertions(+), 305 deletions(-) diff --git a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md index ffe6cb453..45cff9de9 100644 --- a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md +++ b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md @@ -104,8 +104,11 @@ included. Adaptive mode performs one evaluation per request at its first target-decode boundary. A successful evaluation produces the request's confidence score and commits either AR or speculation through EOS/max tokens; a failed evaluation commits AR without inventing a score. The request is never -evaluated again. Later execution observations calibrate activations -for future requests, not the mode of a request already in flight. +evaluated again. Every request is ranked independently from its own initial +confidence, an immutable offline yield scale, and the measured cost for the +current graph shape. Accepted-token history from this or any prior request is +never used. The only runtime EWMA tracks hardware latency by graph shape, not +request content or user/chat identity. The prefill logits produce the first sampled output token before a target-decode step exists. The one-time draft and activation therefore finish @@ -160,7 +163,7 @@ Read the report in this order: never use the dense prompt label as the final activation answer. 3. Inspect `Activation outcome against matched pure AR`. It compares each `(live, k, path)` shape to pure AR at the same live concurrency. -4. Inspect `Gate prediction calibration` for predicted-versus-realized +4. Inspect `Initial prediction accuracy` for predicted-versus-realized goodput and realized-versus-AR regret. 5. Use phase attribution to choose the next optimization: unexpected draft work on k=0, verify, or the structural replay forward. diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 6db4a0d64..64b5c6798 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -27,8 +27,10 @@ The default fresh-process matrix is: proposal is reused immediately; a request that chooses AR never drafts again. A target-decode step exists only after prefill has produced the first sampled token, so one-token requests that retire directly from prefill never enter - adaptive activation. Later execution evidence calibrates activations for new - requests and never re-evaluates an active request. The last arm sets + adaptive activation. Each request is ranked only from its own initial score; + accepted-token history and user/chat identity are not inputs. A shape-keyed + timing EWMA may refine hardware cost estimates without carrying content + history. The last arm sets `DFLASH_SPEC_CONFIDENCE=0`, suppressing activation scoring so it remains a draft-free AR ablation. There is no eager-draft matrix axis. - A Spec request runs the speculative executor immediately even below diff --git a/harness/benchmarks/concurrency/analyze_gate_decisions.py b/harness/benchmarks/concurrency/analyze_gate_decisions.py index 7995ac79d..1e764fc2c 100644 --- a/harness/benchmarks/concurrency/analyze_gate_decisions.py +++ b/harness/benchmarks/concurrency/analyze_gate_decisions.py @@ -97,13 +97,13 @@ def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: raise ValueError( f"{path}:{line_no}: spec-activation fallback_reason is required" ) - for key in ("initial_confidence", "calibrated_yield"): + for key in ("initial_confidence", "expected_yield"): if key not in row: raise ValueError( f"{path}:{line_no}: spec-activation {key} is required" ) if evaluation == "scored": - for key in ("initial_confidence", "calibrated_yield"): + for key in ("initial_confidence", "expected_yield"): value = row.get(key) if ( type(value) not in (int, float) @@ -121,7 +121,7 @@ def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: ) return if row.get("initial_confidence") is not None or ( - row.get("calibrated_yield") is not None + row.get("expected_yield") is not None ): raise ValueError( f"{path}:{line_no}: failed spec-activation scores must be null" @@ -622,8 +622,8 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A activation.get("initial_confidence") if activation is not None else None ), - "calibrated_yield": ( - activation.get("calibrated_yield") + "expected_yield": ( + activation.get("expected_yield") if activation is not None else None ), "activation_decision": ( @@ -787,7 +787,7 @@ def compare_prompts(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: "activation_fallback_reason" ), "initial_confidence": request.get("initial_confidence"), - "calibrated_yield": request.get("calibrated_yield"), + "expected_yield": request.get("expected_yield"), "mean_confidence_yield": request.get( "mean_confidence_yield" ), @@ -1039,7 +1039,7 @@ def render_markdown(report: dict[str, Any]) -> str: if gate_rows: lines += [ "", - "## Gate prediction calibration", + "## Initial prediction accuracy", "", "| Workload | C | Variant | k | Rounds | Pred tok/s | " "Realized tok/s | Realized/pred | Realized/AR | Pred cost us | " @@ -1072,7 +1072,7 @@ def build_report( ] activation_comparisons = compare_activation_shapes(cases) return { - "schema_version": 5, + "schema_version": 6, "cases": cases, "prompt_comparisons": compare_prompts(cases), "activation_comparisons": activation_comparisons, diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index 29414cf29..12f896543 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -375,7 +375,7 @@ def test_measured_step_timing_is_joined_and_summarized(self) -> None: } activation = { "request_id": 7, "slot": 0, - "initial_confidence": 3.25, "calibrated_yield": 1.75, + "initial_confidence": 3.25, "expected_yield": 1.75, "evaluation": "scored", "fallback_reason": None, "decision": "ar", } @@ -407,7 +407,7 @@ def test_measured_step_timing_is_joined_and_summarized(self) -> None: request = report["requests"][0] self.assertEqual(request["activation_slot"], 0) self.assertEqual(request["initial_confidence"], 3.25) - self.assertEqual(request["calibrated_yield"], 1.75) + self.assertEqual(request["expected_yield"], 1.75) self.assertEqual(request["activation_decision"], "ar") self.assertEqual(request["activation_evaluation"], "scored") self.assertIsNone(request["activation_fallback_reason"]) @@ -415,13 +415,13 @@ def test_measured_step_timing_is_joined_and_summarized(self) -> None: def test_adaptive_on_activation_proof_fails_closed(self) -> None: activation_7 = { "request_id": 7, "slot": 0, - "initial_confidence": 2.0, "calibrated_yield": 1.25, + "initial_confidence": 2.0, "expected_yield": 1.25, "evaluation": "scored", "fallback_reason": None, "decision": "speculation", } activation_8 = { "request_id": 8, "slot": 1, - "initial_confidence": 1.0, "calibrated_yield": 1.0, + "initial_confidence": 1.0, "expected_yield": 1.0, "evaluation": "scored", "fallback_reason": None, "decision": "ar", } @@ -462,7 +462,7 @@ def test_confidence_off_is_exempt_from_activation_coverage(self) -> None: def test_adaptive_on_activation_proof_enforces_sticky_execution(self) -> None: activation = { "request_id": 7, "slot": 0, - "initial_confidence": 2.0, "calibrated_yield": 1.5, + "initial_confidence": 2.0, "expected_yield": 1.5, "evaluation": "scored", "fallback_reason": None, "decision": "speculation", } @@ -503,14 +503,14 @@ def test_failed_evaluation_activation_is_request_local_and_explicit( ) -> None: failed = { "request_id": 7, "slot": 0, - "initial_confidence": None, "calibrated_yield": None, + "initial_confidence": None, "expected_yield": None, "evaluation": "failed", "fallback_reason": "confidence_evaluation_failed", "decision": "ar", } scored = { "request_id": 8, "slot": 1, - "initial_confidence": 2.0, "calibrated_yield": 1.5, + "initial_confidence": 2.0, "expected_yield": 1.5, "evaluation": "scored", "fallback_reason": None, "decision": "speculation", } @@ -541,13 +541,13 @@ def test_failed_evaluation_activation_is_request_local_and_explicit( def test_activation_record_fields_are_strictly_validated(self) -> None: valid = { "request_id": 7, "slot": 0, - "initial_confidence": 1.0, "calibrated_yield": 1.0, + "initial_confidence": 1.0, "expected_yield": 1.0, "evaluation": "scored", "fallback_reason": None, "decision": "ar", } failed = { "request_id": 7, "slot": 0, - "initial_confidence": None, "calibrated_yield": None, + "initial_confidence": None, "expected_yield": None, "evaluation": "failed", "fallback_reason": "confidence_evaluation_failed", "decision": "ar", @@ -562,7 +562,7 @@ def test_activation_record_fields_are_strictly_validated(self) -> None: ({**valid, "request_id": True}, "request_id"), ({**valid, "slot": -1}, "slot"), ({**valid, "initial_confidence": 0.99}, "initial_confidence"), - ({**valid, "calibrated_yield": math.nan}, "calibrated_yield"), + ({**valid, "expected_yield": math.nan}, "expected_yield"), ({**valid, "decision": "undecided"}, "decision"), (missing_evaluation, "evaluation"), (missing_fallback_reason, "fallback_reason"), diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h index 4256b487b..b389055cf 100644 --- a/server/src/common/concurrency/speculation_gate.h +++ b/server/src/common/concurrency/speculation_gate.h @@ -21,12 +21,9 @@ namespace dflash::common { struct SpecGateConfig { - // Offline calibration for the first request cohort, before any admitted - // request has produced online yield evidence. Generic callers default to - // neutral; a concrete speculator may install its fitted prior. - double initial_yield_scale = 1.0; - double yield_ema_alpha = 0.20; - double calibration_ema_alpha = 0.20; + // Immutable offline fit applied independently to every request's one-time + // confidence score. It never learns from request execution history. + double fixed_yield_scale = 1.0; double cost_ema_alpha = 0.20; double adaptive_gain_margin = 0.02; }; @@ -101,8 +98,8 @@ struct SpecCandidate { // bootstrap and a finite value is the preferred activation measurement. // Evaluation failure explicitly falls back to sticky AR without inventing // a score. Otherwise the gate commits exactly one mode from this - // survival-product expected yield, including the root; the gate owns - // calibration and clamping. + // survival-product expected yield, including the root; the gate applies + // only its fixed offline scale and clamping. double confidence_yield = std::numeric_limits::quiet_NaN(); }; @@ -195,9 +192,9 @@ struct SpecPlan { double profiled_cost = 0.0; double cost_scale = 1.0; double predicted_cost = 0.0; - // Final calibrated expected yield for admitted confidence-scored lanes. - // This is directly comparable with realized emitted tokens in telemetry. - double calibration_predicted_tokens = 0.0; + // Fixed-scale expected yield for admitted confidence-scored lanes. This is + // directly comparable with realized emitted tokens in telemetry. + double initial_predicted_tokens = 0.0; double goodput = 0.0; double ar_goodput = 0.0; int unavailable_count = 0; @@ -245,15 +242,6 @@ class SpeculationGate { std::numeric_limits::quiet_NaN(); SpecDecision decision = SpecDecision::Undecided; bool confidence_evaluation_failed = false; - double yield_ema = std::numeric_limits::quiet_NaN(); - uint64_t yield_observations = 0; - double calibration_ratio_ema = 1.0; - uint64_t calibration_observations = 0; - }; - - struct PendingConfidence { - double confidence_yield = - std::numeric_limits::quiet_NaN(); }; struct CostShape { @@ -296,8 +284,6 @@ class SpeculationGate { struct CandidateScore { double expected_yield = 1.0; - double uncalibrated_confidence = - std::numeric_limits::quiet_NaN(); SpecScoreSource source = SpecScoreSource::Unavailable; }; @@ -323,12 +309,10 @@ class SpeculationGate { auto valid_alpha = [](double value) { return std::isfinite(value) && value > 0.0 && value <= 1.0; }; - return valid_alpha(config_.yield_ema_alpha) && - valid_alpha(config_.calibration_ema_alpha) && - valid_alpha(config_.cost_ema_alpha) && - std::isfinite(config_.initial_yield_scale) && - config_.initial_yield_scale >= kCalibrationScaleMin && - config_.initial_yield_scale <= kCalibrationScaleMax && + return valid_alpha(config_.cost_ema_alpha) && + std::isfinite(config_.fixed_yield_scale) && + config_.fixed_yield_scale >= kFixedYieldScaleMin && + config_.fixed_yield_scale <= kFixedYieldScaleMax && std::isfinite(config_.adaptive_gain_margin) && config_.adaptive_gain_margin >= 0.0 && costs_.valid() && geometry_.tree_width >= 1 && @@ -349,17 +333,9 @@ class SpeculationGate { return out; } - // A plan is consumed synchronously by the engine. Drop any abandoned - // prediction from a failed previous execution before recording this one. - for (const SpecCandidate & candidate : candidates) { - pending_confidence_.erase(candidate.request_id); - } - struct Ranked { const SpecCandidate * candidate = nullptr; double score = 1.0; - double uncalibrated_confidence = - std::numeric_limits::quiet_NaN(); SpecScoreSource source = SpecScoreSource::Unavailable; SpecDecision decision = SpecDecision::Undecided; bool forced = false; @@ -402,13 +378,11 @@ class SpeculationGate { // Record the mandatory initial score in activation telemetry, // but permanently unsupported execution commits directly to AR. forced_ar.push_back({ - &candidate, score.expected_yield, - score.uncalibrated_confidence, score.source, + &candidate, score.expected_yield, score.source, prior_decision, false, true}); continue; } - Ranked ranked{&candidate, score.expected_yield, - score.uncalibrated_confidence, score.source, + Ranked ranked{&candidate, score.expected_yield, score.source, prior_decision, candidate.policy == SpeculationPolicy::Always || prior_decision == SpecDecision::Speculation, @@ -565,12 +539,8 @@ class SpeculationGate { out.ordered[(size_t)i].admitted = true; out.admitted_request_ids.push_back(ranked[(size_t)i].candidate->request_id); out.admitted_slots.push_back(ranked[(size_t)i].candidate->slot); - if (std::isfinite( - ranked[(size_t)i].uncalibrated_confidence)) { - out.calibration_predicted_tokens += - ranked[(size_t)i].score; - pending_confidence_[ranked[(size_t)i].candidate->request_id] = - {ranked[(size_t)i].uncalibrated_confidence}; + if (ranked[(size_t)i].source != SpecScoreSource::Unavailable) { + out.initial_predicted_tokens += ranked[(size_t)i].score; } } if (out.decisions_committed) { @@ -590,55 +560,6 @@ class SpeculationGate { return out; } - // `confidence_yield` may describe the completed speculative block. It - // calibrates later request activations (including forced executions) but - // cannot change this request's activation score or mode. `generated_tokens` - // remains only for source compatibility with existing adapters. - void observe(uint64_t request_id, double emitted_tokens, - int /* generated_tokens */, - double confidence_yield = - std::numeric_limits::quiet_NaN()) { - auto pending = pending_confidence_.find(request_id); - if (!std::isfinite(emitted_tokens) || emitted_tokens < 1.0 || - emitted_tokens > static_cast(max_accept_)) { - if (pending != pending_confidence_.end()) - pending_confidence_.erase(pending); - return; - } - - RequestState & state = request_states_[request_id]; - update_ema(state.yield_ema, state.yield_observations, - emitted_tokens, config_.yield_ema_alpha); - - double raw = std::numeric_limits::quiet_NaN(); - if (std::isfinite(confidence_yield)) { - raw = std::clamp(confidence_yield, 1.0, - static_cast(max_accept_)); - if (!std::isfinite(state.initial_confidence)) - state.initial_confidence = raw; - } else if (pending != pending_confidence_.end()) { - raw = pending->second.confidence_yield; - } - if (std::isfinite(raw)) { - const double ratio = std::clamp( - emitted_tokens / raw, - kCalibrationScaleMin, kCalibrationScaleMax); - update_ema(state.calibration_ratio_ema, - state.calibration_observations, ratio, - config_.calibration_ema_alpha); - update_ema(global_calibration_ratio_ema_, - global_calibration_observations_, ratio, - config_.calibration_ema_alpha); - } - if (pending != pending_confidence_.end()) - pending_confidence_.erase(pending); - } - - // Compatibility for adapters that do not publish a current-block score. - void observe(uint64_t request_id, double emitted_tokens) { - observe(request_id, emitted_tokens, 0); - } - void observe_cost(const SpecPlan & plan, double measured_us) { if (!plan.valid || !std::isfinite(measured_us) || measured_us <= 0.0 || !std::isfinite(plan.profiled_cost) || plan.profiled_cost <= 0.0) { @@ -665,13 +586,11 @@ class SpeculationGate { std::numeric_limits::quiet_NaN(); state.confidence_evaluation_failed = true; state.decision = SpecDecision::AR; - pending_confidence_.erase(request_id); return true; } void forget(uint64_t request_id) { request_states_.erase(request_id); - pending_confidence_.erase(request_id); } bool has_state(uint64_t request_id) const { @@ -698,42 +617,7 @@ class SpeculationGate { return state == request_states_.end() ? SpecDecision::Undecided : state->second.decision; } - double yield_ema(uint64_t request_id) const { - auto state = request_states_.find(request_id); - return state == request_states_.end() - ? std::numeric_limits::quiet_NaN() - : state->second.yield_ema; - } - - double calibration_scale(uint64_t request_id) const { - auto state = request_states_.find(request_id); - if (state != request_states_.end() && - state->second.calibration_observations > 0) { - double scale = state->second.calibration_ratio_ema; - // The emitted-yield EWMA is a second request-local calibration - // view. Blend it as a ratio against the initial raw score so it - // refines that score instead of becoming a score by itself. - if (state->second.yield_observations > 0 && - std::isfinite(state->second.yield_ema) && - std::isfinite(state->second.initial_confidence)) { - const double yield_scale = std::clamp( - state->second.yield_ema / - state->second.initial_confidence, - kCalibrationScaleMin, kCalibrationScaleMax); - scale = 0.5 * (scale + yield_scale); - } - return std::clamp( - scale, kCalibrationScaleMin, kCalibrationScaleMax); - } - return calibration_scale(); - } - double calibration_scale() const { - return global_calibration_observations_ == 0 - ? config_.initial_yield_scale : global_calibration_ratio_ema_; - } - uint64_t calibration_observations() const { - return global_calibration_observations_; - } + double fixed_yield_scale() const { return config_.fixed_yield_scale; } const SpecCostTables & costs() const { return costs_; } private: @@ -750,10 +634,8 @@ class SpeculationGate { } return { std::clamp( - calibration_scale(candidate.request_id) * - state.initial_confidence, + config_.fixed_yield_scale * state.initial_confidence, 1.0, static_cast(max_accept_)), - state.initial_confidence, accepted_initial_score ? SpecScoreSource::Confidence : SpecScoreSource::InitialConfidence, }; @@ -763,18 +645,13 @@ class SpeculationGate { std::isfinite(state->second.initial_confidence)) { return { std::clamp( - calibration_scale(candidate.request_id) * + config_.fixed_yield_scale * state->second.initial_confidence, 1.0, static_cast(max_accept_)), - state->second.initial_confidence, SpecScoreSource::InitialConfidence, }; } - return { - 1.0, - std::numeric_limits::quiet_NaN(), - SpecScoreSource::Unavailable, - }; + return {1.0, SpecScoreSource::Unavailable}; } static void update_ema(double & value, uint64_t & observations, @@ -802,14 +679,11 @@ class SpeculationGate { SpecCostTables costs_; SpecStepGeometry geometry_; int max_accept_ = 1; - static constexpr double kCalibrationScaleMin = 0.25; - static constexpr double kCalibrationScaleMax = 4.0; + static constexpr double kFixedYieldScaleMin = 0.25; + static constexpr double kFixedYieldScaleMax = 4.0; static constexpr double kCostScaleMin = 0.25; static constexpr double kCostScaleMax = 4.0; - double global_calibration_ratio_ema_ = 1.0; - uint64_t global_calibration_observations_ = 0; std::unordered_map request_states_; - std::unordered_map pending_confidence_; std::unordered_map cost_states_; ClampLogger clamp_logger_; }; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index f9aa51473..ef85225a4 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -40,7 +40,7 @@ int decode_bucket_width(int live_count) { return 64; } -double confidence_realized_tokens( +double initial_prediction_realized_tokens( const SpecPlan & plan, const SeqEngine::StepResult & result) { double realized = 0.0; for (const SpecPlanScore & score : plan.ordered) { @@ -60,10 +60,8 @@ double confidence_realized_tokens( return realized; } -void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, - uint64_t calibration_rounds, - double calibration_realized_tokens, - double measured_us) { +void log_spec_gate_plan(const SpecPlan & plan, double fixed_yield_scale, + double initial_realized_tokens, double measured_us) { int confidence = 0; int initial = 0; int unavailable = plan.unavailable_count; @@ -96,12 +94,11 @@ void log_spec_gate_plan(const SpecPlan & plan, double calibration_scale, } std::fprintf(stderr, "] sources=confidence:%d,initial:%d,unavailable:%d " - "calibration=%.3f rounds=%llu calib_tokens=%.3f/", - confidence, initial, unavailable, calibration_scale, - (unsigned long long)calibration_rounds, - plan.calibration_predicted_tokens); - if (std::isfinite(calibration_realized_tokens)) { - std::fprintf(stderr, "%.3f", calibration_realized_tokens); + "fixed_yield_scale=%.3f initial_tokens=%.3f/", + confidence, initial, unavailable, fixed_yield_scale, + plan.initial_predicted_tokens); + if (std::isfinite(initial_realized_tokens)) { + std::fprintf(stderr, "%.3f", initial_realized_tokens); } else { std::fprintf(stderr, "n/a"); } @@ -124,7 +121,7 @@ void log_spec_activations(const SpecPlan & plan, const double initial = gate.initial_confidence(score.request_id); std::fprintf(stderr, "[spec-activation] {\"request_id\":%llu,\"slot\":%d," - "\"initial_confidence\":%.6f,\"calibrated_yield\":%.6f," + "\"initial_confidence\":%.6f,\"expected_yield\":%.6f," "\"evaluation\":\"scored\",\"fallback_reason\":null," "\"decision\":\"%s\"}\n", (unsigned long long)score.request_id, score.slot, @@ -136,7 +133,7 @@ void log_spec_activations(const SpecPlan & plan, void log_spec_evaluation_fallback(uint64_t request_id, int slot) { std::fprintf(stderr, "[spec-activation] {\"request_id\":%llu,\"slot\":%d," - "\"initial_confidence\":null,\"calibrated_yield\":null," + "\"initial_confidence\":null,\"expected_yield\":null," "\"evaluation\":\"failed\"," "\"fallback_reason\":\"confidence_evaluation_failed\"," "\"decision\":\"ar\"}\n", @@ -1718,18 +1715,6 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( out.target_forwards = 2; out.committed_tokens.assign( proposal.path.begin() + 1, proposal.path.end()); - if (speculation_gate_) { - const double block_confidence = - proposal.slot >= 0 && - proposal.slot < (int)last_survival_score_.size() - ? last_survival_score_[(size_t)proposal.slot] - : std::numeric_limits::quiet_NaN(); - speculation_gate_->observe( - slots_.slot(proposal.slot).request_id, - (double)proposal.path.size(), - slots_.slot(proposal.slot).generated_tokens(), - block_confidence); - } } else { ArLane & ar = ar_lanes[static_cast(ar_for_input[i])]; @@ -2745,11 +2730,10 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (have_gate_plan && spec_gate_debug_enabled() && spec_completed && !proposal_failed) { const double realized_tokens = spec_completed - ? confidence_realized_tokens(gate_plan, speculative) + ? initial_prediction_realized_tokens(gate_plan, speculative) : std::numeric_limits::quiet_NaN(); log_spec_gate_plan( - gate_plan, speculation_gate_->calibration_scale(), - speculation_gate_->calibration_observations(), + gate_plan, speculation_gate_->fixed_yield_scale(), realized_tokens, cost_sample_valid ? measured_us @@ -3163,8 +3147,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (spec_gate_debug_enabled()) { log_spec_gate_plan( - *pending_ar_gate_plan, speculation_gate_->calibration_scale(), - speculation_gate_->calibration_observations(), 0.0, + *pending_ar_gate_plan, + speculation_gate_->fixed_yield_scale(), 0.0, cost_sample_valid ? measured_us : std::numeric_limits::quiet_NaN()); diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index efa9795ab..59070e6e6 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -291,123 +291,45 @@ int main() { CHECK(plan.cost_lookup_clamped); CHECK(clamp_logs > 0); - // Calibration is request-local and active after one observation. An - // explicit current-block confidence calibrates that observation without - // replacing the immutable initial score; the global ratio helps only new - // undecided requests. Telemetry sums the final calibrated admitted score. - SpecGateConfig fast; - fast.yield_ema_alpha = 0.5; - fast.calibration_ema_alpha = 0.5; - SpeculationGate calibrated(fast, crossover, geometry(), 4); - plan = calibrated.plan(1, {candidate(700, 0, 4.0)}, 1); - CHECK(calibrated.decision(700) == SpecDecision::Speculation); - CHECK(calibrated.initial_confidence(700) == 4.0); - calibrated.observe(700, 2.0, 1, 4.0); - CHECK(calibrated.calibration_observations() == 1); - CHECK(std::abs(calibrated.calibration_scale() - 0.5) < 1e-12); - CHECK(std::abs(calibrated.calibration_scale(700) - 0.5) < 1e-12); - CHECK(std::abs(calibrated.yield_ema(700) - 2.0) < 1e-12); - - plan = calibrated.plan(1, {candidate(700, 0, 1.0)}, 1); - CHECK(plan.admitted_count == 1); - CHECK(plan.ordered[0].forced); - CHECK(plan.ordered[0].source == SpecScoreSource::InitialConfidence); - CHECK(plan.ordered[0].expected_yield == 2.0); - CHECK(plan.calibration_predicted_tokens == 2.0); - CHECK(calibrated.initial_confidence(700) == 4.0); - - calibrated.observe(700, 4.0, 2, 2.0); - CHECK(std::abs(calibrated.calibration_scale() - 1.25) < 1e-12); - CHECK(std::abs(calibrated.calibration_scale(700) - 1.0) < 1e-12); - CHECK(std::abs(calibrated.yield_ema(700) - 3.0) < 1e-12); - CHECK(calibrated.initial_confidence(700) == 4.0); - plan = calibrated.plan(1, {candidate(700, 0, NAN)}, 1); - CHECK(plan.ordered[0].expected_yield == 4.0); - CHECK(plan.calibration_predicted_tokens == 4.0); - - plan = calibrated.plan(1, {candidate(701, 1, 2.0)}, 1); - CHECK(std::abs(plan.ordered[0].expected_yield - 2.5) < 1e-12); - CHECK(std::abs(calibrated.calibration_scale(701) - 1.25) < 1e-12); - calibrated.observe(701, 1.0, 1, 2.0); - CHECK(std::abs(calibrated.calibration_scale(701) - 0.5) < 1e-12); - CHECK(std::abs(calibrated.calibration_scale(700) - 1.0) < 1e-12); - - // A deployment may install a request-level offline fit for the cold - // cohort. It calibrates the first irreversible choice immediately, while - // the immutable raw confidence remains available for telemetry. + // Every request is ranked only from its own immutable first score and the + // fixed offline scale. Re-presenting a different score cannot change the + // request, and one request never supplies a prior for another. SpecGateConfig fitted_config; - fitted_config.initial_yield_scale = 0.5; + fitted_config.fixed_yield_scale = 0.5; SpeculationGate fitted( fitted_config, constant_costs(1.0, 10.0, 1.0), geometry(), 4); - plan = fitted.plan(1, {candidate(702, 0, 4.0)}, 1); - CHECK(fitted.calibration_scale() == 0.5); - CHECK(fitted.initial_confidence(702) == 4.0); + plan = fitted.plan(1, {candidate(700, 0, 4.0)}, 1); + CHECK(fitted.fixed_yield_scale() == 0.5); + CHECK(fitted.initial_confidence(700) == 4.0); CHECK(plan.ordered[0].expected_yield == 2.0); + CHECK(plan.initial_predicted_tokens == 2.0); CHECK(plan.ordered[0].newly_decided); - // Yield evidence alone does not invent an activation score. Forced/probe - // execution can publish a score and train calibration immediately while - // leaving the configured policy outside the adaptive decision state. - SpeculationGate isolated(constant_costs(1.0, 10.0, 1.0), - geometry(), 4); - isolated.observe(800, 3.0, 1); - CHECK(isolated.calibration_observations() == 0); - CHECK(!isolated.has_confidence(800)); - plan = isolated.plan(1, {candidate(800, 0, NAN)}, 1); - CHECK(plan.pending_evaluations.size() == 1); - CHECK(plan.pending_evaluations[0].slot == 0); - CHECK(plan.pending_evaluations[0].action == - SpecEvaluationAction::Score); - CHECK(isolated.decision(800) == SpecDecision::Undecided); - - plan = isolated.plan(1, { - candidate(801, 0, NAN, SpeculationPolicy::Always)}, 1); - CHECK(plan.admitted_count == 1); - isolated.observe(801, 2.0, 1, 4.0); - CHECK(isolated.has_confidence(801)); - CHECK(isolated.initial_confidence(801) == 4.0); - CHECK(isolated.calibration_observations() == 1); - CHECK(std::abs(isolated.calibration_scale(801) - 0.5) < 1e-12); - CHECK(isolated.decision(801) == SpecDecision::Undecided); - plan = isolated.plan(1, {candidate(801, 0, NAN)}, 1); - CHECK(plan.pending_evaluations.empty()); - CHECK(isolated.decision(801) == SpecDecision::Speculation); + plan = fitted.plan(1, {candidate(700, 0, 1.0)}, 1); + CHECK(plan.ordered[0].forced); + CHECK(plan.ordered[0].source == SpecScoreSource::InitialConfidence); + CHECK(plan.ordered[0].expected_yield == 2.0); + CHECK(fitted.initial_confidence(700) == 4.0); - // Forget is the only adaptive state reset. It drops the sticky decision, - // immutable score, yield, and local calibration while preserving the - // deployment-global calibration prior for future requests. - const double global_before_forget = calibrated.calibration_scale(); - calibrated.forget(700); - CHECK(!calibrated.has_state(700)); - CHECK(!calibrated.has_confidence(700)); - CHECK(std::isnan(calibrated.initial_confidence(700))); - CHECK(calibrated.decision(700) == SpecDecision::Undecided); - CHECK(calibrated.calibration_scale() == global_before_forget); - plan = calibrated.plan(1, {candidate(700, 0, NAN)}, 1); + plan = fitted.plan(1, {candidate(701, 0, 4.0)}, 1); + CHECK(plan.ordered[0].expected_yield == 2.0); + CHECK(fitted.initial_confidence(701) == 4.0); + CHECK(fitted.initial_confidence(700) == 4.0); + + // forget() removes only this request's activation state. A repeated ID is + // cold again; fixed configuration and shape-cost feedback are independent. + fitted.forget(700); + CHECK(!fitted.has_state(700)); + CHECK(!fitted.has_confidence(700)); + CHECK(std::isnan(fitted.initial_confidence(700))); + CHECK(fitted.decision(700) == SpecDecision::Undecided); + plan = fitted.plan(1, {candidate(700, 0, NAN)}, 1); CHECK(!plan.decisions_committed); CHECK(plan.pending_evaluations.size() == 1); CHECK(plan.pending_evaluations[0].slot == 0); CHECK(plan.pending_evaluations[0].action == SpecEvaluationAction::Score); - // Calibration clamps are active on the first forced/probe observation. - SpeculationGate lower_bound(constant_costs(1.0, 1.0, 1.0), - geometry(), 16); - plan = lower_bound.plan(1, { - candidate(900, 0, 16.0, SpeculationPolicy::Always)}, 1); - lower_bound.observe(900, 1.0, 1); - CHECK(lower_bound.calibration_observations() == 1); - CHECK(lower_bound.calibration_scale() == 0.25); - lower_bound.forget(900); - CHECK(lower_bound.calibration_scale() == 0.25); - SpeculationGate upper_bound(constant_costs(1.0, 1.0, 1.0), - geometry(), 16); - plan = upper_bound.plan(1, { - candidate(901, 0, 4.0, SpeculationPolicy::Always)}, 1); - upper_bound.observe(901, 16.0, 1); - CHECK(upper_bound.calibration_observations() == 1); - CHECK(upper_bound.calibration_scale() == 4.0); - // Shape-local total-cost feedback changes only future undecided choices. // It cannot flip a request whose one-shot decision is already sticky. SpeculationGate cost_feedback( @@ -426,6 +348,8 @@ int main() { CHECK(plan.predicted_cost == 48.0); CHECK(cost_feedback.decision(950) == SpecDecision::Speculation); + cost_feedback.forget(950); + CHECK(cost_feedback.decision(950) == SpecDecision::Undecided); plan = cost_feedback.plan(1, {candidate(951, 0, 4.0)}, 1); CHECK(plan.admitted_count == 0); CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); From 79ac1b48b0c13a28230b18ac64f694faeae95cfc Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 19:44:51 +0000 Subject: [PATCH 35/42] feat(concurrency): add adaptive DFlash2 activation Add batched DFlash2 chain selection and a request-local benefit adapter feeding the sticky joint goodput gate, with explicit AR fallback and typed activation telemetry. Fix fractional gate pricing and executed-shape cost feedback, add variable positive chain depth, replay-journal prototypes, artifact conversion hardening, and fail-closed subset/refill analysis harnesses. --- harness/benchmarks/concurrency/README.md | 87 ++ .../concurrency/analyze_dflash2_selector.py | 815 ++++++++++++++++++ .../concurrency/analyze_gate_decisions.py | 177 +++- .../concurrency/forced_subset_benchmark.py | 597 +++++++++++++ .../concurrency/refill_subset_benchmark.py | 638 ++++++++++++++ .../concurrency/run_qwen38_dflash2_subsets.sh | 283 ++++++ .../test_analyze_dflash2_selector.py | 237 +++++ .../concurrency/test_feature_tools.py | 110 +++ .../test_forced_subset_benchmark.py | 268 ++++++ .../test_refill_subset_benchmark.py | 255 ++++++ server/CMakeLists.txt | 44 + .../deps/llama.cpp/ggml/include/ggml-cuda.h | 13 + server/deps/llama.cpp/ggml/include/ggml.h | 8 + .../ggml/src/ggml-cuda/gated_delta_net.cu | 72 +- .../src/ggml-cuda/gdn-transition-journal.cu | 184 ++++ server/deps/llama.cpp/ggml/src/ggml.c | 26 + server/scripts/quantize_draft_q8.py | 131 ++- .../common/concurrency/chain_spec_shapes.h | 23 + .../src/common/concurrency/speculation_gate.h | 249 ++++-- server/src/common/dflash2_batch.cpp | 440 ++++++++++ server/src/common/dflash2_benefit.cpp | 198 +++++ server/src/common/dflash2_benefit.h | 92 ++ server/src/common/dflash2_head.cpp | 73 +- server/src/common/dflash2_head.h | 38 +- .../src/common/dflash2_selector_validation.h | 93 ++ server/src/common/dflash_draft_kv.cpp | 66 +- server/src/common/dflash_draft_kv.h | 6 +- .../src/common/geometric_draft_topk_cuda.cu | 18 +- server/src/common/geometric_draft_topk_cuda.h | 9 + server/src/common/speculation_policy.h | 6 +- server/src/draft/draft_gguf_loader.cpp | 25 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 566 +++++++++--- .../qwen35/concurrency/qwen35_seq_engine.h | 29 +- .../src/qwen35/delta_transition_journal.cpp | 142 +++ server/src/qwen35/delta_transition_journal.h | 74 ++ server/src/qwen35/qwen35_backend.cpp | 63 +- server/test/test_chain_spec_shapes.cpp | 18 + server/test/test_delta_transition_journal.cpp | 186 ++++ server/test/test_dflash2_benefit.cpp | 230 +++++ .../test/test_dflash2_selector_validation.cpp | 86 ++ server/test/test_draft_topk_cuda.cpp | 57 +- server/test/test_gdn_transition_journal.cpp | 563 ++++++++++++ server/test/test_speculation_gate.cpp | 61 +- server/tests/test_quantize_draft_q8.py | 54 ++ 44 files changed, 7034 insertions(+), 376 deletions(-) create mode 100644 harness/benchmarks/concurrency/analyze_dflash2_selector.py create mode 100755 harness/benchmarks/concurrency/forced_subset_benchmark.py create mode 100644 harness/benchmarks/concurrency/refill_subset_benchmark.py create mode 100755 harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh create mode 100644 harness/benchmarks/concurrency/test_analyze_dflash2_selector.py create mode 100644 harness/benchmarks/concurrency/test_forced_subset_benchmark.py create mode 100644 harness/benchmarks/concurrency/test_refill_subset_benchmark.py create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu create mode 100644 server/src/common/dflash2_batch.cpp create mode 100644 server/src/common/dflash2_benefit.cpp create mode 100644 server/src/common/dflash2_benefit.h create mode 100644 server/src/common/dflash2_selector_validation.h create mode 100644 server/src/qwen35/delta_transition_journal.cpp create mode 100644 server/src/qwen35/delta_transition_journal.h create mode 100644 server/test/test_delta_transition_journal.cpp create mode 100644 server/test/test_dflash2_benefit.cpp create mode 100644 server/test/test_dflash2_selector_validation.cpp create mode 100644 server/test/test_gdn_transition_journal.cpp diff --git a/harness/benchmarks/concurrency/README.md b/harness/benchmarks/concurrency/README.md index b39db289e..ae9dfdae2 100644 --- a/harness/benchmarks/concurrency/README.md +++ b/harness/benchmarks/concurrency/README.md @@ -68,3 +68,90 @@ 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. + +## Forced DFlash2 subset/depth diagnostics + +`forced_subset_benchmark.py` is the fail-closed client for the DFlash2 +concurrency bring-up. It is intentionally limited to forced controls: every +request receives an explicit `decode_mode` of `ar` or `speculation`, and the +artifact explicitly forbids interpreting the result as adaptive activation. +The repository-owned runner generates DFlash2-specific metadata and keeps one +server process alive across the complete positional mask set: + +```bash +MODEL=/path/Qwen3.8-27B-target.gguf \ +DRAFT_MODEL=/path/Qwen3.8-27B-DFlash2-q8_0.gguf \ +PROMPT_FILE=/path/prompts.jsonl \ +CLIENTS=2 MASKS=AA,AS,SA,SS SPEC_DEPTH=4 \ +harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh +``` + +Run a fresh process/output directory for each depth in a `{2,4,8}` screen. +The runner performs all-AR and all-SPEC warmups before measuring any mask. +It starts the backend in global `speculation` mode so the chain is allocated; +each `A` request then overrides that default to request-local AR. The host +default is `VISIBLE_DEVICES=0` (the R9700), and can be overridden explicitly. + +Launch the server with `DFLASH_SPEC_CHAIN_DEPTH=`, +`DFLASH_STEP_TIMING=1`, and, whenever any lane is speculative, +`DFLASH_DFLASH2_SELECTOR_LOG=1`. Those values must also appear in the metadata +file's `launch_environment`; the client rejects a mismatch. The server log may +contain startup and warmup output: the client hashes and parses only bytes +appended after its synchronized request cohort is ready to run. + +The JSON retains each request, its positional forced mode, prompt/payload and +exact content/reasoning hashes, every measured `[step-timing]`, +`[spec-selector]`, `[spec-activation]`, and `[concurrency-metrics]` record, plus +an exact hash and byte range for the measured server-log span. A case fails if +the requested modes did not execute, the inferred tree depth differs, output +or token accounting is incomplete, request starts exceed the configured skew, +or live concurrency `C` is not sustained for at least two consecutive decode +rounds. Use `--min-full-live-rounds` only to state a different threshold +explicitly; do not waive the requirement for at least one full-live round. + +### Refill/saturated service diagnostic + +Set `REFILL_WAVES` to at least 3 to run `refill_subset_benchmark.py` through +the same persistent-server runner: + +```bash +MODEL=/path/Qwen3.8-27B-target.gguf \ +DRAFT_MODEL=/path/Qwen3.8-27B-DFlash2-q8_0.gguf \ +PROMPT_FILE=/path/prompts.jsonl \ +CLIENTS=6 MASKS=AAAAAA,SSSSSS SPEC_DEPTH=8 REFILL_WAVES=4 \ +MAX_TOKENS=64 \ +harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh +``` + +Each positional lane reuses its deterministic prompt and forced A/S mode, and +the first `C * (waves - 1)` successful completions submit another request on +the lane that completed. Faster lanes therefore receive more requests instead +of exhausting a fixed per-lane quota, and the active A/S mask is preserved +until the final drain. This closed-loop workload answers whether the server can +profit while work is replenished. The default `REFILL_WAVES=1` instead measures +a single closed cohort whose faster lanes create a terminal tail. Refill mode +requires at least three waves: the last C scheduled refills form a guard cohort +around the saturation proof. Run both protocols when comparing scheduler policy. + +The refill report deliberately separates two metrics: + +- `aggregate_refill_tok_s` is exact completion tokens divided by the entire + multi-wave wall interval. It includes initial prefill, client/server handoff + gaps, later prefills, and the final drain. Increasing the declared wave count + amortizes the one-time boundaries; never relabel it as a closed-cohort rate. +- `validation.full_live.engine_round_goodput_tok_s` is emitted tokens divided by + summed server `total_us` for only `live=C` timed engine rounds. It isolates + saturated round economics and includes work accounted inside `total_us`, but + excludes time between timing records such as client handoff gaps and is not + end-to-end throughput. + +The client retains every raw timing, selector, activation, and per-request +metric record and its exact measured log byte range. It also retains each +request payload/prompt/output hash. It fails closed unless all `C * waves` +requests finish the fixed token count, each forced mode and DFlash2 selector +mapping is proven by per-request telemetry, depth matches, identical repeated +lane inputs produce identical output hashes, refill handoff gaps meet the +stated bound, and a later `live=C` round occurs after at least +`C * (waves - 2)` completions while one C-request guard cohort remains before +the final drain. The artifact explicitly permits neither an adaptive activation +claim nor a closed-cohort makespan claim. diff --git a/harness/benchmarks/concurrency/analyze_dflash2_selector.py b/harness/benchmarks/concurrency/analyze_dflash2_selector.py new file mode 100644 index 000000000..69ea1148d --- /dev/null +++ b/harness/benchmarks/concurrency/analyze_dflash2_selector.py @@ -0,0 +1,815 @@ +#!/usr/bin/env python3 +"""Offline analysis for forced DFlash2 subset/depth artifacts. + +The forced-subset client retains exact JSON records from ``[spec-selector]``, +``[concurrency-metrics]``, and ``[step-timing]`` inside each ``bench.json``. +This analyzer validates those records, joins selector engine IDs to wire +requests, and keeps two questions separate: + +* can proposal-local raw signals predict accepted yield; and +* can independent request rankings predict the best concurrent cohort? + +It deliberately does not fit or emit a runtime activation threshold. A +calibrator is only justified when held-out data contains both positive and +negative acceptance/benefit labels. +""" + +from __future__ import annotations + +import argparse +import itertools +import json +import math +import statistics +from collections import defaultdict +from pathlib import Path +from typing import Any, Iterable + + +FEATURE_DIRECTIONS = { + "chain_lm_logp": 1, + "chain_lm_probability": 1, + "mean_selected_logp": 1, + "min_selected_logp": 1, + "mean_lm_margin": 1, + "min_lm_margin": 1, + "mean_topk_mass": 1, + "min_topk_mass": 1, + "lm_top1_fraction": 1, + "mean_selector_margin": 1, + "min_selector_margin": 1, + "selector_chain_probability": 1, + "mean_selector_mass": 1, + "min_selector_mass": 1, + "mean_selector_entropy": -1, + "max_selector_entropy": -1, + "mean_rank": -1, + "max_rank": -1, +} + + +def _finite_number(value: Any, label: str) -> float: + if type(value) not in (int, float) or not math.isfinite(value): + raise ValueError(f"{label} must be a finite number") + return float(value) + + +def _non_negative_int(value: Any, label: str) -> int: + if type(value) is not int or value < 0: + raise ValueError(f"{label} must be a non-negative integer") + return value + + +def _record(wrapper: Any, label: str) -> dict[str, Any]: + if not isinstance(wrapper, dict) or not isinstance(wrapper.get("record"), dict): + raise ValueError(f"{label} must wrap a JSON object in record") + record = wrapper["record"] + raw = wrapper.get("raw_json") + if not isinstance(raw, str): + raise ValueError(f"{label} must retain raw_json") + try: + parsed = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"{label} contains invalid raw_json: {exc}") from exc + if parsed != record: + raise ValueError(f"{label} raw_json and parsed record disagree") + return record + + +def parse_profile_lines(data: bytes) -> dict[str, list[dict[str, Any]]]: + """Parse the three profiling records used by this analyzer. + + This helper is intentionally small; the benchmark client remains the + authority for log-span capture. It is useful for parser tests and for + diagnosing an artifact whose retained record is malformed. + """ + prefixes = { + "selectors": b"[spec-selector] ", + "requests": b"[concurrency-metrics] ", + "rounds": b"[step-timing] ", + } + out = {key: [] for key in prefixes} + for line_index, line in enumerate(data.splitlines(), 1): + for key, prefix in prefixes.items(): + position = line.find(prefix) + if position < 0: + continue + raw = line[position + len(prefix):].decode("utf-8", errors="strict") + try: + record = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError( + f"line {line_index}: invalid {prefix.decode().strip()} JSON: {exc}" + ) from exc + if not isinstance(record, dict): + raise ValueError(f"line {line_index}: profiling JSON must be an object") + out[key].append({ + "line_index": line_index, + "raw_json": raw, + "record": record, + }) + break + return out + + +def first_block_features(selector: dict[str, Any]) -> dict[str, float]: + """Collapse pre-verification raw signals from one proposal block.""" + depths = selector.get("depths") + if not isinstance(depths, list) or not depths: + raise ValueError("selector depths must be a non-empty array") + numeric: dict[str, list[float]] = defaultdict(list) + top1: list[float] = [] + expected_depth = 1 + for row in depths: + if not isinstance(row, dict): + raise ValueError("selector depth entry must be an object") + if row.get("depth") != expected_depth: + raise ValueError("selector depths must be contiguous and one-based") + expected_depth += 1 + for key in ( + "selected_logp", "lm_margin", "topk_mass", "selector_margin", + "selector_mass", "selector_entropy", "rank", + ): + numeric[key].append(_finite_number(row.get(key), f"selector {key}")) + if type(row.get("lm_top1")) is not bool: + raise ValueError("selector lm_top1 must be boolean") + top1.append(float(row["lm_top1"])) + + logp = numeric["selected_logp"] + selector_mass = numeric["selector_mass"] + chain_lm_logp = sum(logp) + selector_chain_logp = sum(math.log(max(value, 1e-300)) for value in selector_mass) + return { + "chain_lm_logp": chain_lm_logp, + "chain_lm_probability": math.exp(chain_lm_logp), + "mean_selected_logp": statistics.fmean(logp), + "min_selected_logp": min(logp), + "mean_lm_margin": statistics.fmean(numeric["lm_margin"]), + "min_lm_margin": min(numeric["lm_margin"]), + "mean_topk_mass": statistics.fmean(numeric["topk_mass"]), + "min_topk_mass": min(numeric["topk_mass"]), + "lm_top1_fraction": statistics.fmean(top1), + "mean_selector_margin": statistics.fmean(numeric["selector_margin"]), + "min_selector_margin": min(numeric["selector_margin"]), + "selector_chain_probability": math.exp(selector_chain_logp), + "mean_selector_mass": statistics.fmean(selector_mass), + "min_selector_mass": min(selector_mass), + "mean_selector_entropy": statistics.fmean(numeric["selector_entropy"]), + "max_selector_entropy": max(numeric["selector_entropy"]), + "mean_rank": statistics.fmean(numeric["rank"]), + "max_rank": max(numeric["rank"]), + } + + +def _validate_selector(selector: dict[str, Any], label: str) -> None: + _non_negative_int(selector.get("request_id"), f"{label} request_id") + _non_negative_int(selector.get("slot"), f"{label} slot") + _non_negative_int(selector.get("generated"), f"{label} generated") + accepted_depth = _non_negative_int( + selector.get("accepted_depth"), f"{label} accepted_depth", + ) + features = first_block_features(selector) + del features + depths = selector["depths"] + if accepted_depth > len(depths): + raise ValueError(f"{label} accepted_depth exceeds proposal depth") + for index, row in enumerate(depths, 1): + if type(row.get("accepted")) is not bool: + raise ValueError(f"{label} depth {index} accepted must be boolean") + if row["accepted"] != (index <= accepted_depth): + raise ValueError(f"{label} accepted flags are not a prefix") + + +def _round_summary(rows: list[dict[str, Any]]) -> dict[str, Any]: + total_us = sum(float(row["total_us"]) for row in rows) + emitted = sum(int(row["emitted_tokens"]) for row in rows) + accepted = sum(int(row["accepted_tokens"]) for row in rows) + lane_steps = sum(int(row["k"]) for row in rows) + live_histogram: dict[int, int] = defaultdict(int) + for row in rows: + live_histogram[int(row["live"])] += 1 + return { + "rounds": len(rows), + "total_us": total_us, + "emitted_tokens": emitted, + "accepted_tokens": accepted, + "spec_lane_steps": lane_steps, + "goodput_tok_s": emitted * 1e6 / total_us if total_us else None, + "accepted_per_spec_lane_step": ( + accepted / lane_steps if lane_steps else None + ), + "live_histogram": dict(sorted(live_histogram.items())), + } + + +def analyze_artifact(path: Path) -> dict[str, Any]: + """Validate and join one forced-subset ``bench.json`` artifact.""" + bench = json.loads(path.read_text(encoding="utf-8")) + if bench.get("kind") != "dflash2-forced-subset-diagnostic": + raise ValueError(f"{path}: not a DFlash2 forced-subset artifact") + validation = bench.get("validation") + if not isinstance(validation, dict) or validation.get("passed") is not True: + raise ValueError(f"{path}: benchmark evidence did not pass validation") + level = bench.get("level") + if not isinstance(level, dict): + raise ValueError(f"{path}: level must be an object") + clients = _non_negative_int(level.get("clients"), f"{path} clients") + if clients < 1: + raise ValueError(f"{path}: clients must be positive") + spec_depth = _non_negative_int(bench.get("spec_depth"), f"{path} spec_depth") + if spec_depth < 2: + raise ValueError(f"{path}: spec_depth must be at least two") + mask = level.get("request_mode_mask") + if not isinstance(mask, str) or len(mask) != clients or set(mask) - {"A", "S"}: + raise ValueError(f"{path}: invalid request mode mask") + + details = level.get("requests_detail") + if not isinstance(details, list) or len(details) != clients: + raise ValueError(f"{path}: requests_detail count does not match clients") + detail_by_wire: dict[str, dict[str, Any]] = {} + for detail in details: + if not isinstance(detail, dict) or not isinstance(detail.get("request_id"), str): + raise ValueError(f"{path}: request detail lacks request_id") + wire = detail["request_id"] + if wire in detail_by_wire: + raise ValueError(f"{path}: duplicate wire request {wire}") + detail_by_wire[wire] = detail + + retained = bench.get("server_records") + if not isinstance(retained, dict): + raise ValueError(f"{path}: server_records must be an object") + timings = [] + for index, wrapper in enumerate(retained.get("rounds") or []): + row = _record(wrapper, f"{path} timing record {index}") + for key in ("live", "k", "accepted_tokens", "emitted_tokens"): + _non_negative_int(row.get(key), f"{path} timing {key}") + if row.get("path") not in ("ar", "spec"): + raise ValueError(f"{path}: timing path must be ar or spec") + total_us = _finite_number(row.get("total_us"), f"{path} timing total_us") + if total_us <= 0.0: + raise ValueError(f"{path}: timing total_us must be positive") + if int(row["live"]) < 1 or int(row["live"]) > clients: + raise ValueError(f"{path}: timing live exceeds request cohort") + timings.append(row) + + full_live = [row for row in timings if int(row["live"]) == clients] + tail = [row for row in timings if int(row["live"]) < clients] + round_timing = { + "all": _round_summary(timings), + "full_live": _round_summary(full_live), + "tail": _round_summary(tail), + } + + metric_by_engine: dict[int, dict[str, Any]] = {} + metric_by_wire: dict[str, dict[str, Any]] = {} + for index, wrapper in enumerate(retained.get("requests") or []): + metric = _record(wrapper, f"{path} request record {index}") + wire = metric.get("request_id") + engine = metric.get("engine_request_id") + if not isinstance(wire, str) or wire not in detail_by_wire: + raise ValueError(f"{path}: metric references unknown wire request {wire!r}") + engine = _non_negative_int(engine, f"{path} engine_request_id") + if engine in metric_by_engine or wire in metric_by_wire: + raise ValueError(f"{path}: duplicate request metric mapping") + metric_by_engine[engine] = metric + metric_by_wire[wire] = metric + if set(metric_by_wire) != set(detail_by_wire): + missing = sorted(set(detail_by_wire) - set(metric_by_wire)) + raise ValueError(f"{path}: missing concurrency metrics for {missing}") + + selectors_by_engine: dict[int, list[dict[str, Any]]] = defaultdict(list) + for index, wrapper in enumerate(retained.get("selectors") or []): + selector = _record(wrapper, f"{path} selector record {index}") + _validate_selector(selector, f"{path} selector record {index}") + engine = int(selector["request_id"]) + if engine not in metric_by_engine: + raise ValueError(f"{path}: selector references unknown engine request {engine}") + if len(selector["depths"]) != spec_depth - 1: + raise ValueError(f"{path}: selector proposal depth disagrees with spec_depth") + selectors_by_engine[engine].append(selector) + + requests = [] + for position, detail in enumerate(details): + wire = detail["request_id"] + metric = metric_by_wire[wire] + engine = int(metric["engine_request_id"]) + spec_steps = _non_negative_int(metric.get("spec_steps"), f"{path} spec_steps") + accepted = _non_negative_int( + metric.get("spec_accepted_tokens"), f"{path} spec_accepted_tokens", + ) + selectors = sorted( + selectors_by_engine.get(engine, []), key=lambda row: row["generated"], + ) + expected_mode = "speculation" if mask[position] == "S" else "ar" + if detail.get("decode_mode") != expected_mode: + raise ValueError(f"{path}: request mode mask and detail disagree") + if expected_mode == "ar": + if spec_steps or accepted or selectors: + raise ValueError(f"{path}: forced AR request contains speculation") + first = None + lifetime_yield = None + yield_fraction = None + else: + if spec_steps < 1 or len(selectors) != spec_steps: + raise ValueError(f"{path}: selector count does not match spec_steps") + if sum(int(row["accepted_depth"]) for row in selectors) != accepted: + raise ValueError(f"{path}: selector acceptance does not match request metric") + generated = [int(row["generated"]) for row in selectors] + if len(set(generated)) != len(generated) or generated[0] != 0: + raise ValueError(f"{path}: selector sequence lacks a unique first block") + first = first_block_features(selectors[0]) + lifetime_yield = accepted / spec_steps + yield_fraction = lifetime_yield / (spec_depth - 1) + requests.append({ + "position": position, + "wire_request_id": wire, + "engine_request_id": engine, + "mode": expected_mode, + "prompt_index": detail.get("prompt_index"), + "prompt_sha256": detail.get("prompt_sha256"), + "request_decode_tok_s": detail.get("request_decode_tok_s"), + "spec_steps": spec_steps, + "accepted_tokens": accepted, + "lifetime_accepted_yield": lifetime_yield, + "lifetime_yield_fraction": yield_fraction, + "first_accepted_depth": selectors[0]["accepted_depth"] if selectors else None, + "first_features": first, + "selectors": selectors, + }) + + return { + "path": str(path), + "clients": clients, + "spec_depth": spec_depth, + "mask": mask, + "repeat": (bench.get("server_metadata") or {}).get("repeat"), + "prompt_set_sha256": level.get("selected_prompt_set_sha256"), + "aggregate_tok_s": _finite_number( + level.get("aggregate_tok_s"), f"{path} aggregate_tok_s", + ), + "wall_s": _finite_number(level.get("wall_s"), f"{path} wall_s"), + "round_timing": round_timing, + "requests": requests, + } + + +def _average_ranks(values: list[float]) -> list[float]: + ordered = sorted(range(len(values)), key=values.__getitem__) + ranks = [0.0] * len(values) + offset = 0 + while offset < len(ordered): + end = offset + 1 + while end < len(ordered) and values[ordered[end]] == values[ordered[offset]]: + end += 1 + rank = (offset + end - 1) / 2.0 + 1.0 + for index in ordered[offset:end]: + ranks[index] = rank + offset = end + return ranks + + +def spearman_correlation(x: list[float], y: list[float]) -> float | None: + if len(x) != len(y) or len(x) < 2: + return None + rx = _average_ranks(x) + ry = _average_ranks(y) + mx = statistics.fmean(rx) + my = statistics.fmean(ry) + numerator = sum((a - mx) * (b - my) for a, b in zip(rx, ry)) + dx = sum((a - mx) ** 2 for a in rx) + dy = sum((b - my) ** 2 for b in ry) + if dx == 0.0 or dy == 0.0: + return None + return numerator / math.sqrt(dx * dy) + + +def _correlations( + rows: list[dict[str, Any]], outcome: str, +) -> dict[str, Any]: + eligible = [row for row in rows if row.get(outcome) is not None] + grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) + for index, row in enumerate(eligible): + prompt = row.get("prompt_sha256") + grouped[str(prompt) if prompt is not None else f"row-{index}"].append(row) + + independent = [] + for prompt, prompt_rows in sorted(grouped.items()): + independent.append({ + "prompt_sha256": prompt, + outcome: statistics.fmean(float(row[outcome]) for row in prompt_rows), + "first_features": { + feature: statistics.fmean( + float(row["first_features"][feature]) for row in prompt_rows + ) + for feature in FEATURE_DIRECTIONS + }, + }) + + def correlations(sample: list[dict[str, Any]]) -> dict[str, float | None]: + y = [float(row[outcome]) for row in sample] + values = {} + for feature in FEATURE_DIRECTIONS: + x = [float(row["first_features"][feature]) for row in sample] + raw = spearman_correlation(x, y) + values[feature] = ( + raw * FEATURE_DIRECTIONS[feature] if raw is not None else None + ) + return values + + independent_y = [float(row[outcome]) for row in independent] + full_yield = None + non_full_yield = None + if outcome == "lifetime_yield_fraction" and independent_y: + maximum = max(independent_y) + full_yield = sum(abs(value - maximum) <= 1e-12 for value in independent_y) + non_full_yield = len(independent_y) - full_yield + label_support = full_yield >= 3 and non_full_yield >= 3 + else: + label_support = len(set(independent_y)) >= 3 + identifiable = len(independent) >= 10 and label_support + observation_y = [float(row[outcome]) for row in eligible] + return { + "observations": len(eligible), + "unique_prompts": len(independent), + "outcome_min": min(independent_y) if independent_y else None, + "outcome_max": max(independent_y) if independent_y else None, + "identifiable": identifiable, + "minimum_recommended_unique_prompts": 10, + "full_yield_prompts": full_yield, + "non_full_yield_prompts": non_full_yield, + "observation_outcome_min": min(observation_y) if observation_y else None, + "observation_outcome_max": max(observation_y) if observation_y else None, + "observation_weighted_spearman": correlations(eligible), + "spearman_higher_is_better": correlations(independent), + } + + +def _calibration( + selectors: Iterable[tuple[int, int, str, dict[str, Any]]], +) -> list[dict[str, Any]]: + grouped: dict[ + tuple[int, int, int], + list[tuple[bool, float, float, float, float, str]], + ] = defaultdict(list) + for clients, spec_depth, prompt, selector in selectors: + lm_chain = 1.0 + selector_chain = 1.0 + for depth in selector["depths"]: + lm_chain *= math.exp(float(depth["selected_logp"])) + selector_chain *= float(depth["selector_mass"]) + grouped[(clients, spec_depth, int(depth["depth"]))].append(( + bool(depth["accepted"]), + math.exp(float(depth["selected_logp"])), + float(depth["selector_mass"]), + lm_chain, + selector_chain, + prompt, + )) + out = [] + for (clients, spec_depth, depth), rows in sorted(grouped.items()): + labels = [float(row[0]) for row in rows] + observed = statistics.fmean(labels) + by_prompt: dict[str, list[bool]] = defaultdict(list) + for row in rows: + by_prompt[row[5]].append(bool(row[0])) + + def metric(index: int) -> dict[str, float]: + predicted = [row[index] for row in rows] + mean = statistics.fmean(predicted) + return { + "mean_raw_probability": mean, + "observed_minus_raw": observed - mean, + "brier": statistics.fmean( + (prediction - label) ** 2 + for prediction, label in zip(predicted, labels) + ), + } + + rejected_prompts = sum( + any(not label for label in prompt_labels) + for prompt_labels in by_prompt.values() + ) + out.append({ + "clients": clients, + "spec_depth": spec_depth, + "proposal_depth": depth, + "observations": len(rows), + "unique_prompts": len(by_prompt), + "prompts_with_rejection": rejected_prompts, + "prompts_always_accepted": len(by_prompt) - rejected_prompts, + "accepted": sum(int(value) for value in labels), + "observed_survival": observed, + "selected_token_probability": metric(1), + "selector_mass": metric(2), + "lm_chain_probability": metric(3), + "selector_chain_probability": metric(4), + "has_both_labels": len(set(labels)) > 1, + }) + return out + + +def _mask_for_positions(clients: int, positions: Iterable[int]) -> str: + chosen = set(positions) + return "".join("S" if index in chosen else "A" for index in range(clients)) + + +def _shapley_values(values: dict[str, float], clients: int) -> list[float]: + denominator = math.factorial(clients) + out = [] + for player in range(clients): + contribution = 0.0 + others = [index for index in range(clients) if index != player] + for size in range(clients): + weight = ( + math.factorial(size) * math.factorial(clients - size - 1) + / denominator + ) + for subset in itertools.combinations(others, size): + without = _mask_for_positions(clients, subset) + with_player = _mask_for_positions(clients, (*subset, player)) + contribution += weight * (values[with_player] - values[without]) + out.append(contribution) + return out + + +def _relative_regret(oracle: float, value: float | None) -> float | None: + if value is None or oracle <= 0: + return None + return max(0.0, (oracle - value) / oracle) + + +def _subset_group(cases: list[dict[str, Any]]) -> dict[str, Any]: + clients = int(cases[0]["clients"]) + depth = int(cases[0]["spec_depth"]) + by_mask: dict[str, list[float]] = defaultdict(list) + case_by_mask: dict[str, list[dict[str, Any]]] = defaultdict(list) + for case in cases: + by_mask[case["mask"]].append(float(case["aggregate_tok_s"])) + case_by_mask[case["mask"]].append(case) + values = {mask: statistics.median(rows) for mask, rows in by_mask.items()} + walls = { + mask: statistics.median(float(case["wall_s"]) for case in rows) + for mask, rows in case_by_mask.items() + } + + def timing_values(scope: str) -> dict[str, float]: + out = {} + for mask, rows in case_by_mask.items(): + samples = [ + case["round_timing"][scope]["goodput_tok_s"] for case in rows + if case["round_timing"][scope]["goodput_tok_s"] is not None + ] + if samples: + out[mask] = statistics.median(float(value) for value in samples) + return dict(sorted(out.items())) + + expected = { + "".join(bits) for bits in itertools.product("AS", repeat=clients) + } + missing = sorted(expected - set(values)) + result: dict[str, Any] = { + "clients": clients, + "spec_depth": depth, + "prompt_set_sha256": cases[0].get("prompt_set_sha256"), + "mask_goodput_tok_s": dict(sorted(values.items())), + "mask_makespan_s": dict(sorted(walls.items())), + "mask_all_round_goodput_tok_s": timing_values("all"), + "mask_full_live_goodput_tok_s": timing_values("full_live"), + "mask_tail_goodput_tok_s": timing_values("tail"), + "complete_exhaustive": not missing, + "missing_masks": missing, + } + if missing: + return result + + oracle_mask = max(values, key=values.get) + oracle = values[oracle_mask] + all_ar_mask = "A" * clients + all_spec_mask = "S" * clients + mixed = {mask: value for mask, value in values.items() if "A" in mask and "S" in mask} + best_mixed_mask = max(mixed, key=mixed.get) if mixed else None + homogeneous_mask = max((all_ar_mask, all_spec_mask), key=values.get) + homogeneous = values[homogeneous_mask] + result.update({ + "oracle_mask": oracle_mask, + "oracle_goodput_tok_s": oracle, + "all_ar_ratio": values[all_ar_mask] / oracle, + "all_spec_ratio": values[all_spec_mask] / oracle, + "best_homogeneous_mask": homogeneous_mask, + "best_homogeneous_regret": _relative_regret(oracle, homogeneous), + "best_mixed_mask": best_mixed_mask, + "best_mixed_regret": _relative_regret( + oracle, values.get(best_mixed_mask) if best_mixed_mask else None, + ), + "oracle_is_homogeneous": oracle_mask in (all_ar_mask, all_spec_mask), + "homogeneous_dominates_every_mixed": ( + bool(mixed) and homogeneous > max(mixed.values()) + ), + }) + + shapley = _shapley_values(values, clients) + marginals = [] + sign_flips = 0 + for player in range(clients): + deltas = [] + others = [index for index in range(clients) if index != player] + for size in range(clients): + for subset in itertools.combinations(others, size): + without = _mask_for_positions(clients, subset) + with_player = _mask_for_positions(clients, (*subset, player)) + deltas.append(values[with_player] - values[without]) + flip = min(deltas) < 0.0 < max(deltas) + sign_flips += int(flip) + marginals.append({ + "position": player, + "shapley_goodput_tok_s": shapley[player], + "marginal_min_tok_s": min(deltas), + "marginal_max_tok_s": max(deltas), + "marginal_sign_flips_with_peer_modes": flip, + }) + result["per_request_marginals"] = marginals + result["requests_with_contextual_sign_flip"] = sign_flips + + source_cases = case_by_mask[all_spec_mask] + source = source_cases[0] + by_position = {row["position"]: row for row in source["requests"]} + prefix_results = {} + benefit_rows = [] + for position in range(clients): + request = by_position[position] + benefit_rows.append({ + **request, + "shapley_goodput_tok_s": shapley[position], + }) + for feature, direction in FEATURE_DIRECTIONS.items(): + ranking = sorted( + range(clients), + key=lambda position: direction * float( + by_position[position]["first_features"][feature] + ), + reverse=True, + ) + by_k = [] + for count in range(clients + 1): + prefix_mask = _mask_for_positions(clients, ranking[:count]) + candidates = { + mask: value for mask, value in values.items() if mask.count("S") == count + } + best_mask = max(candidates, key=candidates.get) + by_k.append({ + "spec_requests": count, + "ranked_prefix_mask": prefix_mask, + "best_same_size_mask": best_mask, + "same_size_regret": _relative_regret( + candidates[best_mask], values[prefix_mask], + ), + }) + prefix_results[feature] = { + "ranked_positions": ranking, + "by_subset_size": by_k, + } + result["raw_feature_prefix_rankings"] = prefix_results + result["first_feature_vs_shapley"] = _correlations( + benefit_rows, "shapley_goodput_tok_s", + ) + return result + + +def analyze(paths: Iterable[Path]) -> dict[str, Any]: + bench_paths: set[Path] = set() + for path in paths: + if path.is_file(): + bench_paths.add(path) + elif path.is_dir(): + bench_paths.update(path.rglob("bench.json")) + else: + raise ValueError(f"input does not exist: {path}") + cases = [] + skipped = [] + for path in sorted(bench_paths): + try: + header = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"cannot read {path}: {exc}") from exc + if header.get("kind") != "dflash2-forced-subset-diagnostic": + skipped.append(str(path)) + continue + # Warmups have the same artifact kind but live below a warmup directory. + if "warmup" in path.parts: + skipped.append(str(path)) + continue + cases.append(analyze_artifact(path)) + if not cases: + raise ValueError("no measured DFlash2 forced-subset artifacts found") + + request_rows = [ + {**row, "clients": case["clients"], "spec_depth": case["spec_depth"], + "mask": case["mask"], "artifact": case["path"]} + for case in cases for row in case["requests"] if row["mode"] == "speculation" + ] + all_selectors = [ + (int(row["clients"]), int(row["spec_depth"]), + str(row.get("prompt_sha256")), selector) + for row in request_rows for selector in row["selectors"] + ] + first_selectors = [ + (int(row["clients"]), int(row["spec_depth"]), + str(row.get("prompt_sha256")), row["selectors"][0]) + for row in request_rows + ] + + grouped_cases: dict[tuple[int, int, Any], list[dict[str, Any]]] = defaultdict(list) + for case in cases: + grouped_cases[( + int(case["clients"]), int(case["spec_depth"]), + case.get("prompt_set_sha256"), + )].append(case) + subset_groups = [ + _subset_group(rows) for _key, rows in sorted(grouped_cases.items()) + ] + + correlation_groups = [] + grouped_requests: dict[tuple[int, int], list[dict[str, Any]]] = defaultdict(list) + for row in request_rows: + grouped_requests[(int(row["clients"]), int(row["spec_depth"]))].append(row) + for (clients, depth), rows in sorted(grouped_requests.items()): + correlation_groups.append({ + "clients": clients, + "spec_depth": depth, + **_correlations(rows, "lifetime_yield_fraction"), + }) + + complete = [row for row in subset_groups if row["complete_exhaustive"]] + sign_flips = sum(row.get("requests_with_contextual_sign_flip", 0) for row in complete) + homogeneous_wins = sum( + int(row.get("homogeneous_dominates_every_mixed", False)) for row in complete + ) + calibration_all = _calibration(all_selectors) + calibration_first = _calibration(first_selectors) + has_negative_acceptance = any( + row["observed_survival"] < 1.0 for row in calibration_all + ) + calibration_label_support = any( + row["prompts_with_rejection"] >= 3 + and row["prompts_always_accepted"] >= 3 + for row in calibration_all + ) + benefit_identifiable = any( + bool(row.get("first_feature_vs_shapley", {}).get("identifiable")) + for row in complete + ) + yield_identifiable = any(row["identifiable"] for row in correlation_groups) + return { + "schema_version": 1, + "artifact_count": len(cases), + "skipped_artifacts": skipped, + "spec_request_observations": len(request_rows), + "unique_spec_prompts": len({row.get("prompt_sha256") for row in request_rows}), + "acceptance_calibration": { + "all_blocks": calibration_all, + "first_blocks": calibration_first, + "contains_negative_acceptance_labels": has_negative_acceptance, + "independent_label_support": calibration_label_support, + }, + "first_feature_vs_lifetime_yield": correlation_groups, + "subset_oracle": subset_groups, + "evidence_assessment": { + "yield_rank_identifiable": yield_identifiable, + "benefit_rank_identifiable": benefit_identifiable, + "complete_exhaustive_groups": len(complete), + "groups_where_homogeneous_beats_every_mixed_subset": homogeneous_wins, + "request_positions_with_context_dependent_marginal_sign": sign_flips, + "runtime_calibrator_ready": ( + calibration_label_support and yield_identifiable + and benefit_identifiable + ), + "interpretation": ( + "Raw selector values are proposal features, not confidence. " + "Fit only on held-out request-level benefit labels, stratified " + "by executed concurrency/depth shape; keep cohort utility in the " + "activation objective." + ), + }, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("inputs", nargs="+", type=Path) + parser.add_argument("--out", type=Path) + parser.add_argument("--compact", action="store_true") + args = parser.parse_args(argv) + report = analyze(args.inputs) + text = json.dumps( + report, sort_keys=True, indent=None if args.compact else 2, + separators=(",", ":") if args.compact else None, + ) + "\n" + if args.out is None: + print(text, end="") + else: + args.out.write_text(text, encoding="utf-8") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/analyze_gate_decisions.py b/harness/benchmarks/concurrency/analyze_gate_decisions.py index 1e764fc2c..c70e9e26d 100644 --- a/harness/benchmarks/concurrency/analyze_gate_decisions.py +++ b/harness/benchmarks/concurrency/analyze_gate_decisions.py @@ -30,7 +30,8 @@ ) SCORE_RE = re.compile( rf"(?P\d+):(?P{NUMBER}|nan)/" - r"(?P[a-z_-]+)(?P\*?)" + r"(?P[a-z_-]+)" + r"(?:/(?P[a-z0-9_-]+))?(?P\*?)" ) METRIC_RE = re.compile(r"\[concurrency-metrics\] (?P\{.*\})") TIMING_RE = re.compile(r"\[step-timing\] (?P\{.*\})") @@ -38,6 +39,20 @@ TIMING_COUNT_FIELDS = ( "live", "k", "emitted_tokens", "accepted_tokens", "target_forwards", ) +ACTIVATION_SCORE_KINDS = { + "dspark_confidence", + "dflash2_selector_benefit_v1", + "unspecified", +} +ACTIVATION_FALLBACK_REASONS = { + "confidence_evaluation_failed", + "benefit_evaluation_failed", + "benefit_adapter_unavailable", + "benefit_adapter_invalid_config", + "cost_profile_unavailable", + "activation_evaluation_failed", +} +TYPED_REASON_RE = re.compile(r"^[a-z][a-z0-9]*(?:_[a-z0-9]+)*$") def _json_object(match: re.Match[str], path: Path, line_no: int) -> dict[str, Any]: @@ -73,6 +88,94 @@ def _validate_timing(row: dict[str, Any], path: Path, line_no: int) -> None: ) +def _validate_typed_activation( + row: dict[str, Any], path: Path, line_no: int, +) -> None: + required = ( + "initial_confidence", "activation_score", "request_benefit", + "score_kind", "expected_yield", "decision_reason", + ) + for key in required: + if key not in row: + raise ValueError( + f"{path}:{line_no}: typed spec-activation {key} is required" + ) + kind = row["score_kind"] + if kind not in ACTIVATION_SCORE_KINDS: + raise ValueError( + f"{path}:{line_no}: spec-activation score_kind is unsupported" + ) + reason = row["decision_reason"] + if not isinstance(reason, str) or not TYPED_REASON_RE.fullmatch(reason): + raise ValueError( + f"{path}:{line_no}: spec-activation decision_reason must be a " + "nonempty snake-case reason" + ) + + evaluation = row["evaluation"] + score_fields = ("activation_score", "request_benefit", "expected_yield") + if evaluation == "scored": + if kind == "unspecified": + raise ValueError( + f"{path}:{line_no}: scored spec-activation score_kind must " + "identify its scoring model" + ) + for key in score_fields: + value = row[key] + if ( + type(value) not in (int, float) + or not math.isfinite(value) + or value < 1.0 + ): + raise ValueError( + f"{path}:{line_no}: spec-activation {key} must be finite " + "and at least 1 for a scored evaluation" + ) + confidence = row["initial_confidence"] + if kind == "dspark_confidence": + if ( + type(confidence) not in (int, float) + or not math.isfinite(confidence) + or confidence < 1.0 + ): + raise ValueError( + f"{path}:{line_no}: DSpark initial_confidence must be " + "finite and at least 1" + ) + elif confidence is not None: + raise ValueError( + f"{path}:{line_no}: non-DSpark initial_confidence must be null" + ) + if row["fallback_reason"] is not None: + raise ValueError( + f"{path}:{line_no}: scored spec-activation fallback_reason " + "must be null" + ) + return + + if row["initial_confidence"] is not None or any( + row[key] is not None for key in score_fields + ): + raise ValueError( + f"{path}:{line_no}: failed spec-activation scores must be null" + ) + if row["decision"] != "ar": + raise ValueError( + f"{path}:{line_no}: failed spec-activation decision must be ar" + ) + if reason != "evaluation_failed": + raise ValueError( + f"{path}:{line_no}: failed spec-activation decision_reason must " + "be evaluation_failed" + ) + fallback = row["fallback_reason"] + if fallback not in ACTIVATION_FALLBACK_REASONS: + raise ValueError( + f"{path}:{line_no}: failed spec-activation fallback_reason must " + "be a recognized typed reason" + ) + + def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: for key in ("request_id", "slot"): value = row.get(key) @@ -97,6 +200,13 @@ def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: raise ValueError( f"{path}:{line_no}: spec-activation fallback_reason is required" ) + typed_keys = ( + "score_kind", "activation_score", "request_benefit", "decision_reason", + ) + if any(key in row for key in typed_keys): + _validate_typed_activation(row, path, line_no) + return + for key in ("initial_confidence", "expected_yield"): if key not in row: raise ValueError( @@ -164,6 +274,7 @@ def parse_server_log( else math.nan ), "source": score.group("source"), + "score_kind": score.group("score_kind"), "admitted": score.group("admitted") == "*", }) rounds.append({ @@ -530,6 +641,15 @@ def _activation_proof( ) for evaluation in ("scored", "failed") } + score_kind_counts: dict[str, int] = defaultdict(int) + fallback_reason_counts: dict[str, int] = defaultdict(int) + decision_reason_counts: dict[str, int] = defaultdict(int) + for row in activations: + score_kind_counts[str(row.get("score_kind") or "legacy_confidence")] += 1 + if isinstance(row.get("fallback_reason"), str): + fallback_reason_counts[row["fallback_reason"]] += 1 + if isinstance(row.get("decision_reason"), str): + decision_reason_counts[row["decision_reason"]] += 1 return ({ "required": required, "validation": "passed" if required else "not-required", @@ -540,6 +660,9 @@ def _activation_proof( "matched_requests": len(set(by_id) & expected_ids), "decision_counts": decisions, "evaluation_counts": evaluations, + "score_kind_counts": dict(sorted(score_kind_counts.items())), + "fallback_reason_counts": dict(sorted(fallback_reason_counts.items())), + "decision_reason_counts": dict(sorted(decision_reason_counts.items())), }, unique) @@ -582,7 +705,11 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A } per_request: dict[int, dict[str, Any]] = defaultdict( - lambda: {"rounds": 0, "admitted": 0, "score_sum": 0.0, "scored": 0} + lambda: { + "rounds": 0, "admitted": 0, + "activation_score_sum": 0.0, "activation_scored": 0, + "confidence_score_sum": 0.0, "confidence_scored": 0, + } ) k_histogram: dict[int, int] = defaultdict(int) for entry in rounds: @@ -592,9 +719,17 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A stats["rounds"] += 1 if score["admitted"]: stats["admitted"] += 1 - if score["source"] == "confidence" and math.isfinite(score["score"]): - stats["score_sum"] += score["score"] - stats["scored"] += 1 + if math.isfinite(score["score"]): + if score["source"] in ("confidence", "current", "initial"): + stats["activation_score_sum"] += score["score"] + stats["activation_scored"] += 1 + if ( + (score["source"] == "confidence" and + score["score_kind"] is None) + or score["score_kind"] == "dspark_confidence" + ): + stats["confidence_score_sum"] += score["score"] + stats["confidence_scored"] += 1 requests = [] for engine_id, prompt in sorted( @@ -612,8 +747,13 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A "admitted_fraction": ( stats["admitted"] / stats["rounds"] if stats["rounds"] else 0.0 ), + "mean_activation_score": ( + stats["activation_score_sum"] / stats["activation_scored"] + if stats["activation_scored"] else None + ), "mean_confidence_yield": ( - stats["score_sum"] / stats["scored"] if stats["scored"] else None + stats["confidence_score_sum"] / stats["confidence_scored"] + if stats["confidence_scored"] else None ), "activation_slot": ( activation.get("slot") if activation is not None else None @@ -622,6 +762,18 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A activation.get("initial_confidence") if activation is not None else None ), + "activation_score": ( + activation.get("activation_score") + if activation is not None else None + ), + "request_benefit": ( + activation.get("request_benefit") + if activation is not None else None + ), + "activation_score_kind": ( + activation.get("score_kind") + if activation is not None else None + ), "expected_yield": ( activation.get("expected_yield") if activation is not None else None @@ -637,6 +789,10 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A activation.get("fallback_reason") if activation is not None else None ), + "activation_decision_reason": ( + activation.get("decision_reason") + if activation is not None else None + ), "commit_per_spec_step": ( (steps + accepted) / steps if steps else None ), @@ -786,6 +942,15 @@ def compare_prompts(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: "activation_fallback_reason": request.get( "activation_fallback_reason" ), + "activation_decision_reason": request.get( + "activation_decision_reason" + ), + "activation_score_kind": request.get( + "activation_score_kind" + ), + "activation_score": request.get("activation_score"), + "request_benefit": request.get("request_benefit"), + "mean_activation_score": request.get("mean_activation_score"), "initial_confidence": request.get("initial_confidence"), "expected_yield": request.get("expected_yield"), "mean_confidence_yield": request.get( diff --git a/harness/benchmarks/concurrency/forced_subset_benchmark.py b/harness/benchmarks/concurrency/forced_subset_benchmark.py new file mode 100755 index 000000000..17ed03c1b --- /dev/null +++ b/harness/benchmarks/concurrency/forced_subset_benchmark.py @@ -0,0 +1,597 @@ +#!/usr/bin/env python3 +"""Forced AR/speculation subset diagnostic for concurrent DFlash2. + +This is deliberately not an adaptive benchmark. Every request carries an +explicit ``decode_mode`` and the report fails closed unless server telemetry +proves that the requested mixed mode, active chain depth, and live concurrency +actually executed during the measured window. +""" + +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 + +import concurrent_benchmark as base + + +CLIENT_SCRIPT = Path(__file__).resolve() +PROFILE_PREFIXES = { + "rounds": "[step-timing] ", + "selectors": "[spec-selector] ", + "activations": "[spec-activation] ", + "requests": "[concurrency-metrics] ", +} +MODES = ("ar", "speculation") + + +def digest_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def canonical_digest(value: Any) -> str: + wire = json.dumps(value, ensure_ascii=False, separators=(",", ":")) + return base.sha256_text(wire) + + +def client_provenance(argv: list[str] | None = None) -> dict[str, Any]: + process_argv = list(sys.orig_argv if argv is None else argv) + if not process_argv or not all(isinstance(item, str) for item 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": digest_bytes(CLIENT_SCRIPT.read_bytes()), + } + + +def parse_request_modes(raw: str, clients: int) -> list[str]: + modes = [item.strip() for item in raw.split(",")] + if len(modes) != clients or any(item not in MODES for item in modes): + raise ValueError( + "--request-modes must provide exactly one ar/speculation mode " + "per client; adaptive is intentionally out of scope" + ) + return modes + + +def validate_server_metadata( + metadata: dict[str, Any], clients: int, spec_depth: int, + prompt_offset: int, require_selector: bool, +) -> None: + if type(metadata.get("clients")) is not int or metadata["clients"] != clients: + raise ValueError("server metadata clients does not match --clients") + launch = metadata.get("launch_environment") + if not isinstance(launch, dict): + raise ValueError("server metadata lacks launch_environment") + expected = { + "DFLASH_SPEC_CHAIN_DEPTH": str(spec_depth), + "DFLASH_STEP_TIMING": "1", + "PROMPT_OFFSET": str(prompt_offset), + } + if require_selector: + expected["DFLASH_DFLASH2_SELECTOR_LOG"] = "1" + for key, value in expected.items(): + if launch.get(key) != value: + raise ValueError( + f"server metadata must record {key}={value}; got " + f"{launch.get(key)!r}" + ) + + +def stream_request( + args: argparse.Namespace, prompt: str, decode_mode: 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": 0.0, + "seed": args.seed, + "ignore_eos": True, + "decode_mode": decode_mode, + } + 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"]) + 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 evidence 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, + "decode_mode": decode_mode, + "request_payload_sha256": canonical_digest(payload), + "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"), + "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), + "combined_output_sha256": canonical_digest([output, reasoning_output]), + "content_chars": len(output), + "reasoning_content_chars": len(reasoning_output), + "request_output_tok_s": ( + completion_tokens / (ended - started) + if type(completion_tokens) is int and ended > started else None + ), + "request_decode_tok_s": request_decode_tok_s, + } + + +def run_level( + args: argparse.Namespace, prompts: list[str], modes: list[str], +) -> dict[str, Any]: + selected = base.request_prompts(prompts, args.clients, args.prompt_offset) + barrier = threading.Barrier(args.clients + 1) + records: list[dict[str, Any] | None] = [None] * args.clients + worker_errors: list[BaseException | None] = [None] * args.clients + + def worker(index: int) -> None: + try: + barrier.wait(timeout=min(args.timeout, 60.0)) + record = stream_request(args, selected[index], modes[index]) + record["request_index"] = index + record["prompt_index"] = args.prompt_offset + index + record["prompt_sha256"] = base.sha256_text(selected[index]) + records[index] = record + except BaseException as exc: # surfaced in the main thread + worker_errors[index] = exc + + threads = [ + threading.Thread(target=worker, args=(index,), daemon=True) + for index in range(args.clients) + ] + for thread in threads: + thread.start() + barrier.wait(timeout=min(args.timeout, 60.0)) + barrier_released = time.perf_counter() + deadline = time.monotonic() + args.timeout + 30.0 + for thread in threads: + thread.join(max(0.0, deadline - time.monotonic())) + hung = sum(thread.is_alive() for thread in threads) + if hung: + raise TimeoutError(f"{hung} request worker(s) exceeded the level deadline") + first_worker_error = next((error for error in worker_errors if error), None) + if first_worker_error is not None: + raise RuntimeError(f"request worker failed: {first_worker_error}") + completed = [record for record in records if record is not None] + if len(completed) != args.clients: + raise RuntimeError("not every synchronized request worker returned a record") + + starts = [float(record["t_start"]) for record in completed] + ends = [float(record["t_end"]) for record in completed] + level_start = min(starts) + wall = max(ends) - level_start + for record in completed: + record["start_offset_s"] = float(record["t_start"]) - level_start + record["barrier_release_offset_s"] = ( + float(record["t_start"]) - barrier_released + ) + ok = [record for record in completed if record["error"] is None] + completion = [record["completion_tokens"] for record in ok] + prompt_counts = [record["prompt_tokens"] for record in ok] + complete_tokens = bool(ok) and all(type(value) is int for value in completion) + complete_prompts = bool(ok) and all(type(value) is int for value in prompt_counts) + decode_rates = [ + record["request_decode_tok_s"] for record in ok + if type(record.get("request_decode_tok_s")) in (int, float) + ] + prompt_hashes = [record["prompt_sha256"] for record in completed] + output_hashes = [ + [record["content_sha256"], record["reasoning_content_sha256"]] + for record in completed + ] + return { + "clients": args.clients, + "request_modes": modes, + "request_mode_mask": "".join( + "A" if mode == "ar" else "S" for mode in modes + ), + "requests": args.clients, + "requests_ok": len(ok), + "failures": args.clients - len(ok), + "wall_s": wall, + "start_skew_s": max(starts) - min(starts), + "completion_tokens_total": sum(completion) if complete_tokens else None, + "token_count_complete": complete_tokens, + "prompt_token_count_complete": complete_prompts, + "fixed_token_workload_valid": ( + len(ok) == args.clients and complete_tokens + and all(value == args.max_tokens for value in completion) + ), + "aggregate_tok_s": ( + sum(completion) / wall if complete_tokens and wall > 0 else None + ), + "request_decode_tok_s_median": ( + statistics.median(decode_rates) + if len(decode_rates) == len(ok) and ok else None + ), + "prompt_tokens_total": sum(prompt_counts) if complete_prompts else None, + "selected_prompt_set_sha256": canonical_digest(prompt_hashes), + "selected_output_set_sha256": canonical_digest(output_hashes), + "requests_detail": completed, + } + + +def parse_profile_records(data: bytes) -> dict[str, list[dict[str, Any]]]: + records = {key: [] for key in PROFILE_PREFIXES} + text = data.decode("utf-8", errors="replace") + for line_index, line in enumerate(text.splitlines(), 1): + for key, prefix in PROFILE_PREFIXES.items(): + marker = line.find(prefix) + if marker < 0: + continue + raw = line[marker + len(prefix):] + try: + value = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError( + f"measured server log line {line_index}: invalid {prefix.strip()} " + f"JSON: {exc}" + ) from exc + if not isinstance(value, dict): + raise ValueError( + f"measured server log line {line_index}: {prefix.strip()} " + "record must be an object" + ) + records[key].append({ + "line_index": line_index, + "raw_json": raw, + "record": value, + }) + break + return records + + +def read_log_span(path: Path, start: int) -> tuple[bytes, int]: + size = path.stat().st_size + if size < start: + raise ValueError("server log was truncated or rotated during the benchmark") + with path.open("rb") as handle: + handle.seek(start) + return handle.read(size - start), size + + +def _longest_full_live_streak(rounds: Iterable[dict[str, Any]], clients: int) -> int: + longest = 0 + current = 0 + for wrapped in rounds: + row = wrapped["record"] + if type(row.get("live")) is int and row["live"] == clients: + current += 1 + longest = max(longest, current) + else: + current = 0 + return longest + + +def validate_evidence( + level: dict[str, Any], records: dict[str, list[dict[str, Any]]], + clients: int, modes: list[str], spec_depth: int, + max_start_skew_ms: float, min_full_live_rounds: int, +) -> dict[str, Any]: + errors: list[str] = [] + if level["failures"] or level["requests_ok"] != clients: + errors.append("one or more requests failed") + if not level["token_count_complete"] or not level["prompt_token_count_complete"]: + errors.append("wire token accounting is incomplete") + if level["fixed_token_workload_valid"] is not True: + errors.append("ignore-eos fixed-token workload was not completed exactly") + if float(level["start_skew_s"]) * 1000.0 > max_start_skew_ms: + errors.append( + f"request start skew exceeds {max_start_skew_ms:g} ms" + ) + request_ids = [row.get("request_id") for row in level["requests_detail"]] + if ( + any(not isinstance(value, str) or not value for value in request_ids) + or len(set(request_ids)) != clients + ): + errors.append("wire request IDs are missing or not unique") + + rounds = records["rounds"] + full_live_rounds = sum( + wrapped["record"].get("live") == clients for wrapped in rounds + ) + longest_streak = _longest_full_live_streak(rounds, clients) + if longest_streak < min_full_live_rounds: + errors.append( + f"sustained live=C proof absent: longest live={clients} streak " + f"is {longest_streak}, need {min_full_live_rounds}" + ) + invalid_live = [ + wrapped["record"].get("live") for wrapped in rounds + if type(wrapped["record"].get("live")) is not int + or wrapped["record"]["live"] < 1 + or wrapped["record"]["live"] > clients + ] + if invalid_live: + errors.append(f"step-timing contains invalid live values: {invalid_live}") + + spec_requested = any(mode == "speculation" for mode in modes) + spec_rounds = [ + wrapped["record"] for wrapped in rounds + if wrapped["record"].get("path") == "spec" + and type(wrapped["record"].get("k")) is int + and wrapped["record"]["k"] > 0 + ] + inferred_depths: list[int] = [] + for row in spec_rounds: + bucket = row.get("tree_bucket") + tree_rows = row.get("tree_rows") + if ( + type(bucket) is not int or bucket <= 0 + or type(tree_rows) is not int or tree_rows <= 0 + or tree_rows % bucket != 0 + ): + errors.append("spec step-timing lacks a valid tree_rows/tree_bucket shape") + continue + inferred_depths.append(tree_rows // bucket) + if spec_requested: + if not spec_rounds: + errors.append("speculation was requested but no spec round executed") + if not records["selectors"]: + errors.append("speculation was requested but no DFlash2 selector records exist") + wrong_depths = sorted(set(depth for depth in inferred_depths if depth != spec_depth)) + if wrong_depths: + errors.append( + f"executed chain depths {wrong_depths} do not match requested " + f"depth {spec_depth}" + ) + elif spec_rounds: + errors.append("all-AR mask unexpectedly executed speculation") + + metrics_by_id: dict[str, dict[str, Any]] = {} + for wrapped in records["requests"]: + row = wrapped["record"] + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + errors.append("concurrency metric has no wire request_id") + elif request_id in metrics_by_id: + errors.append(f"duplicate concurrency metric for {request_id}") + else: + metrics_by_id[request_id] = row + missing = sorted(set(request_ids) - set(metrics_by_id)) + extra = sorted(set(metrics_by_id) - set(request_ids)) + if missing: + errors.append(f"missing per-request concurrency metrics: {missing}") + if extra: + errors.append(f"unmatched per-request concurrency metrics: {extra}") + for request, requested_mode in zip(level["requests_detail"], modes): + metric = metrics_by_id.get(request.get("request_id")) + if metric is None: + continue + steps = metric.get("spec_steps") + if type(steps) is not int or steps < 0: + errors.append(f"{request.get('request_id')}: invalid spec_steps") + elif requested_mode == "speculation" and steps == 0: + errors.append(f"{request.get('request_id')}: forced speculation never executed") + elif requested_mode == "ar" and steps != 0: + errors.append(f"{request.get('request_id')}: forced AR executed speculation") + + output_hashes_complete = all( + isinstance(request.get(key), str) and len(request[key]) == 64 + for request in level["requests_detail"] + for key in ( + "content_sha256", "reasoning_content_sha256", + "combined_output_sha256", + ) + ) + if not output_hashes_complete: + errors.append("exact request output hashes are incomplete") + return { + "passed": not errors, + "errors": errors, + "adaptive_claims_permitted": False, + "full_live_rounds": full_live_rounds, + "longest_full_live_streak": longest_streak, + "min_full_live_rounds": min_full_live_rounds, + "executed_spec_depths": sorted(set(inferred_depths)), + "round_records": len(rounds), + "selector_records": len(records["selectors"]), + "request_metric_records": len(records["requests"]), + } + + +def markdown(report: dict[str, Any]) -> str: + level = report["level"] + validation = report["validation"] + status = "PASS" if validation["passed"] else "FAIL" + return ( + f"# Forced DFlash2 subset diagnostic — {report['label']}\n\n" + "This is a forced-control diagnostic; it makes no adaptive result claim.\n\n" + "| C | Mask | Depth | Ok | Goodput tok/s | Full-live streak | Status |\n" + "| ---: | :--- | ---: | ---: | ---: | ---: | :--- |\n" + f"| {level['clients']} | {level['request_mode_mask']} | " + f"{report['spec_depth']} | {level['requests_ok']}/{level['requests']} | " + f"{base.fmt(level['aggregate_tok_s'])} | " + f"{validation['longest_full_live_streak']} | {status} |\n" + ) + + +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, required=True) + parser.add_argument("--request-modes", required=True) + parser.add_argument("--spec-depth", type=int, required=True) + parser.add_argument("--prompt-file", type=Path, required=True) + parser.add_argument("--prompt-offset", type=int, default=0) + parser.add_argument("--max-tokens", type=int, default=256) + parser.add_argument("--seed", type=int, default=1) + parser.add_argument("--timeout", type=float, default=1800.0) + parser.add_argument("--max-start-skew-ms", type=float, default=100.0) + parser.add_argument("--min-full-live-rounds", type=int, default=2) + parser.add_argument("--log-settle-ms", type=float, default=100.0) + parser.add_argument("--server-metadata-json", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--label", default="") + return parser + + +def run(args: argparse.Namespace) -> int: + if args.clients < 1: + raise ValueError("--clients must be positive") + if args.spec_depth < 2: + raise ValueError("--spec-depth must be at least 2; depth 1 is AR-equivalent") + if args.prompt_offset < 0 or args.max_tokens < 1 or args.timeout <= 0: + raise ValueError("invalid prompt offset, max tokens, or timeout") + if args.max_start_skew_ms < 0 or args.min_full_live_rounds < 1: + raise ValueError("invalid synchronization proof threshold") + if args.log_settle_ms < 0: + raise ValueError("--log-settle-ms must be non-negative") + modes = parse_request_modes(args.request_modes, args.clients) + prompts = base.load_prompts(args.prompt_file) + metadata_bytes = args.server_metadata_json.read_bytes() + metadata = json.loads(metadata_bytes) + if not isinstance(metadata, dict): + raise ValueError("server metadata must be a JSON object") + validate_server_metadata( + metadata, args.clients, args.spec_depth, args.prompt_offset, + require_selector=any(mode == "speculation" for mode in modes), + ) + log_start = args.server_log.stat().st_size + level = run_level(args, prompts, modes) + if args.log_settle_ms: + time.sleep(args.log_settle_ms / 1000.0) + log_span, log_end = read_log_span(args.server_log, log_start) + records = parse_profile_records(log_span) + validation = validate_evidence( + level, records, args.clients, modes, args.spec_depth, + args.max_start_skew_ms, args.min_full_live_rounds, + ) + report = { + "schema_version": 1, + "kind": "dflash2-forced-subset-diagnostic", + "label": args.label, + "scope": { + "forced_controls_only": True, + "adaptive_evaluation": False, + "interpretation": ( + "This artifact may compare forced AR/speculation subsets and " + "depths; it is not an adaptive activation result." + ), + }, + "base_url": args.base_url, + "model": args.model, + "max_tokens": args.max_tokens, + "temperature": 0.0, + "seed": args.seed, + "ignore_eos": True, + "spec_depth": args.spec_depth, + "prompt_offset": args.prompt_offset, + "prompt_file": str(args.prompt_file.resolve()), + "prompt_file_sha256": digest_bytes(args.prompt_file.read_bytes()), + "server_metadata": metadata, + "server_metadata_sha256": digest_bytes(metadata_bytes), + "server_log": { + "path": str(args.server_log.resolve()), + "start_offset": log_start, + "end_offset": log_end, + "span_bytes": len(log_span), + "span_sha256": digest_bytes(log_span), + }, + "level": level, + "server_records": records, + "validation": validation, + **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="") + if not validation["passed"]: + for error in validation["errors"]: + print(f"[forced-subset] validation: {error}", file=sys.stderr) + return 0 if validation["passed"] else 1 + + +def main() -> int: + try: + return run(build_parser().parse_args()) + except Exception as exc: + print(f"[forced-subset] error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/refill_subset_benchmark.py b/harness/benchmarks/concurrency/refill_subset_benchmark.py new file mode 100644 index 000000000..8b8b0d248 --- /dev/null +++ b/harness/benchmarks/concurrency/refill_subset_benchmark.py @@ -0,0 +1,638 @@ +#!/usr/bin/env python3 +"""Sustained/refill forced AR/speculation diagnostic for concurrent DFlash2. + +Each positional lane keeps one request in flight. The first +``clients * (waves - 1)`` completions immediately refill the lane that +completed, preserving the active A/S mask until the final C-request drain. +This measures a closed-loop saturated service workload, not adaptive +activation. The report keeps end-to-end refill goodput separate from +full-live engine-round goodput and fails closed unless telemetry proves every +request-local mode, chain depth, refill recovery, token count, and exact +output hash. +""" + + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import threading +import time +from pathlib import Path +from typing import Any + +import concurrent_benchmark as base +import forced_subset_benchmark as forced + + +CLIENT_SCRIPT = Path(__file__).resolve() + + +def client_provenance(argv: list[str] | None = None) -> dict[str, Any]: + process_argv = list(sys.orig_argv if argv is None else argv) + if not process_argv or not all(isinstance(item, str) for item 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": forced.digest_bytes(CLIENT_SCRIPT.read_bytes()), + } + + +def run_refill( + args: argparse.Namespace, prompts: list[str], modes: list[str], +) -> dict[str, Any]: + """Keep C positional lanes full for a fixed global request budget.""" + selected = base.request_prompts(prompts, args.clients, args.prompt_offset) + barrier = threading.Barrier(args.clients + 1) + lane_records: list[list[dict[str, Any]]] = [ + [] for _ in range(args.clients) + ] + worker_errors: list[BaseException | None] = [None] * args.clients + refill_lock = threading.Lock() + refill_budget = args.clients * (args.waves - 1) + completed_ok = 0 + next_request_index = args.clients + abort_refills = False + + def worker(lane_index: int) -> None: + nonlocal completed_ok, next_request_index, abort_refills + try: + barrier.wait(timeout=min(args.timeout, 60.0)) + previous_end: float | None = None + lane_request_index = 0 + request_index: int | None = lane_index + while request_index is not None: + record = forced.stream_request( + args, selected[lane_index], modes[lane_index], + ) + record["lane_index"] = lane_index + record["lane_request_index"] = lane_request_index + record["request_index"] = request_index + record["admission_group_index"] = request_index // args.clients + record["prompt_index"] = args.prompt_offset + lane_index + record["prompt_sha256"] = base.sha256_text(selected[lane_index]) + record["refill_gap_s"] = ( + float(record["t_start"]) - previous_end + if previous_end is not None else None + ) + lane_records[lane_index].append(record) + previous_end = float(record["t_end"]) + lane_request_index += 1 + with refill_lock: + if record["error"] is not None: + abort_refills = True + request_index = None + else: + completed_ok += 1 + if not abort_refills and completed_ok <= refill_budget: + request_index = next_request_index + next_request_index += 1 + else: + request_index = None + except BaseException as exc: # surfaced in the main thread + with refill_lock: + abort_refills = True + worker_errors[lane_index] = exc + + threads = [ + threading.Thread(target=worker, args=(index,), daemon=True) + for index in range(args.clients) + ] + for thread in threads: + thread.start() + barrier.wait(timeout=min(args.timeout, 60.0)) + barrier_released = time.perf_counter() + deadline = time.monotonic() + args.timeout * args.clients * args.waves + 30.0 + for thread in threads: + thread.join(max(0.0, deadline - time.monotonic())) + hung = sum(thread.is_alive() for thread in threads) + if hung: + raise TimeoutError(f"{hung} refill lane(s) exceeded the workload deadline") + first_worker_error = next((error for error in worker_errors if error), None) + if first_worker_error is not None: + raise RuntimeError(f"refill request worker failed: {first_worker_error}") + + completed = sorted( + (record for lane in lane_records for record in lane), + key=lambda record: int(record["request_index"]), + ) + expected_requests = args.clients * args.waves + starts = [float(record["t_start"]) for record in completed] + ends = [float(record["t_end"]) for record in completed] + if not starts: + raise RuntimeError("no refill request returned a record") + workload_start = min(starts) + workload_end = max(ends) + wall = workload_end - workload_start + for record in completed: + record["start_offset_s"] = float(record["t_start"]) - workload_start + record["end_offset_s"] = float(record["t_end"]) - workload_start + if record["lane_request_index"] == 0: + record["barrier_release_offset_s"] = ( + float(record["t_start"]) - barrier_released + ) + + ok = [record for record in completed if record["error"] is None] + completion = [record["completion_tokens"] for record in ok] + prompt_counts = [record["prompt_tokens"] for record in ok] + complete_tokens = bool(ok) and all(type(value) is int for value in completion) + complete_prompts = bool(ok) and all( + type(value) is int for value in prompt_counts + ) + first_outputs = [ + float(record["t_first"]) for record in ok + if type(record.get("t_first")) in (int, float) + ] + output_wall = ( + workload_end - min(first_outputs) + if len(first_outputs) == len(ok) and first_outputs else None + ) + first_wave_starts = [ + float(lane[0]["t_start"]) for lane in lane_records if lane + ] + refill_gaps = [ + float(record["refill_gap_s"]) + for record in completed + if type(record.get("refill_gap_s")) in (int, float) + ] + output_hashes = [ + [ + record["request_index"], record["lane_index"], + record["lane_request_index"], + record["content_sha256"], record["reasoning_content_sha256"], + ] + for record in completed + ] + lane_summaries = [] + for lane_index, lane in enumerate(lane_records): + lane_completion = [ + record["completion_tokens"] for record in lane + if type(record.get("completion_tokens")) is int + ] + signatures = { + ( + record.get("content_sha256"), + record.get("reasoning_content_sha256"), + ) + for record in lane + } + lane_summaries.append({ + "lane_index": lane_index, + "decode_mode": modes[lane_index], + "prompt_index": args.prompt_offset + lane_index, + "requests": len(lane), + "requests_ok": sum(record["error"] is None for record in lane), + "completion_tokens_total": sum(lane_completion), + "exact_output_stable": bool(lane) and len(signatures) == 1, + "ordered_output_set_sha256": forced.canonical_digest([ + [record["content_sha256"], record["reasoning_content_sha256"]] + for record in lane + ]), + }) + return { + "clients": args.clients, + "waves": args.waves, + "scheduled_refills": refill_budget, + "terminal_guard_requests": args.clients, + "request_modes": modes, + "request_mode_mask": "".join( + "A" if mode == "ar" else "S" for mode in modes + ), + "expected_requests": expected_requests, + "requests": len(completed), + "requests_ok": len(ok), + "failures": len(completed) - len(ok), + "missing_requests": expected_requests - len(completed), + "wall_s": wall, + "initial_start_skew_s": ( + max(first_wave_starts) - min(first_wave_starts) + if len(first_wave_starts) == args.clients else None + ), + "completion_tokens_total": sum(completion) if complete_tokens else None, + "token_count_complete": complete_tokens, + "prompt_token_count_complete": complete_prompts, + "fixed_token_workload_valid": ( + len(ok) == expected_requests and complete_tokens + and all(value == args.max_tokens for value in completion) + ), + "aggregate_refill_tok_s": ( + sum(completion) / wall if complete_tokens and wall > 0 else None + ), + "output_window_s": output_wall, + "output_window_tok_s": ( + sum(completion) / output_wall + if complete_tokens and output_wall is not None and output_wall > 0 + else None + ), + "prompt_tokens_total": sum(prompt_counts) if complete_prompts else None, + "refill_handoffs": len(refill_gaps), + "refill_gap_s_median": ( + statistics.median(refill_gaps) if refill_gaps else None + ), + "refill_gap_s_max": max(refill_gaps) if refill_gaps else None, + "selected_prompt_set_sha256": forced.canonical_digest([ + base.sha256_text(prompt) for prompt in selected + ]), + "ordered_output_set_sha256": forced.canonical_digest(output_hashes), + "exact_output_stable_per_lane": all( + lane["exact_output_stable"] for lane in lane_summaries + ), + "lanes": lane_summaries, + "requests_detail": completed, + } + + +def _full_live_round_summary( + rounds: list[dict[str, Any]], clients: int, +) -> tuple[dict[str, Any], list[str]]: + errors: list[str] = [] + full_live: list[dict[str, Any]] = [] + all_timed_us = 0.0 + all_emitted = 0 + longest = 0 + current = 0 + for wrapped in rounds: + row = wrapped["record"] + live = row.get("live") + if type(live) is not int or live < 1 or live > clients: + errors.append(f"step-timing line {wrapped['line_index']} has invalid live={live!r}") + current = 0 + continue + total_us = row.get("total_us") + emitted = row.get("emitted_tokens") + if ( + type(total_us) not in (int, float) or total_us <= 0 + or type(emitted) is not int or emitted < 1 + ): + errors.append( + f"step-timing line {wrapped['line_index']} lacks positive " + "total_us/emitted_tokens" + ) + current = 0 + continue + all_timed_us += float(total_us) + all_emitted += emitted + if live == clients: + full_live.append(wrapped) + current += 1 + longest = max(longest, current) + else: + current = 0 + full_us = sum(float(row["record"]["total_us"]) for row in full_live) + full_emitted = sum(int(row["record"]["emitted_tokens"]) for row in full_live) + return ({ + "rounds": len(full_live), + "longest_streak": longest, + "timed_us": full_us, + "emitted_tokens": full_emitted, + "engine_round_goodput_tok_s": ( + full_emitted * 1_000_000.0 / full_us if full_us > 0 else None + ), + "timed_fraction": full_us / all_timed_us if all_timed_us > 0 else None, + "emitted_fraction": ( + full_emitted / all_emitted if all_emitted > 0 else None + ), + "first_line_index": full_live[0]["line_index"] if full_live else None, + "last_line_index": full_live[-1]["line_index"] if full_live else None, + "path_counts": { + path: sum(row["record"].get("path") == path for row in full_live) + for path in ("ar", "spec") + }, + }, errors) + + +def validate_evidence( + workload: dict[str, Any], records: dict[str, list[dict[str, Any]]], + clients: int, modes: list[str], spec_depth: int, waves: int, + max_start_skew_ms: float, max_refill_gap_ms: float, + min_full_live_rounds: int, +) -> dict[str, Any]: + errors: list[str] = [] + expected_requests = clients * waves + if workload["requests"] != expected_requests or workload["missing_requests"] != 0: + errors.append( + f"completed request records {workload['requests']} do not match " + f"clients*waves={expected_requests}" + ) + if workload["failures"] or workload["requests_ok"] != expected_requests: + errors.append("one or more refill requests failed") + if not workload["token_count_complete"] or not workload["prompt_token_count_complete"]: + errors.append("wire token accounting is incomplete") + if workload["fixed_token_workload_valid"] is not True: + errors.append("ignore-eos fixed-token refill workload was not completed exactly") + start_skew = workload.get("initial_start_skew_s") + if type(start_skew) not in (int, float): + errors.append("initial synchronized request start skew is unavailable") + elif float(start_skew) * 1000.0 > max_start_skew_ms: + errors.append(f"initial request start skew exceeds {max_start_skew_ms:g} ms") + expected_handoffs = clients * (waves - 1) + if workload.get("refill_handoffs") != expected_handoffs: + errors.append( + f"refill handoff count {workload.get('refill_handoffs')} does not " + f"match {expected_handoffs}" + ) + max_gap = workload.get("refill_gap_s_max") + if type(max_gap) not in (int, float): + errors.append("refill handoff latency is unavailable") + elif float(max_gap) < 0: + errors.append("a refill request started before its predecessor completed") + elif float(max_gap) * 1000.0 > max_refill_gap_ms: + errors.append(f"refill handoff gap exceeds {max_refill_gap_ms:g} ms") + if workload.get("exact_output_stable_per_lane") is not True: + errors.append("deterministic exact output hashes changed across refill waves") + + requests = workload["requests_detail"] + request_ids = [row.get("request_id") for row in requests] + if ( + any(not isinstance(value, str) or not value for value in request_ids) + or len(set(request_ids)) != expected_requests + ): + errors.append("wire request IDs are missing or not unique") + for lane_index, lane in enumerate(workload.get("lanes") or []): + if lane.get("lane_index") != lane_index or not lane.get("requests"): + errors.append(f"lane {lane_index} did not execute an initial request") + if lane.get("decode_mode") != modes[lane_index]: + errors.append(f"lane {lane_index} mode changed during refill") + output_hashes_complete = all( + isinstance(request.get(key), str) and len(request[key]) == 64 + for request in requests + for key in ( + "content_sha256", "reasoning_content_sha256", + "combined_output_sha256", + ) + ) + if not output_hashes_complete: + errors.append("exact request output hashes are incomplete") + + rounds = records["rounds"] + full_live, timing_errors = _full_live_round_summary(rounds, clients) + errors.extend(timing_errors) + if full_live["longest_streak"] < min_full_live_rounds: + errors.append( + f"sustained live=C proof absent: longest live={clients} streak is " + f"{full_live['longest_streak']}, need {min_full_live_rounds}" + ) + + spec_requested = any(mode == "speculation" for mode in modes) + spec_rounds = [ + wrapped for wrapped in rounds + if wrapped["record"].get("path") == "spec" + and type(wrapped["record"].get("k")) is int + and wrapped["record"]["k"] > 0 + ] + inferred_depths: list[int] = [] + for wrapped in spec_rounds: + row = wrapped["record"] + bucket = row.get("tree_bucket") + tree_rows = row.get("tree_rows") + if ( + type(bucket) is not int or bucket <= 0 + or type(tree_rows) is not int or tree_rows <= 0 + or tree_rows % bucket != 0 + ): + errors.append("spec step-timing lacks a valid tree_rows/tree_bucket shape") + continue + inferred_depths.append(tree_rows // bucket) + if spec_requested: + if not spec_rounds: + errors.append("speculation was requested but no spec round executed") + wrong_depths = sorted(set(depth for depth in inferred_depths if depth != spec_depth)) + if wrong_depths: + errors.append( + f"executed chain depths {wrong_depths} do not match requested " + f"depth {spec_depth}" + ) + elif spec_rounds: + errors.append("all-AR refill mask unexpectedly executed speculation") + + metrics_by_id: dict[str, dict[str, Any]] = {} + metric_wrappers_by_id: dict[str, dict[str, Any]] = {} + for wrapped in records["requests"]: + row = wrapped["record"] + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + errors.append("concurrency metric has no wire request_id") + elif request_id in metrics_by_id: + errors.append(f"duplicate concurrency metric for {request_id}") + else: + metrics_by_id[request_id] = row + metric_wrappers_by_id[request_id] = wrapped + missing = sorted(set(request_ids) - set(metrics_by_id)) + extra = sorted(set(metrics_by_id) - set(request_ids)) + if missing: + errors.append(f"missing per-request concurrency metrics: {missing}") + if extra: + errors.append(f"unmatched per-request concurrency metrics: {extra}") + + selector_engine_ids = { + wrapped["record"].get("request_id") for wrapped in records["selectors"] + if type(wrapped["record"].get("request_id")) is int + } + for request in requests: + request_id = request.get("request_id") + metric = metrics_by_id.get(request_id) + if metric is None: + continue + expected_mode = modes[int(request["lane_index"])] + steps = metric.get("spec_steps") + if type(steps) is not int or steps < 0: + errors.append(f"{request_id}: invalid spec_steps") + elif expected_mode == "speculation" and steps == 0: + errors.append(f"{request_id}: forced speculation never executed") + elif expected_mode == "ar" and steps != 0: + errors.append(f"{request_id}: forced AR executed speculation") + engine_id = metric.get("engine_request_id") + if type(engine_id) is not int: + errors.append(f"{request_id}: missing integer engine_request_id") + elif expected_mode == "speculation" and engine_id not in selector_engine_ids: + errors.append(f"{request_id}: no DFlash2 selector evidence") + elif expected_mode == "ar" and engine_id in selector_engine_ids: + errors.append(f"{request_id}: forced AR emitted DFlash2 selector evidence") + + last_full_live_line = full_live["last_line_index"] + completed_before_last_full_live = ( + sum( + wrapped["line_index"] < last_full_live_line + for wrapped in metric_wrappers_by_id.values() + ) + if type(last_full_live_line) is int else 0 + ) + scheduled_refills = expected_requests - clients + required_refill_recoveries = expected_requests - 2 * clients + refill_recovery_proved = ( + completed_before_last_full_live >= required_refill_recoveries + ) + if not refill_recovery_proved: + errors.append( + "full live=C was not recovered inside the guarded refill window: " + f"completions: proved {completed_before_last_full_live}, need " + f"{required_refill_recoveries}" + ) + return { + "passed": not errors, + "errors": errors, + "adaptive_claims_permitted": False, + "closed_cohort_claims_permitted": False, + "full_live": full_live, + "min_full_live_rounds": min_full_live_rounds, + "executed_spec_depths": sorted(set(inferred_depths)), + "round_records": len(rounds), + "selector_records": len(records["selectors"]), + "request_metric_records": len(records["requests"]), + "required_refill_recoveries": required_refill_recoveries, + "scheduled_refills": scheduled_refills, + "terminal_guard_requests": clients, + "completed_before_last_full_live": completed_before_last_full_live, + "refill_recovery_proved": refill_recovery_proved, + } + + +def markdown(report: dict[str, Any]) -> str: + workload = report["workload"] + validation = report["validation"] + full_live = validation["full_live"] + status = "PASS" if validation["passed"] else "FAIL" + return ( + f"# Forced DFlash2 refill diagnostic — {report['label']}\n\n" + "This is a forced closed-loop refill diagnostic. It makes neither an " + "adaptive claim nor a closed-cohort makespan claim.\n\n" + "| C | Mask | Waves | Depth | Requests | Refill tok/s | " + "Full-live round tok/s | Recovery | Status |\n" + "| ---: | :--- | ---: | ---: | ---: | ---: | ---: | :--- | :--- |\n" + f"| {workload['clients']} | {workload['request_mode_mask']} | " + f"{workload['waves']} | {report['spec_depth']} | " + f"{workload['requests_ok']}/{workload['expected_requests']} | " + f"{base.fmt(workload['aggregate_refill_tok_s'])} | " + f"{base.fmt(full_live['engine_round_goodput_tok_s'])} | " + f"{validation['completed_before_last_full_live']}/" + f"{validation['required_refill_recoveries']} | {status} |\n" + ) + + +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, required=True) + parser.add_argument("--request-modes", required=True) + parser.add_argument("--waves", type=int, required=True) + parser.add_argument("--spec-depth", type=int, required=True) + parser.add_argument("--prompt-file", type=Path, required=True) + parser.add_argument("--prompt-offset", type=int, default=0) + parser.add_argument("--max-tokens", type=int, default=256) + parser.add_argument("--seed", type=int, default=1) + parser.add_argument("--timeout", type=float, default=1800.0) + parser.add_argument("--max-start-skew-ms", type=float, default=100.0) + parser.add_argument("--max-refill-gap-ms", type=float, default=100.0) + parser.add_argument("--min-full-live-rounds", type=int, default=2) + parser.add_argument("--log-settle-ms", type=float, default=100.0) + parser.add_argument("--server-metadata-json", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--label", default="") + return parser + + +def run(args: argparse.Namespace) -> int: + if args.clients < 1: + raise ValueError("--clients must be positive") + if args.waves < 3: + raise ValueError("--waves must be at least 3 for a guarded refill workload") + if args.spec_depth < 2: + raise ValueError("--spec-depth must be at least 2; depth 1 is AR-equivalent") + if args.prompt_offset < 0 or args.max_tokens < 1 or args.timeout <= 0: + raise ValueError("invalid prompt offset, max tokens, or timeout") + if ( + args.max_start_skew_ms < 0 or args.max_refill_gap_ms < 0 + or args.min_full_live_rounds < 1 + ): + raise ValueError("invalid synchronization/refill proof threshold") + if args.log_settle_ms < 0: + raise ValueError("--log-settle-ms must be non-negative") + modes = forced.parse_request_modes(args.request_modes, args.clients) + prompts = base.load_prompts(args.prompt_file) + metadata_bytes = args.server_metadata_json.read_bytes() + metadata = json.loads(metadata_bytes) + if not isinstance(metadata, dict): + raise ValueError("server metadata must be a JSON object") + forced.validate_server_metadata( + metadata, args.clients, args.spec_depth, args.prompt_offset, + require_selector=any(mode == "speculation" for mode in modes), + ) + log_start = args.server_log.stat().st_size + workload = run_refill(args, prompts, modes) + if args.log_settle_ms: + time.sleep(args.log_settle_ms / 1000.0) + log_span, log_end = forced.read_log_span(args.server_log, log_start) + records = forced.parse_profile_records(log_span) + validation = validate_evidence( + workload, records, args.clients, modes, args.spec_depth, args.waves, + args.max_start_skew_ms, args.max_refill_gap_ms, + args.min_full_live_rounds, + ) + report = { + "schema_version": 1, + "kind": "dflash2-forced-refill-diagnostic", + "label": args.label, + "scope": { + "forced_controls_only": True, + "adaptive_evaluation": False, + "closed_cohort_makespan": False, + "interpretation": ( + "aggregate_refill_tok_s is end-to-end closed-loop goodput " + "across persistent client lanes and includes startup, request " + "handoffs, prefill, and terminal drain. full_live decode-round " + "goodput uses only timed live=C decode rounds and excludes " + "handoff/prefill gaps. Neither is an adaptive result or a " + "single closed-cohort makespan result." + ), + }, + "base_url": args.base_url, + "model": args.model, + "max_tokens": args.max_tokens, + "temperature": 0.0, + "seed": args.seed, + "ignore_eos": True, + "spec_depth": args.spec_depth, + "prompt_offset": args.prompt_offset, + "prompt_file": str(args.prompt_file.resolve()), + "prompt_file_sha256": forced.digest_bytes(args.prompt_file.read_bytes()), + "server_metadata": metadata, + "server_metadata_sha256": forced.digest_bytes(metadata_bytes), + "server_log": { + "path": str(args.server_log.resolve()), + "start_offset": log_start, + "end_offset": log_end, + "span_bytes": len(log_span), + "span_sha256": forced.digest_bytes(log_span), + }, + "workload": workload, + "server_records": records, + "validation": validation, + **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="") + if not validation["passed"]: + for error in validation["errors"]: + print(f"[forced-refill] validation: {error}", file=sys.stderr) + return 0 if validation["passed"] else 1 + + +def main() -> int: + try: + return run(build_parser().parse_args()) + except Exception as exc: + print(f"[forced-refill] error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh b/harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh new file mode 100755 index 000000000..1c34a22c8 --- /dev/null +++ b/harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh @@ -0,0 +1,283 @@ +#!/usr/bin/env bash +# Persistent-server forced AR/speculation subset screen for concurrent DFlash2. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +REPO="${REPO:-$(cd -- "$SCRIPT_DIR/../../.." && pwd -P)}" +CLIENT="${CLIENT:-$SCRIPT_DIR/forced_subset_benchmark.py}" +REFILL_CLIENT="${REFILL_CLIENT:-$SCRIPT_DIR/refill_subset_benchmark.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:-}" +PROMPT_FILE="${PROMPT_FILE:-}" +PROMPT_OFFSET="${PROMPT_OFFSET:-0}" +LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" +OUT="${OUT:-$REPO/.harness-runs/qwen38-dflash2-subsets-$(date -u +%Y%m%dT%H%M%SZ)}" +CLIENTS="${CLIENTS:-2}" +MASKS="${MASKS:-AA,AS,SA,SS}" +SPEC_DEPTH="${SPEC_DEPTH:-4}" +REPEATS="${REPEATS:-1}" +REFILL_WAVES="${REFILL_WAVES:-1}" +MAX_TOKENS="${MAX_TOKENS:-256}" +WARMUP_TOKENS="${WARMUP_TOKENS:-16}" +MIN_FULL_LIVE_ROUNDS="${MIN_FULL_LIVE_ROUNDS:-2}" +MAX_START_SKEW_MS="${MAX_START_SKEW_MS:-100}" +MAX_REFILL_GAP_MS="${MAX_REFILL_GAP_MS:-100}" +SLOTS="${SLOTS:-8}" +MAX_CTX="${MAX_CTX:-8192}" +MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" +CACHE_TYPE_K="${CACHE_TYPE_K:-q8_0}" +CACHE_TYPE_V="${CACHE_TYPE_V:-q8_0}" +FA_WINDOW="${FA_WINDOW:-0}" +PORT="${PORT:-18139}" +HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-900}" +REQUEST_TIMEOUT_SECONDS="${REQUEST_TIMEOUT_SECONDS:-1800}" +COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-1}" +TARGET_DEVICE="${TARGET_DEVICE:-hip:0}" +DRAFT_DEVICE="${DRAFT_DEVICE:-hip:0}" +VISIBLE_DEVICES="${VISIBLE_DEVICES:-0}" + +usage() { + cat <<'EOF' +Usage: + MODEL=/path/Qwen3.8-27B-target.gguf \ + DRAFT_MODEL=/path/Qwen3.8-27B-DFlash2-q8_0.gguf \ + PROMPT_FILE=/path/prompts.jsonl \ + CLIENTS=2 MASKS=AA,AS,SA,SS SPEC_DEPTH=4 PROMPT_OFFSET=0 \ + REFILL_WAVES=1 \ + harness/benchmarks/concurrency/run_qwen38_dflash2_subsets.sh + +Mask characters are positional: A forces request-local AR and S forces +request-local speculation. Adaptive is intentionally unsupported. One server +is launched for the selected depth, warmed once in all-AR and all-SPEC modes, +then kept alive for every mask and repeat. Invoke a fresh OUT/process for each +depth in a {2,4,8} screen. REFILL_WAVES=1 runs the original synchronized, +closed-cohort diagnostic. REFILL_WAVES>=3 instead keeps each positional lane +full with repeated deterministic requests and reports refill goodput separately +from full-live engine-round goodput. The report fails unless at least +MIN_FULL_LIVE_ROUNDS consecutive engine rounds execute at live=CLIENTS; refill +also reserves one guard cohort to prove saturation away from the final drain. + +The runner records the target, DFlash2 draft, binary, libraries, git revision, +exact command, depth/timing/selector environment, prompt hash, request masks, +per-request output hashes, and measured server records. These are forced +controls only and cannot be reported as adaptive activation results. +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 Qwen3.8 target GGUF" >&2; exit 2; } +[[ -r "$DRAFT_MODEL" ]] || { echo "set DRAFT_MODEL to a readable DFlash2 GGUF" >&2; exit 2; } +[[ -r "$PROMPT_FILE" ]] || { echo "set PROMPT_FILE to a readable JSONL/text prompt file" >&2; exit 2; } +[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +[[ -r "$CLIENT" && -r "$REFILL_CLIENT" && -r "$METADATA_TOOL" && -r "$RUNTIME_METADATA_TOOL" ]] || { + echo "missing DFlash2 subset harness tool" >&2; exit 2; +} +for value_name in CLIENTS REPEATS REFILL_WAVES MAX_TOKENS WARMUP_TOKENS MIN_FULL_LIVE_ROUNDS SLOTS MAX_CTX MAX_CONCURRENT_PREFILLS HEALTH_TIMEOUT_SECONDS REQUEST_TIMEOUT_SECONDS; do + value="${!value_name}" + [[ "$value" =~ ^[1-9][0-9]*$ ]] || { echo "$value_name must be positive" >&2; exit 2; } +done +(( REFILL_WAVES == 1 || REFILL_WAVES >= 3 )) || { echo "REFILL_WAVES must be 1 (closed cohort) or at least 3 (guarded refill)" >&2; exit 2; } +[[ "$SPEC_DEPTH" =~ ^[2-9][0-9]*$ ]] || { echo "SPEC_DEPTH must be at least 2" >&2; exit 2; } +[[ "$PROMPT_OFFSET" =~ ^[0-9]+$ ]] || { echo "PROMPT_OFFSET must be non-negative" >&2; exit 2; } +[[ "$MAX_REFILL_GAP_MS" =~ ^[0-9]+([.][0-9]+)?$ ]] || { echo "MAX_REFILL_GAP_MS must be non-negative" >&2; exit 2; } +[[ "$PORT" =~ ^[1-9][0-9]*$ ]] || { echo "PORT must be positive" >&2; exit 2; } +[[ "$COOLDOWN_SECONDS" =~ ^[0-9]+$ ]] || { echo "COOLDOWN_SECONDS must be non-negative" >&2; exit 2; } +[[ "$MAX_START_SKEW_MS" =~ ^[0-9]+([.][0-9]+)?$ ]] || { echo "MAX_START_SKEW_MS must be non-negative" >&2; exit 2; } +(( CLIENTS <= SLOTS )) || { echo "CLIENTS exceeds SLOTS" >&2; exit 2; } +[[ "$FA_WINDOW" == 0 ]] || { echo "paged concurrency requires FA_WINDOW=0" >&2; exit 2; } +[[ ! -e "$OUT" ]] || { echo "refusing to overwrite OUT=$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 mask_list <<< "$MASKS" +(( ${#mask_list[@]} > 0 )) || { echo "MASKS must not be empty" >&2; exit 2; } +declare -A seen_masks=() +for mask in "${mask_list[@]}"; do + [[ ${#mask} -eq CLIENTS && "$mask" =~ ^[AS]+$ ]] || { + echo "mask $mask must contain exactly CLIENTS=$CLIENTS A/S characters" >&2 + exit 2 + } + [[ -z "${seen_masks[$mask]+present}" ]] || { echo "duplicate mask $mask" >&2; exit 2; } + seen_masks[$mask]=1 +done + +mode_csv() { + local mask="$1" result="" mode index + for ((index=0; index/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 +} + +MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" +DRAFT_MODEL_SHA256="$(sha256sum "$DRAFT_MODEL" | awk '{print $1}')" +capacity=$((SLOTS * MAX_CTX)) +model_id=qwen38-dflash2 +mkdir -p "$OUT" + +command=( + "$LUCE_SERVER_BIN" "$MODEL" --draft "$DRAFT_MODEL" + --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_DEVICE" + --paged-attention --max-concurrency "$SLOTS" + --kv-pool-tokens "$capacity" --max-ctx "$MAX_CTX" + --cache-type-k "$CACHE_TYPE_K" --cache-type-v "$CACHE_TYPE_V" + --fa-window "$FA_WINDOW" --prefix-cache-slots 0 --prefill-cache-slots 0 + --admission-coalesce-ms 20 --draft-residency persistent + --decode-mode speculation --host 127.0.0.1 --port "$PORT" --model-name "$model_id" +) +launch_env=( + "HIP_VISIBLE_DEVICES=$VISIBLE_DEVICES" + "DFLASH_MAX_CONCURRENT_PREFILLS=$MAX_CONCURRENT_PREFILLS" + "DFLASH_SPEC_BATCHED_DRAFT=1" + "DFLASH_SPEC_CHAIN_DEPTH=$SPEC_DEPTH" + "DFLASH_STEP_TIMING=1" + "DFLASH_DFLASH2_SELECTOR_LOG=1" + "PROMPT_OFFSET=$PROMPT_OFFSET" +) +printf 'env ' > "$OUT/server-command.txt" +printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$OUT/server-command.txt" +printf '\n' >> "$OUT/server-command.txt" + +write_case_metadata() { + local case_dir="$1" variant="$2" repeat="$3" workload="${4:-dflash2-forced-subsets}" + mkdir -p "$case_dir" + metadata=( + python3 "$METADATA_TOOL" --out "$case_dir/server-metadata.json" + --variant "$variant" --workload "$workload" + --clients "$CLIENTS" --repeat "$repeat" + --binary "$LUCE_SERVER_BIN" --model "$MODEL" + --model-sha256 "$MODEL_SHA256" --prompt-file "$PROMPT_FILE" + --command-file "$OUT/server-command.txt" --repo "$REPO" + --max-concurrent-prefills "$MAX_CONCURRENT_PREFILLS" + --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_DEVICE" + --draft-model "$DRAFT_MODEL" --draft-model-sha256 "$DRAFT_MODEL_SHA256" + --decode-mode speculation --cache-type-k "$CACHE_TYPE_K" --cache-type-v "$CACHE_TYPE_V" + --fa-window "$FA_WINDOW" --draft-residency persistent + ) + local item + for item in "${launch_env[@]}"; do metadata+=(--launch-env "$item"); done + "${metadata[@]}" + python3 "$RUNTIME_METADATA_TOOL" \ + --metadata "$case_dir/server-metadata.json" --server-log "$OUT/server.log" +} + +run_client_case() { + local case_dir="$1" mask="$2" max_tokens="$3" repeat="$4" min_rounds="$5" waves="${6:-1}" + local modes client="$CLIENT" workload=dflash2-forced-subsets variant + modes="$(mode_csv "$mask")" + variant="dflash2-depth-$SPEC_DEPTH-mask-$mask" + if (( waves > 1 )); then + client="$REFILL_CLIENT" + workload=dflash2-forced-refill + variant+="-refill-w$waves" + fi + write_case_metadata "$case_dir" "$variant" "$repeat" "$workload" + client_cmd=( + python3 "$client" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" + --clients "$CLIENTS" --request-modes "$modes" --spec-depth "$SPEC_DEPTH" + --prompt-file "$PROMPT_FILE" --prompt-offset "$PROMPT_OFFSET" --max-tokens "$max_tokens" + --timeout "$REQUEST_TIMEOUT_SECONDS" --max-start-skew-ms "$MAX_START_SKEW_MS" + --min-full-live-rounds "$min_rounds" + --server-metadata-json "$case_dir/server-metadata.json" + --server-log "$OUT/server.log" --out "$case_dir/bench.json" + --label "DFlash2 depth=$SPEC_DEPTH mask=$mask repeat=$repeat" + ) + if (( waves > 1 )); then + client_cmd+=(--waves "$waves" --max-refill-gap-ms "$MAX_REFILL_GAP_MS") + fi + printf '%q ' "${client_cmd[@]}" > "$case_dir/client-command.txt" + printf '\n' >> "$case_dir/client-command.txt" + "${client_cmd[@]}" | tee "$case_dir/bench.txt" +} + +port_is_available +env "${launch_env[@]}" "${command[@]}" > "$OUT/server.log" 2>&1 & +server_pid=$! +if ! wait_health; then + tail -n 160 "$OUT/server.log" >&2 || true + exit 1 +fi + +all_ar="$(repeat_char A)" +all_spec="$(repeat_char S)" +run_client_case "$OUT/warmup/ar" "$all_ar" "$WARMUP_TOKENS" 0 1 +run_client_case "$OUT/warmup/speculation" "$all_spec" "$WARMUP_TOKENS" 0 1 +sleep "$COOLDOWN_SECONDS" + +for ((repeat=1; repeat<=REPEATS; repeat++)); do + shift_by=$(((repeat - 1) % ${#mask_list[@]})) + for ((index=0; index<${#mask_list[@]}; index++)); do + mask="${mask_list[$(((index + shift_by) % ${#mask_list[@]}))]}" + case_dir="$OUT/c$CLIENTS/depth$SPEC_DEPTH/r$repeat/$mask" + echo "[run] C=$CLIENTS depth=$SPEC_DEPTH repeat=$repeat mask=$mask refill_waves=$REFILL_WAVES" + run_client_case "$case_dir" "$mask" "$MAX_TOKENS" "$repeat" "$MIN_FULL_LIVE_ROUNDS" "$REFILL_WAVES" + sleep "$COOLDOWN_SECONDS" + done +done + +stop_server +echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py b/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py new file mode 100644 index 000000000..efd0f9cff --- /dev/null +++ b/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +"""Tests for the offline DFlash2 selector/subset analyzer.""" + +from __future__ import annotations + +import importlib.util +import json +import math +import tempfile +import unittest +from pathlib import Path + + +HERE = Path(__file__).parent +SCRIPT = HERE / "analyze_dflash2_selector.py" +SPEC = importlib.util.spec_from_file_location("analyze_dflash2_selector", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +analyzer = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(analyzer) + + +def wrapped(record: dict, line: int = 1) -> dict: + return { + "line_index": line, + "raw_json": json.dumps(record, separators=(",", ":")), + "record": record, + } + + +def depth_row(depth: int = 1, accepted: bool = True) -> dict: + return { + "depth": depth, + "accepted": accepted, + "selected_logp": -0.2 * depth, + "lm_margin": 2.0 / depth, + "topk_mass": 0.99, + "rank": 0, + "lm_top1": True, + "selector_margin": 3.0 / depth, + "selector_mass": 0.9, + "selector_entropy": 0.1, + } + + +def selector(engine: int, generated: int, accepted: int = 1) -> dict: + return { + "request_id": engine, + "slot": 1, + "generated": generated, + "accepted_depth": accepted, + "depths": [depth_row(1, accepted >= 1)], + } + + +def artifact() -> dict: + details = [ + { + "request_id": "wire-a", "request_index": 0, "prompt_index": 0, + "prompt_sha256": "a" * 64, "decode_mode": "ar", + "request_decode_tok_s": 20.0, + }, + { + "request_id": "wire-s", "request_index": 1, "prompt_index": 1, + "prompt_sha256": "b" * 64, "decode_mode": "speculation", + "request_decode_tok_s": 30.0, + }, + ] + metrics = [ + { + "request_id": "wire-a", "engine_request_id": 10, + "spec_steps": 0, "spec_accepted_tokens": 0, + }, + { + "request_id": "wire-s", "engine_request_id": 11, + "spec_steps": 2, "spec_accepted_tokens": 2, + }, + ] + selectors = [selector(11, 0), selector(11, 2)] + timings = [ + {"path": "spec", "live": 2, "k": 1, "accepted_tokens": 1, + "emitted_tokens": 3, "total_us": 1000.0}, + {"path": "ar", "live": 1, "k": 0, "accepted_tokens": 0, + "emitted_tokens": 1, "total_us": 1000.0}, + ] + return { + "kind": "dflash2-forced-subset-diagnostic", + "spec_depth": 2, + "validation": {"passed": True}, + "server_metadata": {"repeat": 1}, + "level": { + "clients": 2, + "request_mode_mask": "AS", + "requests_detail": details, + "selected_prompt_set_sha256": "c" * 64, + "aggregate_tok_s": 50.0, + "wall_s": 2.0, + }, + "server_records": { + "requests": [wrapped(row, index + 1) for index, row in enumerate(metrics)], + "selectors": [wrapped(row, index + 3) for index, row in enumerate(selectors)], + "rounds": [wrapped(row, index + 5) for index, row in enumerate(timings)], + "activations": [], + }, + } + + +def all_features(value: float) -> dict[str, float]: + return {key: value for key in analyzer.FEATURE_DIRECTIONS} + + +def subset_case(mask: str, goodput: float) -> dict: + requests = [] + for position, mode in enumerate(mask): + requests.append({ + "position": position, + "mode": "speculation" if mode == "S" else "ar", + "prompt_sha256": str(position), + "first_features": all_features(float(2 - position)) if mode == "S" else None, + "lifetime_yield_fraction": 1.0 if mode == "S" else None, + }) + return { + "clients": 2, + "spec_depth": 8, + "prompt_set_sha256": "set", + "wall_s": 1.0, + "round_timing": { + "all": {"goodput_tok_s": goodput}, + "full_live": {"goodput_tok_s": goodput}, + "tail": {"goodput_tok_s": None}, + }, + "mask": mask, + "aggregate_tok_s": goodput, + "requests": requests, + } + + +class DFlash2SelectorAnalyzerTests(unittest.TestCase): + def test_profile_parser_retains_selector_metric_and_timing_json(self) -> None: + data = ( + b'noise [spec-selector] {"request_id":7,"depths":[]}\n' + b'[concurrency-metrics] {"request_id":"wire","engine_request_id":7}\n' + b'[step-timing] {"path":"spec","live":2}\n' + ) + parsed = analyzer.parse_profile_lines(data) + self.assertEqual(parsed["selectors"][0]["record"]["request_id"], 7) + self.assertEqual( + parsed["requests"][0]["record"]["engine_request_id"], 7, + ) + self.assertEqual(parsed["rounds"][0]["record"]["live"], 2) + self.assertIn('"request_id":7', parsed["selectors"][0]["raw_json"]) + + def test_profile_parser_rejects_malformed_selector_json(self) -> None: + with self.assertRaisesRegex(ValueError, "invalid \\[spec-selector\\]"): + analyzer.parse_profile_lines(b"[spec-selector] {bad}\n") + + def test_artifact_joins_engine_selector_to_wire_request(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bench.json" + path.write_text(json.dumps(artifact()), encoding="utf-8") + case = analyzer.analyze_artifact(path) + request = case["requests"][1] + self.assertEqual(request["wire_request_id"], "wire-s") + self.assertEqual(request["engine_request_id"], 11) + self.assertEqual(request["spec_steps"], 2) + self.assertEqual(request["lifetime_accepted_yield"], 1.0) + self.assertEqual(request["lifetime_yield_fraction"], 1.0) + self.assertAlmostEqual( + request["first_features"]["chain_lm_probability"], math.exp(-0.2), + ) + self.assertEqual(case["round_timing"]["full_live"]["goodput_tok_s"], 3000.0) + self.assertEqual(case["round_timing"]["tail"]["goodput_tok_s"], 1000.0) + self.assertEqual(case["round_timing"]["all"]["goodput_tok_s"], 2000.0) + + def test_artifact_rejects_unknown_selector_engine_id(self) -> None: + value = artifact() + bad = selector(99, 0) + value["server_records"]["selectors"][0] = wrapped(bad) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bench.json" + path.write_text(json.dumps(value), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "unknown engine request 99"): + analyzer.analyze_artifact(path) + + def test_artifact_rejects_acceptance_counter_mismatch(self) -> None: + value = artifact() + metric = value["server_records"]["requests"][1]["record"] + metric["spec_accepted_tokens"] = 1 + value["server_records"]["requests"][1] = wrapped(metric) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bench.json" + path.write_text(json.dumps(value), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "acceptance does not match"): + analyzer.analyze_artifact(path) + + def test_calibration_reports_raw_probability_gap_per_depth(self) -> None: + rows = analyzer._calibration([(1, 2, "prompt", selector(1, 0, accepted=1))]) + self.assertEqual(len(rows), 1) + row = rows[0] + self.assertEqual(row["observed_survival"], 1.0) + self.assertFalse(row["has_both_labels"]) + self.assertEqual(row["unique_prompts"], 1) + self.assertEqual(row["prompts_with_rejection"], 0) + self.assertAlmostEqual( + row["selected_token_probability"]["mean_raw_probability"], + math.exp(-0.2), + ) + self.assertGreater( + row["selected_token_probability"]["observed_minus_raw"], 0.0, + ) + + def test_subset_oracle_exposes_homogeneous_synergy_and_sign_flip(self) -> None: + result = analyzer._subset_group([ + subset_case("AA", 100.0), + subset_case("AS", 70.0), + subset_case("SA", 80.0), + subset_case("SS", 130.0), + ]) + self.assertTrue(result["complete_exhaustive"]) + self.assertEqual(result["oracle_mask"], "SS") + self.assertTrue(result["oracle_is_homogeneous"]) + self.assertTrue(result["homogeneous_dominates_every_mixed"]) + self.assertEqual(result["requests_with_contextual_sign_flip"], 2) + self.assertAlmostEqual( + result["per_request_marginals"][0]["shapley_goodput_tok_s"], 20.0, + ) + self.assertAlmostEqual( + result["per_request_marginals"][1]["shapley_goodput_tok_s"], 10.0, + ) + ranked = result["raw_feature_prefix_rankings"]["chain_lm_logp"] + self.assertEqual(ranked["ranked_positions"], [0, 1]) + self.assertEqual(ranked["by_subset_size"][1]["ranked_prefix_mask"], "SA") + self.assertEqual(ranked["by_subset_size"][1]["same_size_regret"], 0.0) + self.assertFalse(result["first_feature_vs_shapley"]["identifiable"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index 12f896543..173025057 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -538,6 +538,116 @@ def test_failed_evaluation_activation_is_request_local_and_explicit( by_id[8]["activation_evaluation"], "scored", ) + def test_dflash2_activation_fields_and_gate_score_kind_are_preserved( + self, + ) -> None: + scored = { + "request_id": 7, "slot": 0, + "initial_confidence": None, + "activation_score": 5.3306, + "request_benefit": 5.3306, + "score_kind": "dflash2_selector_benefit_v1", + "expected_yield": 5.3306, + "evaluation": "scored", "fallback_reason": None, + "decision_reason": "selected_by_joint_goodput", + "decision": "speculation", + } + failed = { + "request_id": 8, "slot": 1, + "initial_confidence": None, + "activation_score": None, + "request_benefit": None, + "score_kind": "dflash2_selector_benefit_v1", + "expected_yield": None, + "evaluation": "failed", + "fallback_reason": "benefit_evaluation_failed", + "decision_reason": "evaluation_failed", + "decision": "ar", + } + with tempfile.TemporaryDirectory() as tmp: + case = self._write_activation_case( + Path(tmp), "adaptive-on", [scored, failed], + execution={7: (2, 4), 8: (0, 1)}, + ) + log_path = case / "benchmark-server.log" + log_path.write_text( + "[spec-gate] C=2 k=1 " + "scores=[7:5.331/confidence/dflash2_selector_benefit_v1*," + "8:2.685/initial/dflash2_selector_benefit_v1] " + "G(k)=0.010 G(0)=0.009 predicted_cost=1us measured=2us\n" + + log_path.read_text(encoding="utf-8"), + encoding="utf-8", + ) + report = gate_analysis.analyze_case(case) + self.assertEqual( + report["activation"]["score_kind_counts"], + {"dflash2_selector_benefit_v1": 2}, + ) + self.assertEqual( + report["activation"]["fallback_reason_counts"], + {"benefit_evaluation_failed": 1}, + ) + by_id = { + row["engine_request_id"]: row for row in report["requests"] + } + request = by_id[7] + self.assertIsNone(request["initial_confidence"]) + self.assertAlmostEqual(request["activation_score"], 5.3306) + self.assertAlmostEqual(request["request_benefit"], 5.3306) + self.assertEqual( + request["activation_score_kind"], + "dflash2_selector_benefit_v1", + ) + self.assertEqual( + request["activation_decision_reason"], + "selected_by_joint_goodput", + ) + self.assertAlmostEqual(request["mean_activation_score"], 5.331) + self.assertIsNone(request["mean_confidence_yield"]) + rounds, _, _, _ = gate_analysis.parse_server_log(log_path) + self.assertEqual( + rounds[0]["entries"][0]["score_kind"], + "dflash2_selector_benefit_v1", + ) + + + def test_typed_fallback_reasons_are_nonempty_and_recognized(self) -> None: + base = { + "request_id": 7, "slot": 0, + "initial_confidence": None, + "activation_score": None, + "request_benefit": None, + "score_kind": "unspecified", + "expected_yield": None, + "evaluation": "failed", + "decision_reason": "evaluation_failed", + "decision": "ar", + } + accepted = ( + "benefit_adapter_unavailable", + "benefit_adapter_invalid_config", + "cost_profile_unavailable", + ) + for reason in accepted: + with self.subTest(reason=reason), tempfile.TemporaryDirectory() as tmp: + row = {**base, "fallback_reason": reason} + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(row)}\n", encoding="utf-8", + ) + _, _, _, activations = gate_analysis.parse_server_log(path) + self.assertEqual(activations, [row]) + for reason in ("", "unknown_failure"): + with self.subTest(reason=reason), tempfile.TemporaryDirectory() as tmp: + row = {**base, "fallback_reason": reason} + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(row)}\n", encoding="utf-8", + ) + with self.assertRaisesRegex(ValueError, "fallback_reason"): + gate_analysis.parse_server_log(path) + + def test_activation_record_fields_are_strictly_validated(self) -> None: valid = { "request_id": 7, "slot": 0, diff --git a/harness/benchmarks/concurrency/test_forced_subset_benchmark.py b/harness/benchmarks/concurrency/test_forced_subset_benchmark.py new file mode 100644 index 000000000..c5deb4351 --- /dev/null +++ b/harness/benchmarks/concurrency/test_forced_subset_benchmark.py @@ -0,0 +1,268 @@ +#!/usr/bin/env python3 +"""Tests for the forced DFlash2 subset/depth diagnostic client.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +SCRIPT = HERE / "forced_subset_benchmark.py" +SPEC = importlib.util.spec_from_file_location("forced_subset_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) + + +def wrapped(record: dict, line: int = 1) -> dict: + raw = json.dumps(record, separators=(",", ":")) + return {"line_index": line, "raw_json": raw, "record": record} + + +def request(request_id: str, mode: str) -> dict: + return { + "request_id": request_id, + "decode_mode": mode, + "error": None, + "content_sha256": "a" * 64, + "reasoning_content_sha256": "b" * 64, + "combined_output_sha256": "c" * 64, + } + + +def valid_level() -> dict: + return { + "requests": 2, + "requests_ok": 2, + "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": True, + "fixed_token_workload_valid": True, + "start_skew_s": 0.001, + "requests_detail": [request("wire-ar", "ar"), request("wire-spec", "speculation")], + } + + +def valid_records() -> dict: + return { + "rounds": [ + wrapped({"path": "ar", "live": 2, "k": 0}, 1), + wrapped({ + "path": "spec", "live": 2, "k": 1, + "tree_bucket": 2, "tree_rows": 8, + }, 2), + ], + "selectors": [wrapped({"request_id": 7, "depths": []}, 3)], + "activations": [], + "requests": [ + wrapped({"request_id": "wire-ar", "spec_steps": 0}, 4), + wrapped({"request_id": "wire-spec", "spec_steps": 3}, 5), + ], + } + + +class ForcedSubsetTests(unittest.TestCase): + def test_request_payload_forces_mode_zero_temperature_and_ignore_eos(self) -> None: + captured: dict = {} + + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"id":"wire-1","choices":[{"delta":{"content":"x"}}]}\n', + b"\n", + b'data: {"id":"wire-1","choices":[{"delta":{},' + b'"finish_reason":"length"}]}\n', + b"\n", + b'data: {"id":"wire-1","choices":[],"usage":' + b'{"prompt_tokens":3,"completion_tokens":8}}\n', + b"\n", + b"data: [DONE]\n", b"\n", + ]) + + def open_request(http_request, timeout): + captured.update(json.loads(http_request.data)) + self.assertEqual(timeout, 2.0) + return Response() + + args = argparse.Namespace( + model="m", max_tokens=8, seed=1, api_key="", + base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object( + benchmark.urllib.request, "urlopen", side_effect=open_request, + ): + row = benchmark.stream_request(args, "prompt", "speculation") + self.assertEqual(captured["decode_mode"], "speculation") + self.assertEqual(captured["temperature"], 0.0) + self.assertIs(captured["ignore_eos"], True) + self.assertEqual(row["request_id"], "wire-1") + self.assertEqual(row["decode_mode"], "speculation") + self.assertEqual(len(row["combined_output_sha256"]), 64) + self.assertIsNone(row["error"]) + + def test_modes_are_positional_and_reject_adaptive(self) -> None: + self.assertEqual( + benchmark.parse_request_modes("ar,speculation", 2), + ["ar", "speculation"], + ) + with self.assertRaisesRegex(ValueError, "adaptive is intentionally out of scope"): + benchmark.parse_request_modes("adaptive", 1) + with self.assertRaisesRegex(ValueError, "exactly one"): + benchmark.parse_request_modes("ar", 2) + + def test_depth_and_profiling_environment_are_required_in_metadata(self) -> None: + metadata = { + "clients": 4, + "launch_environment": { + "DFLASH_SPEC_CHAIN_DEPTH": "8", + "DFLASH_STEP_TIMING": "1", + "DFLASH_DFLASH2_SELECTOR_LOG": "1", + "PROMPT_OFFSET": "5", + }, + } + benchmark.validate_server_metadata(metadata, 4, 8, 5, True) + metadata["launch_environment"]["DFLASH_SPEC_CHAIN_DEPTH"] = "4" + with self.assertRaisesRegex(ValueError, "DFLASH_SPEC_CHAIN_DEPTH=8"): + benchmark.validate_server_metadata(metadata, 4, 8, 5, True) + metadata["launch_environment"]["DFLASH_SPEC_CHAIN_DEPTH"] = "8" + with self.assertRaisesRegex(ValueError, "PROMPT_OFFSET=4"): + benchmark.validate_server_metadata(metadata, 4, 8, 4, True) + + def test_ar_only_metadata_does_not_require_selector_logging(self) -> None: + metadata = { + "clients": 1, + "launch_environment": { + "DFLASH_SPEC_CHAIN_DEPTH": "4", + "DFLASH_STEP_TIMING": "1", + "PROMPT_OFFSET": "0", + }, + } + benchmark.validate_server_metadata(metadata, 1, 4, 0, False) + + def test_profile_parser_retains_raw_round_and_request_records(self) -> None: + data = ( + b'prefix [step-timing] {"path":"spec","live":2,"k":1}\n' + b'[spec-selector] {"request_id":7,"depths":[]}\n' + b'[concurrency-metrics] {"request_id":"wire","spec_steps":1}\n' + ) + records = benchmark.parse_profile_records(data) + self.assertEqual(records["rounds"][0]["record"]["live"], 2) + self.assertEqual(records["selectors"][0]["record"]["request_id"], 7) + self.assertEqual(records["requests"][0]["record"]["request_id"], "wire") + self.assertIn('"path":"spec"', records["rounds"][0]["raw_json"]) + + def test_profile_parser_rejects_malformed_measured_json(self) -> None: + with self.assertRaisesRegex(ValueError, r"invalid \[step-timing\]"): + benchmark.parse_profile_records(b"[step-timing] {bad}\n") + + def test_valid_mixed_subset_proves_depth_and_sustained_full_live(self) -> None: + result = benchmark.validate_evidence( + valid_level(), valid_records(), 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertEqual(result["executed_spec_depths"], [4]) + self.assertEqual(result["longest_full_live_streak"], 2) + self.assertIs(result["adaptive_claims_permitted"], False) + + def test_missing_requested_live_concurrency_fails_closed(self) -> None: + records = valid_records() + for row in records["rounds"]: + row["record"]["live"] = 1 + result = benchmark.validate_evidence( + valid_level(), records, 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=1, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("sustained live=C proof absent" in error for error in result["errors"])) + + def test_executed_depth_mismatch_fails_closed(self) -> None: + records = valid_records() + records["rounds"][1]["record"]["tree_rows"] = 16 + result = benchmark.validate_evidence( + valid_level(), records, 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("do not match requested depth 4" in error for error in result["errors"])) + + def test_request_mode_execution_mismatch_fails_closed(self) -> None: + records = valid_records() + records["requests"][0]["record"]["spec_steps"] = 2 + result = benchmark.validate_evidence( + valid_level(), records, 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("forced AR executed speculation" in error for error in result["errors"])) + + def test_log_span_is_exact_and_rejects_truncation(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_bytes(b"warmup\nmeasured\n") + data, end = benchmark.read_log_span(path, len(b"warmup\n")) + self.assertEqual(data, b"measured\n") + self.assertEqual(end, path.stat().st_size) + with self.assertRaisesRegex(ValueError, "truncated or rotated"): + benchmark.read_log_span(path, end + 1) + + def test_ar_only_evidence_passes_without_selector_records(self) -> None: + level = { + "requests": 1, "requests_ok": 1, "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": True, + "fixed_token_workload_valid": True, + "start_skew_s": 0.0, + "requests_detail": [request("wire-ar", "ar")], + } + records = { + "rounds": [ + wrapped({"path": "ar", "live": 1, "k": 0}, 1), + wrapped({"path": "ar", "live": 1, "k": 0}, 2), + ], + "selectors": [], "activations": [], + "requests": [wrapped({"request_id": "wire-ar", "spec_steps": 0}, 3)], + } + result = benchmark.validate_evidence( + level, records, 1, ["ar"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + + def test_runner_pins_one_server_and_dflash2_metadata_across_masks(self) -> None: + runner = (HERE / "run_qwen38_dflash2_subsets.sh").read_text( + encoding="utf-8", + ) + self.assertEqual( + runner.count( + 'env "${launch_env[@]}" "${command[@]}" > "$OUT/server.log"' + ), + 1, + ) + self.assertIn('"DFLASH_SPEC_CHAIN_DEPTH=$SPEC_DEPTH"', runner) + self.assertIn('"DFLASH_DFLASH2_SELECTOR_LOG=1"', runner) + self.assertIn('PROMPT_OFFSET="${PROMPT_OFFSET:-0}"', runner) + self.assertIn('"PROMPT_OFFSET=$PROMPT_OFFSET"', runner) + self.assertIn('--prompt-offset "$PROMPT_OFFSET"', runner) + self.assertNotIn('--prompt-offset 0', runner) + self.assertIn('--draft-model "$DRAFT_MODEL"', runner) + self.assertIn('--decode-mode speculation --host', runner) + self.assertIn('--decode-mode speculation --cache-type-k', runner) + self.assertIn('VISIBLE_DEVICES="${VISIBLE_DEVICES:-0}"', runner) + self.assertNotIn('--decode-mode ar', runner) + self.assertIn('run_client_case "$case_dir" "$mask"', runner) + self.assertNotIn("dspark", runner.lower()) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_refill_subset_benchmark.py b/harness/benchmarks/concurrency/test_refill_subset_benchmark.py new file mode 100644 index 000000000..5210b4428 --- /dev/null +++ b/harness/benchmarks/concurrency/test_refill_subset_benchmark.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python3 +"""Tests for the fail-closed DFlash2 refill diagnostic.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +import threading +import time +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +SCRIPT = HERE / "refill_subset_benchmark.py" +SPEC = importlib.util.spec_from_file_location("refill_subset_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) + + +def wrapped(record: dict, line: int) -> dict: + raw = json.dumps(record, separators=(",", ":")) + return {"line_index": line, "raw_json": raw, "record": record} + + +def valid_workload() -> dict: + requests = [] + for wave in range(3): + for lane in range(2): + digest = ("a" if lane == 0 else "b") * 64 + requests.append({ + "request_id": f"wire-{wave}-{lane}", + "lane_index": lane, + "wave_index": wave, + "decode_mode": "ar" if lane == 0 else "speculation", + "error": None, + "content_sha256": digest, + "reasoning_content_sha256": "c" * 64, + "combined_output_sha256": "d" * 64, + }) + return { + "clients": 2, + "waves": 3, + "requests": 6, + "requests_ok": 6, + "failures": 0, + "missing_requests": 0, + "token_count_complete": True, + "prompt_token_count_complete": True, + "fixed_token_workload_valid": True, + "initial_start_skew_s": 0.001, + "refill_handoffs": 4, + "refill_gap_s_max": 0.001, + "exact_output_stable_per_lane": True, + "lanes": [ + {"lane_index": 0, "decode_mode": "ar", "requests": 3}, + {"lane_index": 1, "decode_mode": "speculation", "requests": 3}, + ], + "requests_detail": requests, + } + + +def valid_records() -> dict: + rounds = [ + wrapped({ + "path": "spec", "live": 2, "k": 1, + "tree_bucket": 1, "tree_rows": 4, + "total_us": 20_000.0, "emitted_tokens": 2, + }, 1), + wrapped({ + "path": "spec", "live": 2, "k": 1, + "tree_bucket": 1, "tree_rows": 4, + "total_us": 20_000.0, "emitted_tokens": 2, + }, 4), + wrapped({ + "path": "spec", "live": 2, "k": 1, + "tree_bucket": 1, "tree_rows": 4, + "total_us": 20_000.0, "emitted_tokens": 2, + }, 7), + ] + metrics = [] + selectors = [] + line_by_wave = (2, 5, 8) + for wave, line in enumerate(line_by_wave): + for lane in range(2): + engine_id = wave * 2 + lane + 1 + metrics.append(wrapped({ + "request_id": f"wire-{wave}-{lane}", + "engine_request_id": engine_id, + "spec_steps": 0 if lane == 0 else 3, + }, line + lane)) + if lane == 1: + selectors.append(wrapped({ + "request_id": engine_id, + "accepted_depth": 3, + "depths": [], + }, line + lane)) + return { + "rounds": rounds, + "selectors": selectors, + "activations": [], + "requests": metrics, + } + + +class RefillSubsetTests(unittest.TestCase): + def test_refill_client_keeps_positional_modes_and_sequences_each_lane(self) -> None: + calls: dict[str, int] = {"prompt-a": 0, "prompt-b": 0} + lock = threading.Lock() + + def fake_request(args, prompt, mode): + with lock: + wave = calls[prompt] + calls[prompt] += 1 + started = time.perf_counter() + digest = benchmark.base.sha256_text(prompt) + return { + "request_id": f"{prompt}-{wave}", + "decode_mode": mode, + "t_start": started, + "t_first": started, + "t_end": time.perf_counter(), + "completion_tokens": args.max_tokens, + "prompt_tokens": 4, + "error": None, + "content_sha256": digest, + "reasoning_content_sha256": "e" * 64, + "combined_output_sha256": "f" * 64, + } + + args = argparse.Namespace( + clients=2, waves=3, prompt_offset=0, timeout=2.0, max_tokens=8, + ) + with mock.patch.object( + benchmark.forced, "stream_request", side_effect=fake_request, + ): + workload = benchmark.run_refill( + args, ["prompt-a", "prompt-b"], ["ar", "speculation"], + ) + self.assertEqual(sum(calls.values()), 6) + self.assertTrue(all(count >= 1 for count in calls.values())) + self.assertEqual(workload["requests"], 6) + self.assertEqual(workload["refill_handoffs"], 4) + self.assertEqual(workload["request_mode_mask"], "AS") + self.assertTrue(workload["exact_output_stable_per_lane"]) + self.assertEqual( + [row["request_index"] for row in workload["requests_detail"]], + list(range(6)), + ) + self.assertEqual( + [row["admission_group_index"] for row in workload["requests_detail"]], + [0, 0, 1, 1, 2, 2], + ) + + def test_valid_refill_proves_every_handoff_and_full_live_round_rate(self) -> None: + result = benchmark.validate_evidence( + valid_workload(), valid_records(), 2, ["ar", "speculation"], + 4, 3, max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertTrue(result["refill_recovery_proved"]) + self.assertEqual(result["completed_before_last_full_live"], 4) + self.assertEqual(result["required_refill_recoveries"], 2) + self.assertEqual(result["scheduled_refills"], 4) + self.assertEqual(result["terminal_guard_requests"], 2) + self.assertEqual(result["executed_spec_depths"], [4]) + self.assertAlmostEqual( + result["full_live"]["engine_round_goodput_tok_s"], 100.0, + ) + self.assertIs(result["adaptive_claims_permitted"], False) + self.assertIs(result["closed_cohort_claims_permitted"], False) + + def test_missing_post_refill_full_live_proof_fails_closed(self) -> None: + records = valid_records() + del records["rounds"][1:] + result = benchmark.validate_evidence( + valid_workload(), records, 2, ["ar", "speculation"], 4, 3, + 100.0, 100.0, 2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any( + "full live=C was not recovered" in error for error in result["errors"] + )) + + def test_every_request_mode_and_selector_mapping_are_checked(self) -> None: + records = valid_records() + records["requests"][0]["record"]["spec_steps"] = 1 + records["selectors"].append(wrapped({"request_id": 1}, 3)) + result = benchmark.validate_evidence( + valid_workload(), records, 2, ["ar", "speculation"], 4, 3, + 100.0, 100.0, 2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any( + "forced AR executed speculation" in error for error in result["errors"] + )) + self.assertTrue(any( + "forced AR emitted DFlash2 selector" in error for error in result["errors"] + )) + + def test_exact_output_instability_and_slow_handoff_fail_closed(self) -> None: + workload = valid_workload() + workload["exact_output_stable_per_lane"] = False + workload["refill_gap_s_max"] = 0.2 + result = benchmark.validate_evidence( + workload, valid_records(), 2, ["ar", "speculation"], 4, 3, + 100.0, 100.0, 2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("exact output hashes changed" in e for e in result["errors"])) + self.assertTrue(any("handoff gap exceeds" in e for e in result["errors"])) + + def test_timing_rows_must_have_exact_positive_fields(self) -> None: + records = valid_records() + records["rounds"][0]["record"]["emitted_tokens"] = 0 + result = benchmark.validate_evidence( + valid_workload(), records, 2, ["ar", "speculation"], 4, 3, + 100.0, 100.0, 2, + ) + self.assertFalse(result["passed"]) + self.assertTrue(any("positive total_us/emitted_tokens" in e for e in result["errors"])) + + def test_runner_refill_hook_is_opt_in_and_warmups_remain_closed(self) -> None: + runner = (HERE / "run_qwen38_dflash2_subsets.sh").read_text( + encoding="utf-8", + ) + self.assertIn('REFILL_WAVES="${REFILL_WAVES:-1}"', runner) + self.assertIn('REFILL_CLIENT="${REFILL_CLIENT:-$SCRIPT_DIR/refill_subset_benchmark.py}"', runner) + self.assertIn('client_cmd+=(--waves "$waves" --max-refill-gap-ms "$MAX_REFILL_GAP_MS")', runner) + self.assertIn('workload=dflash2-forced-refill', runner) + self.assertIn( + 'run_client_case "$OUT/warmup/ar" "$all_ar" "$WARMUP_TOKENS" 0 1', + runner, + ) + self.assertIn( + 'run_client_case "$case_dir" "$mask" "$MAX_TOKENS" "$repeat" ' + '"$MIN_FULL_LIVE_ROUNDS" "$REFILL_WAVES"', + runner, + ) + self.assertIn('> "$case_dir/client-command.txt"', runner) + + def test_refill_requires_three_waves_for_a_guard_cohort(self) -> None: + args = argparse.Namespace(clients=1, waves=2) + with self.assertRaisesRegex(ValueError, "at least 3"): + benchmark.run(args) + + +if __name__ == "__main__": + unittest.main() diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 57f2ff462..74a7ae187 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -439,7 +439,9 @@ add_library(dflash_common STATIC src/common/dynamic_backend.cpp src/common/domino_head.cpp src/common/dspark_head.cpp + src/common/dflash2_benefit.cpp src/common/dflash2_head.cpp + src/common/dflash2_batch.cpp src/common/target_shard_ipc.cpp src/common/target_shard_ipc_daemon.cpp src/common/dflash_feature_ring.cpp @@ -1453,6 +1455,37 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/test) list(APPEND _raw_unit_test_targets test_chain_spec_shapes) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dflash2_selector_validation.cpp") + # Pure host-side selector metadata/layout validation: no GPU. + add_executable(test_dflash2_selector_validation + test/test_dflash2_selector_validation.cpp) + target_include_directories(test_dflash2_selector_validation PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_dflash2_selector_validation) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dflash2_benefit.cpp") + # Pure host-side, versioned selector-to-benefit adapter: no GPU. + add_executable(test_dflash2_benefit + test/test_dflash2_benefit.cpp + src/common/dflash2_benefit.cpp) + target_include_directories(test_dflash2_benefit PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/test + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + list(APPEND _raw_unit_test_targets test_dflash2_benefit) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_delta_transition_journal.cpp") + # Pure host-side proof for replay-free DeltaNet transition commits. + add_executable(test_delta_transition_journal + test/test_delta_transition_journal.cpp + src/qwen35/delta_transition_journal.cpp) + target_include_directories(test_delta_transition_journal PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_delta_transition_journal) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dspark_batched_head.cpp") add_executable(test_dspark_batched_head test/test_dspark_batched_head.cpp) @@ -1969,6 +2002,17 @@ if(DFLASH27B_TESTS) add_dependencies(check test_batched_gdn) endif() endif() + if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR + DFLASH27B_GPU_BACKEND STREQUAL "hip") + AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_gdn_transition_journal.cpp") + dflash_add_ggml_gpu_executable( + test_gdn_transition_journal + test/test_gdn_transition_journal.cpp) + add_test(NAME gdn_transition_journal COMMAND test_gdn_transition_journal) + if(TARGET check) + add_dependencies(check test_gdn_transition_journal) + endif() + endif() if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR DFLASH27B_GPU_BACKEND STREQUAL "hip") AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_concat_transpose.cpp") diff --git a/server/deps/llama.cpp/ggml/include/ggml-cuda.h b/server/deps/llama.cpp/ggml/include/ggml-cuda.h index bb850e2a9..c8d62e21a 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -93,6 +93,19 @@ GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void); GGML_BACKEND_API bool ggml_backend_cuda_topk_rows(const struct ggml_tensor * logits, int k, float * probs_out, int32_t * ids_out); +// Apply compact GDN journal prefixes to persistent F32 state. The journal is +// [J,H,T,B] (see ggml_gated_delta_net_set_transition_journal); accepted and +// active slots are contiguous I32 [B]. Negative/out-of-range slots are +// padding. Phase 1 is synchronous, single-device, and requires unique slots. +// Each live state slot must still contain the same base state from which its +// journal row was captured, because delta is state-dependent. Call only after +// the synchronous graph compute that produced the journal has returned. +GGML_BACKEND_API bool ggml_backend_cuda_gdn_transition_journal_commit( + const struct ggml_tensor * journal, + struct ggml_tensor * state, + const struct ggml_tensor * accepted_prefixes, + const struct ggml_tensor * active_slot_ids); + // Attach learned per-expert decode tables to a mixed-precision tensor. The // host variants copy the tables to the device that owns `base`. Call the // matching unregister function before releasing the tensor's backing buffer. diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 567e29094..fd01a3a17 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2923,6 +2923,14 @@ extern "C" { struct ggml_tensor * tensor, bool skip_intermediate); + // CUDA/HIP linear-chain journal in compact F32 [J,H,T,B] layout: + // scalar gate J=2*S_v+1 stores [g | k | delta], while KDA J=3*S_v + // stores [g[S_v] | k | delta]. Delta is captured after the + // state-dependent reduction. Tree mode is deliberately rejected. + GGML_API void ggml_gated_delta_net_set_transition_journal( + struct ggml_tensor * tensor, + struct ggml_tensor * journal); + // dflash extension: let the kernel derive the gates from the raw // projections instead of graph-side sigmoid/softplus ops: // beta_val = sigmoid(beta_raw) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu index e9f416dd8..1489877cd 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu @@ -82,6 +82,7 @@ gated_delta_net_cuda(const float * q, float * state_out, const int * parent_ids, // TREE_MODE only; else ignored InterT * persist_inter, // optional external buffer for per-token intermediates + float * transition_journal, int64_t H, int64_t n_tokens, int64_t n_seqs, @@ -198,6 +199,13 @@ gated_delta_net_cuda(const float * q, const float * beta_t = beta + gb_offset; const float * g_t = g + gb_offset * (KDA ? S_v : 1); + constexpr int journal_gate_values = KDA ? S_v : 1; + constexpr int journal_width = journal_gate_values + 2*S_v; + float * journal_t = transition_journal + ? transition_journal + + ((sequence * n_tokens + t) * H + h_idx) * journal_width + : nullptr; + // raw-gate mode: beta = sigmoid(beta_raw); g = softplus(alpha_raw + bias) * A const bool raw_gates = gate_bias != nullptr; const float beta_val = raw_gates ? 1.0f / (1.0f + expf(-(*beta_t))) : *beta_t; @@ -212,6 +220,17 @@ gated_delta_net_cuda(const float * q, q_reg[r] = q_t[i]; } + if (journal_t && blockIdx.z == 0 && threadIdx.y == 0) { +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int i = r * warp_size + lane; + journal_t[journal_gate_values + i] = k_reg[r]; + if constexpr (KDA) { + journal_t[i] = expf(g_t[i]); + } + } + } + if constexpr (!KDA) { float g_log = *g_t; if (raw_gates) { @@ -219,6 +238,9 @@ gated_delta_net_cuda(const float * q, g_log = ((a > 20.0f) ? a : logf(1.0f + expf(a))) * gate_A[h_idx]; } const float g_val = expf(g_log); + if (journal_t && lane == 0 && col == 0) { + journal_t[0] = g_val; + } // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] float kv_shard = 0.0f; @@ -230,6 +252,9 @@ gated_delta_net_cuda(const float * q, // delta[col] = (v[col] - g * kv[col]) * beta float delta_col = (v_t[col] - g_val * kv_col) * beta_val; + if (journal_t && lane == 0) { + journal_t[journal_gate_values + S_v + col] = delta_col; + } // fused: S[i][col] = g * S[i][col] + k[i] * delta[col] // attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i] @@ -258,6 +283,9 @@ gated_delta_net_cuda(const float * q, // delta[col] = (v[col] - kv[col]) * beta float delta_col = (v_t[col] - kv_col) * beta_val; + if (journal_t && lane == 0) { + journal_t[journal_gate_values + S_v + col] = delta_col; + } // fused: S[i][col] = g[i] * S[i][col] + k[i] * delta[col] // attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i] @@ -517,7 +545,8 @@ static void launch_gated_delta_net( int64_t sb1, int64_t sb2, int64_t sb3, int64_t neqk1, int64_t rq3, float scale, cudaStream_t stream, - const float * gate_bias = nullptr, const float * gate_A = nullptr) { + const float * gate_bias = nullptr, const float * gate_A = nullptr, + float * transition_journal_d = nullptr) { //TODO: Add chunked kernel for even faster pre-fill const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; const int num_warps = 4; @@ -539,26 +568,26 @@ static void launch_gated_delta_net( switch (S_v) { case 16: gated_delta_net_cuda<16, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; case 32: gated_delta_net_cuda<32, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; case 64: { gated_delta_net_cuda<64, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); break; } case 128: { if constexpr (!KDA && !TREE_MODE) { - if (use_grouped_cols && + if (transition_journal_d == nullptr && use_grouped_cols && ((GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) || GGML_CUDA_CC_IS_AMD(cc))) { constexpr int cols = 4; @@ -583,19 +612,19 @@ static void launch_gated_delta_net( sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } } else { gated_delta_net_cuda<128, KDA, TREE_MODE, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, H, + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } @@ -625,6 +654,8 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * // Optional 9th source maps compact sequence rows to physical recurrent // state slabs. Negative ids are graph-bucket padding rows. ggml_tensor * src_active_slots = dst->src[8]; + // Optional compact transition journal [J,H,T,B]. Linear-chain only. + ggml_tensor * src_transition_journal = dst->src[11]; GGML_TENSOR_LOCALS(int64_t, neq, src_q, ne); GGML_TENSOR_LOCALS(size_t , nbq, src_q, nb); @@ -667,6 +698,9 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * void * persist_inter_d = src_persist_inter ? src_persist_inter->data : nullptr; + float * transition_journal_d = src_transition_journal + ? (float *) src_transition_journal->data + : nullptr; const bool persist_is_f16 = src_persist_inter && src_persist_inter->type == GGML_TYPE_F16; if (src_persist_inter) { @@ -695,6 +729,16 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * GGML_ASSERT(ggml_is_contiguous(src_active_slots)); GGML_ASSERT(ggml_nelements(src_active_slots) == n_seqs); } + if (src_transition_journal) { + const int64_t journal_width = kda ? 3*S_v : 2*S_v + 1; + GGML_ASSERT(!src_parent); + GGML_ASSERT(src_transition_journal->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src_transition_journal)); + GGML_ASSERT(src_transition_journal->ne[0] == journal_width); + GGML_ASSERT(src_transition_journal->ne[1] == H); + GGML_ASSERT(src_transition_journal->ne[2] == n_tokens); + GGML_ASSERT(src_transition_journal->ne[3] == n_seqs); + } // strides in floats (beta strides used for both g and beta offset computation) const int64_t sq1 = nbq1 / sizeof(float); @@ -737,34 +781,34 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } else if (write_intermediate) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } else { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } \ } else { \ if (tree_mode) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, parent_ids_d, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } else if (write_intermediate) { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } else { \ launch_gated_delta_net( \ q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, nullptr, persist_typed, \ S_v, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, \ - sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d); \ + sb1, sb2, sb3, neqk1, rq3, scale, stream, gate_bias_d, gate_A_d, transition_journal_d); \ } \ } \ } while (0) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu new file mode 100644 index 000000000..114b760fb --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu @@ -0,0 +1,184 @@ +#include "common.cuh" +#include "ggml-cuda.h" + +#include +#include +#include + +#if defined(GGML_USE_HIP) +#ifndef cudaPointerAttributes +#define cudaPointerAttributes hipPointerAttribute_t +#define cudaPointerGetAttributes hipPointerGetAttributes +#define cudaMemoryTypeDevice hipMemoryTypeDevice +#define cudaMemoryTypeManaged hipMemoryTypeManaged +#endif +#endif + +namespace { + +__global__ void gdn_transition_journal_commit_kernel( + const float * journal, + float * state, + const int32_t * accepted_prefixes, + const int32_t * active_slot_ids, + int state_size, + int n_heads, + int n_tokens, + int n_seqs, + int n_state_slots, + int journal_width, + int gate_values) { + const int sequence = blockIdx.z; + const int head = blockIdx.y; + const int element = blockIdx.x * blockDim.x + threadIdx.x; + const int state_elements = state_size * state_size; + if (sequence >= n_seqs || head >= n_heads || + element >= state_elements) { + return; + } + + const int slot = active_slot_ids[sequence]; + const int accepted = accepted_prefixes[sequence]; + if (slot < 0 || slot >= n_state_slots || + accepted < 0 || accepted > n_tokens) { + return; + } + + const int row = element % state_size; + const int col = element / state_size; + const size_t state_offset = + (((size_t) slot*n_heads + head)*state_size + col)*state_size + row; + float current = state[state_offset]; + + for (int token = 0; token < accepted; ++token) { + const float * transition = journal + + (((size_t) sequence*n_tokens + token)*n_heads + head) * + journal_width; + const float gate = gate_values == 1 + ? transition[0] + : transition[row]; + const float key = transition[gate_values + row]; + const float delta = + transition[gate_values + state_size + col]; + current = fmaf(key, delta, gate * current); + } + + state[state_offset] = current; +} + +bool device_pointer(const void * pointer, int & device) { + if (pointer == nullptr) return false; + cudaPointerAttributes attributes{}; + if (cudaPointerGetAttributes(&attributes, pointer) != cudaSuccess) { + (void) cudaGetLastError(); + return false; + } + if (attributes.type != cudaMemoryTypeDevice && + attributes.type != cudaMemoryTypeManaged) { + return false; + } + device = attributes.device; + return true; +} + +} // namespace + +extern "C" bool ggml_backend_cuda_gdn_transition_journal_commit( + const ggml_tensor * journal, + ggml_tensor * state, + const ggml_tensor * accepted_prefixes, + const ggml_tensor * active_slot_ids) { + if (!journal || !state || !accepted_prefixes || !active_slot_ids || + journal->type != GGML_TYPE_F32 || + state->type != GGML_TYPE_F32 || + accepted_prefixes->type != GGML_TYPE_I32 || + active_slot_ids->type != GGML_TYPE_I32 || + !ggml_is_contiguous(journal) || + !ggml_is_contiguous(state) || + !ggml_is_contiguous(accepted_prefixes) || + !ggml_is_contiguous(active_slot_ids)) { + return false; + } + + const int64_t state_size = state->ne[0]; + const int64_t n_heads = state->ne[2]; + const int64_t n_state_slots = state->ne[3]; + const int64_t journal_width = journal->ne[0]; + const int64_t n_tokens = journal->ne[2]; + const int64_t n_seqs = journal->ne[3]; + const bool supported_state_size = + state_size == 16 || state_size == 32 || + state_size == 64 || state_size == 128; + if (!supported_state_size || state->ne[1] != state_size || + n_heads < 1 || journal->ne[1] != n_heads || + n_tokens < 1 || n_seqs < 1 || n_state_slots < 1 || + ggml_nelements(accepted_prefixes) != n_seqs || + ggml_nelements(active_slot_ids) != n_seqs || + (journal_width != 2*state_size + 1 && + journal_width != 3*state_size) || + state_size > std::numeric_limits::max() || + n_heads > 65535 || + n_tokens > std::numeric_limits::max() || + n_seqs > 65535 || + n_state_slots > std::numeric_limits::max() || + journal_width > std::numeric_limits::max()) { + return false; + } + + int device = -1; + int pointer_device = -1; + const void * pointers[] = { + journal->data, state->data, + accepted_prefixes->data, active_slot_ids->data, + }; + for (const void * pointer : pointers) { + if (!device_pointer(pointer, pointer_device)) return false; + if (device < 0) device = pointer_device; + if (pointer_device != device) return false; + } + ggml_cuda_set_device(device); + + std::vector accepted((size_t) n_seqs); + std::vector slots((size_t) n_seqs); + const size_t map_bytes = (size_t) n_seqs * sizeof(int32_t); + if (cudaMemcpy(accepted.data(), accepted_prefixes->data, map_bytes, + cudaMemcpyDeviceToHost) != cudaSuccess || + cudaMemcpy(slots.data(), active_slot_ids->data, map_bytes, + cudaMemcpyDeviceToHost) != cudaSuccess) { + return false; + } + + std::vector seen((size_t) n_state_slots, 0); + for (int64_t sequence = 0; sequence < n_seqs; ++sequence) { + if (accepted[(size_t) sequence] < 0 || + accepted[(size_t) sequence] > n_tokens) { + return false; + } + const int32_t slot = slots[(size_t) sequence]; + if (slot < 0 || slot >= n_state_slots) continue; + if (seen[(size_t) slot]) return false; + seen[(size_t) slot] = 1; + } + + constexpr int threads = 256; + const int64_t state_elements = state_size * state_size; + const dim3 grid( + (unsigned int) ((state_elements + threads - 1) / threads), + (unsigned int) n_heads, + (unsigned int) n_seqs); + (void) cudaGetLastError(); + gdn_transition_journal_commit_kernel<<>>( + (const float *) journal->data, + (float *) state->data, + (const int32_t *) accepted_prefixes->data, + (const int32_t *) active_slot_ids->data, + (int) state_size, + (int) n_heads, + (int) n_tokens, + (int) n_seqs, + (int) n_state_slots, + (int) journal_width, + journal_width == 2*state_size + 1 ? 1 : (int) state_size); + if (cudaGetLastError() != cudaSuccess) return false; + return cudaDeviceSynchronize() == cudaSuccess; +} diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index b46a6df2a..6f51347a2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -6801,6 +6801,32 @@ void ggml_gated_delta_net_set_skip_intermediate( tensor->nb[3] = tensor->nb[2]*tensor->ne[2]; } +void ggml_gated_delta_net_set_transition_journal( + struct ggml_tensor * tensor, + struct ggml_tensor * journal) { + GGML_ASSERT(tensor != NULL && journal != NULL); + GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); + GGML_ASSERT(tensor->src[6] == NULL); + GGML_ASSERT(journal->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(journal)); + + const struct ggml_tensor * v = tensor->src[2]; + const struct ggml_tensor * g = tensor->src[3]; + GGML_ASSERT(v != NULL && g != NULL); + const int64_t S_v = v->ne[0]; + const int64_t H = v->ne[1]; + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + const bool kda = g->ne[0] == S_v; + const int64_t journal_width = kda ? 3*S_v : 2*S_v + 1; + GGML_ASSERT(journal->ne[0] == journal_width && + journal->ne[1] == H && + journal->ne[2] == n_tokens && + journal->ne[3] == n_seqs); + + tensor->src[11] = journal; +} + // dflash: raw-gate mode (see ggml.h). src[8] is reserved for the optional // active-slot map; dt_bias -> src[9], A -> src[10], // op_params[2] = 1. diff --git a/server/scripts/quantize_draft_q8.py b/server/scripts/quantize_draft_q8.py index b8f72fdbc..df1793065 100644 --- a/server/scripts/quantize_draft_q8.py +++ b/server/scripts/quantize_draft_q8.py @@ -19,35 +19,24 @@ """ import argparse -import json -import struct import sys from pathlib import Path import numpy as np -sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "deps" / "llama.cpp" / "gguf-py")) +SCRIPT_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(SCRIPT_DIR.parent / "deps" / "llama.cpp" / "gguf-py")) +sys.path.insert(0, str(SCRIPT_DIR)) import gguf +import convert_dflash_to_gguf as canonical # ────────────────────────────────────────────────────────────────────── -# DFlash 27B draft architecture constants (must match dflash27b.h) +# Legacy fallback constants. The canonical converter resolves these from +# config.json + tensor shapes when the source ships model metadata. # ────────────────────────────────────────────────────────────────────── -ARCH = "qwen35-dflash-draft" -HIDDEN = 5120 -N_LAYER = 5 -N_HEAD = 32 -N_HEAD_KV = 8 -HEAD_DIM = 128 -INTERMEDIATE = 17408 -VOCAB = 248320 -N_TARGET_LAYERS = 5 -ROPE_THETA = 1_000_000.0 -RMS_EPS = 1e-6 -MASK_TOKEN_ID = 248070 -BLOCK_SIZE = 16 -CTX_LEN = 32768 +ARCH = canonical.ARCH Q8_0_BLOCK_SIZE = 32 # elements per Q8_0 block @@ -67,60 +56,28 @@ def add_qwen36_swa_metadata(writer, enabled: bool) -> None: # ────────────────────────────────────────────────────────────────────── -# Tensor name mapping — DFlash safetensors -> llama.cpp GGUF -# (Identical to convert_dflash_to_gguf.py) +# Tensor name mapping — share the DFlash2-aware canonical converter. # ────────────────────────────────────────────────────────────────────── def map_name(name: str) -> str | None: - if name == "fc.weight": return "dflash.fc.weight" - if name == "hidden_norm.weight": return "dflash.hidden_norm.weight" - if name == "norm.weight": return "output_norm.weight" - if name.startswith("layers."): - parts = name.split(".", 2) - if len(parts) < 3: return None - i = int(parts[1]) - rest = parts[2] - layer_map = { - "input_layernorm.weight": f"blk.{i}.attn_norm.weight", - "post_attention_layernorm.weight": f"blk.{i}.ffn_norm.weight", - "self_attn.q_proj.weight": f"blk.{i}.attn_q.weight", - "self_attn.k_proj.weight": f"blk.{i}.attn_k.weight", - "self_attn.v_proj.weight": f"blk.{i}.attn_v.weight", - "self_attn.o_proj.weight": f"blk.{i}.attn_output.weight", - "self_attn.q_norm.weight": f"blk.{i}.attn_q_norm.weight", - "self_attn.k_norm.weight": f"blk.{i}.attn_k_norm.weight", - "mlp.gate_proj.weight": f"blk.{i}.ffn_gate.weight", - "mlp.up_proj.weight": f"blk.{i}.ffn_up.weight", - "mlp.down_proj.weight": f"blk.{i}.ffn_down.weight", - } - return layer_map.get(rest) - return None + return canonical.map_name(name) def is_norm_tensor(gguf_name: str) -> bool: return ( gguf_name.endswith("_norm.weight") or gguf_name == "output_norm.weight" or - gguf_name == "dflash.hidden_norm.weight" + gguf_name == "dflash.hidden_norm.weight" or + gguf_name.endswith("_conv.base") ) # ────────────────────────────────────────────────────────────────────── -# safetensors reader +# safetensors reader aliases — kept public for existing tests/importers. # ────────────────────────────────────────────────────────────────────── -def load_safetensors_header(path: Path): - with open(path, "rb") as f: - header_size = struct.unpack(" bytes: - start, end = info["data_offsets"] - with open(path, "rb") as f: - f.seek(8 + header_size + start) - return f.read(end - start) +load_safetensors_header = canonical.load_safetensors_header +read_tensor_bytes = canonical.read_tensor_bytes def bf16_bytes_to_f32(raw: bytes, shape: list[int]) -> np.ndarray: @@ -128,6 +85,44 @@ def bf16_bytes_to_f32(raw: bytes, shape: list[int]) -> np.ndarray: u32 = (u16.astype(np.uint32) << 16) return u32.view(" None: + """Write the same resolved architecture profile as the F16 converter.""" + writer.add_string("general.name", f"DFlash-Draft-{a['hidden']}h-{a['n_layer']}L-Q8_0") + writer.add_quantization_version(gguf.GGML_QUANT_VERSION) + writer.add_uint32(f"{ARCH}.context_length", a["ctx_len"]) + writer.add_uint32(f"{ARCH}.embedding_length", a["hidden"]) + writer.add_uint32(f"{ARCH}.block_count", a["n_layer"]) + writer.add_uint32(f"{ARCH}.feed_forward_length", a["intermediate"]) + writer.add_uint32(f"{ARCH}.attention.head_count", a["n_head"]) + writer.add_uint32(f"{ARCH}.attention.head_count_kv", a["n_head_kv"]) + writer.add_uint32(f"{ARCH}.attention.key_length", a["head_dim"]) + writer.add_uint32(f"{ARCH}.attention.value_length", a["head_dim"]) + writer.add_uint32(f"{ARCH}.vocab_size", a["vocab"]) + writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", a["rms_eps"]) + writer.add_float32(f"{ARCH}.rope.freq_base", a["rope_theta"]) + + if qwen36_swa: + add_qwen36_swa_metadata(writer, True) + elif a.get("swa_pattern"): + writer.add_uint32(f"{ARCH}.attention.sliding_window", a["swa_window"]) + writer.add_array(f"{ARCH}.attention.sliding_window_pattern", [bool(x) for x in a["swa_pattern"]]) + + writer.add_uint32(f"{ARCH}.dflash.n_target_layers", a["n_target_layers"]) + writer.add_uint32(f"{ARCH}.dflash.block_size", a["block_size"]) + writer.add_uint32(f"{ARCH}.dflash.mask_token_id", a["mask_token_id"]) + capture_ids = a.get("capture_layer_ids") + if capture_ids and len(capture_ids) == a["n_target_layers"]: + writer.add_array(f"{ARCH}.dflash.target_layer_ids", [int(x) for x in capture_ids]) + elif capture_ids: + print(f"[warn] capture_layer_ids len {len(capture_ids)} != n_target_layers {a['n_target_layers']}; not embedding ids", file=sys.stderr) + + if a.get("conv_kernel_size"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_kernel_size", a["conv_kernel_size"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.conv_group_size", a["conv_group_size"]) + if a.get("selector_rank"): + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_rank", a["selector_rank"]) + writer.add_uint32(f"{ARCH}.dflash.dflash2.selector_top_k", a["selector_top_k"]) + # ────────────────────────────────────────────────────────────────────── # Main @@ -160,32 +155,14 @@ def main(): header_size, header = load_safetensors_header(args.safetensors) n_entries = sum(1 for k in header if k != "__metadata__") print(f"[info] {n_entries} tensor entries") + arch = canonical.load_arch(args.safetensors, header) writer = gguf.GGUFWriter(args.out_gguf, ARCH) - # Architecture metadata (identical to convert_dflash_to_gguf.py) - writer.add_string("general.name", "Qwen3.5-27B-DFlash-Draft-Q8_0") - writer.add_quantization_version(gguf.GGML_QUANT_VERSION) - writer.add_uint32(f"{ARCH}.context_length", CTX_LEN) - writer.add_uint32(f"{ARCH}.embedding_length", HIDDEN) - writer.add_uint32(f"{ARCH}.block_count", N_LAYER) - writer.add_uint32(f"{ARCH}.feed_forward_length", INTERMEDIATE) - writer.add_uint32(f"{ARCH}.attention.head_count", N_HEAD) - writer.add_uint32(f"{ARCH}.attention.head_count_kv", N_HEAD_KV) - writer.add_uint32(f"{ARCH}.attention.key_length", HEAD_DIM) - writer.add_uint32(f"{ARCH}.attention.value_length", HEAD_DIM) - writer.add_uint32(f"{ARCH}.vocab_size", VOCAB) - writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", RMS_EPS) - writer.add_float32(f"{ARCH}.rope.freq_base", ROPE_THETA) - add_qwen36_swa_metadata(writer, args.qwen36_swa) + add_arch_metadata(writer, arch, args.qwen36_swa) if args.qwen36_swa: print("[info] Qwen3.6 draft SWA: layers 0-3 window=2048; layer 4 full attention") - # DFlash-specific hyperparameters - writer.add_uint32(f"{ARCH}.dflash.n_target_layers", N_TARGET_LAYERS) - writer.add_uint32(f"{ARCH}.dflash.block_size", BLOCK_SIZE) - writer.add_uint32(f"{ARCH}.dflash.mask_token_id", MASK_TOKEN_ID) - # Collect and sort tensors (same order as convert_dflash_to_gguf.py) pending = [] for st_name, info in header.items(): diff --git a/server/src/common/concurrency/chain_spec_shapes.h b/server/src/common/concurrency/chain_spec_shapes.h index 0b392d32d..b283ef966 100644 --- a/server/src/common/concurrency/chain_spec_shapes.h +++ b/server/src/common/concurrency/chain_spec_shapes.h @@ -51,6 +51,29 @@ inline DDTree make_dspark_chain_tree( return tree; } +// A chain verify always includes the pending root. Depth 1 would therefore +// be an AR-equivalent target step, which is forbidden once a request has +// committed to sticky speculation. `requested == 0` means use the configured +// drafter maximum; any other invalid value fails closed. +inline int resolve_chain_verify_depth(int requested, int maximum) { + if (maximum < 2) return 0; + if (requested == 0) return maximum; + return requested >= 2 && requested <= maximum ? requested : 0; +} + +// Keep proposal generation at its configured maximum while allowing one +// common verify depth to be selected per round. Failure leaves the proposal +// untouched, which makes malformed controller/config output non-destructive. +inline bool truncate_chain_proposal( + std::vector & draft_tokens, int verify_depth) { + if (verify_depth < 2 || + verify_depth > static_cast(draft_tokens.size())) { + return false; + } + draft_tokens.resize(static_cast(verify_depth)); + return true; +} + struct ChainLaunchShape { int spec_lanes = 0; int tree_bucket = 0; diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h index b389055cf..e0eb51c56 100644 --- a/server/src/common/concurrency/speculation_gate.h +++ b/server/src/common/concurrency/speculation_gate.h @@ -1,6 +1,6 @@ // Per-request adaptive speculation policy over startup-profiled costs. Every // adaptive request receives one AR or speculation decision and keeps that -// decision until forget(). A failed confidence evaluation is represented +// decision until forget(). A failed activation-score evaluation is represented // explicitly and commits sticky AR without inventing a score. // Pure host code: no graph, backend, or scheduler types belong here. @@ -22,7 +22,7 @@ namespace dflash::common { struct SpecGateConfig { // Immutable offline fit applied independently to every request's one-time - // confidence score. It never learns from request execution history. + // activation score. It never learns from request execution history. double fixed_yield_scale = 1.0; double cost_ema_alpha = 0.20; double adaptive_gain_margin = 0.02; @@ -82,6 +82,22 @@ struct SpecCostTables { } }; +enum class SpecScoreKind : uint8_t { + Unspecified, + DSparkConfidence, + DFlash2SelectorBenefitV1, +}; + +inline const char * spec_score_kind_name(SpecScoreKind kind) { + switch (kind) { + case SpecScoreKind::Unspecified: return "unspecified"; + case SpecScoreKind::DSparkConfidence: return "dspark_confidence"; + case SpecScoreKind::DFlash2SelectorBenefitV1: + return "dflash2_selector_benefit_v1"; + } + return "unknown"; +} + struct SpecCandidate { uint64_t request_id = 0; int slot = -1; @@ -101,6 +117,7 @@ struct SpecCandidate { // survival-product expected yield, including the root; the gate applies // only its fixed offline scale and clamping. double confidence_yield = std::numeric_limits::quiet_NaN(); + SpecScoreKind score_kind = SpecScoreKind::Unspecified; }; struct SpecStepGeometry { @@ -115,10 +132,10 @@ struct SpecStepGeometry { int tree_rows(int spec_lanes) const { return bucketed_lanes(spec_lanes) * std::max(1, tree_width); } - int step_rows(int concurrency, int spec_lanes, - double expected_spec_tokens) const { - const int accepted_rows = std::max( - spec_lanes, static_cast(std::lround(expected_spec_tokens))); + double expected_step_rows(int concurrency, int spec_lanes, + double expected_spec_tokens) const { + const double accepted_rows = std::max( + static_cast(spec_lanes), expected_spec_tokens); return accepted_rows + bucketed_lanes(concurrency - spec_lanes); } }; @@ -176,6 +193,27 @@ struct SpecPlanScore { // AR/speculation decision. This supports one activation record per // request without treating later sticky execution as a new decision. bool newly_decided = false; + SpecScoreKind score_kind = SpecScoreKind::Unspecified; + bool execution_unsupported = false; +}; + +// Discrete graph shape that actually ran. Unlike a SpecPlan's fractional +// expected replay rows, every field here comes from executor telemetry and is +// safe to use as an online timing key. +struct SpecExecutionShape { + int concurrency = 0; + int admitted_count = 0; + int tree_rows = 0; + int step_rows = 0; + int draft_lanes = 0; + + bool operator==(const SpecExecutionShape & other) const { + return concurrency == other.concurrency && + admitted_count == other.admitted_count && + tree_rows == other.tree_rows && + step_rows == other.step_rows && + draft_lanes == other.draft_lanes; + } }; struct SpecPlan { @@ -184,7 +222,7 @@ struct SpecPlan { int concurrency = 0; int admitted_count = 0; int tree_rows = 0; - int step_rows = 0; + double expected_step_rows = 0.0; int draft_lanes = 0; double expected_tokens = 0.0; // Startup-profiled cost before online correction, and the shape-local @@ -207,19 +245,17 @@ struct SpecPlan { std::vector ordered; std::vector admitted_request_ids; std::vector admitted_slots; - // Score actions are batched for one-time confidence initialization. + // Score actions are batched for one-time activation-score initialization. // FallbackAR actions cannot attempt scoring and must instead commit sticky // AR with an explicit failed-evaluation activation. One tagged record keeps // request identity and slot inseparable on all failure paths. std::vector pending_evaluations; }; -// Generic confidence contract for every chain speculator: `confidences[i]` -// is a probability-like, monotone-in-acceptance score for accepting position -// i of the current block conditioned on the preceding positions. The gate -// converts that adapter-owned vector into an expected emitted-token yield. It -// does not depend on how a producer obtains the scores (trained head, selector -// softmax, or a future speculator-specific readout). +// DSpark confidence-head contract: `confidences[i]` is a probability-like, +// monotone-in-acceptance score for position i conditioned on its prefix. +// DFlash2 selector evidence must go through its model-specific benefit adapter +// and must never be passed to this helper as though it were confidence. inline double confidence_survival_yield( const std::vector & confidences, int max_accept) { if (max_accept <= 1) return 1.0; @@ -242,26 +278,11 @@ class SpeculationGate { std::numeric_limits::quiet_NaN(); SpecDecision decision = SpecDecision::Undecided; bool confidence_evaluation_failed = false; + SpecScoreKind score_kind = SpecScoreKind::Unspecified; }; - struct CostShape { - int concurrency = 0; - int admitted_count = 0; - int tree_rows = 0; - int step_rows = 0; - int draft_lanes = 0; - - bool operator==(const CostShape & other) const { - return concurrency == other.concurrency && - admitted_count == other.admitted_count && - tree_rows == other.tree_rows && - step_rows == other.step_rows && - draft_lanes == other.draft_lanes; - } - }; - - struct CostShapeHash { - size_t operator()(const CostShape & shape) const { + struct ExecutionShapeHash { + size_t operator()(const SpecExecutionShape & shape) const { size_t seed = 0; auto mix = [&](int value) { seed ^= std::hash{}(value) + @@ -282,9 +303,15 @@ class SpeculationGate { uint64_t observations = 0; }; + struct CostPrice { + double profiled = std::numeric_limits::infinity(); + double predicted = std::numeric_limits::infinity(); + }; + struct CandidateScore { double expected_yield = 1.0; SpecScoreSource source = SpecScoreSource::Unavailable; + SpecScoreKind score_kind = SpecScoreKind::Unspecified; }; public: @@ -337,6 +364,7 @@ class SpeculationGate { const SpecCandidate * candidate = nullptr; double score = 1.0; SpecScoreSource source = SpecScoreSource::Unavailable; + SpecScoreKind score_kind = SpecScoreKind::Unspecified; SpecDecision decision = SpecDecision::Undecided; bool forced = false; bool commit_candidate = false; @@ -379,11 +407,11 @@ class SpeculationGate { // but permanently unsupported execution commits directly to AR. forced_ar.push_back({ &candidate, score.expected_yield, score.source, - prior_decision, false, true}); + score.score_kind, prior_decision, false, true}); continue; } Ranked ranked{&candidate, score.expected_yield, score.source, - prior_decision, + score.score_kind, prior_decision, candidate.policy == SpeculationPolicy::Always || prior_decision == SpecDecision::Speculation, adaptive_undecided}; @@ -433,7 +461,7 @@ class SpeculationGate { item.score, item.source, item.decision, item.forced, false, - item.commit_candidate}); + item.commit_candidate, item.score_kind}); } const int forced_count = static_cast(forced.size()); @@ -441,24 +469,22 @@ class SpeculationGate { double expected_sum = 0.0; for (int i = 0; i < forced_count; ++i) expected_sum += ranked[i].score; - const SpecCostLookup ar_lookup = costs_.step_cost.lookup( - geometry_.bucketed_lanes(concurrency)); - report_clamp("step", ar_lookup, out); - const CostShape ar_shape{ + const SpecExecutionShape ar_shape{ concurrency, 0, 0, geometry_.bucketed_lanes(concurrency), 0}; - const double ar_scale = cost_scale(ar_shape); + const CostPrice ar_price = price_execution_shape(ar_shape, &out); out.ar_goodput = concurrency == 0 ? 0.0 : static_cast(concurrency) / - (ar_lookup.cost * ar_scale); + ar_price.predicted; struct PlanPoint { int k = 0; double goodput = -1.0; double profiled_cost = 0.0; + double predicted_cost = 0.0; double cost_scale = 1.0; double expected_tokens = 0.0; int tree_rows = 0; - int step_rows = 0; + double expected_step_rows = 0.0; int draft_lanes = 0; }; PlanPoint baseline; @@ -468,41 +494,25 @@ class SpeculationGate { if (k > forced_count) expected_sum += ranked[k - 1].score; const double expected_tokens = static_cast(concurrency - k) + expected_sum; - double profiled_cost = 0.0; int tree_rows = 0; - int step_rows = geometry_.bucketed_lanes(concurrency); + double expected_step_rows = geometry_.bucketed_lanes(concurrency); const int draft_lanes = draft_lanes_override >= 0 ? draft_lanes_override : k; - if (k == 0 && draft_lanes == 0) { - profiled_cost = ar_lookup.cost; - } else { - if (k > 0) { - tree_rows = geometry_.tree_rows(k); - const SpecCostLookup tree = costs_.tree_cost.lookup(tree_rows); - report_clamp("tree", tree, out); - profiled_cost += tree.cost; - step_rows = geometry_.step_rows( - concurrency, k, expected_sum); - } - const SpecCostLookup step = costs_.step_cost.lookup(step_rows); - report_clamp("step", step, out); - profiled_cost += step.cost; - if (draft_lanes > 0) { - const SpecCostLookup draft = - costs_.draft_cost.lookup(draft_lanes); - report_clamp("draft", draft, out); - profiled_cost += draft.cost; - } + if (k > 0) { + tree_rows = geometry_.tree_rows(k); + expected_step_rows = geometry_.expected_step_rows( + concurrency, k, expected_sum); } - const CostShape shape{ - concurrency, k, tree_rows, step_rows, draft_lanes}; - const double scale = cost_scale(shape); - const double goodput = - expected_tokens / (profiled_cost * scale); + const CostPrice price = price_expected_shape( + {concurrency, k, tree_rows, 0, draft_lanes}, + expected_step_rows, &out); + const double scale = price.predicted / price.profiled; + const double goodput = expected_tokens / price.predicted; + const PlanPoint point{ - k, goodput, profiled_cost, scale, expected_tokens, - tree_rows, step_rows, draft_lanes}; + k, goodput, price.profiled, price.predicted, scale, + expected_tokens, tree_rows, expected_step_rows, draft_lanes}; if (k == forced_count) baseline = point; if (goodput > best.goodput) { best = point; @@ -522,10 +532,10 @@ class SpeculationGate { out.goodput = std::max(0.0, best.goodput); out.profiled_cost = best.profiled_cost; out.cost_scale = best.cost_scale; - out.predicted_cost = best.profiled_cost * best.cost_scale; + out.predicted_cost = best.predicted_cost; out.expected_tokens = best.expected_tokens; out.tree_rows = best.tree_rows; - out.step_rows = best.step_rows; + out.expected_step_rows = best.expected_step_rows; out.draft_lanes = best.draft_lanes; for (size_t i = 0; i < ranked.size(); ++i) { if (!ranked[i].commit_candidate) continue; @@ -554,30 +564,36 @@ class SpeculationGate { false, false, true, + item.score_kind, + true, }); } } return out; } - void observe_cost(const SpecPlan & plan, double measured_us) { - if (!plan.valid || !std::isfinite(measured_us) || measured_us <= 0.0 || - !std::isfinite(plan.profiled_cost) || plan.profiled_cost <= 0.0) { + void observe_cost(const SpecExecutionShape & executed, + double measured_us) { + if (!std::isfinite(measured_us) || measured_us <= 0.0 || + executed.concurrency < 0 || executed.admitted_count < 0 || + executed.admitted_count > executed.concurrency || + executed.tree_rows < 0 || executed.step_rows <= 0 || + executed.draft_lanes < 0) { return; } - const CostShape shape{ - plan.concurrency, plan.admitted_count, plan.tree_rows, - plan.step_rows, plan.draft_lanes}; + const double profiled_cost = + profile_execution_shape(executed, nullptr); + if (!std::isfinite(profiled_cost) || profiled_cost <= 0.0) return; const double ratio = std::clamp( - measured_us / plan.profiled_cost, + measured_us / profiled_cost, kCostScaleMin, kCostScaleMax); - CostState & state = cost_states_[shape]; + CostState & state = cost_states_[executed]; update_ema(state.scale, state.observations, ratio, config_.cost_ema_alpha); } // Commit the explicit cold-evaluation failure policy. This is a real - // sticky AR decision, but intentionally has no initial confidence value. + // sticky AR decision, but intentionally has no initial activation score. // False means the request had already committed a mode. bool commit_evaluation_fallback_ar(uint64_t request_id) { RequestState & state = request_states_[request_id]; @@ -612,6 +628,11 @@ class SpeculationGate { ? std::numeric_limits::quiet_NaN() : state->second.initial_confidence; } + SpecScoreKind initial_score_kind(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state == request_states_.end() + ? SpecScoreKind::Unspecified : state->second.score_kind; + } SpecDecision decision(uint64_t request_id) const { auto state = request_states_.find(request_id); return state == request_states_.end() @@ -630,6 +651,7 @@ class SpeculationGate { RequestState & state = request_states_[candidate.request_id]; if (!std::isfinite(state.initial_confidence)) { state.initial_confidence = raw; + state.score_kind = candidate.score_kind; accepted_initial_score = true; } return { @@ -638,6 +660,7 @@ class SpeculationGate { 1.0, static_cast(max_accept_)), accepted_initial_score ? SpecScoreSource::Confidence : SpecScoreSource::InitialConfidence, + state.score_kind, }; } auto state = request_states_.find(candidate.request_id); @@ -649,9 +672,11 @@ class SpeculationGate { state->second.initial_confidence, 1.0, static_cast(max_accept_)), SpecScoreSource::InitialConfidence, + state->second.score_kind, }; } - return {1.0, SpecScoreSource::Unavailable}; + return {1.0, SpecScoreSource::Unavailable, + SpecScoreKind::Unspecified}; } static void update_ema(double & value, uint64_t & observations, @@ -661,12 +686,65 @@ class SpeculationGate { ++observations; } - double cost_scale(const CostShape & shape) const { + double cost_scale(const SpecExecutionShape & shape) const { auto state = cost_states_.find(shape); return state == cost_states_.end() || state->second.observations == 0 ? 1.0 : state->second.scale; } + double profile_execution_shape(const SpecExecutionShape & shape, + SpecPlan * plan) const { + double cost = 0.0; + auto add = [&](const char * name, const SpecCostLookup & lookup) { + if (plan) report_clamp(name, lookup, *plan); + cost += lookup.cost; + }; + + if (shape.admitted_count > 0) { + add("tree", costs_.tree_cost.lookup(shape.tree_rows)); + } + add("step", costs_.step_cost.lookup(shape.step_rows)); + if (shape.draft_lanes > 0) { + add("draft", costs_.draft_cost.lookup(shape.draft_lanes)); + } + return cost; + } + + CostPrice price_execution_shape(const SpecExecutionShape & shape, + SpecPlan * plan) const { + const double profiled = profile_execution_shape(shape, plan); + return {profiled, profiled * cost_scale(shape)}; + } + + CostPrice price_expected_shape(SpecExecutionShape shape, + double expected_step_rows, + SpecPlan * plan) const { + if (!std::isfinite(expected_step_rows) || expected_step_rows < 0.0) + return {}; + + // The executor can only launch an integer row count, while the gate + // owns an expected count. Price that expectation continuously across + // the two neighboring executable shapes so nonlinear profile cliffs + // retain their cost without an lround() decision discontinuity. Apply + // each neighbor's own online correction before interpolating it. + const int lower_rows = + static_cast(std::floor(expected_step_rows)); + const int upper_rows = + static_cast(std::ceil(expected_step_rows)); + shape.step_rows = lower_rows; + const CostPrice lower = price_execution_shape(shape, plan); + if (lower_rows == upper_rows) return lower; + + shape.step_rows = upper_rows; + const CostPrice upper = price_execution_shape(shape, plan); + const double upper_weight = expected_step_rows - lower_rows; + return { + lower.profiled + upper_weight * (upper.profiled - lower.profiled), + lower.predicted + + upper_weight * (upper.predicted - lower.predicted), + }; + } + void report_clamp(const char * name, const SpecCostLookup & lookup, SpecPlan & plan) const { if (!lookup.clamped) return; @@ -684,7 +762,8 @@ class SpeculationGate { static constexpr double kCostScaleMin = 0.25; static constexpr double kCostScaleMax = 4.0; std::unordered_map request_states_; - std::unordered_map cost_states_; + std::unordered_map + cost_states_; ClampLogger clamp_logger_; }; diff --git a/server/src/common/dflash2_batch.cpp b/server/src/common/dflash2_batch.cpp new file mode 100644 index 000000000..879564798 --- /dev/null +++ b/server/src/common/dflash2_batch.cpp @@ -0,0 +1,440 @@ +#include "dflash2_head.h" + +#include "dflash2_selector_validation.h" +#include "ddtree.h" +#include "geometric_draft_topk_cuda.h" +#include "ggml-alloc.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { +namespace { + +struct ProjectionGraph { + const DraftWeights * dw = nullptr; + ggml_backend_t backend = nullptr; + ggml_tensor * lm_head = nullptr; + int n_positions = 0; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_tensor * inp_hidden = nullptr; + ggml_tensor * logits = nullptr; +}; + +struct BatchedSelectorGraph { + const DraftWeights * dw = nullptr; + ggml_backend_t backend = nullptr; + int n_lanes = 0; + int n_cand = 0; + int K = 0; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_tensor * inp_hidden = nullptr; + ggml_tensor * inp_succ = nullptr; + ggml_tensor * inp_pred = nullptr; + ggml_tensor * hproj = nullptr; + ggml_tensor * succ = nullptr; + ggml_tensor * pred = nullptr; +}; + +ProjectionGraph & projection_graph() { + static thread_local ProjectionGraph graph; + return graph; +} + +BatchedSelectorGraph & batched_selector_graph() { + static thread_local BatchedSelectorGraph graph; + return graph; +} + +void free_projection_graph(ProjectionGraph & graph) { + if (graph.galloc) { + ggml_gallocr_free(graph.galloc); + graph.galloc = nullptr; + } + if (graph.ctx) { + ggml_free(graph.ctx); + graph.ctx = nullptr; + } + graph = {}; +} + +void free_selector_graph(BatchedSelectorGraph & graph) { + if (graph.galloc) { + ggml_gallocr_free(graph.galloc); + graph.galloc = nullptr; + } + if (graph.ctx) { + ggml_free(graph.ctx); + graph.ctx = nullptr; + } + graph = {}; +} + +bool ensure_projection_graph( + ProjectionGraph & graph, const DraftWeights & dw, + ggml_backend_t backend, ggml_tensor * lm_head, int n_positions) { + if (graph.ctx && graph.dw == &dw && graph.backend == backend && + graph.lm_head == lm_head && graph.n_positions == n_positions) { + return true; + } + free_projection_graph(graph); + if (!backend || !lm_head || n_positions <= 0 || dw.n_embd <= 0 || + lm_head->ne[0] != dw.n_embd || lm_head->ne[1] <= 0) { + return false; + } + + const size_t arena_size = + ggml_tensor_overhead() * 32 + + ggml_graph_overhead_custom(256, false) + 4096; + graph.arena.assign(arena_size, 0); + ggml_init_params params{}; + params.mem_size = graph.arena.size(); + params.mem_buffer = graph.arena.data(); + params.no_alloc = true; + graph.ctx = ggml_init(params); + if (!graph.ctx) return false; + graph.gf = ggml_new_graph_custom(graph.ctx, 256, false); + graph.inp_hidden = ggml_new_tensor_2d( + graph.ctx, GGML_TYPE_F32, dw.n_embd, n_positions); + ggml_set_input(graph.inp_hidden); + graph.logits = ggml_mul_mat(graph.ctx, lm_head, graph.inp_hidden); + ggml_set_output(graph.logits); + ggml_build_forward_expand(graph.gf, graph.logits); + graph.galloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!graph.galloc || !ggml_gallocr_alloc_graph(graph.galloc, graph.gf)) { + std::fprintf(stderr, + "dflash2_select_chains_batched: projection graph alloc failed\n"); + free_projection_graph(graph); + return false; + } + graph.dw = &dw; + graph.backend = backend; + graph.lm_head = lm_head; + graph.n_positions = n_positions; + return true; +} + +bool ensure_selector_graph( + BatchedSelectorGraph & graph, const DraftWeights & dw, + ggml_backend_t backend, int n_lanes, int n_cand, int K) { + if (graph.ctx && graph.dw == &dw && graph.backend == backend && + graph.n_lanes == n_lanes && graph.n_cand == n_cand && + graph.K == K) { + return true; + } + free_selector_graph(graph); + const DraftSelectorWeights & selector = dw.selector; + if (!backend || n_lanes <= 0 || n_cand <= 0 || K <= 0 || + dw.n_embd <= 0 || selector.rank <= 0 || !selector.hproj || + !selector.pred_cb || !selector.succ_cb) { + return false; + } + + const int n_positions = n_lanes * n_cand; + const int n_pred_rows = n_lanes + n_positions * K; + const size_t arena_size = + ggml_tensor_overhead() * 48 + + ggml_graph_overhead_custom(256, false) + 4096; + graph.arena.assign(arena_size, 0); + ggml_init_params params{}; + params.mem_size = graph.arena.size(); + params.mem_buffer = graph.arena.data(); + params.no_alloc = true; + graph.ctx = ggml_init(params); + if (!graph.ctx) return false; + graph.gf = ggml_new_graph_custom(graph.ctx, 256, false); + graph.inp_hidden = ggml_new_tensor_2d( + graph.ctx, GGML_TYPE_F32, dw.n_embd, n_positions); + graph.inp_succ = ggml_new_tensor_1d( + graph.ctx, GGML_TYPE_I32, n_positions * K); + graph.inp_pred = ggml_new_tensor_1d( + graph.ctx, GGML_TYPE_I32, n_pred_rows); + ggml_set_input(graph.inp_hidden); + ggml_set_input(graph.inp_succ); + ggml_set_input(graph.inp_pred); + graph.hproj = + ggml_mul_mat(graph.ctx, selector.hproj, graph.inp_hidden); + graph.succ = + ggml_get_rows(graph.ctx, selector.succ_cb, graph.inp_succ); + graph.pred = + ggml_get_rows(graph.ctx, selector.pred_cb, graph.inp_pred); + ggml_set_output(graph.hproj); + ggml_set_output(graph.succ); + ggml_set_output(graph.pred); + ggml_build_forward_expand(graph.gf, graph.hproj); + ggml_build_forward_expand(graph.gf, graph.succ); + ggml_build_forward_expand(graph.gf, graph.pred); + graph.galloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!graph.galloc || !ggml_gallocr_alloc_graph(graph.galloc, graph.gf)) { + std::fprintf(stderr, + "dflash2_select_chains_batched: selector graph alloc failed\n"); + free_selector_graph(graph); + return false; + } + graph.dw = &dw; + graph.backend = backend; + graph.n_lanes = n_lanes; + graph.n_cand = n_cand; + graph.K = K; + return true; +} + +DFlash2DepthSignal summarize_depth( + const float * log_probs, const std::vector & scores, + int K, int selected) { + DFlash2DepthSignal signal; + if (!log_probs || K <= 0 || selected < 0 || selected >= K || + static_cast(scores.size()) != K) { + return signal; + } + signal.selected_log_prob = log_probs[selected]; + signal.lm_top2_margin = K > 1 ? log_probs[0] - log_probs[1] + : std::numeric_limits::infinity(); + float top_k_mass = 0.0f; + for (int k = 0; k < K; ++k) top_k_mass += std::exp(log_probs[k]); + signal.top_k_mass = std::clamp(top_k_mass, 0.0f, 1.0f); + signal.selected_rank = selected; + signal.agrees_with_lm_top1 = selected == 0; + + float runner_up = -INFINITY; + for (int k = 0; k < K; ++k) { + if (k != selected) { + runner_up = std::max(runner_up, scores[(size_t) k]); + } + } + signal.selector_margin = K > 1 + ? scores[(size_t) selected] - runner_up + : std::numeric_limits::infinity(); + + const float max_score = + *std::max_element(scores.begin(), scores.end()); + float z = 0.0f; + for (float score : scores) z += std::exp(score - max_score); + if (z > 0.0f && std::isfinite(z)) { + signal.selector_winner_mass = + std::exp(scores[(size_t) selected] - max_score) / z; + float entropy = 0.0f; + for (float score : scores) { + const float p = std::exp(score - max_score) / z; + if (p > 0.0f) entropy -= p * std::log(p); + } + signal.selector_entropy = entropy; + } + return signal; +} + +} // namespace + +bool dflash2_select_chains_batched( + const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const std::vector & hidden_by_lane, + int q_len, + const std::vector & last_tokens, + std::vector> & draft_tokens, + std::vector * traces) { + draft_tokens.clear(); + if (traces) traces->clear(); + const DraftSelectorWeights & selector = dw.selector; + const int n_lanes = static_cast(hidden_by_lane.size()); + const int n_cand = q_len - 1; + const int K = selector.top_k; + const int rank = selector.rank; + const int hdim = dw.n_embd; + if (!selector.enabled || !selector.hproj || !selector.pred_cb || + !selector.succ_cb || !backend || !lm_head || n_lanes <= 0 || + static_cast(last_tokens.size()) != n_lanes || + n_cand <= 0 || K <= 0 || rank <= 0 || hdim <= 0) { + return false; + } + DFlash2SelectorLayout selector_layout; + selector_layout.rank = rank; + selector_layout.top_k = K; + selector_layout.hproj_rank = selector.hproj->ne[1]; + selector_layout.pred_rank = selector.pred_cb->ne[0]; + selector_layout.pred_vocab = selector.pred_cb->ne[1]; + selector_layout.succ_rank = selector.succ_cb->ne[0]; + selector_layout.succ_vocab = selector.succ_cb->ne[1]; + selector_layout.target_output_vocab = lm_head->ne[1]; + std::string selector_error; + if (!validate_dflash2_selector_layout( + selector_layout, selector_error)) { + std::fprintf(stderr, "dflash2_select_chains_batched: %s\n", + selector_error.c_str()); + return false; + } + for (const float * hidden : hidden_by_lane) { + if (!hidden) return false; + } + + const int n_positions = n_lanes * n_cand; + std::vector candidate_hidden( + (size_t) hdim * (size_t) n_positions); + for (int lane = 0; lane < n_lanes; ++lane) { + for (int depth = 0; depth < n_cand; ++depth) { + const int position = lane * n_cand + depth; + const float * source = hidden_by_lane[(size_t) lane] + + (size_t) (depth + 1) * (size_t) hdim; + std::memcpy( + candidate_hidden.data() + + (size_t) position * (size_t) hdim, + source, sizeof(float) * (size_t) hdim); + } + } + + ProjectionGraph & projection = projection_graph(); + if (!ensure_projection_graph( + projection, dw, backend, lm_head, n_positions)) { + return false; + } + ggml_backend_tensor_set( + projection.inp_hidden, candidate_hidden.data(), 0, + sizeof(float) * candidate_hidden.size()); + if (ggml_backend_graph_compute(backend, projection.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "dflash2_select_chains_batched: projection compute failed\n"); + return false; + } + + const int vocab = static_cast(lm_head->ne[1]); + std::vector candidate_log_probs( + (size_t) n_positions * (size_t) K); + std::vector candidate_ids( + (size_t) n_positions * (size_t) K); + bool have_top_k = false; +#ifdef DFLASH27B_HAVE_DRAFT_TOPK + static const bool gpu_top_k = []() { + const char * value = std::getenv("DFLASH_GPU_DRAFT_TOPK"); + return value == nullptr || value[0] != '0'; + }(); + if (gpu_top_k && projection.logits && projection.logits->data) { + have_top_k = geometric_extract_draft_topk_cuda( + projection.logits->data, n_positions, vocab, K, + candidate_log_probs.data(), candidate_ids.data(), 1.0f); + } +#endif + if (!have_top_k) { + std::vector logits( + (size_t) vocab * (size_t) n_positions); + ggml_backend_tensor_get( + projection.logits, logits.data(), 0, + sizeof(float) * logits.size()); + extract_draft_topk( + logits.data(), n_positions, vocab, K, + candidate_log_probs.data(), candidate_ids.data(), 1.0f); + } + + BatchedSelectorGraph & graph = batched_selector_graph(); + if (!ensure_selector_graph( + graph, dw, backend, n_lanes, n_cand, K)) { + return false; + } + std::vector predecessor_ids( + (size_t) n_lanes + candidate_ids.size()); + std::copy( + last_tokens.begin(), last_tokens.end(), predecessor_ids.begin()); + std::copy( + candidate_ids.begin(), candidate_ids.end(), + predecessor_ids.begin() + n_lanes); + ggml_backend_tensor_set( + graph.inp_hidden, candidate_hidden.data(), 0, + sizeof(float) * candidate_hidden.size()); + ggml_backend_tensor_set( + graph.inp_succ, candidate_ids.data(), 0, + sizeof(int32_t) * candidate_ids.size()); + ggml_backend_tensor_set( + graph.inp_pred, predecessor_ids.data(), 0, + sizeof(int32_t) * predecessor_ids.size()); + if (ggml_backend_graph_compute(backend, graph.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "dflash2_select_chains_batched: selector compute failed\n"); + return false; + } + + std::vector projected_hidden( + (size_t) rank * (size_t) n_positions); + std::vector successor_codes( + (size_t) rank * candidate_ids.size()); + std::vector predecessor_codes( + (size_t) rank * predecessor_ids.size()); + ggml_backend_tensor_get_async( + backend, graph.hproj, projected_hidden.data(), 0, + sizeof(float) * projected_hidden.size()); + ggml_backend_tensor_get_async( + backend, graph.succ, successor_codes.data(), 0, + sizeof(float) * successor_codes.size()); + ggml_backend_tensor_get_async( + backend, graph.pred, predecessor_codes.data(), 0, + sizeof(float) * predecessor_codes.size()); + ggml_backend_synchronize(backend); + + draft_tokens.assign( + (size_t) n_lanes, + std::vector((size_t) q_len)); + if (traces) traces->resize((size_t) n_lanes); + for (int lane = 0; lane < n_lanes; ++lane) { + draft_tokens[(size_t) lane][0] = last_tokens[(size_t) lane]; + if (traces) { + (*traces)[(size_t) lane].depths.reserve((size_t) n_cand); + } + int predecessor_row = lane; + for (int depth = 0; depth < n_cand; ++depth) { + const int position = lane * n_cand + depth; + const float * predecessor = predecessor_codes.data() + + (size_t) predecessor_row * (size_t) rank; + const float * hidden = projected_hidden.data() + + (size_t) position * (size_t) rank; + float best_score = -INFINITY; + int best_candidate = 0; + std::vector scores((size_t) K); + for (int candidate = 0; candidate < K; ++candidate) { + const int candidate_row = position * K + candidate; + const float * successor = successor_codes.data() + + (size_t) candidate_row * (size_t) rank; + float correction = 0.0f; + for (int r = 0; r < rank; ++r) { + correction += + predecessor[r] * hidden[r] * successor[r]; + } + const float score = + candidate_log_probs[(size_t) candidate_row] + + correction; + scores[(size_t) candidate] = score; + if (score > best_score) { + best_score = score; + best_candidate = candidate; + } + } + const int selected_row = position * K + best_candidate; + draft_tokens[(size_t) lane][(size_t) depth + 1] = + candidate_ids[(size_t) selected_row]; + if (traces) { + (*traces)[(size_t) lane].depths.push_back( + summarize_depth( + candidate_log_probs.data() + + (size_t) position * (size_t) K, + scores, K, best_candidate)); + } + predecessor_row = n_lanes + selected_row; + } + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash2_benefit.cpp b/server/src/common/dflash2_benefit.cpp new file mode 100644 index 000000000..d2d5873b8 --- /dev/null +++ b/server/src/common/dflash2_benefit.cpp @@ -0,0 +1,198 @@ +#include "dflash2_benefit.h" + +#include +#include +#include +#include +#include +#include + +namespace dflash::common { +namespace { + +// Seed artifacts (full hashes retained for offline provenance; startup uses the +// cheap size plus structural signature): target IQ4_XS sha256 +// 4e44edf892af6d57506fcd9eeaf5d0628f8737cfdce89652cb9d9bff82808eae, +// draft Q8_0 sha256 +// bb727abc583498aa4deea8b3cd0c34c2d96553954cbff25b5f7bdd469f0f1306. +constexpr DFlash2BenefitModelSignature kSeededQwen38DFlash2 = { + /*target_layers=*/64, + /*target_hidden=*/5120, + /*target_vocab=*/248320, + /*draft_layers=*/5, + /*draft_hidden=*/5120, + /*draft_block_size=*/8, + /*selector_rank=*/256, + /*selector_top_k=*/16, + /*selector_vocab=*/248320, + /*conv_kernel_size=*/2, + /*conv_group_size=*/16, + /*target_file_size=*/15195272800ULL, + /*draft_file_size=*/2045471776ULL, +}; + +bool same_signature(const DFlash2BenefitModelSignature & a, + const DFlash2BenefitModelSignature & b) { + return a.target_layers == b.target_layers && + a.target_hidden == b.target_hidden && + a.target_vocab == b.target_vocab && + a.draft_layers == b.draft_layers && + a.draft_hidden == b.draft_hidden && + a.draft_block_size == b.draft_block_size && + a.selector_rank == b.selector_rank && + a.selector_top_k == b.selector_top_k && + a.selector_vocab == b.selector_vocab && + a.conv_kernel_size == b.conv_kernel_size && + a.conv_group_size == b.conv_group_size && + a.target_file_size == b.target_file_size && + a.draft_file_size == b.draft_file_size; +} + +bool parse_finite_env(const char * name, double minimum, double maximum, + double & value, std::string & error) { + const char * text = std::getenv(name); + if (!text || !*text) return true; + errno = 0; + char * end = nullptr; + const double parsed = std::strtod(text, &end); + if (errno != 0 || end == text || !end || *end != '\0' || + !std::isfinite(parsed) || parsed < minimum || parsed > maximum) { + error = std::string(name) + " must be finite in [" + + std::to_string(minimum) + "," + std::to_string(maximum) + "]"; + return false; + } + value = parsed; + return true; +} + +void set_error(std::string * output, const std::string & value) { + if (output) *output = value; +} + +} // namespace + +std::string DFlash2BenefitModelSignature::str() const { + std::ostringstream out; + out << "target:l" << target_layers << ":h" << target_hidden + << ":v" << target_vocab + << "/draft:l" << draft_layers << ":h" << draft_hidden + << ":b" << draft_block_size + << "/selector:r" << selector_rank << ":k" << selector_top_k + << ":v" << selector_vocab + << "/conv:k" << conv_kernel_size << ":g" << conv_group_size + << "/files:t" << target_file_size << ":d" << draft_file_size; + return out.str(); +} + +DFlash2BenefitConfig DFlash2BenefitProvider::config_from_environment( + std::string & error) { + error.clear(); + DFlash2BenefitConfig config; + if (const char * version = + std::getenv("DFLASH_DFLASH2_BENEFIT_ADAPTER")) { + config.adapter_version = version; + } + if (!parse_finite_env( + "DFLASH_DFLASH2_BENEFIT_LM_WEIGHT", 0.0, 1.0, + config.lm_log_weight, error)) { + return config; + } + if (!parse_finite_env( + "DFLASH_DFLASH2_BENEFIT_HAZARD_SCALE", 0.0, 1.0, + config.hazard_scale, error) || config.hazard_scale <= 0.0) { + if (error.empty()) { + error = "DFLASH_DFLASH2_BENEFIT_HAZARD_SCALE must be in (0,1]"; + } + } + return config; +} + +DFlash2BenefitProvider::DFlash2BenefitProvider( + DFlash2BenefitModelSignature model_signature, DFlash2BenefitConfig config) + : model_signature_(std::move(model_signature)), config_(std::move(config)) { + if (config_.adapter_version != kDFlash2BenefitAdapterVersion) { + error_ = "unsupported DFlash2 benefit adapter version '" + + config_.adapter_version + "'"; + return; + } + if (!std::isfinite(config_.lm_log_weight) || + config_.lm_log_weight < 0.0 || config_.lm_log_weight > 1.0 || + !std::isfinite(config_.hazard_scale) || + config_.hazard_scale <= 0.0 || config_.hazard_scale > 1.0) { + error_ = "invalid DFlash2 benefit coefficients"; + return; + } + if (!same_signature(model_signature_, kSeededQwen38DFlash2)) { + error_ = "unsupported DFlash2 model signature " + + model_signature_.str(); + } +} + +bool DFlash2BenefitProvider::estimate( + const DFlash2SelectorTrace & trace, int max_accept, + DFlash2BenefitEstimate & out, std::string * error) const { + out = {}; + if (!ready()) { + set_error(error, error_); + return false; + } + if (max_accept < 2 || max_accept > model_signature_.draft_block_size) { + set_error(error, "DFlash2 benefit depth is outside the seeded block"); + return false; + } + const size_t required = static_cast(max_accept - 1); + if (trace.depths.size() < required) { + set_error(error, "DFlash2 selector trace is missing required depths"); + return false; + } + + out.conditional_hazards.reserve(required); + double survival = 1.0; + double expected = 1.0; + const double selector_weight = 1.0 - config_.lm_log_weight; + for (size_t depth = 0; depth < required; ++depth) { + const DFlash2DepthSignal & signal = trace.depths[depth]; + if (!std::isfinite(signal.selected_log_prob) || + signal.selected_log_prob > 1e-6f || + !std::isfinite(signal.selector_winner_mass) || + signal.selector_winner_mass <= 0.0f || + signal.selector_winner_mass > 1.0f + 1e-6f) { + set_error(error, "DFlash2 selector trace contains invalid evidence"); + out = {}; + return false; + } + const double log_hazard = + config_.lm_log_weight * signal.selected_log_prob + + selector_weight * + std::log(std::min(1.0, signal.selector_winner_mass)); + const double hazard = std::clamp( + config_.hazard_scale * std::exp(log_hazard), 0.0, 1.0); + if (!std::isfinite(hazard)) { + set_error(error, "DFlash2 selector trace produced a nonfinite hazard"); + out = {}; + return false; + } + out.conditional_hazards.push_back(hazard); + survival *= hazard; + expected += survival; + } + out.expected_yield = std::clamp( + expected, 1.0, static_cast(max_accept)); + if (error) error->clear(); + return true; +} + +bool DFlash2BenefitProvider::publish_once( + const DFlash2SelectorTrace & trace, int max_accept, + double & destination, std::string * error) const { + if (std::isfinite(destination)) { + if (error) error->clear(); + return true; + } + DFlash2BenefitEstimate estimate_result; + if (!estimate(trace, max_accept, estimate_result, error)) return false; + destination = estimate_result.expected_yield; + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash2_benefit.h b/server/src/common/dflash2_benefit.h new file mode 100644 index 000000000..3e899cb81 --- /dev/null +++ b/server/src/common/dflash2_benefit.h @@ -0,0 +1,92 @@ +#pragma once + +#include "dflash2_head.h" + +#include +#include +#include +#include + +namespace dflash::common { + +// Versioned structural signature of the Qwen3.8/DFlash2 pair on which +// the deliberately simple benefit heuristic below was empirically seeded. +// The architecture plus exact artifact sizes reject known-incompatible +// artifacts; they are not a cryptographic identity or a claim of held-out +// calibration. Deployment must preserve the seeded hashes recorded in the +// implementation, and any different target/selector needs its own offline fit. +struct DFlash2BenefitModelSignature { + int target_layers = 0; + int target_hidden = 0; + int target_vocab = 0; + int draft_layers = 0; + int draft_hidden = 0; + int draft_block_size = 0; + int selector_rank = 0; + int selector_top_k = 0; + int selector_vocab = 0; + int conv_kernel_size = 0; + int conv_group_size = 0; + uint64_t target_file_size = 0; + uint64_t draft_file_size = 0; + + std::string str() const; +}; + +inline constexpr const char * kDFlash2BenefitAdapterVersion = + "qwen38-dflash2-selector-benefit-v1"; + +struct DFlash2BenefitConfig { + std::string adapter_version = kDFlash2BenefitAdapterVersion; + + // Conditional acceptance hazard at each depth: + // exp(lm_log_weight * selected_log_prob + // + (1-lm_log_weight) * log(selector_winner_mass)) + // The selector is the better signal on the seed traces, while the + // LM term conservatively lowers a selector winner that has weak model + // probability. hazard_scale may only lower the estimate. + double lm_log_weight = 0.10; + double hazard_scale = 1.0; +}; + +struct DFlash2BenefitEstimate { + double expected_yield = std::numeric_limits::quiet_NaN(); + std::vector conditional_hazards; +}; + +// Stateless request-local adapter. It does not learn from prior requests. +// Request lifetime is owned by the caller via publish_once(destination): the +// first valid trace fills an empty destination, and subsequent traces cannot +// overwrite that request's activation score. +class DFlash2BenefitProvider { +public: + DFlash2BenefitProvider( + DFlash2BenefitModelSignature model_signature, + DFlash2BenefitConfig config = {}); + + static DFlash2BenefitConfig config_from_environment( + std::string & error); + + bool ready() const { return error_.empty(); } + const std::string & error() const { return error_; } + const DFlash2BenefitModelSignature & model_signature() const { + return model_signature_; + } + const DFlash2BenefitConfig & config() const { return config_; } + const char * score_kind() const { return kDFlash2BenefitAdapterVersion; } + + bool estimate(const DFlash2SelectorTrace & trace, int max_accept, + DFlash2BenefitEstimate & out, + std::string * error = nullptr) const; + + bool publish_once(const DFlash2SelectorTrace & trace, int max_accept, + double & destination, + std::string * error = nullptr) const; + +private: + DFlash2BenefitModelSignature model_signature_; + DFlash2BenefitConfig config_; + std::string error_; +}; + +} // namespace dflash::common diff --git a/server/src/common/dflash2_head.cpp b/server/src/common/dflash2_head.cpp index e4d6325ac..e44cda882 100644 --- a/server/src/common/dflash2_head.cpp +++ b/server/src/common/dflash2_head.cpp @@ -1,10 +1,16 @@ #include "dflash2_head.h" +#include "dflash2_selector_validation.h" +#include "ddtree.h" +#include "geometric_draft_topk_cuda.h" #include "ggml-alloc.h" +#include #include #include +#include #include +#include #include namespace dflash::common { @@ -43,6 +49,47 @@ void selector_graph_free(SelectorGraph & g) { g.K = 0; } +DFlash2DepthSignal make_depth_signal( + const float * log_probs, const std::vector & scores, + int K, int selected) { + DFlash2DepthSignal signal; + if (!log_probs || K <= 0 || selected < 0 || selected >= K || + static_cast(scores.size()) != K) { + return signal; + } + signal.selected_log_prob = log_probs[selected]; + signal.lm_top2_margin = K > 1 ? log_probs[0] - log_probs[1] + : std::numeric_limits::infinity(); + float top_k_mass = 0.0f; + for (int k = 0; k < K; ++k) top_k_mass += std::exp(log_probs[k]); + signal.top_k_mass = std::clamp(top_k_mass, 0.0f, 1.0f); + signal.selected_rank = selected; + signal.agrees_with_lm_top1 = selected == 0; + + float runner_up = -INFINITY; + for (int k = 0; k < K; ++k) { + if (k != selected) runner_up = std::max(runner_up, scores[(size_t)k]); + } + signal.selector_margin = K > 1 + ? scores[(size_t)selected] - runner_up + : std::numeric_limits::infinity(); + + const float max_score = *std::max_element(scores.begin(), scores.end()); + float z = 0.0f; + for (float score : scores) z += std::exp(score - max_score); + if (z > 0.0f && std::isfinite(z)) { + signal.selector_winner_mass = + std::exp(scores[(size_t)selected] - max_score) / z; + float entropy = 0.0f; + for (float score : scores) { + const float p = std::exp(score - max_score) / z; + if (p > 0.0f) entropy -= p * std::log(p); + } + signal.selector_entropy = entropy; + } + return signal; +} + } // namespace bool dflash2_select_chain(const DraftWeights & dw, @@ -51,7 +98,8 @@ bool dflash2_select_chain(const DraftWeights & dw, const float * local_hidden, int q_len, int32_t last_tok, - std::vector & draft_tok) { + std::vector & draft_tok, + DFlash2SelectorTrace * trace) { const DraftSelectorWeights & sel = dw.selector; if (!sel.enabled || !sel.hproj || !sel.pred_cb || !sel.succ_cb) return false; if (q_len <= 1 || !local_hidden || !backend) return false; @@ -60,6 +108,21 @@ bool dflash2_select_chain(const DraftWeights & dw, const int K = sel.top_k; const int n_cand = q_len - 1; if (hdim <= 0 || rank <= 0 || K <= 0) return false; + DFlash2SelectorLayout selector_layout; + selector_layout.rank = rank; + selector_layout.top_k = K; + selector_layout.hproj_rank = sel.hproj->ne[1]; + selector_layout.pred_rank = sel.pred_cb->ne[0]; + selector_layout.pred_vocab = sel.pred_cb->ne[1]; + selector_layout.succ_rank = sel.succ_cb->ne[0]; + selector_layout.succ_vocab = sel.succ_cb->ne[1]; + std::string selector_error; + if (!validate_dflash2_selector_layout( + selector_layout, selector_error)) { + std::fprintf(stderr, "dflash2_select_chain: %s\n", + selector_error.c_str()); + return false; + } // 1. Top-k candidates (log-probs) per block position through the target // lm_head. Position 0 of local_hidden is the seed slot; candidates are @@ -134,20 +197,28 @@ bool dflash2_select_chain(const DraftWeights & dw, // 3. Path search: greedy over the candidates, conditioned on the previous pick. draft_tok.assign((size_t)q_len, last_tok); + if (trace) { + trace->depths.clear(); + trace->depths.reserve((size_t)n_cand); + } int prev_row = 0; // row in h_pred: 0 = seed, 1 + i*K + k = candidate k of position i for (int i = 0; i < n_cand; ++i) { const float * pr = h_pred.data() + (size_t)prev_row * rank; const float * hp = h_hproj.data() + (size_t)i * rank; float best = -INFINITY; int best_k = 0; + std::vector scores((size_t)K); for (int k = 0; k < K; ++k) { const float * sc = h_succ.data() + ((size_t)i * K + k) * rank; float dot = 0.0f; for (int r = 0; r < rank; ++r) dot += pr[r] * hp[r] * sc[r]; const float score = cand_lp[(size_t)i * K + k] + dot; + scores[(size_t)k] = score; if (score > best) { best = score; best_k = k; } } draft_tok[(size_t)i + 1] = cand_ids[(size_t)i * K + best_k]; + if (trace) trace->depths.push_back(make_depth_signal( + cand_lp.data() + (size_t)i * K, scores, K, best_k)); prev_row = 1 + i * K + best_k; } return true; diff --git a/server/src/common/dflash2_head.h b/server/src/common/dflash2_head.h index 446a8646c..3b47bc2aa 100644 --- a/server/src/common/dflash2_head.h +++ b/server/src/common/dflash2_head.h @@ -8,6 +8,26 @@ namespace dflash::common { +// Raw selector diagnostics for one proposed depth. These are deliberately +// not called confidence: the DFlash2 selector is trained to rank its top-K +// candidates, not to emit calibrated target-acceptance probabilities. An +// offline, model-specific adapter may later map these values to survival +// probabilities for the adaptive gate. +struct DFlash2DepthSignal { + float selected_log_prob = 0.0f; + float lm_top2_margin = 0.0f; + float top_k_mass = 0.0f; + int selected_rank = 0; + bool agrees_with_lm_top1 = false; + float selector_margin = 0.0f; + float selector_winner_mass = 0.0f; + float selector_entropy = 0.0f; +}; + +struct DFlash2SelectorTrace { + std::vector depths; +}; + // DFlash 2 candidate selector for greedy chain drafting. // // For every drafted block position the target lm_head logits are reduced to @@ -24,6 +44,22 @@ bool dflash2_select_chain(const DraftWeights & dw, const float * local_hidden, int q_len, int32_t last_tok, - std::vector & draft_tok); + std::vector & draft_tok, + DFlash2SelectorTrace * trace = nullptr); + +// Same selector, batched over host-resident drafter hidden blocks and using a +// local target lm_head tensor. The expensive lm_head projection covers every +// (lane, depth) in one graph, GPU top-K is invoked once, and selector +// projections/readback are shared across the cohort. This is the concurrent +// paged-engine entry point; no non-paged DFlashTarget adapter is required. +bool dflash2_select_chains_batched( + const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const std::vector & hidden_by_lane, + int q_len, + const std::vector & last_tokens, + std::vector> & draft_tokens, + std::vector * traces = nullptr); } // namespace dflash::common diff --git a/server/src/common/dflash2_selector_validation.h b/server/src/common/dflash2_selector_validation.h new file mode 100644 index 000000000..6bcd50403 --- /dev/null +++ b/server/src/common/dflash2_selector_validation.h @@ -0,0 +1,93 @@ +#pragma once + +#include "geometric_draft_topk_cuda.h" + +#include +#include + +namespace dflash::common { + +// Host-only description of the selector tensors. Keeping validation in terms +// of dimensions makes it usable both while GGUF tensor descriptors are being +// loaded and when a concrete target lm_head is attached to the batched path. +struct DFlash2SelectorLayout { + int rank = 0; + int top_k = 0; + int64_t hproj_rank = 0; + int64_t pred_rank = 0; + int64_t pred_vocab = 0; + int64_t succ_rank = 0; + int64_t succ_vocab = 0; + // Zero means that source is unavailable at this validation point. A + // partial target shard, for example, can declare n_vocab without owning + // the final output tensor; the concrete lm_head is checked again at use. + int64_t target_output_vocab = 0; + int64_t target_declared_vocab = 0; +}; + +inline bool validate_dflash2_selector_layout( + const DFlash2SelectorLayout & layout, std::string & error) { + error.clear(); + if (layout.rank <= 0) { + error = "DFlash 2 selector rank must be positive (got " + + std::to_string(layout.rank) + ")"; + return false; + } + if (!geometric_draft_topk_cuda_supports_k(layout.top_k)) { + error = "DFlash 2 selector top_k=" + std::to_string(layout.top_k) + + " is unsupported; expected one of 1..8, 12, or 16"; + return false; + } + if (layout.hproj_rank != layout.rank || + layout.pred_rank != layout.rank || + layout.succ_rank != layout.rank) { + error = "DFlash 2 selector rank mismatch: metadata=" + + std::to_string(layout.rank) + " hproj=" + + std::to_string(layout.hproj_rank) + " pred_cb=" + + std::to_string(layout.pred_rank) + " succ_cb=" + + std::to_string(layout.succ_rank); + return false; + } + if (layout.pred_vocab <= 0 || layout.succ_vocab <= 0) { + error = "DFlash 2 selector codebook vocab must be positive: pred_cb=" + + std::to_string(layout.pred_vocab) + " succ_cb=" + + std::to_string(layout.succ_vocab); + return false; + } + if (layout.pred_vocab != layout.succ_vocab) { + error = "DFlash 2 selector codebook vocab mismatch: pred_cb=" + + std::to_string(layout.pred_vocab) + " succ_cb=" + + std::to_string(layout.succ_vocab); + return false; + } + if (layout.top_k > layout.pred_vocab) { + error = "DFlash 2 selector top_k=" + std::to_string(layout.top_k) + + " exceeds codebook vocab=" + std::to_string(layout.pred_vocab); + return false; + } + if (layout.target_output_vocab > 0 && + layout.target_declared_vocab > 0 && + layout.target_output_vocab != layout.target_declared_vocab) { + error = "DFlash 2 target vocab mismatch: output/lm_head=" + + std::to_string(layout.target_output_vocab) + " target.n_vocab=" + + std::to_string(layout.target_declared_vocab); + return false; + } + if (layout.target_output_vocab > 0 && + layout.pred_vocab != layout.target_output_vocab) { + error = "DFlash 2 selector vocab mismatch: codebook=" + + std::to_string(layout.pred_vocab) + " target output/lm_head=" + + std::to_string(layout.target_output_vocab); + return false; + } + if (layout.target_declared_vocab > 0 && + layout.pred_vocab != layout.target_declared_vocab) { + error = "DFlash 2 selector vocab mismatch: codebook=" + + std::to_string(layout.pred_vocab) + " target.n_vocab=" + + std::to_string(layout.target_declared_vocab); + return false; + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash_draft_kv.cpp b/server/src/common/dflash_draft_kv.cpp index 9077da6a8..b4366de84 100644 --- a/server/src/common/dflash_draft_kv.cpp +++ b/server/src/common/dflash_draft_kv.cpp @@ -333,6 +333,7 @@ void draft_kv_batch_free(DraftKvBatchGraph & batch) { } batch.gf = nullptr; batch.seed_tokens = nullptr; + batch.hidden_by_lane.clear(); batch.token_depths.clear(); batch.confidence_depths.clear(); batch.lane_states.clear(); @@ -340,6 +341,7 @@ void draft_kv_batch_free(DraftKvBatchGraph & batch) { batch.n_lanes = 0; batch.q_len = 0; batch.has_confidence = false; + batch.uses_dflash2 = false; batch.built_for = nullptr; batch.built_lm_head = nullptr; } @@ -413,12 +415,18 @@ static bool draft_kv_batch_build( prenorm.push_back(outputs.hidden_prenorm); } + const bool uses_dflash2 = dw.selector.enabled; DSparkBatchedChainOutputs chain; - if (!build_dspark_markov_batched_chain( - batch.g_ctx, batch.gf, dw, lm_head, hidden, prenorm, - batch.seed_tokens, dw.block_size, true, chain) || - chain.n_lanes != n_lanes || - static_cast(chain.tokens.size()) != dw.block_size - 1) { + if (uses_dflash2) { + for (ggml_tensor * lane_hidden : hidden) { + ggml_set_output(lane_hidden); + ggml_build_forward_expand(batch.gf, lane_hidden); + } + } else if (!build_dspark_markov_batched_chain( + batch.g_ctx, batch.gf, dw, lm_head, hidden, prenorm, + batch.seed_tokens, dw.block_size, true, chain) || + chain.n_lanes != n_lanes || + static_cast(chain.tokens.size()) != dw.block_size - 1) { draft_kv_batch_free(batch); return false; } @@ -435,17 +443,22 @@ static bool draft_kv_batch_build( batch.n_lanes = n_lanes; batch.q_len = dw.block_size; + batch.uses_dflash2 = uses_dflash2; batch.has_confidence = - !chain.confidence.empty() && chain.confidence[0] != nullptr; + !uses_dflash2 && !chain.confidence.empty() && + chain.confidence[0] != nullptr; batch.built_for = &dw; batch.built_lm_head = lm_head; batch.lane_states = lane_states; + batch.hidden_by_lane = uses_dflash2 ? hidden + : std::vector{}; batch.token_depths = std::move(chain.tokens); batch.confidence_depths = std::move(chain.confidence); std::fprintf(stderr, "[draft-kv-batch] packed graph ready lanes=%d q_len=%d " - "confidence=%s\n", + "head=%s confidence=%s\n", n_lanes, dw.block_size, + uses_dflash2 ? "dflash2" : "dspark", batch.has_confidence ? "on" : "off"); return true; } @@ -458,9 +471,11 @@ bool draft_kv_batch_compute( const std::vector & lane_states, const std::vector & seed_tokens, std::vector> & draft_tokens, - std::vector> & confidences) { + std::vector> & confidences, + std::vector * selector_traces) { draft_tokens.clear(); confidences.clear(); + if (selector_traces) selector_traces->clear(); if (lane_states.empty() || seed_tokens.size() != lane_states.size()) { return false; @@ -475,9 +490,11 @@ bool draft_kv_batch_compute( return false; } - ggml_backend_tensor_set( - batch.seed_tokens, seed_tokens.data(), 0, - sizeof(int32_t) * seed_tokens.size()); + if (!batch.uses_dflash2) { + ggml_backend_tensor_set( + batch.seed_tokens, seed_tokens.data(), 0, + sizeof(int32_t) * seed_tokens.size()); + } if (ggml_backend_graph_compute(backend, batch.gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, @@ -486,6 +503,33 @@ bool draft_kv_batch_compute( return false; } + if (batch.uses_dflash2) { + if (static_cast(batch.hidden_by_lane.size()) != + batch.n_lanes) { + return false; + } + const int hidden = dw.n_embd; + std::vector> hidden_host( + (size_t) batch.n_lanes, + std::vector( + (size_t) hidden * (size_t) batch.q_len)); + for (int lane = 0; lane < batch.n_lanes; ++lane) { + ggml_backend_tensor_get_async( + backend, batch.hidden_by_lane[(size_t) lane], + hidden_host[(size_t) lane].data(), 0, + sizeof(float) * hidden_host[(size_t) lane].size()); + } + ggml_backend_synchronize(backend); + std::vector hidden_ptrs((size_t) batch.n_lanes); + for (int lane = 0; lane < batch.n_lanes; ++lane) { + hidden_ptrs[(size_t) lane] = + hidden_host[(size_t) lane].data(); + } + confidences.assign((size_t) batch.n_lanes, {}); + return dflash2_select_chains_batched( + dw, backend, lm_head, hidden_ptrs, batch.q_len, + seed_tokens, draft_tokens, selector_traces); + } const int depths = batch.q_len - 1; std::vector depth_tokens( (size_t)depths * batch.n_lanes); diff --git a/server/src/common/dflash_draft_kv.h b/server/src/common/dflash_draft_kv.h index 04842b6e6..58270c35c 100644 --- a/server/src/common/dflash_draft_kv.h +++ b/server/src/common/dflash_draft_kv.h @@ -19,6 +19,7 @@ #pragma once #include "dflash_feature_ring.h" +#include "dflash2_head.h" #include "draft/draft_graph.h" #include "internal.h" // DraftWeights @@ -108,6 +109,7 @@ struct DraftKvBatchGraph { int n_lanes = 0; int q_len = 0; bool has_confidence = false; + bool uses_dflash2 = false; const void * built_for = nullptr; ggml_tensor * built_lm_head = nullptr; std::vector lane_states; @@ -117,6 +119,7 @@ struct DraftKvBatchGraph { ggml_cgraph * gf = nullptr; ggml_gallocr_t galloc = nullptr; ggml_tensor * seed_tokens = nullptr; + std::vector hidden_by_lane; std::vector token_depths; std::vector confidence_depths; }; @@ -134,6 +137,7 @@ bool draft_kv_batch_compute( const std::vector & lane_states, const std::vector & seed_tokens, std::vector> & draft_tokens, - std::vector> & confidences); + std::vector> & confidences, + std::vector * selector_traces = nullptr); } // namespace dflash::common diff --git a/server/src/common/geometric_draft_topk_cuda.cu b/server/src/common/geometric_draft_topk_cuda.cu index ba287656d..45a9477f2 100644 --- a/server/src/common/geometric_draft_topk_cuda.cu +++ b/server/src/common/geometric_draft_topk_cuda.cu @@ -13,7 +13,7 @@ namespace dflash::common { namespace { -constexpr int kMaxK = 16; // ddtree_K is 8, the DFlash 2 selector uses 16; K>kMaxK → CPU fallback +constexpr int kMaxK = 16; // largest instantiated K; supported set is declared in the header constexpr int kBlock = 256; // threads per block (power of two for the reduction) constexpr int kMaxSplit = 128; // max vocab splits per position (combine-block cap) @@ -331,7 +331,13 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, float * out_log_probs, int32_t * out_token_ids, float temperature) { - if (!d_logits || n_positions <= 0 || vocab <= 0 || K <= 0 || K > kMaxK) return false; + // Reject before touching CUDA or scratch. In particular, K values in the + // holes between instantiated templates (9-11 and 13-15) must fall back to + // the CPU path instead of copying stale data from a previous invocation. + if (!d_logits || !out_log_probs || !out_token_ids || n_positions <= 0 || + vocab <= 0 || K > vocab || !geometric_draft_topk_cuda_supports_k(K)) { + return false; + } cudaPointerAttributes attr{}; if (cudaPointerGetAttributes(&attr, d_logits) != cudaSuccess) { @@ -362,9 +368,11 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, // the tensor base aligned and a vocab stride that is a multiple of 4. const bool use_vec = (vocab % 4 == 0) && (reinterpret_cast(lp_in) % 16 == 0); + bool dispatched = false; // K (and the vectorization flag) are compile-time template parameters // so the per-thread/per-partial top-K stays register-resident; dispatch - // the runtime K to its instantiation. K>kMaxK is already rejected above. + // the runtime K to its instantiation. Unsupported K is rejected before + // scratch allocation above, so the default is unreachable hardening. #define DFLASH_TOPK_LAUNCH(KV, VEC) \ geometric_draft_topk_partial<<>>( \ lp_in, vocab, inv_t, split, \ @@ -374,6 +382,7 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, split, g_scratch.d_lp, g_scratch.d_ids); #define DFLASH_TOPK_CASE(KV) \ case KV: \ + dispatched = true; \ if (use_vec) { DFLASH_TOPK_LAUNCH(KV, true) } \ else { DFLASH_TOPK_LAUNCH(KV, false) } \ break; @@ -387,7 +396,8 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, #undef DFLASH_TOPK_LAUNCH if (kProfile) cudaEventRecord(e_k1); - if (cudaGetLastError() == cudaSuccess && cudaDeviceSynchronize() == cudaSuccess) { + if (dispatched && cudaGetLastError() == cudaSuccess && + cudaDeviceSynchronize() == cudaSuccess) { const cudaError_t e1 = cudaMemcpy(out_log_probs, g_scratch.d_lp, n * sizeof(float), cudaMemcpyDeviceToHost); const cudaError_t e2 = cudaMemcpy(out_token_ids, g_scratch.d_ids, diff --git a/server/src/common/geometric_draft_topk_cuda.h b/server/src/common/geometric_draft_topk_cuda.h index b926dbefc..0edd48aa5 100644 --- a/server/src/common/geometric_draft_topk_cuda.h +++ b/server/src/common/geometric_draft_topk_cuda.h @@ -25,6 +25,15 @@ namespace dflash::common { +// Keep the public capability predicate in lockstep with the template +// instantiations dispatched by geometric_draft_topk_cuda.cu. Callers use this +// to choose the CPU fallback without entering CUDA, and loader validation uses +// it to reject selector metadata that cannot be executed consistently across +// the host and device paths. +inline constexpr bool geometric_draft_topk_cuda_supports_k(int K) noexcept { + return (K >= 1 && K <= 8) || K == 12 || K == 16; +} + // d_logits: device pointer to row-major [n_positions][vocab] f32 logits (the // position stride is `vocab` floats — pass an offset pointer to skip // leading positions). out_* are HOST buffers of size n_positions*K. diff --git a/server/src/common/speculation_policy.h b/server/src/common/speculation_policy.h index aa6b25c1d..1f56d8ffe 100644 --- a/server/src/common/speculation_policy.h +++ b/server/src/common/speculation_policy.h @@ -14,8 +14,10 @@ enum class SpeculationPolicy { }; // Runtime capabilities for the concurrent decode path. Forced speculation -// only needs an executable draft/verify chain; adaptive additionally needs -// the activation gate and its cost profile. AR is always supported. +// needs an executable draft/verify chain. Adaptive means the server can honor +// the one-shot activation contract; a configured chain may satisfy that by +// recording a request-local sticky-AR fallback when scoring or profiling is +// unavailable. AR is always supported. struct ConcurrentDecodeCapabilities { bool forced_speculation = false; bool adaptive = false; diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index 58c203195..43304f9f3 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -25,6 +25,7 @@ // blk..ffn_down.weight [hidden, intermediate] Q8_0 / F16 #include "internal.h" +#include "common/dflash2_selector_validation.h" #include "common/derived_scalars.h" #include "common/gguf_mmap.h" #include "common/gguf_bounds.h" @@ -542,13 +543,31 @@ bool load_draft_gguf(const std::string & path, out.selector.top_k = (int)read_u32("dflash.dflash2.selector_top_k", 16); char shape_err[192]; const int64_t R = out.selector.rank; - if (!check_shape_2d(out.selector.hproj, out.n_embd, R, "selector.hproj", shape_err, sizeof(shape_err)) || - !check_shape_2d(out.selector.pred_cb, R, out.selector.pred_cb->ne[1], "selector.pred_cb", shape_err, sizeof(shape_err)) || - !check_shape_2d(out.selector.succ_cb, R, out.selector.pred_cb->ne[1], "selector.succ_cb", shape_err, sizeof(shape_err))) { + if (!check_shape_2d(out.selector.hproj, out.n_embd, R, + "selector.hproj", shape_err, sizeof(shape_err))) { set_last_error(shape_err); ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); return false; } + DFlash2SelectorLayout selector_layout; + selector_layout.rank = out.selector.rank; + selector_layout.top_k = out.selector.top_k; + selector_layout.hproj_rank = out.selector.hproj->ne[1]; + selector_layout.pred_rank = out.selector.pred_cb->ne[0]; + selector_layout.pred_vocab = out.selector.pred_cb->ne[1]; + selector_layout.succ_rank = out.selector.succ_cb->ne[0]; + selector_layout.succ_vocab = out.selector.succ_cb->ne[1]; + selector_layout.target_output_vocab = + target && target->output ? target->output->ne[1] : 0; + selector_layout.target_declared_vocab = + target ? target->n_vocab : 0; + std::string selector_error; + if (!validate_dflash2_selector_layout( + selector_layout, selector_error)) { + set_last_error("draft GGUF: " + selector_error); + ggml_free(meta_ctx); out.ctx = nullptr; gguf_free(gctx); + return false; + } out.selector.enabled = true; std::fprintf(stderr, "[draft GGUF] DFlash 2 selector enabled: rank=%d top_k=%d vocab=%lld\n", out.selector.rank, out.selector.top_k, (long long)out.selector.pred_cb->ne[1]); diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index ef85225a4..6afce5052 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include #include @@ -62,12 +63,12 @@ double initial_prediction_realized_tokens( void log_spec_gate_plan(const SpecPlan & plan, double fixed_yield_scale, double initial_realized_tokens, double measured_us) { - int confidence = 0; + int current = 0; int initial = 0; int unavailable = plan.unavailable_count; for (const SpecPlanScore & score : plan.ordered) { switch (score.source) { - case SpecScoreSource::Confidence: ++confidence; break; + case SpecScoreSource::Confidence: ++current; break; case SpecScoreSource::InitialConfidence: ++initial; break; case SpecScoreSource::Unavailable: ++unavailable; break; } @@ -77,11 +78,12 @@ void log_spec_gate_plan(const SpecPlan & plan, double fixed_yield_scale, plan.concurrency, plan.admitted_count); for (size_t i = 0; i < plan.ordered.size(); ++i) { const SpecPlanScore & score = plan.ordered[i]; - std::fprintf(stderr, "%s%llu:%.3f/%s%s", + std::fprintf(stderr, "%s%llu:%.3f/%s/%s%s", i == 0 ? "" : ",", (unsigned long long)score.request_id, score.expected_yield, spec_score_source_name(score.source), + spec_score_kind_name(score.score_kind), score.admitted ? "*" : ""); } std::fprintf(stderr, "] decisions=["); @@ -93,9 +95,9 @@ void log_spec_gate_plan(const SpecPlan & plan, double fixed_yield_scale, spec_decision_name(score.decision)); } std::fprintf(stderr, - "] sources=confidence:%d,initial:%d,unavailable:%d " + "] sources=current:%d,initial:%d,unavailable:%d " "fixed_yield_scale=%.3f initial_tokens=%.3f/", - confidence, initial, unavailable, fixed_yield_scale, + current, initial, unavailable, fixed_yield_scale, plan.initial_predicted_tokens); if (std::isfinite(initial_realized_tokens)) { std::fprintf(stderr, "%.3f", initial_realized_tokens); @@ -119,25 +121,81 @@ void log_spec_activations(const SpecPlan & plan, for (const SpecPlanScore & score : plan.ordered) { if (!score.newly_decided) continue; const double initial = gate.initial_confidence(score.request_id); + const SpecScoreKind kind = gate.initial_score_kind(score.request_id); + const char * decision_reason = score.execution_unsupported + ? "execution_unsupported" + : score.decision == SpecDecision::Speculation + ? "selected_by_joint_goodput" + : "ar_counterfactual_won"; std::fprintf(stderr, - "[spec-activation] {\"request_id\":%llu,\"slot\":%d," - "\"initial_confidence\":%.6f,\"expected_yield\":%.6f," + "[spec-activation] {\"request_id\":%llu,\"slot\":%d,", + (unsigned long long)score.request_id, score.slot); + if (kind == SpecScoreKind::DSparkConfidence) { + // Legacy DSpark field retained for harness compatibility. + std::fprintf(stderr, + "\"initial_confidence\":%.6f,", initial); + } else { + // Selector evidence is adapted to request benefit; it is not a + // trained confidence value. + std::fprintf(stderr, "\"initial_confidence\":null,"); + } + std::fprintf(stderr, + "\"activation_score\":%.6f," + "\"request_benefit\":%.6f," + "\"score_kind\":\"%s\",\"expected_yield\":%.6f," "\"evaluation\":\"scored\",\"fallback_reason\":null," + "\"decision_reason\":\"%s\"," "\"decision\":\"%s\"}\n", - (unsigned long long)score.request_id, score.slot, - initial, score.expected_yield, + initial, initial, spec_score_kind_name(kind), + score.expected_yield, decision_reason, spec_decision_name(score.decision)); } } -void log_spec_evaluation_fallback(uint64_t request_id, int slot) { +void log_spec_evaluation_fallback(uint64_t request_id, int slot, + SpecScoreKind kind, + const char * reason) { std::fprintf(stderr, "[spec-activation] {\"request_id\":%llu,\"slot\":%d," - "\"initial_confidence\":null,\"expected_yield\":null," - "\"evaluation\":\"failed\"," - "\"fallback_reason\":\"confidence_evaluation_failed\"," + "\"initial_confidence\":null,\"activation_score\":null," + "\"request_benefit\":null,\"score_kind\":\"%s\"," + "\"expected_yield\":null,\"evaluation\":\"failed\"," + "\"fallback_reason\":\"%s\"," + "\"decision_reason\":\"evaluation_failed\"," "\"decision\":\"ar\"}\n", - (unsigned long long)request_id, slot); + (unsigned long long)request_id, slot, spec_score_kind_name(kind), + reason ? reason : "activation_evaluation_failed"); +} + +uint64_t file_size_or_zero(const char * path) { + if (!path || !*path) return 0; + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) return 0; + const std::streamoff size = file.tellg(); + return size > 0 ? static_cast(size) : 0; +} + +int configured_chain_verify_depth(int maximum) { + const char * value = std::getenv("DFLASH_SPEC_CHAIN_DEPTH"); + if (!value || !*value) { + return resolve_chain_verify_depth(0, maximum); + } + + char * end = nullptr; + const long parsed = std::strtol(value, &end, 10); + const bool integer = end != value && end && *end == '\0' && + parsed >= std::numeric_limits::min() && + parsed <= std::numeric_limits::max(); + const int resolved = integer + ? resolve_chain_verify_depth(static_cast(parsed), maximum) + : 0; + if (resolved != 0) return resolved; + + std::fprintf(stderr, + "[parallel-chain] ignoring invalid DFLASH_SPEC_CHAIN_DEPTH=%s; " + "expected root-inclusive depth 2..%d, using %d\n", + value, maximum, maximum); + return resolve_chain_verify_depth(0, maximum); } } // namespace @@ -158,14 +216,69 @@ Qwen35SeqEngine::Qwen35SeqEngine( slots_(pool, max_ctx, std::max(1, tree_width), backend.paged_kv_residency_.get()), scratch_row_(scratch_row), tree_width_(tree_width), + chain_verify_depth_(configured_chain_verify_depth(tree_width)), tree_scratch_base_(tree_scratch_base), tree_scratch_stride_(tree_scratch_stride), spec_mode_(spec_mode) { + if (spec_mode_ == SpecMode::chain && + chain_verify_depth_ >= 2 && chain_verify_depth_ != tree_width_) { + std::fprintf(stderr, + "[parallel-chain] verify_depth=%d draft_width=%d " + "(DFLASH_SPEC_CHAIN_DEPTH)\n", + chain_verify_depth_, tree_width_); + } const int n_slots = slots_.slot_count(); slot_draft_kv_.resize((size_t)n_slots); prepared_chain_drafts_.resize((size_t)n_slots); last_survival_score_.assign( (size_t)n_slots, std::numeric_limits::quiet_NaN()); + adaptive_fallback_ar_.assign((size_t)n_slots, 0); + if (spec_mode_ == SpecMode::chain && b_.dw_.selector.enabled) { + DFlash2BenefitModelSignature signature; + signature.target_layers = b_.w_.n_layer; + signature.target_hidden = b_.w_.n_embd; + signature.target_vocab = b_.w_.n_vocab; + signature.draft_layers = b_.dw_.n_layer; + signature.draft_hidden = b_.dw_.n_embd; + signature.draft_block_size = b_.dw_.block_size; + signature.selector_rank = b_.dw_.selector.rank; + signature.selector_top_k = b_.dw_.selector.top_k; + signature.selector_vocab = b_.dw_.selector.pred_cb + ? static_cast(b_.dw_.selector.pred_cb->ne[1]) : 0; + signature.conv_kernel_size = b_.dw_.conv_kernel_size; + signature.conv_group_size = b_.dw_.conv_group_size; + signature.target_file_size = file_size_or_zero(b_.cfg_.target_path); + signature.draft_file_size = file_size_or_zero(b_.cfg_.draft_path); + + std::string config_error; + DFlash2BenefitConfig config = + DFlash2BenefitProvider::config_from_environment(config_error); + if (!config_error.empty()) { + adaptive_fallback_reason_ = "benefit_adapter_invalid_config"; + std::fprintf(stderr, + "[parallel-chain] DFlash2 benefit adapter disabled: %s; " + "adaptive requests will use sticky AR\n", + config_error.c_str()); + } else { + dflash2_benefit_provider_ = + std::make_unique(signature, config); + if (!dflash2_benefit_provider_->ready()) { + adaptive_fallback_reason_ = "benefit_adapter_unavailable"; + std::fprintf(stderr, + "[parallel-chain] DFlash2 benefit adapter disabled: %s; " + "adaptive requests will use sticky AR\n", + dflash2_benefit_provider_->error().c_str()); + } else { + std::fprintf(stderr, + "[parallel-chain] DFlash2 request-benefit adapter=%s " + "lm_weight=%.3f hazard_scale=%.3f signature=%s\n", + dflash2_benefit_provider_->score_kind(), + config.lm_log_weight, config.hazard_scale, + signature.str().c_str()); + } + } + } + // 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 @@ -238,9 +351,12 @@ bool Qwen35SeqEngine::step_timing_enabled() { } bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { + adaptive_fallback_reason_ = "cost_profile_unavailable"; speculation_gate_.reset(); - if (spec_mode_ != SpecMode::dspark_chain || !capture_features_ || - tree_width_ <= 1 || tree_width_ > 16 || slots_.residency_active()) { + if (spec_mode_ != SpecMode::chain || !capture_features_ || + !activation_scoring_available() || tree_width_ <= 1 || tree_width_ > 16 || + resolve_chain_verify_depth(chain_verify_depth_, tree_width_) == 0 || + slots_.residency_active()) { std::fprintf(stderr, "[spec-profile] disabled: chain/features unavailable or " "concurrent KVFlash residency active; adaptive capability " @@ -250,6 +366,7 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { const int n_slots = slots_.slot_count(); const int T = tree_width_; + const int V = chain_verify_depth_for_round(); const int hidden = b_.w_.n_embd; const int n_head_kv = b_.w_.n_head_kv; const int max_profile_ctx = std::min( @@ -323,7 +440,7 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { ggml_backend_synchronize(b_.target_backend_); const SpecProfileGrid grid = build_spec_profile_grid( - n_slots, T, T, [](int lanes) { + n_slots, V, V, [](int lanes) { return chain_decode_bucket_width(lanes); }); @@ -332,16 +449,16 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { if (!profile_error.empty()) return std::numeric_limits::infinity(); if (prepared_tree_rows != total_rows) { - if (total_rows <= 0 || total_rows % T != 0) { + if (total_rows <= 0 || total_rows % V != 0) { profile_error = "invalid tree profiling shape"; return std::numeric_limits::infinity(); } - const int bucket = total_rows / T; + const int bucket = total_rows / V; const int live = std::min(bucket, n_slots); StepGraph & sg = b_.sg_; if (!build_target_step_paged_tree( sg, b_.w_, b_.cache_, b_.target_backend_, - T, bucket, ctx_tokens, + V, bucket, ctx_tokens, tree_scratch_base_, tree_scratch_stride_, b_.cfg_.kq_stride_pad)) { profile_error = "tree profiling graph build failed"; @@ -365,11 +482,11 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { } for (int lane = 0; lane < live; ++lane) { const int slot = synthetic_slots[(size_t)lane]; - sizes[(size_t)lane] = T; + sizes[(size_t)lane] = V; active[(size_t)lane] = slot; state[(size_t)lane] = slot; - for (int node = 0; node < T; ++node) { - const int row = lane * T + node; + for (int node = 0; node < V; ++node) { + const int row = lane * V + node; parents[(size_t)row] = node == 0 ? -1 : node - 1; queries[(size_t)row] = slot; for (int axis = 0; axis < 3; ++axis) { @@ -423,14 +540,14 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { if (!profile_error.empty()) return std::numeric_limits::infinity(); if (prepared_step_rows != total_rows) { - if (total_rows <= 0 || total_rows > n_slots * T) { + if (total_rows <= 0 || total_rows > n_slots * V) { profile_error = "invalid durable-step profiling shape"; return std::numeric_limits::infinity(); } std::vector segments; int offset = 0; for (int lane = 0; offset < total_rows; ++lane) { - const int length = std::min(T, total_rows - offset); + const int length = std::min(V, total_rows - offset); segments.push_back({ offset, length, synthetic_slots[(size_t)lane]}); offset += length; @@ -440,7 +557,7 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { sg, b_.w_, b_.cache_, b_.target_backend_, 0, total_rows, false, true, false, 0, 0, b_.cfg_.kq_stride_pad, false, false, false, true, - 1, 0, ctx_tokens + T, + 1, 0, ctx_tokens + V, total_rows, segments.data(), (int)segments.size(), (int)segments.size(), false) || !sg.kv_write_rows || !sg.target_feat_rows || @@ -616,14 +733,28 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { prenorm_hidden.data(), 0, sizeof(float) * prenorm_hidden.size()); ggml_backend_synchronize(b_.draft_backend_); - std::vector draft_tokens; - std::vector confidence; - if (!dspark_markov_correct_greedy_chain_fused( + bool chain_ok = false; + if (b_.dw_.selector.enabled) { + std::vector hidden_ptrs{ + local_hidden.data()}; + std::vector one_seed{profile_token}; + std::vector> draft_tokens; + chain_ok = dflash2_select_chains_batched( + b_.dw_, b_.draft_backend_, b_.w_.output, + hidden_ptrs, T, one_seed, draft_tokens) && + draft_tokens.size() == 1 && + (int)draft_tokens[0].size() == T; + } else { + std::vector draft_tokens; + std::vector confidence; + chain_ok = dspark_markov_correct_greedy_chain_fused( b_.dw_, b_.draft_backend_, b_.w_.output, local_hidden.data(), T, profile_token, draft_tokens, &confidence, - prenorm_hidden.data()) || - (int)draft_tokens.size() != T) { + prenorm_hidden.data()) && + (int)draft_tokens.size() == T; + } + if (!chain_ok) { profile_error = "draft profiling chain failed"; return std::numeric_limits::infinity(); } @@ -650,12 +781,12 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { SpecCostTables tables{ std::move(tree.table), std::move(step.table), std::move(draft.table)}; SpecStepGeometry geometry; - geometry.tree_width = T; + geometry.tree_width = V; geometry.bucket = [](int lanes) { return chain_decode_bucket_width(lanes); }; speculation_gate_ = std::make_unique( - tables, geometry, T, + tables, geometry, V, [](const char * table, int requested, int profiled) { std::fprintf(stderr, "[spec-gate] %s_cost index %d outside profile; clamped to %d\n", @@ -681,6 +812,7 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { print_table("tree_cost", tables.tree_cost); print_table("step_cost", tables.step_cost); print_table("draft_cost", tables.draft_cost); + adaptive_fallback_reason_.clear(); return true; } @@ -717,8 +849,33 @@ bool Qwen35SeqEngine::batched_drafting_enabled() const { return !value || std::atoi(value) != 0; } -bool Qwen35SeqEngine::confidence_scoring_enabled() const { - const char * value = std::getenv("DFLASH_SPEC_CONFIDENCE"); +bool Qwen35SeqEngine::activation_scoring_available() const { + if (spec_mode_ != SpecMode::chain || !capture_features_) return false; + if (b_.dw_.selector.enabled) { + return dflash2_benefit_provider_ && + dflash2_benefit_provider_->ready(); + } + const int hidden = b_.dw_.n_embd; + return b_.dw_.dspark.enabled && b_.dw_.dspark.confidence_w && + b_.dw_.dspark.confidence_b && + (b_.dw_.dspark.confidence_dim == hidden || + b_.dw_.dspark.confidence_dim == + hidden + b_.dw_.dspark.markov_rank); +} + +SpecScoreKind Qwen35SeqEngine::chain_activation_score_kind() const { + if (b_.dw_.selector.enabled && dflash2_benefit_provider_ && + dflash2_benefit_provider_->ready()) { + return SpecScoreKind::DFlash2SelectorBenefitV1; + } + return activation_scoring_available() + ? SpecScoreKind::DSparkConfidence + : SpecScoreKind::Unspecified; +} + +bool Qwen35SeqEngine::activation_scoring_enabled() const { + const char * value = std::getenv("DFLASH_SPEC_ACTIVATION_SCORE"); + if (!value) value = std::getenv("DFLASH_SPEC_CONFIDENCE"); return !value || std::atoi(value) != 0; } @@ -779,7 +936,7 @@ bool Qwen35SeqEngine::prepare_chain_drafts( for (size_t i = 0; i < inputs.size(); ++i) { if (!selected[i]) continue; const StepInput & in = inputs[i]; - if (!chain_confidence_input_scoreable(in)) return false; + if (!chain_proposal_input_capable(in)) return false; DraftKvState * state = ensure_slot_draft_kv(in.slot); DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); if (!state || !mirror || @@ -804,6 +961,7 @@ bool Qwen35SeqEngine::prepare_chain_drafts( std::vector> drafts; std::vector> confidences; + std::vector selector_traces; bool used_batch = false; const bool try_batched = !force_serial && batched_drafting_enabled(); if (try_batched) { @@ -856,7 +1014,8 @@ bool Qwen35SeqEngine::prepare_chain_drafts( used_batch = draft_kv_batch_compute( batch_draft_graph_, b_.dw_, b_.draft_backend_, b_.w_.output, - batch_states, seeds, drafts, confidences); + batch_states, seeds, drafts, confidences, + &selector_traces); } } if (!used_batch) { @@ -898,31 +1057,63 @@ bool Qwen35SeqEngine::prepare_chain_drafts( if (!used_batch) { drafts.resize(lanes.size()); confidences.resize(lanes.size()); - std::vector local_hidden((size_t)hidden * T); - std::vector prenorm_hidden((size_t)hidden * T); - for (size_t lane = 0; lane < lanes.size(); ++lane) { - DraftKvState * state = lanes[lane].state; - if (ggml_backend_graph_compute( - b_.draft_backend_, state->gf) != - GGML_STATUS_SUCCESS) { - return false; + if (b_.dw_.selector.enabled) { + selector_traces.resize(lanes.size()); + std::vector> hidden_blocks( + lanes.size(), + std::vector((size_t) hidden * (size_t) T)); + std::vector seeds; + seeds.reserve(lanes.size()); + for (size_t lane = 0; lane < lanes.size(); ++lane) { + DraftKvState * state = lanes[lane].state; + if (ggml_backend_graph_compute( + b_.draft_backend_, state->gf) != + GGML_STATUS_SUCCESS) { + return false; + } + ggml_backend_tensor_get_async( + b_.draft_backend_, state->hidden_states, + hidden_blocks[lane].data(), 0, + sizeof(float) * hidden_blocks[lane].size()); + seeds.push_back(lanes[lane].seed); } - ggml_backend_tensor_get_async( - b_.draft_backend_, state->hidden_states, - local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); - ggml_backend_tensor_get_async( - b_.draft_backend_, state->hidden_prenorm, - prenorm_hidden.data(), 0, - sizeof(float) * prenorm_hidden.size()); ggml_backend_synchronize(b_.draft_backend_); - if (!dspark_markov_correct_greedy_chain_fused( + std::vector hidden_ptrs(lanes.size()); + for (size_t lane = 0; lane < lanes.size(); ++lane) { + hidden_ptrs[lane] = hidden_blocks[lane].data(); + } + if (!dflash2_select_chains_batched( b_.dw_, b_.draft_backend_, b_.w_.output, - local_hidden.data(), T, lanes[lane].seed, - drafts[lane], &confidences[lane], - prenorm_hidden.data())) { + hidden_ptrs, T, seeds, drafts, &selector_traces)) { return false; } + } else { + std::vector local_hidden((size_t) hidden * T); + std::vector prenorm_hidden((size_t) hidden * T); + for (size_t lane = 0; lane < lanes.size(); ++lane) { + DraftKvState * state = lanes[lane].state; + if (ggml_backend_graph_compute( + b_.draft_backend_, state->gf) != + GGML_STATUS_SUCCESS) { + return false; + } + ggml_backend_tensor_get_async( + b_.draft_backend_, state->hidden_states, + local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + ggml_backend_tensor_get_async( + b_.draft_backend_, state->hidden_prenorm, + prenorm_hidden.data(), 0, + sizeof(float) * prenorm_hidden.size()); + ggml_backend_synchronize(b_.draft_backend_); + if (!dspark_markov_correct_greedy_chain_fused( + b_.dw_, b_.draft_backend_, b_.w_.output, + local_hidden.data(), T, lanes[lane].seed, + drafts[lane], &confidences[lane], + prenorm_hidden.data())) { + return false; + } + } } } @@ -930,7 +1121,6 @@ bool Qwen35SeqEngine::prepare_chain_drafts( confidences.size() < lanes.size()) { return false; } - static bool missing_confidence_warned = false; for (size_t lane = 0; lane < lanes.size(); ++lane) { if ((int)drafts[lane].size() != T) return false; const Lane & info = lanes[lane]; @@ -942,18 +1132,39 @@ bool Qwen35SeqEngine::prepare_chain_drafts( prepared.root = info.seed; prepared.tokens = std::move(drafts[lane]); prepared.confidence = std::move(confidences[lane]); - - double score = std::numeric_limits::quiet_NaN(); - if (!prepared.confidence.empty()) { - score = confidence_survival_yield( - prepared.confidence, tree_width_); - } else if (!missing_confidence_warned) { - missing_confidence_warned = true; - std::fprintf(stderr, - "[spec-gate] current confidence unavailable; " - "cannot score adaptive activation\n"); + prepared.selector_trace = lane < selector_traces.size() + ? std::move(selector_traces[lane]) + : DFlash2SelectorTrace{}; + + double & published = last_survival_score_[(size_t)info.slot]; + if (!std::isfinite(published)) { + if (b_.dw_.selector.enabled) { + std::string adapter_error; + if (!dflash2_benefit_provider_ || + !dflash2_benefit_provider_->publish_once( + prepared.selector_trace, + chain_verify_depth_for_round(), published, + &adapter_error)) { + std::fprintf(stderr, + "[spec-gate] DFlash2 request-benefit evaluation " + "failed request=%llu slot=%d: %s\n", + (unsigned long long) + slots_.slot(info.slot).request_id, + info.slot, adapter_error.empty() + ? "adapter unavailable" : adapter_error.c_str()); + } + } else if (!prepared.confidence.empty()) { + published = confidence_survival_yield( + prepared.confidence, chain_verify_depth_for_round()); + } else if (chain_activation_input_scoreable( + inputs[info.input_index])) { + std::fprintf(stderr, + "[spec-gate] DSpark confidence evaluation unavailable " + "request=%llu slot=%d\n", + (unsigned long long) + slots_.slot(info.slot).request_id, info.slot); + } } - last_survival_score_[(size_t)info.slot] = score; } return true; } @@ -988,28 +1199,36 @@ bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { } return true; } -bool Qwen35SeqEngine::chain_confidence_input_scoreable( +bool Qwen35SeqEngine::chain_proposal_input_capable( const StepInput & in) const { - const int hidden = b_.dw_.n_embd; - const bool have_confidence = - b_.dw_.dspark.confidence_w && - b_.dw_.dspark.confidence_b && - (b_.dw_.dspark.confidence_dim == hidden || - b_.dw_.dspark.confidence_dim == - hidden + b_.dw_.dspark.markov_rank); - return spec_mode_ == SpecMode::dspark_chain && capture_features_ && + const bool have_dflash2 = + b_.dw_.selector.enabled && b_.dw_.selector.hproj && + b_.dw_.selector.pred_cb && b_.dw_.selector.succ_cb && + b_.dw_.selector.rank > 0 && b_.dw_.selector.top_k > 0; + const bool have_dspark = + b_.dw_.dspark.enabled && b_.dw_.dspark.markov_w1 && + b_.dw_.dspark.markov_w2; + return spec_mode_ == SpecMode::chain && capture_features_ && tree_width_ > 1 && tree_width_ <= 16 && - b_.dw_.block_size == tree_width_ && b_.dw_.dspark.enabled && - have_confidence && + resolve_chain_verify_depth( + chain_verify_depth_for_round(), tree_width_) != 0 && + b_.dw_.block_size == tree_width_ && + (have_dflash2 || have_dspark) && in.slot >= 0 && in.slot < slots_.slot_count() && slots_.slot(in.slot).decoding() && slots_.slot(in.slot).cur_pos >= 1 && slots_.slot(in.slot).cur_pos < slots_.max_context(); } +bool Qwen35SeqEngine::chain_activation_input_scoreable( + const StepInput & in) const { + return chain_proposal_input_capable(in) && + activation_scoring_available(); +} + bool Qwen35SeqEngine::chain_spec_request_capable( const StepInput & in) const { - return chain_confidence_input_scoreable(in) && + return chain_proposal_input_capable(in) && in.allow_speculation && in.speculation_policy != SpeculationPolicy::Never && !slots_.slot(in.slot).sampler.needs_logit_processing(); @@ -1030,6 +1249,12 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( } const int T = tree_width_; + const int V = chain_verify_depth_for_round(); + if (resolve_chain_verify_depth(V, T) == 0) { + result.error = + "invalid root-inclusive speculative chain verify depth"; + return result; + } const int hidden = b_.w_.n_embd; const int n_head_kv = b_.w_.n_head_kv; const int n_slots = slots_.slot_count(); @@ -1068,6 +1293,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( std::vector accepted; std::vector path; std::vector confidence; + DFlash2SelectorTrace selector_trace; int32_t verify_bonus = -1; int32_t pending = -1; }; @@ -1173,11 +1399,18 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( next.input_index = i; next.slot = in.slot; next.root = in.token; + next.selector_trace = std::move(prepared.selector_trace); next.flat = std::move(prepared.tokens); next.confidence = std::move(prepared.confidence); prepared.valid = false; + if (!truncate_chain_proposal(next.flat, V)) return false; + const size_t verified_signal_depths = + static_cast(V - 1); + if (next.selector_trace.depths.size() > verified_signal_depths) { + next.selector_trace.depths.resize(verified_signal_depths); + } next.tree = make_dspark_chain_tree(next.flat); - if (next.tree.n_nodes + 1 != T) return false; + if (next.tree.n_nodes + 1 != V) return false; proposal = std::move(next); return true; }; @@ -1221,7 +1454,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( t_verify_build_start = timing_clock::now(); if (!build_target_step_paged_tree( tree_sg, b_.w_, b_.cache_, b_.target_backend_, - T, tree_bucket, max_prefix, + V, tree_bucket, max_prefix, tree_scratch_base_, tree_scratch_stride_, b_.cfg_.kq_stride_pad)) { result.error = "packed DSpark chain verify graph build failed"; @@ -1229,7 +1462,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( } t_verify_build_end = timing_clock::now(); - const int total_tree = T * tree_bucket; + const int total_tree = V * tree_bucket; std::vector flat_tokens(static_cast(total_tree), 0); std::vector parents(static_cast(total_tree), -1); std::vector sizes(static_cast(tree_bucket), 0); @@ -1247,13 +1480,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( for (int lane = 0; lane < spec_count; ++lane) { const Proposal & proposal = proposals[static_cast(lane)]; - const int base = lane * T; - sizes[static_cast(lane)] = T; + const int base = lane * V; + sizes[static_cast(lane)] = V; tree_slots[static_cast(lane)] = proposal.slot; tree_state_slots[static_cast(lane)] = proposal.slot; seq_lens_[static_cast(proposal.slot)] = slots_.slot(proposal.slot).cur_pos; - for (int node = 0; node < T; ++node) { + for (int node = 0; node < V; ++node) { const int row = base + node; flat_tokens[static_cast(row)] = proposal.flat[static_cast(node)]; @@ -1324,7 +1557,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( for (int lane = 0; lane < spec_count; ++lane) { Proposal & proposal = proposals[static_cast(lane)]; const int32_t * lane_posterior = - posterior.data() + static_cast(lane) * T; + posterior.data() + static_cast(lane) * V; proposal.accepted = follow_verified_tree( proposal.tree, lane_posterior, proposal.verify_bonus); const int room = @@ -1350,6 +1583,44 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( [&](int32_t token) { return token_is_eos(token); }); proposal.path.resize(safe_prefix); proposal.accepted.resize(safe_prefix); + if (!proposal.selector_trace.depths.empty()) { + static const bool selector_log_enabled = []() { + const char * value = + std::getenv("DFLASH_DFLASH2_SELECTOR_LOG"); + return value && std::atoi(value) != 0; + }(); + if (selector_log_enabled) { + const Qwen35Slot & sequence = slots_.slot(proposal.slot); + const size_t accepted_depth = + proposal.path.empty() ? 0 : proposal.path.size() - 1; + std::fprintf(stderr, + "[spec-selector] {\"request_id\":%llu,\"slot\":%d," + "\"generated\":%d,\"accepted_depth\":%zu,\"depths\":[", + (unsigned long long) sequence.request_id, + proposal.slot, sequence.generated_tokens(), + accepted_depth); + for (size_t depth = 0; + depth < proposal.selector_trace.depths.size(); ++depth) { + const DFlash2DepthSignal & signal = + proposal.selector_trace.depths[depth]; + std::fprintf(stderr, + "%s{\"depth\":%zu,\"accepted\":%s," + "\"selected_logp\":%.8g,\"lm_margin\":%.8g," + "\"topk_mass\":%.8g,\"rank\":%d," + "\"lm_top1\":%s,\"selector_margin\":%.8g," + "\"selector_mass\":%.8g,\"selector_entropy\":%.8g}", + depth == 0 ? "" : ",", depth + 1, + depth < accepted_depth ? "true" : "false", + signal.selected_log_prob, signal.lm_top2_margin, + signal.top_k_mass, signal.selected_rank, + signal.agrees_with_lm_top1 ? "true" : "false", + signal.selector_margin, + signal.selector_winner_mass, + signal.selector_entropy); + } + std::fprintf(stderr, "]}\n"); + } + } replay_total += static_cast(proposal.path.size()); } } else { @@ -1739,7 +2010,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( "\"accepted_tokens\":%d,\"emitted_tokens\":%d," "\"target_forwards\":%d}\n", spec_count + ar_count, spec_count, - tree_bucket, T * tree_bucket, replay_total, + tree_bucket, V * tree_bucket, replay_total, ar_count, ar_bucket, max_kv_len, round_draft_us_, round_draft_lanes_, span_us(t_round_start, t_verify_build_start), @@ -2203,6 +2474,10 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( prepared_chain_drafts_[(size_t)result.slot].valid = false; } } + if (result.slot >= 0 && + result.slot < (int)adaptive_fallback_ar_.size()) { + adaptive_fallback_ar_[(size_t)result.slot] = 0; + } } return result; } @@ -2449,13 +2724,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const bool timing = step_timing_enabled(); using timing_clock = std::chrono::steady_clock; const bool gate_cost_timing = - spec_mode_ == SpecMode::dspark_chain && + spec_mode_ == SpecMode::chain && speculation_gate_ != nullptr; const auto decode_round_started = timing || gate_cost_timing ? timing_clock::now() : timing_clock::time_point{}; std::optional pending_ar_gate_plan; - if (spec_mode_ == SpecMode::dspark_chain && !inputs.empty()) { + if (spec_mode_ == SpecMode::chain && !inputs.empty()) { // New chain round: restart the [step-timing] draft attribution. round_draft_us_ = 0.0; round_draft_lanes_ = 0; @@ -2470,7 +2745,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (speculation_gate_) { std::vector candidates; candidates.reserve(inputs.size()); - const bool use_confidence = confidence_scoring_enabled(); + const bool use_activation_score = activation_scoring_enabled(); for (const StepInput & in : inputs) { const Qwen35Slot & seq = slots_.slot(in.slot); SpeculationPolicy policy = in.speculation_policy; @@ -2478,27 +2753,28 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (force == "all") policy = SpeculationPolicy::Always; if (force == "none") policy = SpeculationPolicy::Never; } - // The confidence-off arm is an explicit AR ablation. It must + // The activation-score-off arm is an explicit AR ablation. It must // not pay the one-time activation draft. - if (!use_confidence && + if (!use_activation_score && policy == SpeculationPolicy::Adaptive) { policy = SpeculationPolicy::Never; } const bool scoreable = - chain_confidence_input_scoreable(in); + chain_activation_input_scoreable(in); const bool can_speculate = chain_spec_request_capable(in); - double confidence = + double activation_score = std::numeric_limits::quiet_NaN(); - if (use_confidence && scoreable && in.slot >= 0 && + if (use_activation_score && scoreable && in.slot >= 0 && in.slot < (int)last_survival_score_.size() && std::isfinite( last_survival_score_[(size_t)in.slot])) { - confidence = last_survival_score_[(size_t)in.slot]; + activation_score = last_survival_score_[(size_t)in.slot]; } candidates.push_back({ seq.request_id, in.slot, policy, - scoreable, can_speculate, confidence, + scoreable, can_speculate, activation_score, + chain_activation_score_kind(), }); } gate_plan = speculation_gate_->plan( @@ -2509,7 +2785,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { return fail_step(gate_plan.error.empty() ? "adaptive speculation gate failed" : gate_plan.error); } - auto replan_with_published_confidence = [&]() { + auto replan_with_published_score = [&]() { for (SpecCandidate & candidate : candidates) { if (!candidate.scoreable || candidate.policy == SpeculationPolicy::Never) { @@ -2550,8 +2826,15 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { reset_evaluation_lane(evaluation.slot); if (speculation_gate_->commit_evaluation_fallback_ar( evaluation.request_id)) { + const SpecScoreKind kind = + chain_activation_score_kind(); + const char * reason = kind == + SpecScoreKind::DFlash2SelectorBenefitV1 + ? "benefit_evaluation_failed" + : "confidence_evaluation_failed"; log_spec_evaluation_fallback( - evaluation.request_id, evaluation.slot); + evaluation.request_id, evaluation.slot, + kind, reason); } }; @@ -2577,7 +2860,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } } - if (use_confidence && !score_evaluations.empty()) { + if (use_activation_score && !score_evaluations.empty()) { const bool batch_scored = prepare_chain_drafts( inputs, bootstrap, /*force_serial=*/false, /*fail_fast_batch=*/true); @@ -2623,7 +2906,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } } - if (!replan_with_published_confidence()) { + if (!replan_with_published_score()) { return fail_step(gate_plan.error.empty() ? "adaptive speculation gate failed" : gate_plan.error); } @@ -2636,7 +2919,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { for (const SpecPendingEvaluation & evaluation : unresolved) { commit_evaluation_fallback(evaluation); } - if (!replan_with_published_confidence()) { + if (!replan_with_published_score()) { return fail_step(gate_plan.error.empty() ? "adaptive speculation gate failed" : gate_plan.error); @@ -2650,6 +2933,24 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } log_spec_activations(gate_plan, *speculation_gate_); + // A one-shot AR decision discards the evaluation proposal. SPEC + // keeps that exact first proposal so the bootstrap is useful work. + for (const SpecPlanScore & score : gate_plan.ordered) { + if (!score.newly_decided || + score.decision != SpecDecision::AR) { + continue; + } + const int slot = score.slot; + if (slot >= 0 && + slot < (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)slot] = {}; + } + if (slot >= 0 && slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)slot]); + } + } + // Preserve every sticky Spec admission. Prefills are deferred // below, the min-token floor is enforced inside the speculative // path, and a later capability invariant failure becomes a @@ -2668,9 +2969,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } } } else { - // Explicit forced speculation remains usable without a cost - // profile. Adaptive mode fails explicitly: silently serving an - // unscored request as AR violates its activation contract. + // Forced modes remain available without a cost profile. A normal + // Adaptive request fails closed once, request-locally, and then + // remains AR for its entire slot lifetime without failing UX. for (size_t i = 0; i < inputs.size(); ++i) { SpeculationPolicy policy = inputs[i].speculation_policy; if (policy == SpeculationPolicy::Adaptive && @@ -2680,8 +2981,22 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { force == "none") { policy = SpeculationPolicy::Never; } else if (policy == SpeculationPolicy::Adaptive) { - return fail_step( - "adaptive speculation gate unavailable"); + policy = SpeculationPolicy::Never; + const int slot = inputs[i].slot; + if (slot >= 0 && + slot < (int)adaptive_fallback_ar_.size() && + !adaptive_fallback_ar_[(size_t)slot]) { + adaptive_fallback_ar_[(size_t)slot] = 1; + const uint64_t request_id = + slots_.slot(slot).request_id; + const char * reason = + adaptive_fallback_reason_.empty() + ? "cost_profile_unavailable" + : adaptive_fallback_reason_.c_str(); + log_spec_evaluation_fallback( + request_id, slot, + chain_activation_score_kind(), reason); + } } admitted[i] = policy == SpeculationPolicy::Always; } @@ -2700,6 +3015,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { std::chrono::steady_clock::now() - chain_started).count(); const bool spec_completed = speculative.error.empty(); bool proposal_failed = false; + std::vector accepted_lengths(inputs.size(), 0); if (spec_completed) { for (size_t i = 0; i < inputs.size(); ++i) { if (!admitted[i]) continue; @@ -2708,10 +3024,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { [&](const DecodeOutput & item) { return item.slot == inputs[i].slot; }); - if (output == speculative.decode.end() || output->failed) { + if (output == speculative.decode.end() || output->failed || + output->spec_steps == 0) { proposal_failed = true; break; } + accepted_lengths[i] = + 1 + static_cast(output->spec_accepted_tokens); } } if (spec_completed && !plan.prefills.empty()) { @@ -2722,10 +3041,16 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } } const bool cost_sample_valid = - have_gate_plan && spec_completed && !proposal_failed && - round_draft_lanes_ == gate_plan.draft_lanes; + have_gate_plan && spec_completed && !proposal_failed; if (cost_sample_valid) { - speculation_gate_->observe_cost(gate_plan, measured_us); + const ChainLaunchShape executed = chain_launch_shape( + admitted, accepted_lengths, + chain_verify_depth_for_round()); + speculation_gate_->observe_cost( + {static_cast(inputs.size()), executed.spec_lanes, + executed.tree_rows, executed.commit_rows, + round_draft_lanes_}, + measured_us); } if (have_gate_plan && spec_gate_debug_enabled() && spec_completed && !proposal_failed) { @@ -3140,18 +3465,14 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (pending_ar_gate_plan) { const double measured_us = span_us(decode_round_started, t_ar_end); - const bool cost_sample_valid = round_draft_lanes_ == 0; - if (cost_sample_valid) { - speculation_gate_->observe_cost( - *pending_ar_gate_plan, measured_us); - } + speculation_gate_->observe_cost( + {live_count, 0, 0, decode_bucket, round_draft_lanes_}, + measured_us); if (spec_gate_debug_enabled()) { log_spec_gate_plan( *pending_ar_gate_plan, speculation_gate_->fixed_yield_scale(), 0.0, - cost_sample_valid - ? measured_us - : std::numeric_limits::quiet_NaN()); + measured_us); } } if (timing) { @@ -3190,6 +3511,9 @@ void Qwen35SeqEngine::retire(int slot) { prepared_chain_drafts_[(size_t)slot].valid = false; } } + if (slot >= 0 && slot < (int)adaptive_fallback_ar_.size()) { + adaptive_fallback_ar_[(size_t)slot] = 0; + } slots_.retire(slot); } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index b0d5f8751..ed2c79564 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -23,6 +23,7 @@ #include "common/concurrency/seq_engine.h" #include "common/concurrency/speculation_gate.h" +#include "common/dflash2_benefit.h" #include "common/dflash_draft_kv.h" #include "common/dflash_feature_ring.h" #include "common/ddtree.h" @@ -34,6 +35,7 @@ #include #include #include +#include #include namespace dflash::common { @@ -45,7 +47,7 @@ class Qwen35SeqEngine final : public SeqEngine { enum class SpecMode { none, ddtree, - dspark_chain, + chain, }; // `pool` and `backend` must outlive the engine. `scratch_row` is the @@ -76,6 +78,10 @@ class Qwen35SeqEngine final : public SeqEngine { // Fabricate a steady-state paged context and profile the three launch // families used by the DSpark gate. Called once from backend init. bool profile_spec_costs(int context_tokens); + // True only when a trained DSpark head or the guarded DFlash2 benefit + // adapter can produce a first-request activation score. A configured chain + // may still accept Adaptive requests and serve sticky AR when this is false. + bool activation_scoring_available() const; StepPlanLimits step_plan_limits(int decode_rows) const override { const bool mixed = decode_rows > 0; const int per_sequence = mixed ? 512 : 2048; @@ -133,7 +139,8 @@ class Qwen35SeqEngine final : public SeqEngine { DraftFeatureMirror * slot_feature_mirror(int slot); DraftKvState * ensure_slot_draft_kv(int slot); bool ddtree_eligible(const StepPlan & plan) const; - bool chain_confidence_input_scoreable(const StepInput & input) const; + bool chain_proposal_input_capable(const StepInput & input) const; + bool chain_activation_input_scoreable(const StepInput & input) const; bool chain_spec_request_capable(const StepInput & input) const; bool chain_spec_input_eligible(const StepInput & input) const; bool spec_gate_debug_enabled() const; @@ -143,6 +150,7 @@ class Qwen35SeqEngine final : public SeqEngine { int32_t root = -1; std::vector tokens; std::vector confidence; + DFlash2SelectorTrace selector_trace; }; bool prepare_chain_drafts( const std::vector & inputs, @@ -150,7 +158,13 @@ class Qwen35SeqEngine final : public SeqEngine { bool force_serial = false, bool fail_fast_batch = false); bool batched_drafting_enabled() const; - bool confidence_scoring_enabled() const; + bool activation_scoring_enabled() const; + SpecScoreKind chain_activation_score_kind() const; + // Deliberate seam for a later per-round cohort controller. Request mode + // remains sticky; every returned depth must stay in [2, tree_width_]. + int chain_verify_depth_for_round() const { + return chain_verify_depth_; + } // DFLASH_STEP_TIMING=1 emits one [step-timing] JSON line per decode // round attributing wall time to draft, verify, readback, CPU commit, // replay, and packed-AR phases. Diagnostic only; off by default. @@ -167,6 +181,10 @@ class Qwen35SeqEngine final : public SeqEngine { Qwen35SlotManager slots_; int64_t scratch_row_ = 0; int tree_width_ = 0; + // Root-inclusive verification depth. The drafter still produces + // tree_width_ tokens; this common cohort depth may vary between 2 and + // tree_width_ without changing a request's sticky SPEC mode. + int chain_verify_depth_ = 0; int tree_scratch_base_ = 0; int tree_scratch_stride_ = 0; bool capture_features_ = false; @@ -179,6 +197,11 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector> dummy_draft_kv_; std::vector prepared_chain_drafts_; std::unique_ptr speculation_gate_; + std::unique_ptr dflash2_benefit_provider_; + // Startup profile/adapter failure is a request-local sticky AR outcome, + // never an admission or step error for a configured chain. + std::string adaptive_fallback_reason_ = "cost_profile_unavailable"; + std::vector adaptive_fallback_ar_; std::vector last_survival_score_; // Per-round draft cost accumulator for [step-timing]; reset at the top // of each dspark_chain round, accumulated by prepare_chain_drafts. diff --git a/server/src/qwen35/delta_transition_journal.cpp b/server/src/qwen35/delta_transition_journal.cpp new file mode 100644 index 000000000..8a1ff1638 --- /dev/null +++ b/server/src/qwen35/delta_transition_journal.cpp @@ -0,0 +1,142 @@ +#include "qwen35/delta_transition_journal.h" + +#include +#include +#include + +namespace dflash::qwen35 { +namespace { + +bool matrix_size(size_t rows, size_t cols, size_t & elements) { + if (rows == 0 || cols == 0 || + rows > std::numeric_limits::max() / cols) { + return false; + } + elements = rows * cols; + return true; +} + +bool transition_shape_valid( + const DeltaTransition & transition, + size_t rows, + size_t cols) { + const size_t expected_gate = + transition.gate_mode == DeltaTransitionGateMode::Scalar ? 1 : rows; + return transition.gate.size() == expected_gate && + transition.key.size() == rows && + transition.delta.size() == cols; +} + +float gate_at(const DeltaTransition & transition, size_t row) { + return transition.gate_mode == DeltaTransitionGateMode::Scalar + ? transition.gate[0] + : transition.gate[row]; +} + +} // namespace + +bool capture_delta_transition( + const std::vector & state, + size_t rows, + size_t cols, + const std::vector & key, + const std::vector & value, + const std::vector & gate, + float beta, + DeltaTransitionGateMode gate_mode, + DeltaTransition & output) { + size_t state_elements = 0; + const size_t expected_gate = + gate_mode == DeltaTransitionGateMode::Scalar ? 1 : rows; + if (!matrix_size(rows, cols, state_elements) || + state.size() != state_elements || key.size() != rows || + value.size() != cols || gate.size() != expected_gate) { + return false; + } + + DeltaTransition next; + next.gate_mode = gate_mode; + next.gate = gate; + next.key = key; + next.delta.resize(cols); + + for (size_t col = 0; col < cols; ++col) { + float projection = 0.0f; + for (size_t row = 0; row < rows; ++row) { + const float row_gate = gate_mode == DeltaTransitionGateMode::Scalar + ? 1.0f + : gate[row]; + projection += row_gate * state[col * rows + row] * key[row]; + } + const float scalar_gate = gate_mode == DeltaTransitionGateMode::Scalar + ? gate[0] + : 1.0f; + next.delta[col] = + (value[col] - scalar_gate * projection) * beta; + } + + output = std::move(next); + return true; +} + +bool apply_delta_transition( + const DeltaTransition & transition, + size_t rows, + size_t cols, + std::vector & state) { + size_t state_elements = 0; + if (!matrix_size(rows, cols, state_elements) || + state.size() != state_elements || + !transition_shape_valid(transition, rows, cols)) { + return false; + } + + for (size_t col = 0; col < cols; ++col) { + for (size_t row = 0; row < rows; ++row) { + const size_t index = col * rows + row; + state[index] = std::fma( + transition.key[row], transition.delta[col], + gate_at(transition, row) * state[index]); + } + } + return true; +} + +bool commit_delta_transition_prefix( + const DeltaTransitionJournal & journal, + size_t accepted, + std::vector & state) { + size_t state_elements = 0; + if (!matrix_size(journal.rows, journal.cols, state_elements) || + state.size() != state_elements || + accepted > journal.transitions.size()) { + return false; + } + for (size_t i = 0; i < accepted; ++i) { + if (!transition_shape_valid( + journal.transitions[i], journal.rows, journal.cols)) { + return false; + } + } + for (size_t i = 0; i < accepted; ++i) { + // Already validated, so this cannot partially fail. + apply_delta_transition( + journal.transitions[i], journal.rows, journal.cols, state); + } + return true; +} + +size_t delta_transition_float_count( + size_t rows, + size_t cols, + DeltaTransitionGateMode gate_mode) { + const size_t gate_values = + gate_mode == DeltaTransitionGateMode::Scalar ? 1 : rows; + if (rows > std::numeric_limits::max() - cols || + rows + cols > std::numeric_limits::max() - gate_values) { + return 0; + } + return rows + cols + gate_values; +} + +} // namespace dflash::qwen35 diff --git a/server/src/qwen35/delta_transition_journal.h b/server/src/qwen35/delta_transition_journal.h new file mode 100644 index 000000000..5dbe887ba --- /dev/null +++ b/server/src/qwen35/delta_transition_journal.h @@ -0,0 +1,74 @@ +#pragma once + +#include +#include + +namespace dflash::qwen35 { + +// Host-side contract for the compact journal emitted by a future GDN verify +// kernel. The persistent state uses the kernel's transposed layout: +// state[col * rows + row]. A transition contains exactly the values needed to +// repeat the state update, but none of the model projections: +// +// state' = gate * state + key (outer-product) delta +// +// In scalar mode gate has one value. In row-wise mode it has one value per +// state row. key and delta are the normalized/resolved values produced by +// verification; in particular, delta is captured after the state-dependent +// k^T S reduction. Therefore this journal is valid only for a chain replayed +// from the same base recurrent state. +enum class DeltaTransitionGateMode { + Scalar, + RowWise, +}; + +struct DeltaTransition { + DeltaTransitionGateMode gate_mode = DeltaTransitionGateMode::Scalar; + std::vector gate; + std::vector key; + std::vector delta; +}; + +struct DeltaTransitionJournal { + size_t rows = 0; + size_t cols = 0; + std::vector transitions; +}; + +// Resolve the state-dependent delta exactly once, as the verification kernel +// would. gate contains already-exponentiated multipliers (not raw/log gates). +// The output is unchanged when validation fails. +bool capture_delta_transition( + const std::vector & state, + size_t rows, + size_t cols, + const std::vector & key, + const std::vector & value, + const std::vector & gate, + float beta, + DeltaTransitionGateMode gate_mode, + DeltaTransition & output); + +// Apply one captured transition without evaluating projections or recomputing +// delta. The state is unchanged when validation fails. +bool apply_delta_transition( + const DeltaTransition & transition, + size_t rows, + size_t cols, + std::vector & state); + +// Commit transitions [0, accepted) to a persistent state. accepted == 0 is a +// no-op. Oversized or malformed prefixes fail before modifying state. +bool commit_delta_transition_prefix( + const DeltaTransitionJournal & journal, + size_t accepted, + std::vector & state); + +// Compact recurrent journal footprint per head/token, excluding allocator +// alignment. Runtime storage may further share keys across grouped V heads. +size_t delta_transition_float_count( + size_t rows, + size_t cols, + DeltaTransitionGateMode gate_mode); + +} // namespace dflash::qwen35 diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 46e87575c..03336699a 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -448,21 +448,16 @@ bool Qwen35Backend::init() { const bool concurrent_local_ddtree = concurrent_local_draft && cfg_.ddtree_mode; const bool concurrent_local_chain = - concurrent_local_draft && !cfg_.ddtree_mode && dw_.dspark.enabled && - qwen35_dspark_enabled() && + concurrent_local_draft && !cfg_.ddtree_mode && + (dw_.selector.enabled || + (dw_.dspark.enabled && qwen35_dspark_enabled())) && cfg_.speculation_policy != SpeculationPolicy::Never; - const bool concurrent_confidence_head = - concurrent_local_chain && dw_.dspark.confidence_w && - dw_.dspark.confidence_b && - (dw_.dspark.confidence_dim == dw_.n_embd || - dw_.dspark.confidence_dim == - dw_.n_embd + dw_.dspark.markov_rank); const bool concurrent_spec_tree = concurrent_local_ddtree || concurrent_local_chain; const Qwen35SeqEngine::SpecMode spec_mode = concurrent_local_ddtree ? Qwen35SeqEngine::SpecMode::ddtree : concurrent_local_chain - ? Qwen35SeqEngine::SpecMode::dspark_chain + ? Qwen35SeqEngine::SpecMode::chain : Qwen35SeqEngine::SpecMode::none; const int tree_width = concurrent_local_ddtree ? cfg_.ddtree_budget + 1 @@ -605,14 +600,17 @@ bool Qwen35Backend::init() { long_mixed_prefill_tokens, long_prefill_threshold, idle_prefill_tokens, prefill_quantum); concurrent_decode_capabilities_.forced_speculation = - concurrent_local_ddtree || concurrent_confidence_head; + concurrent_local_ddtree || concurrent_local_chain; concurrent_decode_capabilities_.adaptive = - concurrent_local_ddtree; + concurrent_local_ddtree || concurrent_local_chain; // Per-request decode_mode may select Adaptive even when the - // server default is forced speculation. Build the gate whenever - // the concurrent DSpark chain exists so such requests cannot - // silently fall through to unscored AR. - if (concurrent_confidence_head) { + // server default is forced speculation. A configured chain always + // accepts Adaptive: if activation scoring or startup profiling is + // unavailable, the engine records a request-local sticky-AR + // fallback instead of rejecting the request or failing its peers. + bool adaptive_scored = false; + if (concurrent_local_chain && + seq_engine_->activation_scoring_available()) { int profile_ctx = 4096; if (const char * value = std::getenv("DFLASH_SPEC_PROFILE_CONTEXT")) { @@ -620,19 +618,20 @@ bool Qwen35Backend::init() { } const bool profile_ready = seq_engine_->profile_spec_costs(profile_ctx); - concurrent_decode_capabilities_.adaptive = profile_ready; + adaptive_scored = profile_ready; if (!profile_ready) { std::fprintf(stderr, - "[parallel-dspark] adaptive capability unavailable: " - "cost profile failed; adaptive requests will be " - "rejected at admission (forced speculation remains " + "[parallel-chain] adaptive scoring unavailable: " + "cost profile failed; adaptive requests will use " + "request-local sticky AR (forced speculation remains " "available)\n"); } } else if (concurrent_local_chain) { std::fprintf(stderr, - "[parallel-dspark] activation unavailable: drafter has " - "no compatible confidence head; speculation/adaptive " - "requests will be rejected at admission\n"); + "[parallel-chain] adaptive activation unavailable: " + "drafter has no compatible request-benefit adapter; " + "adaptive requests will use request-local sticky AR " + "(forced speculation remains available)\n"); } if (concurrent_local_ddtree) { const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); @@ -641,23 +640,25 @@ bool Qwen35Backend::init() { cfg_.ddtree_budget, tree_width, adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); } - if (concurrent_confidence_head) { + if (concurrent_local_chain) { std::fprintf(stderr, - "[parallel-dspark] enabled width=%d mode=packed-chain-verify " - "decode_mode=%s draft=q4-mix-compatible\n", + "[parallel-chain] enabled producer=%s width=%d " + "mode=packed-chain-verify decode_mode=%s adaptive=%s\n", + dw_.selector.enabled ? "dflash2" : "dspark", tree_width, - speculation_policy_name(cfg_.speculation_policy)); + speculation_policy_name(cfg_.speculation_policy), + adaptive_scored ? "scored" : "fallback-ar"); } else if (!cfg_.ddtree_mode && - cfg_.speculation_policy != SpeculationPolicy::Never && - !concurrent_local_chain) { + cfg_.speculation_policy != SpeculationPolicy::Never) { std::fprintf(stderr, - "[parallel-dspark] unavailable for this concurrent " + "[parallel-chain] unavailable for this concurrent " "configuration; speculation/adaptive requests will be " "rejected at admission\n"); } else if (!cfg_.ddtree_mode && concurrent_local_draft && - dw_.dspark.enabled && qwen35_dspark_enabled()) { + (dw_.selector.enabled || + (dw_.dspark.enabled && qwen35_dspark_enabled()))) { std::fprintf(stderr, - "[parallel-dspark] disabled by decode_mode=ar; " + "[parallel-chain] disabled by decode_mode=ar; " "per-request speculation/adaptive overrides will be " "rejected at admission\n"); } diff --git a/server/test/test_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp index 88bd6ccf7..e7a63dbda 100644 --- a/server/test/test_chain_spec_shapes.cpp +++ b/server/test/test_chain_spec_shapes.cpp @@ -15,6 +15,24 @@ int main() { CHECK((tree.token_ids == std::vector{11, 12, 13})); CHECK((tree.depths == std::vector{1, 2, 3})); CHECK((tree.parents == std::vector{-1, 0, 1, 2})); + + CHECK(resolve_chain_verify_depth(0, 4) == 4); + CHECK(resolve_chain_verify_depth(2, 4) == 2); + CHECK(resolve_chain_verify_depth(4, 4) == 4); + CHECK(resolve_chain_verify_depth(1, 4) == 0); + CHECK(resolve_chain_verify_depth(5, 4) == 0); + CHECK(resolve_chain_verify_depth(0, 1) == 0); + std::vector short_draft = draft; + CHECK(truncate_chain_proposal(short_draft, 2)); + CHECK((short_draft == std::vector{10, 11})); + const DDTree short_tree = make_dspark_chain_tree(short_draft); + CHECK(short_tree.n_nodes == 1); + CHECK((short_tree.parents == std::vector{-1, 0})); + const std::vector before_invalid = short_draft; + CHECK(!truncate_chain_proposal(short_draft, 1)); + CHECK(short_draft == before_invalid); + CHECK(!truncate_chain_proposal(short_draft, 3)); + CHECK(short_draft == before_invalid); const std::vector bucket_inputs = { 1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 13, 16, 17, }; diff --git a/server/test/test_delta_transition_journal.cpp b/server/test/test_delta_transition_journal.cpp new file mode 100644 index 000000000..a5859b2e6 --- /dev/null +++ b/server/test/test_delta_transition_journal.cpp @@ -0,0 +1,186 @@ +#include "qwen35/delta_transition_journal.h" +#include "host_check.h" + +#include +#include +#include +#include + +using dflash::qwen35::DeltaTransition; +using dflash::qwen35::DeltaTransitionGateMode; +using dflash::qwen35::DeltaTransitionJournal; +using dflash::qwen35::apply_delta_transition; +using dflash::qwen35::capture_delta_transition; +using dflash::qwen35::commit_delta_transition_prefix; +using dflash::qwen35::delta_transition_float_count; + +static int g_checks = 0; + +namespace { + +struct RawStep { + std::vector key; + std::vector value; + std::vector gate; + float beta = 0.0f; +}; + +bool near(const std::vector & lhs, const std::vector & rhs) { + if (lhs.size() != rhs.size()) return false; + for (size_t i = 0; i < lhs.size(); ++i) { + const float scale = std::max( + 1.0f, std::max(std::fabs(lhs[i]), std::fabs(rhs[i]))); + if (std::fabs(lhs[i] - rhs[i]) > 4e-6f * scale) return false; + } + return true; +} + +// Independent spelling of the existing gated_delta_net_cuda recurrence. It +// deliberately consumes raw step inputs rather than a captured transition. +void replay_reference( + std::vector & state, + size_t rows, + size_t cols, + const RawStep & step, + DeltaTransitionGateMode gate_mode) { + for (size_t col = 0; col < cols; ++col) { + float projection = 0.0f; + for (size_t row = 0; row < rows; ++row) { + const float projection_gate = + gate_mode == DeltaTransitionGateMode::RowWise + ? step.gate[row] + : 1.0f; + projection += projection_gate * + state[col * rows + row] * step.key[row]; + } + const float scalar_gate = + gate_mode == DeltaTransitionGateMode::Scalar + ? step.gate[0] + : 1.0f; + const float delta = + (step.value[col] - scalar_gate * projection) * step.beta; + for (size_t row = 0; row < rows; ++row) { + const size_t index = col * rows + row; + const float update_gate = + gate_mode == DeltaTransitionGateMode::Scalar + ? step.gate[0] + : step.gate[row]; + state[index] = std::fma( + step.key[row], delta, update_gate * state[index]); + } + } +} + +std::vector initial_state(size_t rows, size_t cols, int layer) { + std::vector state(rows * cols); + for (size_t i = 0; i < state.size(); ++i) { + state[i] = 0.013f * static_cast(i + 1) - + 0.07f * static_cast(layer + 1); + } + return state; +} + +std::vector make_steps( + size_t rows, + size_t cols, + size_t count, + int layer, + DeltaTransitionGateMode gate_mode) { + std::vector steps(count); + for (size_t t = 0; t < count; ++t) { + RawStep & step = steps[t]; + step.key.resize(rows); + step.value.resize(cols); + step.gate.resize( + gate_mode == DeltaTransitionGateMode::Scalar ? 1 : rows); + for (size_t row = 0; row < rows; ++row) { + step.key[row] = 0.021f * static_cast(row + 1) - + 0.009f * static_cast(t + layer); + } + for (size_t col = 0; col < cols; ++col) { + step.value[col] = 0.031f * static_cast(col + 1) + + 0.017f * static_cast(t + 2 * layer); + } + for (size_t row = 0; row < step.gate.size(); ++row) { + step.gate[row] = 0.78f + + 0.011f * static_cast((row + t + layer) % 7); + } + step.beta = 0.42f + 0.03f * static_cast(t % 4); + } + return steps; +} + +void prove_all_prefixes(DeltaTransitionGateMode gate_mode) { + constexpr size_t rows = 8; + constexpr size_t cols = 7; + constexpr size_t tokens = 6; + constexpr int layers = 2; + + for (int layer = 0; layer < layers; ++layer) { + const std::vector base = initial_state(rows, cols, layer); + const std::vector steps = + make_steps(rows, cols, tokens, layer, gate_mode); + + DeltaTransitionJournal journal; + journal.rows = rows; + journal.cols = cols; + std::vector verify_state = base; + for (const RawStep & step : steps) { + DeltaTransition transition; + CHECK(capture_delta_transition( + verify_state, rows, cols, step.key, step.value, step.gate, + step.beta, gate_mode, transition)); + CHECK(apply_delta_transition( + transition, rows, cols, verify_state)); + journal.transitions.push_back(transition); + } + + for (size_t accepted = 1; accepted <= tokens; ++accepted) { + std::vector replayed = base; + for (size_t t = 0; t < accepted; ++t) { + replay_reference( + replayed, rows, cols, steps[t], gate_mode); + } + + std::vector committed = base; + CHECK(commit_delta_transition_prefix( + journal, accepted, committed)); + CHECK(near(committed, replayed)); + } + + std::vector full_replay = base; + for (const RawStep & step : steps) { + replay_reference(full_replay, rows, cols, step, gate_mode); + } + CHECK(near(verify_state, full_replay)); + } +} + +} // namespace + +int main() { + prove_all_prefixes(DeltaTransitionGateMode::Scalar); + prove_all_prefixes(DeltaTransitionGateMode::RowWise); + + CHECK(delta_transition_float_count( + 128, 128, DeltaTransitionGateMode::Scalar) == 257); + CHECK(delta_transition_float_count( + 128, 128, DeltaTransitionGateMode::RowWise) == 384); + + // Contract guards are fail-closed and transactional. + DeltaTransitionJournal malformed; + malformed.rows = 2; + malformed.cols = 2; + malformed.transitions.push_back(DeltaTransition{}); + std::vector state = {1.0f, 2.0f, 3.0f, 4.0f}; + const std::vector original = state; + CHECK(!commit_delta_transition_prefix(malformed, 1, state)); + CHECK(state == original); + CHECK(!commit_delta_transition_prefix(malformed, 2, state)); + CHECK(state == original); + CHECK(commit_delta_transition_prefix(malformed, 0, state)); + CHECK(state == original); + + std::printf("delta transition journal tests passed: %d checks\n", g_checks); + return 0; +} diff --git a/server/test/test_dflash2_benefit.cpp b/server/test/test_dflash2_benefit.cpp new file mode 100644 index 000000000..23023205c --- /dev/null +++ b/server/test/test_dflash2_benefit.cpp @@ -0,0 +1,230 @@ +#include "common/dflash2_benefit.h" +#include "common/concurrency/speculation_gate.h" +#include "host_check.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +static DFlash2BenefitModelSignature seeded_signature() { + DFlash2BenefitModelSignature value; + value.target_layers = 64; + value.target_hidden = 5120; + value.target_vocab = 248320; + value.draft_layers = 5; + value.draft_hidden = 5120; + value.draft_block_size = 8; + value.selector_rank = 256; + value.selector_top_k = 16; + value.selector_vocab = 248320; + value.conv_kernel_size = 2; + value.conv_group_size = 16; + value.target_file_size = 15195272800ULL; + value.draft_file_size = 2045471776ULL; + return value; +} + +static DFlash2SelectorTrace trace_from( + std::initializer_list> values) { + DFlash2SelectorTrace trace; + for (const auto & value : values) { + DFlash2DepthSignal signal; + signal.selected_log_prob = value.first; + signal.selector_winner_mass = value.second; + trace.depths.push_back(signal); + } + return trace; +} + +int main() { + DFlash2BenefitProvider provider(seeded_signature()); + CHECK(provider.ready()); + CHECK(std::string(provider.score_kind()) == + "qwen38-dflash2-selector-benefit-v1"); + CHECK(provider.config().lm_log_weight == 0.10); + CHECK(provider.config().hazard_scale == 1.0); + + // Retained C1 first-block traces for he08 code and prose. The adapter is + // continuous and content-agnostic: code has the higher expected yield, while + // prose remains lower and cohort-dependent. Values are (selected LM + // log-prob, selector winner mass). + const DFlash2SelectorTrace code_like = trace_from({ + {0.0f, 0.99999988f}, {-7.6293945e-06f, 1.0f}, + {-0.44184685f, 0.64686394f}, {-2.0253334f, 1.0f}, + {-1.2170925f, 1.0f}, {-0.84755516f, 0.99942774f}, + {-3.7030106f, 0.86279351f}, + }); + const DFlash2SelectorTrace prose_like = trace_from({ + {-0.046934128f, 0.95176214f}, {-1.527895f, 0.37046611f}, + {-3.398098f, 0.98939508f}, {-3.0824165f, 0.41014573f}, + {-0.30518532f, 0.47964928f}, {-4.0057392f, 0.99999905f}, + {-1.676815f, 0.89187294f}, + }); + + DFlash2BenefitEstimate code; + DFlash2BenefitEstimate prose; + std::string error; + CHECK(provider.estimate(code_like, 8, code, &error)); + CHECK(error.empty()); + CHECK(provider.estimate(prose_like, 8, prose, &error)); + CHECK(code.conditional_hazards.size() == 7); + CHECK(prose.conditional_hazards.size() == 7); + CHECK(std::abs(code.expected_yield - 5.330600824) < 1e-6); + CHECK(std::abs(prose.expected_yield - 2.684506084) < 1e-6); + CHECK(code.expected_yield > prose.expected_yield + 2.5); + + // Feed the retained request benefits into representative C2 profile + // costs. These are joint economics, not a topic classifier: code wins + // alone, prose loses alone and as an all-prose C2 cohort, while a mixed + // cohort must rank code first and may admit prose only as a prefix-k + // amortization result. + const SpecCostTables observed_costs{ + {{8, 16}, {42912.6, 47572.8}}, + {{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + {29322.0, 34968.2, 38639.9, 38639.9, + 40077.6, 40077.6, 40077.6, 40077.6, + 42360.3, 42360.3, 42360.3, 42360.3, + 42360.3, 42360.3, 42360.3, 44412.2}}, + {{1, 2}, {7934.4, 13967.3}}, + }; + SpecStepGeometry observed_geometry; + observed_geometry.tree_width = 8; + observed_geometry.bucket = [](int lanes) { return lanes; }; + auto benefit_candidate = [](uint64_t id, int slot, double score, + SpeculationPolicy policy = + SpeculationPolicy::Adaptive) { + return SpecCandidate{ + id, slot, policy, true, true, score, + SpecScoreKind::DFlash2SelectorBenefitV1}; + }; + + SpeculationGate code_c1(observed_costs, observed_geometry, 8); + SpecPlan gate_plan = code_c1.plan( + 1, {benefit_candidate(100, 5, code.expected_yield)}, 1); + CHECK(gate_plan.admitted_count == 1); + CHECK(code_c1.decision(100) == SpecDecision::Speculation); + CHECK(code_c1.initial_score_kind(100) == + SpecScoreKind::DFlash2SelectorBenefitV1); + + SpeculationGate prose_c1(observed_costs, observed_geometry, 8); + gate_plan = prose_c1.plan( + 1, {benefit_candidate(200, 3, prose.expected_yield)}, 1); + CHECK(gate_plan.admitted_count == 0); + CHECK(prose_c1.decision(200) == SpecDecision::AR); + + SpeculationGate prose_c2(observed_costs, observed_geometry, 8); + gate_plan = prose_c2.plan(2, { + benefit_candidate(300, 4, prose.expected_yield), + benefit_candidate(301, 1, prose.expected_yield)}, 2); + CHECK(gate_plan.admitted_count == 0); + CHECK(prose_c2.decision(300) == SpecDecision::AR); + CHECK(prose_c2.decision(301) == SpecDecision::AR); + + SpeculationGate mixed_c2(observed_costs, observed_geometry, 8); + gate_plan = mixed_c2.plan(2, { + benefit_candidate(401, 3, prose.expected_yield), + benefit_candidate(400, 5, code.expected_yield)}, 2); + CHECK(gate_plan.admitted_count >= 1); + CHECK(!gate_plan.admitted_request_ids.empty()); + CHECK(gate_plan.admitted_request_ids[0] == 400); + CHECK(mixed_c2.decision(400) == SpecDecision::Speculation); + CHECK(gate_plan.ordered.size() == 2); + CHECK(gate_plan.ordered[0].request_id == 400); + CHECK(gate_plan.ordered[0].slot == 5); + CHECK(gate_plan.ordered[1].request_id == 401); + CHECK(gate_plan.ordered[1].slot == 3); + CHECK((gate_plan.admitted_count == 1 && + mixed_c2.decision(401) == SpecDecision::AR) || + (gate_plan.admitted_count == 2 && + mixed_c2.decision(401) == SpecDecision::Speculation)); + + SpeculationGate code_with_ar_peer( + observed_costs, observed_geometry, 8); + gate_plan = code_with_ar_peer.plan(2, { + benefit_candidate(500, 2, code.expected_yield), + benefit_candidate(501, 7, prose.expected_yield, + SpeculationPolicy::Never)}, 2); + CHECK(gate_plan.admitted_count == 1); + CHECK(code_with_ar_peer.decision(500) == SpecDecision::Speculation); + CHECK(code_with_ar_peer.decision(501) == SpecDecision::Undecided); + + // Maximum depth consumes exactly block_size-1 signals; diagnostic tail + // values beyond that depth cannot affect the score. + DFlash2SelectorTrace with_tail = code_like; + with_tail.depths.push_back(trace_from({{-100.0f, 0.001f}}).depths[0]); + DFlash2BenefitEstimate tail; + CHECK(provider.estimate(with_tail, 8, tail, &error)); + CHECK(tail.conditional_hazards.size() == 7); + CHECK(std::abs(tail.expected_yield - code.expected_yield) < 1e-12); + CHECK(!provider.estimate(code_like, 9, tail, &error)); + CHECK(error.find("outside") != std::string::npos); + + // A partial, malformed, or nonfinite first trace fails closed and never + // publishes a synthetic request score. + DFlash2SelectorTrace missing = code_like; + missing.depths.pop_back(); + CHECK(!provider.estimate(missing, 8, tail, &error)); + CHECK(error.find("missing") != std::string::npos); + DFlash2SelectorTrace malformed = code_like; + malformed.depths[2].selector_winner_mass = 0.0f; + CHECK(!provider.estimate(malformed, 8, tail, &error)); + malformed = code_like; + malformed.depths[2].selected_log_prob = 0.1f; + CHECK(!provider.estimate(malformed, 8, tail, &error)); + malformed = code_like; + malformed.depths[2].selected_log_prob = + std::numeric_limits::quiet_NaN(); + CHECK(!provider.estimate(malformed, 8, tail, &error)); + double unpublished = std::numeric_limits::quiet_NaN(); + CHECK(!provider.publish_once(malformed, 8, unpublished, &error)); + CHECK(std::isnan(unpublished)); + + // Request lifetime is external and explicit: once the first valid score + // is published, neither a lower later trace nor a malformed trace can + // overwrite it. + double first_score = std::numeric_limits::quiet_NaN(); + CHECK(provider.publish_once(code_like, 8, first_score, &error)); + CHECK(std::abs(first_score - code.expected_yield) < 1e-12); + CHECK(provider.publish_once(prose_like, 8, first_score, &error)); + CHECK(std::abs(first_score - code.expected_yield) < 1e-12); + CHECK(provider.publish_once(malformed, 8, first_score, &error)); + CHECK(std::abs(first_score - code.expected_yield) < 1e-12); + + // Unknown models and invalid adapter versions/coefficients are not + // silently generalized. A conservative scale may only reduce yield. + DFlash2BenefitModelSignature unknown = seeded_signature(); + unknown.selector_rank = 128; + DFlash2BenefitProvider unknown_provider(unknown); + CHECK(!unknown_provider.ready()); + CHECK(unknown_provider.error().find("unsupported") != std::string::npos); + CHECK(!unknown_provider.estimate(code_like, 8, tail, &error)); + + DFlash2BenefitConfig bad_version; + bad_version.adapter_version = "future-unfitted-adapter"; + DFlash2BenefitProvider version_provider( + seeded_signature(), bad_version); + CHECK(!version_provider.ready()); + + DFlash2BenefitConfig conservative; + conservative.hazard_scale = 0.9; + DFlash2BenefitProvider conservative_provider( + seeded_signature(), conservative); + CHECK(conservative_provider.ready()); + CHECK(conservative_provider.estimate(code_like, 8, tail, &error)); + CHECK(tail.expected_yield < code.expected_yield); + + DFlash2BenefitConfig invalid_coefficients; + invalid_coefficients.lm_log_weight = + std::numeric_limits::infinity(); + DFlash2BenefitProvider invalid_provider( + seeded_signature(), invalid_coefficients); + CHECK(!invalid_provider.ready()); + + std::printf("DFlash2 benefit adapter: %d checks passed\n", g_checks); + return 0; +} diff --git a/server/test/test_dflash2_selector_validation.cpp b/server/test/test_dflash2_selector_validation.cpp new file mode 100644 index 000000000..bbc46acc8 --- /dev/null +++ b/server/test/test_dflash2_selector_validation.cpp @@ -0,0 +1,86 @@ +#include "common/dflash2_selector_validation.h" +#include "host_check.h" + +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +static DFlash2SelectorLayout valid_layout() { + DFlash2SelectorLayout layout; + layout.rank = 32; + layout.top_k = 16; + layout.hproj_rank = 32; + layout.pred_rank = 32; + layout.pred_vocab = 151936; + layout.succ_rank = 32; + layout.succ_vocab = 151936; + layout.target_output_vocab = 151936; + layout.target_declared_vocab = 151936; + return layout; +} + +int main() { + std::string error; + DFlash2SelectorLayout layout = valid_layout(); + CHECK(validate_dflash2_selector_layout(layout, error)); + CHECK(error.empty()); + + for (int K = 1; K <= 8; ++K) { + layout = valid_layout(); + layout.top_k = K; + CHECK(validate_dflash2_selector_layout(layout, error)); + } + for (int K : {12, 16}) { + layout = valid_layout(); + layout.top_k = K; + CHECK(validate_dflash2_selector_layout(layout, error)); + } + for (int K : {0, 9, 10, 11, 13, 14, 15, 17}) { + layout = valid_layout(); + layout.top_k = K; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("top_k=") != std::string::npos); + CHECK(error.find("unsupported") != std::string::npos); + } + + layout = valid_layout(); + layout.succ_vocab--; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("codebook vocab mismatch") != std::string::npos); + + layout = valid_layout(); + layout.target_output_vocab--; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("target vocab mismatch") != std::string::npos); + + layout = valid_layout(); + layout.target_declared_vocab = 0; + layout.target_output_vocab--; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("target output/lm_head") != std::string::npos); + + layout = valid_layout(); + layout.target_output_vocab = 0; + layout.target_declared_vocab--; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("target.n_vocab") != std::string::npos); + + layout = valid_layout(); + layout.pred_vocab = 8; + layout.succ_vocab = 8; + layout.target_output_vocab = 0; + layout.target_declared_vocab = 0; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("exceeds codebook vocab") != std::string::npos); + + layout = valid_layout(); + layout.succ_rank = 31; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("rank mismatch") != std::string::npos); + + std::printf("dflash2 selector validation: %d checks passed\n", g_checks); + return 0; +} diff --git a/server/test/test_draft_topk_cuda.cpp b/server/test/test_draft_topk_cuda.cpp index a1defe7c7..dd507e6ee 100644 --- a/server/test/test_draft_topk_cuda.cpp +++ b/server/test/test_draft_topk_cuda.cpp @@ -30,6 +30,7 @@ using dflash::common::extract_draft_topk; using dflash::common::geometric_extract_draft_topk_cuda; +using dflash::common::geometric_draft_topk_cuda_supports_k; namespace { @@ -124,6 +125,30 @@ namespace { struct DraftTopkCudaFixture {}; } +TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_dispatch_contract_host_only) { + for (int K = -1; K <= 18; ++K) { + const bool expected = (K >= 1 && K <= 8) || K == 12 || K == 16; + CHECK(geometric_draft_topk_cuda_supports_k(K) == expected); + } + + // Unsupported K must be rejected before pointer inspection or any CUDA + // call. This makes the fallback contract testable on a host with no GPU. + const void * invalid_device_pointer = + reinterpret_cast(uintptr_t{1}); + std::vector log_probs(64, 123.0f); + std::vector token_ids(64, 456); + for (int K : {0, 9, 10, 11, 13, 14, 15, 17, 64}) { + CHECK(!geometric_extract_draft_topk_cuda( + invalid_device_pointer, 1, 128, K, + log_probs.data(), token_ids.data(), 1.0f)); + CHECK(log_probs[0] == 123.0f); + CHECK(token_ids[0] == 456); + } + CHECK(!geometric_extract_draft_topk_cuda( + invalid_device_pointer, 1, 8, 16, + log_probs.data(), token_ids.data(), 1.0f)); +} + TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_suite) { int dev_count = 0; if (cudaGetDeviceCount(&dev_count) != cudaSuccess || dev_count == 0) { @@ -131,8 +156,7 @@ TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_suite) { return; } - // The kernel supports K up to kMaxK (=8 in geometric_draft_topk_cuda.cu); larger K is - // handled by a documented CPU fallback (returns false), checked separately. + // Exercise every instantiated dispatch family, including DFlash 2's K=16. const Case cases[] = { // Realistic decode shape: Qwen3.5 vocab, small position batch. {15, 151936, 8, 1.0f}, @@ -145,6 +169,8 @@ TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_suite) { {3, 257, 8, 1.0f}, // vocab barely above K, non-power-of-two {1, 151936, 1, 1.0f}, // K=1 (argmax + log_z) {15, 151936, 4, 1.0f}, + {3, 4096, 12, 1.0f}, + {3, 4096, 16, 1.0f}, }; int failures = 0; @@ -154,24 +180,27 @@ TEST_CASE(DraftTopkCudaFixture, draft_topk_cuda_suite) { idx++; } - // Fallback contract: K beyond the kernel's supported range must return false - // (not silently produce wrong output) so the caller can use the CPU path. + // Fallback contract: both in-range dispatch holes and K beyond the maximum + // must return false so the caller can use the CPU path. { - const int n = 4, vocab = 4096, big_K = 64; + const int n = 4, vocab = 4096; std::vector h(n * vocab, 0.f); float * d = nullptr; if (cudaMalloc(&d, h.size() * sizeof(float)) == cudaSuccess) { cudaMemcpy(d, h.data(), h.size() * sizeof(float), cudaMemcpyHostToDevice); - std::vector lp(n * big_K); - std::vector ids(n * big_K); - bool ret = geometric_extract_draft_topk_cuda(d, n, vocab, big_K, - lp.data(), ids.data(), 1.0f); + for (int K : {9, 10, 11, 13, 14, 15, 64}) { + std::vector lp((size_t)n * K); + std::vector ids((size_t)n * K); + bool ret = geometric_extract_draft_topk_cuda( + d, n, vocab, K, lp.data(), ids.data(), 1.0f); + const bool pass = !ret; + printf(" [%s] fallback contract: K=%d returned %s\n", + pass ? "PASS" : "FAIL", K, + ret ? "true" : "false"); + if (!pass) failures++; + idx++; + } cudaFree(d); - const bool pass = !ret; // expect false - printf(" [%s] fallback contract: K=%d (>kMaxK) returned %s\n", - pass ? "PASS" : "FAIL", big_K, ret ? "true" : "false"); - if (!pass) failures++; - idx++; } } diff --git a/server/test/test_gdn_transition_journal.cpp b/server/test/test_gdn_transition_journal.cpp new file mode 100644 index 000000000..2748af966 --- /dev/null +++ b/server/test/test_gdn_transition_journal.cpp @@ -0,0 +1,563 @@ +// Standalone GPU proof for the compact linear-chain GDN transition journal. +// The captured [gate, key, state-dependent delta] tuples must reconstruct +// every accepted prefix without rerunning the target recurrence. +#include "ggml-backend.h" +#include "ggml-cuda.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int S = 128; +constexpr int H = 48; +constexpr int KEY_HEADS = 16; +constexpr int T = 6; +constexpr int B = 4; +constexpr int PHYSICAL_SLOTS = 3; +constexpr float FIELD_TOLERANCE = 5.0e-5f; +constexpr float STATE_TOLERANCE = 2.0e-4f; + +size_t qkv_index(int sequence, int token, int head, int value) { + return (((size_t) sequence*T + token)*H + head)*S + value; +} + +size_t key_index(int sequence, int token, int head, int value) { + return (((size_t) sequence*T + token)*KEY_HEADS + + head%KEY_HEADS)*S + value; +} + +size_t scalar_index(int sequence, int token, int head) { + return ((size_t) sequence*T + token)*H + head; +} + +size_t state_index(int slot, int head, int col, int row) { + return (((size_t) slot*H + head)*S + col)*S + row; +} + +size_t journal_index( + int sequence, int token, int head, int width, int value) { + return ((((size_t) sequence*T + token)*H + head)*width) + value; +} + +float sigmoid(float x) { + return 1.0f/(1.0f + std::exp(-x)); +} + +float softplus(float x) { + return x > 20.0f ? x : std::log1p(std::exp(x)); +} + +bool compare_vectors( + const char * label, + const std::vector & actual, + const std::vector & expected, + float tolerance) { + if (actual.size() != expected.size()) { + std::fprintf(stderr, "%s: size mismatch %zu != %zu\n", label, + actual.size(), expected.size()); + return false; + } + float max_error = 0.0f; + size_t worst = 0; + for (size_t i = 0; i < actual.size(); ++i) { + if (!std::isfinite(actual[i]) || !std::isfinite(expected[i])) { + std::fprintf(stderr, + "%s: non-finite value at %zu (actual %.9g expected %.9g)\n", + label, i, actual[i], expected[i]); + return false; + } + const float error = std::fabs(actual[i] - expected[i]); + if (error > max_error) { + max_error = error; + worst = i; + } + } + if (max_error > tolerance || !std::isfinite(max_error)) { + std::fprintf(stderr, + "%s: max error %.9g at %zu (actual %.9g expected %.9g, tolerance %.9g)\n", + label, max_error, worst, actual[worst], expected[worst], + tolerance); + return false; + } + return true; +} + +struct Inputs { + std::vector q; + std::vector k; + std::vector v; + std::vector g; + std::vector beta; + std::vector state; + std::vector dt_bias; + std::vector gate_A; +}; + +Inputs make_inputs(bool kda, bool raw_gates) { + std::mt19937 rng(20260819 + 17*kda + 31*raw_gates); + std::uniform_real_distribution small(-0.25f, 0.25f); + std::uniform_real_distribution state_dist(-0.06f, 0.06f); + std::uniform_real_distribution gate_dist(0.82f, 0.98f); + std::uniform_real_distribution beta_dist(0.15f, 0.85f); + std::uniform_real_distribution raw_dist(-1.5f, 1.5f); + + Inputs in; + const size_t qkv_elements = (size_t) S*H*T*B; + const size_t qk_elements = (size_t) S*KEY_HEADS*T*B; + in.q.resize(qk_elements); + in.k.resize(qk_elements); + in.v.resize(qkv_elements); + in.g.resize((size_t) (kda ? S : 1)*H*T*B); + in.beta.resize((size_t) H*T*B); + in.state.resize((size_t) S*S*H*B); + in.dt_bias.resize(H); + in.gate_A.resize(H); + + for (float & value : in.q) value = small(rng); + for (float & value : in.v) value = small(rng); + for (float & value : in.state) value = state_dist(rng); + + // The target feeds normalized/shared keys to GDN. Normalize every + // sequence/token/head vector before capture so the proof exercises that + // exact resolved input rather than an arbitrary projection. + for (int sequence = 0; sequence < B; ++sequence) { + for (int token = 0; token < T; ++token) { + for (int head = 0; head < KEY_HEADS; ++head) { + float norm2 = 0.0f; + for (int row = 0; row < S; ++row) { + const float value = small(rng); + in.k[key_index(sequence, token, head, row)] = value; + norm2 += value*value; + } + const float inverse_norm = 1.0f/std::sqrt(norm2); + for (int row = 0; row < S; ++row) { + in.k[key_index(sequence, token, head, row)] *= inverse_norm; + } + } + } + } + + if (raw_gates) { + for (float & value : in.g) value = raw_dist(rng); + for (float & value : in.beta) value = raw_dist(rng); + for (int head = 0; head < H; ++head) { + in.dt_bias[head] = -0.35f + 0.12f*head; + in.gate_A[head] = -0.12f - 0.07f*head; + } + } else { + for (float & value : in.g) value = std::log(gate_dist(rng)); + for (float & value : in.beta) value = beta_dist(rng); + } + return in; +} + +float resolved_gate( + const Inputs & in, bool kda, bool raw_gates, + int sequence, int token, int head, int row) { + if (kda) { + return std::exp(in.g[qkv_index(sequence, token, head, row)]); + } + const float raw_or_log = in.g[scalar_index(sequence, token, head)]; + if (!raw_gates) return std::exp(raw_or_log); + return std::exp( + softplus(raw_or_log + in.dt_bias[head])*in.gate_A[head]); +} + +float resolved_beta( + const Inputs & in, bool raw_gates, + int sequence, int token, int head) { + const float value = in.beta[scalar_index(sequence, token, head)]; + return raw_gates ? sigmoid(value) : value; +} + +std::vector ordinary_recurrence( + const Inputs & in, bool kda, bool raw_gates, + int accepted_prefix) { + std::vector state = in.state; + for (int sequence = 0; sequence < B; ++sequence) { + for (int token = 0; token < accepted_prefix; ++token) { + for (int head = 0; head < H; ++head) { + const float beta = resolved_beta( + in, raw_gates, sequence, token, head); + for (int col = 0; col < S; ++col) { + float projection = 0.0f; + for (int row = 0; row < S; ++row) { + const float gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, row); + const float state_value = + state[state_index(sequence, head, col, row)]; + const float key = + in.k[key_index(sequence, token, head, row)]; + projection += (kda ? gate : 1.0f)*state_value*key; + } + const float scalar_gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, 0); + const float delta = + (in.v[qkv_index(sequence, token, head, col)] - + (kda ? projection : scalar_gate*projection))*beta; + for (int row = 0; row < S; ++row) { + const float gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, row); + const float key = + in.k[key_index(sequence, token, head, row)]; + float & state_value = + state[state_index(sequence, head, col, row)]; + state_value = std::fma(key, delta, gate*state_value); + } + } + } + } + } + return state; +} + +std::vector expected_journal( + const Inputs & in, bool kda, bool raw_gates) { + const int gate_values = kda ? S : 1; + const int width = gate_values + 2*S; + std::vector journal((size_t) width*H*T*B); + std::vector state = in.state; + for (int sequence = 0; sequence < B; ++sequence) { + for (int token = 0; token < T; ++token) { + for (int head = 0; head < H; ++head) { + for (int row = 0; row < gate_values; ++row) { + journal[journal_index(sequence, token, head, width, row)] = + resolved_gate(in, kda, raw_gates, + sequence, token, head, row); + } + for (int row = 0; row < S; ++row) { + journal[journal_index( + sequence, token, head, width, + gate_values + row)] = + in.k[key_index(sequence, token, head, row)]; + } + const float beta = resolved_beta( + in, raw_gates, sequence, token, head); + for (int col = 0; col < S; ++col) { + float projection = 0.0f; + for (int row = 0; row < S; ++row) { + const float gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, row); + projection += (kda ? gate : 1.0f)* + state[state_index(sequence, head, col, row)]* + in.k[key_index(sequence, token, head, row)]; + } + const float scalar_gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, 0); + const float delta = + (in.v[qkv_index(sequence, token, head, col)] - + (kda ? projection : scalar_gate*projection))*beta; + journal[journal_index( + sequence, token, head, width, + gate_values + S + col)] = delta; + for (int row = 0; row < S; ++row) { + const float gate = resolved_gate( + in, kda, raw_gates, + sequence, token, head, row); + float & value = + state[state_index(sequence, head, col, row)]; + value = std::fma( + in.k[key_index(sequence, token, head, row)], + delta, gate*value); + } + } + } + } + } + return journal; +} + +struct CaseTensors { + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buffer = nullptr; + ggml_tensor * journal = nullptr; + ggml_tensor * identity_state = nullptr; + ggml_tensor * mapped_state = nullptr; + ggml_tensor * accepted = nullptr; + ggml_tensor * slots = nullptr; +}; + +void destroy(CaseTensors & tensors) { + if (tensors.buffer) ggml_backend_buffer_free(tensors.buffer); + if (tensors.ctx) ggml_free(tensors.ctx); + tensors = {}; +} + +bool run_case( + ggml_backend_t backend, bool kda, bool raw_gates, + bool report_timing) { + const char * name = raw_gates ? "scalar-raw" : kda ? "kda" : "scalar"; + const int gate_values = kda ? S : 1; + const int width = gate_values + 2*S; + const Inputs inputs = make_inputs(kda, raw_gates); + + ggml_init_params params{}; + params.mem_size = 8*1024*1024; + params.no_alloc = true; + CaseTensors tensors; + tensors.ctx = ggml_init(params); + if (!tensors.ctx) return false; + + ggml_tensor * q = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, KEY_HEADS, T, B); + ggml_tensor * k = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, KEY_HEADS, T, B); + ggml_tensor * v = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, H, T, B); + ggml_tensor * g = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, kda ? S : 1, H, T, B); + ggml_tensor * beta = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, 1, H, T, B); + ggml_tensor * capture_state = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, S, H, B); + tensors.journal = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, width, H, T, B); + tensors.identity_state = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, S, H, B); + tensors.mapped_state = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, S, S, H, PHYSICAL_SLOTS); + tensors.accepted = ggml_new_tensor_1d( + tensors.ctx, GGML_TYPE_I32, B); + tensors.slots = ggml_new_tensor_1d( + tensors.ctx, GGML_TYPE_I32, B); + ggml_tensor * dt_bias = nullptr; + ggml_tensor * gate_A = nullptr; + if (raw_gates) { + dt_bias = ggml_new_tensor_1d(tensors.ctx, GGML_TYPE_F32, H); + gate_A = ggml_new_tensor_1d(tensors.ctx, GGML_TYPE_F32, H); + } + + ggml_tensor * result = ggml_gated_delta_net( + tensors.ctx, q, k, v, g, beta, capture_state); + ggml_gated_delta_net_set_skip_intermediate(result, true); + if (raw_gates) { + ggml_gated_delta_net_set_raw_gates(result, dt_bias, gate_A); + } + ggml_gated_delta_net_set_transition_journal(result, tensors.journal); + ggml_set_output(result); + ggml_cgraph * graph = ggml_new_graph(tensors.ctx); + ggml_build_forward_expand(graph, result); + + tensors.buffer = ggml_backend_alloc_ctx_tensors(tensors.ctx, backend); + if (!tensors.buffer) { + std::fprintf(stderr, "%s: GPU tensor allocation failed\n", name); + destroy(tensors); + return false; + } + auto upload_f32 = [](ggml_tensor * tensor, + const std::vector & values) { + ggml_backend_tensor_set(tensor, values.data(), 0, + values.size()*sizeof(float)); + }; + upload_f32(q, inputs.q); + upload_f32(k, inputs.k); + upload_f32(v, inputs.v); + upload_f32(g, inputs.g); + upload_f32(beta, inputs.beta); + upload_f32(capture_state, inputs.state); + if (raw_gates) { + upload_f32(dt_bias, inputs.dt_bias); + upload_f32(gate_A, inputs.gate_A); + } + + bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + std::vector actual_journal((size_t) width*H*T*B); + if (ok) { + ggml_backend_tensor_get( + tensors.journal, actual_journal.data(), 0, + actual_journal.size()*sizeof(float)); + ok = compare_vectors( + name, actual_journal, + expected_journal(inputs, kda, raw_gates), FIELD_TOLERANCE); + } + + const std::vector identity_slots{0, 1, 2, 3}; + std::vector accepted(B); + std::vector actual_state(inputs.state.size()); + for (int prefix = 0; ok && prefix <= T; ++prefix) { + std::fill(accepted.begin(), accepted.end(), prefix); + upload_f32(tensors.identity_state, inputs.state); + ggml_backend_tensor_set(tensors.accepted, accepted.data(), 0, + accepted.size()*sizeof(accepted[0])); + ggml_backend_tensor_set(tensors.slots, identity_slots.data(), 0, + identity_slots.size()*sizeof(identity_slots[0])); + ok = ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.identity_state, + tensors.accepted, tensors.slots); + if (ok) { + ggml_backend_tensor_get( + tensors.identity_state, actual_state.data(), 0, + actual_state.size()*sizeof(float)); + char label[64]; + std::snprintf(label, sizeof(label), "%s prefix %d", name, prefix); + ok = compare_vectors( + label, actual_state, + ordinary_recurrence(inputs, kda, raw_gates, prefix), + STATE_TOLERANCE); + } + } + + // Compact lanes {0,2,3} map to physical slots {2,0,1}; lane 1 is bucket + // padding. Each physical base must match the state used to capture that + // lane's state-dependent deltas. + const std::vector mapped_slots{2, -1, 0, 1}; + const std::vector mapped_prefixes{T, T, 2, 4}; + std::vector mapped_base((size_t) S*S*H*PHYSICAL_SLOTS); + for (int sequence : {0, 2, 3}) { + const int slot = mapped_slots[(size_t) sequence]; + for (int head = 0; head < H; ++head) { + for (int col = 0; col < S; ++col) { + for (int row = 0; row < S; ++row) { + mapped_base[state_index(slot, head, col, row)] = + inputs.state[state_index(sequence, head, col, row)]; + } + } + } + } + std::vector mapped_expected = mapped_base; + for (int sequence : {0, 2, 3}) { + const int slot = mapped_slots[(size_t) sequence]; + const std::vector lane_state = ordinary_recurrence( + inputs, kda, raw_gates, mapped_prefixes[(size_t) sequence]); + for (int head = 0; head < H; ++head) { + for (int col = 0; col < S; ++col) { + for (int row = 0; row < S; ++row) { + mapped_expected[state_index(slot, head, col, row)] = + lane_state[state_index(sequence, head, col, row)]; + } + } + } + } + std::vector mapped_actual(mapped_base.size()); + if (ok) { + upload_f32(tensors.mapped_state, mapped_base); + ggml_backend_tensor_set(tensors.accepted, mapped_prefixes.data(), 0, + mapped_prefixes.size()*sizeof(mapped_prefixes[0])); + ggml_backend_tensor_set(tensors.slots, mapped_slots.data(), 0, + mapped_slots.size()*sizeof(mapped_slots[0])); + ok = ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.mapped_state, + tensors.accepted, tensors.slots); + if (ok) { + ggml_backend_tensor_get( + tensors.mapped_state, mapped_actual.data(), 0, + mapped_actual.size()*sizeof(float)); + ok = compare_vectors( + "permuted/padded slots", mapped_actual, mapped_expected, + STATE_TOLERANCE); + } + } + // Out-of-range ids are padding too. + if (ok) { + const std::vector out_of_range_slots{2, 99, 0, 1}; + upload_f32(tensors.mapped_state, mapped_base); + ggml_backend_tensor_set(tensors.slots, out_of_range_slots.data(), 0, + out_of_range_slots.size()*sizeof(out_of_range_slots[0])); + ok = ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.mapped_state, + tensors.accepted, tensors.slots); + if (ok) { + ggml_backend_tensor_get( + tensors.mapped_state, mapped_actual.data(), 0, + mapped_actual.size()*sizeof(float)); + ok = compare_vectors( + "out-of-range padded slot", mapped_actual, mapped_expected, + STATE_TOLERANCE); + } + } + + // Host validation is transactional: malformed prefixes and duplicate live + // slots are rejected before the state kernel can launch. + if (ok) { + const std::vector unchanged = inputs.state; + std::vector invalid_prefix(B, 1); + invalid_prefix[0] = T + 1; + upload_f32(tensors.identity_state, unchanged); + ggml_backend_tensor_set(tensors.accepted, invalid_prefix.data(), 0, + invalid_prefix.size()*sizeof(invalid_prefix[0])); + ggml_backend_tensor_set(tensors.slots, identity_slots.data(), 0, + identity_slots.size()*sizeof(identity_slots[0])); + ok = !ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.identity_state, + tensors.accepted, tensors.slots); + const std::vector duplicate_slots{0, 0, 2, 3}; + accepted.assign(B, 1); + ggml_backend_tensor_set(tensors.accepted, accepted.data(), 0, + accepted.size()*sizeof(accepted[0])); + ggml_backend_tensor_set(tensors.slots, duplicate_slots.data(), 0, + duplicate_slots.size()*sizeof(duplicate_slots[0])); + ok = ok && !ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.identity_state, + tensors.accepted, tensors.slots); + ggml_backend_tensor_get( + tensors.identity_state, actual_state.data(), 0, + actual_state.size()*sizeof(float)); + ok = ok && compare_vectors( + "transactional validation", actual_state, unchanged, 0.0f); + } + + if (ok && report_timing) { + constexpr int repetitions = 25; + std::vector elapsed_us; + elapsed_us.reserve(repetitions); + accepted.assign(B, T); + ggml_backend_tensor_set(tensors.accepted, accepted.data(), 0, + accepted.size()*sizeof(accepted[0])); + ggml_backend_tensor_set(tensors.slots, identity_slots.data(), 0, + identity_slots.size()*sizeof(identity_slots[0])); + for (int repetition = 0; repetition < repetitions; ++repetition) { + upload_f32(tensors.identity_state, inputs.state); + const auto start = std::chrono::steady_clock::now(); + const bool committed = + ggml_backend_cuda_gdn_transition_journal_commit( + tensors.journal, tensors.identity_state, + tensors.accepted, tensors.slots); + const auto stop = std::chrono::steady_clock::now(); + if (!committed) { + ok = false; + break; + } + elapsed_us.push_back( + std::chrono::duration(stop - start).count()); + } + if (ok) { + std::sort(elapsed_us.begin(), elapsed_us.end()); + std::printf( + "gdn journal commit S=%d H=%d T=%d B=%d median %.1f us (synchronous Phase 1)\n", + S, H, T, B, elapsed_us[elapsed_us.size()/2]); + } + } + + std::printf("gdn transition journal %-10s: %s\n", name, + ok ? "PASS" : "FAIL"); + destroy(tensors); + return ok; +} + +} // namespace + +int main() { + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, "GPU backend unavailable\n"); + return 1; + } + bool ok = run_case(backend, false, false, true); + ok = run_case(backend, true, false, false) && ok; + ok = run_case(backend, false, true, false) && ok; + ggml_backend_free(backend); + return ok ? 0 : 1; +} diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index 59070e6e6..adc6836b1 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -170,6 +170,7 @@ int main() { CHECK(plan.admitted_count == 0); CHECK(plan.ordered.size() == 1); CHECK(plan.ordered[0].decision == SpecDecision::AR); + CHECK(plan.ordered[0].execution_unsupported); CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); CHECK(support.decision(30) == SpecDecision::AR); CHECK(support.initial_confidence(30) == 4.0); @@ -330,6 +331,62 @@ int main() { CHECK(plan.pending_evaluations[0].action == SpecEvaluationAction::Score); + // Interpolate the nonlinear profile across fractional expected rows. A + // tiny score increase across 16.5 rows must preserve the 16 -> 17 cost + // cliff without inheriting lround()'s full 100-us decision discontinuity. + SpecCostSeries cliff_step = series(32, 100.0); + for (size_t i = 16; i < cliff_step.costs.size(); ++i) { + cliff_step.costs[i] = 200.0; + } + const SpecCostTables cliff_costs{ + series(128, 1.0), cliff_step, series(16, 1.0)}; + SpeculationGate cliff_gate(cliff_costs, geometry(), 8); + const SpecPlan below_cliff = cliff_gate.plan(4, { + candidate(800, 0, 4.12475, SpeculationPolicy::Always), + candidate(801, 1, 4.12475, SpeculationPolicy::Always), + candidate(802, 2, 4.12475, SpeculationPolicy::Always), + candidate(803, 3, 4.12475, SpeculationPolicy::Always)}, 4); + const SpecPlan above_cliff = cliff_gate.plan(4, { + candidate(804, 0, 4.12525, SpeculationPolicy::Always), + candidate(805, 1, 4.12525, SpeculationPolicy::Always), + candidate(806, 2, 4.12525, SpeculationPolicy::Always), + candidate(807, 3, 4.12525, SpeculationPolicy::Always)}, 4); + CHECK(std::abs(below_cliff.expected_step_rows - 16.499) < 1e-9); + CHECK(std::abs(above_cliff.expected_step_rows - 16.501) < 1e-9); + CHECK(std::abs(below_cliff.profiled_cost - 151.9) < 1e-9); + CHECK(std::abs(above_cliff.profiled_cost - 152.1) < 1e-9); + CHECK(above_cliff.profiled_cost > below_cliff.profiled_cost); + CHECK(above_cliff.profiled_cost - below_cliff.profiled_cost < 1.0); + + // A realized row-4 sample belongs only to the row-4 executable shape. It + // must not be written under the gate's earlier row-1 expectation. + SpeculationGate shape_feedback( + constant_costs(1.0, 10.0, 1.0), geometry(), 4); + SpecPlan expected_row_one = shape_feedback.plan(1, { + candidate(900, 0, 1.0, SpeculationPolicy::Always)}, 1); + CHECK(expected_row_one.expected_step_rows == 1.0); + CHECK(expected_row_one.cost_scale == 1.0); + shape_feedback.observe_cost({1, 1, 4, 4, 1}, 48.0); + expected_row_one = shape_feedback.plan(1, { + candidate(901, 0, 1.0, SpeculationPolicy::Always)}, 1); + const SpecPlan observed_row_four = shape_feedback.plan(1, { + candidate(902, 0, 4.0, SpeculationPolicy::Always)}, 1); + CHECK(expected_row_one.cost_scale == 1.0); + CHECK(expected_row_one.predicted_cost == 12.0); + CHECK(observed_row_four.cost_scale == 4.0); + CHECK(observed_row_four.predicted_cost == 48.0); + + // One-time confidence drafting on a k=0 round is likewise isolated from + // the pure-AR shape used by the gate's future counterfactual. + SpeculationGate draft_shape_feedback( + constant_costs(100.0, 10.0, 100.0), geometry(), 4); + draft_shape_feedback.observe_cost({1, 0, 0, 1, 1}, 440.0); + const SpecPlan pure_ar_after_bootstrap = draft_shape_feedback.plan( + 1, {candidate(903, 0, 4.0)}, 1); + CHECK(pure_ar_after_bootstrap.admitted_count == 0); + CHECK(pure_ar_after_bootstrap.cost_scale == 1.0); + CHECK(pure_ar_after_bootstrap.predicted_cost == 10.0); + // Shape-local total-cost feedback changes only future undecided choices. // It cannot flip a request whose one-shot decision is already sticky. SpeculationGate cost_feedback( @@ -339,7 +396,7 @@ int main() { CHECK(plan.profiled_cost == 12.0); CHECK(plan.cost_scale == 1.0); CHECK(cost_feedback.decision(950) == SpecDecision::Speculation); - cost_feedback.observe_cost(plan, 48.0); + cost_feedback.observe_cost({1, 1, 4, 4, 1}, 48.0); plan = cost_feedback.plan(1, {candidate(950, 0, NAN)}, 1); CHECK(plan.admitted_count == 1); CHECK(plan.ordered[0].forced); @@ -364,7 +421,7 @@ int main() { SpecPlan ar_plan = ar_feedback.plan(1, {candidate(960, 0, 4.0)}, 1); CHECK(ar_plan.admitted_count == 0); CHECK(ar_feedback.decision(960) == SpecDecision::AR); - ar_feedback.observe_cost(ar_plan, 20.0); + ar_feedback.observe_cost({1, 0, 0, 1, 0}, 20.0); ar_plan = ar_feedback.plan(1, {candidate(960, 0, NAN)}, 1); CHECK(ar_plan.admitted_count == 0); CHECK(ar_plan.profiled_cost == 10.0); diff --git a/server/tests/test_quantize_draft_q8.py b/server/tests/test_quantize_draft_q8.py index 4fd201eda..01158c0ea 100644 --- a/server/tests/test_quantize_draft_q8.py +++ b/server/tests/test_quantize_draft_q8.py @@ -29,6 +29,15 @@ def add_uint32(self, key, value): def add_array(self, key, value): self.calls.append(("array", key, value)) + def add_string(self, key, value): + self.calls.append(("string", key, value)) + + def add_float32(self, key, value): + self.calls.append(("float32", key, value)) + + def add_quantization_version(self, value): + self.calls.append(("quantization_version", value)) + class Qwen36SwaMetadataTest(unittest.TestCase): @staticmethod @@ -108,6 +117,51 @@ def test_converter_cli_writes_profile_only_when_requested(self): self.assertIsNone(window) self.assertIsNone(pattern) + def test_dflash2_tensor_mapping_is_complete(self): + expected = { + "layers.0.attention_conv.base_kernel": "blk.0.attn_conv.base", + "layers.0.attention_conv.kernel_projection.weight": "blk.0.attn_conv.proj.weight", + "layers.0.mlp_conv.base_kernel": "blk.0.ffn_conv.base", + "layers.0.mlp_conv.kernel_projection.weight": "blk.0.ffn_conv.proj.weight", + "candidate_selector.hidden_projection.weight": "dflash.selector.hproj.weight", + "candidate_selector.predecessor_codebook": "dflash.selector.pred_cb", + "candidate_selector.successor_codebook": "dflash.selector.succ_cb", + } + for source, output in expected.items(): + self.assertEqual(MODULE.map_name(source), output) + self.assertTrue(MODULE.is_norm_tensor("blk.0.attn_conv.base")) + + def test_dflash2_metadata_is_emitted_from_resolved_profile(self): + profile = dict( + hidden=32, + n_layer=1, + n_head=2, + n_head_kv=1, + head_dim=16, + intermediate=64, + vocab=64, + n_target_layers=1, + rope_theta=10_000_000.0, + rms_eps=1e-6, + mask_token_id=63, + block_size=8, + ctx_len=4096, + capture_layer_ids=[7], + conv_kernel_size=2, + conv_group_size=16, + selector_rank=32, + selector_top_k=16, + ) + writer = RecordingWriter() + MODULE.add_arch_metadata(writer, profile) + prefix = "qwen35-dflash-draft.dflash." + self.assertIn(("array", prefix + "target_layer_ids", [7]), writer.calls) + self.assertIn(("uint32", prefix + "block_size", 8), writer.calls) + self.assertIn(("uint32", prefix + "dflash2.conv_kernel_size", 2), writer.calls) + self.assertIn(("uint32", prefix + "dflash2.conv_group_size", 16), writer.calls) + self.assertIn(("uint32", prefix + "dflash2.selector_rank", 32), writer.calls) + self.assertIn(("uint32", prefix + "dflash2.selector_top_k", 16), writer.calls) + def test_gguf_round_trip_preserves_types_and_values(self): with tempfile.TemporaryDirectory() as tmp: path = Path(tmp) / "metadata.gguf" From 460c2dc1a398390d82461763be81eba398c5052b Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 20:48:43 +0000 Subject: [PATCH 36/42] refactor(speculation): extract generic activation primitives --- server/CMakeLists.txt | 4 +- .../spec_cost_profile.cpp} | 153 +++++++++++------- .../common/speculation/spec_cost_profile.h | 42 +++++ .../speculation_gate.h | 131 ++++++--------- server/src/common/speculation/speculator.h | 91 +++++++++++ .../src/common/speculation/survival_score.h | 49 ++++++ server/test/test_spec_cost_profile.cpp | 84 ++++++---- server/test/test_speculation_gate.cpp | 107 ++++++------ 8 files changed, 436 insertions(+), 225 deletions(-) rename server/src/common/{concurrency/spec_cost_profile.h => speculation/spec_cost_profile.cpp} (52%) create mode 100644 server/src/common/speculation/spec_cost_profile.h rename server/src/common/{concurrency => speculation}/speculation_gate.h (86%) create mode 100644 server/src/common/speculation/speculator.h create mode 100644 server/src/common/speculation/survival_score.h diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 74a7ae187..a2f425062 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -440,6 +440,7 @@ add_library(dflash_common STATIC src/common/domino_head.cpp src/common/dspark_head.cpp src/common/dflash2_benefit.cpp + src/common/speculation/spec_cost_profile.cpp src/common/dflash2_head.cpp src/common/dflash2_batch.cpp src/common/target_shard_ipc.cpp @@ -1507,7 +1508,8 @@ if(DFLASH27B_TESTS) endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_spec_cost_profile.cpp") add_executable(test_spec_cost_profile - test/test_spec_cost_profile.cpp) + test/test_spec_cost_profile.cpp + src/common/speculation/spec_cost_profile.cpp) target_include_directories(test_spec_cost_profile PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src ${CMAKE_CURRENT_SOURCE_DIR}/test) diff --git a/server/src/common/concurrency/spec_cost_profile.h b/server/src/common/speculation/spec_cost_profile.cpp similarity index 52% rename from server/src/common/concurrency/spec_cost_profile.h rename to server/src/common/speculation/spec_cost_profile.cpp index 19f9d7b1e..8bcc33a82 100644 --- a/server/src/common/concurrency/spec_cost_profile.h +++ b/server/src/common/speculation/spec_cost_profile.cpp @@ -1,38 +1,85 @@ -// Pure startup profiling protocol for monotone speculation cost tables. - -#pragma once - -#include "common/concurrency/speculation_gate.h" +#include "common/speculation/spec_cost_profile.h" #include #include -#include -#include -#include -#include +#include namespace dflash::common { +namespace { -struct SpecProfileGrid { - std::vector tree_rows; - std::vector step_rows; - std::vector draft_lanes; -}; - -inline void sort_unique_positive(std::vector & values) { - values.erase(std::remove_if(values.begin(), values.end(), - [](int value) { return value <= 0; }), - values.end()); +void sort_unique_positive(std::vector & values) { + values.erase( + std::remove_if( + values.begin(), values.end(), + [](int value) { return value <= 0; }), + values.end()); std::sort(values.begin(), values.end()); values.erase(std::unique(values.begin(), values.end()), values.end()); } -inline SpecProfileGrid build_spec_profile_grid( +struct SeriesResult { + SpecCostSeries table; + std::string error; +}; + +SeriesResult profile_monotonic_costs( + std::vector indices, + const SpecCostProfiler::Runner & runner, + int repetitions) { + SeriesResult result; + if (!runner) { + result.error = "profiling runner is missing"; + return result; + } + if (repetitions <= 0) { + result.error = "profiling repetitions must be positive"; + return result; + } + sort_unique_positive(indices); + if (indices.empty()) { + result.error = "profiling grid is empty"; + return result; + } + + result.table.indices = indices; + result.table.costs.reserve(indices.size()); + for (int index : indices) { + (void)runner(index); + std::vector samples; + samples.reserve(static_cast(repetitions)); + for (int rep = 0; rep < repetitions; ++rep) { + const double sample = runner(index); + if (!std::isfinite(sample) || sample <= 0.0) { + result.error = "profiling runner returned an invalid cost"; + result.table = {}; + return result; + } + samples.push_back(sample); + } + std::sort(samples.begin(), samples.end()); + double median = samples[static_cast(repetitions) / 2]; + if (repetitions % 2 == 0) { + median = 0.5 * ( + samples[static_cast(repetitions) / 2 - 1] + + median); + } + if (!result.table.costs.empty()) { + median = std::max(median, result.table.costs.back()); + } + result.table.costs.push_back(median); + } + return result; +} + +} // namespace + +SpecProfileGrid build_spec_profile_grid( int max_concurrency, int tree_width, int max_accept, const std::function & bucket) { SpecProfileGrid grid; - if (max_concurrency <= 0 || tree_width <= 0 || max_accept <= 0) + if (max_concurrency <= 0 || tree_width <= 0 || max_accept <= 0) { return grid; + } auto bucketed = [&](int lanes) { if (lanes <= 0) return 0; return std::max(lanes, bucket ? bucket(lanes) : lanes); @@ -57,50 +104,34 @@ inline SpecProfileGrid build_spec_profile_grid( return grid; } -struct SpecProfileResult { - SpecCostSeries table; - std::string error; - bool ok() const { return error.empty() && table.valid(); } -}; - -template -SpecProfileResult profile_monotonic_costs( - std::vector indices, Runner && runner, int reps = 5) { - SpecProfileResult result; - if (reps <= 0) { - result.error = "profiling repetitions must be positive"; +SpecCostProfileResult SpecCostProfiler::profile( + const SpecProfileGrid & grid, + Runner tree_runner, + Runner step_runner, + Runner draft_runner, + std::string speculator_id, + int repetitions) const { + SpecCostProfileResult result; + if (speculator_id.empty()) { + result.error = "speculator id is empty"; return result; } - sort_unique_positive(indices); - if (indices.empty()) { - result.error = "profiling grid is empty"; + + SeriesResult tree = profile_monotonic_costs( + grid.tree_rows, tree_runner, repetitions); + SeriesResult step = profile_monotonic_costs( + grid.step_rows, step_runner, repetitions); + SeriesResult draft = profile_monotonic_costs( + grid.draft_lanes, draft_runner, repetitions); + if (!tree.error.empty() || !step.error.empty() || !draft.error.empty()) { + result.error = tree.error + step.error + draft.error; return result; } - result.table.indices = indices; - result.table.costs.reserve(indices.size()); - for (int index : indices) { - (void)runner(index); // graph capture / allocator warmup - std::vector samples; - samples.reserve(static_cast(reps)); - for (int rep = 0; rep < reps; ++rep) { - const double sample = runner(index); - if (!std::isfinite(sample) || sample <= 0.0) { - result.error = "profiling runner returned an invalid cost"; - result.table = {}; - return result; - } - samples.push_back(sample); - } - std::sort(samples.begin(), samples.end()); - double median = samples[(size_t)reps / 2]; - if (reps % 2 == 0) { - median = 0.5 * (samples[(size_t)reps / 2 - 1] + median); - } - if (!result.table.costs.empty()) { - median = std::max(median, result.table.costs.back()); - } - result.table.costs.push_back(median); - } + + result.tables.tree_cost = std::move(tree.table); + result.tables.step_cost = std::move(step.table); + result.tables.draft_cost = std::move(draft.table); + result.tables.speculator_id = std::move(speculator_id); return result; } diff --git a/server/src/common/speculation/spec_cost_profile.h b/server/src/common/speculation/spec_cost_profile.h new file mode 100644 index 000000000..1e9333200 --- /dev/null +++ b/server/src/common/speculation/spec_cost_profile.h @@ -0,0 +1,42 @@ +// Generic startup profiling protocol for monotone speculation cost tables. +#pragma once + +#include "common/speculation/speculation_gate.h" + +#include +#include +#include + +namespace dflash::common { + +struct SpecProfileGrid { + std::vector tree_rows; + std::vector step_rows; + std::vector draft_lanes; +}; + +SpecProfileGrid build_spec_profile_grid( + int max_concurrency, int tree_width, int max_accept, + const std::function & bucket); + +struct SpecCostProfileResult { + SpecCostTables tables; + std::string error; + + bool ok() const { return error.empty() && tables.valid(); } +}; + +class SpecCostProfiler { +public: + using Runner = std::function; + + SpecCostProfileResult profile( + const SpecProfileGrid & grid, + Runner tree_runner, + Runner step_runner, + Runner draft_runner, + std::string speculator_id, + int repetitions = 5) const; +}; + +} // namespace dflash::common diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/speculation/speculation_gate.h similarity index 86% rename from server/src/common/concurrency/speculation_gate.h rename to server/src/common/speculation/speculation_gate.h index e0eb51c56..3efd879a0 100644 --- a/server/src/common/concurrency/speculation_gate.h +++ b/server/src/common/speculation/speculation_gate.h @@ -21,9 +21,6 @@ namespace dflash::common { struct SpecGateConfig { - // Immutable offline fit applied independently to every request's one-time - // activation score. It never learns from request execution history. - double fixed_yield_scale = 1.0; double cost_ema_alpha = 0.20; double adaptive_gain_margin = 0.02; }; @@ -76,27 +73,14 @@ struct SpecCostTables { SpecCostSeries tree_cost; SpecCostSeries step_cost; SpecCostSeries draft_cost; + std::string speculator_id; bool valid() const { return tree_cost.valid() && step_cost.valid() && draft_cost.valid(); } }; -enum class SpecScoreKind : uint8_t { - Unspecified, - DSparkConfidence, - DFlash2SelectorBenefitV1, -}; - -inline const char * spec_score_kind_name(SpecScoreKind kind) { - switch (kind) { - case SpecScoreKind::Unspecified: return "unspecified"; - case SpecScoreKind::DSparkConfidence: return "dspark_confidence"; - case SpecScoreKind::DFlash2SelectorBenefitV1: - return "dflash2_selector_benefit_v1"; - } - return "unknown"; -} +inline constexpr const char * kUnspecifiedScoreKind = "unspecified"; struct SpecCandidate { uint64_t request_id = 0; @@ -114,10 +98,11 @@ struct SpecCandidate { // bootstrap and a finite value is the preferred activation measurement. // Evaluation failure explicitly falls back to sticky AR without inventing // a score. Otherwise the gate commits exactly one mode from this - // survival-product expected yield, including the root; the gate applies - // only its fixed offline scale and clamping. - double confidence_yield = std::numeric_limits::quiet_NaN(); - SpecScoreKind score_kind = SpecScoreKind::Unspecified; + // adapter-provided expected yield, including the root. Adapter estimates + // are already calibrated; the gate only clamps executor bounds. + double activation_yield = std::numeric_limits::quiet_NaN(); + std::vector conditional_hazards; + std::string score_kind = kUnspecifiedScoreKind; }; struct SpecStepGeometry { @@ -141,15 +126,15 @@ struct SpecStepGeometry { }; enum class SpecScoreSource : uint8_t { - Confidence, - InitialConfidence, + Fresh, + Initial, Unavailable, }; inline const char * spec_score_source_name(SpecScoreSource source) { switch (source) { - case SpecScoreSource::Confidence: return "confidence"; - case SpecScoreSource::InitialConfidence: return "initial"; + case SpecScoreSource::Fresh: return "fresh"; + case SpecScoreSource::Initial: return "initial"; case SpecScoreSource::Unavailable: return "unavailable"; } return "unknown"; @@ -193,7 +178,7 @@ struct SpecPlanScore { // AR/speculation decision. This supports one activation record per // request without treating later sticky execution as a new decision. bool newly_decided = false; - SpecScoreKind score_kind = SpecScoreKind::Unspecified; + std::string score_kind = kUnspecifiedScoreKind; bool execution_unsupported = false; }; @@ -230,7 +215,7 @@ struct SpecPlan { double profiled_cost = 0.0; double cost_scale = 1.0; double predicted_cost = 0.0; - // Fixed-scale expected yield for admitted confidence-scored lanes. This is + // Fixed-scale expected yield for admitted activation-scored lanes. This is // directly comparable with realized emitted tokens in telemetry. double initial_predicted_tokens = 0.0; double goodput = 0.0; @@ -252,33 +237,15 @@ struct SpecPlan { std::vector pending_evaluations; }; -// DSpark confidence-head contract: `confidences[i]` is a probability-like, -// monotone-in-acceptance score for position i conditioned on its prefix. -// DFlash2 selector evidence must go through its model-specific benefit adapter -// and must never be passed to this helper as though it were confidence. -inline double confidence_survival_yield( - const std::vector & confidences, int max_accept) { - if (max_accept <= 1) return 1.0; - double expected = 1.0; - double survival = 1.0; - const int depth = std::min( - static_cast(confidences.size()), max_accept - 1); - for (int i = 0; i < depth; ++i) { - const double c = std::clamp(confidences[(size_t)i], 0.0, 1.0); - survival *= c; - expected += survival; - } - return std::clamp(expected, 1.0, static_cast(max_accept)); -} - class SpeculationGate { private: struct RequestState { - double initial_confidence = + double initial_score = std::numeric_limits::quiet_NaN(); SpecDecision decision = SpecDecision::Undecided; - bool confidence_evaluation_failed = false; - SpecScoreKind score_kind = SpecScoreKind::Unspecified; + bool evaluation_failed = false; + std::string score_kind = kUnspecifiedScoreKind; + std::vector conditional_hazards; }; struct ExecutionShapeHash { @@ -311,7 +278,7 @@ class SpeculationGate { struct CandidateScore { double expected_yield = 1.0; SpecScoreSource source = SpecScoreSource::Unavailable; - SpecScoreKind score_kind = SpecScoreKind::Unspecified; + std::string score_kind = kUnspecifiedScoreKind; }; public: @@ -337,9 +304,6 @@ class SpeculationGate { return std::isfinite(value) && value > 0.0 && value <= 1.0; }; return valid_alpha(config_.cost_ema_alpha) && - std::isfinite(config_.fixed_yield_scale) && - config_.fixed_yield_scale >= kFixedYieldScaleMin && - config_.fixed_yield_scale <= kFixedYieldScaleMax && std::isfinite(config_.adaptive_gain_margin) && config_.adaptive_gain_margin >= 0.0 && costs_.valid() && geometry_.tree_width >= 1 && @@ -364,7 +328,7 @@ class SpeculationGate { const SpecCandidate * candidate = nullptr; double score = 1.0; SpecScoreSource source = SpecScoreSource::Unavailable; - SpecScoreKind score_kind = SpecScoreKind::Unspecified; + std::string score_kind = kUnspecifiedScoreKind; SpecDecision decision = SpecDecision::Undecided; bool forced = false; bool commit_candidate = false; @@ -598,9 +562,9 @@ class SpeculationGate { bool commit_evaluation_fallback_ar(uint64_t request_id) { RequestState & state = request_states_[request_id]; if (state.decision != SpecDecision::Undecided) return false; - state.initial_confidence = + state.initial_score = std::numeric_limits::quiet_NaN(); - state.confidence_evaluation_failed = true; + state.evaluation_failed = true; state.decision = SpecDecision::AR; return true; } @@ -612,71 +576,72 @@ class SpeculationGate { bool has_state(uint64_t request_id) const { return request_states_.find(request_id) != request_states_.end(); } - bool has_confidence(uint64_t request_id) const { + bool has_score(uint64_t request_id) const { auto state = request_states_.find(request_id); return state != request_states_.end() && - std::isfinite(state->second.initial_confidence); + std::isfinite(state->second.initial_score); } - bool confidence_evaluation_failed(uint64_t request_id) const { + bool evaluation_failed(uint64_t request_id) const { auto state = request_states_.find(request_id); return state != request_states_.end() && - state->second.confidence_evaluation_failed; + state->second.evaluation_failed; } - double initial_confidence(uint64_t request_id) const { + double initial_score(uint64_t request_id) const { auto state = request_states_.find(request_id); return state == request_states_.end() ? std::numeric_limits::quiet_NaN() - : state->second.initial_confidence; + : state->second.initial_score; + } + std::string initial_score_kind(uint64_t request_id) const { + auto state = request_states_.find(request_id); + return state == request_states_.end() + ? kUnspecifiedScoreKind : state->second.score_kind; } - SpecScoreKind initial_score_kind(uint64_t request_id) const { + const std::vector & initial_hazards(uint64_t request_id) const { + static const std::vector empty; auto state = request_states_.find(request_id); return state == request_states_.end() - ? SpecScoreKind::Unspecified : state->second.score_kind; + ? empty : state->second.conditional_hazards; } SpecDecision decision(uint64_t request_id) const { auto state = request_states_.find(request_id); return state == request_states_.end() ? SpecDecision::Undecided : state->second.decision; } - double fixed_yield_scale() const { return config_.fixed_yield_scale; } const SpecCostTables & costs() const { return costs_; } private: CandidateScore score_candidate(const SpecCandidate & candidate) { bool accepted_initial_score = false; - if (std::isfinite(candidate.confidence_yield)) { + if (std::isfinite(candidate.activation_yield)) { const double raw = std::clamp( - candidate.confidence_yield, 1.0, + candidate.activation_yield, 1.0, static_cast(max_accept_)); RequestState & state = request_states_[candidate.request_id]; - if (!std::isfinite(state.initial_confidence)) { - state.initial_confidence = raw; + if (!std::isfinite(state.initial_score)) { + state.initial_score = raw; state.score_kind = candidate.score_kind; + state.conditional_hazards = candidate.conditional_hazards; accepted_initial_score = true; } return { - std::clamp( - config_.fixed_yield_scale * state.initial_confidence, - 1.0, static_cast(max_accept_)), - accepted_initial_score ? SpecScoreSource::Confidence - : SpecScoreSource::InitialConfidence, + state.initial_score, + accepted_initial_score ? SpecScoreSource::Fresh + : SpecScoreSource::Initial, state.score_kind, }; } auto state = request_states_.find(candidate.request_id); if (state != request_states_.end() && - std::isfinite(state->second.initial_confidence)) { + std::isfinite(state->second.initial_score)) { return { - std::clamp( - config_.fixed_yield_scale * - state->second.initial_confidence, - 1.0, static_cast(max_accept_)), - SpecScoreSource::InitialConfidence, + state->second.initial_score, + SpecScoreSource::Initial, state->second.score_kind, }; } return {1.0, SpecScoreSource::Unavailable, - SpecScoreKind::Unspecified}; + kUnspecifiedScoreKind}; } static void update_ema(double & value, uint64_t & observations, @@ -757,8 +722,6 @@ class SpeculationGate { SpecCostTables costs_; SpecStepGeometry geometry_; int max_accept_ = 1; - static constexpr double kFixedYieldScaleMin = 0.25; - static constexpr double kFixedYieldScaleMax = 4.0; static constexpr double kCostScaleMin = 0.25; static constexpr double kCostScaleMax = 4.0; std::unordered_map request_states_; diff --git a/server/src/common/speculation/speculator.h b/server/src/common/speculation/speculator.h new file mode 100644 index 000000000..6a6d47fc4 --- /dev/null +++ b/server/src/common/speculation/speculator.h @@ -0,0 +1,91 @@ +// Model-agnostic speculation adapter contract. +#pragma once + +#include +#include +#include +#include + +namespace dflash::common { + +enum SpeculatorInputRequirement : uint32_t { + SpeculatorInputNone = 0, + SpeculatorInputHidden = 1u << 0, + SpeculatorInputPrenorm = 1u << 1, +}; + +struct ActivationEstimate { + double expected_yield = std::numeric_limits::quiet_NaN(); + std::vector conditional_hazards; +}; + +struct SpeculatorBatchInput { + int lane_count = 0; + int requested_depth = 0; + std::vector hidden_by_lane; + std::vector prenorm_by_lane; + std::vector seed_tokens; +}; + +struct SpecProposal { + std::vector tokens; + ActivationEstimate estimate; + std::string error; + // Optional per-depth JSON field fragments used only by debug telemetry. + std::vector debug_depth_fields; +}; + +inline bool speculator_input_satisfies( + const SpeculatorBatchInput & input, uint32_t requirements) { + if (input.lane_count <= 0 || input.requested_depth < 2 || + static_cast(input.seed_tokens.size()) != input.lane_count) { + return false; + } + auto has_lanes = [&](const std::vector & lanes) { + if (static_cast(lanes.size()) != input.lane_count) return false; + for (const float * lane : lanes) { + if (!lane) return false; + } + return true; + }; + if ((requirements & SpeculatorInputHidden) != 0 && + !has_lanes(input.hidden_by_lane)) { + return false; + } + if ((requirements & SpeculatorInputPrenorm) != 0 && + !has_lanes(input.prenorm_by_lane)) { + return false; + } + return true; +} + +class Speculator { +public: + virtual ~Speculator() = default; + + // Opaque, versioned identity. The activation engine never enumerates it. + virtual const std::string & score_kind() const = 0; + virtual int max_block_size() const = 0; + virtual uint32_t input_requirements() const = 0; + virtual bool ready() const = 0; + virtual const std::string & error() const = 0; + + // Draft and score each lane in one adapter call. A false return means the + // whole batch failed; lane-local failures use SpecProposal::error. + virtual bool propose(const SpeculatorBatchInput & input, + std::vector & output) = 0; +}; + +inline constexpr const char * kNoSpeculatorAdapterReason = + "no_speculator_adapter"; + +inline bool speculator_is_ready(const Speculator * speculator) { + return speculator != nullptr && speculator->ready(); +} + +inline const char * speculator_fallback_reason(const Speculator * speculator) { + return speculator_is_ready(speculator) + ? nullptr : kNoSpeculatorAdapterReason; +} + +} // namespace dflash::common diff --git a/server/src/common/speculation/survival_score.h b/server/src/common/speculation/survival_score.h new file mode 100644 index 000000000..4d1edf64e --- /dev/null +++ b/server/src/common/speculation/survival_score.h @@ -0,0 +1,49 @@ +// Reusable scoring contract for conditional acceptance hazards. +#pragma once + +#include "common/speculation/speculator.h" + +#include +#include + +namespace dflash::common { + +// Each hazard is a probability-like, monotone-in-acceptance score for one +// position conditioned on accepting its prefix. The yield includes the root. +inline double hazard_survival_yield( + const std::vector & hazards, int max_accept) { + if (max_accept <= 1) return 1.0; + double expected = 1.0; + double survival = 1.0; + const int depth = std::min( + static_cast(hazards.size()), max_accept - 1); + for (int i = 0; i < depth; ++i) { + const double hazard = + std::clamp(hazards[static_cast(i)], 0.0, 1.0); + survival *= hazard; + expected += survival; + } + return std::clamp(expected, 1.0, static_cast(max_accept)); +} + +class ConfidenceVectorScorer { +public: + ActivationEstimate score( + const std::vector & confidences, int max_accept) const { + ActivationEstimate estimate; + const int depth = std::max( + 0, std::min( + static_cast(confidences.size()), max_accept - 1)); + estimate.conditional_hazards.reserve(static_cast(depth)); + for (int i = 0; i < depth; ++i) { + estimate.conditional_hazards.push_back( + std::clamp( + confidences[static_cast(i)], 0.0, 1.0)); + } + estimate.expected_yield = hazard_survival_yield( + estimate.conditional_hazards, max_accept); + return estimate; + } +}; + +} // namespace dflash::common diff --git a/server/test/test_spec_cost_profile.cpp b/server/test/test_spec_cost_profile.cpp index 6a2e4e87d..d65b31df9 100644 --- a/server/test/test_spec_cost_profile.cpp +++ b/server/test/test_spec_cost_profile.cpp @@ -1,4 +1,4 @@ -#include "common/concurrency/spec_cost_profile.h" +#include "common/speculation/spec_cost_profile.h" #include "host_check.h" #include @@ -34,46 +34,62 @@ int main() { CHECK(bucketed.step_rows.back() == 35); CHECK(build_spec_profile_grid(0, 16, 4, {}).tree_rows.empty()); - std::unordered_map calls; - const std::unordered_map intended{{1, 10.0}, {2, 5.0}}; + SpecProfileGrid small; + small.tree_rows = {2, 1, 2}; + small.step_rows = {4}; + small.draft_lanes = {3}; + std::unordered_map tree_calls; const double noise[] = {-2.0, 1.0, 0.0, 2.0, -1.0}; - SpecProfileResult profiled = profile_monotonic_costs( - {2, 1, 2}, [&](int index) { - const int call = calls[index]++; - if (call == 0) return 10000.0; // discarded warmup - return intended.at(index) + noise[(call - 1) % 5]; - }); - CHECK(profiled.ok()); - CHECK((profiled.table.indices == std::vector{1, 2})); - CHECK(profiled.table.costs.size() == 2); - CHECK(profiled.table.costs[0] == 10.0); - CHECK(profiled.table.costs[1] == 10.0); // monotone clamp - CHECK(calls[1] == 6 && calls[2] == 6); - - int even_calls = 0; - profiled = profile_monotonic_costs({4}, [&](int) { - ++even_calls; - static const double samples[] = {999.0, 1.0, 2.0, 3.0, 4.0}; - return samples[(even_calls - 1) % 5]; - }, 4); + auto tree_runner = [&](int index) { + const int call = tree_calls[index]++; + if (call == 0) return 10000.0; + const double intended = index == 1 ? 10.0 : 5.0; + return intended + noise[(call - 1) % 5]; + }; + int step_calls = 0; + int draft_calls = 0; + const SpecCostProfileResult profiled = SpecCostProfiler{}.profile( + small, + tree_runner, + [&](int index) { + ++step_calls; + return index == 4 ? 20.0 : 0.0; + }, + [&](int index) { + ++draft_calls; + return index == 3 ? 30.0 : 0.0; + }, + "adapter-score-v1"); CHECK(profiled.ok()); - CHECK(profiled.table.costs[0] == 2.5); - CHECK(even_calls == 5); + CHECK(profiled.tables.speculator_id == "adapter-score-v1"); + CHECK((profiled.tables.tree_cost.indices == std::vector{1, 2})); + CHECK(profiled.tables.tree_cost.costs[0] == 10.0); + CHECK(profiled.tables.tree_cost.costs[1] == 10.0); + CHECK(tree_calls[1] == 6 && tree_calls[2] == 6); + CHECK(step_calls == 6); + CHECK(draft_calls == 6); - SpecProfileResult bad = profile_monotonic_costs( - {}, [](int) { return 1.0; }); + SpecCostProfileResult bad = SpecCostProfiler{}.profile( + {}, [](int) { return 1.0; }, [](int) { return 1.0; }, + [](int) { return 1.0; }, "adapter-score-v1"); CHECK(!bad.ok() && !bad.error.empty()); - bad = profile_monotonic_costs( - {1}, [](int) { return 1.0; }, 0); + bad = SpecCostProfiler{}.profile( + small, [](int) { return 1.0; }, [](int) { return 1.0; }, + [](int) { return 1.0; }, "", 5); CHECK(!bad.ok()); int invalid_calls = 0; - bad = profile_monotonic_costs({1}, [&](int) { - ++invalid_calls; - return invalid_calls == 2 - ? std::numeric_limits::quiet_NaN() : 1.0; - }); + bad = SpecCostProfiler{}.profile( + {{1}, {1}, {1}}, + [&](int) { + ++invalid_calls; + return invalid_calls == 2 + ? std::numeric_limits::quiet_NaN() : 1.0; + }, + [](int) { return 1.0; }, + [](int) { return 1.0; }, + "adapter-score-v1"); CHECK(!bad.ok()); - CHECK(bad.table.indices.empty()); + CHECK(bad.tables.tree_cost.indices.empty()); std::printf("spec cost profile tests passed: %d checks\n", g_checks); return 0; diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index adc6836b1..11c913b56 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -1,4 +1,6 @@ -#include "common/concurrency/speculation_gate.h" +#include "common/speculation/speculation_gate.h" +#include "common/speculation/speculator.h" +#include "common/speculation/survival_score.h" #include "host_check.h" #include @@ -31,16 +33,26 @@ static SpecStepGeometry geometry() { } static SpecCandidate candidate( - uint64_t id, int slot, double confidence, + uint64_t id, int slot, double activation_score, SpeculationPolicy policy = SpeculationPolicy::Adaptive, - bool scoreable = true, bool can_speculate = true) { - return {id, slot, policy, scoreable, can_speculate, confidence}; + bool scoreable = true, bool can_speculate = true, + std::vector hazards = {}, + std::string score_kind = "test-score-v1") { + return {id, slot, policy, scoreable, can_speculate, activation_score, + std::move(hazards), std::move(score_kind)}; } int main() { - CHECK(std::abs(confidence_survival_yield({0.5f, 0.5f}, 4) - 1.75) < 1e-9); - CHECK(confidence_survival_yield({2.0f, -1.0f}, 4) == 2.0); - CHECK(confidence_survival_yield({}, 4) == 1.0); + // A drafter without a registered adapter (including a DSpark-only GGUF) + // is not scoreable and reports the generic sticky-AR fallback reason. + CHECK(!speculator_is_ready(nullptr)); + CHECK(std::string(speculator_fallback_reason(nullptr)) == + "no_speculator_adapter"); + const ConfidenceVectorScorer confidence_scorer; + CHECK(std::abs(confidence_scorer.score({0.5f, 0.5f}, 4) + .expected_yield - 1.75) < 1e-9); + CHECK(confidence_scorer.score({2.0f, -1.0f}, 4).expected_yield == 2.0); + CHECK(confidence_scorer.score({}, 4).expected_yield == 1.0); CHECK(std::string(spec_decision_name(SpecDecision::Undecided)) == "undecided"); CHECK(std::string(spec_decision_name(SpecDecision::AR)) == "ar"); @@ -65,7 +77,7 @@ int main() { CHECK(plan.predicted_cost == plan.profiled_cost); CHECK(costly.decision(1) == SpecDecision::AR); CHECK(costly.decision(2) == SpecDecision::AR); - CHECK(costly.initial_confidence(1) == 4.0); + CHECK(costly.initial_score(1) == 4.0); plan = costly.plan(2, { candidate(1, 0, NAN), candidate(2, 1, NAN)}, 2); CHECK(plan.pending_evaluations.empty()); @@ -87,9 +99,9 @@ int main() { CHECK(plan.ordered[0].admitted); CHECK(plan.ordered[0].decision == SpecDecision::Speculation); CHECK(plan.ordered[1].decision == SpecDecision::AR); - CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); + CHECK(plan.ordered[0].source == SpecScoreSource::Fresh); CHECK(std::string(spec_score_source_name(plan.ordered[0].source)) == - "confidence"); + "fresh"); CHECK(prefix.decision(10) == SpecDecision::Speculation); CHECK(prefix.decision(11) == SpecDecision::AR); plan = prefix.plan(3, { @@ -98,9 +110,9 @@ int main() { CHECK((plan.admitted_request_ids == std::vector{10})); CHECK(plan.ordered.size() == 1); CHECK(plan.ordered[0].forced); - CHECK(plan.ordered[0].source == SpecScoreSource::InitialConfidence); - CHECK(prefix.initial_confidence(10) == 4.0); - CHECK(prefix.initial_confidence(11) == 1.0); + CHECK(plan.ordered[0].source == SpecScoreSource::Initial); + CHECK(prefix.initial_score(10) == 4.0); + CHECK(prefix.initial_score(11) == 1.0); // A cold batch is atomic: every scoreable undecided lane without a score // is returned for bootstrap, and no scored-but-undecided peer commits until @@ -121,7 +133,7 @@ int main() { SpecEvaluationAction::Score); CHECK(bootstrap.decision(20) == SpecDecision::Undecided); CHECK(bootstrap.decision(21) == SpecDecision::Undecided); - CHECK(bootstrap.initial_confidence(21) == 4.0); + CHECK(bootstrap.initial_score(21) == 4.0); plan = bootstrap.plan(2, { candidate(20, 0, 1.0), candidate(21, 1, NAN)}, 2); @@ -132,8 +144,8 @@ int main() { CHECK((plan.admitted_request_ids == std::vector{21})); CHECK(bootstrap.decision(20) == SpecDecision::AR); CHECK(bootstrap.decision(21) == SpecDecision::Speculation); - CHECK(bootstrap.initial_confidence(20) == 1.0); - CHECK(bootstrap.initial_confidence(21) == 4.0); + CHECK(bootstrap.initial_score(20) == 1.0); + CHECK(bootstrap.initial_score(21) == 4.0); plan = bootstrap.plan(2, { candidate(20, 0, 4.0), candidate(21, 1, 1.0)}, 2); @@ -144,12 +156,14 @@ int main() { CHECK(plan.ordered.size() == 1); CHECK(plan.ordered[0].forced); CHECK(plan.ordered[0].decision == SpecDecision::Speculation); - CHECK(bootstrap.initial_confidence(20) == 1.0); - CHECK(bootstrap.initial_confidence(21) == 4.0); + CHECK(bootstrap.initial_score(20) == 1.0); + CHECK(bootstrap.initial_score(21) == 4.0); + // The gate side of the no-adapter contract emits FallbackAR and commits it + // once; later finite values cannot reopen the request. // Request-lifetime scoreability is separate from permanent executor - // support. Unsupported requests still bootstrap and retain a confidence, - // then commit directly to AR. A request that cannot evaluate confidence + // support. Unsupported requests still bootstrap and retain an activation score, + // then commit directly to AR. A request that cannot evaluate an activation score // receives an explicit failed-evaluation action and sticky AR with no // synthetic score. SpeculationGate support(crossover, geometry(), 4); @@ -171,9 +185,9 @@ int main() { CHECK(plan.ordered.size() == 1); CHECK(plan.ordered[0].decision == SpecDecision::AR); CHECK(plan.ordered[0].execution_unsupported); - CHECK(plan.ordered[0].source == SpecScoreSource::Confidence); + CHECK(plan.ordered[0].source == SpecScoreSource::Fresh); CHECK(support.decision(30) == SpecDecision::AR); - CHECK(support.initial_confidence(30) == 4.0); + CHECK(support.initial_score(30) == 4.0); plan = support.plan(1, { candidate(31, 0, NAN, SpeculationPolicy::Adaptive, false, false)}, 1); CHECK(plan.valid); @@ -187,9 +201,9 @@ int main() { CHECK(support.commit_evaluation_fallback_ar(31)); CHECK(!support.commit_evaluation_fallback_ar(31)); CHECK(support.decision(31) == SpecDecision::AR); - CHECK(support.confidence_evaluation_failed(31)); - CHECK(!support.has_confidence(31)); - CHECK(std::isnan(support.initial_confidence(31))); + CHECK(support.evaluation_failed(31)); + CHECK(!support.has_score(31)); + CHECK(std::isnan(support.initial_score(31))); plan = support.plan(1, {candidate(31, 0, 4.0)}, 1); CHECK(plan.valid); CHECK(plan.decisions_committed); @@ -220,9 +234,9 @@ int main() { CHECK(plan.decisions_committed); CHECK(plan.pending_evaluations.empty()); CHECK(mixed_activation.decision(32) == SpecDecision::AR); - CHECK(mixed_activation.confidence_evaluation_failed(32)); + CHECK(mixed_activation.evaluation_failed(32)); CHECK(mixed_activation.decision(33) == SpecDecision::Speculation); - CHECK(mixed_activation.initial_confidence(33) == 4.0); + CHECK(mixed_activation.initial_score(33) == 4.0); CHECK((plan.admitted_request_ids == std::vector{33})); // Explicit Always/Never are configured execution policies, not adaptive @@ -292,37 +306,40 @@ int main() { CHECK(plan.cost_lookup_clamped); CHECK(clamp_logs > 0); - // Every request is ranked only from its own immutable first score and the - // fixed offline scale. Re-presenting a different score cannot change the - // request, and one request never supplies a prior for another. - SpecGateConfig fitted_config; - fitted_config.fixed_yield_scale = 0.5; + // Every request is ranked only from its own immutable, already-calibrated + // adapter score. Re-presenting a different score cannot change the request, + // and one request never supplies a prior for another. SpeculationGate fitted( - fitted_config, constant_costs(1.0, 10.0, 1.0), geometry(), 4); - plan = fitted.plan(1, {candidate(700, 0, 4.0)}, 1); - CHECK(fitted.fixed_yield_scale() == 0.5); - CHECK(fitted.initial_confidence(700) == 4.0); + constant_costs(1.0, 10.0, 1.0), geometry(), 4); + plan = fitted.plan(1, { + candidate( + 700, 0, 2.0, SpeculationPolicy::Adaptive, true, true, + {0.5, 0.25}, "adapter-score-v1")}, 1); + CHECK(fitted.initial_score(700) == 2.0); + CHECK(fitted.initial_score_kind(700) == "adapter-score-v1"); + CHECK((fitted.initial_hazards(700) == + std::vector{0.5, 0.25})); CHECK(plan.ordered[0].expected_yield == 2.0); CHECK(plan.initial_predicted_tokens == 2.0); CHECK(plan.ordered[0].newly_decided); plan = fitted.plan(1, {candidate(700, 0, 1.0)}, 1); CHECK(plan.ordered[0].forced); - CHECK(plan.ordered[0].source == SpecScoreSource::InitialConfidence); + CHECK(plan.ordered[0].source == SpecScoreSource::Initial); CHECK(plan.ordered[0].expected_yield == 2.0); - CHECK(fitted.initial_confidence(700) == 4.0); + CHECK(fitted.initial_score(700) == 2.0); plan = fitted.plan(1, {candidate(701, 0, 4.0)}, 1); - CHECK(plan.ordered[0].expected_yield == 2.0); - CHECK(fitted.initial_confidence(701) == 4.0); - CHECK(fitted.initial_confidence(700) == 4.0); + CHECK(plan.ordered[0].expected_yield == 4.0); + CHECK(fitted.initial_score(701) == 4.0); + CHECK(fitted.initial_score(700) == 2.0); // forget() removes only this request's activation state. A repeated ID is // cold again; fixed configuration and shape-cost feedback are independent. fitted.forget(700); CHECK(!fitted.has_state(700)); - CHECK(!fitted.has_confidence(700)); - CHECK(std::isnan(fitted.initial_confidence(700))); + CHECK(!fitted.has_score(700)); + CHECK(std::isnan(fitted.initial_score(700))); CHECK(fitted.decision(700) == SpecDecision::Undecided); plan = fitted.plan(1, {candidate(700, 0, NAN)}, 1); CHECK(!plan.decisions_committed); @@ -376,7 +393,7 @@ int main() { CHECK(observed_row_four.cost_scale == 4.0); CHECK(observed_row_four.predicted_cost == 48.0); - // One-time confidence drafting on a k=0 round is likewise isolated from + // One-time activation drafting on a k=0 round is likewise isolated from // the pure-AR shape used by the gate's future counterfactual. SpeculationGate draft_shape_feedback( constant_costs(100.0, 10.0, 100.0), geometry(), 4); @@ -440,7 +457,7 @@ int main() { plan = margin_gate.plan(1, {candidate(970, 0, 4.0)}, 1); CHECK(plan.admitted_count == 0); CHECK(plan.ordered.empty()); - CHECK(margin_gate.initial_confidence(970) == 1.04); + CHECK(margin_gate.initial_score(970) == 1.04); SpecGateConfig zero_margin; zero_margin.adaptive_gain_margin = 0.0; From a34aca2391db6eb26e8415639a4a809a6de4fe5a Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 20:48:58 +0000 Subject: [PATCH 37/42] feat(speculation): add speculator interface and DFlash2 adapter --- server/CMakeLists.txt | 1 + server/src/common/dflash2_benefit.cpp | 13 ++- server/src/common/dflash2_benefit.h | 2 + .../adapters/dflash2_speculator.cpp | 102 ++++++++++++++++++ .../speculation/adapters/dflash2_speculator.h | 42 ++++++++ server/test/test_dflash2_benefit.cpp | 9 +- 6 files changed, 163 insertions(+), 6 deletions(-) create mode 100644 server/src/common/speculation/adapters/dflash2_speculator.cpp create mode 100644 server/src/common/speculation/adapters/dflash2_speculator.h diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index a2f425062..06739ed43 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -441,6 +441,7 @@ add_library(dflash_common STATIC src/common/dspark_head.cpp src/common/dflash2_benefit.cpp src/common/speculation/spec_cost_profile.cpp + src/common/speculation/adapters/dflash2_speculator.cpp src/common/dflash2_head.cpp src/common/dflash2_batch.cpp src/common/target_shard_ipc.cpp diff --git a/server/src/common/dflash2_benefit.cpp b/server/src/common/dflash2_benefit.cpp index d2d5873b8..98dc138ca 100644 --- a/server/src/common/dflash2_benefit.cpp +++ b/server/src/common/dflash2_benefit.cpp @@ -103,6 +103,12 @@ DFlash2BenefitConfig DFlash2BenefitProvider::config_from_environment( if (error.empty()) { error = "DFLASH_DFLASH2_BENEFIT_HAZARD_SCALE must be in (0,1]"; } + return config; + } + if (!parse_finite_env( + "DFLASH_DFLASH2_BENEFIT_YIELD_SCALE", 0.25, 4.0, + config.yield_scale, error)) { + return config; } return config; } @@ -118,7 +124,9 @@ DFlash2BenefitProvider::DFlash2BenefitProvider( if (!std::isfinite(config_.lm_log_weight) || config_.lm_log_weight < 0.0 || config_.lm_log_weight > 1.0 || !std::isfinite(config_.hazard_scale) || - config_.hazard_scale <= 0.0 || config_.hazard_scale > 1.0) { + config_.hazard_scale <= 0.0 || config_.hazard_scale > 1.0 || + !std::isfinite(config_.yield_scale) || + config_.yield_scale < 0.25 || config_.yield_scale > 4.0) { error_ = "invalid DFlash2 benefit coefficients"; return; } @@ -177,7 +185,8 @@ bool DFlash2BenefitProvider::estimate( expected += survival; } out.expected_yield = std::clamp( - expected, 1.0, static_cast(max_accept)); + config_.yield_scale * expected, 1.0, + static_cast(max_accept)); if (error) error->clear(); return true; } diff --git a/server/src/common/dflash2_benefit.h b/server/src/common/dflash2_benefit.h index 3e899cb81..99b8b3d92 100644 --- a/server/src/common/dflash2_benefit.h +++ b/server/src/common/dflash2_benefit.h @@ -47,6 +47,8 @@ struct DFlash2BenefitConfig { // probability. hazard_scale may only lower the estimate. double lm_log_weight = 0.10; double hazard_scale = 1.0; + // Offline per-adapter calibration; the generic gate never rescales yield. + double yield_scale = 1.0; }; struct DFlash2BenefitEstimate { diff --git a/server/src/common/speculation/adapters/dflash2_speculator.cpp b/server/src/common/speculation/adapters/dflash2_speculator.cpp new file mode 100644 index 000000000..d7c7709b3 --- /dev/null +++ b/server/src/common/speculation/adapters/dflash2_speculator.cpp @@ -0,0 +1,102 @@ +#include "common/speculation/adapters/dflash2_speculator.h" + +#include "common/dflash2_head.h" + +#include +#include +#include + +namespace dflash::common { +namespace { + +std::string depth_debug_fields(const DFlash2DepthSignal & signal) { + std::ostringstream out; + out << std::setprecision(9) + << "\"selected_logp\":" << signal.selected_log_prob + << ",\"lm_margin\":" << signal.lm_top2_margin + << ",\"topk_mass\":" << signal.top_k_mass + << ",\"rank\":" << signal.selected_rank + << ",\"lm_top1\":" + << (signal.agrees_with_lm_top1 ? "true" : "false") + << ",\"selector_margin\":" << signal.selector_margin + << ",\"selector_mass\":" << signal.selector_winner_mass + << ",\"selector_entropy\":" << signal.selector_entropy; + return out.str(); +} + +} // namespace + +DFlash2Speculator::DFlash2Speculator( + const DraftWeights & weights, + ggml_backend_t backend, + ggml_tensor * lm_head, + DFlash2BenefitModelSignature signature, + DFlash2BenefitConfig config) + : weights_(weights), backend_(backend), lm_head_(lm_head), + benefit_(std::move(signature), std::move(config)), + score_kind_(benefit_.score_kind()) { + const DraftSelectorWeights & selector = weights_.selector; + if (!backend_ || !lm_head_ || !selector.enabled || !selector.hproj || + !selector.pred_cb || !selector.succ_cb || selector.rank <= 0 || + selector.top_k <= 0 || weights_.block_size <= 1) { + error_ = "DFlash2 selector inputs are unavailable"; + } else if (!benefit_.ready()) { + error_ = benefit_.error(); + } +} + +int DFlash2Speculator::max_block_size() const { + return weights_.block_size; +} + +bool DFlash2Speculator::propose( + const SpeculatorBatchInput & input, + std::vector & output) { + output.clear(); + if (!ready() || + !speculator_input_satisfies(input, input_requirements()) || + input.requested_depth > max_block_size()) { + return false; + } + + std::vector> draft_tokens; + std::vector traces; + if (!dflash2_select_chains_batched( + weights_, backend_, lm_head_, input.hidden_by_lane, + input.requested_depth, input.seed_tokens, + draft_tokens, &traces) || + static_cast(draft_tokens.size()) != input.lane_count || + static_cast(traces.size()) != input.lane_count) { + return false; + } + + output.resize(static_cast(input.lane_count)); + for (int lane = 0; lane < input.lane_count; ++lane) { + SpecProposal & proposal = output[static_cast(lane)]; + proposal.tokens = std::move(draft_tokens[static_cast(lane)]); + + DFlash2BenefitEstimate estimate; + std::string estimate_error; + if (!benefit_.estimate( + traces[static_cast(lane)], + input.requested_depth, estimate, &estimate_error)) { + proposal.error = estimate_error.empty() + ? "DFlash2 activation estimate failed" + : std::move(estimate_error); + continue; + } + proposal.estimate.expected_yield = estimate.expected_yield; + proposal.estimate.conditional_hazards = + std::move(estimate.conditional_hazards); + proposal.debug_depth_fields.reserve( + traces[static_cast(lane)].depths.size()); + for (const DFlash2DepthSignal & signal : + traces[static_cast(lane)].depths) { + proposal.debug_depth_fields.push_back( + depth_debug_fields(signal)); + } + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/speculation/adapters/dflash2_speculator.h b/server/src/common/speculation/adapters/dflash2_speculator.h new file mode 100644 index 000000000..d9ba6dcec --- /dev/null +++ b/server/src/common/speculation/adapters/dflash2_speculator.h @@ -0,0 +1,42 @@ +#pragma once + +#include "common/dflash2_benefit.h" +#include "common/speculation/speculator.h" +#include "internal.h" + +#include "ggml-backend.h" + +#include + +namespace dflash::common { + +class DFlash2Speculator final : public Speculator { +public: + DFlash2Speculator( + const DraftWeights & weights, + ggml_backend_t backend, + ggml_tensor * lm_head, + DFlash2BenefitModelSignature signature, + DFlash2BenefitConfig config = {}); + + const std::string & score_kind() const override { return score_kind_; } + int max_block_size() const override; + uint32_t input_requirements() const override { + return SpeculatorInputHidden; + } + bool ready() const override { return error_.empty(); } + const std::string & error() const override { return error_; } + + bool propose(const SpeculatorBatchInput & input, + std::vector & output) override; + +private: + const DraftWeights & weights_; + ggml_backend_t backend_ = nullptr; + ggml_tensor * lm_head_ = nullptr; + DFlash2BenefitProvider benefit_; + std::string score_kind_; + std::string error_; +}; + +} // namespace dflash::common diff --git a/server/test/test_dflash2_benefit.cpp b/server/test/test_dflash2_benefit.cpp index 23023205c..4b7e12bc2 100644 --- a/server/test/test_dflash2_benefit.cpp +++ b/server/test/test_dflash2_benefit.cpp @@ -1,5 +1,5 @@ #include "common/dflash2_benefit.h" -#include "common/concurrency/speculation_gate.h" +#include "common/speculation/speculation_gate.h" #include "host_check.h" #include @@ -48,6 +48,7 @@ int main() { "qwen38-dflash2-selector-benefit-v1"); CHECK(provider.config().lm_log_weight == 0.10); CHECK(provider.config().hazard_scale == 1.0); + CHECK(provider.config().yield_scale == 1.0); // Retained C1 first-block traces for he08 code and prose. The adapter is // continuous and content-agnostic: code has the higher expected yield, while @@ -99,8 +100,8 @@ int main() { SpeculationPolicy policy = SpeculationPolicy::Adaptive) { return SpecCandidate{ - id, slot, policy, true, true, score, - SpecScoreKind::DFlash2SelectorBenefitV1}; + id, slot, policy, true, true, score, {}, + kDFlash2BenefitAdapterVersion}; }; SpeculationGate code_c1(observed_costs, observed_geometry, 8); @@ -109,7 +110,7 @@ int main() { CHECK(gate_plan.admitted_count == 1); CHECK(code_c1.decision(100) == SpecDecision::Speculation); CHECK(code_c1.initial_score_kind(100) == - SpecScoreKind::DFlash2SelectorBenefitV1); + kDFlash2BenefitAdapterVersion); SpeculationGate prose_c1(observed_costs, observed_geometry, 8); gate_plan = prose_c1.plan( From 6f1486a127af02eb6334e08e897fcb7c969ba964 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 20:49:16 +0000 Subject: [PATCH 38/42] refactor(speculation): remove concurrent DSpark drafting --- .../common/concurrency/chain_spec_shapes.h | 2 +- server/src/common/dflash_draft_kv.cpp | 206 ++--- server/src/common/dflash_draft_kv.h | 21 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 708 +++++++----------- .../qwen35/concurrency/qwen35_seq_engine.h | 24 +- server/test/test_chain_spec_shapes.cpp | 4 +- 6 files changed, 368 insertions(+), 597 deletions(-) diff --git a/server/src/common/concurrency/chain_spec_shapes.h b/server/src/common/concurrency/chain_spec_shapes.h index b283ef966..380babb61 100644 --- a/server/src/common/concurrency/chain_spec_shapes.h +++ b/server/src/common/concurrency/chain_spec_shapes.h @@ -23,7 +23,7 @@ inline int chain_decode_bucket_width(int lanes) { // draft_tokens[0] is the already-pending root; positions 1.. form the // proposal. DDTree's flat indices then coincide with chain depth. -inline DDTree make_dspark_chain_tree( +inline DDTree make_chain_verify_tree( const std::vector & draft_tokens) { DDTree tree; if (draft_tokens.size() <= 1) return tree; diff --git a/server/src/common/dflash_draft_kv.cpp b/server/src/common/dflash_draft_kv.cpp index b4366de84..1ff4cb531 100644 --- a/server/src/common/dflash_draft_kv.cpp +++ b/server/src/common/dflash_draft_kv.cpp @@ -1,5 +1,4 @@ #include "dflash_draft_kv.h" -#include "dspark_head.h" #include #include @@ -332,33 +331,28 @@ void draft_kv_batch_free(DraftKvBatchGraph & batch) { batch.g_ctx = nullptr; } batch.gf = nullptr; - batch.seed_tokens = nullptr; batch.hidden_by_lane.clear(); - batch.token_depths.clear(); - batch.confidence_depths.clear(); + batch.prenorm_by_lane.clear(); batch.lane_states.clear(); batch.meta_arena.clear(); batch.n_lanes = 0; batch.q_len = 0; - batch.has_confidence = false; - batch.uses_dflash2 = false; + batch.outputs_prenorm = false; batch.built_for = nullptr; - batch.built_lm_head = nullptr; } static bool draft_kv_batch_build( DraftKvBatchGraph & batch, const DraftWeights & dw, ggml_backend_t backend, - ggml_tensor * lm_head, - const std::vector & lane_states) { - if (!backend || !lm_head || lane_states.empty() || - dw.block_size <= 1) { + const std::vector & lane_states, + bool need_prenorm) { + if (!backend || lane_states.empty() || dw.block_size <= 1) { return false; } for (DraftKvState * state : lane_states) { if (!state || !state->mem_buf || state->q_len != dw.block_size || - state->built_for != (const void *)&dw) { + state->built_for != static_cast(&dw)) { return false; } } @@ -366,27 +360,24 @@ static bool draft_kv_batch_build( draft_kv_batch_free(batch); const int n_lanes = static_cast(lane_states.size()); const size_t arena_size = - (32u + 16u * (size_t)n_lanes) * 1024u * 1024u; + (32u + 16u * static_cast(n_lanes)) * 1024u * 1024u; batch.meta_arena.resize(arena_size); - ggml_init_params gp{}; - gp.mem_size = batch.meta_arena.size(); - gp.mem_buffer = batch.meta_arena.data(); - gp.no_alloc = true; - batch.g_ctx = ggml_init(gp); + ggml_init_params params{}; + params.mem_size = batch.meta_arena.size(); + params.mem_buffer = batch.meta_arena.data(); + params.no_alloc = true; + batch.g_ctx = ggml_init(params); if (!batch.g_ctx) { draft_kv_batch_free(batch); return false; } batch.gf = ggml_new_graph_custom( batch.g_ctx, 4096 * n_lanes + 2048, false); - batch.seed_tokens = - ggml_new_tensor_1d(batch.g_ctx, GGML_TYPE_I32, n_lanes); - ggml_set_input(batch.seed_tokens); - - std::vector hidden; - std::vector prenorm; - hidden.reserve((size_t)n_lanes); - prenorm.reserve((size_t)n_lanes); + + batch.hidden_by_lane.reserve(static_cast(n_lanes)); + if (need_prenorm) { + batch.prenorm_by_lane.reserve(static_cast(n_lanes)); + } for (DraftKvState * state : lane_states) { DraftKvAppendInputs append{}; append.n_rows = state->a_step; @@ -405,61 +396,42 @@ static bool draft_kv_batch_build( step.noise_rows = state->noise_rows; step.mask_full = state->mask_full; step.mask_swa = state->mask_swa; - DraftGraphOutputs outputs = build_draft_kv_step( + DraftGraphOutputs output = build_draft_kv_step( batch.g_ctx, batch.gf, dw, state->cache, step); - if (!outputs.hidden_prenorm || !outputs.hidden_states) { + if (!output.hidden_states || + (need_prenorm && !output.hidden_prenorm)) { draft_kv_batch_free(batch); return false; } - hidden.push_back(outputs.hidden_states); - prenorm.push_back(outputs.hidden_prenorm); - } - - const bool uses_dflash2 = dw.selector.enabled; - DSparkBatchedChainOutputs chain; - if (uses_dflash2) { - for (ggml_tensor * lane_hidden : hidden) { - ggml_set_output(lane_hidden); - ggml_build_forward_expand(batch.gf, lane_hidden); + ggml_set_output(output.hidden_states); + ggml_build_forward_expand(batch.gf, output.hidden_states); + batch.hidden_by_lane.push_back(output.hidden_states); + if (need_prenorm) { + ggml_set_output(output.hidden_prenorm); + ggml_build_forward_expand(batch.gf, output.hidden_prenorm); + batch.prenorm_by_lane.push_back(output.hidden_prenorm); } - } else if (!build_dspark_markov_batched_chain( - batch.g_ctx, batch.gf, dw, lm_head, hidden, prenorm, - batch.seed_tokens, dw.block_size, true, chain) || - chain.n_lanes != n_lanes || - static_cast(chain.tokens.size()) != dw.block_size - 1) { - draft_kv_batch_free(batch); - return false; } batch.galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); if (!batch.galloc || !ggml_gallocr_alloc_graph(batch.galloc, batch.gf)) { - std::fprintf(stderr, "[draft-kv-batch] graph alloc failed lanes=%d\n", - n_lanes); + std::fprintf(stderr, + "[draft-kv-batch] graph alloc failed lanes=%d\n", n_lanes); draft_kv_batch_free(batch); return false; } batch.n_lanes = n_lanes; batch.q_len = dw.block_size; - batch.uses_dflash2 = uses_dflash2; - batch.has_confidence = - !uses_dflash2 && !chain.confidence.empty() && - chain.confidence[0] != nullptr; + batch.outputs_prenorm = need_prenorm; batch.built_for = &dw; - batch.built_lm_head = lm_head; batch.lane_states = lane_states; - batch.hidden_by_lane = uses_dflash2 ? hidden - : std::vector{}; - batch.token_depths = std::move(chain.tokens); - batch.confidence_depths = std::move(chain.confidence); std::fprintf(stderr, - "[draft-kv-batch] packed graph ready lanes=%d q_len=%d " - "head=%s confidence=%s\n", - n_lanes, dw.block_size, - uses_dflash2 ? "dflash2" : "dspark", - batch.has_confidence ? "on" : "off"); + "[draft-kv-batch] packed backbone ready lanes=%d q_len=%d " + "prenorm=%s\n", + n_lanes, dw.block_size, need_prenorm ? "on" : "off"); return true; } @@ -467,34 +439,23 @@ bool draft_kv_batch_compute( DraftKvBatchGraph & batch, const DraftWeights & dw, ggml_backend_t backend, - ggml_tensor * lm_head, const std::vector & lane_states, - const std::vector & seed_tokens, - std::vector> & draft_tokens, - std::vector> & confidences, - std::vector * selector_traces) { - draft_tokens.clear(); - confidences.clear(); - if (selector_traces) selector_traces->clear(); - if (lane_states.empty() || - seed_tokens.size() != lane_states.size()) { - return false; - } + bool need_prenorm, + std::vector> & hidden_by_lane, + std::vector> & prenorm_by_lane) { + hidden_by_lane.clear(); + prenorm_by_lane.clear(); + if (lane_states.empty()) return false; + const bool reusable = - batch.gf && batch.built_for == (const void *)&dw && - batch.built_lm_head == lm_head && - batch.lane_states == lane_states; + batch.gf && batch.built_for == static_cast(&dw) && + batch.lane_states == lane_states && + batch.outputs_prenorm == need_prenorm; if (!reusable && !draft_kv_batch_build( - batch, dw, backend, lm_head, lane_states)) { + batch, dw, backend, lane_states, need_prenorm)) { return false; } - - if (!batch.uses_dflash2) { - ggml_backend_tensor_set( - batch.seed_tokens, seed_tokens.data(), 0, - sizeof(int32_t) * seed_tokens.size()); - } if (ggml_backend_graph_compute(backend, batch.gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, @@ -503,74 +464,29 @@ bool draft_kv_batch_compute( return false; } - if (batch.uses_dflash2) { - if (static_cast(batch.hidden_by_lane.size()) != - batch.n_lanes) { - return false; - } - const int hidden = dw.n_embd; - std::vector> hidden_host( - (size_t) batch.n_lanes, - std::vector( - (size_t) hidden * (size_t) batch.q_len)); - for (int lane = 0; lane < batch.n_lanes; ++lane) { - ggml_backend_tensor_get_async( - backend, batch.hidden_by_lane[(size_t) lane], - hidden_host[(size_t) lane].data(), 0, - sizeof(float) * hidden_host[(size_t) lane].size()); - } - ggml_backend_synchronize(backend); - std::vector hidden_ptrs((size_t) batch.n_lanes); - for (int lane = 0; lane < batch.n_lanes; ++lane) { - hidden_ptrs[(size_t) lane] = - hidden_host[(size_t) lane].data(); - } - confidences.assign((size_t) batch.n_lanes, {}); - return dflash2_select_chains_batched( - dw, backend, lm_head, hidden_ptrs, batch.q_len, - seed_tokens, draft_tokens, selector_traces); + const size_t elements = + static_cast(dw.n_embd) * static_cast(batch.q_len); + hidden_by_lane.assign( + static_cast(batch.n_lanes), + std::vector(elements)); + if (need_prenorm) { + prenorm_by_lane.assign( + static_cast(batch.n_lanes), + std::vector(elements)); } - const int depths = batch.q_len - 1; - std::vector depth_tokens( - (size_t)depths * batch.n_lanes); - std::vector depth_confidence( - (size_t)depths * batch.n_lanes); - for (int depth = 0; depth < depths; ++depth) { + for (int lane = 0; lane < batch.n_lanes; ++lane) { ggml_backend_tensor_get_async( - backend, batch.token_depths[(size_t)depth], - depth_tokens.data() + (size_t)depth * batch.n_lanes, - 0, sizeof(int32_t) * (size_t)batch.n_lanes); - if (batch.has_confidence && - batch.confidence_depths[(size_t)depth]) { + backend, batch.hidden_by_lane[static_cast(lane)], + hidden_by_lane[static_cast(lane)].data(), 0, + sizeof(float) * elements); + if (need_prenorm) { ggml_backend_tensor_get_async( - backend, batch.confidence_depths[(size_t)depth], - depth_confidence.data() + - (size_t)depth * batch.n_lanes, - 0, sizeof(float) * (size_t)batch.n_lanes); + backend, batch.prenorm_by_lane[static_cast(lane)], + prenorm_by_lane[static_cast(lane)].data(), 0, + sizeof(float) * elements); } } ggml_backend_synchronize(backend); - - draft_tokens.assign( - (size_t)batch.n_lanes, - std::vector((size_t)batch.q_len)); - confidences.assign((size_t)batch.n_lanes, {}); - for (int lane = 0; lane < batch.n_lanes; ++lane) { - draft_tokens[(size_t)lane][0] = - seed_tokens[(size_t)lane]; - if (batch.has_confidence) { - confidences[(size_t)lane].resize((size_t)depths); - } - for (int depth = 0; depth < depths; ++depth) { - draft_tokens[(size_t)lane][(size_t)depth + 1] = - depth_tokens[(size_t)depth * batch.n_lanes + lane]; - if (batch.has_confidence) { - confidences[(size_t)lane][(size_t)depth] = - depth_confidence[ - (size_t)depth * batch.n_lanes + lane]; - } - } - } return true; } diff --git a/server/src/common/dflash_draft_kv.h b/server/src/common/dflash_draft_kv.h index 58270c35c..fb9c46f43 100644 --- a/server/src/common/dflash_draft_kv.h +++ b/server/src/common/dflash_draft_kv.h @@ -19,7 +19,6 @@ #pragma once #include "dflash_feature_ring.h" -#include "dflash2_head.h" #include "draft/draft_graph.h" #include "internal.h" // DraftWeights @@ -108,36 +107,30 @@ struct DraftKvBatchGraph { int n_lanes = 0; int q_len = 0; - bool has_confidence = false; - bool uses_dflash2 = false; + bool outputs_prenorm = false; const void * built_for = nullptr; - ggml_tensor * built_lm_head = nullptr; std::vector lane_states; std::vector meta_arena; ggml_context * g_ctx = nullptr; ggml_cgraph * gf = nullptr; ggml_gallocr_t galloc = nullptr; - ggml_tensor * seed_tokens = nullptr; std::vector hidden_by_lane; - std::vector token_depths; - std::vector confidence_depths; + std::vector prenorm_by_lane; }; void draft_kv_batch_free(DraftKvBatchGraph & batch); // All lane states must already have draft_kv_begin_step() inputs and -// inp_embed uploaded. The graph is rebuilt only when the ordered state cohort -// changes, then replays as one backend compute and one synchronization. +// inp_embed uploaded. The packed graph only computes the shared backbone; +// adapters consume the returned per-lane host views. bool draft_kv_batch_compute( DraftKvBatchGraph & batch, const DraftWeights & dw, ggml_backend_t backend, - ggml_tensor * lm_head, const std::vector & lane_states, - const std::vector & seed_tokens, - std::vector> & draft_tokens, - std::vector> & confidences, - std::vector * selector_traces = nullptr); + bool need_prenorm, + std::vector> & hidden_by_lane, + std::vector> & prenorm_by_lane); } // namespace dflash::common diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 6afce5052..2433f14e3 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -14,9 +14,9 @@ #include "common/sampler.h" #include "common/ddtree.h" #include "common/geometric_draft_topk_cuda.h" -#include "common/dspark_head.h" +#include "common/speculation/adapters/dflash2_speculator.h" #include "common/concurrency/chain_spec_shapes.h" -#include "common/concurrency/spec_cost_profile.h" +#include "common/speculation/spec_cost_profile.h" #include "internal.h" #include @@ -61,15 +61,17 @@ double initial_prediction_realized_tokens( return realized; } -void log_spec_gate_plan(const SpecPlan & plan, double fixed_yield_scale, - double initial_realized_tokens, double measured_us) { - int current = 0; +void log_spec_gate_plan( + const SpecPlan & plan, + double initial_realized_tokens, + double measured_us) { + int fresh = 0; int initial = 0; int unavailable = plan.unavailable_count; for (const SpecPlanScore & score : plan.ordered) { switch (score.source) { - case SpecScoreSource::Confidence: ++current; break; - case SpecScoreSource::InitialConfidence: ++initial; break; + case SpecScoreSource::Fresh: ++fresh; break; + case SpecScoreSource::Initial: ++initial; break; case SpecScoreSource::Unavailable: ++unavailable; break; } } @@ -80,10 +82,10 @@ void log_spec_gate_plan(const SpecPlan & plan, double fixed_yield_scale, const SpecPlanScore & score = plan.ordered[i]; std::fprintf(stderr, "%s%llu:%.3f/%s/%s%s", i == 0 ? "" : ",", - (unsigned long long)score.request_id, + static_cast(score.request_id), score.expected_yield, spec_score_source_name(score.source), - spec_score_kind_name(score.score_kind), + score.score_kind.c_str(), score.admitted ? "*" : ""); } std::fprintf(stderr, "] decisions=["); @@ -91,14 +93,13 @@ void log_spec_gate_plan(const SpecPlan & plan, double fixed_yield_scale, const SpecPlanScore & score = plan.ordered[i]; std::fprintf(stderr, "%s%llu:%s", i == 0 ? "" : ",", - (unsigned long long)score.request_id, + static_cast(score.request_id), spec_decision_name(score.decision)); } std::fprintf(stderr, - "] sources=current:%d,initial:%d,unavailable:%d " - "fixed_yield_scale=%.3f initial_tokens=%.3f/", - current, initial, unavailable, fixed_yield_scale, - plan.initial_predicted_tokens); + "] sources=fresh:%d,initial:%d,unavailable:%d " + "initial_tokens=%.3f/", + fresh, initial, unavailable, plan.initial_predicted_tokens); if (std::isfinite(initial_realized_tokens)) { std::fprintf(stderr, "%.3f", initial_realized_tokens); } else { @@ -116,55 +117,54 @@ void log_spec_gate_plan(const SpecPlan & plan, double fixed_yield_scale, } } -void log_spec_activations(const SpecPlan & plan, - const SpeculationGate & gate) { +void log_spec_activations( + const SpecPlan & plan, + const SpeculationGate & gate) { for (const SpecPlanScore & score : plan.ordered) { if (!score.newly_decided) continue; - const double initial = gate.initial_confidence(score.request_id); - const SpecScoreKind kind = gate.initial_score_kind(score.request_id); + const double initial = gate.initial_score(score.request_id); + const std::string kind = gate.initial_score_kind(score.request_id); + const std::vector & hazards = + gate.initial_hazards(score.request_id); const char * decision_reason = score.execution_unsupported ? "execution_unsupported" : score.decision == SpecDecision::Speculation ? "selected_by_joint_goodput" : "ar_counterfactual_won"; std::fprintf(stderr, - "[spec-activation] {\"request_id\":%llu,\"slot\":%d,", - (unsigned long long)score.request_id, score.slot); - if (kind == SpecScoreKind::DSparkConfidence) { - // Legacy DSpark field retained for harness compatibility. - std::fprintf(stderr, - "\"initial_confidence\":%.6f,", initial); - } else { - // Selector evidence is adapted to request benefit; it is not a - // trained confidence value. - std::fprintf(stderr, "\"initial_confidence\":null,"); + "[spec-activation] {\"request_id\":%llu,\"slot\":%d," + "\"activation_score\":%.6f,\"score_kind\":\"%s\"," + "\"expected_yield\":%.6f,\"hazards\":[", + static_cast(score.request_id), + score.slot, initial, kind.c_str(), score.expected_yield); + for (size_t i = 0; i < hazards.size(); ++i) { + std::fprintf(stderr, "%s%.8g", i == 0 ? "" : ",", hazards[i]); } std::fprintf(stderr, - "\"activation_score\":%.6f," - "\"request_benefit\":%.6f," - "\"score_kind\":\"%s\",\"expected_yield\":%.6f," - "\"evaluation\":\"scored\",\"fallback_reason\":null," - "\"decision_reason\":\"%s\"," - "\"decision\":\"%s\"}\n", - initial, initial, spec_score_kind_name(kind), - score.expected_yield, decision_reason, - spec_decision_name(score.decision)); + "],\"evaluation\":\"scored\",\"fallback_reason\":null," + "\"decision_reason\":\"%s\",\"decision\":\"%s\"}\n", + decision_reason, spec_decision_name(score.decision)); } } -void log_spec_evaluation_fallback(uint64_t request_id, int slot, - SpecScoreKind kind, - const char * reason) { +void log_spec_evaluation_fallback( + uint64_t request_id, + int slot, + const std::string & kind, + const char * reason) { + const std::string cause = reason + ? reason : "activation_evaluation_failed"; + const char * decision_reason = cause == "activation_evaluation_failed" + ? "evaluation_failed" : cause.c_str(); std::fprintf(stderr, "[spec-activation] {\"request_id\":%llu,\"slot\":%d," - "\"initial_confidence\":null,\"activation_score\":null," - "\"request_benefit\":null,\"score_kind\":\"%s\"," - "\"expected_yield\":null,\"evaluation\":\"failed\"," - "\"fallback_reason\":\"%s\"," - "\"decision_reason\":\"evaluation_failed\"," - "\"decision\":\"ar\"}\n", - (unsigned long long)request_id, slot, spec_score_kind_name(kind), - reason ? reason : "activation_evaluation_failed"); + "\"activation_score\":null,\"score_kind\":\"%s\"," + "\"expected_yield\":null,\"hazards\":null," + "\"evaluation\":\"failed\"," + "\"fallback_reason\":\"activation_evaluation_failed\"," + "\"decision_reason\":\"%s\",\"decision\":\"ar\"}\n", + static_cast(request_id), slot, kind.c_str(), + decision_reason); } uint64_t file_size_or_zero(const char * path) { @@ -229,8 +229,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( const int n_slots = slots_.slot_count(); slot_draft_kv_.resize((size_t)n_slots); prepared_chain_drafts_.resize((size_t)n_slots); - last_survival_score_.assign( - (size_t)n_slots, std::numeric_limits::quiet_NaN()); + last_activation_estimate_.resize((size_t)n_slots); adaptive_fallback_ar_.assign((size_t)n_slots, 0); if (spec_mode_ == SpecMode::chain && b_.dw_.selector.enabled) { @@ -253,30 +252,37 @@ Qwen35SeqEngine::Qwen35SeqEngine( std::string config_error; DFlash2BenefitConfig config = DFlash2BenefitProvider::config_from_environment(config_error); - if (!config_error.empty()) { - adaptive_fallback_reason_ = "benefit_adapter_invalid_config"; + if (config_error.empty()) { + speculator_ = std::make_unique( + b_.dw_, b_.draft_backend_, b_.w_.output, + signature, config); + } + if (!config_error.empty() || !speculator_is_ready(speculator_.get())) { + const std::string error = !config_error.empty() + ? config_error + : speculator_ ? speculator_->error() + : "adapter construction failed"; + speculator_.reset(); + adaptive_fallback_reason_ = + speculator_fallback_reason(speculator_.get()); std::fprintf(stderr, - "[parallel-chain] DFlash2 benefit adapter disabled: %s; " + "[parallel-chain] no speculator adapter: %s; " "adaptive requests will use sticky AR\n", - config_error.c_str()); + error.c_str()); } else { - dflash2_benefit_provider_ = - std::make_unique(signature, config); - if (!dflash2_benefit_provider_->ready()) { - adaptive_fallback_reason_ = "benefit_adapter_unavailable"; - std::fprintf(stderr, - "[parallel-chain] DFlash2 benefit adapter disabled: %s; " - "adaptive requests will use sticky AR\n", - dflash2_benefit_provider_->error().c_str()); - } else { - std::fprintf(stderr, - "[parallel-chain] DFlash2 request-benefit adapter=%s " - "lm_weight=%.3f hazard_scale=%.3f signature=%s\n", - dflash2_benefit_provider_->score_kind(), - config.lm_log_weight, config.hazard_scale, - signature.str().c_str()); - } - } + std::fprintf(stderr, + "[parallel-chain] speculator=%s lm_weight=%.3f " + "hazard_scale=%.3f yield_scale=%.3f signature=%s\n", + speculator_->score_kind().c_str(), + config.lm_log_weight, config.hazard_scale, + config.yield_scale, signature.str().c_str()); + } + } else if (spec_mode_ == SpecMode::chain) { + adaptive_fallback_reason_ = + speculator_fallback_reason(speculator_.get()); + std::fprintf(stderr, + "[parallel-chain] no speculator adapter for loaded drafter; " + "adaptive requests will use sticky AR\n"); } // The concurrent DDTree stack is gated to a local same-device drafter. @@ -384,9 +390,8 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { draft_kv_reset(*slot_draft_kv_[(size_t)slot]); } if (slots_.is_active(slot)) slots_.retire(slot); - if (slot >= 0 && slot < (int)last_survival_score_.size()) { - last_survival_score_[(size_t)slot] = - std::numeric_limits::quiet_NaN(); + if (slot >= 0 && slot < (int)last_activation_estimate_.size()) { + last_activation_estimate_[(size_t)slot] = {}; } } }; @@ -628,158 +633,55 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { std::chrono::steady_clock::now() - start).count(); }; - std::vector noise((size_t)T, b_.w_.mask_token_id); - noise[0] = profile_token; - std::vector noise_embed((size_t)hidden * T); - std::vector local_hidden((size_t)hidden * T); - std::vector prenorm_hidden((size_t)hidden * T); const bool profile_batched = batched_drafting_enabled(); - if (!b_.w_.embedder.embed(noise.data(), T, noise_embed.data())) { - cleanup(); - return false; - } auto draft_runner = [&](int lanes) -> double { - if (!profile_error.empty()) + if (!profile_error.empty()) { return std::numeric_limits::infinity(); - const auto start = std::chrono::steady_clock::now(); - std::vector states; - std::vector seeds; - states.reserve((size_t)chain_decode_bucket_width(lanes)); - seeds.reserve(states.capacity()); + } + std::vector profile_inputs; + std::vector selected; + profile_inputs.reserve(static_cast(lanes)); + selected.assign(static_cast(lanes), 1); for (int lane = 0; lane < lanes; ++lane) { - const int slot = synthetic_slots[(size_t)lane]; - DraftKvState * draft = ensure_slot_draft_kv(slot); - DraftFeatureMirror * mirror = slot_feature_mirror(slot); - if (!draft || !mirror || - !draft_kv_begin_step( - *draft, b_.dw_, b_.draft_backend_, - *mirror, ctx_tokens)) { - profile_error = "draft profiling setup failed"; - return std::numeric_limits::infinity(); - } - ggml_backend_tensor_set( - draft->inp_embed, noise_embed.data(), 0, - sizeof(float) * noise_embed.size()); - states.push_back(draft); - seeds.push_back(profile_token); - } - - if (profile_batched) { - const int bucket = chain_decode_bucket_width(lanes); - const int dummy_count = bucket - lanes; - const int cap = std::min( - slot_feature_mirrors_[0].cap, - std::max(1, b_.cfg_.draft_ctx_max)); - while ((int)dummy_draft_kv_.size() < dummy_count) { - auto dummy = std::make_unique(); - if (!draft_kv_init( - *dummy, b_.dw_, b_.draft_backend_, - cap, nullptr)) { - draft_kv_free(*dummy); - break; - } - dummy_draft_kv_.push_back(std::move(dummy)); - } - if ((int)dummy_draft_kv_.size() < dummy_count) { - profile_error = "draft profiling dummy allocation failed"; - return std::numeric_limits::infinity(); - } - for (int i = 0; i < dummy_count; ++i) { - DraftKvState * dummy = - dummy_draft_kv_[(size_t)i].get(); - if (!draft_kv_begin_step( - *dummy, b_.dw_, b_.draft_backend_, - slot_feature_mirrors_[0], 1)) { - profile_error = "draft profiling dummy setup failed"; - return std::numeric_limits::infinity(); - } - ggml_backend_tensor_set( - dummy->inp_embed, noise_embed.data(), 0, - sizeof(float) * noise_embed.size()); - states.push_back(dummy); - seeds.push_back(profile_token); - } - std::vector> draft_tokens; - std::vector> confidence; - if (!draft_kv_batch_compute( - batch_draft_graph_, b_.dw_, b_.draft_backend_, - b_.w_.output, states, seeds, - draft_tokens, confidence) || - draft_tokens.size() < (size_t)lanes) { - profile_error = "batched draft profiling compute failed"; + profile_inputs.push_back({ + synthetic_slots[static_cast(lane)], + profile_token, + true, + SpeculationPolicy::Always, + }); + } + const auto start = std::chrono::steady_clock::now(); + if (!prepare_chain_drafts( + profile_inputs, selected, + /*force_serial=*/!profile_batched, + /*fail_fast_batch=*/profile_batched)) { + profile_error = "draft profiling adapter proposal failed"; + return std::numeric_limits::infinity(); + } + for (const StepInput & input : profile_inputs) { + const PreparedChainDraft & prepared = + prepared_chain_drafts_[static_cast(input.slot)]; + if (!prepared.valid || + static_cast(prepared.tokens.size()) != T) { + profile_error = "draft profiling proposal shape failed"; return std::numeric_limits::infinity(); } - for (int lane = 0; lane < lanes; ++lane) { - if ((int)draft_tokens[(size_t)lane].size() != T) { - profile_error = "batched draft profiling shape failed"; - return std::numeric_limits::infinity(); - } - } - } else { - for (int lane = 0; lane < lanes; ++lane) { - DraftKvState * draft = states[(size_t)lane]; - if (ggml_backend_graph_compute( - b_.draft_backend_, draft->gf) != - GGML_STATUS_SUCCESS) { - profile_error = "draft profiling backbone failed"; - return std::numeric_limits::infinity(); - } - ggml_backend_tensor_get_async( - b_.draft_backend_, draft->hidden_states, - local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); - ggml_backend_tensor_get_async( - b_.draft_backend_, draft->hidden_prenorm, - prenorm_hidden.data(), 0, - sizeof(float) * prenorm_hidden.size()); - ggml_backend_synchronize(b_.draft_backend_); - bool chain_ok = false; - if (b_.dw_.selector.enabled) { - std::vector hidden_ptrs{ - local_hidden.data()}; - std::vector one_seed{profile_token}; - std::vector> draft_tokens; - chain_ok = dflash2_select_chains_batched( - b_.dw_, b_.draft_backend_, b_.w_.output, - hidden_ptrs, T, one_seed, draft_tokens) && - draft_tokens.size() == 1 && - (int)draft_tokens[0].size() == T; - } else { - std::vector draft_tokens; - std::vector confidence; - chain_ok = dspark_markov_correct_greedy_chain_fused( - b_.dw_, b_.draft_backend_, b_.w_.output, - local_hidden.data(), T, profile_token, - draft_tokens, &confidence, - prenorm_hidden.data()) && - (int)draft_tokens.size() == T; - } - if (!chain_ok) { - profile_error = "draft profiling chain failed"; - return std::numeric_limits::infinity(); - } - } } return std::chrono::duration( std::chrono::steady_clock::now() - start).count(); }; - SpecProfileResult tree = profile_monotonic_costs( - grid.tree_rows, tree_runner); - SpecProfileResult step = profile_monotonic_costs( - grid.step_rows, step_runner); - SpecProfileResult draft = profile_monotonic_costs( - grid.draft_lanes, draft_runner); + SpecCostProfileResult profiled = SpecCostProfiler{}.profile( + grid, tree_runner, step_runner, draft_runner, + speculator_->score_kind(), 5); cleanup(); - if (!tree.ok() || !step.ok() || !draft.ok() || !profile_error.empty()) { - std::fprintf(stderr, "[spec-profile] failed: %s%s%s%s\n", - profile_error.c_str(), tree.error.c_str(), step.error.c_str(), - draft.error.c_str()); + if (!profiled.ok() || !profile_error.empty()) { + std::fprintf(stderr, "[spec-profile] failed: %s%s\n", + profile_error.c_str(), profiled.error.c_str()); return false; } - SpecCostTables tables{ - std::move(tree.table), std::move(step.table), std::move(draft.table)}; + SpecCostTables tables = std::move(profiled.tables); SpecStepGeometry geometry; geometry.tree_width = V; geometry.bucket = [](int lanes) { @@ -806,9 +708,10 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { std::fprintf(stderr, "\n"); }; std::fprintf(stderr, - "[spec-profile] context=%d reps=5 mode=%s-draft\n", + "[spec-profile] context=%d reps=5 mode=%s-draft speculator=%s\n", ctx_tokens, - profile_batched ? "batched" : "serial"); + profile_batched ? "batched" : "serial", + tables.speculator_id.c_str()); print_table("tree_cost", tables.tree_cost); print_table("step_cost", tables.step_cost); print_table("draft_cost", tables.draft_cost); @@ -850,32 +753,17 @@ bool Qwen35SeqEngine::batched_drafting_enabled() const { } bool Qwen35SeqEngine::activation_scoring_available() const { - if (spec_mode_ != SpecMode::chain || !capture_features_) return false; - if (b_.dw_.selector.enabled) { - return dflash2_benefit_provider_ && - dflash2_benefit_provider_->ready(); - } - const int hidden = b_.dw_.n_embd; - return b_.dw_.dspark.enabled && b_.dw_.dspark.confidence_w && - b_.dw_.dspark.confidence_b && - (b_.dw_.dspark.confidence_dim == hidden || - b_.dw_.dspark.confidence_dim == - hidden + b_.dw_.dspark.markov_rank); + return spec_mode_ == SpecMode::chain && capture_features_ && + speculator_is_ready(speculator_.get()); } -SpecScoreKind Qwen35SeqEngine::chain_activation_score_kind() const { - if (b_.dw_.selector.enabled && dflash2_benefit_provider_ && - dflash2_benefit_provider_->ready()) { - return SpecScoreKind::DFlash2SelectorBenefitV1; - } - return activation_scoring_available() - ? SpecScoreKind::DSparkConfidence - : SpecScoreKind::Unspecified; +std::string Qwen35SeqEngine::chain_activation_score_kind() const { + return speculator_is_ready(speculator_.get()) + ? speculator_->score_kind() : kUnspecifiedScoreKind; } bool Qwen35SeqEngine::activation_scoring_enabled() const { const char * value = std::getenv("DFLASH_SPEC_ACTIVATION_SCORE"); - if (!value) value = std::getenv("DFLASH_SPEC_CONFIDENCE"); return !value || std::atoi(value) != 0; } @@ -884,7 +772,8 @@ bool Qwen35SeqEngine::prepare_chain_drafts( const std::vector & selected, bool force_serial, bool fail_fast_batch) { - if (selected.size() != inputs.size()) return false; + if (selected.size() != inputs.size() || + !speculator_is_ready(speculator_.get())) return false; // Accumulate the full drafting wall (draft graph compute + fused // selector + readbacks) into the round's [step-timing] attribution, @@ -910,6 +799,9 @@ bool Qwen35SeqEngine::prepare_chain_drafts( const int T = tree_width_; const int hidden = b_.w_.n_embd; + const uint32_t requirements = speculator_->input_requirements(); + const bool need_prenorm = + (requirements & SpeculatorInputPrenorm) != 0; struct Lane { size_t input_index = 0; int slot = -1; @@ -923,7 +815,7 @@ bool Qwen35SeqEngine::prepare_chain_drafts( std::vector noise_embed((size_t)hidden * T); // Proposal validity is current-block-specific. Do not clear the last - // published confidence before the draft succeeds: bootstrap must either + // published activation score before the draft succeeds: bootstrap must either // publish a finite score or fail, and a later draft failure must not erase // the immutable activation score already owned by the gate. for (size_t i = 0; i < inputs.size(); ++i) { @@ -959,9 +851,26 @@ bool Qwen35SeqEngine::prepare_chain_drafts( } if (lanes.empty()) return true; - std::vector> drafts; - std::vector> confidences; - std::vector selector_traces; + std::vector proposals; + auto invoke_adapter = [&]( + const std::vector> & hidden_blocks, + const std::vector> & prenorm_blocks, + const std::vector & seeds) { + SpeculatorBatchInput adapter_input; + adapter_input.lane_count = static_cast(seeds.size()); + adapter_input.requested_depth = T; + adapter_input.seed_tokens = seeds; + for (const std::vector & block : hidden_blocks) { + adapter_input.hidden_by_lane.push_back(block.data()); + } + if (need_prenorm) { + for (const std::vector & block : prenorm_blocks) { + adapter_input.prenorm_by_lane.push_back(block.data()); + } + } + return speculator_input_satisfies(adapter_input, requirements) && + speculator_->propose(adapter_input, proposals); + }; bool used_batch = false; const bool try_batched = !force_serial && batched_drafting_enabled(); if (try_batched) { @@ -1011,11 +920,14 @@ bool Qwen35SeqEngine::prepare_chain_drafts( } } if (dummy_ok) { + std::vector> hidden_blocks; + std::vector> prenorm_blocks; used_batch = draft_kv_batch_compute( - batch_draft_graph_, b_.dw_, - b_.draft_backend_, b_.w_.output, - batch_states, seeds, drafts, confidences, - &selector_traces); + batch_draft_graph_, b_.dw_, b_.draft_backend_, + batch_states, need_prenorm, + hidden_blocks, prenorm_blocks) && + invoke_adapter(hidden_blocks, prenorm_blocks, seeds) && + proposals.size() >= lanes.size(); } } if (!used_batch) { @@ -1055,115 +967,80 @@ bool Qwen35SeqEngine::prepare_chain_drafts( } if (!used_batch) { - drafts.resize(lanes.size()); - confidences.resize(lanes.size()); - if (b_.dw_.selector.enabled) { - selector_traces.resize(lanes.size()); - std::vector> hidden_blocks( + std::vector> hidden_blocks( + lanes.size(), + std::vector( + static_cast(hidden) * static_cast(T))); + std::vector> prenorm_blocks; + if (need_prenorm) { + prenorm_blocks.assign( lanes.size(), - std::vector((size_t) hidden * (size_t) T)); - std::vector seeds; - seeds.reserve(lanes.size()); - for (size_t lane = 0; lane < lanes.size(); ++lane) { - DraftKvState * state = lanes[lane].state; - if (ggml_backend_graph_compute( - b_.draft_backend_, state->gf) != - GGML_STATUS_SUCCESS) { - return false; - } - ggml_backend_tensor_get_async( - b_.draft_backend_, state->hidden_states, - hidden_blocks[lane].data(), 0, - sizeof(float) * hidden_blocks[lane].size()); - seeds.push_back(lanes[lane].seed); - } - ggml_backend_synchronize(b_.draft_backend_); - std::vector hidden_ptrs(lanes.size()); - for (size_t lane = 0; lane < lanes.size(); ++lane) { - hidden_ptrs[lane] = hidden_blocks[lane].data(); - } - if (!dflash2_select_chains_batched( - b_.dw_, b_.draft_backend_, b_.w_.output, - hidden_ptrs, T, seeds, drafts, &selector_traces)) { + std::vector( + static_cast(hidden) * static_cast(T))); + } + std::vector seeds; + seeds.reserve(lanes.size()); + for (size_t lane = 0; lane < lanes.size(); ++lane) { + DraftKvState * state = lanes[lane].state; + if (ggml_backend_graph_compute( + b_.draft_backend_, state->gf) != + GGML_STATUS_SUCCESS) { return false; } - } else { - std::vector local_hidden((size_t) hidden * T); - std::vector prenorm_hidden((size_t) hidden * T); - for (size_t lane = 0; lane < lanes.size(); ++lane) { - DraftKvState * state = lanes[lane].state; - if (ggml_backend_graph_compute( - b_.draft_backend_, state->gf) != - GGML_STATUS_SUCCESS) { - return false; - } - ggml_backend_tensor_get_async( - b_.draft_backend_, state->hidden_states, - local_hidden.data(), 0, - sizeof(float) * local_hidden.size()); + ggml_backend_tensor_get_async( + b_.draft_backend_, state->hidden_states, + hidden_blocks[lane].data(), 0, + sizeof(float) * hidden_blocks[lane].size()); + if (need_prenorm) { ggml_backend_tensor_get_async( b_.draft_backend_, state->hidden_prenorm, - prenorm_hidden.data(), 0, - sizeof(float) * prenorm_hidden.size()); - ggml_backend_synchronize(b_.draft_backend_); - if (!dspark_markov_correct_greedy_chain_fused( - b_.dw_, b_.draft_backend_, b_.w_.output, - local_hidden.data(), T, lanes[lane].seed, - drafts[lane], &confidences[lane], - prenorm_hidden.data())) { - return false; - } + prenorm_blocks[lane].data(), 0, + sizeof(float) * prenorm_blocks[lane].size()); } + seeds.push_back(lanes[lane].seed); + } + ggml_backend_synchronize(b_.draft_backend_); + if (!invoke_adapter(hidden_blocks, prenorm_blocks, seeds) || + proposals.size() != lanes.size()) { + return false; } } - if (drafts.size() < lanes.size() || - confidences.size() < lanes.size()) { - return false; - } + if (proposals.size() < lanes.size()) return false; for (size_t lane = 0; lane < lanes.size(); ++lane) { - if ((int)drafts[lane].size() != T) return false; + SpecProposal & proposal = proposals[lane]; + if (!proposal.error.empty() || + static_cast(proposal.tokens.size()) != T || + !std::isfinite(proposal.estimate.expected_yield) || + static_cast( + proposal.estimate.conditional_hazards.size()) < T - 1) { + if (!proposal.error.empty()) { + std::fprintf(stderr, + "[spec-gate] activation evaluation failed " + "request=%llu slot=%d kind=%s: %s\n", + static_cast( + slots_.slot(lanes[lane].slot).request_id), + lanes[lane].slot, speculator_->score_kind().c_str(), + proposal.error.c_str()); + } + return false; + } + const Lane & info = lanes[lane]; PreparedChainDraft & prepared = - prepared_chain_drafts_[(size_t)info.slot]; + prepared_chain_drafts_[static_cast(info.slot)]; prepared.valid = true; - prepared.generated = - slots_.slot(info.slot).generated_tokens(); + prepared.generated = slots_.slot(info.slot).generated_tokens(); prepared.root = info.seed; - prepared.tokens = std::move(drafts[lane]); - prepared.confidence = std::move(confidences[lane]); - prepared.selector_trace = lane < selector_traces.size() - ? std::move(selector_traces[lane]) - : DFlash2SelectorTrace{}; - - double & published = last_survival_score_[(size_t)info.slot]; - if (!std::isfinite(published)) { - if (b_.dw_.selector.enabled) { - std::string adapter_error; - if (!dflash2_benefit_provider_ || - !dflash2_benefit_provider_->publish_once( - prepared.selector_trace, - chain_verify_depth_for_round(), published, - &adapter_error)) { - std::fprintf(stderr, - "[spec-gate] DFlash2 request-benefit evaluation " - "failed request=%llu slot=%d: %s\n", - (unsigned long long) - slots_.slot(info.slot).request_id, - info.slot, adapter_error.empty() - ? "adapter unavailable" : adapter_error.c_str()); - } - } else if (!prepared.confidence.empty()) { - published = confidence_survival_yield( - prepared.confidence, chain_verify_depth_for_round()); - } else if (chain_activation_input_scoreable( - inputs[info.input_index])) { - std::fprintf(stderr, - "[spec-gate] DSpark confidence evaluation unavailable " - "request=%llu slot=%d\n", - (unsigned long long) - slots_.slot(info.slot).request_id, info.slot); - } + prepared.tokens = std::move(proposal.tokens); + prepared.estimate = proposal.estimate; + prepared.debug_depth_fields = + std::move(proposal.debug_depth_fields); + + ActivationEstimate & published = + last_activation_estimate_[static_cast(info.slot)]; + if (!std::isfinite(published.expected_yield)) { + published = std::move(proposal.estimate); } } return true; @@ -1201,19 +1078,17 @@ bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { } bool Qwen35SeqEngine::chain_proposal_input_capable( const StepInput & in) const { - const bool have_dflash2 = - b_.dw_.selector.enabled && b_.dw_.selector.hproj && - b_.dw_.selector.pred_cb && b_.dw_.selector.succ_cb && - b_.dw_.selector.rank > 0 && b_.dw_.selector.top_k > 0; - const bool have_dspark = - b_.dw_.dspark.enabled && b_.dw_.dspark.markov_w1 && - b_.dw_.dspark.markov_w2; + const uint32_t supported_inputs = + SpeculatorInputHidden | SpeculatorInputPrenorm; + const bool adapter_capable = + speculator_is_ready(speculator_.get()) && + speculator_->max_block_size() == tree_width_ && + (speculator_->input_requirements() & ~supported_inputs) == 0; return spec_mode_ == SpecMode::chain && capture_features_ && tree_width_ > 1 && tree_width_ <= 16 && resolve_chain_verify_depth( chain_verify_depth_for_round(), tree_width_) != 0 && - b_.dw_.block_size == tree_width_ && - (have_dflash2 || have_dspark) && + b_.dw_.block_size == tree_width_ && adapter_capable && in.slot >= 0 && in.slot < slots_.slot_count() && slots_.slot(in.slot).decoding() && slots_.slot(in.slot).cur_pos >= 1 && @@ -1244,7 +1119,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( StepResult result; const std::vector & inputs = plan.decode; if (admitted.size() != inputs.size() || !plan.prefills.empty()) { - result.error = "invalid DSpark chain admission plan"; + result.error = "invalid chain speculation admission plan"; return result; } @@ -1266,7 +1141,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( admitted.begin(), admitted.end(), [](uint8_t value) { return value != 0; })); if (requested_spec_count == 0) { - result.error = "empty DSpark chain admission plan"; + result.error = "empty chain speculation admission plan"; return result; } @@ -1292,8 +1167,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( std::vector flat; std::vector accepted; std::vector path; - std::vector confidence; - DFlash2SelectorTrace selector_trace; + std::vector debug_depth_fields; int32_t verify_bonus = -1; int32_t pending = -1; }; @@ -1342,7 +1216,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( return true; } fail_proposal_lane( - i, "DSpark proposal preparation failed after clean retry"); + i, "chain proposal preparation failed after clean retry"); return false; }; @@ -1399,17 +1273,17 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( next.input_index = i; next.slot = in.slot; next.root = in.token; - next.selector_trace = std::move(prepared.selector_trace); + next.debug_depth_fields = + std::move(prepared.debug_depth_fields); next.flat = std::move(prepared.tokens); - next.confidence = std::move(prepared.confidence); prepared.valid = false; if (!truncate_chain_proposal(next.flat, V)) return false; const size_t verified_signal_depths = static_cast(V - 1); - if (next.selector_trace.depths.size() > verified_signal_depths) { - next.selector_trace.depths.resize(verified_signal_depths); + if (next.debug_depth_fields.size() > verified_signal_depths) { + next.debug_depth_fields.resize(verified_signal_depths); } - next.tree = make_dspark_chain_tree(next.flat); + next.tree = make_chain_verify_tree(next.flat); if (next.tree.n_nodes + 1 != V) return false; proposal = std::move(next); return true; @@ -1425,7 +1299,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( } else { if (active_admitted[i]) { fail_proposal_lane( - i, "DSpark proposal remained invalid after clean retry"); + i, "chain proposal remained invalid after clean retry"); } continue; } @@ -1457,7 +1331,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( V, tree_bucket, max_prefix, tree_scratch_base_, tree_scratch_stride_, b_.cfg_.kq_stride_pad)) { - result.error = "packed DSpark chain verify graph build failed"; + result.error = "packed chain speculation verify graph build failed"; return result; } t_verify_build_end = timing_clock::now(); @@ -1514,7 +1388,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( if (!b_.w_.embedder.embed( flat_tokens.data(), total_tree, tree_embed.data())) { - result.error = "packed DSpark chain embedding failed"; + result.error = "packed chain speculation embedding failed"; return result; } ggml_backend_tensor_set(tree_sg.inp_embed, tree_embed.data(), 0, @@ -1543,7 +1417,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( sizeof(int32_t) * seq_lens_.size()); if (ggml_backend_graph_compute(b_.target_backend_, tree_sg.gf) != GGML_STATUS_SUCCESS) { - result.error = "packed DSpark chain verify compute failed"; + result.error = "packed chain speculation verify compute failed"; return result; } t_verify_exec_end = timing_clock::now(); @@ -1566,7 +1440,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( proposal.accepted, static_cast(std::max(0, room)), lane_posterior, proposal.verify_bonus); if (proposal.accepted.empty()) { - result.error = "DSpark accepted path has no context headroom"; + result.error = "chain accepted path has no context headroom"; return result; } proposal.path.reserve(proposal.accepted.size()); @@ -1583,7 +1457,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( [&](int32_t token) { return token_is_eos(token); }); proposal.path.resize(safe_prefix); proposal.accepted.resize(safe_prefix); - if (!proposal.selector_trace.depths.empty()) { + if (!proposal.debug_depth_fields.empty()) { static const bool selector_log_enabled = []() { const char * value = std::getenv("DFLASH_DFLASH2_SELECTOR_LOG"); @@ -1595,28 +1469,18 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( proposal.path.empty() ? 0 : proposal.path.size() - 1; std::fprintf(stderr, "[spec-selector] {\"request_id\":%llu,\"slot\":%d," - "\"generated\":%d,\"accepted_depth\":%zu,\"depths\":[", - (unsigned long long) sequence.request_id, - proposal.slot, sequence.generated_tokens(), - accepted_depth); + "\"score_kind\":\"%s\",\"generated\":%d," + "\"accepted_depth\":%zu,\"depths\":[", + static_cast(sequence.request_id), + proposal.slot, chain_activation_score_kind().c_str(), + sequence.generated_tokens(), accepted_depth); for (size_t depth = 0; - depth < proposal.selector_trace.depths.size(); ++depth) { - const DFlash2DepthSignal & signal = - proposal.selector_trace.depths[depth]; + depth < proposal.debug_depth_fields.size(); ++depth) { std::fprintf(stderr, - "%s{\"depth\":%zu,\"accepted\":%s," - "\"selected_logp\":%.8g,\"lm_margin\":%.8g," - "\"topk_mass\":%.8g,\"rank\":%d," - "\"lm_top1\":%s,\"selector_margin\":%.8g," - "\"selector_mass\":%.8g,\"selector_entropy\":%.8g}", + "%s{\"depth\":%zu,\"accepted\":%s,%s}", depth == 0 ? "" : ",", depth + 1, depth < accepted_depth ? "true" : "false", - signal.selected_log_prob, signal.lm_top2_margin, - signal.top_k_mass, signal.selected_rank, - signal.agrees_with_lm_top1 ? "true" : "false", - signal.selector_margin, - signal.selector_winner_mass, - signal.selector_entropy); + proposal.debug_depth_fields[depth].c_str()); } std::fprintf(stderr, "]}\n"); } @@ -1658,8 +1522,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( if (!app.ok || app.physical_rows.size() != proposal.path.size() || !table_ok) { result.error = app.busy - ? "paged KV pool exhausted during DSpark chain commit" - : "DSpark accepted-path K/V append failed"; + ? "paged KV pool exhausted during chain speculation commit" + : "chain accepted-path K/V append failed"; return result; } replay_segments.push_back({ @@ -1720,14 +1584,14 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( out.slot = inputs[i].slot; out.failed = true; out.error = proposal_errors[i].empty() - ? "DSpark proposal lane made no progress" + ? "chain proposal lane made no progress" : proposal_errors[i]; result.decode.push_back(std::move(out)); } return result; } if (!upload_all_active_block_tables()) { - result.error = "DSpark mixed-step block-table refresh failed"; + result.error = "chain mixed-step block-table refresh failed"; return result; } t_commit_end = timing_clock::now(); @@ -1754,7 +1618,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( (ar_bucket > 0 && (!durable_sg.active_slot_ids || !durable_sg.state_slot_ids)) || !durable_sg.logits_row_indices || !durable_sg.argmax_tokens) { - result.error = "DSpark mixed commit/AR graph build failed"; + result.error = "chain mixed commit/AR graph build failed"; return result; } t_replay_build_end = timing_clock::now(); @@ -1769,7 +1633,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( embed_buf_.resize(static_cast(hidden) * n_total); if (!b_.w_.embedder.embed( durable_tokens.data(), n_total, embed_buf_.data())) { - result.error = "DSpark mixed commit/AR embedding failed"; + result.error = "chain mixed commit/AR embedding failed"; return result; } ggml_backend_tensor_set( @@ -1887,7 +1751,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( if (ggml_backend_graph_compute(b_.target_backend_, durable_sg.gf) != GGML_STATUS_SUCCESS) { - result.error = "DSpark mixed commit/AR compute failed"; + result.error = "chain mixed commit/AR compute failed"; return result; } t_replay_exec_end = timing_clock::now(); @@ -1900,7 +1764,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( ggml_backend_synchronize(b_.target_backend_); for (int lane = 0; lane < spec_count; ++lane) { if (argmax_buf_[static_cast(lane)] < 0) { - result.error = "DSpark durable replay produced invalid token"; + result.error = "chain durable replay produced invalid token"; return result; } } @@ -1919,7 +1783,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( } } if (!commit_residency_writes(write_slots)) { - result.error = "DSpark mixed-step KV write commit failed"; + result.error = "chain mixed-step KV write commit failed"; return result; } @@ -1936,7 +1800,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( proposal.slot, lane, &argmax_buf_[static_cast(lane)], &logits_buf_); if (proposal.pending < 0) { - result.error = "DSpark durable replay sampling failed"; + result.error = "chain durable replay sampling failed"; return result; } } @@ -2466,9 +2330,8 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( draft_kv_reset(*slot_draft_kv_[(size_t)result.slot]); } if (result.slot >= 0 && - result.slot < (int)last_survival_score_.size()) { - last_survival_score_[(size_t)result.slot] = - std::numeric_limits::quiet_NaN(); + result.slot < (int)last_activation_estimate_.size()) { + last_activation_estimate_[(size_t)result.slot] = {}; if (result.slot < (int)prepared_chain_drafts_.size()) { prepared_chain_drafts_[(size_t)result.slot].valid = false; @@ -2765,15 +2628,19 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { chain_spec_request_capable(in); double activation_score = std::numeric_limits::quiet_NaN(); + std::vector activation_hazards; if (use_activation_score && scoreable && in.slot >= 0 && - in.slot < (int)last_survival_score_.size() && - std::isfinite( - last_survival_score_[(size_t)in.slot])) { - activation_score = last_survival_score_[(size_t)in.slot]; + in.slot < (int)last_activation_estimate_.size() && + std::isfinite(last_activation_estimate_[(size_t)in.slot].expected_yield)) { + activation_score = + last_activation_estimate_[(size_t)in.slot].expected_yield; + activation_hazards = last_activation_estimate_[ + (size_t)in.slot].conditional_hazards; } candidates.push_back({ seq.request_id, in.slot, policy, scoreable, can_speculate, activation_score, + std::move(activation_hazards), chain_activation_score_kind(), }); } @@ -2793,11 +2660,12 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } const int slot = candidate.slot; if (slot >= 0 && - slot < (int)last_survival_score_.size() && - std::isfinite( - last_survival_score_[(size_t)slot])) { - candidate.confidence_yield = - last_survival_score_[(size_t)slot]; + slot < (int)last_activation_estimate_.size() && + std::isfinite(last_activation_estimate_[(size_t)slot].expected_yield)) { + candidate.activation_yield = + last_activation_estimate_[(size_t)slot].expected_yield; + candidate.conditional_hazards = + last_activation_estimate_[(size_t)slot].conditional_hazards; } } gate_plan = speculation_gate_->plan( @@ -2816,9 +2684,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { draft_kv_reset(*slot_draft_kv_[(size_t)slot]); } if (slot >= 0 && - slot < (int)last_survival_score_.size()) { - last_survival_score_[(size_t)slot] = - std::numeric_limits::quiet_NaN(); + slot < (int)last_activation_estimate_.size()) { + last_activation_estimate_[(size_t)slot] = {}; } }; auto commit_evaluation_fallback = @@ -2826,12 +2693,10 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { reset_evaluation_lane(evaluation.slot); if (speculation_gate_->commit_evaluation_fallback_ar( evaluation.request_id)) { - const SpecScoreKind kind = + const std::string kind = chain_activation_score_kind(); - const char * reason = kind == - SpecScoreKind::DFlash2SelectorBenefitV1 - ? "benefit_evaluation_failed" - : "confidence_evaluation_failed"; + const char * reason = + "activation_evaluation_failed"; log_spec_evaluation_fallback( evaluation.request_id, evaluation.slot, kind, reason); @@ -2881,9 +2746,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { inputs, one, /*force_serial=*/true); if (!lane_scored || evaluation.slot < 0 || evaluation.slot >= - (int)last_survival_score_.size() || - !std::isfinite(last_survival_score_[ - (size_t)evaluation.slot])) { + (int)last_activation_estimate_.size() || + !std::isfinite(last_activation_estimate_[ + (size_t)evaluation.slot].expected_yield)) { commit_evaluation_fallback(evaluation); } } @@ -2892,9 +2757,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { score_evaluations) { if (evaluation.slot < 0 || evaluation.slot >= - (int)last_survival_score_.size() || - !std::isfinite(last_survival_score_[ - (size_t)evaluation.slot])) { + (int)last_activation_estimate_.size() || + !std::isfinite(last_activation_estimate_[ + (size_t)evaluation.slot].expected_yield)) { commit_evaluation_fallback(evaluation); } } @@ -3058,8 +2923,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { ? initial_prediction_realized_tokens(gate_plan, speculative) : std::numeric_limits::quiet_NaN(); log_spec_gate_plan( - gate_plan, speculation_gate_->fixed_yield_scale(), - realized_tokens, + gate_plan, realized_tokens, cost_sample_valid ? measured_us : std::numeric_limits::quiet_NaN()); @@ -3470,8 +3334,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { measured_us); if (spec_gate_debug_enabled()) { log_spec_gate_plan( - *pending_ar_gate_plan, - speculation_gate_->fixed_yield_scale(), 0.0, + *pending_ar_gate_plan, 0.0, measured_us); } } @@ -3504,9 +3367,8 @@ void Qwen35SeqEngine::retire(int slot) { if (!slots_.is_active(slot)) return; const uint64_t request_id = slots_.slot(slot).request_id; if (speculation_gate_) speculation_gate_->forget(request_id); - if (slot >= 0 && slot < (int)last_survival_score_.size()) { - last_survival_score_[(size_t)slot] = - std::numeric_limits::quiet_NaN(); + if (slot >= 0 && slot < (int)last_activation_estimate_.size()) { + last_activation_estimate_[(size_t)slot] = {}; if (slot < (int)prepared_chain_drafts_.size()) { prepared_chain_drafts_[(size_t)slot].valid = false; } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index ed2c79564..fcd72ac0c 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -22,8 +22,8 @@ #pragma once #include "common/concurrency/seq_engine.h" -#include "common/concurrency/speculation_gate.h" -#include "common/dflash2_benefit.h" +#include "common/speculation/speculator.h" +#include "common/speculation/speculation_gate.h" #include "common/dflash_draft_kv.h" #include "common/dflash_feature_ring.h" #include "common/ddtree.h" @@ -76,10 +76,10 @@ class Qwen35SeqEngine final : public SeqEngine { StepResult step(const StepPlan & plan) override; // Fabricate a steady-state paged context and profile the three launch - // families used by the DSpark gate. Called once from backend init. + // series used by the activation gate. Called once from backend init. bool profile_spec_costs(int context_tokens); - // True only when a trained DSpark head or the guarded DFlash2 benefit - // adapter can produce a first-request activation score. A configured chain + // True only when the registered speculator adapter can produce a + // first-request activation score. A configured chain // may still accept Adaptive requests and serve sticky AR when this is false. bool activation_scoring_available() const; StepPlanLimits step_plan_limits(int decode_rows) const override { @@ -149,8 +149,8 @@ class Qwen35SeqEngine final : public SeqEngine { int generated = -1; int32_t root = -1; std::vector tokens; - std::vector confidence; - DFlash2SelectorTrace selector_trace; + ActivationEstimate estimate; + std::vector debug_depth_fields; }; bool prepare_chain_drafts( const std::vector & inputs, @@ -159,7 +159,7 @@ class Qwen35SeqEngine final : public SeqEngine { bool fail_fast_batch = false); bool batched_drafting_enabled() const; bool activation_scoring_enabled() const; - SpecScoreKind chain_activation_score_kind() const; + std::string chain_activation_score_kind() const; // Deliberate seam for a later per-round cohort controller. Request mode // remains sticky; every returned depth must stay in [2, tree_width_]. int chain_verify_depth_for_round() const { @@ -169,7 +169,7 @@ class Qwen35SeqEngine final : public SeqEngine { // round attributing wall time to draft, verify, readback, CPU commit, // replay, and packed-AR phases. Diagnostic only; off by default. static bool step_timing_enabled(); - // DDTree preserves its legacy best-effort AR fallback. DSpark chain + // DDTree preserves its legacy best-effort AR fallback. Chain // proposal failures are instead returned as lane-local DecodeOutput // failures so a sticky speculation decision can never execute as AR. std::optional step_ddtree(const StepPlan & plan); @@ -197,14 +197,14 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector> dummy_draft_kv_; std::vector prepared_chain_drafts_; std::unique_ptr speculation_gate_; - std::unique_ptr dflash2_benefit_provider_; + std::unique_ptr speculator_; // Startup profile/adapter failure is a request-local sticky AR outcome, // never an admission or step error for a configured chain. std::string adaptive_fallback_reason_ = "cost_profile_unavailable"; std::vector adaptive_fallback_ar_; - std::vector last_survival_score_; + std::vector last_activation_estimate_; // Per-round draft cost accumulator for [step-timing]; reset at the top - // of each dspark_chain round, accumulated by prepare_chain_drafts. + // of each chain-speculation round, accumulated by prepare_chain_drafts. double round_draft_us_ = 0.0; int round_draft_lanes_ = 0; // Hoisted per-step buffers (reused across step() calls). diff --git a/server/test/test_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp index e7a63dbda..6814dd05e 100644 --- a/server/test/test_chain_spec_shapes.cpp +++ b/server/test/test_chain_spec_shapes.cpp @@ -10,7 +10,7 @@ static int g_checks = 0; int main() { const std::vector draft = {10, 11, 12, 13}; - const DDTree tree = make_dspark_chain_tree(draft); + const DDTree tree = make_chain_verify_tree(draft); CHECK(tree.n_nodes == 3); CHECK((tree.token_ids == std::vector{11, 12, 13})); CHECK((tree.depths == std::vector{1, 2, 3})); @@ -25,7 +25,7 @@ int main() { std::vector short_draft = draft; CHECK(truncate_chain_proposal(short_draft, 2)); CHECK((short_draft == std::vector{10, 11})); - const DDTree short_tree = make_dspark_chain_tree(short_draft); + const DDTree short_tree = make_chain_verify_tree(short_draft); CHECK(short_tree.n_nodes == 1); CHECK((short_tree.parents == std::vector{-1, 0})); const std::vector before_invalid = short_draft; From 375543e7c9ca4ca29ee3713da6b6b5731b32fe10 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 19 Aug 2026 20:49:35 +0000 Subject: [PATCH 39/42] chore(speculation): consolidate activation harness and telemetry --- .../QWEN38_DSPARK_ADAPTIVE_SELECTION.md | 121 ++---- .../benchmarks/concurrency/FEATURE_MATRIX.md | 88 +--- .../concurrency/analyze_gate_decisions.py | 207 +++------ .../concurrency/concurrent_benchmark.py | 118 ++++- .../feature_concurrent_benchmark.py | 257 ----------- .../concurrency/generate_dspark_prompts.py | 168 -------- .../concurrency/generate_feature_prompts.py | 57 --- .../concurrency/generate_ragged_prompts.py | 16 +- .../concurrency/run_qwen36_feature_matrix.sh | 4 +- .../concurrency/run_qwen38_dspark_matrix.sh | 371 ---------------- .../concurrency/summarize_feature_matrix.py | 320 +------------- .../test_feature_concurrent_benchmark.py | 6 +- .../concurrency/test_feature_tools.py | 407 +++++------------- .../concurrency/verify_feature_metrics.py | 10 +- 14 files changed, 332 insertions(+), 1818 deletions(-) delete mode 100755 harness/benchmarks/concurrency/feature_concurrent_benchmark.py delete mode 100755 harness/benchmarks/concurrency/generate_dspark_prompts.py delete mode 100644 harness/benchmarks/concurrency/generate_feature_prompts.py delete mode 100755 harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh diff --git a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md index 45cff9de9..ff685211f 100644 --- a/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md +++ b/harness/benchmarks/QWEN38_DSPARK_ADAPTIVE_SELECTION.md @@ -1,4 +1,4 @@ -# Qwen3.8 DSpark adaptive-selection prompts +# Qwen3.8 adaptive-selection prompts (historical DSpark baseline) This six-request workload is a prompt-selection fixture for the concurrent adaptive speculation gate. It is deliberately balanced rather than @@ -15,11 +15,11 @@ that justified the label. The exact completion lengths and ordered output hashes are retained locally in `QWEN38_DSPARK_ADAPTIVE_SELECTION_R9700.json`. -DSpark is the fixed proposal mechanism for this fixture. The benchmark is -about adaptive activation: which lanes the gate admits at each live -concurrency, whether the selected `k` beats matched pure AR, whether rejected -lanes avoid drafting, and which target phase dominates a bad decision. - +The dense labels below were collected with DSpark and remain useful priors for +request selection. They do not prescribe the current proposal mechanism. +Concurrent runs use whichever single `Speculator` adapter the server registers; +currently that is DFlash2. The activation engine owns scoring, ranking, cost +comparison, sticky AR fallback, and telemetry. ## Dense screening baseline @@ -55,88 +55,33 @@ interpret their throughput as a valid speculative win or loss. ## Concurrent activation benchmark -The smallest adversarial cohort is fixed at C=3: - -- two dense strong-win prompts (`sum_product` and `rolling_max`); -- the strongest dense AR win (`reproducibility`). - -This is intentionally difficult. A useful gate must learn that dense labels -are only priors: concurrency can move the break-even point enough that a -formerly strong speculative sample should run as AR. - -Run three paired fresh-process repeats: - -```bash -MODEL=/path/Qwen3.8-27B-PR625-IQ4_XS.gguf \ -DRAFT_MODEL=/path/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ -LUCE_SERVER_BIN=server/build-pr625-r9700/dflash_server \ -VISIBLE_DEVICES=0 \ -WORKLOADS=adaptive-selection-c3 CLIENTS=3 \ -DECODE_MODES=ar,speculation,adaptive REPEATS=3 \ -harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh -``` - -The paged concurrent executor requires full attention, represented by -`FA_WINDOW=0`. At this fixture's maximum context (at most 134 prompt tokens -plus 256 generated tokens), a 2048-token window would not discard context, so -full attention does not change the attended token set. It may still change -kernel overhead; the paired controls measure that runtime. The runner uses -Q8_0 K/V by default and records both cache types and the window in every case. - -For the wider selection boundary, run all six prompts at C=6: - -Run the selection workload only at C=6: - -```bash -MODEL=/path/Qwen3.8-27B-PR625-IQ4_XS.gguf \ -DRAFT_MODEL=/path/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ -LUCE_SERVER_BIN=server/build-pr625-r9700/dflash_server \ -VISIBLE_DEVICES=0 \ -WORKLOADS=adaptive-selection CLIENTS=6 \ -DECODE_MODES=ar,speculation,adaptive REPEATS=3 \ -harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh -``` - -The dense oracle labels are priors, not a hard assertion about the concurrent -executor. In particular, the two marginal dense wins may correctly become AR -choices after packed-tree, replay, padding, and synchronization costs are -included. Adaptive mode performs one evaluation per request at its first -target-decode boundary. A successful evaluation produces the request's -confidence score and commits either AR or speculation through EOS/max tokens; -a failed evaluation commits AR without inventing a score. The request is never -evaluated again. Every request is ranked independently from its own initial -confidence, an immutable offline yield scale, and the measured cost for the -current graph shape. Accepted-token history from this or any prior request is -never used. The only runtime EWMA tracks hardware latency by graph shape, not -request content or user/chat identity. - -The prefill logits produce the first sampled output token before a -target-decode step exists. The one-time draft and activation therefore finish -before the first target-decode execution. A request that chooses speculation -reuses that proposal immediately, so it has no preliminary AR decode round; a -request that retires directly from prefill (for example, `max_tokens=1`) has no -target-decode mode to activate. - -`DFLASH_MIN_TOKENS` does not delay speculation or switch the request to AR. -Below the floor, verification stops before an accepted EOS and durable replay -samples the replacement non-EOS token from the kept prefix tip. At or above the -floor, the EOS is kept and deeper tokens are discarded. - -Spec decode and prompt prefill use incompatible target graphs. When at least -one request in the live decode cohort chose speculation, the executor advances -that decode cohort and reports every selected prefill slice as deferred; the -scheduler retries those unchanged prompts after the Spec wave drains. This -preserves the literal request mode but can increase TTFT for requests queued -behind a long Spec wave. An all-AR cohort keeps the fused mixed -prefill/decode path and its continuous-batching TTFT behavior. - -The useful adaptive behavior is: - -1. rank the two strong wins above the two known losses; -2. keep the prose loss and normally the HumanEval loss in AR; -3. make the marginal pair expose the actual concurrent break-even boundary; -4. preserve the ordered greedy output hashes across AR, forced speculation, - and adaptive modes. +The smallest adversarial cohort is fixed at C=3: the two dense strong-win +prompts (`sum_product` and `rolling_max`) plus the strongest dense AR win +(`reproducibility`). The wider boundary uses all six prompts at C=6. + +The removed DSpark-only matrix runner must not be used for new measurements. +Launch paired fresh AR, forced-speculation, and adaptive server processes with +the active DFlash2 drafter, and drive each process with +`harness/benchmarks/concurrency/concurrent_benchmark.py`. The checked-in JSONL +is already suitable as `--prompt-file`; select the C=3 IDs explicitly when +running the smaller cohort. `run_qwen38_dflash2_subsets.sh` remains the core +forced-mode/refill control. + +The dense oracle labels are priors, not hard assertions about the concurrent +executor. Adaptive mode evaluates each request once at its first target-decode +boundary. The active adapter returns proposal tokens, an activation score, +expected yield, and optional conditional hazards. The generic engine ranks +those scores against the profiled cost table and commits each request to AR or +speculation through retirement. Evaluation failure emits +`activation_evaluation_failed` and commits sticky AR. Accepted-token history +and user identity are not scoring inputs. + +The prefill logits produce the first sampled output token before a target-decode +step exists. A speculative request reuses its activation proposal immediately; +a request that retires directly from prefill has no target-decode mode to +activate. The useful behavior is still to preserve ordered greedy output hashes, +keep known losses in AR, and let measured concurrent costs decide the marginal +pair. ## Profiling outputs diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 64b5c6798..03f7e1a2b 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -3,81 +3,19 @@ The bounded Strix Halo measurements collected for the draft implementation are recorded in [`STRIX_HALO_RESULTS.md`](STRIX_HALO_RESULTS.md). -## Qwen3.8 DSpark adaptive matrix - -`run_qwen38_dspark_matrix.sh` is the C7 acceptance-gated speculation matrix. -It uses only the [RadixArk Qwen3.8-27B-DSpark](https://huggingface.co/RadixArk/Qwen3.8-27B-DSpark) -source. `DRAFT_MODEL` must point to a GGUF produced from that repository. The -labeled R9700 baseline uses the no-YaRN Q8_0 artifact. - -```bash -MODEL=/opt/models/Qwen3.8-27B-Q4_K_M.gguf \ -DRAFT_MODEL=/opt/models/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ -REPEATS=5 \ -harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh -``` - -The default fresh-process matrix is: - -- `ar`, `speculation`, `adaptive-on`, and `adaptive-confidence-off`. - `adaptive-on` batch-evaluates each new adaptive request exactly once at its - first target-decode boundary. A successful score chooses AR or speculation; - a failed evaluation chooses AR without a synthetic score. That mode remains - fixed through retirement. If speculation is chosen, the one-time bootstrap - proposal is reused immediately; a request that chooses AR never drafts again. - A target-decode step exists only after prefill has produced the first sampled - token, so one-token requests that retire directly from prefill never enter - adaptive activation. Each request is ranked only from its own initial score; - accepted-token history and user/chat identity are not inputs. A shape-keyed - timing EWMA may refine hardware cost estimates without carrying content - history. The last arm sets - `DFLASH_SPEC_CONFIDENCE=0`, suppressing activation scoring so it remains a - draft-free AR ablation. There is no eager-draft matrix axis. -- A Spec request runs the speculative executor immediately even below - `DFLASH_MIN_TOKENS`. An accepted EOS below the floor is excluded from the - committed path and durable replay samples the non-EOS replacement; an allowed - EOS ends the path. The floor never creates an AR warm-up phase. -- Spec decode is graph-exclusive with prompt prefill. If a scheduler plan - contains any sticky-Spec lane, decode advances and selected prefill slices are - reported as deferred with no prompt progress. They are retried after the Spec - wave drains. This can raise TTFT for queued requests, so the inverse-TTFT gate - measures an intentional mode-fidelity tradeoff. All-AR cohorts continue - using the fused mixed prefill/decode path. -- Live concurrency `C ∈ {1,2,3,4,6,8}` over the checked-in HumanEval and - GSM8K cohorts plus deterministic prose prompts. -- A fixed C=6 north-star cohort with two code and four chat requests. -- Batched drafting enabled for DSpark rows; adaptive drafts the one-time cold - batch and thereafter only requests whose sticky mode is speculation. - Startup profiling uses a 4096-token synthetic context. - -Every process records the target, server, shared-library, and drafter hashes, -the literal command and launch environment, startup pool dimensions, request -IDs, and terminal concurrency counters. The proof rejects forced-speculation -rows unless every measured request has positive `spec_steps`. Adaptive rows -may legitimately choose k=0, but must show exactly one `[spec-activation]` -AR/speculation decision for every measured adaptive request. A scored -activation carries finite confidence/yield and no fallback reason; a failed -one carries null scores, sticky AR, and -`fallback_reason=confidence_evaluation_failed`. The proof also requires both -the packed DSpark startup marker and a completed startup cost profile. Chain -rows must keep all -`ddtree_*` counters at zero, preserving the DDTree proof semantics below. - -For every workload/concurrency pair, the summarizer forms a paired oracle: - -```text -goodput oracle = max(ar, speculation) -TTFT oracle = min(ar, speculation) -``` - -`adaptive-on` must reach at least 0.995 of that oracle for mean and median -output goodput and inverse TTFT. The summary fails the run if any of those four -ratios misses the gate; p95 alone is never used as acceptance evidence. A -second table reports paired goodput and inverse-TTFT -deltas between `adaptive-on` and `adaptive-confidence-off` without gating the -ablation. Set `WORKLOADS`, `CLIENTS`, `DECODE_MODES`, -or `CONFIDENCE_ABLATION=0` to select a smaller -diagnostic subset. +## Qwen3.8 adaptive speculation + +The DSpark-only matrix runner was removed with the concurrent DSpark execution +path. The checked-in Qwen3.8 selection prompts remain useful workload fixtures, +but proposal generation and activation scoring now come from the one active +`Speculator` adapter (currently DFlash2). + +Use `concurrent_benchmark.py` as the common request client for AR, forced +speculation, and adaptive server processes. Use +`run_qwen38_dflash2_subsets.sh` for the forced AR/speculation and refill +controls. Adaptive runs must retain the fail-closed `[spec-activation]` proof: +one scored or failed decision per measured request, opaque `score_kind`, +optional hazards, and `activation_evaluation_failed` for sticky-AR fallback. ## Qwen3.6 DDTree/PFlash/KVFlash matrix diff --git a/harness/benchmarks/concurrency/analyze_gate_decisions.py b/harness/benchmarks/concurrency/analyze_gate_decisions.py index c70e9e26d..5752fa05b 100644 --- a/harness/benchmarks/concurrency/analyze_gate_decisions.py +++ b/harness/benchmarks/concurrency/analyze_gate_decisions.py @@ -39,19 +39,6 @@ TIMING_COUNT_FIELDS = ( "live", "k", "emitted_tokens", "accepted_tokens", "target_forwards", ) -ACTIVATION_SCORE_KINDS = { - "dspark_confidence", - "dflash2_selector_benefit_v1", - "unspecified", -} -ACTIVATION_FALLBACK_REASONS = { - "confidence_evaluation_failed", - "benefit_evaluation_failed", - "benefit_adapter_unavailable", - "benefit_adapter_invalid_config", - "cost_profile_unavailable", - "activation_evaluation_failed", -} TYPED_REASON_RE = re.compile(r"^[a-z][a-z0-9]*(?:_[a-z0-9]+)*$") @@ -88,133 +75,71 @@ def _validate_timing(row: dict[str, Any], path: Path, line_no: int) -> None: ) -def _validate_typed_activation( - row: dict[str, Any], path: Path, line_no: int, -) -> None: +def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: required = ( - "initial_confidence", "activation_score", "request_benefit", - "score_kind", "expected_yield", "decision_reason", + "request_id", "slot", "activation_score", "score_kind", + "expected_yield", "evaluation", "fallback_reason", + "decision_reason", "decision", ) for key in required: if key not in row: raise ValueError( - f"{path}:{line_no}: typed spec-activation {key} is required" - ) - kind = row["score_kind"] - if kind not in ACTIVATION_SCORE_KINDS: - raise ValueError( - f"{path}:{line_no}: spec-activation score_kind is unsupported" - ) - reason = row["decision_reason"] - if not isinstance(reason, str) or not TYPED_REASON_RE.fullmatch(reason): - raise ValueError( - f"{path}:{line_no}: spec-activation decision_reason must be a " - "nonempty snake-case reason" - ) - - evaluation = row["evaluation"] - score_fields = ("activation_score", "request_benefit", "expected_yield") - if evaluation == "scored": - if kind == "unspecified": - raise ValueError( - f"{path}:{line_no}: scored spec-activation score_kind must " - "identify its scoring model" - ) - for key in score_fields: - value = row[key] - if ( - type(value) not in (int, float) - or not math.isfinite(value) - or value < 1.0 - ): - raise ValueError( - f"{path}:{line_no}: spec-activation {key} must be finite " - "and at least 1 for a scored evaluation" - ) - confidence = row["initial_confidence"] - if kind == "dspark_confidence": - if ( - type(confidence) not in (int, float) - or not math.isfinite(confidence) - or confidence < 1.0 - ): - raise ValueError( - f"{path}:{line_no}: DSpark initial_confidence must be " - "finite and at least 1" - ) - elif confidence is not None: - raise ValueError( - f"{path}:{line_no}: non-DSpark initial_confidence must be null" - ) - if row["fallback_reason"] is not None: - raise ValueError( - f"{path}:{line_no}: scored spec-activation fallback_reason " - "must be null" + f"{path}:{line_no}: spec-activation {key} is required" ) - return - - if row["initial_confidence"] is not None or any( - row[key] is not None for key in score_fields - ): - raise ValueError( - f"{path}:{line_no}: failed spec-activation scores must be null" - ) - if row["decision"] != "ar": - raise ValueError( - f"{path}:{line_no}: failed spec-activation decision must be ar" - ) - if reason != "evaluation_failed": - raise ValueError( - f"{path}:{line_no}: failed spec-activation decision_reason must " - "be evaluation_failed" - ) - fallback = row["fallback_reason"] - if fallback not in ACTIVATION_FALLBACK_REASONS: - raise ValueError( - f"{path}:{line_no}: failed spec-activation fallback_reason must " - "be a recognized typed reason" - ) - -def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: for key in ("request_id", "slot"): - value = row.get(key) + value = row[key] if type(value) is not int or value < 0: raise ValueError( f"{path}:{line_no}: spec-activation {key} must be a " "non-negative int" ) - decision = row.get("decision") - if decision not in ("ar", "speculation"): + if row["decision"] not in ("ar", "speculation"): raise ValueError( f"{path}:{line_no}: spec-activation decision must be ar or " "speculation" ) - evaluation = row.get("evaluation") - if evaluation not in ("scored", "failed"): + if row["evaluation"] not in ("scored", "failed"): raise ValueError( f"{path}:{line_no}: spec-activation evaluation must be scored " "or failed" ) - if "fallback_reason" not in row: + + kind = row["score_kind"] + if not isinstance(kind, str) or not kind: raise ValueError( - f"{path}:{line_no}: spec-activation fallback_reason is required" + f"{path}:{line_no}: spec-activation score_kind must be a " + "nonempty string" + ) + reason = row["decision_reason"] + if not isinstance(reason, str) or not TYPED_REASON_RE.fullmatch(reason): + raise ValueError( + f"{path}:{line_no}: spec-activation decision_reason must be a " + "nonempty snake-case reason" ) - typed_keys = ( - "score_kind", "activation_score", "request_benefit", "decision_reason", - ) - if any(key in row for key in typed_keys): - _validate_typed_activation(row, path, line_no) - return - for key in ("initial_confidence", "expected_yield"): - if key not in row: + hazards = row.get("hazards") + if hazards is not None: + if not isinstance(hazards, list) or any( + type(value) not in (int, float) + or not math.isfinite(value) + or not 0.0 <= value <= 1.0 + for value in hazards + ): raise ValueError( - f"{path}:{line_no}: spec-activation {key} is required" + f"{path}:{line_no}: spec-activation hazards must be an " + "array of finite probabilities" + ) + + score_fields = ("activation_score", "expected_yield") + if row["evaluation"] == "scored": + if kind == "unspecified": + raise ValueError( + f"{path}:{line_no}: scored spec-activation score_kind must " + "identify its scoring model" ) - if evaluation == "scored": - for key in ("initial_confidence", "expected_yield"): - value = row.get(key) + for key in score_fields: + value = row[key] if ( type(value) not in (int, float) or not math.isfinite(value) @@ -230,20 +155,19 @@ def _validate_activation(row: dict[str, Any], path: Path, line_no: int) -> None: "must be null" ) return - if row.get("initial_confidence") is not None or ( - row.get("expected_yield") is not None - ): + + if any(row[key] is not None for key in score_fields) or hazards is not None: raise ValueError( f"{path}:{line_no}: failed spec-activation scores must be null" ) - if decision != "ar": + if row["decision"] != "ar": raise ValueError( f"{path}:{line_no}: failed spec-activation decision must be ar" ) - if row["fallback_reason"] != "confidence_evaluation_failed": + if row["fallback_reason"] != "activation_evaluation_failed": raise ValueError( f"{path}:{line_no}: failed spec-activation fallback_reason must " - "be confidence_evaluation_failed" + "be activation_evaluation_failed" ) @@ -645,7 +569,7 @@ def _activation_proof( fallback_reason_counts: dict[str, int] = defaultdict(int) decision_reason_counts: dict[str, int] = defaultdict(int) for row in activations: - score_kind_counts[str(row.get("score_kind") or "legacy_confidence")] += 1 + score_kind_counts[row["score_kind"]] += 1 if isinstance(row.get("fallback_reason"), str): fallback_reason_counts[row["fallback_reason"]] += 1 if isinstance(row.get("decision_reason"), str): @@ -708,7 +632,6 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A lambda: { "rounds": 0, "admitted": 0, "activation_score_sum": 0.0, "activation_scored": 0, - "confidence_score_sum": 0.0, "confidence_scored": 0, } ) k_histogram: dict[int, int] = defaultdict(int) @@ -719,17 +642,12 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A stats["rounds"] += 1 if score["admitted"]: stats["admitted"] += 1 - if math.isfinite(score["score"]): - if score["source"] in ("confidence", "current", "initial"): - stats["activation_score_sum"] += score["score"] - stats["activation_scored"] += 1 - if ( - (score["source"] == "confidence" and - score["score_kind"] is None) - or score["score_kind"] == "dspark_confidence" - ): - stats["confidence_score_sum"] += score["score"] - stats["confidence_scored"] += 1 + if ( + math.isfinite(score["score"]) + and score["source"] in ("fresh", "initial") + ): + stats["activation_score_sum"] += score["score"] + stats["activation_scored"] += 1 requests = [] for engine_id, prompt in sorted( @@ -751,29 +669,20 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A stats["activation_score_sum"] / stats["activation_scored"] if stats["activation_scored"] else None ), - "mean_confidence_yield": ( - stats["confidence_score_sum"] / stats["confidence_scored"] - if stats["confidence_scored"] else None - ), "activation_slot": ( activation.get("slot") if activation is not None else None ), - "initial_confidence": ( - activation.get("initial_confidence") - if activation is not None else None - ), "activation_score": ( activation.get("activation_score") if activation is not None else None ), - "request_benefit": ( - activation.get("request_benefit") - if activation is not None else None - ), "activation_score_kind": ( activation.get("score_kind") if activation is not None else None ), + "activation_hazards": ( + activation.get("hazards") if activation is not None else None + ), "expected_yield": ( activation.get("expected_yield") if activation is not None else None @@ -949,13 +858,9 @@ def compare_prompts(cases: list[dict[str, Any]]) -> list[dict[str, Any]]: "activation_score_kind" ), "activation_score": request.get("activation_score"), - "request_benefit": request.get("request_benefit"), "mean_activation_score": request.get("mean_activation_score"), - "initial_confidence": request.get("initial_confidence"), "expected_yield": request.get("expected_yield"), - "mean_confidence_yield": request.get( - "mean_confidence_yield" - ), + "activation_hazards": request.get("activation_hazards"), "commit_per_spec_step": request.get( "commit_per_spec_step" ), @@ -1237,7 +1142,7 @@ def build_report( ] activation_comparisons = compare_activation_shapes(cases) return { - "schema_version": 6, + "schema_version": 7, "cases": cases, "prompt_comparisons": compare_prompts(cases), "activation_comparisons": activation_comparisons, diff --git a/harness/benchmarks/concurrency/concurrent_benchmark.py b/harness/benchmarks/concurrency/concurrent_benchmark.py index 4c3e85774..e5004089a 100755 --- a/harness/benchmarks/concurrency/concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/concurrent_benchmark.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Measure end-to-end output goodput and TTFT under concurrent streaming load.""" +"""Measure concurrent goodput, TTFT, request IDs, and prompt telemetry.""" from __future__ import annotations @@ -14,6 +14,21 @@ from pathlib import Path from typing import Any, Iterable +CLIENT_SCRIPT = Path(__file__).resolve() + + +def client_provenance(argv: list[str] | None = None) -> dict[str, Any]: + """Return the literal process argv and 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 sha256_text(text: str) -> str: return hashlib.sha256(text.encode("utf-8")).hexdigest() @@ -65,12 +80,15 @@ def iter_sse_data(lines: Iterable[bytes]) -> Iterable[str]: 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, @@ -97,11 +115,17 @@ def stream_request(args: argparse.Namespace, prompt: str) -> dict[str, Any]: done_received = True break event = json.loads(data) + if isinstance(event.get("id"), str): + request_id = event["id"] usage = event.get("usage") or {} - if isinstance(usage.get("completion_tokens"), int): + if type(usage.get("completion_tokens")) is int: completion_tokens = usage["completion_tokens"] - if isinstance(usage.get("prompt_tokens"), int): + if type(usage.get("prompt_tokens")) is int: prompt_tokens = usage["prompt_tokens"] + if isinstance(usage.get("timings"), dict): + timings = dict(usage["timings"]) + 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"] @@ -114,7 +138,7 @@ def stream_request(args: argparse.Namespace, prompt: str) -> dict[str, Any]: 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 + 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]" @@ -126,15 +150,25 @@ def stream_request(args: argparse.Namespace, prompt: str) -> dict[str, Any]: 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 + 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": sha256_text(output), "reasoning_content_sha256": sha256_text(reasoning_output), @@ -250,6 +284,39 @@ def worker(index: int) -> None: } +def enrich_level(level: dict[str, Any]) -> None: + ok = [row for row in level["requests_detail"] if row.get("error") is None] + effective = [row.get("effective_prompt_tokens") for row in ok] + effective_complete = bool(ok) and all(type(value) is int for value in effective) + request_ids = [row.get("request_id") for row in ok] + request_ids_complete = ( + bool(ok) + and all(isinstance(value, str) and value for value 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 = [ + row.get(key) for row in ok + if type(row.get(key)) in (int, float) + ] + level[f"{key}_median"] = ( + statistics.median(values) if len(values) == len(ok) and ok else None + ) + + def fmt(value: Any, spec: str = ".2f") -> str: return format(value, spec) if isinstance(value, (int, float)) else "n/a" @@ -258,9 +325,9 @@ 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 |", - "| ---: | ---: | ---: | ---: | ---: | ---: | :--- | ---: | ---: | ---: |", + "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( @@ -269,9 +336,13 @@ def markdown(report: dict[str, Any]) -> str: 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'])} |" + f"{fmt(level['prompt_tokens_min'], '.0f')}–" + f"{fmt(level['prompt_tokens_max'], '.0f')} | " + f"{fmt(level['effective_prompt_tokens_min'], '.0f')}–" + f"{fmt(level['effective_prompt_tokens_max'], '.0f')} | " + f"{fmt(level['effective_to_wire_prompt_ratio'], '.3f')} | " + f"{fmt(level['ttft_median_s'], '.3f')}/" + f"{fmt(level['ttft_max_s'], '.3f')} |" ) return "\n".join(lines) + "\n" @@ -304,6 +375,10 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--server-metadata-json", type=Path) parser.add_argument("--out", type=Path, required=True) parser.add_argument("--label", default="") + parser.add_argument( + "--require-effective-prompt-telemetry", action="store_true", + help="fail when usage.timings.effective_prompt_tokens is absent", + ) return parser @@ -320,24 +395,37 @@ def run(args: argparse.Namespace) -> int: 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)) + level = 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": 2, "label": args.label, "base_url": args.base_url, + "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") + 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) + bad = any( + 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 diff --git a/harness/benchmarks/concurrency/feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/feature_concurrent_benchmark.py deleted file mode 100755 index c9df122e1..000000000 --- a/harness/benchmarks/concurrency/feature_concurrent_benchmark.py +++ /dev/null @@ -1,257 +0,0 @@ -#!/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 | Prompt tok/s to first | Wire prompt range | " - "Effective prompt range | Effective/wire | TTFT median/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_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" - - -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_dspark_prompts.py b/harness/benchmarks/concurrency/generate_dspark_prompts.py deleted file mode 100755 index 2b3830e40..000000000 --- a/harness/benchmarks/concurrency/generate_dspark_prompts.py +++ /dev/null @@ -1,168 +0,0 @@ -#!/usr/bin/env python3 -"""Build deterministic Qwen3.8 DSpark concurrency workload manifests.""" - -from __future__ import annotations - -import argparse -import json -from pathlib import Path - -from generate_ragged_prompts import write_records - - -HERE = Path(__file__).resolve().parent -PROMPT_ROOT = HERE.parent / "prompts" -SOURCE_FILES = { - "humaneval": PROMPT_ROOT / "bench_he.jsonl", - "gsm8k": PROMPT_ROOT / "bench_gsm.jsonl", -} -ADAPTIVE_SELECTION_SOURCE = ( - PROMPT_ROOT / "qwen38_dspark_adaptive_selection.jsonl" -) -# Mixed C=3 cohort: two strong dense speculation wins plus the strongest -# dense speculation loss. The adaptive win condition is a 2-spec+1-AR split -# that beats both uniform modes. -ADAPTIVE_SELECTION_C3_IDS = ( - "adaptive-he-09-sum-product", - "adaptive-he-10-rolling-max", - "adaptive-prose-01-reproducibility", -) -PROSE_TOPICS = ( - "why reproducible benchmarks need immutable inputs and explicit hardware metadata", - "how admission control improves the reliability of a concurrent inference service", - "the tradeoff between latency, throughput, and fairness in token scheduling", - "why calibrated confidence is useful when deciding whether to speculate", - "how bounded caches turn memory pressure into a predictable operating policy", - "the value of fail-closed telemetry when validating an optimization", - "why deterministic greedy output is a strong regression oracle", - "how batching amortizes launch overhead without sharing request state", - "why startup profiling should use monotone cost tables", - "how paired fresh-process repeats reduce benchmark order bias", -) - - -def _message_prompt(row: dict[str, object], source: Path, line_no: int) -> str: - messages = row.get("messages") - if not isinstance(messages, list) or not messages: - raise ValueError(f"{source}:{line_no}: missing messages") - pieces: list[str] = [] - for message in messages: - if not isinstance(message, dict): - raise ValueError(f"{source}:{line_no}: message must be an object") - role = message.get("role") - content = message.get("content") - if not isinstance(role, str) or not isinstance(content, str) or not content: - raise ValueError(f"{source}:{line_no}: invalid message") - pieces.append(content if len(messages) == 1 else f"{role}: {content}") - return "\n\n".join(pieces) - - -def _source_records(profile: str) -> list[dict[str, object]]: - source = SOURCE_FILES[profile] - records: list[dict[str, object]] = [] - for line_no, raw in enumerate(source.read_text(encoding="utf-8").splitlines(), 1): - if not raw.strip(): - continue - row = json.loads(raw) - records.append({ - "id": str(row.get("id") or f"{profile}-{line_no:02d}"), - "suite": profile, - "prompt": _message_prompt(row, source, line_no), - }) - if len(records) < 8: - raise ValueError(f"{source}: need at least eight distinct prompts") - return records - - -def _prose_records() -> list[dict[str, object]]: - return [ - { - "id": f"prose-{index:02d}", - "suite": "prose", - "prompt": ( - "Write a clear, self-contained technical essay of about 500 words on " - f"{topic}. Include one concrete example and end with a concise conclusion." - ), - } - for index, topic in enumerate(PROSE_TOPICS, 1) - ] - - -def _adaptive_selection_records() -> list[dict[str, object]]: - records: list[dict[str, object]] = [] - for line_no, raw in enumerate( - ADAPTIVE_SELECTION_SOURCE.read_text(encoding="utf-8").splitlines(), 1, - ): - if not raw.strip(): - continue - row = json.loads(raw) - if not isinstance(row.get("id"), str) or not isinstance(row.get("prompt"), str): - raise ValueError( - f"{ADAPTIVE_SELECTION_SOURCE}:{line_no}: invalid id or prompt" - ) - if row.get("expected_dense_oracle") not in ("ar", "speculation"): - raise ValueError( - f"{ADAPTIVE_SELECTION_SOURCE}:{line_no}: invalid dense oracle" - ) - baseline = row.get("dense_r9700_baseline") - if not isinstance(baseline, dict) or baseline.get("lossless") is not True: - raise ValueError( - f"{ADAPTIVE_SELECTION_SOURCE}:{line_no}: baseline must be lossless" - ) - records.append(row) - if len(records) != 6: - raise ValueError( - f"{ADAPTIVE_SELECTION_SOURCE}: expected exactly six prompts" - ) - return records - - -def build_records(profile: str) -> list[dict[str, object]]: - if profile in SOURCE_FILES: - return _source_records(profile) - prose = _prose_records() - if profile == "adaptive-selection": - return _adaptive_selection_records() - if profile == "adaptive-selection-c3": - by_id = {row["id"]: row for row in _adaptive_selection_records()} - missing = [pid for pid in ADAPTIVE_SELECTION_C3_IDS if pid not in by_id] - if missing: - raise ValueError( - f"{ADAPTIVE_SELECTION_SOURCE}: missing C=3 prompts {missing}" - ) - return [dict(by_id[pid]) for pid in ADAPTIVE_SELECTION_C3_IDS] - if profile == "prose": - return prose - if profile == "north-star": - code = _source_records("humaneval")[:2] - chat = [dict(row) for row in prose[:4]] - for row in code: - row["cohort"] = "2code+4chat" - for row in chat: - row["cohort"] = "2code+4chat" - return code + chat - raise ValueError(f"unknown DSpark workload {profile!r}") - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--profile", choices=( - "humaneval", "gsm8k", "prose", "north-star", - "adaptive-selection", "adaptive-selection-c3", - ), - 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)} {args.profile} prompts to {args.out}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/generate_feature_prompts.py b/harness/benchmarks/concurrency/generate_feature_prompts.py deleted file mode 100644 index d642b138e..000000000 --- a/harness/benchmarks/concurrency/generate_feature_prompts.py +++ /dev/null @@ -1,57 +0,0 @@ -#!/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 index c04d4c434..3595dd4d8 100755 --- a/harness/benchmarks/concurrency/generate_ragged_prompts.py +++ b/harness/benchmarks/concurrency/generate_ragged_prompts.py @@ -14,6 +14,10 @@ "short": (250, 350, 450, 550), "medium": (650, 850, 1150, 1350), "long": (2000, 2600, 3400, 4000), + # Long-context profiles. Runtime telemetry, not word count, proves that + # PFlash compression or KVFlash pressure actually activated. + "compression": (34000, 36000, 38000, 40000), + "kv-pressure": (12000, 14000, 16000, 18000), } WORD_BANK = ( @@ -83,8 +87,16 @@ def build_profile_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: +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}, + ) if path.exists(): raise FileExistsError(f"refusing to overwrite {path}") path.parent.mkdir(parents=True, exist_ok=True) diff --git a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh index c7a113916..aeb8c1d50 100755 --- a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh @@ -4,8 +4,8 @@ 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}" +CLIENT="${CLIENT:-$SCRIPT_DIR/concurrent_benchmark.py}" +GENERATOR="${GENERATOR:-$SCRIPT_DIR/generate_ragged_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}" diff --git a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh b/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh deleted file mode 100755 index 3444781dc..000000000 --- a/harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh +++ /dev/null @@ -1,371 +0,0 @@ -#!/usr/bin/env bash -# Fresh-process Qwen3.8 DSpark mode matrix with fail-closed chain 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_dspark_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}" -ANALYZER="${ANALYZER:-$SCRIPT_DIR/analyze_gate_decisions.py}" - -MODEL="${MODEL:-}" -DRAFT_MODEL="${DRAFT_MODEL:-}" -LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" -OUT="${OUT:-$REPO/.harness-runs/qwen38-dspark-matrix-$(date -u +%Y%m%dT%H%M%SZ)}" -REPEATS="${REPEATS:-1}" -WORKLOADS="${WORKLOADS:-humaneval,gsm8k,prose,north-star}" -DECODE_MODES="${DECODE_MODES:-ar,speculation,adaptive}" -CONFIDENCE_ABLATION="${CONFIDENCE_ABLATION:-1}" -STEP_TIMING="${STEP_TIMING:-1}" -CLIENTS="${CLIENTS:-1,2,3,4,6,8}" -SLOTS="${SLOTS:-8}" -MAX_CTX="${MAX_CTX:-8192}" -MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" -MAX_TOKENS="${MAX_TOKENS:-256}" -WARMUP_TOKENS="${WARMUP_TOKENS:-16}" -PROFILE_CONTEXT="${PROFILE_CONTEXT:-4096}" -CACHE_TYPE_K="${CACHE_TYPE_K:-q8_0}" -CACHE_TYPE_V="${CACHE_TYPE_V:-q8_0}" -FA_WINDOW="${FA_WINDOW:-0}" -PORT="${PORT:-18138}" -COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-3}" -HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-900}" -REQUEST_TIMEOUT_SECONDS="${REQUEST_TIMEOUT_SECONDS:-1800}" -TARGET_DEVICE="${TARGET_DEVICE:-hip:0}" -DRAFT_DEVICE="${DRAFT_DEVICE:-hip:0}" -VISIBLE_DEVICES="${VISIBLE_DEVICES:-1}" - -usage() { - cat <<'EOF' -Usage: - MODEL=/path/Qwen3.8-27B-PR625-IQ4_XS.gguf \ - DRAFT_MODEL=/path/Qwen3.8-27B-DSpark-RadixArk-no-yarn-q8_0.gguf \ - harness/benchmarks/concurrency/run_qwen38_dspark_matrix.sh - -The only supported drafter source for this matrix is: - https://huggingface.co/RadixArk/Qwen3.8-27B-DSpark - -DRAFT_MODEL is a GGUF produced from that repository. The labeled R9700 -adaptive-selection baseline uses the no-YaRN Q8_0 artifact; changing draft -quantization changes both acceptance and cost, so it invalidates those stored -oracle ratios. -The default fresh-process matrix runs ar, forced speculation, one-shot -adaptive activation with one confidence evaluation per request (a failed -evaluation falls back to sticky AR), and a draft-free confidence-off AR -ablation at live concurrency 1,2,3,4,6,8 over -HumanEval, GSM8K, and prose. An adaptive request keeps its selected -speculation/AR mode until retirement. The 2-code+4-chat north-star row runs -only at C=6. The summary fails if adaptive-on is below 0.995 of the paired -ar/speculation oracle in mean or median goodput or TTFT. The confidence -ablation delta is reported but not gated; there is no eager-draft matrix axis. -The labeled adaptive-selection workload is fixed at C=6. Run it with -WORKLOADS=adaptive-selection CLIENTS=6 to compare adaptive decisions against -two strong speculation wins, two marginal wins, and two AR wins. -The adaptive-selection-c3 workload is fixed at C=3: two strong speculation -wins plus the strongest AR win. Run it with WORKLOADS=adaptive-selection-c3 -CLIENTS=3; the target adaptive outcome is a 2-spec+1-AR split that beats both -uniform decode modes. -STEP_TIMING=1 (the default) records phase-attributed JSON for every measured -decode round and writes profiling.json plus profiling.md. Set STEP_TIMING=0 -only when measuring the small overhead of diagnostic logging itself. -The screened prompt labels used Q8 KV and FA_WINDOW=2048. The concurrent -paged executor requires full attention, so this runner pins Q8 KV but defaults -FA_WINDOW=0 and records all three values in every case. Dense labels are priors; -the paired concurrent AR/speculation controls are the runtime oracle. - -Defaults select the second visible host GPU and address it as hip:0 inside the -process (VISIBLE_DEVICES=1, TARGET_DEVICE=hip:0, DRAFT_DEVICE=hip:0). Override -VISIBLE_DEVICES on a single-GPU Strix Halo host. 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 Qwen3.8 target GGUF" >&2; exit 2; } -[[ -r "$DRAFT_MODEL" ]] || { echo "set DRAFT_MODEL to a readable RadixArk DSpark GGUF" >&2; exit 2; } -[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } -for value_name in REPEATS SLOTS MAX_CTX MAX_CONCURRENT_PREFILLS MAX_TOKENS WARMUP_TOKENS PROFILE_CONTEXT HEALTH_TIMEOUT_SECONDS REQUEST_TIMEOUT_SECONDS; do - value="${!value_name}" - [[ "$value" =~ ^[1-9][0-9]*$ ]] || { echo "$value_name must be positive" >&2; exit 2; } -done -[[ "$PORT" =~ ^[1-9][0-9]*$ ]] || { echo "PORT must be positive" >&2; exit 2; } -[[ "$FA_WINDOW" == 0 ]] || { echo "paged concurrency requires FA_WINDOW=0" >&2; exit 2; } -[[ -n "$CACHE_TYPE_K" && -n "$CACHE_TYPE_V" ]] || { echo "cache types must be non-empty" >&2; exit 2; } -[[ "$COOLDOWN_SECONDS" =~ ^[0-9]+$ ]] || { echo "COOLDOWN_SECONDS must be non-negative" >&2; exit 2; } -[[ "$CONFIDENCE_ABLATION" == 0 || "$CONFIDENCE_ABLATION" == 1 ]] || { echo "CONFIDENCE_ABLATION must be 0 or 1" >&2; exit 2; } -[[ "$STEP_TIMING" == 0 || "$STEP_TIMING" == 1 ]] || { echo "STEP_TIMING must be 0 or 1" >&2; exit 2; } -(( PROFILE_CONTEXT < MAX_CTX )) || { echo "PROFILE_CONTEXT must be below MAX_CTX" >&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 mode_list <<< "$DECODE_MODES" -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 WORKLOADS "${workload_list[@]}" || exit 2 -reject_duplicates DECODE_MODES "${mode_list[@]}" || exit 2 -reject_duplicates CLIENTS "${client_list[@]}" || exit 2 - -for workload in "${workload_list[@]}"; do - case "$workload" in - humaneval|gsm8k|prose|north-star|adaptive-selection|adaptive-selection-c3) ;; - *) echo "unknown workload $workload" >&2; exit 2 ;; - esac -done -for mode in "${mode_list[@]}"; do - case "$mode" in - ar|speculation|adaptive) ;; - *) echo "unknown decode mode $mode" >&2; exit 2 ;; - esac -done -for clients in "${client_list[@]}"; do - [[ "$clients" =~ ^[1-9][0-9]*$ ]] || { echo "CLIENTS entries must be positive" >&2; exit 2; } - (( clients <= SLOTS )) || { echo "CLIENTS=$clients exceeds SLOTS=$SLOTS" >&2; exit 2; } -done - -variants=() -for mode in "${mode_list[@]}"; do - if [[ "$mode" == adaptive ]]; then - variants+=("adaptive-on") - if [[ "$CONFIDENCE_ABLATION" == 1 ]]; then - variants+=("adaptive-confidence-off") - fi - else - variants+=("$mode") - fi -done -(( ${#variants[@]} > 0 )) || { echo "no decode variants selected" >&2; exit 2; } - -MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" -DRAFT_MODEL_SHA256="$(sha256sum "$DRAFT_MODEL" | 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 -} - -case_applicable() { - local workload="$1" clients="$2" - case "$workload" in - north-star|adaptive-selection) [[ "$clients" == 6 ]] ;; - adaptive-selection-c3) [[ "$clients" == 3 ]] ;; - *) return 0 ;; - esac -} - -run_case() { - local repeat="$1" workload="$2" clients="$3" variant="$4" - local decode_mode="$variant" confidence="" - if [[ "$variant" == adaptive-confidence-off ]]; then - decode_mode=adaptive - confidence=off - elif [[ "$variant" == adaptive-on ]]; then - decode_mode=adaptive - confidence=on - fi - - local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" - mkdir -p "$case_dir" - local capacity=$((SLOTS * MAX_CTX)) - local model_id=qwen38-dspark - local -a command=( - "$LUCE_SERVER_BIN" "$MODEL" - --draft "$DRAFT_MODEL" - --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_DEVICE" - --paged-attention --max-concurrency "$SLOTS" - --kv-pool-tokens "$capacity" --max-ctx "$MAX_CTX" - --cache-type-k "$CACHE_TYPE_K" --cache-type-v "$CACHE_TYPE_V" - --fa-window "$FA_WINDOW" - --prefix-cache-slots 0 --prefill-cache-slots 0 - --admission-coalesce-ms 20 --draft-residency persistent - --decode-mode "$decode_mode" - --host 127.0.0.1 --port "$PORT" --model-name "$model_id" - ) - local -a launch_env=( - "HIP_VISIBLE_DEVICES=$VISIBLE_DEVICES" - "DFLASH_MAX_CONCURRENT_PREFILLS=$MAX_CONCURRENT_PREFILLS" - "DFLASH_SPEC_BATCHED_DRAFT=1" - "DFLASH_STEP_TIMING=$STEP_TIMING" - ) - if [[ "$decode_mode" == adaptive ]]; then - launch_env+=( - "DFLASH_SPEC_GATE_LOG=1" - "DFLASH_SPEC_PROFILE_CONTEXT=$PROFILE_CONTEXT" - ) - if [[ "$confidence" == on ]]; then - launch_env+=("DFLASH_SPEC_CONFIDENCE=1") - else - launch_env+=("DFLASH_SPEC_CONFIDENCE=0") - fi - fi - - printf 'env ' > "$case_dir/server-command.txt" - printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$case_dir/server-command.txt" - printf '\n' >> "$case_dir/server-command.txt" - - local -a metadata=( - python3 "$METADATA_TOOL" - --out "$case_dir/server-metadata.json" - --variant "$variant" --workload "$workload" - --clients "$clients" --repeat "$repeat" - --binary "$LUCE_SERVER_BIN" --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_CONCURRENT_PREFILLS" - --target-device "$TARGET_DEVICE" --draft-device "$DRAFT_DEVICE" - --draft-model "$DRAFT_MODEL" --draft-model-sha256 "$DRAFT_MODEL_SHA256" - --decode-mode "$decode_mode" - --cache-type-k "$CACHE_TYPE_K" --cache-type-v "$CACHE_TYPE_V" - --fa-window "$FA_WINDOW" - ) - if [[ "$decode_mode" == adaptive ]]; then - metadata+=(--draft-always off --confidence "$confidence") - fi - local item - for item in "${launch_env[@]}"; do metadata+=(--launch-env "$item"); done - "${metadata[@]}" - - echo "[run] $workload C=$clients repeat=$repeat variant=$variant" - port_is_available || return 1 - env "${launch_env[@]}" "${command[@]}" > "$case_dir/server.log" 2>&1 & - server_pid=$! - if ! wait_health; then tail -n 160 "$case_dir/server.log" >&2 || true; return 1; fi - python3 "$RUNTIME_METADATA_TOOL" \ - --metadata "$case_dir/server-metadata.json" \ - --server-log "$case_dir/server.log" - - local 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 0 - --require-distinct-prompts --temperature 0 - --require-effective-prompt-telemetry - --timeout "$REQUEST_TIMEOUT_SECONDS" --cooldown 0 - ) - local -a warmup_cmd=( - python3 "$CLIENT" "${common_client[@]}" - --max-tokens "$WARMUP_TOKENS" - --out "$case_dir/warmup.json" - --label "$variant $workload C=$clients warmup" - ) - "${warmup_cmd[@]}" > "$case_dir/warmup.txt" - - # Keep a measured-only log alongside the full startup/warmup log. This - # prevents warmup rounds from contaminating gate and phase distributions. - local benchmark_log_offset - benchmark_log_offset="$(wc -c < "$case_dir/server.log")" - local -a benchmark_cmd=( - python3 "$CLIENT" "${common_client[@]}" - --max-tokens "$MAX_TOKENS" - --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 - tail -c "+$((benchmark_log_offset + 1))" "$case_dir/server.log" \ - > "$case_dir/benchmark-server.log" - - local -a proof_cmd=( - python3 "$PROOF_TOOL" - --bench "$case_dir/bench.json" - --server-log "$case_dir/server.log" - --out "$case_dir/feature-proof.json" - ) - [[ "$decode_mode" != ar ]] && proof_cmd+=(--expect chain) - "${proof_cmd[@]}" - 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]}" - if ! case_applicable "$workload" "$clients"; then - echo "[skip] $workload is pinned to one concurrency; C=$clients" - continue - fi - shift_by=$(((repeat + c_index) % ${#variants[@]})) - for ((i=0; i<${#variants[@]}; i++)); do - variant="${variants[$(((i + shift_by) % ${#variants[@]}))]}" - 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 "$ANALYZER" "$OUT" \ - --out "$OUT/profiling.json" --markdown-out "$OUT/profiling.md" -python3 "$SUMMARIZER" "$OUT" --out "$OUT/summary.md" -echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/summarize_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py index aba59d4d3..17076a62e 100755 --- a/harness/benchmarks/concurrency/summarize_feature_matrix.py +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Summarize Qwen3.6 feature ablations or Qwen3.8 DSpark oracle runs.""" +"""Summarize Qwen3.6 concurrent feature ablations.""" from __future__ import annotations @@ -233,327 +233,9 @@ def summarize_qwen36(reports: list[dict]) -> str: return "\n".join(lines) -DSPARK_VARIANTS = { - "ar", - "speculation", - "adaptive-on", - "adaptive-confidence-off", -} -ORACLE_THRESHOLD = 0.995 - - -def _is_dspark_item(item: dict) -> bool: - config = item["meta"].get("feature_config") or {} - return config.get("decode_mode") in ("ar", "speculation", "adaptive") - - -def _positive_metric(item: dict, key: str) -> float: - value = item["level"].get(key) - if type(value) not in (int, float) or value <= 0: - raise ValueError(f"{item['path']}: missing positive {key}") - return float(value) - - -def _dspark_signature(item: dict) -> tuple[object, ...]: - base = run_signature(item) - meta = item["meta"] - config = meta.get("feature_config") or {} - draft_sha256 = config.get("draft_model_sha256") - binary_sha256 = meta.get("server_binary_sha256") - if ( - not isinstance(draft_sha256, str) or not draft_sha256 - or not isinstance(binary_sha256, str) or not binary_sha256 - ): - raise ValueError(f"{item['path']}: incomplete DSpark provenance") - return (*base, draft_sha256, binary_sha256) - - -def _mean(values: list[float]) -> float: - return statistics.fmean(values) - - -def _pair(values: list[float]) -> tuple[float, float]: - return _mean(values), median(values) - - -def _fmt_pair(values: list[float], digits: int = 2) -> str: - mean_value, median_value = _pair(values) - return f"{mean_value:.{digits}f}/{median_value:.{digits}f}" - - -def _fmt_percent_pair(values: list[float]) -> str: - mean_value, median_value = _pair(values) - return f"{mean_value * 100:+.1f}%/{median_value * 100:+.1f}%" - - -def summarize_dspark(reports: list[dict]) -> str: - grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) - for item in reports: - meta, level = item["meta"], item["level"] - variant = str(meta.get("variant", "")) - if variant not in DSPARK_VARIANTS: - raise ValueError(f"{item['path']}: mixed or unknown DSpark variant {variant!r}") - config = meta.get("feature_config") or {} - expected_confidence = { - "ar": None, - "speculation": None, - "adaptive-on": "on", - "adaptive-confidence-off": "off", - } - draft_always = config.get("draft_always") - valid_draft_policy = ( - draft_always in (None, "off") - if variant.startswith("adaptive-") - else draft_always is None - ) - if ( - not valid_draft_policy - or config.get("confidence") != expected_confidence[variant] - ): - raise ValueError( - f"{item['path']}: adaptive metadata does not match variant {variant}" - ) - key = (str(meta["workload"]), int(level["clients"]), 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({_dspark_signature(item) for item in items}) != 1: - raise ValueError(f"{key}: incompatible DSpark run metadata") - prompt_hashes = { - item["level"].get("selected_prompt_set_sha256") for item in items - } - if len(prompt_hashes) != 1 or not all( - isinstance(value, str) and value for value in prompt_hashes - ): - raise ValueError(f"{key}: missing or inconsistent selected prompts") - - lines = [ - "# Qwen3.8 DSpark adaptive concurrency matrix", "", - "Every row requires request-correlated telemetry. Forced speculation must " - "record positive chain steps for every request; adaptive rows must prove " - "the packed DSpark backend and startup profile were active.", - "", - "Mean/median are computed across paired fresh-process repeats. The oracle " - "uses max(AR, speculation) for goodput and min(AR, speculation) for TTFT " - "within each repeat.", - "", - "| Workload | C | Mode | N | Goodput mean/median | Oracle goodput " - "mean/median | Adaptive/oracle goodput mean/median | TTFT mean/median s | " - "Oracle/adaptive TTFT mean/median | Spec accepted/step | Spec steps | " - "Target forwards | Stable output |", - "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " - "---: | ---: | ---: | :---: |", - ] - regressions: list[str] = [] - variant_order = { - "ar": 0, - "speculation": 1, - "adaptive-on": 2, - "adaptive-confidence-off": 3, - } - for workload, clients, variant in sorted( - grouped, key=lambda key: (key[0], key[1], variant_order[key[2]]) - ): - items = grouped[(workload, clients, variant)] - by_repeat = {int(item["meta"]["repeat"]): item for item in items} - goodputs = [_positive_metric(item, "aggregate_tok_s") for item in items] - ttfts = [_positive_metric(item, "ttft_median_s") for item in items] - controls = { - name: grouped.get((workload, clients, name), []) - for name in ("ar", "speculation") - } - oracle_goodputs: list[float] = [] - oracle_ttfts: list[float] = [] - if controls["ar"] and controls["speculation"]: - control_maps = { - name: {int(item["meta"]["repeat"]): item for item in values} - for name, values in controls.items() - } - if control_maps["ar"].keys() != control_maps["speculation"].keys(): - raise ValueError(f"{workload} C={clients}: oracle repeat sets differ") - for repeat in sorted(control_maps["ar"]): - ar = control_maps["ar"][repeat] - speculation = control_maps["speculation"][repeat] - if _dspark_signature(ar) != _dspark_signature(speculation): - raise ValueError( - f"{workload} C={clients} repeat={repeat}: oracle metadata differs" - ) - if ( - ar["level"].get("selected_prompt_set_sha256") - != speculation["level"].get("selected_prompt_set_sha256") - ): - raise ValueError( - f"{workload} C={clients} repeat={repeat}: oracle prompts differ" - ) - ar_output = ar["level"].get("selected_output_set_sha256") - speculation_output = speculation["level"].get( - "selected_output_set_sha256" - ) - if not isinstance(ar_output, str) or ar_output != speculation_output: - raise ValueError( - f"{workload} C={clients} repeat={repeat}: " - "AR/speculation outputs differ" - ) - oracle_goodputs.append(max( - _positive_metric(ar, "aggregate_tok_s"), - _positive_metric(speculation, "aggregate_tok_s"), - )) - oracle_ttfts.append(min( - _positive_metric(ar, "ttft_median_s"), - _positive_metric(speculation, "ttft_median_s"), - )) - - goodput_ratio = "—" - ttft_ratio = "—" - if variant.startswith("adaptive-"): - if not oracle_goodputs: - raise ValueError(f"{workload} C={clients}: adaptive row lacks AR/spec oracle") - oracle_repeats = { - int(item["meta"]["repeat"]) for item in controls["ar"] - } - if by_repeat.keys() != oracle_repeats: - raise ValueError( - f"{workload} C={clients} {variant}: repeat set differs from oracle" - ) - reference = controls["ar"][0] - if any(_dspark_signature(item) != _dspark_signature(reference) for item in items): - raise ValueError( - f"{workload} C={clients} {variant}: metadata differs from oracle" - ) - if any( - item["level"].get("selected_prompt_set_sha256") - != reference["level"].get("selected_prompt_set_sha256") - for item in items - ): - raise ValueError( - f"{workload} C={clients} {variant}: prompts differ from oracle" - ) - ar_by_repeat = { - int(item["meta"]["repeat"]): item for item in controls["ar"] - } - for repeat, item in by_repeat.items(): - adaptive_output = item["level"].get("selected_output_set_sha256") - ar_output = ar_by_repeat[repeat]["level"].get( - "selected_output_set_sha256" - ) - if not isinstance(adaptive_output, str) or adaptive_output != ar_output: - raise ValueError( - f"{workload} C={clients} {variant} repeat={repeat}: " - "adaptive/AR outputs differ" - ) - gp_mean, gp_median = _pair(goodputs) - oracle_gp_mean, oracle_gp_median = _pair(oracle_goodputs) - ttft_mean, ttft_median = _pair(ttfts) - oracle_ttft_mean, oracle_ttft_median = _pair(oracle_ttfts) - gp_ratios = (gp_mean / oracle_gp_mean, gp_median / oracle_gp_median) - ttft_ratios = ( - oracle_ttft_mean / ttft_mean, - oracle_ttft_median / ttft_median, - ) - goodput_ratio = f"{gp_ratios[0]:.3f}/{gp_ratios[1]:.3f}" - ttft_ratio = f"{ttft_ratios[0]:.3f}/{ttft_ratios[1]:.3f}" - if ( - variant == "adaptive-on" - and min(*gp_ratios, *ttft_ratios) < ORACLE_THRESHOLD - ): - regressions.append( - f"{workload} C={clients} {variant}: goodput={goodput_ratio} " - f"ttft={ttft_ratio}" - ) - - aggregates = [item["proof"]["aggregate"] for item in items] - spec_steps = sum(value["spec_steps"] for value in aggregates) - accepted = sum(value["spec_accepted_tokens"] for value in aggregates) - accepted_per_step = accepted / spec_steps if spec_steps else 0.0 - target_forwards = median([value["target_forwards"] for value in aggregates]) - stable = output_stability(items) - lines.append( - f"| {workload} | {clients} | {variant} | {len(items)} | " - f"{_fmt_pair(goodputs)} | " - f"{_fmt_pair(oracle_goodputs) if oracle_goodputs else 'n/a'} | " - f"{goodput_ratio} | {_fmt_pair(ttfts, 3)} | {ttft_ratio} | " - f"{accepted_per_step:.2f} | {spec_steps} | {target_forwards:.0f} | " - f"{stable} |" - ) - - ablation_keys = sorted({ - (workload, clients) - for workload, clients, variant in grouped - if ( - variant == "adaptive-confidence-off" - and (workload, clients, "adaptive-on") in grouped - ) - }) - if ablation_keys: - lines += [ - "", - "## Confidence ablation", - "", - "Paired per-repeat deltas compare one-shot adaptive activation " - "against the draft-free confidence-off AR fallback. Positive values " - "favor per-request activation; these deltas are reported but are not " - "acceptance-gated.", - "", - "| Workload | C | Goodput delta mean/median | " - "Inverse-TTFT delta mean/median |", - "| :--- | ---: | ---: | ---: |", - ] - for workload, clients in ablation_keys: - confidence_items = grouped[(workload, clients, "adaptive-on")] - ablated_items = grouped[ - (workload, clients, "adaptive-confidence-off") - ] - confidence_by_repeat = { - int(item["meta"]["repeat"]): item for item in confidence_items - } - ablated_by_repeat = { - int(item["meta"]["repeat"]): item for item in ablated_items - } - if confidence_by_repeat.keys() != ablated_by_repeat.keys(): - raise ValueError( - f"{workload} C={clients}: confidence ablation repeats differ" - ) - goodput_deltas: list[float] = [] - inverse_ttft_deltas: list[float] = [] - for repeat in sorted(confidence_by_repeat): - enabled = confidence_by_repeat[repeat] - disabled = ablated_by_repeat[repeat] - goodput_deltas.append( - _positive_metric(enabled, "aggregate_tok_s") - / _positive_metric(disabled, "aggregate_tok_s") - - 1.0 - ) - inverse_ttft_deltas.append( - _positive_metric(disabled, "ttft_median_s") - / _positive_metric(enabled, "ttft_median_s") - - 1.0 - ) - lines.append( - f"| {workload} | {clients} | " - f"{_fmt_percent_pair(goodput_deltas)} | " - f"{_fmt_percent_pair(inverse_ttft_deltas)} |" - ) - lines += [ - "", - f"Oracle-relative gate: adaptive-on mean/median " - f"goodput and inverse TTFT ratio must be >= {ORACLE_THRESHOLD:.3f}.", - ] - if regressions: - raise ValueError( - "oracle-relative criterion failed: " + "; ".join(regressions) - ) - return "\n".join(lines) def summarize(reports: list[dict]) -> str: - dspark = [_is_dspark_item(item) for item in reports] - if any(dspark): - if not all(dspark): - raise ValueError("cannot mix Qwen3.6 feature and Qwen3.8 DSpark rows") - return summarize_dspark(reports) return summarize_qwen36(reports) diff --git a/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py index 232057dcc..d21be4450 100644 --- a/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py @@ -14,8 +14,8 @@ 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) +SCRIPT = HERE / "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) @@ -109,7 +109,7 @@ def test_boolean_wire_counts_are_not_accepted_as_integers(self) -> None: 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"] + argv = ["python3", "concurrent_benchmark.py", "--clients", "4"] result = benchmark.client_provenance(argv) self.assertEqual(result["client_argv"], argv) self.assertEqual(result["client_script"], str(SCRIPT.resolve())) diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index 173025057..b61a05f33 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -27,8 +27,7 @@ def load(name: str): return module -generator = load("generate_feature_prompts") -dspark_generator = load("generate_dspark_prompts") +generator = load("generate_ragged_prompts") proof = load("verify_feature_metrics") summary = load("summarize_feature_matrix") gate_analysis = load("analyze_gate_decisions") @@ -106,53 +105,6 @@ def test_activation_profiles_are_disjoint_and_above_thresholds(self) -> None: self.assertGreaterEqual(min(row["target_words"] for row in pressure), 12000) self.assertTrue(all(row["activation_target"] == "pflash-auto" for row in compression)) - def test_dspark_profiles_reuse_public_fixtures_and_build_north_star(self) -> None: - humaneval = dspark_generator.build_records("humaneval") - gsm8k = dspark_generator.build_records("gsm8k") - prose = dspark_generator.build_records("prose") - north_star = dspark_generator.build_records("north-star") - selection = dspark_generator.build_records("adaptive-selection") - self.assertGreaterEqual(len(humaneval), 8) - self.assertGreaterEqual(len(gsm8k), 8) - self.assertGreaterEqual(len(prose), 8) - self.assertEqual(len(north_star), 6) - self.assertEqual([row["suite"] for row in north_star[:2]], ["humaneval"] * 2) - self.assertEqual([row["suite"] for row in north_star[2:]], ["prose"] * 4) - self.assertEqual(len({row["prompt"] for row in north_star}), 6) - self.assertEqual(len(selection), 6) - self.assertEqual( - [row["selection_class"] for row in selection], - [ - "speculation_strong_win", - "speculation_strong_win", - "speculation_marginal_win", - "speculation_marginal_win", - "speculation_loss", - "speculation_loss", - ], - ) - self.assertEqual( - [row["expected_dense_oracle"] for row in selection], - ["speculation"] * 4 + ["ar"] * 2, - ) - selection_c3 = dspark_generator.build_records("adaptive-selection-c3") - self.assertEqual( - [row["id"] for row in selection_c3], - list(dspark_generator.ADAPTIVE_SELECTION_C3_IDS), - ) - self.assertEqual( - [row["selection_class"] for row in selection_c3], - [ - "speculation_strong_win", - "speculation_strong_win", - "speculation_loss", - ], - ) - self.assertEqual( - [row["expected_dense_oracle"] for row in selection_c3], - ["speculation", "speculation", "ar"], - ) - class FeatureRunnerShellTests(unittest.TestCase): def run_invalid_matrix( self, tmp: str, **overrides: str, @@ -216,34 +168,6 @@ def test_signals_exit_and_launch_environment_is_recorded(self) -> None: self.assertIn("'env ' > \"$case_dir/server-command.txt\"", runner) self.assertIn('"${launch_env[@]}" "${command[@]}"', runner) - def test_dspark_runner_has_explicit_confidence_ablation(self) -> None: - runner = (HERE / "run_qwen38_dspark_matrix.sh").read_text( - encoding="utf-8", - ) - self.assertIn('variants+=("adaptive-on")', runner) - self.assertIn('variants+=("adaptive-confidence-off")', runner) - self.assertIn('"DFLASH_SPEC_CONFIDENCE=1"', runner) - self.assertIn('"DFLASH_SPEC_CONFIDENCE=0"', runner) - self.assertIn( - 'metadata+=(--draft-always off --confidence "$confidence")', - runner, - ) - self.assertNotIn("ADAPTIVE_DRAFT_ALWAYS", runner) - self.assertNotIn("DFLASH_SPEC_DRAFT_ALWAYS", runner) - self.assertNotIn("RadixArk q4-mix GGUF", runner) - - def test_dspark_runner_profiles_only_the_measured_window(self) -> None: - runner = (HERE / "run_qwen38_dspark_matrix.sh").read_text( - encoding="utf-8", - ) - self.assertIn('STEP_TIMING="${STEP_TIMING:-1}"', runner) - self.assertIn('"DFLASH_STEP_TIMING=$STEP_TIMING"', runner) - self.assertIn('CACHE_TYPE_K="${CACHE_TYPE_K:-q8_0}"', runner) - self.assertIn('CACHE_TYPE_V="${CACHE_TYPE_V:-q8_0}"', runner) - self.assertIn('FA_WINDOW="${FA_WINDOW:-0}"', runner) - self.assertIn('"$case_dir/benchmark-server.log"', runner) - self.assertIn('--markdown-out "$OUT/profiling.md"', 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") @@ -271,6 +195,49 @@ def test_llama_only_does_not_require_lucebox_binary(self) -> None: class GateAnalysisTests(unittest.TestCase): + @staticmethod + def _scored( + request_id: int, + slot: int, + score: float, + expected_yield: float, + decision: str, + score_kind: str = "test_score", + hazards: list[float] | None = None, + ) -> dict: + return { + "request_id": request_id, + "slot": slot, + "activation_score": score, + "score_kind": score_kind, + "expected_yield": expected_yield, + "hazards": [] if hazards is None else hazards, + "evaluation": "scored", + "fallback_reason": None, + "decision_reason": ( + "selected_by_joint_goodput" + if decision == "speculation" else "ar_counterfactual_won" + ), + "decision": decision, + } + + @staticmethod + def _failed( + request_id: int, slot: int, score_kind: str = "unspecified", + ) -> dict: + return { + "request_id": request_id, + "slot": slot, + "activation_score": None, + "score_kind": score_kind, + "expected_yield": None, + "hazards": None, + "evaluation": "failed", + "fallback_reason": "activation_evaluation_failed", + "decision_reason": "evaluation_failed", + "decision": "ar", + } + @staticmethod def _write_activation_case( root: Path, @@ -373,15 +340,10 @@ def test_measured_step_timing_is_joined_and_summarized(self) -> None: "spec_accepted_tokens": 0, "spec_steps": 0, "target_forwards": 1, "output_tokens": 1, } - activation = { - "request_id": 7, "slot": 0, - "initial_confidence": 3.25, "expected_yield": 1.75, - "evaluation": "scored", "fallback_reason": None, - "decision": "ar", - } + activation = self._scored(7, 0, 3.25, 1.75, "ar") (case / "benchmark-server.log").write_text( - "[spec-gate] C=1 k=0 scores=[7:1.000/confidence] " - "sources=confidence:1,unavailable:0 " + "[spec-gate] C=1 k=0 scores=[7:1.000/fresh/test_score] " + "sources=fresh:1,initial:0,unavailable:0 " "G(k)=0.010000 G(0)=0.020000 " "predicted_cost=50.0us measured=ar-path\n" f"[spec-activation] {json.dumps(activation)}\n" @@ -406,25 +368,15 @@ def test_measured_step_timing_is_joined_and_summarized(self) -> None: self.assertEqual(report["activation"]["records"], 1) request = report["requests"][0] self.assertEqual(request["activation_slot"], 0) - self.assertEqual(request["initial_confidence"], 3.25) + self.assertEqual(request["activation_score"], 3.25) self.assertEqual(request["expected_yield"], 1.75) self.assertEqual(request["activation_decision"], "ar") self.assertEqual(request["activation_evaluation"], "scored") self.assertIsNone(request["activation_fallback_reason"]) def test_adaptive_on_activation_proof_fails_closed(self) -> None: - activation_7 = { - "request_id": 7, "slot": 0, - "initial_confidence": 2.0, "expected_yield": 1.25, - "evaluation": "scored", "fallback_reason": None, - "decision": "speculation", - } - activation_8 = { - "request_id": 8, "slot": 1, - "initial_confidence": 1.0, "expected_yield": 1.0, - "evaluation": "scored", "fallback_reason": None, - "decision": "ar", - } + activation_7 = self._scored(7, 0, 2.0, 1.25, "speculation") + activation_8 = self._scored(8, 1, 1.0, 1.0, "ar") activation_9 = {**activation_8, "request_id": 9} cases = ( ("missing", [activation_7], "missing activations.*8"), @@ -460,12 +412,7 @@ def test_confidence_off_is_exempt_from_activation_coverage(self) -> None: )) def test_adaptive_on_activation_proof_enforces_sticky_execution(self) -> None: - activation = { - "request_id": 7, "slot": 0, - "initial_confidence": 2.0, "expected_yield": 1.5, - "evaluation": "scored", "fallback_reason": None, - "decision": "speculation", - } + activation = self._scored(7, 0, 2.0, 1.5, "speculation") with tempfile.TemporaryDirectory() as tmp: case = self._write_activation_case( Path(tmp), "adaptive-on", [activation], @@ -501,19 +448,8 @@ def test_adaptive_on_activation_proof_enforces_sticky_execution(self) -> None: def test_failed_evaluation_activation_is_request_local_and_explicit( self, ) -> None: - failed = { - "request_id": 7, "slot": 0, - "initial_confidence": None, "expected_yield": None, - "evaluation": "failed", - "fallback_reason": "confidence_evaluation_failed", - "decision": "ar", - } - scored = { - "request_id": 8, "slot": 1, - "initial_confidence": 2.0, "expected_yield": 1.5, - "evaluation": "scored", "fallback_reason": None, - "decision": "speculation", - } + failed = self._failed(7, 0) + scored = self._scored(8, 1, 2.0, 1.5, "speculation") with tempfile.TemporaryDirectory() as tmp: case = self._write_activation_case( Path(tmp), "adaptive-on", [failed, scored], @@ -529,41 +465,22 @@ def test_failed_evaluation_activation_is_request_local_and_explicit( } self.assertEqual(by_id[7]["activation_decision"], "ar") self.assertEqual(by_id[7]["activation_evaluation"], "failed") - self.assertIsNone(by_id[7]["initial_confidence"]) + self.assertIsNone(by_id[7]["activation_score"]) self.assertEqual( by_id[7]["activation_fallback_reason"], - "confidence_evaluation_failed", - ) - self.assertEqual( - by_id[8]["activation_evaluation"], "scored", + "activation_evaluation_failed", ) + self.assertEqual(by_id[8]["activation_evaluation"], "scored") + def test_dflash2_activation_fields_and_gate_score_kind_are_preserved( self, ) -> None: - scored = { - "request_id": 7, "slot": 0, - "initial_confidence": None, - "activation_score": 5.3306, - "request_benefit": 5.3306, - "score_kind": "dflash2_selector_benefit_v1", - "expected_yield": 5.3306, - "evaluation": "scored", "fallback_reason": None, - "decision_reason": "selected_by_joint_goodput", - "decision": "speculation", - } - failed = { - "request_id": 8, "slot": 1, - "initial_confidence": None, - "activation_score": None, - "request_benefit": None, - "score_kind": "dflash2_selector_benefit_v1", - "expected_yield": None, - "evaluation": "failed", - "fallback_reason": "benefit_evaluation_failed", - "decision_reason": "evaluation_failed", - "decision": "ar", - } + score_kind = "qwen38-dflash2-selector-benefit-v1" + scored = self._scored( + 7, 0, 5.3306, 5.3306, "speculation", score_kind, [0.1, 0.2], + ) + failed = self._failed(8, 1, score_kind) with tempfile.TemporaryDirectory() as tmp: case = self._write_activation_case( Path(tmp), "adaptive-on", [scored, failed], @@ -572,72 +489,56 @@ def test_dflash2_activation_fields_and_gate_score_kind_are_preserved( log_path = case / "benchmark-server.log" log_path.write_text( "[spec-gate] C=2 k=1 " - "scores=[7:5.331/confidence/dflash2_selector_benefit_v1*," - "8:2.685/initial/dflash2_selector_benefit_v1] " + "scores=[7:5.331/fresh/qwen38-dflash2-selector-benefit-v1*," + "8:2.685/initial/qwen38-dflash2-selector-benefit-v1] " "G(k)=0.010 G(0)=0.009 predicted_cost=1us measured=2us\n" + log_path.read_text(encoding="utf-8"), encoding="utf-8", ) report = gate_analysis.analyze_case(case) self.assertEqual( - report["activation"]["score_kind_counts"], - {"dflash2_selector_benefit_v1": 2}, + report["activation"]["score_kind_counts"], {score_kind: 2}, ) self.assertEqual( report["activation"]["fallback_reason_counts"], - {"benefit_evaluation_failed": 1}, + {"activation_evaluation_failed": 1}, ) by_id = { row["engine_request_id"]: row for row in report["requests"] } request = by_id[7] - self.assertIsNone(request["initial_confidence"]) self.assertAlmostEqual(request["activation_score"], 5.3306) - self.assertAlmostEqual(request["request_benefit"], 5.3306) - self.assertEqual( - request["activation_score_kind"], - "dflash2_selector_benefit_v1", - ) + self.assertEqual(request["activation_hazards"], [0.1, 0.2]) + self.assertEqual(request["activation_score_kind"], score_kind) self.assertEqual( request["activation_decision_reason"], "selected_by_joint_goodput", ) self.assertAlmostEqual(request["mean_activation_score"], 5.331) - self.assertIsNone(request["mean_confidence_yield"]) rounds, _, _, _ = gate_analysis.parse_server_log(log_path) - self.assertEqual( - rounds[0]["entries"][0]["score_kind"], - "dflash2_selector_benefit_v1", - ) + self.assertEqual(rounds[0]["entries"][0]["score_kind"], score_kind) - def test_typed_fallback_reasons_are_nonempty_and_recognized(self) -> None: - base = { - "request_id": 7, "slot": 0, - "initial_confidence": None, - "activation_score": None, - "request_benefit": None, - "score_kind": "unspecified", - "expected_yield": None, - "evaluation": "failed", - "decision_reason": "evaluation_failed", - "decision": "ar", - } - accepted = ( - "benefit_adapter_unavailable", - "benefit_adapter_invalid_config", - "cost_profile_unavailable", - ) - for reason in accepted: - with self.subTest(reason=reason), tempfile.TemporaryDirectory() as tmp: - row = {**base, "fallback_reason": reason} - path = Path(tmp) / "server.log" - path.write_text( - f"[spec-activation] {json.dumps(row)}\n", encoding="utf-8", - ) - _, _, _, activations = gate_analysis.parse_server_log(path) - self.assertEqual(activations, [row]) - for reason in ("", "unknown_failure"): + + def test_failed_activation_uses_one_generic_fallback_reason(self) -> None: + base = self._failed(7, 0) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(base)}\n", encoding="utf-8", + ) + _, _, _, activations = gate_analysis.parse_server_log(path) + self.assertEqual(activations, [base]) + no_adapter = {**base, "decision_reason": "no_speculator_adapter"} + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + f"[spec-activation] {json.dumps(no_adapter)}\n", + encoding="utf-8", + ) + _, _, _, activations = gate_analysis.parse_server_log(path) + self.assertEqual(activations, [no_adapter]) + for reason in ("", "benefit_evaluation_failed", "unknown_failure"): with self.subTest(reason=reason), tempfile.TemporaryDirectory() as tmp: row = {**base, "fallback_reason": reason} path = Path(tmp) / "server.log" @@ -648,37 +549,30 @@ def test_typed_fallback_reasons_are_nonempty_and_recognized(self) -> None: gate_analysis.parse_server_log(path) + def test_activation_record_fields_are_strictly_validated(self) -> None: - valid = { - "request_id": 7, "slot": 0, - "initial_confidence": 1.0, "expected_yield": 1.0, - "evaluation": "scored", "fallback_reason": None, - "decision": "ar", - } - failed = { - "request_id": 7, "slot": 0, - "initial_confidence": None, "expected_yield": None, - "evaluation": "failed", - "fallback_reason": "confidence_evaluation_failed", - "decision": "ar", - } + valid = self._scored(7, 0, 1.0, 1.0, "ar") + failed = self._failed(7, 0) missing_evaluation = dict(valid) missing_evaluation.pop("evaluation") missing_fallback_reason = dict(valid) missing_fallback_reason.pop("fallback_reason") missing_failed_score = dict(failed) - missing_failed_score.pop("initial_confidence") + missing_failed_score.pop("activation_score") cases = ( ({**valid, "request_id": True}, "request_id"), ({**valid, "slot": -1}, "slot"), - ({**valid, "initial_confidence": 0.99}, "initial_confidence"), + ({**valid, "activation_score": 0.99}, "activation_score"), ({**valid, "expected_yield": math.nan}, "expected_yield"), + ({**valid, "hazards": [1.1]}, "hazards"), + ({**valid, "score_kind": ""}, "score_kind"), ({**valid, "decision": "undecided"}, "decision"), (missing_evaluation, "evaluation"), (missing_fallback_reason, "fallback_reason"), - (missing_failed_score, "initial_confidence"), + (missing_failed_score, "activation_score"), ({**valid, "fallback_reason": "failure"}, "fallback_reason"), - ({**failed, "initial_confidence": 1.0}, "scores must be null"), + ({**failed, "activation_score": 1.0}, "scores must be null"), + ({**failed, "hazards": []}, "scores must be null"), ({**failed, "decision": "speculation"}, "decision must be ar"), ({**failed, "fallback_reason": "draft_failed"}, "fallback_reason"), ) @@ -698,6 +592,7 @@ def test_activation_record_fields_are_strictly_validated(self) -> None: _, _, _, activations = gate_analysis.parse_server_log(path) self.assertEqual(activations, [failed]) + def test_paired_controls_produce_a_per_prompt_concurrent_oracle(self) -> None: base_request = { "prompt_id": "p1", @@ -1040,7 +935,7 @@ def test_forced_chain_requires_steps_and_matching_startup_proof(self) -> None: row["spec_steps"] = 3 row["spec_accepted_tokens"] = 2 startup = ( - "[parallel-dspark] enabled width=7 mode=packed-chain-verify " + "[parallel-chain] speculator=test-score " "decode_mode=speculation draft=q4-mix-compatible" ) result = proof.verify(input_report, rows, {"chain"}, startup) @@ -1070,8 +965,8 @@ def test_adaptive_chain_allows_ar_argmax_but_requires_profile(self) -> None: row["ddtree_accepted_tokens"] = 0 startup = "\n".join(( "[spec-profile] context=4096 reps=5 mode=batched-draft", - "[parallel-dspark] enabled width=7 mode=packed-chain-verify " - "decode_mode=adaptive draft=q4-mix-compatible", + "[parallel-chain] speculator=test-score " + "", )) result = proof.verify(input_report, rows, {"chain"}, startup) self.assertTrue(result["valid"], result["errors"]) @@ -1137,59 +1032,6 @@ def item( }, }, } - @staticmethod - def dspark_item( - variant: str, goodput: float, ttft: float, *, repeat: int = 1, - ) -> dict: - mode = ( - "adaptive" if variant.startswith("adaptive-") else variant - ) - if variant.startswith("adaptive-"): - draft_always = "off" - confidence = ( - "off" if variant == "adaptive-confidence-off" else "on" - ) - else: - draft_always = None - confidence = None - spec_steps = 0 if variant == "ar" else 8 - return { - "path": Path(f"/tmp/{variant}-r{repeat}/bench.json"), - "report": { - "max_tokens": 256, "ignore_eos": True, - "temperature": 0.0, "seed": 1, - }, - "meta": { - "workload": "humaneval", "variant": variant, - "repeat": repeat, "model_sha256": "a" * 64, - "server_binary_sha256": "b" * 64, - "feature_config": { - "decode_mode": mode, - "draft_always": draft_always, - "confidence": confidence, - "draft_model_sha256": "c" * 64, - }, - }, - "level": { - "clients": 6, - "aggregate_tok_s": goodput, - "ttft_median_s": ttft, - "selected_prompt_set_sha256": "same-prompts", - "selected_output_set_sha256": "same-output", - }, - "proof": { - "aggregate": { - "ddtree_steps": 0, - "ddtree_suspensions": 0, - "ddtree_accepted_tokens": 0, - "spec_steps": spec_steps, - "spec_accepted_tokens": spec_steps * 2, - "target_forwards": 16, - }, - }, - } - - 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) @@ -1284,48 +1126,5 @@ def test_unstable_ar_control_suppresses_feature_delta(self) -> None: row = next(line for line in text.splitlines() if "| full |" in line) self.assertEqual(row.split("|")[7].strip(), "n/a") - def test_dspark_summary_enforces_mean_median_oracle_gate(self) -> None: - reports = [ - self.dspark_item("ar", 100.0, 1.0, repeat=1), - self.dspark_item("ar", 102.0, 1.1, repeat=2), - self.dspark_item("speculation", 98.0, 0.9, repeat=1), - self.dspark_item("speculation", 103.0, 1.0, repeat=2), - self.dspark_item("adaptive-on", 100.0, 0.9, repeat=1), - self.dspark_item("adaptive-on", 103.0, 1.0, repeat=2), - self.dspark_item("adaptive-confidence-off", 75.0, 1.8, repeat=1), - self.dspark_item("adaptive-confidence-off", 80.0, 1.7, repeat=2), - ] - text = summary.summarize(reports) - self.assertIn("Qwen3.8 DSpark adaptive concurrency matrix", text) - self.assertIn("Oracle-relative gate", text) - self.assertIn("| adaptive-on |", text) - self.assertIn("| adaptive-confidence-off |", text) - self.assertIn("## Confidence ablation", text) - self.assertIn("Positive values favor per-request activation", text) - - def test_dspark_summary_rejects_eager_draft_metadata(self) -> None: - adaptive = self.dspark_item("adaptive-on", 100.0, 0.9) - adaptive["meta"]["feature_config"]["draft_always"] = "on" - with self.assertRaisesRegex( - ValueError, "adaptive metadata does not match variant adaptive-on", - ): - summary.summarize([ - self.dspark_item("ar", 100.0, 1.0), - self.dspark_item("speculation", 98.0, 0.9), - adaptive, - ]) - - def test_dspark_summary_rejects_oracle_regression(self) -> None: - reports = [ - self.dspark_item("ar", 100.0, 1.0), - self.dspark_item("speculation", 98.0, 0.9), - self.dspark_item("adaptive-on", 90.0, 1.5), - ] - with self.assertRaisesRegex( - ValueError, "oracle-relative criterion failed.*adaptive-on", - ): - summary.summarize(reports) - - if __name__ == "__main__": unittest.main() diff --git a/harness/benchmarks/concurrency/verify_feature_metrics.py b/harness/benchmarks/concurrency/verify_feature_metrics.py index b184b49e1..0bd98cc27 100644 --- a/harness/benchmarks/concurrency/verify_feature_metrics.py +++ b/harness/benchmarks/concurrency/verify_feature_metrics.py @@ -19,7 +19,7 @@ "kvflash_reselects", ) DECODE_MODES = ("ar", "speculation", "adaptive") -DSPARK_STARTUP_PREFIX = "[parallel-dspark] enabled" +SPECULATOR_STARTUP_PREFIX = "[parallel-chain] speculator=" SPEC_PROFILE_PREFIX = "[spec-profile] context=" REQUIRED_KEYS = ( "request_id", "effective_prompt_tokens", *COUNTERS, @@ -200,16 +200,14 @@ def verify( startup_mode = ( isinstance(decode_mode, str) and re.search( - rf"^.*{re.escape(DSPARK_STARTUP_PREFIX)}.*" - rf"decode_mode={re.escape(decode_mode)}\b.*" - r"draft=q4-mix-compatible.*$", + rf"^.*{re.escape(SPECULATOR_STARTUP_PREFIX)}\S+.*$", log_text, flags=re.MULTILINE, ) ) if not startup_mode: errors.append( - "chain requested but matching packed DSpark startup proof is missing" + "chain requested but matching speculator adapter startup proof is missing" ) if decode_mode == "adaptive" and SPEC_PROFILE_PREFIX not in log_text: errors.append( @@ -331,7 +329,7 @@ def verify( ) return { - "schema_version": 4, + "schema_version": 5, "decode_mode": decode_mode, "expected_features": sorted(expected), "valid": not errors, From 07a37a1790ddde5f1ba56bb5179a4d2c1366ffa4 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Thu, 20 Aug 2026 11:21:39 +0000 Subject: [PATCH 40/42] perf(qwen35): commit mixed speculation directly --- .../deps/llama.cpp/ggml/include/ggml-cuda.h | 24 + .../ggml/src/ggml-cuda/gated_delta_net.cu | 75 ++- .../src/ggml-cuda/gdn-transition-journal.cu | 317 +++++++++++ .../ggml/src/ggml-cuda/paged-attn.cu | 37 +- server/deps/llama.cpp/ggml/src/ggml.c | 10 +- server/src/common/concurrency/seq_engine.h | 10 + .../src/common/speculation/speculation_gate.h | 57 +- server/src/common/step_graph.h | 18 + server/src/internal.h | 11 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 515 +++++++++++++++--- server/src/qwen35/graph_builders.cpp | 65 ++- server/src/qwen35/graph_builders.h | 20 +- server/src/qwen35/qwen35_target_graph.cpp | 158 ++++-- server/src/server/scheduler.cpp | 3 + server/test/test_gdn_transition_journal.cpp | 109 ++++ server/test/test_seq_batch_plan.cpp | 13 + server/test/test_speculation_gate.cpp | 76 ++- 17 files changed, 1334 insertions(+), 184 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml-cuda.h b/server/deps/llama.cpp/ggml/include/ggml-cuda.h index c8d62e21a..a418a9ab7 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -106,6 +106,30 @@ GGML_BACKEND_API bool ggml_backend_cuda_gdn_transition_journal_commit( const struct ggml_tensor * accepted_prefixes, const struct ggml_tensor * active_slot_ids); +// Batched concurrent-tree commit. Validation is fail-closed before any kernel +// launches; all layer journals and convolution windows commit on one device +// synchronization. +GGML_BACKEND_API bool ggml_backend_cuda_gdn_transition_journal_commit_many( + const struct ggml_tensor * const * journals, + struct ggml_tensor * const * states, + const struct ggml_tensor * const * conv_inputs, + struct ggml_tensor * const * conv_states, + int n_layers, + const struct ggml_tensor * accepted_prefixes, + const struct ggml_tensor * active_slot_ids); + +// Promote accepted packed-tree K/V scratch rows into pager-owned rows. +GGML_BACKEND_API bool ggml_backend_cuda_tree_cache_commit_many( + struct ggml_tensor * const * caches, int n_caches, + const struct ggml_tensor * commit_rows, + const struct ggml_tensor * active_slot_ids, + int tree_scratch_base, int tree_scratch_stride); + +// Promote accepted BF16 tree feature rows into slot-local feature rings. +GGML_BACKEND_API bool ggml_backend_cuda_tree_feature_commit( + const struct ggml_tensor * source, struct ggml_tensor * destination, + const struct ggml_tensor * destination_rows); + // Attach learned per-expert decode tables to a mixed-precision tensor. The // host variants copy the tables to the device that owns `base`. Call the // matching unregister function before releasing the tensor's backing buffer. diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu index 1489877cd..50a7a4184 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gated_delta_net.cu @@ -328,7 +328,7 @@ gated_delta_net_cuda(const float * q, } } -template +template __global__ void __launch_bounds__(WARP_THREADS * 8, 2) gated_delta_net_cuda_grouped_cols(const float * q, const float * k, @@ -337,9 +337,11 @@ gated_delta_net_cuda_grouped_cols(const float * q, const float * beta, const float * curr_state, const int * active_slot_ids, + const int * parent_ids, float * dst, float * state_out, InterT * persist_inter, + float * transition_journal, int64_t H, int64_t n_tokens, int64_t n_seqs, @@ -391,7 +393,7 @@ gated_delta_net_cuda_grouped_cols(const float * q, n_seqs, n_state_slots, physical_sequence, physical_state_offset); InterT * inter_states = nullptr; InterT * inter_base = nullptr; - if constexpr (WRITE_INTER) { + if constexpr (WRITE_INTER || TREE_MODE) { inter_states = persist_inter ? persist_inter : (InterT *)(dst + attn_score_elems + final_state_elems); @@ -401,6 +403,10 @@ gated_delta_net_cuda_grouped_cols(const float * q, const float * curr_state_seq = physical_sequence >= 0 ? curr_state + physical_state_offset : nullptr; + const int * parent_ids_seq = nullptr; + if constexpr (TREE_MODE) { + parent_ids_seq = parent_ids + sequence * n_tokens; + } attn_data += (sequence * n_tokens * H + h_idx) * S_v; float state_shard[COLS][rows_per_lane]; @@ -417,6 +423,39 @@ gated_delta_net_cuda_grouped_cols(const float * q, } for (int t = 0; t < n_tokens; ++t) { + if constexpr (TREE_MODE) { + if (t > 0) { + const int parent_t = parent_ids_seq[t]; + if (parent_t == GGML_GDN_TREE_ROOT_PARENT) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + state_shard[c][r] = curr_state_seq + ? curr_state_seq[col * S_v + row] + : 0.0f; + } + } + } else if (parent_t != t - 1) { + const InterT * parent_base = inter_states + + ((sequence * n_tokens + parent_t) * H + h_idx) * + S_v * S_v; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + state_shard[c][r] = load_inter_state( + parent_base, col * S_v + row); + } + } + } + } + } + const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; @@ -440,6 +479,12 @@ gated_delta_net_cuda_grouped_cols(const float * q, g_val = __shfl_sync(0xffffffffU, g_val, 0); beta_val = __shfl_sync(0xffffffffU, beta_val, 0); + constexpr int journal_width = 2*S_v + 1; + float * journal_t = transition_journal + ? transition_journal + + ((sequence * n_tokens + t) * H + h_idx) * journal_width + : nullptr; + float k_reg[rows_per_lane]; float q_reg[rows_per_lane]; float kv_partial[COLS]; @@ -456,12 +501,20 @@ gated_delta_net_cuda_grouped_cols(const float * q, const float k_val = k_t[row]; q_reg[r] = q_val; k_reg[r] = k_val; + if (journal_t && blockIdx.z == 0 && threadIdx.y == 0 && + subgroup == 0) { + journal_t[1 + row] = k_val; + } #pragma unroll for (int c = 0; c < COLS; ++c) { kv_partial[c] += state_shard[c][r] * k_val; } } + if (journal_t && blockIdx.z == 0 && threadIdx.y == 0 && + subgroup == 0 && lane == 0) { + journal_t[0] = g_val; + } float delta[COLS]; #pragma unroll @@ -470,6 +523,9 @@ gated_delta_net_cuda_grouped_cols(const float * q, float delta_val = 0.0f; if (lane == 0) { delta_val = (v_t[col_base + c] - g_val * kv_col) * beta_val; + if (journal_t) { + journal_t[1 + S_v + col_base + c] = delta_val; + } } delta[c] = gdn_subgroup_broadcast_lane0(delta_val, WIDTH); } @@ -502,7 +558,7 @@ gated_delta_net_cuda_grouped_cols(const float * q, } } - if constexpr (WRITE_INTER) { + if constexpr (WRITE_INTER || TREE_MODE) { #pragma unroll for (int c = 0; c < COLS; ++c) { const int col = col_base + c; @@ -586,8 +642,8 @@ static void launch_gated_delta_net( break; } case 128: { - if constexpr (!KDA && !TREE_MODE) { - if (transition_journal_d == nullptr && use_grouped_cols && + if constexpr (!KDA) { + if (use_grouped_cols && ((GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) || GGML_CUDA_CC_IS_AMD(cc))) { constexpr int cols = 4; @@ -598,16 +654,16 @@ static void launch_gated_delta_net( constexpr int groups_per_warp = 32 / width; dim3 grouped_grid_dims(H, n_seqs, (groups + column_groups_per_block * groups_per_warp - 1) / (column_groups_per_block * groups_per_warp)); dim3 grouped_block_dims(32, column_groups_per_block, 1); - gated_delta_net_cuda_grouped_cols<128, cols, width, 32, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, persist_inter_d, H, + gated_delta_net_cuda_grouped_cols<128, cols, width, 32, TREE_MODE, WRITE_INTER, InterT><<>>( + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, parent_ids_d, dst_d, state_out_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } else if (warp_size == 64) { constexpr int groups_per_warp = 64 / width; dim3 grouped_grid_dims(H, n_seqs, (groups + column_groups_per_block * groups_per_warp - 1) / (column_groups_per_block * groups_per_warp)); dim3 grouped_block_dims(64, column_groups_per_block, 1); - gated_delta_net_cuda_grouped_cols<128, cols, width, 64, WRITE_INTER, InterT><<>>( - q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, dst_d, state_out_d, persist_inter_d, H, + gated_delta_net_cuda_grouped_cols<128, cols, width, 64, TREE_MODE, WRITE_INTER, InterT><<>>( + q_d, k_d, v_d, g_d, b_d, s_d, active_slot_ids_d, parent_ids_d, dst_d, state_out_d, persist_inter_d, transition_journal_d, H, n_tokens, n_seqs, n_state_slots, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, gate_bias, gate_A); } else { @@ -731,7 +787,6 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * } if (src_transition_journal) { const int64_t journal_width = kda ? 3*S_v : 2*S_v + 1; - GGML_ASSERT(!src_parent); GGML_ASSERT(src_transition_journal->type == GGML_TYPE_F32); GGML_ASSERT(ggml_is_contiguous(src_transition_journal)); GGML_ASSERT(src_transition_journal->ne[0] == journal_width); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu index 114b760fb..9ca1cdd9f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu @@ -182,3 +182,320 @@ extern "C" bool ggml_backend_cuda_gdn_transition_journal_commit( if (cudaGetLastError() != cudaSuccess) return false; return cudaDeviceSynchronize() == cudaSuccess; } + +namespace { + +__global__ void gdn_conv_journal_commit_kernel( + const float * conv_input, + float * conv_state, + const int32_t * accepted_prefixes, + const int32_t * active_slot_ids, + int window, + int channels, + int n_tokens, + int n_seqs, + int n_state_slots) { + const int sequence = blockIdx.y; + const int element = blockIdx.x * blockDim.x + threadIdx.x; + const int count = window * channels; + if (sequence >= n_seqs || element >= count) return; + const int slot = active_slot_ids[sequence]; + const int accepted = accepted_prefixes[sequence]; + if (slot < 0 || slot >= n_state_slots || + accepted < 0 || accepted > n_tokens) return; + const int k = element % window; + const int channel = element / window; + const size_t source = + ((size_t) sequence * channels + channel) * (window + n_tokens) + + accepted + k; + const size_t destination = + ((size_t) slot * channels + channel) * window + k; + conv_state[destination] = conv_input[source]; +} + +__global__ void tree_cache_commit_kernel( + uint8_t * cache, + const int64_t * commit_rows, + const int32_t * active_slot_ids, + size_t row_bytes, + size_t head_stride, + int n_heads, + int tree_width, + int n_seqs, + int n_cache_rows, + int scratch_base, + int scratch_stride) { + const int byte = blockIdx.x * blockDim.x + threadIdx.x; + const int flat = blockIdx.y; + const int head = blockIdx.z; + if ((size_t) byte >= row_bytes || flat >= tree_width*n_seqs || + head >= n_heads) return; + const int lane = flat / tree_width; + const int node = flat % tree_width; + const int slot = active_slot_ids[lane]; + const int64_t destination_row = commit_rows[flat]; + if (slot < 0 || destination_row < 0 || + destination_row >= n_cache_rows) return; + const int64_t source_row = + (int64_t) scratch_base + (int64_t) slot*scratch_stride + node; + if (source_row < 0 || source_row >= n_cache_rows) return; + const size_t source = + (size_t) head*head_stride + (size_t) source_row*row_bytes + byte; + const size_t destination = + (size_t) head*head_stride + (size_t) destination_row*row_bytes + byte; + cache[destination] = cache[source]; +} + +__global__ void tree_feature_commit_kernel( + const uint8_t * source, + uint8_t * destination, + const int32_t * destination_rows, + size_t row_bytes, + int n_rows, + int destination_capacity) { + const int byte = blockIdx.x * blockDim.x + threadIdx.x; + const int source_row = blockIdx.y; + if ((size_t) byte >= row_bytes || source_row >= n_rows) return; + const int destination_row = destination_rows[source_row]; + if (destination_row < 0 || destination_row >= destination_capacity) return; + destination[(size_t) destination_row*row_bytes + byte] = + source[(size_t) source_row*row_bytes + byte]; +} + +bool same_device_pointer(const void * pointer, int expected_device) { + int pointer_device = -1; + return device_pointer(pointer, pointer_device) && + pointer_device == expected_device; +} + +} // namespace + +extern "C" bool ggml_backend_cuda_gdn_transition_journal_commit_many( + const ggml_tensor * const * journals, + ggml_tensor * const * states, + const ggml_tensor * const * conv_inputs, + ggml_tensor * const * conv_states, + int n_layers, + const ggml_tensor * accepted_prefixes, + const ggml_tensor * active_slot_ids) { + if (!journals || !states || !conv_inputs || !conv_states || + n_layers <= 0 || !accepted_prefixes || !active_slot_ids || + accepted_prefixes->type != GGML_TYPE_I32 || + active_slot_ids->type != GGML_TYPE_I32 || + !ggml_is_contiguous(accepted_prefixes) || + !ggml_is_contiguous(active_slot_ids)) return false; + + const int64_t n_seqs = ggml_nelements(accepted_prefixes); + if (n_seqs < 1 || ggml_nelements(active_slot_ids) != n_seqs) return false; + int device = -1; + if (!device_pointer(accepted_prefixes->data, device) || + !same_device_pointer(active_slot_ids->data, device)) return false; + + std::vector accepted((size_t) n_seqs); + std::vector slots((size_t) n_seqs); + const size_t map_bytes = (size_t) n_seqs*sizeof(int32_t); + if (cudaMemcpy(accepted.data(), accepted_prefixes->data, map_bytes, + cudaMemcpyDeviceToHost) != cudaSuccess || + cudaMemcpy(slots.data(), active_slot_ids->data, map_bytes, + cudaMemcpyDeviceToHost) != cudaSuccess) return false; + + int common_tokens = -1; + int common_state_slots = -1; + std::vector seen; + for (int layer = 0; layer < n_layers; ++layer) { + const ggml_tensor * journal = journals[layer]; + ggml_tensor * state = states[layer]; + const ggml_tensor * conv_input = conv_inputs[layer]; + ggml_tensor * conv_state = conv_states[layer]; + if (!journal || !state || !conv_input || !conv_state || + journal->type != GGML_TYPE_F32 || state->type != GGML_TYPE_F32 || + conv_input->type != GGML_TYPE_F32 || conv_state->type != GGML_TYPE_F32 || + !ggml_is_contiguous(journal) || !ggml_is_contiguous(state) || + !ggml_is_contiguous(conv_input) || !ggml_is_contiguous(conv_state)) return false; + const int64_t state_size = state->ne[0]; + const int64_t heads = state->ne[2]; + const int64_t tokens = journal->ne[2]; + const int64_t state_slots = state->ne[3]; + if ((state_size != 16 && state_size != 32 && + state_size != 64 && state_size != 128) || + state->ne[1] != state_size || heads < 1 || + journal->ne[1] != heads || journal->ne[3] != n_seqs || + (journal->ne[0] != 2*state_size + 1 && + journal->ne[0] != 3*state_size) || tokens < 1 || + conv_state->ne[0] < 1 || conv_state->ne[1] < 1 || + conv_state->ne[2] != state_slots || conv_state->ne[3] != 1 || + conv_input->ne[0] != conv_state->ne[0] + tokens || + conv_input->ne[1] != conv_state->ne[1] || + conv_input->ne[2] != n_seqs || conv_input->ne[3] != 1) return false; + if (common_tokens < 0) { + common_tokens = (int) tokens; + common_state_slots = (int) state_slots; + seen.assign((size_t) state_slots, 0); + } else if (tokens != common_tokens || state_slots != common_state_slots) { + return false; + } + const void * pointers[] = { + journal->data, state->data, conv_input->data, conv_state->data, + }; + for (const void * pointer : pointers) { + if (!same_device_pointer(pointer, device)) return false; + } + } + for (int64_t lane = 0; lane < n_seqs; ++lane) { + if (accepted[(size_t) lane] < 0 || + accepted[(size_t) lane] > common_tokens) return false; + const int slot = slots[(size_t) lane]; + if (slot < 0 || slot >= common_state_slots) continue; + if (seen[(size_t) slot]) return false; + seen[(size_t) slot] = 1; + } + + ggml_cuda_set_device(device); + constexpr int threads = 256; + (void) cudaGetLastError(); + for (int layer = 0; layer < n_layers; ++layer) { + const ggml_tensor * journal = journals[layer]; + ggml_tensor * state = states[layer]; + const int state_size = (int) state->ne[0]; + const int heads = (int) state->ne[2]; + const int tokens = (int) journal->ne[2]; + const int journal_width = (int) journal->ne[0]; + const int64_t state_elements = (int64_t) state_size*state_size; + const dim3 state_grid( + (unsigned int) ((state_elements + threads - 1)/threads), + (unsigned int) heads, (unsigned int) n_seqs); + gdn_transition_journal_commit_kernel<<>>( + (const float *) journal->data, (float *) state->data, + (const int32_t *) accepted_prefixes->data, + (const int32_t *) active_slot_ids->data, + state_size, heads, tokens, (int) n_seqs, + (int) state->ne[3], journal_width, + journal_width == 2*state_size + 1 ? 1 : state_size); + + const ggml_tensor * conv_input = conv_inputs[layer]; + ggml_tensor * conv_state = conv_states[layer]; + const int conv_elements = + (int) (conv_state->ne[0]*conv_state->ne[1]); + const dim3 conv_grid( + (unsigned int) ((conv_elements + threads - 1)/threads), + (unsigned int) n_seqs, 1); + gdn_conv_journal_commit_kernel<<>>( + (const float *) conv_input->data, (float *) conv_state->data, + (const int32_t *) accepted_prefixes->data, + (const int32_t *) active_slot_ids->data, + (int) conv_state->ne[0], (int) conv_state->ne[1], tokens, + (int) n_seqs, (int) conv_state->ne[2]); + } + if (cudaGetLastError() != cudaSuccess) return false; + return cudaDeviceSynchronize() == cudaSuccess; +} + +extern "C" bool ggml_backend_cuda_tree_cache_commit_many( + ggml_tensor * const * caches, + int n_caches, + const ggml_tensor * commit_rows, + const ggml_tensor * active_slot_ids, + int tree_scratch_base, + int tree_scratch_stride) { + if (!caches || n_caches <= 0 || !commit_rows || !active_slot_ids || + commit_rows->type != GGML_TYPE_I64 || + active_slot_ids->type != GGML_TYPE_I32 || + !ggml_is_contiguous(commit_rows) || + !ggml_is_contiguous(active_slot_ids) || + commit_rows->ne[0] < 1 || commit_rows->ne[1] < 1 || + ggml_nelements(active_slot_ids) != commit_rows->ne[1] || + tree_scratch_base < 0 || tree_scratch_stride < commit_rows->ne[0]) return false; + const int tree_width = (int) commit_rows->ne[0]; + const int n_seqs = (int) commit_rows->ne[1]; + const int n_rows = tree_width*n_seqs; + int device = -1; + if (!device_pointer(commit_rows->data, device) || + !same_device_pointer(active_slot_ids->data, device)) return false; + + std::vector destinations((size_t) n_rows); + std::vector slots((size_t) n_seqs); + if (cudaMemcpy(destinations.data(), commit_rows->data, + destinations.size()*sizeof(int64_t), cudaMemcpyDeviceToHost) != cudaSuccess || + cudaMemcpy(slots.data(), active_slot_ids->data, + slots.size()*sizeof(int32_t), cudaMemcpyDeviceToHost) != cudaSuccess) return false; + int cache_rows = -1; + for (int index = 0; index < n_caches; ++index) { + ggml_tensor * cache = caches[index]; + if (!cache || !ggml_is_contiguous(cache) || cache->ne[0] < 1 || + cache->ne[1] < 1 || cache->ne[2] < 1 || cache->ne[3] != 1 || + cache->nb[1] < ggml_row_size(cache->type, cache->ne[0]) || + !same_device_pointer(cache->data, device)) return false; + if (cache_rows < 0) cache_rows = (int) cache->ne[1]; + else if (cache->ne[1] != cache_rows) return false; + } + for (int lane = 0; lane < n_seqs; ++lane) { + const int slot = slots[(size_t) lane]; + if (slot < 0) continue; + const int64_t source_end = (int64_t) tree_scratch_base + + (int64_t) slot*tree_scratch_stride + tree_width; + if (source_end > cache_rows) return false; + for (int node = 0; node < tree_width; ++node) { + const int64_t destination = + destinations[(size_t) lane*tree_width + node]; + if (destination < -1 || destination >= cache_rows) return false; + } + } + + ggml_cuda_set_device(device); + constexpr int threads = 256; + (void) cudaGetLastError(); + for (int index = 0; index < n_caches; ++index) { + ggml_tensor * cache = caches[index]; + const dim3 grid( + (unsigned int) ((cache->nb[1] + threads - 1)/threads), + (unsigned int) n_rows, (unsigned int) cache->ne[2]); + tree_cache_commit_kernel<<>>( + (uint8_t *) cache->data, + (const int64_t *) commit_rows->data, + (const int32_t *) active_slot_ids->data, + cache->nb[1], cache->nb[2], (int) cache->ne[2], + tree_width, n_seqs, cache_rows, + tree_scratch_base, tree_scratch_stride); + } + if (cudaGetLastError() != cudaSuccess) return false; + return cudaDeviceSynchronize() == cudaSuccess; +} + +extern "C" bool ggml_backend_cuda_tree_feature_commit( + const ggml_tensor * source, + ggml_tensor * destination, + const ggml_tensor * destination_rows) { + if (!source || !destination || !destination_rows || + source->type != destination->type || source->type != GGML_TYPE_BF16 || + destination_rows->type != GGML_TYPE_I32 || + !ggml_is_contiguous(source) || !ggml_is_contiguous(destination) || + !ggml_is_contiguous(destination_rows) || + source->ne[0] != destination->ne[0] || source->ne[2] != 1 || + source->ne[3] != 1 || destination->ne[2] != 1 || + destination->ne[3] != 1 || + ggml_nelements(destination_rows) != source->ne[1] || + source->nb[1] != destination->nb[1]) return false; + int device = -1; + if (!device_pointer(source->data, device) || + !same_device_pointer(destination->data, device) || + !same_device_pointer(destination_rows->data, device)) return false; + const int n_rows = (int) source->ne[1]; + std::vector rows((size_t) n_rows); + if (cudaMemcpy(rows.data(), destination_rows->data, + rows.size()*sizeof(int32_t), cudaMemcpyDeviceToHost) != cudaSuccess) return false; + for (int row : rows) { + if (row < -1 || row >= destination->ne[1]) return false; + } + ggml_cuda_set_device(device); + constexpr int threads = 256; + const dim3 grid( + (unsigned int) ((source->nb[1] + threads - 1)/threads), + (unsigned int) n_rows, 1); + (void) cudaGetLastError(); + tree_feature_commit_kernel<<>>( + (const uint8_t *) source->data, (uint8_t *) destination->data, + (const int32_t *) destination_rows->data, + source->nb[1], n_rows, (int) destination->ne[1]); + if (cudaGetLastError() != cudaSuccess) return false; + return cudaDeviceSynchronize() == cudaSuccess; +} 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 a881b93d3..ad321d094 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 @@ -237,6 +237,7 @@ static __global__ void paged_attn_decode( int32_t block_size, int32_t min_partitions, int32_t tree_width, + int32_t tree_row_offset, int32_t tree_scratch_base, int32_t tree_scratch_stride, float scale) { @@ -265,28 +266,29 @@ static __global__ void paged_attn_decode( 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; const int32_t query_pos = query_positions ? *(const int32_t *) (query_positions + (int64_t) seq * qpos_nb0) : -1; + const bool tree_query = tree_mode && seq >= tree_row_offset; + const int32_t tree_seq = tree_query + ? (seq - tree_row_offset) / tree_width : 0; + const int32_t query_node = tree_query + ? seq - tree_row_offset - tree_seq * tree_width : -1; + const int32_t tree_size = tree_query + ? *(const int32_t *) ( + tree_sizes + (int64_t) tree_seq * tree_size_nb0) + : 0; // A row is live when its slot id selects a real block-table column and, // 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) && - (!tree_mode || + (!query_positions || tree_query || query_pos >= 0) && + (!tree_query || (tree_size >= 0 && tree_size <= tree_width && query_node < tree_size)); const int32_t physical_seq = valid_query ? physical_seq_raw : 0; @@ -296,7 +298,7 @@ static __global__ void paged_attn_decode( : 0; // 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) { + if (query_positions && !tree_query && query_pos < kv_seq_len_raw) { kv_seq_len_raw = query_pos + 1; } const int64_t table_capacity = @@ -310,7 +312,7 @@ static __global__ void paged_attn_decode( // 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) + ? kv_seq_len + (tree_query ? tree_width : 0) : 0; const int32_t n_logical_blocks = (virtual_tokens + block_size - 1) / block_size; @@ -453,7 +455,7 @@ static __global__ void paged_attn_decode( phys_mine = physical_block * block_size + my_token % block_size; } - } else if (tree_mode && my_token < token_end) { + } else if (tree_query && my_token < token_end) { const int32_t candidate = my_token - kv_seq_len; if (paged_attn_tree_visible( parent_ids, parent_nb0, parent_nb1, @@ -708,7 +710,7 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { } const bool tree_mode = parent_ids || tree_sizes; if ((parent_ids == nullptr) != (tree_sizes == nullptr) || - (tree_mode && (!active_slot_ids || query_positions))) { + (tree_mode && !active_slot_ids)) { return false; } @@ -819,7 +821,9 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { 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 || + q->ne[1] < parent_ids->ne[1] * tree_width || + (!query_positions && + q->ne[1] != parent_ids->ne[1] * tree_width) || (int64_t) max_kv_seq_len + tree_width > INT32_MAX) { return false; } @@ -1074,6 +1078,9 @@ static bool try_launch_paged_attn( block_size, min_partitions, tree_width, + parent_ids + ? (int32_t)(q->ne[1] - parent_ids->ne[1] * tree_width) + : 0, tree_scratch_base, tree_scratch_stride, scale); diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 6f51347a2..d6290babd 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5728,7 +5728,10 @@ struct ggml_tensor * ggml_paged_attn_ext( 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); + // Mixed direct-commit batches use causal positions for a compact AR + // prefix and -1 for the fixed-width tree tail. Pure trees keep this null. + GGML_ASSERT(!tree_mode || query_positions == NULL || + query_positions->ne[0] == q->ne[1]); GGML_ASSERT(!tree_mode || parent_ids->type == GGML_TYPE_I32); GGML_ASSERT(!tree_mode || tree_sizes->type == GGML_TYPE_I32); @@ -5779,7 +5782,9 @@ struct ggml_tensor * ggml_paged_attn_ext( 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); + const int64_t tree_rows = parent_ids->ne[1] * tree_width; + GGML_ASSERT(q->ne[1] >= tree_rows); + GGML_ASSERT(query_positions || q->ne[1] == tree_rows); // Every physical sequence slot owns one non-overlapping scratch slab. // Bound the largest address with int64 arithmetic before the GPU sees @@ -6806,7 +6811,6 @@ void ggml_gated_delta_net_set_transition_journal( struct ggml_tensor * journal) { GGML_ASSERT(tensor != NULL && journal != NULL); GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); - GGML_ASSERT(tensor->src[6] == NULL); GGML_ASSERT(journal->type == GGML_TYPE_F32); GGML_ASSERT(ggml_is_contiguous(journal)); diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 62147caa8..ddd2e2cf6 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -212,6 +212,10 @@ class SeqEngine { uint64_t ddtree_suspensions = 0; uint64_t spec_steps = 0; uint64_t spec_accepted_tokens = 0; + // Sticky-SPEC lane advances performed by an explicit packed + // AR+prefill service round. These are scheduling suspensions, not a + // routing-mode change, and are reported separately from spec_steps. + uint64_t spec_service_ar_steps = 0; uint64_t target_forwards = 0; uint64_t kvflash_page_ins = 0; uint64_t kvflash_page_outs = 0; @@ -363,6 +367,8 @@ inline std::string validate_step_result( return "failed decode carries DDTree suspension telemetry"; if (output.spec_steps != 0 || output.spec_accepted_tokens != 0) return "failed decode carries chain speculation telemetry"; + if (output.spec_service_ar_steps != 0) + return "failed decode carries chain service telemetry"; } else { if (output.token < 0) return "successful decode has no pending token"; @@ -380,6 +386,10 @@ inline std::string validate_step_result( return "DDTree suspension has no successful DDTree step"; if (output.spec_accepted_tokens != 0 && output.spec_steps == 0) return "chain acceptance has no successful speculation step"; + if (output.spec_service_ar_steps > output.target_forwards) + return "chain service steps exceed target forwards"; + if (output.spec_service_ar_steps != 0 && output.spec_steps != 0) + return "step mixes chain speculation and AR service"; } decode_seen[(size_t)output.slot] = 1; } diff --git a/server/src/common/speculation/speculation_gate.h b/server/src/common/speculation/speculation_gate.h index 3efd879a0..c5ecdb8ab 100644 --- a/server/src/common/speculation/speculation_gate.h +++ b/server/src/common/speculation/speculation_gate.h @@ -22,7 +22,7 @@ namespace dflash::common { struct SpecGateConfig { double cost_ema_alpha = 0.20; - double adaptive_gain_margin = 0.02; + double adaptive_gain_margin = 0.01; }; struct SpecCostLookup { @@ -91,8 +91,9 @@ struct SpecCandidate { bool scoreable = false; // can_speculate is also request-lifetime (for example, false for an // unsupported sampler or thinking hook). The min-token EOS policy is - // enforced inside the speculative path, and an incompatible prefill graph - // is deferred, so neither changes this capability or the chosen mode. + // enforced inside the speculative path. An incompatible prefill graph may + // suspend speculative execution for one explicitly telemetered AR service + // round without changing this capability or the chosen request mode. bool can_speculate = false; // While an adaptive request is Undecided, NaN requests its one-time // bootstrap and a finite value is the preferred activation measurement. @@ -287,17 +288,20 @@ class SpeculationGate { SpeculationGate(SpecCostTables costs, SpecStepGeometry geometry, int max_accept, ClampLogger clamp_logger = {}, - SpecGateConfig config = {}) + SpecGateConfig config = {}, + bool direct_commit = false) : config_(config), costs_(std::move(costs)), geometry_(std::move(geometry)), max_accept_(std::max(1, max_accept)), + direct_commit_(direct_commit), clamp_logger_(std::move(clamp_logger)) {} SpeculationGate(SpecGateConfig config, SpecCostTables costs, SpecStepGeometry geometry, int max_accept, - ClampLogger clamp_logger = {}) + ClampLogger clamp_logger = {}, + bool direct_commit = false) : SpeculationGate(std::move(costs), std::move(geometry), max_accept, - std::move(clamp_logger), config) {} + std::move(clamp_logger), config, direct_commit) {} bool valid() const { auto valid_alpha = [](double value) { @@ -311,9 +315,12 @@ class SpeculationGate { } // draft_lanes_override prices always-drafting. -1 means admitted-only. + // Non-committing plans retain a first score while leaving the request + // Undecided so a later, formed cohort can publish the sticky decision. SpecPlan plan(int concurrency, const std::vector & candidates, - int k_cap, int draft_lanes_override = -1) { + int k_cap, int draft_lanes_override = -1, + bool commit_decisions = true) { SpecPlan out; out.concurrency = concurrency; if (!valid() || concurrency < 0 || @@ -390,7 +397,8 @@ class SpeculationGate { if (!out.pending_evaluations.empty()) { undecided.clear(); } - out.decisions_committed = out.pending_evaluations.empty(); + out.decisions_committed = + out.pending_evaluations.empty() && commit_decisions; if (out.decisions_committed) { for (const Ranked & item : forced_ar) { request_states_[item.candidate->request_id].decision = @@ -464,9 +472,17 @@ class SpeculationGate { ? draft_lanes_override : k; if (k > 0) { - tree_rows = geometry_.tree_rows(k); - expected_step_rows = geometry_.expected_step_rows( - concurrency, k, expected_sum); + if (direct_commit_) { + // Direct promotion packs one AR row per peer before the + // bucketed fixed-depth speculative tree and performs no + // target replay. + tree_rows = geometry_.tree_rows(k) + concurrency - k; + expected_step_rows = 0.0; + } else { + tree_rows = geometry_.tree_rows(k); + expected_step_rows = geometry_.expected_step_rows( + concurrency, k, expected_sum); + } } const CostPrice price = price_expected_shape( {concurrency, k, tree_rows, 0, draft_lanes}, @@ -505,9 +521,11 @@ class SpeculationGate { if (!ranked[i].commit_candidate) continue; const SpecDecision committed = static_cast(i) < best.k ? SpecDecision::Speculation : SpecDecision::AR; - request_states_[ranked[i].candidate->request_id].decision = - committed; - out.ordered[i].decision = committed; + if (commit_decisions) { + request_states_[ranked[i].candidate->request_id].decision = + committed; + out.ordered[i].decision = committed; + } } for (int i = 0; i < best.k; ++i) { out.ordered[(size_t)i].admitted = true; @@ -541,7 +559,9 @@ class SpeculationGate { if (!std::isfinite(measured_us) || measured_us <= 0.0 || executed.concurrency < 0 || executed.admitted_count < 0 || executed.admitted_count > executed.concurrency || - executed.tree_rows < 0 || executed.step_rows <= 0 || + executed.tree_rows < 0 || executed.step_rows < 0 || + (executed.step_rows == 0 && + (!direct_commit_ || executed.admitted_count == 0)) || executed.draft_lanes < 0) { return; } @@ -667,8 +687,12 @@ class SpeculationGate { if (shape.admitted_count > 0) { add("tree", costs_.tree_cost.lookup(shape.tree_rows)); + if (!direct_commit_) { + add("step", costs_.step_cost.lookup(shape.step_rows)); + } + } else { + add("step", costs_.step_cost.lookup(shape.step_rows)); } - add("step", costs_.step_cost.lookup(shape.step_rows)); if (shape.draft_lanes > 0) { add("draft", costs_.draft_cost.lookup(shape.draft_lanes)); } @@ -722,6 +746,7 @@ class SpeculationGate { SpecCostTables costs_; SpecStepGeometry geometry_; int max_accept_ = 1; + bool direct_commit_ = false; static constexpr double kCostScaleMin = 0.25; static constexpr double kCostScaleMax = 4.0; std::unordered_map request_states_; diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index 5e1814a72..d11a2a208 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -20,6 +20,8 @@ struct StepGraph { ggml_context * ctx = nullptr; ggml_cgraph * gf = nullptr; ggml_gallocr_t alloc = nullptr; + ggml_context * commit_ctx = nullptr; + ggml_backend_buffer_t commit_buffer = nullptr; // Persistent metadata arena for the draft graph. Reusing the same arena // across rebuilds keeps every ggml_tensor at a stable address, which is @@ -66,6 +68,11 @@ struct StepGraph { // 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; + // Packed-tree direct-commit metadata uploaded after posterior selection. + ggml_tensor * accepted_prefixes = nullptr; // [n_tree_seqs] i32 + ggml_tensor * commit_slot_ids = nullptr; // [n_tree_seqs] i32 + ggml_tensor * commit_rows = nullptr; // [tree_width,n_tree_seqs] i64 + ggml_tensor * feature_commit_rows = nullptr; // same shape, i32 // 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; @@ -83,12 +90,18 @@ struct StepGraph { // Per-delta-net-layer captures (verify only). std::vector delta_captures; + ggml_tensor * tree_features = nullptr; std::vector moe_selected; }; // Reset the per-call graph state (ctx + graph + tensor handles) but KEEP the // persistent CUDA buffer in `sg.alloc` alive across steps. inline void step_graph_free(StepGraph & sg) { + if (sg.commit_buffer) { + ggml_backend_buffer_free(sg.commit_buffer); + sg.commit_buffer = nullptr; + } + if (sg.commit_ctx) { ggml_free(sg.commit_ctx); sg.commit_ctx = nullptr; } if (sg.ctx) { ggml_free(sg.ctx); sg.ctx = nullptr; } sg.gf = nullptr; sg.inp_embed = sg.positions = sg.attn_mask = nullptr; @@ -105,6 +118,10 @@ inline void step_graph_free(StepGraph & sg) { sg.paged_query_seq_ids = nullptr; sg.paged_query_positions = nullptr; sg.target_feat_rows = nullptr; + sg.accepted_prefixes = nullptr; + sg.commit_slot_ids = nullptr; + sg.commit_rows = nullptr; + sg.feature_commit_rows = nullptr; sg.logits_row_indices = nullptr; sg.logits = nullptr; sg.hidden_states = nullptr; @@ -116,6 +133,7 @@ inline void step_graph_free(StepGraph & sg) { sg.hot_local_lut = nullptr; sg.valid_lut = nullptr; sg.delta_captures.clear(); + sg.tree_features = nullptr; sg.moe_selected.clear(); } diff --git a/server/src/internal.h b/server/src/internal.h index e4a6ea878..7db863adf 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -678,6 +678,10 @@ bool migrate_prefill_cache(const TargetWeights & w, struct DeltaNetCapture { ggml_tensor * ssm_intermediate_states = nullptr; ggml_tensor * conv_input = nullptr; + // Concurrent tree direct-commit data. The compact journal plus the + // tree conv input can advance accepted recurrent prefixes without a + // second target-model forward. These are graph-owned outputs. + ggml_tensor * transition_journal = nullptr; }; // One contiguous prompt chunk on the flattened token axis of a concurrent @@ -699,6 +703,7 @@ struct QwenGraphInputs { int kv_start; // position where the new tokens begin bool capture_layers; // if true, write captured layer features into cache.target_feat bool capture_delta_intermediate = false; // if true, populate out_delta_captures + bool capture_tree_commit = false; // compact recurrent journal + tree features 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 @@ -764,6 +769,10 @@ 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; + // Mixed direct-commit tree graphs place this many one-token mapped AR + // sequences before the fixed-width speculative tree segment. Their slot + // IDs share active_slot_ids/state_slot_ids with the tree lanes. + int mapped_ar_seqs = 0; int seq_slot = 0; int paged_max_kv_len = 0; int n_prefill_tokens = 0; @@ -788,6 +797,8 @@ struct QwenGraphOutputs { // views marked as ggml_set_output() so their data persists after // graph_compute; the spec-decode loop reads them host-side for rollback. std::vector delta_captures; + // BF16 [n_capture_layers*n_embd, n_tokens], packed-tree only. + ggml_tensor * tree_features = nullptr; // One entry per target layer. Populated only when capture_moe_router is // true; qwen35 dense layers and non-MoE models leave entries null. std::vector moe_selected; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 2433f14e3..d6a3431d6 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -18,6 +18,7 @@ #include "common/concurrency/chain_spec_shapes.h" #include "common/speculation/spec_cost_profile.h" #include "internal.h" +#include "ggml-cuda.h" #include #include @@ -41,6 +42,11 @@ int decode_bucket_width(int live_count) { return 64; } +bool chain_direct_commit_enabled() { + const char * value = std::getenv("DFLASH_CHAIN_DURABLE_REPLAY"); + return !(value && std::atoi(value) != 0); +} + double initial_prediction_realized_tokens( const SpecPlan & plan, const SeqEngine::StepResult & result) { double realized = 0.0; @@ -693,7 +699,7 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { std::fprintf(stderr, "[spec-gate] %s_cost index %d outside profile; clamped to %d\n", table, requested, profiled); - }); + }, SpecGateConfig{}, chain_direct_commit_enabled()); if (!speculation_gate_->valid()) { speculation_gate_.reset(); return false; @@ -1309,41 +1315,98 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( proposals.push_back(std::move(proposal)); } const int spec_count = static_cast(proposals.size()); - const int tree_bucket = spec_count > 0 - ? chain_decode_bucket_width(spec_count) : 0; + const bool direct_commit = chain_direct_commit_enabled(); auto lane_disposition = [&](size_t i) { return chain_lane_disposition( active_admitted[i] != 0, !proposal_errors[i].empty()); }; + std::vector ar_lanes; + ar_lanes.reserve(inputs.size() - static_cast(spec_count)); + std::vector ar_for_input(inputs.size(), -1); + for (size_t i = 0; i < inputs.size(); ++i) { + if (lane_disposition(i) != ChainLaneDisposition::AR) continue; + ArLane lane; + lane.input_index = i; + lane.slot = inputs[i].slot; + lane.token = inputs[i].token; + lane.position = slots_.slot(lane.slot).cur_pos; + ar_for_input[i] = static_cast(ar_lanes.size()); + ar_lanes.push_back(lane); + } + const int ar_count = static_cast(ar_lanes.size()); + const int tree_lane_count = spec_count; + const int tree_bucket = tree_lane_count > 0 + ? chain_decode_bucket_width(tree_lane_count) : 0; + + // Compact direct graphs write AR K/V and recurrent state durably in the + // fused launch. Allocate their physical rows now, but keep history and + // cur_pos staged until the target graph and all promotions succeed. + if (direct_commit) { + for (ArLane & lane : ar_lanes) { + const Qwen35SlotManager::StepAppend app = + slots_.append_token(lane.slot, lane.token); + if (!app.ok) { + result.error = app.busy + ? "paged KV pool exhausted during compact AR staging" + : "compact AR K/V staging failed"; + return result; + } + const bool table_ok = slots_.residency_active() || + app.new_block < 0 || + upload_block_table_delta( + lane.slot, app.new_block_index, &app.new_block, 1); + if (!table_ok) { + result.error = "compact AR block-table update failed"; + return result; + } + lane.position = app.position; + lane.physical_row = app.physical_row; + } + } + StepGraph & tree_sg = b_.sg_; + std::vector posterior; + int replay_total = 0; - if (spec_count > 0) { + if (tree_lane_count > 0) { // Launch 1: scratch-only packed path-tree verification. - StepGraph & tree_sg = b_.sg_; int max_prefix = 1; for (const Proposal & proposal : proposals) { max_prefix = std::max( max_prefix, slots_.slot(proposal.slot).cur_pos); } + if (direct_commit) { + for (const ArLane & ar : ar_lanes) { + max_prefix = std::max(max_prefix, ar.position + 1); + } + } t_verify_build_start = timing_clock::now(); if (!build_target_step_paged_tree( tree_sg, b_.w_, b_.cache_, b_.target_backend_, V, tree_bucket, max_prefix, tree_scratch_base_, tree_scratch_stride_, - b_.cfg_.kq_stride_pad)) { + b_.cfg_.kq_stride_pad, direct_commit ? ar_count : 0, + direct_commit)) { result.error = "packed chain speculation verify graph build failed"; return result; } t_verify_build_end = timing_clock::now(); - const int total_tree = V * tree_bucket; + const int spec_tree_rows = V * tree_bucket; + const int spec_row_offset = direct_commit ? ar_count : 0; + const int total_tree = spec_row_offset + spec_tree_rows; std::vector flat_tokens(static_cast(total_tree), 0); - std::vector parents(static_cast(total_tree), -1); + std::vector parents( + static_cast(spec_tree_rows), -1); std::vector sizes(static_cast(tree_bucket), 0); - std::vector tree_slots(static_cast(tree_bucket), -1); + const int mapped_slot_count = spec_row_offset + tree_bucket; + std::vector tree_slots( + static_cast(mapped_slot_count), -1); std::vector tree_state_slots( - static_cast(tree_bucket), 0); + static_cast(mapped_slot_count), 0); std::vector query_slots(static_cast(total_tree), -1); + std::vector query_positions( + direct_commit ? static_cast(total_tree) : 0, -1); std::vector tree_rows( static_cast(total_tree) * n_head_kv, scratch_row_); std::vector tree_positions( @@ -1352,19 +1415,45 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( static_cast(hidden) * total_tree, 0.0f); seq_lens_.assign(static_cast(n_slots), 0); + if (direct_commit) { + for (int ar_index = 0; ar_index < ar_count; ++ar_index) { + const ArLane & ar = ar_lanes[static_cast(ar_index)]; + const int row = ar_index; + tree_slots[static_cast(ar_index)] = ar.slot; + tree_state_slots[static_cast(ar_index)] = ar.slot; + seq_lens_[static_cast(ar.slot)] = ar.position + 1; + flat_tokens[static_cast(row)] = ar.token; + query_slots[static_cast(row)] = ar.slot; + query_positions[static_cast(row)] = ar.position; + tree_positions[static_cast(0) * total_tree + row] = + ar.position; + tree_positions[static_cast(1) * total_tree + row] = + ar.position; + tree_positions[static_cast(2) * total_tree + row] = + ar.position; + for (int head = 0; head < n_head_kv; ++head) { + tree_rows[static_cast(head) * total_tree + row] = + ar.physical_row; + } + } + } + for (int lane = 0; lane < spec_count; ++lane) { const Proposal & proposal = proposals[static_cast(lane)]; - const int base = lane * V; + const int tree_base = lane * V; + const int row_base = spec_row_offset + tree_base; + const int mapped_lane = spec_row_offset + lane; sizes[static_cast(lane)] = V; - tree_slots[static_cast(lane)] = proposal.slot; - tree_state_slots[static_cast(lane)] = proposal.slot; + tree_slots[static_cast(mapped_lane)] = proposal.slot; + tree_state_slots[static_cast(mapped_lane)] = proposal.slot; seq_lens_[static_cast(proposal.slot)] = slots_.slot(proposal.slot).cur_pos; for (int node = 0; node < V; ++node) { - const int row = base + node; + const int tree_row = tree_base + node; + const int row = row_base + node; flat_tokens[static_cast(row)] = proposal.flat[static_cast(node)]; - parents[static_cast(row)] = node == 0 + parents[static_cast(tree_row)] = node == 0 ? -1 : proposal.tree.parents[static_cast(node)]; query_slots[static_cast(row)] = proposal.slot; const int depth = node == 0 @@ -1410,6 +1499,11 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( ggml_backend_tensor_set( tree_sg.paged_query_seq_ids, query_slots.data(), 0, sizeof(int32_t) * query_slots.size()); + if (tree_sg.paged_query_positions) { + 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.kv_write_rows, tree_rows.data(), 0, sizeof(int64_t) * tree_rows.size()); ggml_backend_tensor_set( @@ -1422,7 +1516,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( } t_verify_exec_end = timing_clock::now(); - std::vector posterior(static_cast(total_tree), -1); + posterior.assign(static_cast(total_tree), -1); ggml_backend_tensor_get( tree_sg.argmax_tokens, posterior.data(), 0, sizeof(int32_t) * posterior.size()); @@ -1431,7 +1525,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( for (int lane = 0; lane < spec_count; ++lane) { Proposal & proposal = proposals[static_cast(lane)]; const int32_t * lane_posterior = - posterior.data() + static_cast(lane) * V; + posterior.data() + static_cast(spec_row_offset + lane * V); proposal.accepted = follow_verified_tree( proposal.tree, lane_posterior, proposal.verify_bonus); const int room = @@ -1543,40 +1637,32 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( max_kv_len = std::max(max_kv_len, seq_len); } - std::vector ar_lanes; - ar_lanes.reserve(inputs.size() - static_cast(spec_count)); - std::vector ar_for_input(inputs.size(), -1); - for (size_t i = 0; i < inputs.size(); ++i) { - if (lane_disposition(i) != ChainLaneDisposition::AR) continue; - const StepInput & in = inputs[i]; - const Qwen35SlotManager::StepAppend app = - slots_.append_token(in.slot, in.token); - if (!app.ok) { - result.error = app.busy - ? "paged KV pool exhausted during mixed AR commit" - : "mixed AR K/V append failed"; - return result; - } - 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) { - result.error = "mixed AR block-table update failed"; - return result; + for (int ar_index = 0; ar_index < ar_count; ++ar_index) { + ArLane & lane = ar_lanes[static_cast(ar_index)]; + const StepInput & in = inputs[lane.input_index]; + if (!direct_commit) { + const Qwen35SlotManager::StepAppend app = + slots_.append_token(in.slot, in.token); + if (!app.ok) { + result.error = app.busy + ? "paged KV pool exhausted during mixed AR commit" + : "mixed AR K/V append failed"; + return result; + } + 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) { + result.error = "mixed AR block-table update failed"; + return result; + } + lane.position = app.position; + lane.physical_row = app.physical_row; } - ArLane lane; - lane.input_index = i; - lane.slot = in.slot; - lane.token = in.token; - lane.position = app.position; - lane.physical_row = app.physical_row; - ar_for_input[i] = static_cast(ar_lanes.size()); - ar_lanes.push_back(lane); - seq_lens_[static_cast(in.slot)] = app.position + 1; - max_kv_len = std::max(max_kv_len, app.position + 1); + seq_lens_[static_cast(in.slot)] = lane.position + 1; + max_kv_len = std::max(max_kv_len, lane.position + 1); } - const int ar_count = static_cast(ar_lanes.size()); if (spec_count == 0 && ar_count == 0) { result.decode.reserve(inputs.size()); for (size_t i = 0; i < inputs.size(); ++i) { @@ -1596,6 +1682,261 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( } t_commit_end = timing_clock::now(); + if (direct_commit) { + const int spec_tree_rows = V * tree_bucket; + const int spec_row_offset = ar_count; + const int total_tree = spec_row_offset + spec_tree_rows; + std::vector accepted_prefixes( + static_cast(tree_bucket), 0); + std::vector commit_rows( + static_cast(spec_tree_rows), -1); + std::vector feature_commit_rows( + static_cast(total_tree), -1); + const int feature_cap = b_.cache_.target_feat_cap; + int replay_cursor = 0; + for (int lane = 0; lane < spec_count; ++lane) { + const Proposal & proposal = proposals[static_cast(lane)]; + if (proposal.path.size() != proposal.accepted.size()) { + result.error = "direct commit path/acceptance size mismatch"; + return result; + } + accepted_prefixes[static_cast(lane)] = + static_cast(proposal.path.size()); + for (size_t depth = 0; depth < proposal.accepted.size(); ++depth) { + const int node = proposal.accepted[depth]; + // DFlash2 proposals are chains. A branching tree needs an + // indexed journal-commit kernel rather than prefix commit. + if (node != static_cast(depth)) { + result.error = + "direct commit requires a contiguous chain acceptance"; + return result; + } + const int flat = lane * V + node; + const int source_row = spec_row_offset + flat; + if (replay_cursor >= static_cast(replay_physical.size())) { + result.error = "direct commit replay cursor overflow"; + return result; + } + commit_rows[static_cast(flat)] = + replay_physical[static_cast(replay_cursor)]; + feature_commit_rows[static_cast(source_row)] = + proposal.slot * feature_cap + + replay_positions[static_cast(replay_cursor)] % + feature_cap; + ++replay_cursor; + } + } + if (replay_cursor != replay_total) { + result.error = "direct commit replay cursor mismatch"; + return result; + } + for (int ar_index = 0; ar_index < ar_count; ++ar_index) { + const ArLane & ar = ar_lanes[static_cast(ar_index)]; + feature_commit_rows[static_cast(ar_index)] = + ar.slot * feature_cap + ar.position % feature_cap; + } + std::vector commit_slots( + static_cast(tree_bucket), -1); + for (int lane = 0; lane < spec_count; ++lane) { + commit_slots[static_cast(lane)] = + proposals[static_cast(lane)].slot; + } + ggml_backend_tensor_set( + tree_sg.accepted_prefixes, accepted_prefixes.data(), 0, + accepted_prefixes.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + tree_sg.commit_slot_ids, commit_slots.data(), 0, + commit_slots.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + tree_sg.commit_rows, commit_rows.data(), 0, + commit_rows.size() * sizeof(int64_t)); + ggml_backend_tensor_set( + tree_sg.feature_commit_rows, feature_commit_rows.data(), 0, + feature_commit_rows.size() * sizeof(int32_t)); + + 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 || + !tree_sg.tree_features || !b_.cache_.target_feat) { + result.error = "direct commit capture set incomplete"; + return result; + } + std::vector journals; + std::vector states; + std::vector conv_inputs; + std::vector conv_states; + journals.reserve(n_delta); + states.reserve(n_delta); + conv_inputs.reserve(n_delta); + conv_states.reserve(n_delta); + for (size_t layer = 0; layer < n_delta; ++layer) { + const DeltaNetCapture & capture = + tree_sg.delta_captures[layer]; + if (!capture.transition_journal || !capture.conv_input || + !b_.cache_.ssm_state[layer] || + !b_.cache_.conv_state[layer]) { + result.error = "direct commit layer capture incomplete"; + return result; + } + journals.push_back(capture.transition_journal); + states.push_back(b_.cache_.ssm_state[layer]); + conv_inputs.push_back(capture.conv_input); + conv_states.push_back(b_.cache_.conv_state[layer]); + } + std::vector cache_tensors; + cache_tensors.reserve( + b_.cache_.attn_k.size() + b_.cache_.attn_v.size()); + for (ggml_tensor * tensor : b_.cache_.attn_k) { + if (tensor) cache_tensors.push_back(tensor); + } + for (ggml_tensor * tensor : b_.cache_.attn_v) { + if (tensor) cache_tensors.push_back(tensor); + } + if (cache_tensors.empty()) { + result.error = "direct commit K/V cache set empty"; + return result; + } + + t_replay_build_end = timing_clock::now(); + if (!ggml_backend_cuda_tree_cache_commit_many( + cache_tensors.data(), + static_cast(cache_tensors.size()), + tree_sg.commit_rows, tree_sg.commit_slot_ids, + tree_scratch_base_, tree_scratch_stride_)) { + result.error = "direct commit K/V promotion failed"; + return result; + } + if (!ggml_backend_cuda_tree_feature_commit( + tree_sg.tree_features, b_.cache_.target_feat, + tree_sg.feature_commit_rows)) { + result.error = "direct commit target-feature promotion failed"; + return result; + } + if (!ggml_backend_cuda_gdn_transition_journal_commit_many( + journals.data(), states.data(), conv_inputs.data(), + conv_states.data(), static_cast(n_delta), + tree_sg.accepted_prefixes, tree_sg.commit_slot_ids)) { + result.error = "direct commit recurrent-state promotion failed"; + return result; + } + ggml_backend_tensor_set( + b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + t_replay_exec_end = timing_clock::now(); + + std::vector write_slots; + write_slots.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + if (chain_lane_executes(lane_disposition(i))) { + write_slots.push_back(inputs[i].slot); + } + } + if (!commit_residency_writes(write_slots)) { + result.error = "direct commit residency write failed"; + return result; + } + for (size_t i = 0; i < inputs.size(); ++i) { + if (chain_lane_executes(lane_disposition(i))) { + slots_.commit_step(inputs[i].slot); + } + } + for (int lane = 0; lane < spec_count; ++lane) { + Proposal & proposal = proposals[static_cast(lane)]; + const int graph_row = spec_row_offset + lane * V + + proposal.accepted.back(); + proposal.pending = sample_graph_row( + proposal.slot, graph_row, + &posterior[static_cast(graph_row)], &logits_buf_); + if (proposal.pending < 0) { + result.error = "direct commit speculative sampling failed"; + return result; + } + } + for (int ar_index = 0; ar_index < ar_count; ++ar_index) { + ArLane & ar = ar_lanes[static_cast(ar_index)]; + const int graph_row = ar_index; + ar.pending = sample_graph_row( + ar.slot, graph_row, + &posterior[static_cast(graph_row)], &logits_buf_); + if (ar.pending < 0) { + result.error = "direct commit AR sampling failed"; + return result; + } + } + t_sample_end = timing_clock::now(); + + for (size_t i = 0; i < inputs.size(); ++i) { + if (!chain_lane_executes(lane_disposition(i))) continue; + std::string reselect_error; + if (!maybe_reselect_residency(inputs[i].slot, reselect_error)) { + result.error = reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error; + return result; + } + } + result.decode.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + DecodeOutput out; + out.slot = inputs[i].slot; + const ChainLaneDisposition disposition = lane_disposition(i); + if (disposition == ChainLaneDisposition::Failed) { + out.failed = true; + out.error = proposal_errors[i]; + result.decode.push_back(std::move(out)); + continue; + } + if (disposition == ChainLaneDisposition::Speculation) { + Proposal & proposal = + proposals[static_cast(proposal_for_input[i])]; + out.token = proposal.pending; + out.spec_steps = 1; + out.spec_accepted_tokens = proposal.path.size() > 1 + ? static_cast(proposal.path.size() - 1) : 0; + out.target_forwards = 1; + out.committed_tokens.assign( + proposal.path.begin() + 1, proposal.path.end()); + } else { + ArLane & ar = + ar_lanes[static_cast(ar_for_input[i])]; + out.token = ar.pending; + out.target_forwards = 1; + } + attach_residency_telemetry(out); + result.decode.push_back(std::move(out)); + } + if (timing) { + const auto t_round_end = timing_clock::now(); + std::fprintf(stderr, + "[step-timing] {\"path\":\"spec-direct\",\"live\":%d," + "\"k\":%d,\"tree_bucket\":%d,\"tree_rows\":%d," + "\"replay_rows\":0,\"ar_lanes\":%d,\"ar_bucket\":0," + "\"max_kv_len\":%d,\"draft_us\":%.1f," + "\"draft_lanes\":%d,\"pre_us\":%.1f," + "\"verify_build_us\":%.1f,\"verify_exec_us\":%.1f," + "\"posterior_read_us\":%.1f,\"commit_cpu_us\":%.1f," + "\"replay_build_us\":%.1f,\"replay_exec_us\":%.1f," + "\"sample_read_us\":%.1f,\"finish_us\":%.1f," + "\"total_us\":%.1f,\"accepted_tokens\":%d," + "\"emitted_tokens\":%d,\"target_forwards\":%d}\n", + spec_count + ar_count, spec_count, tree_bucket, + total_tree, ar_count, max_kv_len, + round_draft_us_, round_draft_lanes_, + span_us(t_round_start, t_verify_build_start), + span_us(t_verify_build_start, t_verify_build_end), + span_us(t_verify_build_end, t_verify_exec_end), + span_us(t_verify_exec_end, t_posterior_end), + span_us(t_posterior_end, t_commit_end), + span_us(t_commit_end, t_replay_build_end), + span_us(t_replay_build_end, t_replay_exec_end), + span_us(t_replay_exec_end, t_sample_end), + span_us(t_sample_end, t_round_end), + span_us(t_round_start, t_round_end), + replay_total - spec_count, replay_total + ar_count, + spec_count + ar_count); + } + return result; + } + // Launch 2: accepted path segments + compact AR rows in the same builder // combination already used by mixed prefill/decode. const int ar_bucket = chain_decode_bucket_width(ar_count); @@ -2592,6 +2933,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const auto decode_round_started = timing || gate_cost_timing ? timing_clock::now() : timing_clock::time_point{}; std::optional pending_ar_gate_plan; + std::vector spec_service_ar(inputs.size(), 0); if (spec_mode_ == SpecMode::chain && !inputs.empty()) { // New chain round: restart the [step-timing] draft attribution. @@ -2609,6 +2951,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { std::vector candidates; candidates.reserve(inputs.size()); const bool use_activation_score = activation_scoring_enabled(); + const bool commit_gate_decisions = plan.prefills.empty(); for (const StepInput & in : inputs) { const Qwen35Slot & seq = slots_.slot(in.slot); SpeculationPolicy policy = in.speculation_policy; @@ -2646,7 +2989,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } gate_plan = speculation_gate_->plan( (int)inputs.size(), candidates, (int)inputs.size(), - -1); + -1, commit_gate_decisions); have_gate_plan = true; if (!gate_plan.valid) { return fail_step(gate_plan.error.empty() @@ -2670,7 +3013,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } gate_plan = speculation_gate_->plan( (int)inputs.size(), candidates, (int)inputs.size(), - -1); + -1, commit_gate_decisions); return gate_plan.valid; }; @@ -2791,18 +3134,22 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } } if (!gate_plan.pending_evaluations.empty() || - !gate_plan.decisions_committed) { + (!gate_plan.decisions_committed && + commit_gate_decisions)) { return fail_step( "adaptive activation state did not commit modes"); } } - log_spec_activations(gate_plan, *speculation_gate_); + if (gate_plan.decisions_committed) { + log_spec_activations(gate_plan, *speculation_gate_); + } // A one-shot AR decision discards the evaluation proposal. SPEC // keeps that exact first proposal so the bootstrap is useful work. for (const SpecPlanScore & score : gate_plan.ordered) { if (!score.newly_decided || - score.decision != SpecDecision::AR) { + (gate_plan.decisions_committed && + score.decision != SpecDecision::AR)) { continue; } const int slot = score.slot; @@ -2816,11 +3163,17 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } } - // Preserve every sticky Spec admission. Prefills are deferred - // below, the min-token floor is enforced inside the speculative - // path, and a later capability invariant failure becomes a - // lane-local error in step_chain_spec rather than an AR step. - for (int slot : gate_plan.admitted_slots) { + // Preserve every sticky Spec admission. The min-token floor is + // enforced inside the speculative path, capability failures remain + // lane-local errors, and planned prefills use an explicitly + // telemetered AR service round without changing the request's + // sticky routing decision. + for (const SpecPlanScore & score : gate_plan.ordered) { + if (!score.admitted || + (!gate_plan.decisions_committed && !score.forced)) { + continue; + } + const int slot = score.slot; bool found = false; for (size_t i = 0; i < inputs.size(); ++i) { if (inputs[i].slot == slot) { @@ -2870,11 +3223,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const bool any_admitted = std::any_of( admitted.begin(), admitted.end(), [](uint8_t value) { return value != 0; }); - if (any_admitted) { - StepPlan chain_plan = plan; - chain_plan.prefills.clear(); + if (any_admitted && plan.prefills.empty()) { StepResult speculative = - step_chain_spec(chain_plan, admitted, decode_round_started); + step_chain_spec(plan, admitted, decode_round_started); const double measured_us = std::chrono::duration( std::chrono::steady_clock::now() - chain_started).count(); @@ -2898,22 +3249,22 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { 1 + static_cast(output->spec_accepted_tokens); } } - if (spec_completed && !plan.prefills.empty()) { - speculative.prefills.reserve(plan.prefills.size()); - for (const PrefillSlice & slice : plan.prefills) { - speculative.prefills.push_back({ - slice.slot, PrefillOutput::Status::deferred, -1, {}}); - } - } const bool cost_sample_valid = have_gate_plan && spec_completed && !proposal_failed; if (cost_sample_valid) { const ChainLaunchShape executed = chain_launch_shape( admitted, accepted_lengths, chain_verify_depth_for_round()); + const bool direct_commit = chain_direct_commit_enabled(); + const int priced_tree_rows = direct_commit + ? executed.tree_rows + + static_cast(inputs.size()) - + executed.spec_lanes + : executed.tree_rows; speculation_gate_->observe_cost( {static_cast(inputs.size()), executed.spec_lanes, - executed.tree_rows, executed.commit_rows, + priced_tree_rows, + direct_commit ? 0 : executed.commit_rows, round_draft_lanes_}, measured_us); } @@ -2930,6 +3281,26 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } return speculative; } + if (any_admitted) { + // Chain verification cannot share a target graph with prompt work. + // Use the existing packed AR+prefill graph for this service round + // so selected prompts make immediate progress. Routing remains + // sticky SPEC; the per-request metric distinguishes this bounded + // scheduling suspension from speculative execution. + spec_service_ar = admitted; + for (size_t i = 0; i < inputs.size(); ++i) { + if (!admitted[i]) continue; + const int slot = inputs[i].slot; + if (slot >= 0 && + slot < (int)prepared_chain_drafts_.size()) { + prepared_chain_drafts_[(size_t)slot] = {}; + } + if (slot >= 0 && slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)slot]); + } + } + } if (!any_admitted && have_gate_plan) { if (gate_plan.admitted_count == 0 && plan.prefills.empty()) { pending_ar_gate_plan = gate_plan; @@ -3288,6 +3659,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (out.failed) continue; slots_.commit_step(out.slot); const int row = decode_row0 + output_rows_[oi]; + out.spec_service_ar_steps = + spec_service_ar[oi] ? 1 : 0; out.token = sample_graph_row( out.slot, row, &argmax_buf_[(size_t)row], &logits_buf_); } diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 1becd302c..db929b4be 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -740,10 +740,18 @@ 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, + int mapped_ar_seqs, + bool capture_tree_commit) { (void)kq_stride_pad; step_graph_free(sg); + if (mapped_ar_seqs < 0 || + (mapped_ar_seqs > 0 && !capture_tree_commit) || + mapped_ar_seqs + n_tree_seqs > cache.n_seq_slots || + mapped_ar_seqs + n_tree_seqs > 64) { + return false; + } if (!detail::validate_target_paged_tree_layout( cache, tree_width, n_tree_seqs, paged_max_kv_len, tree_scratch_base, tree_scratch_stride)) { @@ -754,7 +762,8 @@ bool build_target_step_paged_tree( tree_width, n_tree_seqs, graph_capacity)) { return false; } - const int n_tokens = tree_width * n_tree_seqs; + const int n_tokens = mapped_ar_seqs + tree_width * n_tree_seqs; + const int n_mapped_seqs = mapped_ar_seqs + n_tree_seqs; ggml_init_params ip{}; ip.mem_size = 512 * 1024 * 1024; @@ -767,7 +776,8 @@ 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 + + mapped_ar_seqs + n_tokens; ++i) { (void)ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 1); } @@ -780,11 +790,15 @@ bool build_target_step_paged_tree( 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); + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_mapped_seqs); sg.state_slot_ids = - ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_mapped_seqs); sg.paged_query_seq_ids = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + if (mapped_ar_seqs > 0) { + sg.paged_query_positions = + 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); @@ -799,9 +813,11 @@ bool build_target_step_paged_tree( {sg.active_slot_ids, "active_slot_ids"}, {sg.state_slot_ids, "state_slot_ids"}, {sg.paged_query_seq_ids, "paged_query_seq_ids"}, + {sg.paged_query_positions, "paged_query_positions"}, {sg.kv_write_rows, "kv_write_rows"}, }; for (const NamedInput & input : inputs) { + if (!input.tensor) continue; ggml_set_name(input.tensor, input.name); ggml_set_input(input.tensor); } @@ -812,8 +828,9 @@ 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_layers = capture_tree_commit; gi.capture_delta_intermediate = false; + gi.capture_tree_commit = capture_tree_commit; gi.parent_ids = sg.parent_ids; gi.tree_sizes = sg.tree_sizes; gi.kv_write_rows = sg.kv_write_rows; @@ -822,7 +839,9 @@ bool build_target_step_paged_tree( 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.paged_query_positions = sg.paged_query_positions; gi.n_seqs = n_tree_seqs; + gi.mapped_ar_seqs = mapped_ar_seqs; gi.paged_max_kv_len = paged_max_kv_len; gi.tree_width = tree_width; gi.tree_scratch_base = tree_scratch_base; @@ -831,6 +850,12 @@ bool build_target_step_paged_tree( QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); if (!go.logits) return false; sg.logits = go.logits; + sg.delta_captures = std::move(go.delta_captures); + sg.tree_features = go.tree_features; + if (capture_tree_commit && + (!sg.tree_features || sg.delta_captures.empty())) { + return false; + } ggml_set_output(sg.logits); sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); ggml_set_name(sg.argmax_tokens, "paged_tree_verify_argmax"); @@ -841,8 +866,32 @@ 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) && - detail::target_paged_tree_uploads_ready(sg); + if (!ggml_gallocr_alloc_graph(sg.alloc, sg.gf) || + !detail::target_paged_tree_uploads_ready(sg)) { + return false; + } + if (!capture_tree_commit) return true; + + ggml_init_params commit_params{}; + commit_params.mem_size = 16 * ggml_tensor_overhead(); + commit_params.no_alloc = true; + sg.commit_ctx = ggml_init(commit_params); + if (!sg.commit_ctx) return false; + sg.accepted_prefixes = ggml_new_tensor_1d( + sg.commit_ctx, GGML_TYPE_I32, n_tree_seqs); + sg.commit_slot_ids = ggml_new_tensor_1d( + sg.commit_ctx, GGML_TYPE_I32, n_tree_seqs); + sg.commit_rows = ggml_new_tensor_2d( + sg.commit_ctx, GGML_TYPE_I64, tree_width, n_tree_seqs); + sg.feature_commit_rows = ggml_new_tensor_1d( + sg.commit_ctx, GGML_TYPE_I32, n_tokens); + ggml_set_name(sg.accepted_prefixes, "accepted_prefixes"); + ggml_set_name(sg.commit_slot_ids, "commit_slot_ids"); + ggml_set_name(sg.commit_rows, "commit_rows"); + ggml_set_name(sg.feature_commit_rows, "feature_commit_rows"); + sg.commit_buffer = ggml_backend_alloc_ctx_tensors( + sg.commit_ctx, backend); + return sg.commit_buffer != nullptr; } diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index cbe58b786..b8da0fd7d 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -65,6 +65,8 @@ inline bool target_paged_tree_uploads_ready(const StepGraph & sg) { allocated(sg.parent_ids) && allocated(sg.tree_sizes) && allocated(sg.state_slot_ids) && allocated(sg.paged_query_seq_ids) && + (!sg.paged_query_positions || + allocated(sg.paged_query_positions)) && allocated(sg.kv_write_rows); } @@ -203,13 +205,13 @@ bool build_target_step_tree( 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. +// flattened after an optional compact one-token AR prefix as +// [mapped_ar_seqs + 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. Speculative K/V is +// written into per-slot scratch slabs. With capture_tree_commit, recurrent +// transitions and target features are exposed for post-verification promotion. bool build_target_step_paged_tree( StepGraph & sg, const TargetWeights & w, @@ -220,7 +222,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, + int mapped_ar_seqs = 0, + bool capture_tree_commit = false); // LM-head projection: project draft hidden states through the target output matrix. bool build_lm_head_projection_step( diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index ccea1c550..fb1a9a779 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -856,8 +856,7 @@ static ggml_tensor * build_full_attn_block( const bool ragged = paged_query_seq_ids != nullptr; 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)); + GGML_ASSERT(!paged_tree || (ragged && 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. @@ -958,14 +957,15 @@ static ggml_tensor * build_full_attn_block( 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. + // committed block-table prefix with only this node.s ancestor chain. + // A mixed graph uses causal positions for the compact AR prefix and + // -1 for the tree tail; a pure tree keeps positions absent. 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, + paged_query_seq_ids, paged_query_positions, /*dense_token_layout=*/n_tokens > 1); } else if (ragged) { // ── Ragged concurrent step: prefill chunk rows and decode rows all @@ -1092,6 +1092,7 @@ static ggml_tensor * build_delta_net_block( int n_prefill_segments = 0, ggml_tensor * active_slot_ids = nullptr, ggml_tensor * state_slot_ids = nullptr, + int mapped_ar_seqs = 0, bool allow_inplace_state = false ) { const int head_k_dim = w.ssm_d_state; @@ -1110,18 +1111,26 @@ static ggml_tensor * build_delta_net_block( } GGML_ASSERT((active_slot_ids == nullptr) == (state_slot_ids == nullptr)); const bool mapped_tree = active_slot_ids && parent_ids; - GGML_ASSERT(!active_slot_ids || !cap); + GGML_ASSERT(mapped_ar_seqs >= 0); + GGML_ASSERT(mapped_ar_seqs == 0 || mapped_tree); + GGML_ASSERT(!active_slot_ids || !cap || + (!cap->ssm_intermediate_states && !cap->conv_input)); GGML_ASSERT(!active_slot_ids || (mapped_tree ? (!ragged && prefill_total == 0 && - n_tokens % n_seqs == 0) - : (prefill_total + n_seqs == n_tokens))); + n_tokens >= mapped_ar_seqs && + (n_tokens - mapped_ar_seqs) % n_seqs == 0 && + active_slot_ids->ne[0] == + mapped_ar_seqs + n_seqs && + state_slot_ids->ne[0] == + mapped_ar_seqs + n_seqs) + : (mapped_ar_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); } GGML_ASSERT(!ragged || (!cap && !parent_ids)); - const bool can_skip_gdn_intermediate = skip_gdn_intermediate && !parent_ids && !cap; // Row slices of stacked projections are strided for multi-token inputs. // Materialize only the small beta/alpha slices; qkv keeps its explicit @@ -1180,6 +1189,8 @@ 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); @@ -1196,16 +1207,34 @@ 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; - segs.push_back({prefill_total, tree_tokens, n_seqs, true, - mapped_tree, conv_state, ssm_state}); + if (mapped_tree && mapped_ar_seqs > 0) { + ggml_tensor * ar_active = ggml_view_1d( + ctx, active_slot_ids, mapped_ar_seqs, 0); + ggml_tensor * ar_state = ggml_view_1d( + ctx, state_slot_ids, mapped_ar_seqs, 0); + segs.push_back({0, 1, mapped_ar_seqs, true, false, + conv_state, ssm_state, ar_active, ar_state}); + } + const int tree_tokens = mapped_tree + ? (n_tokens - mapped_ar_seqs) / n_seqs : 1; + const size_t slot_offset = + (size_t)mapped_ar_seqs * active_slot_ids->nb[0]; + ggml_tensor * segment_active = mapped_ar_seqs > 0 + ? ggml_view_1d(ctx, active_slot_ids, n_seqs, slot_offset) + : active_slot_ids; + ggml_tensor * segment_state = mapped_ar_seqs > 0 + ? ggml_view_1d(ctx, state_slot_ids, n_seqs, slot_offset) + : state_slot_ids; + segs.push_back({prefill_total + mapped_ar_seqs, tree_tokens, + n_seqs, true, mapped_tree, conv_state, ssm_state, + segment_active, segment_state}); } 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(); @@ -1226,6 +1255,11 @@ 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 = mapped_tree + ? (seg_tree ? cap : nullptr) : cap; + ggml_tensor * seg_parent_ids = seg_tree ? parent_ids : nullptr; + const bool 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 @@ -1284,7 +1318,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 { @@ -1297,13 +1331,13 @@ static ggml_tensor * build_delta_net_block( // One kernel: window = [conv_state | x], silu(conv), history // write-back, and (when capturing) the rollback window copy. ggml_tensor * ci_dst = nullptr; - if (cap && cap->conv_input) { + if (seg_cap && seg_cap->conv_input) { const int64_t ci_len = (w.ssm_d_conv - 1) + n_tokens; - ci_dst = (ci_len == cap->conv_input->ne[0]) - ? cap->conv_input - : 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); + ci_dst = (ci_len == seg_cap->conv_input->ne[0]) + ? seg_cap->conv_input + : ggml_view_3d(ctx, seg_cap->conv_input, + ci_len, seg_cap->conv_input->ne[1], seg_cap->conv_input->ne[2], + seg_cap->conv_input->nb[1], seg_cap->conv_input->nb[2], 0); } conv_out = ggml_ssm_conv_step(ctx, qkv_mixed, L.ssm_conv1d, conv_states_r, ci_dst); } else { @@ -1322,23 +1356,28 @@ 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]) { + 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_cap->conv_input->ne[2], + 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)); } + if (seg_cap && seg_tree && !seg_cap->conv_input) { + seg_cap->conv_input = conv_input; + ggml_set_output(seg_cap->conv_input); + } + // ── Save the last (kernel-1) steps back to conv_state ggml_tensor * last_conv = ggml_view_3d(ctx, conv_input, w.ssm_d_conv - 1, conv_channels, seg_seqs, @@ -1353,7 +1392,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)); @@ -1364,8 +1403,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). - conv_out = parent_ids - ? ggml_ssm_conv_tree(ctx, conv_input, L.ssm_conv1d, parent_ids) + 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); } @@ -1429,13 +1468,13 @@ static ggml_tensor * build_delta_net_block( 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. + // untouched; accepted paths are committed by later direct promotion. 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); + 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 { @@ -1448,7 +1487,7 @@ static ggml_tensor * build_delta_net_block( // ── Fused Gated DeltaNet op — returns packed (output | new_state [| intermediates]). // In tree mode, the kernel uses parent_ids to reload state at DFS // branch transitions (ported from sglang's retrieve_parent_token path). - // When `cap->ssm_intermediate_states` is present AND we are in tree + // When `seg_cap->ssm_intermediate_states` is present AND we are in tree // mode, use the _tree_persist variant: the kernel writes per-token // intermediate states DIRECTLY into the persistent cache buffer, // eliminating the downstream ggml_cpy that would otherwise copy them. @@ -1464,10 +1503,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 @@ -1491,12 +1530,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 @@ -1510,6 +1549,16 @@ static ggml_tensor * build_delta_net_block( result->src[7] = persist_inter; } } + if (seg_cap && seg_tree) { + const int64_t journal_width = + g_tensor->ne[0] == head_v_dim ? 3*head_v_dim : 2*head_v_dim + 1; + seg_cap->transition_journal = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, journal_width, num_v_heads, + n_seq_tokens, seg_seqs); + ggml_set_output(seg_cap->transition_journal); + ggml_gated_delta_net_set_transition_journal( + result, seg_cap->transition_journal); + } if (raw_gates) { ggml_gated_delta_net_set_raw_gates(result, L.ssm_dt_bias, L.ssm_a); } @@ -1551,10 +1600,10 @@ static ggml_tensor * build_delta_net_block( // forces gallocr to preserve ~50 MB per layer × 48 layers of otherwise // transient memory and inflates graph_build by ~35 ms), we create a VIEW // into the intermediate region and ggml_cpy it into the persistent cache - // buffer cap->ssm_intermediate_states. The gallocr is unaware of the + // buffer seg_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 @@ -1563,7 +1612,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); } } @@ -1730,7 +1779,7 @@ QwenGraphOutputs build_qwen35_graph( // If the caller requested capture, size the output list to the total delta- // net layer count so we can index by dn_idx as we iterate the layers. QwenGraphOutputs og_early{}; - if (in.capture_delta_intermediate) { + if (in.capture_delta_intermediate || in.capture_tree_commit) { const int n_full_attn = w.n_layer / w.full_attention_interval; const int n_delta = w.n_layer - n_full_attn; og_early.delta_captures.resize(n_delta); @@ -1747,8 +1796,10 @@ QwenGraphOutputs build_qwen35_graph( const float eps = w.rms_eps; const bool capture_with_rows = in.capture_layers && cache.target_feat && in.target_feat_rows; + const bool capture_tree_features = + in.capture_layers && in.capture_tree_commit && cache.target_feat; std::vector capture_slices; - if (capture_with_rows) { + if (capture_with_rows || capture_tree_features) { capture_slices.assign((size_t)N_CAPTURE, nullptr); } @@ -1805,15 +1856,17 @@ QwenGraphOutputs build_qwen35_graph( fa_idx++; } else { DeltaNetCapture * cap_ptr = nullptr; - if (in.capture_delta_intermediate) { + if (in.capture_delta_intermediate || in.capture_tree_commit) { cap_ptr = &og_early.delta_captures[dn_idx]; // Point at the persistent per-layer cache buffers so // build_delta_net_block can ggml_cpy into them during graph // execution. The caller (test_dflash.cpp spec loop) reads from // these tensors post-compute; their ->data pointers are always // valid because they're cache-resident, not gallocr-managed. + if (in.capture_delta_intermediate) { cap_ptr->ssm_intermediate_states = cache.ssm_intermediate[dn_idx]; cap_ptr->conv_input = cache.conv_input_cache[dn_idx]; + } } ggml_tensor * conv_st = cache.conv_state[dn_idx]; ggml_tensor * ssm_st = cache.ssm_state[dn_idx]; @@ -1844,6 +1897,7 @@ QwenGraphOutputs build_qwen35_graph( in.n_prefill_segments, in.active_slot_ids, in.state_slot_ids, + in.mapped_ar_seqs, /*allow_inplace_state=*/ in.n_prefill_tokens == 0); dn_idx++; @@ -1880,7 +1934,7 @@ QwenGraphOutputs build_qwen35_graph( if (capture_idx >= 0) { ggml_tensor * cur_2d = ggml_reshape_2d(ctx, cur, hidden, n_tokens); - if (capture_with_rows) { + if (capture_with_rows || capture_tree_features) { capture_slices[(size_t)capture_idx] = cur_2d; inpL = cur; continue; @@ -1921,7 +1975,7 @@ QwenGraphOutputs build_qwen35_graph( inpL = cur; } - if (capture_with_rows) { + if (capture_with_rows || capture_tree_features) { GGML_ASSERT(!capture_slices.empty()); ggml_tensor * feat_cat = capture_slices[0]; GGML_ASSERT(feat_cat); @@ -1931,9 +1985,15 @@ QwenGraphOutputs build_qwen35_graph( ctx, feat_cat, capture_slices[(size_t)k], 0); } feat_cat = ggml_cont(ctx, feat_cat); + if (capture_tree_features) { + og_early.tree_features = ggml_cast(ctx, feat_cat, GGML_TYPE_BF16); + ggml_set_output(og_early.tree_features); + ggml_build_forward_expand(gf, og_early.tree_features); + } else { ggml_build_forward_expand( gf, ggml_set_rows( ctx, cache.target_feat, feat_cat, in.target_feat_rows)); + } } // 2. Final norm diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index a11098fae..7c6b85614 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -44,6 +44,7 @@ struct SchedSlot { uint64_t ddtree_suspensions = 0; uint64_t spec_steps = 0; uint64_t spec_accepted_tokens = 0; + uint64_t spec_service_ar_steps = 0; uint64_t target_forwards = 0; uint64_t kvflash_page_ins = 0; uint64_t kvflash_page_outs = 0; @@ -332,6 +333,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { {"ddtree_suspensions", s.ddtree_suspensions}, {"spec_steps", s.spec_steps}, {"spec_accepted_tokens", s.spec_accepted_tokens}, + {"spec_service_ar_steps", s.spec_service_ar_steps}, {"target_forwards", s.target_forwards}, {"kvflash_page_ins", s.kvflash_page_ins}, {"kvflash_page_outs", s.kvflash_page_outs}, @@ -779,6 +781,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.ddtree_suspensions += out.ddtree_suspensions; s.spec_steps += out.spec_steps; s.spec_accepted_tokens += out.spec_accepted_tokens; + s.spec_service_ar_steps += out.spec_service_ar_steps; 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_gdn_transition_journal.cpp b/server/test/test_gdn_transition_journal.cpp index 2748af966..fd25d7fec 100644 --- a/server/test/test_gdn_transition_journal.cpp +++ b/server/test/test_gdn_transition_journal.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -547,9 +548,116 @@ bool run_case( return ok; } + +bool run_grouped_tree_case(ggml_backend_t backend) { + const Inputs inputs = make_inputs(/*kda=*/false, /*raw_gates=*/false); + constexpr int width = 2*S + 1; + + ggml_init_params params{}; + params.mem_size = 8*1024*1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + + ggml_tensor * q = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, KEY_HEADS, T, B); + ggml_tensor * k = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, KEY_HEADS, T, B); + ggml_tensor * v = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, H, T, B); + ggml_tensor * g = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, 1, H, T, B); + ggml_tensor * beta = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, 1, H, T, B); + ggml_tensor * base_state = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, S, H, B); + ggml_tensor * parents = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, T, B); + ggml_tensor * journal = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, width, H, T, B); + ggml_tensor * committed_state = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S, S, H, B); + ggml_tensor * accepted = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, B); + ggml_tensor * slots = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, B); + + ggml_tensor * result = ggml_gated_delta_net_tree( + ctx, q, k, v, g, beta, base_state, parents); + ggml_gated_delta_net_set_transition_journal(result, journal); + ggml_set_output(result); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, result); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buffer) { + ggml_free(ctx); + return false; + } + auto upload_f32 = [](ggml_tensor * tensor, + const std::vector & values) { + ggml_backend_tensor_set( + tensor, values.data(), 0, values.size()*sizeof(float)); + }; + upload_f32(q, inputs.q); + upload_f32(k, inputs.k); + upload_f32(v, inputs.v); + upload_f32(g, inputs.g); + upload_f32(beta, inputs.beta); + upload_f32(base_state, inputs.state); + + std::vector parent_ids((size_t) T*B, -1); + for (int sequence : {0, 2}) { + for (int token = 1; token < T; ++token) { + parent_ids[(size_t) sequence*T + token] = token - 1; + } + } + ggml_backend_tensor_set( + parents, parent_ids.data(), 0, + parent_ids.size()*sizeof(parent_ids[0])); + + bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + const std::vector prefixes{T, 1, T, 1}; + const std::vector identity_slots{0, 1, 2, 3}; + upload_f32(committed_state, inputs.state); + ggml_backend_tensor_set( + accepted, prefixes.data(), 0, prefixes.size()*sizeof(prefixes[0])); + ggml_backend_tensor_set( + slots, identity_slots.data(), 0, + identity_slots.size()*sizeof(identity_slots[0])); + ok = ok && ggml_backend_cuda_gdn_transition_journal_commit( + journal, committed_state, accepted, slots); + + std::vector expected = inputs.state; + const size_t slot_elements = (size_t) S*S*H; + for (int sequence = 0; sequence < B; ++sequence) { + const std::vector lane = ordinary_recurrence( + inputs, /*kda=*/false, /*raw_gates=*/false, + prefixes[(size_t) sequence]); + const size_t offset = (size_t) sequence*slot_elements; + std::copy_n(lane.begin() + offset, slot_elements, + expected.begin() + offset); + } + std::vector actual(expected.size()); + if (ok) { + ggml_backend_tensor_get( + committed_state, actual.data(), 0, + actual.size()*sizeof(float)); + ok = compare_vectors( + "grouped tree chain/root commit", actual, expected, + STATE_TOLERANCE); + } + + std::printf("gdn grouped tree journal : %s\n", + ok ? "PASS" : "FAIL"); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ok; +} } // namespace int main() { + setenv("DFLASH_GDN_FORCE_GROUPED_COLS", "1", 1); ggml_backend_t backend = ggml_backend_cuda_init(0); if (!backend) { std::fprintf(stderr, "GPU backend unavailable\n"); @@ -558,6 +666,7 @@ int main() { bool ok = run_case(backend, false, false, true); ok = run_case(backend, true, false, false) && ok; ok = run_case(backend, false, true, false) && ok; + ok = run_grouped_tree_case(backend) && ok; ggml_backend_free(backend); return ok ? 0 : 1; } diff --git a/server/test/test_seq_batch_plan.cpp b/server/test/test_seq_batch_plan.cpp index dff2122fa..3755dc741 100644 --- a/server/test/test_seq_batch_plan.cpp +++ b/server/test/test_seq_batch_plan.cpp @@ -143,6 +143,19 @@ int main() { chain_burst.decode[0].target_forwards = 2; CHECK(validate_step_result(work, chain_burst, 2).empty()); + SeqEngine::StepResult chain_service = good; + chain_service.decode[0].spec_service_ar_steps = 1; + chain_service.decode[0].target_forwards = 1; + CHECK(validate_step_result(work, chain_service, 2).empty()); + + SeqEngine::StepResult mixed_chain_paths = chain_burst; + mixed_chain_paths.decode[0].spec_service_ar_steps = 1; + CHECK(!validate_step_result(work, mixed_chain_paths, 2).empty()); + + SeqEngine::StepResult orphan_chain_service = good; + orphan_chain_service.decode[0].spec_service_ar_steps = 1; + CHECK(!validate_step_result(work, orphan_chain_service, 2).empty()); + SeqEngine::StepResult orphan_chain_acceptance = good; orphan_chain_acceptance.decode[0].spec_accepted_tokens = 1; CHECK(!validate_step_result(work, orphan_chain_acceptance, 2).empty()); diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index 11c913b56..015a5880b 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -78,6 +78,21 @@ int main() { CHECK(costly.decision(1) == SpecDecision::AR); CHECK(costly.decision(2) == SpecDecision::AR); CHECK(costly.initial_score(1) == 4.0); + + SpeculationGate provisional( + constant_costs(1.0, 10.0, 1.0), geometry(), 4); + SpecPlan provisional_plan = provisional.plan( + 1, {candidate(3, 0, 4.0)}, 1, -1, false); + CHECK(provisional_plan.valid); + CHECK(!provisional_plan.decisions_committed); + CHECK(provisional_plan.admitted_count == 1); + CHECK(provisional_plan.ordered[0].decision == SpecDecision::Undecided); + CHECK(provisional.decision(3) == SpecDecision::Undecided); + CHECK(provisional.initial_score(3) == 4.0); + provisional_plan = provisional.plan( + 1, {candidate(3, 0, NAN)}, 1); + CHECK(provisional_plan.decisions_committed); + CHECK(provisional.decision(3) == SpecDecision::Speculation); plan = costly.plan(2, { candidate(1, 0, NAN), candidate(2, 1, NAN)}, 2); CHECK(plan.pending_evaluations.empty()); @@ -375,6 +390,59 @@ int main() { CHECK(above_cliff.profiled_cost > below_cliff.profiled_cost); CHECK(above_cliff.profiled_cost - below_cliff.profiled_cost < 1.0); + // Direct promotion prices one full-live verification graph and no replay. + // With an expensive replay-row cliff, the legacy executor therefore picks + // AR while the one-launch executor correctly picks all four Spec lanes. + SpecCostSeries replay_step = series(128, 200.0); + for (size_t i = 0; i < 4; ++i) replay_step.costs[i] = 10.0; + const SpecCostTables replay_priced{ + series(128, 20.0), replay_step, series(16, 1.0)}; + SpeculationGate legacy_replay_gate( + replay_priced, geometry(), 4); + SpeculationGate direct_commit_gate( + replay_priced, geometry(), 4, {}, {}, true); + std::vector direct_candidates; + for (int i = 0; i < 4; ++i) + direct_candidates.push_back(candidate(850 + i, i, 4.0)); + const SpecPlan legacy_replay_plan = legacy_replay_gate.plan( + 4, direct_candidates, 4); + const SpecPlan direct_commit_plan = direct_commit_gate.plan( + 4, direct_candidates, 4); + CHECK(legacy_replay_plan.admitted_count == 0); + CHECK(direct_commit_plan.admitted_count == 4); + CHECK(direct_commit_plan.tree_rows == geometry().tree_rows(4)); + CHECK(direct_commit_plan.expected_step_rows == 0.0); + CHECK(direct_commit_plan.profiled_cost == 21.0); + + // A compact mixed graph carries one row per AR peer before the + // bucketed speculative tree. It must not price every AR peer at the + // full verification depth. + SpeculationGate direct_mixed_gate( + replay_priced, geometry(), 4, {}, {}, true); + const SpecPlan direct_mixed_plan = direct_mixed_gate.plan(4, { + candidate(870, 0, 4.0, SpeculationPolicy::Always), + candidate(871, 1, 4.0, SpeculationPolicy::Always), + candidate(872, 2, 0.0, SpeculationPolicy::Never), + candidate(873, 3, 0.0, SpeculationPolicy::Never)}, 4); + CHECK(direct_mixed_plan.admitted_count == 2); + CHECK(direct_mixed_plan.tree_rows == geometry().tree_rows(2) + 2); + CHECK(direct_mixed_plan.expected_step_rows == 0.0); + CHECK(direct_mixed_plan.profiled_cost == 21.0); + + // A measured direct shape with zero step rows is valid online feedback and + // remains isolated under that exact execution-shape key. + direct_commit_gate.observe_cost({4, 4, 16, 0, 4}, 84.0); + std::vector forced_direct_candidates; + for (int i = 0; i < 4; ++i) { + forced_direct_candidates.push_back(candidate(860 + i, i, 4.0, + SpeculationPolicy::Always)); + } + const SpecPlan observed_direct_plan = direct_commit_gate.plan( + 4, forced_direct_candidates, 4); + CHECK(observed_direct_plan.admitted_count == 4); + CHECK(observed_direct_plan.cost_scale == 4.0); + CHECK(observed_direct_plan.predicted_cost == 84.0); + // A realized row-4 sample belongs only to the row-4 executable shape. It // must not be written under the gate's earlier row-1 expectation. SpeculationGate shape_feedback( @@ -446,24 +514,24 @@ int main() { CHECK(ar_plan.predicted_cost == 20.0); CHECK(std::abs(ar_plan.goodput - ar_plan.ar_goodput) < 1e-12); - // Adaptive gains below the default 2% safety margin commit AR. A zero + // Adaptive gains below the default 1% safety margin commit AR. A zero // margin admits the same new request, while explicit Always is unchanged. const SpecCostTables near_break_even = constant_costs(1.0, 100.0, 1.0); SpeculationGate margin_gate(near_break_even, geometry(), 4); - plan = margin_gate.plan(1, {candidate(970, 0, 1.04)}, 1); + plan = margin_gate.plan(1, {candidate(970, 0, 1.03)}, 1); CHECK(plan.admitted_count == 0); CHECK(margin_gate.decision(970) == SpecDecision::AR); plan = margin_gate.plan(1, {candidate(970, 0, 4.0)}, 1); CHECK(plan.admitted_count == 0); CHECK(plan.ordered.empty()); - CHECK(margin_gate.initial_score(970) == 1.04); + CHECK(margin_gate.initial_score(970) == 1.03); SpecGateConfig zero_margin; zero_margin.adaptive_gain_margin = 0.0; SpeculationGate no_margin( zero_margin, near_break_even, geometry(), 4); - plan = no_margin.plan(1, {candidate(971, 0, 1.04)}, 1); + plan = no_margin.plan(1, {candidate(971, 0, 1.03)}, 1); CHECK(plan.admitted_count == 1); CHECK(no_margin.decision(971) == SpecDecision::Speculation); plan = margin_gate.plan(1, { From 7e876efae3067915805ca186332e2f0acced7a1e Mon Sep 17 00:00:00 2001 From: Graffioh Date: Thu, 20 Aug 2026 11:21:55 +0000 Subject: [PATCH 41/42] bench(concurrency): validate direct adaptive refill --- .../concurrency/analyze_dflash2_selector.py | 6 +- .../concurrency/analyze_gate_decisions.py | 26 +- .../concurrency/forced_subset_benchmark.py | 31 +- .../concurrency/refill_subset_benchmark.py | 440 ++++++++++++++++-- .../test_analyze_dflash2_selector.py | 15 + .../concurrency/test_feature_tools.py | 2 +- .../test_forced_subset_benchmark.py | 16 + .../test_refill_subset_benchmark.py | 197 ++++++++ 8 files changed, 684 insertions(+), 49 deletions(-) diff --git a/harness/benchmarks/concurrency/analyze_dflash2_selector.py b/harness/benchmarks/concurrency/analyze_dflash2_selector.py index 69ea1148d..760a5f5d0 100644 --- a/harness/benchmarks/concurrency/analyze_dflash2_selector.py +++ b/harness/benchmarks/concurrency/analyze_dflash2_selector.py @@ -243,8 +243,10 @@ def analyze_artifact(path: Path) -> dict[str, Any]: row = _record(wrapper, f"{path} timing record {index}") for key in ("live", "k", "accepted_tokens", "emitted_tokens"): _non_negative_int(row.get(key), f"{path} timing {key}") - if row.get("path") not in ("ar", "spec"): - raise ValueError(f"{path}: timing path must be ar or spec") + if row.get("path") not in ("ar", "spec", "spec-direct"): + raise ValueError( + f"{path}: timing path must be ar, spec, or spec-direct" + ) total_us = _finite_number(row.get("total_us"), f"{path} timing total_us") if total_us <= 0.0: raise ValueError(f"{path}: timing total_us must be positive") diff --git a/harness/benchmarks/concurrency/analyze_gate_decisions.py b/harness/benchmarks/concurrency/analyze_gate_decisions.py index 5752fa05b..bef65f9d9 100644 --- a/harness/benchmarks/concurrency/analyze_gate_decisions.py +++ b/harness/benchmarks/concurrency/analyze_gate_decisions.py @@ -53,7 +53,7 @@ def _json_object(match: re.Match[str], path: Path, line_no: int) -> dict[str, An def _validate_timing(row: dict[str, Any], path: Path, line_no: int) -> None: - if row.get("path") not in ("ar", "spec"): + if row.get("path") not in ("ar", "spec", "spec-direct"): raise ValueError(f"{path}:{line_no}: invalid step-timing path") for key in TIMING_COUNT_FIELDS: value = row.get(key) @@ -519,10 +519,13 @@ def _activation_proof( continue metric = metrics[engine_to_wire[request_id]] spec_steps = metric.get("spec_steps") + service_steps = metric.get("spec_service_ar_steps", 0) target_forwards = metric.get("target_forwards") if ( type(spec_steps) is not int or spec_steps < 0 + or type(service_steps) is not int + or service_steps < 0 or type(target_forwards) is not int or target_forwards < 0 ): @@ -532,18 +535,23 @@ def _activation_proof( ) continue decision = rows[0]["decision"] - if decision == "ar" and spec_steps != 0: + if decision == "ar" and ( + spec_steps != 0 or service_steps != 0 + ): errors.append( f"AR activation for engine request {request_id} " - f"executed {spec_steps} speculation steps" + f"executed {spec_steps} speculation steps and " + f"{service_steps} service AR steps" ) if decision == "speculation" and ( - spec_steps == 0 or target_forwards != 2 * spec_steps + spec_steps == 0 + or target_forwards != 2 * spec_steps + service_steps ): errors.append( f"Spec activation for engine request {request_id} " "contains a non-speculative target step " f"(spec_steps={spec_steps}, " + f"spec_service_ar_steps={service_steps}, " f"target_forwards={target_forwards})" ) if errors: @@ -715,7 +723,7 @@ def analyze_case(case_dir: Path, prompt_file: Path | None = None) -> dict[str, A path: _timing_summary([ row for row in timings if row["path"] == path ]) - for path in ("ar", "spec") + for path in ("ar", "spec", "spec-direct") if any(row["path"] == path for row in timings) } shape_groups: dict[tuple[str, int, int], list[dict[str, Any]]] = defaultdict(list) @@ -932,9 +940,9 @@ def compare_activation_shapes( shape.get("phase_mean_us") or {} ).get("draft_us"), "activation_outcome": ( - "profitable" if path == "spec" and ratio is not None - and ratio > 1.0 - else "unprofitable" if path == "spec" + "profitable" if path in ("spec", "spec-direct") + and ratio is not None and ratio > 1.0 + else "unprofitable" if path in ("spec", "spec-direct") else "ar-with-draft-tax" if ( (shape.get("phase_mean_us") or {}).get( "draft_us", 0.0 @@ -972,7 +980,7 @@ def render_markdown(report: dict[str, Any]) -> str: for case in report["cases"]: paths = case["timing"]["by_path"] ar = paths.get("ar", {}) - spec = paths.get("spec", {}) + spec = paths.get("spec-direct") or paths.get("spec", {}) total_tokens = sum( value.get("emitted_tokens", 0) for value in paths.values() ) diff --git a/harness/benchmarks/concurrency/forced_subset_benchmark.py b/harness/benchmarks/concurrency/forced_subset_benchmark.py index 17ed03c1b..78824f109 100755 --- a/harness/benchmarks/concurrency/forced_subset_benchmark.py +++ b/harness/benchmarks/concurrency/forced_subset_benchmark.py @@ -92,6 +92,7 @@ def stream_request( ) -> dict[str, Any]: started = time.perf_counter() first = None + second = None request_id = None content: list[str] = [] reasoning: list[str] = [] @@ -144,11 +145,20 @@ def stream_request( delta = choice.get("delta") or {} piece = delta.get("content") thought = delta.get("reasoning_content") + output_event = ( + isinstance(piece, str) and bool(piece) + ) or ( + isinstance(thought, str) and bool(thought) + ) + if output_event: + now = time.perf_counter() + if first is None: + first = now + elif second is None: + second = now 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 evidence for diagnosis error = f"{type(exc).__name__}: {exc}" @@ -172,6 +182,9 @@ def stream_request( "t_start": started, "t_first": first, "t_end": ended, "duration_s": ended - started, "ttft_s": first - started if first is not None else None, + "first_to_second_output_event_s": ( + second - first if first is not None and second is not None else None + ), "decode_duration_s": decode_duration, "completion_tokens": completion_tokens, "prompt_tokens": prompt_tokens, @@ -386,7 +399,7 @@ def validate_evidence( spec_requested = any(mode == "speculation" for mode in modes) spec_rounds = [ wrapped["record"] for wrapped in rounds - if wrapped["record"].get("path") == "spec" + if wrapped["record"].get("path") in ("spec", "spec-direct") and type(wrapped["record"].get("k")) is int and wrapped["record"]["k"] > 0 ] @@ -394,14 +407,22 @@ def validate_evidence( for row in spec_rounds: bucket = row.get("tree_bucket") tree_rows = row.get("tree_rows") + ar_lanes = row.get("ar_lanes", 0) + spec_rows = ( + tree_rows - ar_lanes + if row.get("path") == "spec-direct" + and type(tree_rows) is int and type(ar_lanes) is int + else tree_rows + ) if ( type(bucket) is not int or bucket <= 0 or type(tree_rows) is not int or tree_rows <= 0 - or tree_rows % bucket != 0 + or type(spec_rows) is not int or spec_rows <= 0 + or spec_rows % bucket != 0 ): errors.append("spec step-timing lacks a valid tree_rows/tree_bucket shape") continue - inferred_depths.append(tree_rows // bucket) + inferred_depths.append(spec_rows // bucket) if spec_requested: if not spec_rounds: errors.append("speculation was requested but no spec round executed") diff --git a/harness/benchmarks/concurrency/refill_subset_benchmark.py b/harness/benchmarks/concurrency/refill_subset_benchmark.py index 8b8b0d248..ba5e3ca2f 100644 --- a/harness/benchmarks/concurrency/refill_subset_benchmark.py +++ b/harness/benchmarks/concurrency/refill_subset_benchmark.py @@ -1,14 +1,13 @@ #!/usr/bin/env python3 -"""Sustained/refill forced AR/speculation diagnostic for concurrent DFlash2. +"""Sustained/refill AR/speculation/adaptive diagnostic for concurrent DFlash2. Each positional lane keeps one request in flight. The first ``clients * (waves - 1)`` completions immediately refill the lane that -completed, preserving the active A/S mask until the final C-request drain. -This measures a closed-loop saturated service workload, not adaptive -activation. The report keeps end-to-end refill goodput separate from -full-live engine-round goodput and fails closed unless telemetry proves every -request-local mode, chain depth, refill recovery, token count, and exact -output hash. +completed, preserving each positional prompt and requested mode until the +final C-request drain. This measures a closed-loop saturated service workload. +Forced A/S controls remain supported; adaptive mode additionally fails closed +unless every activation, expected initial route, execution counter, chain +depth, refill recovery, token count, and exact output hash is proven. """ @@ -23,11 +22,37 @@ from pathlib import Path from typing import Any +import analyze_gate_decisions as gate_analysis import concurrent_benchmark as base import forced_subset_benchmark as forced CLIENT_SCRIPT = Path(__file__).resolve() +MODES = ("ar", "speculation", "adaptive") +MODE_CHARS = {"ar": "A", "speculation": "S", "adaptive": "D"} + + +def parse_request_modes(raw: str, clients: int) -> list[str]: + modes = [item.strip() for item in raw.split(",")] + if len(modes) != clients or any(item not in MODES for item in modes): + raise ValueError( + "--request-modes must provide exactly one ar/speculation/adaptive " + "mode per client" + ) + adaptive = [mode == "adaptive" for mode in modes] + if any(adaptive) and not all(adaptive): + raise ValueError( + "adaptive refill must use adaptive mode for every positional lane" + ) + return modes + + +def latency_summary(values: list[float]) -> dict[str, float | int | None]: + return { + "count": len(values), + "median_s": statistics.median(values) if values else None, + "max_s": max(values) if values else None, + } def client_provenance(argv: list[str] | None = None) -> dict[str, Any]: @@ -150,6 +175,24 @@ def worker(lane_index: int) -> None: workload_end - min(first_outputs) if len(first_outputs) == len(ok) and first_outputs else None ) + ttfts = [ + float(record["ttft_s"]) for record in ok + if type(record.get("ttft_s")) in (int, float) + ] + refill_ttfts = [ + float(record["ttft_s"]) for record in ok + if record["lane_request_index"] > 0 + and type(record.get("ttft_s")) in (int, float) + ] + first_decode_gaps = [ + float(record["first_to_second_output_event_s"]) for record in ok + if type(record.get("first_to_second_output_event_s")) in (int, float) + ] + refill_first_decode_gaps = [ + float(record["first_to_second_output_event_s"]) for record in ok + if record["lane_request_index"] > 0 + and type(record.get("first_to_second_output_event_s")) in (int, float) + ] first_wave_starts = [ float(lane[0]["t_start"]) for lane in lane_records if lane ] @@ -199,7 +242,7 @@ def worker(lane_index: int) -> None: "terminal_guard_requests": args.clients, "request_modes": modes, "request_mode_mask": "".join( - "A" if mode == "ar" else "S" for mode in modes + MODE_CHARS[mode] for mode in modes ), "expected_requests": expected_requests, "requests": len(completed), @@ -227,6 +270,10 @@ def worker(lane_index: int) -> None: if complete_tokens and output_wall is not None and output_wall > 0 else None ), + "ttft": latency_summary(ttfts), + "refill_ttft": latency_summary(refill_ttfts), + "first_to_second_output_event": latency_summary(first_decode_gaps), + "refill_first_to_second_output_event": latency_summary(refill_first_decode_gaps), "prompt_tokens_total": sum(prompt_counts) if complete_prompts else None, "refill_handoffs": len(refill_gaps), "refill_gap_s_median": ( @@ -299,19 +346,180 @@ def _full_live_round_summary( "last_line_index": full_live[-1]["line_index"] if full_live else None, "path_counts": { path: sum(row["record"].get("path") == path for row in full_live) - for path in ("ar", "spec") + for path in ("ar", "spec", "spec-direct") }, }, errors) +def validate_adaptive_activations( + records: dict[str, list[dict[str, Any]]], + requests: list[dict[str, Any]], + metrics_by_id: dict[str, dict[str, Any]], + expected_mask: str, + errors: list[str], +) -> dict[str, Any]: + direct_executor = any( + wrapped["record"].get("path") == "spec-direct" + for wrapped in records["rounds"] + ) + by_engine: dict[int, list[dict[str, Any]]] = {} + valid_rows: list[dict[str, Any]] = [] + for wrapped in records["activations"]: + row = wrapped["record"] + try: + gate_analysis._validate_activation( # shared fail-closed schema + row, Path(""), wrapped["line_index"], + ) + except ValueError as exc: + errors.append(str(exc)) + continue + engine_id = int(row["request_id"]) + by_engine.setdefault(engine_id, []).append(row) + valid_rows.append(row) + + expected_engine_ids: set[int] = set() + initial_routes: dict[int, str] = {} + matched = 0 + for request in requests: + wire_id = request.get("request_id") + metric = metrics_by_id.get(wire_id) + if metric is None: + continue + engine_id = metric.get("engine_request_id") + if type(engine_id) is not int or engine_id < 0: + continue + expected_engine_ids.add(engine_id) + rows = by_engine.get(engine_id, []) + if len(rows) != 1: + errors.append( + f"{wire_id}: expected exactly one adaptive activation for " + f"engine request {engine_id}, got {len(rows)}" + ) + continue + matched += 1 + activation = rows[0] + lane = int(request["lane_index"]) + if activation["evaluation"] != "scored": + errors.append( + f"{wire_id}: adaptive activation evaluation was " + f"{activation['evaluation']!r}, expected 'scored'" + ) + steps = metric.get("spec_steps") + service_steps = metric.get("spec_service_ar_steps", 0) + target_forwards = metric.get("target_forwards") + if type(steps) is not int or steps < 0: + errors.append(f"{wire_id}: adaptive request has invalid spec_steps") + elif type(service_steps) is not int or service_steps < 0: + errors.append( + f"{wire_id}: adaptive request has invalid " + "spec_service_ar_steps" + ) + elif activation["decision"] == "ar" and ( + steps != 0 or service_steps != 0 + ): + errors.append( + f"{wire_id}: AR activation executed spec_steps={steps}, " + f"spec_service_ar_steps={service_steps}" + ) + elif activation["decision"] == "speculation" and ( + steps == 0 + or type(target_forwards) is not int + or target_forwards < + (1 if direct_executor else 2) * steps + service_steps + ): + errors.append( + f"{wire_id}: speculation activation has invalid execution " + f"counters spec_steps={steps!r}, " + f"spec_service_ar_steps={service_steps!r}, " + f"target_forwards={target_forwards!r}" + ) + if request.get("lane_request_index") == 0: + initial_routes[lane] = ( + "S" if activation["decision"] == "speculation" else "A" + ) + expected_decision = ( + "speculation" if expected_mask[lane] == "S" else "ar" + ) + if activation["decision"] != expected_decision: + errors.append( + f"{wire_id}: initial adaptive decision " + f"{activation['decision']!r} does not match lane {lane} " + f"expected {expected_decision!r}" + ) + + duplicate_ids = sorted( + engine_id for engine_id, rows in by_engine.items() if len(rows) != 1 + ) + unknown_ids = sorted(set(by_engine) - expected_engine_ids) + missing_ids = sorted(expected_engine_ids - set(by_engine)) + if duplicate_ids: + errors.append( + "duplicate adaptive activations for engine requests " + + ",".join(str(value) for value in duplicate_ids) + ) + if unknown_ids: + errors.append( + "adaptive activations reference unknown engine requests " + + ",".join(str(value) for value in unknown_ids) + ) + if missing_ids: + errors.append( + "missing adaptive activations for engine requests " + + ",".join(str(value) for value in missing_ids) + ) + + decision_counts = { + decision: sum(row["decision"] == decision for row in valid_rows) + for decision in ("ar", "speculation") + } + evaluation_counts = { + evaluation: sum(row["evaluation"] == evaluation for row in valid_rows) + for evaluation in ("scored", "failed") + } + score_kinds = sorted({str(row["score_kind"]) for row in valid_rows}) + return { + "required": True, + "records": len(records["activations"]), + "valid_records": len(valid_rows), + "matched_requests": matched, + "decision_counts": decision_counts, + "evaluation_counts": evaluation_counts, + "score_kinds": score_kinds, + "expected_initial_route_mask": expected_mask, + "observed_initial_route_mask": "".join( + initial_routes.get(lane, "?") for lane in range(len(expected_mask)) + ), + } + + def validate_evidence( workload: dict[str, Any], records: dict[str, list[dict[str, Any]]], clients: int, modes: list[str], spec_depth: int, waves: int, max_start_skew_ms: float, max_refill_gap_ms: float, min_full_live_rounds: int, + expected_adaptive_mask: str | None = None, + adaptive_scoring_enabled: bool = True, ) -> dict[str, Any]: errors: list[str] = [] + adaptive_requested = all(mode == "adaptive" for mode in modes) + adaptive_activation_required = adaptive_requested and adaptive_scoring_enabled expected_requests = clients * waves + adaptive_mask_valid = ( + isinstance(expected_adaptive_mask, str) + and len(expected_adaptive_mask) == clients + and set(expected_adaptive_mask) <= {"A", "S"} + ) + if adaptive_requested and not adaptive_mask_valid: + errors.append( + "adaptive refill requires an expected A/S route for every lane" + ) + elif not adaptive_requested and expected_adaptive_mask is not None: + errors.append("expected adaptive mask supplied for a forced refill") + if ( + adaptive_requested and not adaptive_scoring_enabled + and expected_adaptive_mask != "A" * clients + ): + errors.append("adaptive scoring-off control must expect all-AR routing") if workload["requests"] != expected_requests or workload["missing_requests"] != 0: errors.append( f"completed request records {workload['requests']} do not match " @@ -368,6 +576,10 @@ def validate_evidence( errors.append("exact request output hashes are incomplete") rounds = records["rounds"] + direct_executor = any( + wrapped["record"].get("path") == "spec-direct" + for wrapped in rounds + ) full_live, timing_errors = _full_live_round_summary(rounds, clients) errors.extend(timing_errors) if full_live["longest_streak"] < min_full_live_rounds: @@ -376,10 +588,15 @@ def validate_evidence( f"{full_live['longest_streak']}, need {min_full_live_rounds}" ) - spec_requested = any(mode == "speculation" for mode in modes) + spec_requested = ( + any(mode == "speculation" for mode in modes) + or adaptive_activation_required + and expected_adaptive_mask is not None + and "S" in expected_adaptive_mask + ) spec_rounds = [ wrapped for wrapped in rounds - if wrapped["record"].get("path") == "spec" + if wrapped["record"].get("path") in ("spec", "spec-direct") and type(wrapped["record"].get("k")) is int and wrapped["record"]["k"] > 0 ] @@ -388,14 +605,22 @@ def validate_evidence( row = wrapped["record"] bucket = row.get("tree_bucket") tree_rows = row.get("tree_rows") + ar_lanes = row.get("ar_lanes", 0) + spec_rows = ( + tree_rows - ar_lanes + if row.get("path") == "spec-direct" + and type(tree_rows) is int and type(ar_lanes) is int + else tree_rows + ) if ( type(bucket) is not int or bucket <= 0 or type(tree_rows) is not int or tree_rows <= 0 - or tree_rows % bucket != 0 + or type(spec_rows) is not int or spec_rows <= 0 + or spec_rows % bucket != 0 ): errors.append("spec step-timing lacks a valid tree_rows/tree_bucket shape") continue - inferred_depths.append(tree_rows // bucket) + inferred_depths.append(spec_rows // bucket) if spec_requested: if not spec_rounds: errors.append("speculation was requested but no spec round executed") @@ -438,12 +663,37 @@ def validate_evidence( continue expected_mode = modes[int(request["lane_index"])] steps = metric.get("spec_steps") + service_steps = metric.get("spec_service_ar_steps", 0) + target_forwards = metric.get("target_forwards") if type(steps) is not int or steps < 0: errors.append(f"{request_id}: invalid spec_steps") + elif type(service_steps) is not int or service_steps < 0: + errors.append(f"{request_id}: invalid spec_service_ar_steps") elif expected_mode == "speculation" and steps == 0: errors.append(f"{request_id}: forced speculation never executed") - elif expected_mode == "ar" and steps != 0: - errors.append(f"{request_id}: forced AR executed speculation") + elif expected_mode == "speculation" and ( + type(target_forwards) is not int + or target_forwards < + (1 if direct_executor else 2) * steps + service_steps + ): + errors.append( + f"{request_id}: forced speculation has invalid execution " + f"counters spec_steps={steps!r}, " + f"spec_service_ar_steps={service_steps!r}, " + f"target_forwards={target_forwards!r}" + ) + elif expected_mode == "ar" and (steps != 0 or service_steps != 0): + errors.append( + f"{request_id}: forced AR executed speculation/service" + ) + elif ( + expected_mode == "adaptive" + and not adaptive_scoring_enabled + and (steps != 0 or service_steps != 0) + ): + errors.append( + f"{request_id}: scoring-off control executed speculation/service" + ) engine_id = metric.get("engine_request_id") if type(engine_id) is not int: errors.append(f"{request_id}: missing integer engine_request_id") @@ -452,6 +702,32 @@ def validate_evidence( elif expected_mode == "ar" and engine_id in selector_engine_ids: errors.append(f"{request_id}: forced AR emitted DFlash2 selector evidence") + activation_summary: dict[str, Any] = { + "required": adaptive_activation_required, + "records": len(records["activations"]), + } + if adaptive_activation_required and adaptive_mask_valid: + assert expected_adaptive_mask is not None + activation_summary = validate_adaptive_activations( + records, requests, metrics_by_id, expected_adaptive_mask, errors, + ) + elif adaptive_activation_required: + activation_summary["validation"] = "invalid-expected-mask" + elif adaptive_requested: + activation_summary.update({ + "scoring_disabled": True, + "expected_initial_route_mask": expected_adaptive_mask, + "observed_initial_route_mask": "A" * clients, + }) + if records["activations"]: + errors.append( + "adaptive scoring-off control emitted activation records" + ) + elif not adaptive_requested and records["activations"]: + errors.append( + "forced refill unexpectedly emitted adaptive activation records" + ) + last_full_live_line = full_live["last_line_index"] completed_before_last_full_live = ( sum( @@ -474,7 +750,15 @@ def validate_evidence( return { "passed": not errors, "errors": errors, - "adaptive_claims_permitted": False, + "adaptive_claims_permitted": ( + adaptive_activation_required and not errors + ), + "adaptive_stack_control_permitted": ( + adaptive_requested + and not adaptive_scoring_enabled + and not errors + ), + "activation": activation_summary, "closed_cohort_claims_permitted": False, "full_live": full_live, "min_full_live_rounds": min_full_live_rounds, @@ -495,18 +779,41 @@ def markdown(report: dict[str, Any]) -> str: validation = report["validation"] full_live = validation["full_live"] status = "PASS" if validation["passed"] else "FAIL" + adaptive_requests = bool(report["scope"].get("adaptive_requests")) + adaptive_evaluation = bool(report["scope"]["adaptive_evaluation"]) + if adaptive_evaluation: + title = "Adaptive" + scope = "This is a fail-closed adaptive closed-loop refill diagnostic." + elif adaptive_requests: + title = "Adaptive OFF" + scope = ( + "This is a matched-stack adaptive-scoring-off refill control." + ) + else: + title = "Forced" + scope = "This is a forced closed-loop refill diagnostic." + activation = validation["activation"] + route_mask = ( + activation.get("observed_initial_route_mask") + if adaptive_requests else workload["request_mode_mask"] + ) + refill_ttft = workload["refill_ttft"] + refill_decode_gap = workload["refill_first_to_second_output_event"] return ( - f"# Forced DFlash2 refill diagnostic — {report['label']}\n\n" - "This is a forced closed-loop refill diagnostic. It makes neither an " - "adaptive claim nor a closed-cohort makespan claim.\n\n" - "| C | Mask | Waves | Depth | Requests | Refill tok/s | " - "Full-live round tok/s | Recovery | Status |\n" - "| ---: | :--- | ---: | ---: | ---: | ---: | ---: | :--- | :--- |\n" - f"| {workload['clients']} | {workload['request_mode_mask']} | " + f"# {title} DFlash2 refill diagnostic — {report['label']}\n\n" + f"{scope} It makes no closed-cohort makespan claim.\n\n" + "| C | Route | Waves | Depth | Requests | Refill tok/s | " + "Full-live tok/s | Refill TTFT med/max s | " + "First-output gap med/max s | Recovery | Status |\n" + "| ---: | :--- | ---: | ---: | ---: | ---: | ---: | :--- | :--- | :--- | :--- |\n" + f"| {workload['clients']} | {route_mask} | " f"{workload['waves']} | {report['spec_depth']} | " f"{workload['requests_ok']}/{workload['expected_requests']} | " f"{base.fmt(workload['aggregate_refill_tok_s'])} | " f"{base.fmt(full_live['engine_round_goodput_tok_s'])} | " + f"{base.fmt(refill_ttft['median_s'])}/{base.fmt(refill_ttft['max_s'])} | " + f"{base.fmt(refill_decode_gap['median_s'])}/" + f"{base.fmt(refill_decode_gap['max_s'])} | " f"{validation['completed_before_last_full_live']}/" f"{validation['required_refill_recoveries']} | {status} |\n" ) @@ -519,6 +826,19 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--model", default="luce-dflash") parser.add_argument("--clients", type=int, required=True) parser.add_argument("--request-modes", required=True) + parser.add_argument( + "--expected-adaptive-mask", + help=( + "Required initial per-lane A/S route for adaptive requests" + ), + ) + parser.add_argument( + "--adaptive-scoring-disabled", + action="store_true", + help=( + "Validate the DFLASH_SPEC_ACTIVATION_SCORE=0 matched-stack control" + ), + ) parser.add_argument("--waves", type=int, required=True) parser.add_argument("--spec-depth", type=int, required=True) parser.add_argument("--prompt-file", type=Path, required=True) @@ -553,7 +873,31 @@ def run(args: argparse.Namespace) -> int: raise ValueError("invalid synchronization/refill proof threshold") if args.log_settle_ms < 0: raise ValueError("--log-settle-ms must be non-negative") - modes = forced.parse_request_modes(args.request_modes, args.clients) + modes = parse_request_modes(args.request_modes, args.clients) + adaptive_requested = all(mode == "adaptive" for mode in modes) + adaptive_scoring_enabled = not bool( + getattr(args, "adaptive_scoring_disabled", False) + ) + if not adaptive_requested and not adaptive_scoring_enabled: + raise ValueError( + "--adaptive-scoring-disabled requires adaptive request modes" + ) + raw_expected_mask = getattr(args, "expected_adaptive_mask", None) + expected_adaptive_mask = ( + raw_expected_mask.strip().upper() + if isinstance(raw_expected_mask, str) and raw_expected_mask.strip() + else None + ) + if adaptive_requested and ( + expected_adaptive_mask is None + or len(expected_adaptive_mask) != args.clients + or set(expected_adaptive_mask) - {"A", "S"} + ): + raise ValueError( + "--expected-adaptive-mask must provide one A/S route per client" + ) + if not adaptive_requested and expected_adaptive_mask is not None: + raise ValueError("--expected-adaptive-mask requires adaptive request modes") prompts = base.load_prompts(args.prompt_file) metadata_bytes = args.server_metadata_json.read_bytes() metadata = json.loads(metadata_bytes) @@ -561,8 +905,25 @@ def run(args: argparse.Namespace) -> int: raise ValueError("server metadata must be a JSON object") forced.validate_server_metadata( metadata, args.clients, args.spec_depth, args.prompt_offset, - require_selector=any(mode == "speculation" for mode in modes), + require_selector=any( + mode in ("speculation", "adaptive") for mode in modes + ), ) + declared_decode_mode = (metadata.get("feature_config") or {}).get("decode_mode") + if adaptive_requested and declared_decode_mode != "adaptive": + raise ValueError( + "adaptive refill requires server metadata decode_mode=adaptive" + ) + if adaptive_requested: + expected_score_switch = "1" if adaptive_scoring_enabled else "0" + actual_score_switch = (metadata.get("launch_environment") or {}).get( + "DFLASH_SPEC_ACTIVATION_SCORE" + ) + if actual_score_switch != expected_score_switch: + raise ValueError( + "adaptive refill metadata must record " + f"DFLASH_SPEC_ACTIVATION_SCORE={expected_score_switch}" + ) log_start = args.server_log.stat().st_size workload = run_refill(args, prompts, modes) if args.log_settle_ms: @@ -572,23 +933,37 @@ def run(args: argparse.Namespace) -> int: validation = validate_evidence( workload, records, args.clients, modes, args.spec_depth, args.waves, args.max_start_skew_ms, args.max_refill_gap_ms, - args.min_full_live_rounds, + args.min_full_live_rounds, expected_adaptive_mask, + adaptive_scoring_enabled, ) report = { - "schema_version": 1, - "kind": "dflash2-forced-refill-diagnostic", + "schema_version": 2, + "kind": ( + "dflash2-adaptive-refill-diagnostic" + if adaptive_requested and adaptive_scoring_enabled else + "dflash2-adaptive-scoring-off-refill-control" + if adaptive_requested else + "dflash2-forced-refill-diagnostic" + ), "label": args.label, "scope": { - "forced_controls_only": True, - "adaptive_evaluation": False, + "forced_controls_only": not adaptive_requested, + "adaptive_requests": adaptive_requested, + "adaptive_evaluation": ( + adaptive_requested and adaptive_scoring_enabled + ), + "adaptive_stack_control": ( + adaptive_requested and not adaptive_scoring_enabled + ), "closed_cohort_makespan": False, "interpretation": ( "aggregate_refill_tok_s is end-to-end closed-loop goodput " "across persistent client lanes and includes startup, request " "handoffs, prefill, and terminal drain. full_live decode-round " "goodput uses only timed live=C decode rounds and excludes " - "handoff/prefill gaps. Neither is an adaptive result or a " - "single closed-cohort makespan result." + "handoff/prefill gaps. Adaptive claims require fail-closed " + "activation validation. Neither metric is a single " + "closed-cohort makespan result." ), }, "base_url": args.base_url, @@ -598,6 +973,7 @@ def run(args: argparse.Namespace) -> int: "seed": args.seed, "ignore_eos": True, "spec_depth": args.spec_depth, + "expected_adaptive_mask": expected_adaptive_mask, "prompt_offset": args.prompt_offset, "prompt_file": str(args.prompt_file.resolve()), "prompt_file_sha256": forced.digest_bytes(args.prompt_file.read_bytes()), diff --git a/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py b/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py index efd0f9cff..d84475889 100644 --- a/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py +++ b/harness/benchmarks/concurrency/test_analyze_dflash2_selector.py @@ -171,6 +171,21 @@ def test_artifact_joins_engine_selector_to_wire_request(self) -> None: self.assertEqual(case["round_timing"]["tail"]["goodput_tok_s"], 1000.0) self.assertEqual(case["round_timing"]["all"]["goodput_tok_s"], 2000.0) + def test_artifact_accepts_direct_spec_timing(self) -> None: + value = artifact() + value["server_records"]["rounds"][0]["record"]["path"] = "spec-direct" + value["server_records"]["rounds"][0]["raw_json"] = json.dumps( + value["server_records"]["rounds"][0]["record"], + separators=(",", ":"), + ) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "bench.json" + path.write_text(json.dumps(value), encoding="utf-8") + case = analyzer.analyze_artifact(path) + self.assertEqual( + case["round_timing"]["full_live"]["goodput_tok_s"], 3000.0, + ) + def test_artifact_rejects_unknown_selector_engine_id(self) -> None: value = artifact() bad = selector(99, 0) diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index b61a05f33..c93207dda 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -649,7 +649,7 @@ def shape(path: str, live: int, k: int, rate: float, draft: float) -> dict: "variant": "adaptive-on", "aggregate_tok_s": 42.0, "timing": {"by_shape": [ shape("ar", 2, 0, 40.0, 10.0), - shape("spec", 3, 2, 49.0, 15.0), + shape("spec-direct", 3, 2, 49.0, 15.0), ]}, }, ] diff --git a/harness/benchmarks/concurrency/test_forced_subset_benchmark.py b/harness/benchmarks/concurrency/test_forced_subset_benchmark.py index c5deb4351..6f18e6f03 100644 --- a/harness/benchmarks/concurrency/test_forced_subset_benchmark.py +++ b/harness/benchmarks/concurrency/test_forced_subset_benchmark.py @@ -80,6 +80,8 @@ def __iter__(self): return iter([ b'data: {"id":"wire-1","choices":[{"delta":{"content":"x"}}]}\n', b"\n", + b'data: {"id":"wire-1","choices":[{"delta":{"content":"y"}}]}\n', + b"\n", b'data: {"id":"wire-1","choices":[{"delta":{},' b'"finish_reason":"length"}]}\n', b"\n", @@ -108,6 +110,8 @@ def open_request(http_request, timeout): self.assertEqual(row["request_id"], "wire-1") self.assertEqual(row["decode_mode"], "speculation") self.assertEqual(len(row["combined_output_sha256"]), 64) + self.assertIsNotNone(row["first_to_second_output_event_s"]) + self.assertGreaterEqual(row["first_to_second_output_event_s"], 0.0) self.assertIsNone(row["error"]) def test_modes_are_positional_and_reject_adaptive(self) -> None: @@ -175,6 +179,18 @@ def test_valid_mixed_subset_proves_depth_and_sustained_full_live(self) -> None: self.assertEqual(result["longest_full_live_streak"], 2) self.assertIs(result["adaptive_claims_permitted"], False) + def test_direct_mixed_subset_infers_depth_without_ar_prefix_rows(self) -> None: + records = valid_records() + records["rounds"][1]["record"].update({ + "path": "spec-direct", "ar_lanes": 1, "tree_rows": 9, + }) + result = benchmark.validate_evidence( + valid_level(), records, 2, ["ar", "speculation"], 4, + max_start_skew_ms=100.0, min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertEqual(result["executed_spec_depths"], [4]) + def test_missing_requested_live_concurrency_fails_closed(self) -> None: records = valid_records() for row in records["rounds"]: diff --git a/harness/benchmarks/concurrency/test_refill_subset_benchmark.py b/harness/benchmarks/concurrency/test_refill_subset_benchmark.py index 5210b4428..b4e30884e 100644 --- a/harness/benchmarks/concurrency/test_refill_subset_benchmark.py +++ b/harness/benchmarks/concurrency/test_refill_subset_benchmark.py @@ -93,6 +93,7 @@ def valid_records() -> dict: "request_id": f"wire-{wave}-{lane}", "engine_request_id": engine_id, "spec_steps": 0 if lane == 0 else 3, + "target_forwards": 3 if lane == 0 else 6, }, line + lane)) if lane == 1: selectors.append(wrapped({ @@ -108,6 +109,49 @@ def valid_records() -> dict: } +def valid_adaptive_case() -> tuple[dict, dict]: + workload = valid_workload() + for request in workload["requests_detail"]: + request["decode_mode"] = "adaptive" + request["lane_request_index"] = request["wave_index"] + for lane in workload["lanes"]: + lane["decode_mode"] = "adaptive" + + records = valid_records() + records["selectors"] = [] + records["activations"] = [] + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + wire_parts = metric["request_id"].split("-") + lane = int(wire_parts[-1]) + engine_id = metric["engine_request_id"] + decision = "speculation" if lane == 0 else "ar" + metric["spec_steps"] = 3 if lane == 0 else 0 + metric["target_forwards"] = 6 if lane == 0 else 3 + if lane == 0: + records["selectors"].append(wrapped({ + "request_id": engine_id, + "accepted_depth": 3, + "depths": [], + }, wrapped_metric["line_index"])) + records["activations"].append(wrapped({ + "request_id": engine_id, + "slot": lane, + "activation_score": 6.0 if lane == 0 else 2.0, + "score_kind": "qwen38-dflash2-selector-benefit-v1", + "expected_yield": 6.0 if lane == 0 else 2.0, + "evaluation": "scored", + "fallback_reason": None, + "decision_reason": ( + "selected_by_joint_goodput" + if lane == 0 else "ar_counterfactual_won" + ), + "decision": decision, + "hazards": [0.1, 0.2], + }, wrapped_metric["line_index"] - 1)) + return workload, records + + class RefillSubsetTests(unittest.TestCase): def test_refill_client_keeps_positional_modes_and_sequences_each_lane(self) -> None: calls: dict[str, int] = {"prompt-a": 0, "prompt-b": 0} @@ -176,6 +220,159 @@ def test_valid_refill_proves_every_handoff_and_full_live_round_rate(self) -> Non self.assertIs(result["adaptive_claims_permitted"], False) self.assertIs(result["closed_cohort_claims_permitted"], False) + def test_direct_refill_validates_compact_depth_and_one_pass_counters(self) -> None: + workload, records = valid_adaptive_case() + for wrapped_round in records["rounds"]: + wrapped_round["record"].update({ + "path": "spec-direct", "ar_lanes": 1, "tree_rows": 5, + }) + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + if metric["spec_steps"]: + # One forward per direct Spec step plus an ordinary AR step + # before the sticky activation was committed. + metric["target_forwards"] = metric["spec_steps"] + 1 + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, expected_adaptive_mask="SA", + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertEqual(result["executed_spec_depths"], [4]) + self.assertEqual(result["full_live"]["path_counts"]["spec-direct"], 3) + + def test_direct_forced_refill_validates_one_pass_counters(self) -> None: + workload = valid_workload() + records = valid_records() + for wrapped_round in records["rounds"]: + wrapped_round["record"].update({ + "path": "spec-direct", "ar_lanes": 1, "tree_rows": 5, + }) + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + if metric["spec_steps"]: + metric["target_forwards"] = metric["spec_steps"] + result = benchmark.validate_evidence( + workload, records, 2, ["ar", "speculation"], 4, 3, + max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, + ) + self.assertTrue(result["passed"], result["errors"]) + + def test_adaptive_refill_proves_route_and_execution_for_every_request( + self, + ) -> None: + workload, records = valid_adaptive_case() + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, expected_adaptive_mask="SA", + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertIs(result["adaptive_claims_permitted"], True) + self.assertEqual(result["activation"]["matched_requests"], 6) + self.assertEqual( + result["activation"]["observed_initial_route_mask"], "SA", + ) + self.assertEqual( + result["activation"]["decision_counts"], + {"ar": 3, "speculation": 3}, + ) + + def test_adaptive_refill_allows_later_joint_gate_route_changes(self) -> None: + workload, records = valid_adaptive_case() + later_engine_id = 3 + for wrapped_activation in records["activations"]: + activation = wrapped_activation["record"] + if activation["request_id"] == later_engine_id: + activation.update({ + "decision": "ar", + "decision_reason": "ar_counterfactual_won", + }) + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + if metric["engine_request_id"] == later_engine_id: + metric.update({"spec_steps": 0, "target_forwards": 3}) + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + 100.0, 100.0, 2, expected_adaptive_mask="SA", + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertEqual( + result["activation"]["observed_initial_route_mask"], "SA", + ) + self.assertEqual( + result["activation"]["decision_counts"], + {"ar": 4, "speculation": 2}, + ) + + def test_adaptive_spec_service_rounds_are_explicit_and_valid(self) -> None: + workload, records = valid_adaptive_case() + for wrapped_metric in records["requests"]: + metric = wrapped_metric["record"] + lane = int(metric["request_id"].split("-")[-1]) + if lane == 0: + metric["spec_service_ar_steps"] = 1 + metric["target_forwards"] = 2 * metric["spec_steps"] + 1 + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + max_start_skew_ms=100.0, max_refill_gap_ms=100.0, + min_full_live_rounds=2, expected_adaptive_mask="SA", + ) + self.assertTrue(result["passed"], result["errors"]) + + records["requests"][1]["record"]["spec_service_ar_steps"] = 1 + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + 100.0, 100.0, 2, expected_adaptive_mask="SA", + ) + self.assertFalse(result["passed"]) + + def test_adaptive_refill_route_mismatch_fails_closed(self) -> None: + workload, records = valid_adaptive_case() + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + 100.0, 100.0, 2, expected_adaptive_mask="AS", + ) + self.assertFalse(result["passed"]) + self.assertIs(result["adaptive_claims_permitted"], False) + self.assertTrue(any( + "does not match lane" in error for error in result["errors"] + )) + + def test_adaptive_scoring_off_is_proven_as_all_ar_control(self) -> None: + workload, records = valid_adaptive_case() + records["activations"] = [] + records["selectors"] = [] + for wrapped_round in records["rounds"]: + wrapped_round["record"].update({ + "path": "ar", "k": 0, "emitted_tokens": 2, + }) + for wrapped_metric in records["requests"]: + wrapped_metric["record"].update({ + "spec_steps": 0, "target_forwards": 3, + }) + result = benchmark.validate_evidence( + workload, records, 2, ["adaptive", "adaptive"], 4, 3, + 100.0, 100.0, 2, expected_adaptive_mask="AA", + adaptive_scoring_enabled=False, + ) + self.assertTrue(result["passed"], result["errors"]) + self.assertIs(result["adaptive_claims_permitted"], False) + self.assertIs(result["adaptive_stack_control_permitted"], True) + self.assertEqual(result["activation"]["records"], 0) + self.assertEqual( + result["activation"]["observed_initial_route_mask"], "AA", + ) + + def test_refill_mode_parser_requires_uniform_adaptive_lanes(self) -> None: + self.assertEqual( + benchmark.parse_request_modes("adaptive,adaptive", 2), + ["adaptive", "adaptive"], + ) + with self.assertRaisesRegex(ValueError, "every positional lane"): + benchmark.parse_request_modes("adaptive,ar", 2) + def test_missing_post_refill_full_live_proof_fails_closed(self) -> None: records = valid_records() del records["rounds"][1:] From 8eeb4e4669720b7b263cb6f64dde7de5cdc7909f Mon Sep 17 00:00:00 2001 From: Graffioh Date: Thu, 20 Aug 2026 16:23:38 +0000 Subject: [PATCH 42/42] fix: replan adaptive speculation by cohort --- .../common/speculation/spec_cost_profile.cpp | 154 ++++++++++ .../common/speculation/spec_cost_profile.h | 11 + .../src/common/speculation/speculation_gate.h | 244 ++++++--------- .../qwen35/concurrency/qwen35_seq_engine.cpp | 284 ++++++++++++------ .../qwen35/concurrency/qwen35_seq_engine.h | 16 +- server/src/qwen35/qwen35_backend.cpp | 8 +- server/test/test_dflash2_benefit.cpp | 16 +- server/test/test_spec_cost_profile.cpp | 30 ++ server/test/test_speculation_gate.cpp | 146 +++------ 9 files changed, 557 insertions(+), 352 deletions(-) diff --git a/server/src/common/speculation/spec_cost_profile.cpp b/server/src/common/speculation/spec_cost_profile.cpp index 8bcc33a82..9e4cd9dca 100644 --- a/server/src/common/speculation/spec_cost_profile.cpp +++ b/server/src/common/speculation/spec_cost_profile.cpp @@ -1,9 +1,18 @@ #include "common/speculation/spec_cost_profile.h" +#include "common/sha1.h" #include #include +#include +#include +#include +#include +#include +#include #include +#include + namespace dflash::common { namespace { @@ -71,8 +80,153 @@ SeriesResult profile_monotonic_costs( return result; } + +constexpr int kProfileCacheVersion = 1; +constexpr size_t kMaxSeriesEntries = 4096; + +bool read_series( + std::istream & input, const char * expected, + SpecCostSeries & series) { + std::string name; + size_t count = 0; + if (!(input >> name >> count) || name != expected || + count == 0 || count > kMaxSeriesEntries) { + return false; + } + series.indices.resize(count); + series.costs.resize(count); + for (size_t i = 0; i < count; ++i) { + if (!(input >> series.indices[i] >> series.costs[i])) return false; + } + return true; +} + +void write_series( + std::ostream & output, const char * name, + const SpecCostSeries & series) { + output << name << ' ' << series.indices.size() << '\n'; + output << std::setprecision(17); + for (size_t i = 0; i < series.indices.size(); ++i) { + output << series.indices[i] << ' ' << series.costs[i] << '\n'; + } +} + +std::string hex_sha1(const std::string & value) { + uint8_t digest[20]; + sha1_hash(value.data(), value.size(), digest); + std::ostringstream out; + out << std::hex << std::setfill('0'); + for (uint8_t byte : digest) out << std::setw(2) << (unsigned)byte; + return out.str(); +} + } // namespace +std::string spec_cost_profile_cache_path(const std::string & identity) { + if (identity.empty()) return {}; + if (const char * configured = std::getenv("DFLASH_SPEC_PROFILE_PATH")) { + if (std::string(configured) == "0") return {}; + if (*configured) return configured; + } + std::filesystem::path root; + if (const char * xdg = std::getenv("XDG_CACHE_HOME"); xdg && *xdg) { + root = xdg; + } else if (const char * home = std::getenv("HOME"); home && *home) { + root = std::filesystem::path(home) / ".cache"; + } else { + return {}; + } + return (root / "lucebox" / + ("spec-cost-v1-" + hex_sha1(identity) + ".profile")).string(); +} + +bool load_spec_cost_profile( + const std::string & path, const std::string & identity, + SpecCostTables & tables, std::string & error) { + tables = {}; + error.clear(); + if (path.empty()) { + error = "profile cache disabled"; + return false; + } + std::ifstream input(path); + if (!input) { + error = "profile cache miss"; + return false; + } + std::string magic; + int version = 0; + std::string stored_identity; + SpecCostTables loaded; + if (!(input >> magic >> version) || magic != "dflash-spec-cost-profile" || + version != kProfileCacheVersion || + !(input >> std::quoted(stored_identity)) || + !(input >> std::quoted(loaded.speculator_id)) || + !read_series(input, "tree", loaded.tree_cost) || + !read_series(input, "step", loaded.step_cost) || + !read_series(input, "draft", loaded.draft_cost)) { + error = "invalid profile cache"; + return false; + } + input >> std::ws; + if (!input.eof() || stored_identity != identity || !loaded.valid()) { + error = stored_identity != identity + ? "profile cache identity mismatch" : "invalid profile cache"; + return false; + } + tables = std::move(loaded); + return true; +} + +bool save_spec_cost_profile( + const std::string & path, const std::string & identity, + const SpecCostTables & tables, std::string & error) { + error.clear(); + if (path.empty()) return true; + if (identity.empty() || !tables.valid()) { + error = "refusing to save invalid profile cache"; + return false; + } + const std::filesystem::path destination(path); + std::error_code ec; + if (!destination.parent_path().empty()) { + std::filesystem::create_directories(destination.parent_path(), ec); + if (ec) { + error = "could not create profile cache directory"; + return false; + } + } + const std::filesystem::path temporary = + destination.string() + ".tmp." + std::to_string(getpid()); + { + std::ofstream output(temporary, std::ios::trunc); + if (!output) { + error = "could not open temporary profile cache"; + return false; + } + output << "dflash-spec-cost-profile " << kProfileCacheVersion << '\n' + << std::quoted(identity) << '\n' + << std::quoted(tables.speculator_id) << '\n'; + write_series(output, "tree", tables.tree_cost); + write_series(output, "step", tables.step_cost); + write_series(output, "draft", tables.draft_cost); + output.flush(); + if (!output) { + error = "could not write profile cache"; + output.close(); + std::filesystem::remove(temporary, ec); + return false; + } + } + std::filesystem::rename(temporary, destination, ec); + if (ec) { + error = "could not publish profile cache"; + std::filesystem::remove(temporary, ec); + return false; + } + return true; +} + SpecProfileGrid build_spec_profile_grid( int max_concurrency, int tree_width, int max_accept, const std::function & bucket) { diff --git a/server/src/common/speculation/spec_cost_profile.h b/server/src/common/speculation/spec_cost_profile.h index 1e9333200..924c527d4 100644 --- a/server/src/common/speculation/spec_cost_profile.h +++ b/server/src/common/speculation/spec_cost_profile.h @@ -26,6 +26,17 @@ struct SpecCostProfileResult { bool ok() const { return error.empty() && tables.valid(); } }; +// Returns an empty path when disk caching is disabled or no cache root exists. +std::string spec_cost_profile_cache_path(const std::string & identity); + +bool load_spec_cost_profile( + const std::string & path, const std::string & identity, + SpecCostTables & tables, std::string & error); + +bool save_spec_cost_profile( + const std::string & path, const std::string & identity, + const SpecCostTables & tables, std::string & error); + class SpecCostProfiler { public: using Runner = std::function; diff --git a/server/src/common/speculation/speculation_gate.h b/server/src/common/speculation/speculation_gate.h index c5ecdb8ab..1e4d4b9fa 100644 --- a/server/src/common/speculation/speculation_gate.h +++ b/server/src/common/speculation/speculation_gate.h @@ -1,7 +1,6 @@ -// Per-request adaptive speculation policy over startup-profiled costs. Every -// adaptive request receives one AR or speculation decision and keeps that -// decision until forget(). A failed activation-score evaluation is represented -// explicitly and commits sticky AR without inventing a score. +// Cohort-planned adaptive speculation policy over startup-profiled costs. +// Request-local state contains only immutable activation knowledge. The engine +// owns the current cohort epoch and decides when to run plan() again. // Pure host code: no graph, backend, or scheduler types belong here. #pragma once @@ -95,12 +94,10 @@ struct SpecCandidate { // suspend speculative execution for one explicitly telemetered AR service // round without changing this capability or the chosen request mode. bool can_speculate = false; - // While an adaptive request is Undecided, NaN requests its one-time - // bootstrap and a finite value is the preferred activation measurement. - // Evaluation failure explicitly falls back to sticky AR without inventing - // a score. Otherwise the gate commits exactly one mode from this - // adapter-provided expected yield, including the root. Adapter estimates - // are already calibrated; the gate only clamps executor bounds. + // NaN requests a one-time bootstrap when no cached score exists. A finite + // value is the preferred activation measurement. Evaluation failure keeps + // the request AR without inventing a score. Adapter estimates are already + // calibrated; the gate only clamps executor bounds. double activation_yield = std::numeric_limits::quiet_NaN(); std::vector conditional_hazards; std::string score_kind = kUnspecifiedScoreKind; @@ -141,12 +138,6 @@ inline const char * spec_score_source_name(SpecScoreSource source) { return "unknown"; } -enum class SpecDecision : uint8_t { - Undecided, - AR, - Speculation, -}; - enum class SpecEvaluationAction : uint8_t { Score, FallbackAR, @@ -158,27 +149,13 @@ struct SpecPendingEvaluation { SpecEvaluationAction action = SpecEvaluationAction::Score; }; -inline const char * spec_decision_name(SpecDecision decision) { - switch (decision) { - case SpecDecision::Undecided: return "undecided"; - case SpecDecision::AR: return "ar"; - case SpecDecision::Speculation: return "speculation"; - } - return "unknown"; -} - struct SpecPlanScore { uint64_t request_id = 0; int slot = -1; double expected_yield = 1.0; SpecScoreSource source = SpecScoreSource::Unavailable; - SpecDecision decision = SpecDecision::Undecided; bool forced = false; bool admitted = false; - // True only in the plan that commits this request's immutable adaptive - // AR/speculation decision. This supports one activation record per - // request without treating later sticky execution as a new decision. - bool newly_decided = false; std::string score_kind = kUnspecifiedScoreKind; bool execution_unsupported = false; }; @@ -222,28 +199,44 @@ struct SpecPlan { double goodput = 0.0; double ar_goodput = 0.0; int unavailable_count = 0; - // False means at least one adaptive request still needs its one-time - // evaluation action. No new adaptive decisions are committed in that - // plan. The engine resolves every tagged action request-locally, then one - // immediate replan commits all remaining undecided lanes atomically. - bool decisions_committed = false; + // The engine resolves pending actions, immediately replans, and caches + // only the completed result as the current cohort epoch. bool cost_lookup_clamped = false; std::vector ordered; std::vector admitted_request_ids; std::vector admitted_slots; // Score actions are batched for one-time activation-score initialization. - // FallbackAR actions cannot attempt scoring and must instead commit sticky - // AR with an explicit failed-evaluation activation. One tagged record keeps + // FallbackAR actions cannot attempt scoring and instead record an + // explicit failed evaluation. One tagged record keeps // request identity and slot inseparable on all failure paths. std::vector pending_evaluations; }; +struct SpecCohortEpoch { + uint64_t id = 0; + std::vector request_ids; + SpecPlan plan; + + bool matches(const std::vector & candidates) const { + return request_ids == ids(candidates); + } + + static std::vector ids( + const std::vector & candidates) { + std::vector out; + out.reserve(candidates.size()); + for (const SpecCandidate & candidate : candidates) + out.push_back(candidate.request_id); + std::sort(out.begin(), out.end()); + return out; + } +}; + class SpeculationGate { private: struct RequestState { double initial_score = std::numeric_limits::quiet_NaN(); - SpecDecision decision = SpecDecision::Undecided; bool evaluation_failed = false; std::string score_kind = kUnspecifiedScoreKind; std::vector conditional_hazards; @@ -315,12 +308,11 @@ class SpeculationGate { } // draft_lanes_override prices always-drafting. -1 means admitted-only. - // Non-committing plans retain a first score while leaving the request - // Undecided so a later, formed cohort can publish the sticky decision. + // The engine calls this only for a new cohort epoch and once more after + // resolving any cold-score actions. SpecPlan plan(int concurrency, const std::vector & candidates, - int k_cap, int draft_lanes_override = -1, - bool commit_decisions = true) { + int k_cap, int draft_lanes_override = -1) { SpecPlan out; out.concurrency = concurrency; if (!valid() || concurrency < 0 || @@ -336,33 +328,27 @@ class SpeculationGate { double score = 1.0; SpecScoreSource source = SpecScoreSource::Unavailable; std::string score_kind = kUnspecifiedScoreKind; - SpecDecision decision = SpecDecision::Undecided; bool forced = false; - bool commit_candidate = false; }; std::vector forced; - std::vector undecided; + std::vector adaptive; std::vector forced_ar; forced.reserve(candidates.size()); - undecided.reserve(candidates.size()); + adaptive.reserve(candidates.size()); forced_ar.reserve(candidates.size()); for (const SpecCandidate & candidate : candidates) { - if (candidate.policy == SpeculationPolicy::Never) { - continue; - } - const SpecDecision prior_decision = decision(candidate.request_id); + if (candidate.policy == SpeculationPolicy::Never) continue; if (candidate.policy == SpeculationPolicy::Adaptive && - prior_decision == SpecDecision::AR) { - // A one-shot AR decision removes the request from all future - // adaptive rankings until forget(). + evaluation_failed(candidate.request_id)) { + forced_ar.push_back({ + &candidate, 1.0, SpecScoreSource::Unavailable, + initial_score_kind(candidate.request_id), false}); continue; } + const CandidateScore score = score_candidate(candidate); - const bool adaptive_undecided = - candidate.policy == SpeculationPolicy::Adaptive && - prior_decision == SpecDecision::Undecided; - if (adaptive_undecided && + if (candidate.policy == SpeculationPolicy::Adaptive && score.source == SpecScoreSource::Unavailable) { ++out.unavailable_count; out.pending_evaluations.push_back({ @@ -373,45 +359,31 @@ class SpeculationGate { }); continue; } - if (adaptive_undecided && !candidate.can_speculate) { - // Record the mandatory initial score in activation telemetry, - // but permanently unsupported execution commits directly to AR. + if (candidate.policy == SpeculationPolicy::Adaptive && + !candidate.can_speculate) { forced_ar.push_back({ &candidate, score.expected_yield, score.source, - score.score_kind, prior_decision, false, true}); + score.score_kind, false}); continue; } - Ranked ranked{&candidate, score.expected_yield, score.source, - score.score_kind, prior_decision, - candidate.policy == SpeculationPolicy::Always || - prior_decision == SpecDecision::Speculation, - adaptive_undecided}; - (ranked.forced ? forced : undecided).push_back(ranked); + Ranked ranked{ + &candidate, score.expected_yield, score.source, + score.score_kind, + candidate.policy == SpeculationPolicy::Always}; + (ranked.forced ? forced : adaptive).push_back(ranked); } - // Selection is atomic across the newly arriving batch. If any cold - // adaptive request still needs its score, defer every undecided lane - // until the engine publishes the batched bootstrap and immediately - // replans. Previously decided speculation and explicit Always lanes - // can still execute in the bootstrap plan. - if (!out.pending_evaluations.empty()) { - undecided.clear(); - } - out.decisions_committed = - out.pending_evaluations.empty() && commit_decisions; - if (out.decisions_committed) { - for (const Ranked & item : forced_ar) { - request_states_[item.candidate->request_id].decision = - SpecDecision::AR; - } - } + // A cohort plan is publishable only after every cold adaptive request + // has resolved. Explicit Always lanes may still run in the bootstrap + // service plan, but the engine never caches that partial result. + if (!out.pending_evaluations.empty()) adaptive.clear(); auto request_order = [](const Ranked & a, const Ranked & b) { return a.candidate->request_id < b.candidate->request_id; }; std::sort(forced.begin(), forced.end(), request_order); std::sort(forced_ar.begin(), forced_ar.end(), request_order); - std::sort(undecided.begin(), undecided.end(), + std::sort(adaptive.begin(), adaptive.end(), [](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; @@ -424,16 +396,20 @@ class SpeculationGate { } std::vector ranked; - ranked.reserve(forced.size() + undecided.size()); + ranked.reserve(forced.size() + adaptive.size()); ranked.insert(ranked.end(), forced.begin(), forced.end()); - ranked.insert(ranked.end(), undecided.begin(), undecided.end()); + ranked.insert(ranked.end(), adaptive.begin(), adaptive.end()); for (const Ranked & item : ranked) { - out.ordered.push_back({item.candidate->request_id, - item.candidate->slot, - item.score, item.source, - item.decision, - item.forced, false, - item.commit_candidate, item.score_kind}); + out.ordered.push_back({ + item.candidate->request_id, + item.candidate->slot, + item.score, + item.source, + item.forced, + false, + item.score_kind, + false, + }); } const int forced_count = static_cast(forced.size()); @@ -445,8 +421,7 @@ class SpeculationGate { concurrency, 0, 0, geometry_.bucketed_lanes(concurrency), 0}; const CostPrice ar_price = price_execution_shape(ar_shape, &out); out.ar_goodput = concurrency == 0 ? 0.0 - : static_cast(concurrency) / - ar_price.predicted; + : static_cast(concurrency) / ar_price.predicted; struct PlanPoint { int k = 0; @@ -473,9 +448,6 @@ class SpeculationGate { if (k > 0) { if (direct_commit_) { - // Direct promotion packs one AR row per peer before the - // bucketed fixed-depth speculative tree and performs no - // target replay. tree_rows = geometry_.tree_rows(k) + concurrency - k; expected_step_rows = 0.0; } else { @@ -489,19 +461,15 @@ class SpeculationGate { expected_step_rows, &out); const double scale = price.predicted / price.profiled; const double goodput = expected_tokens / price.predicted; - const PlanPoint point{ k, goodput, price.profiled, price.predicted, scale, expected_tokens, tree_rows, expected_step_rows, draft_lanes}; if (k == forced_count) baseline = point; - if (goodput > best.goodput) { - best = point; - } + if (goodput > best.goodput) best = point; } - // Explicit Always and sticky-Speculation lanes establish the - // non-negotiable baseline. The safety margin applies only to the - // one-shot admission of additional undecided adaptive lanes. + // Explicit Always lanes establish the non-negotiable baseline. The + // safety margin applies to this epoch's adaptive subset selection. if (best.k > forced_count && best.goodput < baseline.goodput * (1.0 + config_.adaptive_gain_margin)) { @@ -517,39 +485,29 @@ class SpeculationGate { out.tree_rows = best.tree_rows; out.expected_step_rows = best.expected_step_rows; out.draft_lanes = best.draft_lanes; - for (size_t i = 0; i < ranked.size(); ++i) { - if (!ranked[i].commit_candidate) continue; - const SpecDecision committed = static_cast(i) < best.k - ? SpecDecision::Speculation : SpecDecision::AR; - if (commit_decisions) { - request_states_[ranked[i].candidate->request_id].decision = - committed; - out.ordered[i].decision = committed; - } - } for (int i = 0; i < best.k; ++i) { - out.ordered[(size_t)i].admitted = true; - out.admitted_request_ids.push_back(ranked[(size_t)i].candidate->request_id); - out.admitted_slots.push_back(ranked[(size_t)i].candidate->slot); - if (ranked[(size_t)i].source != SpecScoreSource::Unavailable) { - out.initial_predicted_tokens += ranked[(size_t)i].score; + out.ordered[static_cast(i)].admitted = true; + out.admitted_request_ids.push_back( + ranked[static_cast(i)].candidate->request_id); + out.admitted_slots.push_back( + ranked[static_cast(i)].candidate->slot); + if (ranked[static_cast(i)].source != + SpecScoreSource::Unavailable) { + out.initial_predicted_tokens += + ranked[static_cast(i)].score; } } - if (out.decisions_committed) { - for (const Ranked & item : forced_ar) { - out.ordered.push_back({ - item.candidate->request_id, - item.candidate->slot, - item.score, - item.source, - SpecDecision::AR, - false, - false, - true, - item.score_kind, - true, - }); - } + for (const Ranked & item : forced_ar) { + out.ordered.push_back({ + item.candidate->request_id, + item.candidate->slot, + item.score, + item.source, + false, + false, + item.score_kind, + true, + }); } return out; } @@ -576,16 +534,13 @@ class SpeculationGate { config_.cost_ema_alpha); } - // Commit the explicit cold-evaluation failure policy. This is a real - // sticky AR decision, but intentionally has no initial activation score. - // False means the request had already committed a mode. - bool commit_evaluation_fallback_ar(uint64_t request_id) { + // Record a cold evaluation failure without inventing a score. The + // request remains AR in every later cohort because it cannot be ranked. + bool record_evaluation_failure(uint64_t request_id) { RequestState & state = request_states_[request_id]; - if (state.decision != SpecDecision::Undecided) return false; - state.initial_score = - std::numeric_limits::quiet_NaN(); + if (state.evaluation_failed || + std::isfinite(state.initial_score)) return false; state.evaluation_failed = true; - state.decision = SpecDecision::AR; return true; } @@ -623,11 +578,6 @@ class SpeculationGate { return state == request_states_.end() ? empty : state->second.conditional_hazards; } - SpecDecision decision(uint64_t request_id) const { - auto state = request_states_.find(request_id); - return state == request_states_.end() - ? SpecDecision::Undecided : state->second.decision; - } const SpecCostTables & costs() const { return costs_; } private: diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index d6a3431d6..2c1c2af9d 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -25,7 +25,9 @@ #include #include #include +#include #include +#include #include #include @@ -100,7 +102,7 @@ void log_spec_gate_plan( std::fprintf(stderr, "%s%llu:%s", i == 0 ? "" : ",", static_cast(score.request_id), - spec_decision_name(score.decision)); + score.admitted ? "speculation" : "ar"); } std::fprintf(stderr, "] sources=fresh:%d,initial:%d,unavailable:%d " @@ -123,34 +125,69 @@ void log_spec_gate_plan( } } -void log_spec_activations( - const SpecPlan & plan, +void log_spec_epoch( + const SpecCohortEpoch & epoch, const SpeculationGate & gate) { - for (const SpecPlanScore & score : plan.ordered) { - if (!score.newly_decided) continue; + const SpecPlan & plan = epoch.plan; + std::fprintf(stderr, + "[spec-epoch] {\"epoch_id\":%llu,\"request_ids\":[", + static_cast(epoch.id)); + for (size_t i = 0; i < epoch.request_ids.size(); ++i) { + std::fprintf(stderr, "%s%llu", i == 0 ? "" : ",", + static_cast(epoch.request_ids[i])); + } + std::fprintf(stderr, "],\"selected_request_ids\":["); + for (size_t i = 0; i < plan.admitted_request_ids.size(); ++i) { + std::fprintf(stderr, "%s%llu", i == 0 ? "" : ",", + static_cast(plan.admitted_request_ids[i])); + } + std::fprintf(stderr, "],\"requests\":["); + for (size_t i = 0; i < plan.ordered.size(); ++i) { + const SpecPlanScore & score = plan.ordered[i]; + const bool failed = gate.evaluation_failed(score.request_id); const double initial = gate.initial_score(score.request_id); const std::string kind = gate.initial_score_kind(score.request_id); const std::vector & hazards = gate.initial_hazards(score.request_id); - const char * decision_reason = score.execution_unsupported - ? "execution_unsupported" - : score.decision == SpecDecision::Speculation - ? "selected_by_joint_goodput" - : "ar_counterfactual_won"; std::fprintf(stderr, - "[spec-activation] {\"request_id\":%llu,\"slot\":%d," - "\"activation_score\":%.6f,\"score_kind\":\"%s\"," - "\"expected_yield\":%.6f,\"hazards\":[", - static_cast(score.request_id), - score.slot, initial, kind.c_str(), score.expected_yield); - for (size_t i = 0; i < hazards.size(); ++i) { - std::fprintf(stderr, "%s%.8g", i == 0 ? "" : ",", hazards[i]); - } + "%s{\"request_id\":%llu,\"slot\":%d," + "\"activation_score\":", + i == 0 ? "" : ",", + static_cast(score.request_id), score.slot); + if (std::isfinite(initial)) std::fprintf(stderr, "%.6f", initial); + else std::fprintf(stderr, "null"); + std::fprintf(stderr, + ",\"score_kind\":\"%s\",\"expected_yield\":", + kind.c_str()); + if (std::isfinite(initial)) + std::fprintf(stderr, "%.6f", score.expected_yield); + else std::fprintf(stderr, "null"); + std::fprintf(stderr, ",\"hazards\":"); + if (std::isfinite(initial)) { + std::fprintf(stderr, "["); + for (size_t j = 0; j < hazards.size(); ++j) + std::fprintf(stderr, "%s%.8g", j == 0 ? "" : ",", hazards[j]); + std::fprintf(stderr, "]"); + } else { + std::fprintf(stderr, "null"); + } + const char * evaluation = failed ? "failed" + : std::isfinite(initial) ? "scored" : "unavailable"; + const char * reason = failed ? "activation_evaluation_failed" + : score.execution_unsupported ? "execution_unsupported" + : score.admitted ? "selected_by_joint_goodput" + : "ar_counterfactual_won"; std::fprintf(stderr, - "],\"evaluation\":\"scored\",\"fallback_reason\":null," - "\"decision_reason\":\"%s\",\"decision\":\"%s\"}\n", - decision_reason, spec_decision_name(score.decision)); + ",\"evaluation\":\"%s\",\"route\":\"%s\"," + "\"reason\":\"%s\"}", + evaluation, score.admitted ? "speculation" : "ar", reason); } + std::fprintf(stderr, + "],\"profiled_cost_us\":%.1f,\"cost_scale\":%.6f," + "\"predicted_cost_us\":%.1f,\"goodput\":%.9f," + "\"ar_goodput\":%.9f}\n", + plan.profiled_cost, plan.cost_scale, plan.predicted_cost, + plan.goodput, plan.ar_goodput); } void log_spec_evaluation_fallback( @@ -160,17 +197,12 @@ void log_spec_evaluation_fallback( const char * reason) { const std::string cause = reason ? reason : "activation_evaluation_failed"; - const char * decision_reason = cause == "activation_evaluation_failed" - ? "evaluation_failed" : cause.c_str(); std::fprintf(stderr, - "[spec-activation] {\"request_id\":%llu,\"slot\":%d," - "\"activation_score\":null,\"score_kind\":\"%s\"," - "\"expected_yield\":null,\"hazards\":null," - "\"evaluation\":\"failed\"," - "\"fallback_reason\":\"activation_evaluation_failed\"," - "\"decision_reason\":\"%s\",\"decision\":\"ar\"}\n", + "[spec-evaluation] {\"request_id\":%llu,\"slot\":%d," + "\"score_kind\":\"%s\",\"evaluation\":\"failed\"," + "\"reason\":\"%s\"}\n", static_cast(request_id), slot, kind.c_str(), - decision_reason); + cause.c_str()); } uint64_t file_size_or_zero(const char * path) { @@ -273,7 +305,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( speculator_fallback_reason(speculator_.get()); std::fprintf(stderr, "[parallel-chain] no speculator adapter: %s; " - "adaptive requests will use sticky AR\n", + "adaptive requests will use AR fallback\n", error.c_str()); } else { std::fprintf(stderr, @@ -288,7 +320,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( speculator_fallback_reason(speculator_.get()); std::fprintf(stderr, "[parallel-chain] no speculator adapter for loaded drafter; " - "adaptive requests will use sticky AR\n"); + "adaptive requests will use AR fallback\n"); } // The concurrent DDTree stack is gated to a local same-device drafter. @@ -365,6 +397,8 @@ bool Qwen35SeqEngine::step_timing_enabled() { bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { adaptive_fallback_reason_ = "cost_profile_unavailable"; speculation_gate_.reset(); + spec_cohort_epoch_.reset(); + next_spec_cohort_epoch_id_ = 1; if (spec_mode_ != SpecMode::chain || !capture_features_ || !activation_scoring_available() || tree_width_ <= 1 || tree_width_ > 16 || resolve_chain_verify_depth(chain_verify_depth_, tree_width_) == 0 || @@ -385,6 +419,99 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { slots_.max_context() - T, b_.cache_.target_feat_cap); if (n_slots < 1 || max_profile_ctx < 1) return false; const int ctx_tokens = std::clamp(context_tokens, 1, max_profile_ctx); + const SpecProfileGrid grid = build_spec_profile_grid( + n_slots, V, V, [](int lanes) { + return chain_decode_bucket_width(lanes); + }); + const bool profile_batched = batched_drafting_enabled(); + + std::ostringstream identity; + identity << "qwen35-spec-cost-v1" + << "|slots=" << n_slots + << "|tree=" << T + << "|verify=" << V + << "|ctx=" << ctx_tokens + << "|max_ctx=" << slots_.max_context() + << "|kq_pad=" << b_.cfg_.kq_stride_pad + << "|draft_mode=" << (profile_batched ? "batched" : "serial") + << "|commit_mode=" + << (chain_direct_commit_enabled() ? "direct" : "replay") + << "|speculator=" << speculator_->score_kind() + << "|target_device=" << placement_device_name(b_.cfg_.device) + << "|draft_device=" << b_.cfg_.draft_gpu; + if (ggml_backend_dev_t device = + ggml_backend_get_device(b_.target_backend_)) { + identity << "|device_name=" << ggml_backend_dev_name(device) + << "|device_description=" + << ggml_backend_dev_description(device); + } + auto append_file_identity = [&](const char * label, const char * path) { + identity << '|' << label << '=' << (path ? path : ""); + if (!path || !*path) return; + std::error_code ec; + const uintmax_t size = std::filesystem::file_size(path, ec); + identity << ":size=" << (ec ? 0 : size); + ec.clear(); + const auto modified = std::filesystem::last_write_time(path, ec); + identity << ":mtime=" << (ec ? 0 : modified.time_since_epoch().count()); + }; + append_file_identity("target", b_.cfg_.target_path); + append_file_identity("draft", b_.cfg_.draft_path); + auto append_grid = [&](const char * label, const std::vector & values) { + identity << '|' << label << '='; + for (size_t i = 0; i < values.size(); ++i) { + if (i) identity << ','; + identity << values[i]; + } + }; + append_grid("tree_rows", grid.tree_rows); + append_grid("step_rows", grid.step_rows); + append_grid("draft_lanes", grid.draft_lanes); + const std::string profile_identity = identity.str(); + const std::string profile_cache_path = + spec_cost_profile_cache_path(profile_identity); + + auto install_profile = [&](const SpecCostTables & tables) { + SpecStepGeometry geometry; + geometry.tree_width = V; + geometry.bucket = [](int lanes) { + return chain_decode_bucket_width(lanes); + }; + speculation_gate_ = std::make_unique( + tables, geometry, V, + [](const char * table, int requested, int profiled) { + std::fprintf(stderr, + "[spec-gate] %s_cost index %d outside profile; " + "clamped to %d\n", + table, requested, profiled); + }, SpecGateConfig{}, chain_direct_commit_enabled()); + if (!speculation_gate_->valid()) { + speculation_gate_.reset(); + return false; + } + adaptive_fallback_reason_.clear(); + return true; + }; + + SpecCostTables cached_tables; + std::string cache_error; + if (load_spec_cost_profile( + profile_cache_path, profile_identity, + cached_tables, cache_error) && + install_profile(cached_tables)) { + std::fprintf(stderr, + "[spec-profile] loaded %s context=%d mode=%s-draft " + "speculator=%s\n", + profile_cache_path.c_str(), ctx_tokens, + profile_batched ? "batched" : "serial", + cached_tables.speculator_id.c_str()); + return true; + } + if (!profile_cache_path.empty() && + cache_error != "profile cache miss") { + std::fprintf(stderr, "[spec-profile] cache ignored: %s\n", + cache_error.c_str()); + } std::string profile_error; std::vector synthetic_slots; @@ -450,11 +577,6 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { sizeof(int32_t) * seq_lens_.size()); ggml_backend_synchronize(b_.target_backend_); - const SpecProfileGrid grid = build_spec_profile_grid( - n_slots, V, V, [](int lanes) { - return chain_decode_bucket_width(lanes); - }); - int prepared_tree_rows = -1; auto tree_runner = [&](int total_rows) -> double { if (!profile_error.empty()) @@ -639,7 +761,6 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { std::chrono::steady_clock::now() - start).count(); }; - const bool profile_batched = batched_drafting_enabled(); auto draft_runner = [&](int lanes) -> double { if (!profile_error.empty()) { return std::numeric_limits::infinity(); @@ -688,21 +809,14 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { } SpecCostTables tables = std::move(profiled.tables); - SpecStepGeometry geometry; - geometry.tree_width = V; - geometry.bucket = [](int lanes) { - return chain_decode_bucket_width(lanes); - }; - speculation_gate_ = std::make_unique( - tables, geometry, V, - [](const char * table, int requested, int profiled) { - std::fprintf(stderr, - "[spec-gate] %s_cost index %d outside profile; clamped to %d\n", - table, requested, profiled); - }, SpecGateConfig{}, chain_direct_commit_enabled()); - if (!speculation_gate_->valid()) { - speculation_gate_.reset(); - return false; + if (!install_profile(tables)) return false; + if (!save_spec_cost_profile( + profile_cache_path, profile_identity, tables, cache_error)) { + std::fprintf(stderr, "[spec-profile] cache save failed: %s\n", + cache_error.c_str()); + } else if (!profile_cache_path.empty()) { + std::fprintf(stderr, "[spec-profile] saved %s\n", + profile_cache_path.c_str()); } auto print_table = [](const char * name, const SpecCostSeries & series) { @@ -721,7 +835,6 @@ bool Qwen35SeqEngine::profile_spec_costs(int context_tokens) { print_table("tree_cost", tables.tree_cost); print_table("step_cost", tables.step_cost); print_table("draft_cost", tables.draft_cost); - adaptive_fallback_reason_.clear(); return true; } @@ -1233,7 +1346,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( if (!admitted[i]) continue; if (!hard_eligible) { fail_proposal_lane( - i, "sticky speculation request became ineligible"); + i, "epoch-selected speculation request became ineligible"); continue; } const PreparedChainDraft & prepared = @@ -2951,7 +3064,6 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { std::vector candidates; candidates.reserve(inputs.size()); const bool use_activation_score = activation_scoring_enabled(); - const bool commit_gate_decisions = plan.prefills.empty(); for (const StepInput & in : inputs) { const Qwen35Slot & seq = slots_.slot(in.slot); SpeculationPolicy policy = in.speculation_policy; @@ -2987,9 +3099,16 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { chain_activation_score_kind(), }); } - gate_plan = speculation_gate_->plan( - (int)inputs.size(), candidates, (int)inputs.size(), - -1, commit_gate_decisions); + const bool cohort_changed = + !spec_cohort_epoch_.has_value() || + !spec_cohort_epoch_->matches(candidates); + bool published_epoch = false; + if (cohort_changed) { + gate_plan = speculation_gate_->plan( + (int)inputs.size(), candidates, (int)inputs.size()); + } else { + gate_plan = spec_cohort_epoch_->plan; + } have_gate_plan = true; if (!gate_plan.valid) { return fail_step(gate_plan.error.empty() @@ -3012,8 +3131,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } } gate_plan = speculation_gate_->plan( - (int)inputs.size(), candidates, (int)inputs.size(), - -1, commit_gate_decisions); + (int)inputs.size(), candidates, (int)inputs.size()); return gate_plan.valid; }; @@ -3034,7 +3152,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { auto commit_evaluation_fallback = [&](const SpecPendingEvaluation & evaluation) { reset_evaluation_lane(evaluation.slot); - if (speculation_gate_->commit_evaluation_fallback_ar( + if (speculation_gate_->record_evaluation_failure( evaluation.request_id)) { const std::string kind = chain_activation_score_kind(); @@ -3048,7 +3166,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // Resolve every one-time evaluation action before one immediate // replan. A packed bootstrap failure is retried once per lane so - // one broken request falls back to sticky AR without poisoning a + // one broken request is marked evaluation-failed without poisoning a // healthy scored peer or the cohort. if (!gate_plan.pending_evaluations.empty()) { std::vector score_evaluations; @@ -3133,25 +3251,25 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { : gate_plan.error); } } - if (!gate_plan.pending_evaluations.empty() || - (!gate_plan.decisions_committed && - commit_gate_decisions)) { + if (!gate_plan.pending_evaluations.empty()) { return fail_step( - "adaptive activation state did not commit modes"); + "adaptive activation score did not resolve"); } } - if (gate_plan.decisions_committed) { - log_spec_activations(gate_plan, *speculation_gate_); + if (cohort_changed) { + SpecCohortEpoch epoch; + epoch.id = next_spec_cohort_epoch_id_++; + epoch.request_ids = SpecCohortEpoch::ids(candidates); + epoch.plan = gate_plan; + spec_cohort_epoch_ = std::move(epoch); + log_spec_epoch(*spec_cohort_epoch_, *speculation_gate_); + published_epoch = true; } - // A one-shot AR decision discards the evaluation proposal. SPEC - // keeps that exact first proposal so the bootstrap is useful work. + // Discard bootstrap proposals for lanes routed to AR in the new + // epoch. Selected lanes keep that first proposal as useful work. for (const SpecPlanScore & score : gate_plan.ordered) { - if (!score.newly_decided || - (gate_plan.decisions_committed && - score.decision != SpecDecision::AR)) { - continue; - } + if (!published_epoch || score.admitted) continue; const int slot = score.slot; if (slot >= 0 && slot < (int)prepared_chain_drafts_.size()) { @@ -3163,16 +3281,10 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } } - // Preserve every sticky Spec admission. The min-token floor is - // enforced inside the speculative path, capability failures remain - // lane-local errors, and planned prefills use an explicitly - // telemetered AR service round without changing the request's - // sticky routing decision. + // Apply the cached epoch subset. Planned prefills use a + // telemetered AR service round without changing the epoch plan. for (const SpecPlanScore & score : gate_plan.ordered) { - if (!score.admitted || - (!gate_plan.decisions_committed && !score.forced)) { - continue; - } + if (!score.admitted) continue; const int slot = score.slot; bool found = false; for (size_t i = 0; i < inputs.size(); ++i) { @@ -3285,7 +3397,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // Chain verification cannot share a target graph with prompt work. // Use the existing packed AR+prefill graph for this service round // so selected prompts make immediate progress. Routing remains - // sticky SPEC; the per-request metric distinguishes this bounded + // fixed for the current epoch; the per-request metric distinguishes this bounded // scheduling suspension from speculative execution. spec_service_ar = admitted; for (size_t i = 0; i < inputs.size(); ++i) { diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index fcd72ac0c..793d38e75 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -80,7 +80,7 @@ class Qwen35SeqEngine final : public SeqEngine { bool profile_spec_costs(int context_tokens); // True only when the registered speculator adapter can produce a // first-request activation score. A configured chain - // may still accept Adaptive requests and serve sticky AR when this is false. + // may still accept Adaptive requests and serve AR fallback when this is false. bool activation_scoring_available() const; StepPlanLimits step_plan_limits(int decode_rows) const override { const bool mixed = decode_rows > 0; @@ -160,8 +160,8 @@ class Qwen35SeqEngine final : public SeqEngine { bool batched_drafting_enabled() const; bool activation_scoring_enabled() const; std::string chain_activation_score_kind() const; - // Deliberate seam for a later per-round cohort controller. Request mode - // remains sticky; every returned depth must stay in [2, tree_width_]. + // Verification depth may vary between rounds while the cohort route stays + // fixed; every returned depth must stay in [2, tree_width_]. int chain_verify_depth_for_round() const { return chain_verify_depth_; } @@ -171,7 +171,7 @@ class Qwen35SeqEngine final : public SeqEngine { static bool step_timing_enabled(); // DDTree preserves its legacy best-effort AR fallback. Chain // proposal failures are instead returned as lane-local DecodeOutput - // failures so a sticky speculation decision can never execute as AR. + // failures so an epoch-selected speculation lane never silently executes as AR. std::optional step_ddtree(const StepPlan & plan); StepResult step_chain_spec( const StepPlan & plan, const std::vector & admitted, @@ -183,7 +183,7 @@ class Qwen35SeqEngine final : public SeqEngine { int tree_width_ = 0; // Root-inclusive verification depth. The drafter still produces // tree_width_ tokens; this common cohort depth may vary between 2 and - // tree_width_ without changing a request's sticky SPEC mode. + // tree_width_ without changing the current cohort plan. int chain_verify_depth_ = 0; int tree_scratch_base_ = 0; int tree_scratch_stride_ = 0; @@ -197,9 +197,11 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector> dummy_draft_kv_; std::vector prepared_chain_drafts_; std::unique_ptr speculation_gate_; + std::optional spec_cohort_epoch_; + uint64_t next_spec_cohort_epoch_id_ = 1; std::unique_ptr speculator_; - // Startup profile/adapter failure is a request-local sticky AR outcome, - // never an admission or step error for a configured chain. + // Startup profile/adapter failure is a request-local AR outcome, never an + // admission or step error for a configured chain. std::string adaptive_fallback_reason_ = "cost_profile_unavailable"; std::vector adaptive_fallback_ar_; std::vector last_activation_estimate_; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 03336699a..3213c1cdf 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -606,8 +606,8 @@ bool Qwen35Backend::init() { // Per-request decode_mode may select Adaptive even when the // server default is forced speculation. A configured chain always // accepts Adaptive: if activation scoring or startup profiling is - // unavailable, the engine records a request-local sticky-AR - // fallback instead of rejecting the request or failing its peers. + // unavailable, the engine uses an AR fallback instead of rejecting + // the request or failing its peers. bool adaptive_scored = false; if (concurrent_local_chain && seq_engine_->activation_scoring_available()) { @@ -623,14 +623,14 @@ bool Qwen35Backend::init() { std::fprintf(stderr, "[parallel-chain] adaptive scoring unavailable: " "cost profile failed; adaptive requests will use " - "request-local sticky AR (forced speculation remains " + "AR fallback (forced speculation remains " "available)\n"); } } else if (concurrent_local_chain) { std::fprintf(stderr, "[parallel-chain] adaptive activation unavailable: " "drafter has no compatible request-benefit adapter; " - "adaptive requests will use request-local sticky AR " + "adaptive requests will use AR fallback " "(forced speculation remains available)\n"); } if (concurrent_local_ddtree) { diff --git a/server/test/test_dflash2_benefit.cpp b/server/test/test_dflash2_benefit.cpp index 4b7e12bc2..c6bbbaf09 100644 --- a/server/test/test_dflash2_benefit.cpp +++ b/server/test/test_dflash2_benefit.cpp @@ -108,7 +108,6 @@ int main() { SpecPlan gate_plan = code_c1.plan( 1, {benefit_candidate(100, 5, code.expected_yield)}, 1); CHECK(gate_plan.admitted_count == 1); - CHECK(code_c1.decision(100) == SpecDecision::Speculation); CHECK(code_c1.initial_score_kind(100) == kDFlash2BenefitAdapterVersion); @@ -116,15 +115,12 @@ int main() { gate_plan = prose_c1.plan( 1, {benefit_candidate(200, 3, prose.expected_yield)}, 1); CHECK(gate_plan.admitted_count == 0); - CHECK(prose_c1.decision(200) == SpecDecision::AR); SpeculationGate prose_c2(observed_costs, observed_geometry, 8); gate_plan = prose_c2.plan(2, { benefit_candidate(300, 4, prose.expected_yield), benefit_candidate(301, 1, prose.expected_yield)}, 2); CHECK(gate_plan.admitted_count == 0); - CHECK(prose_c2.decision(300) == SpecDecision::AR); - CHECK(prose_c2.decision(301) == SpecDecision::AR); SpeculationGate mixed_c2(observed_costs, observed_geometry, 8); gate_plan = mixed_c2.plan(2, { @@ -133,16 +129,14 @@ int main() { CHECK(gate_plan.admitted_count >= 1); CHECK(!gate_plan.admitted_request_ids.empty()); CHECK(gate_plan.admitted_request_ids[0] == 400); - CHECK(mixed_c2.decision(400) == SpecDecision::Speculation); CHECK(gate_plan.ordered.size() == 2); CHECK(gate_plan.ordered[0].request_id == 400); CHECK(gate_plan.ordered[0].slot == 5); CHECK(gate_plan.ordered[1].request_id == 401); CHECK(gate_plan.ordered[1].slot == 3); - CHECK((gate_plan.admitted_count == 1 && - mixed_c2.decision(401) == SpecDecision::AR) || - (gate_plan.admitted_count == 2 && - mixed_c2.decision(401) == SpecDecision::Speculation)); + CHECK(gate_plan.ordered[0].admitted); + CHECK(gate_plan.ordered[1].admitted == + (gate_plan.admitted_count == 2)); SpeculationGate code_with_ar_peer( observed_costs, observed_geometry, 8); @@ -151,8 +145,8 @@ int main() { benefit_candidate(501, 7, prose.expected_yield, SpeculationPolicy::Never)}, 2); CHECK(gate_plan.admitted_count == 1); - CHECK(code_with_ar_peer.decision(500) == SpecDecision::Speculation); - CHECK(code_with_ar_peer.decision(501) == SpecDecision::Undecided); + CHECK(gate_plan.admitted_request_ids.size() == 1); + CHECK(gate_plan.admitted_request_ids[0] == 500); // Maximum depth consumes exactly block_size-1 signals; diagnostic tail // values beyond that depth cannot affect the score. diff --git a/server/test/test_spec_cost_profile.cpp b/server/test/test_spec_cost_profile.cpp index d65b31df9..990c4603c 100644 --- a/server/test/test_spec_cost_profile.cpp +++ b/server/test/test_spec_cost_profile.cpp @@ -4,10 +4,14 @@ #include #include #include +#include +#include #include #include #include +#include + using namespace dflash::common; static int g_checks = 0; @@ -91,6 +95,32 @@ int main() { CHECK(!bad.ok()); CHECK(bad.tables.tree_cost.indices.empty()); + const std::filesystem::path cache_dir = + std::filesystem::temp_directory_path() / + ("dflash-spec-profile-test-" + std::to_string(getpid())); + const std::filesystem::path cache_path = cache_dir / "profile"; + std::string cache_error; + CHECK(save_spec_cost_profile( + cache_path.string(), "identity-a", profiled.tables, cache_error)); + SpecCostTables loaded; + CHECK(load_spec_cost_profile( + cache_path.string(), "identity-a", loaded, cache_error)); + CHECK(loaded.speculator_id == profiled.tables.speculator_id); + CHECK(loaded.tree_cost.indices == profiled.tables.tree_cost.indices); + CHECK(loaded.tree_cost.costs == profiled.tables.tree_cost.costs); + CHECK(!load_spec_cost_profile( + cache_path.string(), "identity-b", loaded, cache_error)); + CHECK(loaded.tree_cost.indices.empty()); + + { + std::ofstream corrupt(cache_path, std::ios::trunc); + corrupt << "not a profile\n"; + } + CHECK(!load_spec_cost_profile( + cache_path.string(), "identity-a", loaded, cache_error)); + std::error_code remove_error; + std::filesystem::remove_all(cache_dir, remove_error); + std::printf("spec cost profile tests passed: %d checks\n", g_checks); return 0; } diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index 015a5880b..de67ffc57 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -44,7 +44,7 @@ static SpecCandidate candidate( int main() { // A drafter without a registered adapter (including a DSpark-only GGUF) - // is not scoreable and reports the generic sticky-AR fallback reason. + // is not scoreable and reports the generic fallback reason. CHECK(!speculator_is_ready(nullptr)); CHECK(std::string(speculator_fallback_reason(nullptr)) == "no_speculator_adapter"); @@ -53,14 +53,9 @@ int main() { .expected_yield - 1.75) < 1e-9); CHECK(confidence_scorer.score({2.0f, -1.0f}, 4).expected_yield == 2.0); CHECK(confidence_scorer.score({}, 4).expected_yield == 1.0); - CHECK(std::string(spec_decision_name(SpecDecision::Undecided)) == - "undecided"); - CHECK(std::string(spec_decision_name(SpecDecision::AR)) == "ar"); - CHECK(std::string(spec_decision_name(SpecDecision::Speculation)) == - "speculation"); - - // A first finite score is a complete one-shot evaluation. Expensive - // speculation commits both requests to sticky AR and prices k=0 as the + + // A first finite score is a complete cached evaluation. Expensive + // speculation routes both requests to AR and prices k=0 as the // exact pure-AR candidate (no draft tax). SpeculationGate costly(constant_costs(100.0, 10.0, 100.0), geometry(), 4); @@ -68,41 +63,22 @@ int main() { SpecPlan plan = costly.plan(2, { candidate(1, 0, 4.0), candidate(2, 1, 4.0)}, 2); CHECK(plan.valid); - CHECK(plan.decisions_committed); CHECK(plan.admitted_count == 0); CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); CHECK(plan.draft_lanes == 0); CHECK(plan.profiled_cost == 10.0); CHECK(plan.cost_scale == 1.0); CHECK(plan.predicted_cost == plan.profiled_cost); - CHECK(costly.decision(1) == SpecDecision::AR); - CHECK(costly.decision(2) == SpecDecision::AR); CHECK(costly.initial_score(1) == 4.0); - SpeculationGate provisional( - constant_costs(1.0, 10.0, 1.0), geometry(), 4); - SpecPlan provisional_plan = provisional.plan( - 1, {candidate(3, 0, 4.0)}, 1, -1, false); - CHECK(provisional_plan.valid); - CHECK(!provisional_plan.decisions_committed); - CHECK(provisional_plan.admitted_count == 1); - CHECK(provisional_plan.ordered[0].decision == SpecDecision::Undecided); - CHECK(provisional.decision(3) == SpecDecision::Undecided); - CHECK(provisional.initial_score(3) == 4.0); - provisional_plan = provisional.plan( - 1, {candidate(3, 0, NAN)}, 1); - CHECK(provisional_plan.decisions_committed); - CHECK(provisional.decision(3) == SpecDecision::Speculation); plan = costly.plan(2, { candidate(1, 0, NAN), candidate(2, 1, NAN)}, 2); CHECK(plan.pending_evaluations.empty()); - CHECK(plan.ordered.empty()); + CHECK(plan.ordered.size() == 2); CHECK(plan.admitted_count == 0); CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); - // The initial ranking is a prefix argmax. The selected request remains - // speculative and the rejected request remains AR even if later inputs - // present reversed scores; the first score is immutable. + // The initial ranking is a prefix argmax. The first score remains immutable when later inputs present reversed scores. SpeculationGate prefix(constant_costs(1.0, 10.0, 1.0), geometry(), 4); plan = prefix.plan(3, { @@ -112,32 +88,27 @@ int main() { CHECK((plan.admitted_request_ids == std::vector{10})); CHECK(plan.ordered.size() == 2); CHECK(plan.ordered[0].admitted); - CHECK(plan.ordered[0].decision == SpecDecision::Speculation); - CHECK(plan.ordered[1].decision == SpecDecision::AR); CHECK(plan.ordered[0].source == SpecScoreSource::Fresh); CHECK(std::string(spec_score_source_name(plan.ordered[0].source)) == "fresh"); - CHECK(prefix.decision(10) == SpecDecision::Speculation); - CHECK(prefix.decision(11) == SpecDecision::AR); plan = prefix.plan(3, { candidate(10, 0, 1.0), candidate(11, 1, 4.0), candidate(12, 2, 4.0, SpeculationPolicy::Never)}, 3); CHECK((plan.admitted_request_ids == std::vector{10})); - CHECK(plan.ordered.size() == 1); - CHECK(plan.ordered[0].forced); + CHECK(plan.ordered.size() == 2); + CHECK(!plan.ordered[0].forced); CHECK(plan.ordered[0].source == SpecScoreSource::Initial); CHECK(prefix.initial_score(10) == 4.0); CHECK(prefix.initial_score(11) == 1.0); - // A cold batch is atomic: every scoreable undecided lane without a score + // A cold cohort is evaluated before planning: every scoreable undecided lane without a score // is returned for bootstrap, and no scored-but-undecided peer commits until - // the immediate replan. The replan commits every lane exactly once. + // the immediate replan. The immediate replan ranks every lane from cached scores. SpecCostTables crossover = constant_costs(1.0, 4.0, 1.0); SpeculationGate bootstrap(crossover, geometry(), 4); plan = bootstrap.plan(2, { candidate(20, 0, NAN), candidate(21, 1, 4.0)}, 2); CHECK(plan.valid); - CHECK(!plan.decisions_committed); CHECK(plan.admitted_count == 0); CHECK(plan.ordered.empty()); CHECK(plan.unavailable_count == 1); @@ -146,111 +117,90 @@ int main() { CHECK(plan.pending_evaluations[0].slot == 0); CHECK(plan.pending_evaluations[0].action == SpecEvaluationAction::Score); - CHECK(bootstrap.decision(20) == SpecDecision::Undecided); - CHECK(bootstrap.decision(21) == SpecDecision::Undecided); CHECK(bootstrap.initial_score(21) == 4.0); plan = bootstrap.plan(2, { candidate(20, 0, 1.0), candidate(21, 1, NAN)}, 2); CHECK(plan.valid); - CHECK(plan.decisions_committed); CHECK(plan.pending_evaluations.empty()); CHECK(plan.admitted_count == 1); CHECK((plan.admitted_request_ids == std::vector{21})); - CHECK(bootstrap.decision(20) == SpecDecision::AR); - CHECK(bootstrap.decision(21) == SpecDecision::Speculation); CHECK(bootstrap.initial_score(20) == 1.0); CHECK(bootstrap.initial_score(21) == 4.0); plan = bootstrap.plan(2, { candidate(20, 0, 4.0), candidate(21, 1, 1.0)}, 2); - CHECK(plan.decisions_committed); CHECK(plan.pending_evaluations.empty()); CHECK(plan.admitted_count == 1); CHECK((plan.admitted_request_ids == std::vector{21})); - CHECK(plan.ordered.size() == 1); - CHECK(plan.ordered[0].forced); - CHECK(plan.ordered[0].decision == SpecDecision::Speculation); + CHECK(plan.ordered.size() == 2); + CHECK(!plan.ordered[0].forced); + CHECK(plan.ordered[0].source == SpecScoreSource::Initial); CHECK(bootstrap.initial_score(20) == 1.0); CHECK(bootstrap.initial_score(21) == 4.0); - // The gate side of the no-adapter contract emits FallbackAR and commits it - // once; later finite values cannot reopen the request. + // The gate side of the no-adapter contract emits FallbackAR and records it once; later finite values cannot invent a score. // Request-lifetime scoreability is separate from permanent executor // support. Unsupported requests still bootstrap and retain an activation score, // then commit directly to AR. A request that cannot evaluate an activation score - // receives an explicit failed-evaluation action and sticky AR with no + // receives an explicit failed-evaluation action and AR with no // synthetic score. SpeculationGate support(crossover, geometry(), 4); plan = support.plan(1, { candidate(30, 0, NAN, SpeculationPolicy::Adaptive, true, false)}, 1); CHECK(plan.valid); - CHECK(!plan.decisions_committed); CHECK(plan.pending_evaluations.size() == 1); CHECK(plan.pending_evaluations[0].request_id == 30); CHECK(plan.pending_evaluations[0].slot == 0); CHECK(plan.pending_evaluations[0].action == SpecEvaluationAction::Score); - CHECK(support.decision(30) == SpecDecision::Undecided); plan = support.plan(1, { candidate(30, 0, 4.0, SpeculationPolicy::Adaptive, true, false)}, 1); CHECK(plan.valid); - CHECK(plan.decisions_committed); CHECK(plan.admitted_count == 0); CHECK(plan.ordered.size() == 1); - CHECK(plan.ordered[0].decision == SpecDecision::AR); CHECK(plan.ordered[0].execution_unsupported); CHECK(plan.ordered[0].source == SpecScoreSource::Fresh); - CHECK(support.decision(30) == SpecDecision::AR); CHECK(support.initial_score(30) == 4.0); plan = support.plan(1, { candidate(31, 0, NAN, SpeculationPolicy::Adaptive, false, false)}, 1); CHECK(plan.valid); - CHECK(!plan.decisions_committed); CHECK(plan.pending_evaluations.size() == 1); CHECK(plan.pending_evaluations[0].request_id == 31); CHECK(plan.pending_evaluations[0].slot == 0); CHECK(plan.pending_evaluations[0].action == SpecEvaluationAction::FallbackAR); - CHECK(support.decision(31) == SpecDecision::Undecided); - CHECK(support.commit_evaluation_fallback_ar(31)); - CHECK(!support.commit_evaluation_fallback_ar(31)); - CHECK(support.decision(31) == SpecDecision::AR); + CHECK(support.record_evaluation_failure(31)); + CHECK(!support.record_evaluation_failure(31)); CHECK(support.evaluation_failed(31)); CHECK(!support.has_score(31)); CHECK(std::isnan(support.initial_score(31))); plan = support.plan(1, {candidate(31, 0, 4.0)}, 1); CHECK(plan.valid); - CHECK(plan.decisions_committed); CHECK(plan.pending_evaluations.empty()); - CHECK(plan.ordered.empty()); - CHECK(support.decision(31) == SpecDecision::AR); + CHECK(plan.ordered.size() == 1); + CHECK(plan.ordered[0].execution_unsupported); + CHECK(plan.ordered[0].source == SpecScoreSource::Unavailable); // A failed lane and an already-scored healthy lane activate atomically: - // the former becomes sticky AR while the latter still receives its + // the former becomes AR while the latter still receives its // score-based mode on the immediate replan. SpeculationGate mixed_activation(crossover, geometry(), 4); plan = mixed_activation.plan(2, { candidate(32, 0, NAN, SpeculationPolicy::Adaptive, false, false), candidate(33, 1, 4.0)}, 2); CHECK(plan.valid); - CHECK(!plan.decisions_committed); CHECK(plan.pending_evaluations.size() == 1); CHECK(plan.pending_evaluations[0].request_id == 32); CHECK(plan.pending_evaluations[0].action == SpecEvaluationAction::FallbackAR); - CHECK(mixed_activation.decision(32) == SpecDecision::Undecided); - CHECK(mixed_activation.decision(33) == SpecDecision::Undecided); - CHECK(mixed_activation.commit_evaluation_fallback_ar(32)); + CHECK(mixed_activation.record_evaluation_failure(32)); plan = mixed_activation.plan(2, { candidate(32, 0, NAN, SpeculationPolicy::Adaptive, false, false), candidate(33, 1, NAN)}, 2); CHECK(plan.valid); - CHECK(plan.decisions_committed); CHECK(plan.pending_evaluations.empty()); - CHECK(mixed_activation.decision(32) == SpecDecision::AR); CHECK(mixed_activation.evaluation_failed(32)); - CHECK(mixed_activation.decision(33) == SpecDecision::Speculation); CHECK(mixed_activation.initial_score(33) == 4.0); CHECK((plan.admitted_request_ids == std::vector{33})); @@ -267,8 +217,6 @@ int main() { CHECK(plan.admitted_request_ids.front() == 41); CHECK(plan.ordered.front().forced); CHECK(plan.ordered.front().source == SpecScoreSource::Unavailable); - CHECK(policies.decision(41) == SpecDecision::Undecided); - CHECK(policies.decision(42) == SpecDecision::Speculation); plan = policies.plan(2, { candidate(43, 0, 4.0, SpeculationPolicy::Always), candidate(44, 1, 4.0, SpeculationPolicy::Always)}, 1); @@ -287,13 +235,11 @@ int main() { candidate(45, 0, NAN, SpeculationPolicy::Never)}, 1); CHECK(plan.admitted_count == 0); CHECK(plan.ordered.empty()); - CHECK(never.decision(45) == SpecDecision::Undecided); // Capacity, malformed shapes, always-draft pricing, and lookup clamps. SpeculationGate capacity(constant_costs(1.0, 10.0, 1.0), geometry(), 4); plan = capacity.plan(1, {candidate(50, 0, 4.0)}, 0); CHECK(plan.admitted_count == 0); - CHECK(capacity.decision(50) == SpecDecision::AR); plan = capacity.plan(2, {candidate(51, 0, 4.0)}, 1); CHECK(!plan.valid); SpeculationGate always_draft( @@ -336,10 +282,9 @@ int main() { std::vector{0.5, 0.25})); CHECK(plan.ordered[0].expected_yield == 2.0); CHECK(plan.initial_predicted_tokens == 2.0); - CHECK(plan.ordered[0].newly_decided); plan = fitted.plan(1, {candidate(700, 0, 1.0)}, 1); - CHECK(plan.ordered[0].forced); + CHECK(!plan.ordered[0].forced); CHECK(plan.ordered[0].source == SpecScoreSource::Initial); CHECK(plan.ordered[0].expected_yield == 2.0); CHECK(fitted.initial_score(700) == 2.0); @@ -355,9 +300,7 @@ int main() { CHECK(!fitted.has_state(700)); CHECK(!fitted.has_score(700)); CHECK(std::isnan(fitted.initial_score(700))); - CHECK(fitted.decision(700) == SpecDecision::Undecided); plan = fitted.plan(1, {candidate(700, 0, NAN)}, 1); - CHECK(!plan.decisions_committed); CHECK(plan.pending_evaluations.size() == 1); CHECK(plan.pending_evaluations[0].slot == 0); CHECK(plan.pending_evaluations[0].action == @@ -472,30 +415,25 @@ int main() { CHECK(pure_ar_after_bootstrap.cost_scale == 1.0); CHECK(pure_ar_after_bootstrap.predicted_cost == 10.0); - // Shape-local total-cost feedback changes only future undecided choices. - // It cannot flip a request whose one-shot decision is already sticky. + // Shape-local total-cost feedback can change the next cohort epoch's + // route for an existing request. The engine, not request state in the + // gate, keeps the current epoch stable between membership changes. SpeculationGate cost_feedback( constant_costs(1.0, 10.0, 1.0), geometry(), 4); plan = cost_feedback.plan(1, {candidate(950, 0, 4.0)}, 1); CHECK(plan.admitted_count == 1); CHECK(plan.profiled_cost == 12.0); CHECK(plan.cost_scale == 1.0); - CHECK(cost_feedback.decision(950) == SpecDecision::Speculation); cost_feedback.observe_cost({1, 1, 4, 4, 1}, 48.0); plan = cost_feedback.plan(1, {candidate(950, 0, NAN)}, 1); - CHECK(plan.admitted_count == 1); - CHECK(plan.ordered[0].forced); - CHECK(plan.profiled_cost == 12.0); - CHECK(plan.cost_scale == 4.0); - CHECK(plan.predicted_cost == 48.0); - CHECK(cost_feedback.decision(950) == SpecDecision::Speculation); + CHECK(plan.admitted_count == 0); + CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); + CHECK(plan.draft_lanes == 0); cost_feedback.forget(950); - CHECK(cost_feedback.decision(950) == SpecDecision::Undecided); plan = cost_feedback.plan(1, {candidate(951, 0, 4.0)}, 1); CHECK(plan.admitted_count == 0); CHECK(std::abs(plan.goodput - plan.ar_goodput) < 1e-12); - CHECK(cost_feedback.decision(951) == SpecDecision::AR); plan = cost_feedback.plan(2, { candidate(952, 0, 4.0), candidate(953, 1, 4.0)}, 2); CHECK(plan.admitted_count == 2); @@ -505,7 +443,6 @@ int main() { constant_costs(100.0, 10.0, 100.0), geometry(), 4); SpecPlan ar_plan = ar_feedback.plan(1, {candidate(960, 0, 4.0)}, 1); CHECK(ar_plan.admitted_count == 0); - CHECK(ar_feedback.decision(960) == SpecDecision::AR); ar_feedback.observe_cost({1, 0, 0, 1, 0}, 20.0); ar_plan = ar_feedback.plan(1, {candidate(960, 0, NAN)}, 1); CHECK(ar_plan.admitted_count == 0); @@ -521,10 +458,10 @@ int main() { SpeculationGate margin_gate(near_break_even, geometry(), 4); plan = margin_gate.plan(1, {candidate(970, 0, 1.03)}, 1); CHECK(plan.admitted_count == 0); - CHECK(margin_gate.decision(970) == SpecDecision::AR); plan = margin_gate.plan(1, {candidate(970, 0, 4.0)}, 1); CHECK(plan.admitted_count == 0); - CHECK(plan.ordered.empty()); + CHECK(plan.ordered.size() == 1); + CHECK(!plan.ordered[0].admitted); CHECK(margin_gate.initial_score(970) == 1.03); SpecGateConfig zero_margin; @@ -533,7 +470,6 @@ int main() { zero_margin, near_break_even, geometry(), 4); plan = no_margin.plan(1, {candidate(971, 0, 1.03)}, 1); CHECK(plan.admitted_count == 1); - CHECK(no_margin.decision(971) == SpecDecision::Speculation); plan = margin_gate.plan(1, { candidate(972, 0, 1.0, SpeculationPolicy::Always)}, 1); CHECK(plan.admitted_count == 1); @@ -551,8 +487,24 @@ int main() { eight.push_back(candidate(990 + i, i, 4.0)); plan = cannot.plan(8, eight, 8); CHECK(plan.admitted_count == 0); - for (int i = 0; i < 8; ++i) - CHECK(cannot.decision(990 + i) == SpecDecision::AR); + + // Epoch identity follows the live decode request set, not slot + // occupancy or changing scores. Refill in the same slot is a new epoch. + const std::vector cohort = { + candidate(1000, 3, 4.0), candidate(1001, 7, 2.0)}; + SpecCohortEpoch epoch; + epoch.id = 7; + epoch.request_ids = SpecCohortEpoch::ids(cohort); + epoch.plan = plan; + CHECK(epoch.matches(cohort)); + CHECK(epoch.matches({ + candidate(1000, 9, NAN), candidate(1001, 2, NAN)})); + CHECK(!epoch.matches({candidate(1000, 3, NAN)})); + CHECK(epoch.matches({ + candidate(1001, 7, NAN), candidate(1000, 3, NAN)})); + CHECK(!epoch.matches({ + candidate(1000, 3, NAN), candidate(1002, 7, NAN)})); + CHECK((epoch.request_ids == std::vector{1000, 1001})); std::printf("speculation gate tests passed: %d checks\n", g_checks); return 0;