From fc50b81d9c539d36acda053d235d996ed368dbfc Mon Sep 17 00:00:00 2001 From: Graffioh Date: Fri, 21 Aug 2026 15:16:28 +0000 Subject: [PATCH 01/11] concurrency: add ordered burst and staged slot commits --- server/src/common/concurrency/seq_engine.h | 41 +++++++++++-- .../concurrency/qwen35_slot_manager.cpp | 61 ++++++++++++++----- .../qwen35/concurrency/qwen35_slot_manager.h | 14 ++++- server/src/server/scheduler.cpp | 9 ++- server/test/test_seq_engine_contract.cpp | 17 ++++++ server/test/test_seq_slot_manager.cpp | 51 ++++++++++++++-- server/test/test_server_unit.cpp | 16 +++++ 7 files changed, 182 insertions(+), 27 deletions(-) diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index fa4dbba35..0e22269b4 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -181,14 +181,20 @@ class SeqEngine { struct StepInput { int slot = -1; int32_t token = -1; // token to commit at this slot's next position + // Scheduler-side hooks may replace the next token before it is fed + // back. A speculative burst is unsafe while that authority is live. + 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. Durable speculative children precede it. + 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; }; struct PrefillOutput { @@ -250,6 +256,16 @@ class SeqEngine { virtual bool token_is_eos(int32_t token) const = 0; }; +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 +281,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 +289,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 +309,22 @@ 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.token >= 0 || !output.committed_tokens.empty() || + output.error.empty()) + return "failed decode has invalid payload"; + } else { + if (output.token < 0 || !output.error.empty()) + return "successful decode has invalid payload"; + 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/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index b8dcc36ec..74c74ee1e 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -230,37 +230,70 @@ 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) { + 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)) { + 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); out.ok = true; - out.physical_row = (int64_t)app.last.physical_token_index; + out.count = n_tokens; + out.physical_row = out.physical_rows.front(); 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 (!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(); } void Qwen35SlotManager::retire(int slot) { diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index 1da009f69..78c2b4b00 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -45,6 +45,9 @@ 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 reserved by the current target graph. They become durable + // only after the graph and any speculative promotion succeed. + std::vector staged_tokens; int generated_tokens() const { return sample_history.size() > (size_t)prompt_len @@ -94,13 +97,20 @@ class Qwen35SlotManager { bool ok = false; bool busy = false; // no physical block available right now int64_t physical_row = -1; + std::vector physical_rows; + int count = 0; int position = -1; // logical position the fed token is written at int32_t new_block = -1; int new_block_index = -1; + std::vector new_blocks; + int first_new_block = -1; }; - // Allocate the next decode token's cache row, report any new block-table - // entry, and log it to sample_history. cur_pos waits for commit_step(). + StepAppend append_tokens(int slot, const int32_t * fed_tokens, + int n_tokens); + + // Allocate decode cache rows and stage the fed tokens. Both history and + // cur_pos wait for commit_step(). StepAppend append_token(int slot, int32_t fed_token); // The batched step's compute succeeded: cur_pos++. diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 0575d4c2e..4ee87f349 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -618,7 +618,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { 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}); + {i, slots[(size_t)i].pending_tok, + slots[(size_t)i].hook.close_token_ids.empty()}); } else if (slots[(size_t)i].job) { prefill_candidates.push_back( {i, slots[(size_t)i].admission_order}); @@ -666,7 +667,11 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.finished = true; continue; } - advance_slot(s, out.token); + consume_decode_output_tokens(out, [&](int32_t token) { + if (s.finished) return false; + advance_slot(s, token); + return !s.finished; + }); } using PrefillStatus = SeqEngine::PrefillOutput::Status; for (const auto & out : step_result.prefills) { diff --git a/server/test/test_seq_engine_contract.cpp b/server/test/test_seq_engine_contract.cpp index be3f031b4..56f270f41 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; @@ -318,6 +323,18 @@ int main() { CHECK(mentions(violations, test.expected)); } + { + SeqEngine::StepPlan plan; + plan.decode.push_back({0, 7, false}); + SeqEngine::StepResult result; + SeqEngine::DecodeOutput output; + output.slot = 0; + output.token = 8; + output.committed_tokens = {9}; + result.decode.push_back(output); + CHECK(!validate_step_result(plan, result, 1).empty()); + } + std::printf("test_seq_engine_contract: %d checks passed\n", g_checks); return 0; } diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index a64329e5f..5539c3d45 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -85,18 +85,21 @@ 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 rows stage until 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 +336,46 @@ int main() { CHECK(pool.free_block_count() == 0); } + // A multi-token stage owns its rows until one atomic commit. + { + PagedKvPool pool(8, 1, /*block_size=*/4); + Qwen35SlotManager mgr(pool, /*max_ctx=*/32); + 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.count == 6); + CHECK(staged.position == 3 && staged.physical_rows.size() == 6); + CHECK(staged.first_new_block == 1 && 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); + + 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 694d857bb..c7907f4c9 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" @@ -5943,5 +5944,20 @@ TEST_CASE(ServerUnitFixture, test_emitter_function_calls_param_with_literal_thin TEST_ASSERT(em.emit_token_count() - em.first_content_token_index() == 1); } +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; + const bool consumed_all = consume_decode_output_tokens( + burst, [&](int32_t token) { + emitted.push_back(token); + return token != 2; + }); + TEST_ASSERT(!consumed_all); + TEST_ASSERT((emitted == std::vector{101, 2})); +} From aec43dc9efc600dc8189cb72691932c126d611f6 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Fri, 21 Aug 2026 16:33:44 +0000 Subject: [PATCH 02/11] concurrency: add fixed Qwen DFlash2 chain path --- server/CMakeLists.txt | 40 + .../deps/llama.cpp/ggml/include/ggml-cuda.h | 24 + server/deps/llama.cpp/ggml/include/ggml.h | 50 +- .../llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp | 4 +- .../ggml/src/ggml-cuda/gated_delta_net.cu | 104 ++- .../src/ggml-cuda/gdn-transition-journal.cu | 402 +++++++++ .../ggml/src/ggml-cuda/paged-attn.cu | 201 ++++- .../ggml/src/ggml-metal/ggml-metal-device.m | 3 +- .../ggml/src/ggml-sycl/ggml-sycl.cpp | 3 +- .../ggml/src/ggml-vulkan/ggml-vulkan.cpp | 3 +- server/deps/llama.cpp/ggml/src/ggml.c | 83 +- .../common/concurrency/chain_spec_shapes.h | 103 +++ server/src/common/ddtree.cpp | 10 + server/src/common/ddtree.h | 10 + server/src/common/dflash2_batch.cpp | 376 ++++++++ server/src/common/dflash2_head.cpp | 16 + server/src/common/dflash2_head.h | 13 + .../src/common/dflash2_selector_validation.h | 93 ++ server/src/common/dflash_draft_kv.cpp | 140 +++ server/src/common/dflash_draft_kv.h | 29 + server/src/common/feature_gate.cpp | 30 +- .../src/common/geometric_draft_topk_cuda.cu | 13 +- server/src/common/geometric_draft_topk_cuda.h | 4 + server/src/common/gpu_runtime_compat.h | 1 + server/src/common/step_graph.h | 26 + server/src/draft/draft_gguf_loader.cpp | 20 + server/src/internal.h | 48 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 835 +++++++++++++++++- .../qwen35/concurrency/qwen35_seq_engine.h | 57 +- server/src/qwen35/graph_builders.cpp | 253 +++++- server/src/qwen35/graph_builders.h | 79 ++ server/src/qwen35/qwen35_backend.cpp | 62 +- server/src/qwen35/qwen35_target_graph.cpp | 506 +++++++---- server/test/test_chain_spec_shapes.cpp | 89 ++ server/test/test_ddtree_path.cpp | 44 + .../test/test_dflash2_selector_validation.cpp | 86 ++ server/test/test_draft_topk_cuda.cpp | 52 +- server/test/test_feature_gate.cpp | 31 +- server/test/test_gdn_transition_journal.cpp | 789 +++++++++++++++++ server/test/test_paged_attention.cpp | 228 ++++- server/test/test_recurrent_snapshot.cpp | 142 +++ server/test/test_seq_slot_manager.cpp | 2 - 42 files changed, 4790 insertions(+), 314 deletions(-) create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu create mode 100644 server/src/common/concurrency/chain_spec_shapes.h create mode 100644 server/src/common/dflash2_batch.cpp create mode 100644 server/src/common/dflash2_selector_validation.h create mode 100644 server/test/test_chain_spec_shapes.cpp create mode 100644 server/test/test_ddtree_path.cpp create mode 100644 server/test/test_dflash2_selector_validation.cpp create mode 100644 server/test/test_gdn_transition_journal.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index cf1e4ee8b..85dddf537 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -455,6 +455,7 @@ add_library(dflash_common STATIC src/common/domino_head.cpp src/common/dspark_head.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 @@ -1435,6 +1436,34 @@ 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_chain_spec_shapes.cpp") + # Pure host-side DFlash2 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_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_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) @@ -1935,6 +1964,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 0662f0f1b..1281cd061 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -99,6 +99,30 @@ 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); +// 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/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index a4db13a1f..ef3806dc3 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 @@ -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,32 @@ 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 +#ifdef __cplusplus + = nullptr +#endif + , + struct ggml_tensor * tree_sizes +#ifdef __cplusplus + = nullptr +#endif + , + int tree_width +#ifdef __cplusplus + = 0 +#endif + , + int tree_scratch_base +#ifdef __cplusplus + = 0 +#endif + , + int tree_scratch_stride +#ifdef __cplusplus + = 0 +#endif + ); // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse. // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse). @@ -2925,6 +2962,15 @@ extern "C" { struct ggml_tensor * tensor, bool skip_intermediate); + // CUDA/HIP fixed-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. A tree-form op may attach the journal when + // its parent table describes one root-inclusive chain. + 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-cpu/ggml-cpu.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp index 5d6e93a68..00d05b3ee 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -480,7 +480,9 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st case GGML_OP_GATED_DELTA_NET: // The Specla GDN variant (op param 2 == 1) is stateful via HLD and is // only supported by the CUDA kernel. - return ggml_get_op_params_i32(op, 2) != 1; + return ggml_get_op_params_i32(op, 2) != 1 && + ggml_get_op_params_i32(op, 10) != 1 && + op->src[11] == nullptr; case GGML_OP_OUT_PROD: return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; 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 64156d5bc..a43cf10fb 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] @@ -313,6 +341,7 @@ gated_delta_net_cuda_grouped_cols(const float * q, float * state_out, const int * parent_ids, // TREE_MODE only; else ignored InterT * persist_inter, + float * transition_journal, int64_t H, int64_t n_tokens, int64_t n_seqs, @@ -452,6 +481,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]; @@ -468,12 +503,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 @@ -482,6 +525,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); } @@ -557,7 +603,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; @@ -579,19 +626,19 @@ 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; @@ -610,7 +657,7 @@ static void launch_gated_delta_net( 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, 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 if (warp_size == 64) { @@ -618,24 +665,24 @@ static void launch_gated_delta_net( 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, 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); } } 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); } @@ -884,6 +931,9 @@ 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]. The packed tree caller + // uses it only with a root-inclusive linear parent chain. + 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); @@ -926,6 +976,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) { @@ -954,6 +1007,15 @@ 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_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); @@ -995,35 +1057,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, gate_bias_d, gate_A_d); \ + 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, 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); \ + 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, 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); \ + 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, 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); \ + 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, 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); \ + 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, 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); \ + 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, 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..d9a81a68d --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/gdn-transition-journal.cu @@ -0,0 +1,402 @@ +#include "common.cuh" +#include "ggml-cuda.h" + +#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 + +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 == -1) continue; + if (slot < 0 || slot >= common_state_slots) return false; + 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 == -1) continue; + if (slot < 0) return false; + 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 c76d97b36..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 @@ -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,10 @@ 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_row_offset, + int32_t tree_scratch_base, + int32_t tree_scratch_stride, float scale) { constexpr int nthreads = WARP_SIZE; constexpr int values_per_load = 4; @@ -222,29 +265,40 @@ 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 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. 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 || 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; 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. - if (query_positions && query_pos < kv_seq_len_raw) { + // 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 && !tree_query && query_pos < kv_seq_len_raw) { kv_seq_len_raw = query_pos + 1; } const int64_t table_capacity = @@ -254,8 +308,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_query ? 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 +349,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 +423,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 +445,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 +455,19 @@ static __global__ void paged_attn_decode( phys_mine = physical_block * block_size + my_token % block_size; } + } 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, + 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 +697,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 +708,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)) { + return false; + } if (dst->type != GGML_TYPE_F32 || q->type != GGML_TYPE_F32 || @@ -636,7 +721,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 +734,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 +790,49 @@ 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 || + (!query_positions && + 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 +892,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 +971,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 +1001,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 +1054,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 +1066,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 +1077,12 @@ static bool try_launch_paged_attn( (int32_t) block_table->ne[0], 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); if (n_partitions > 1) { diff --git a/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m b/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m index 330daec28..b8d710771 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m +++ b/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m @@ -1186,7 +1186,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te return true; case GGML_OP_GATED_DELTA_NET: return has_simdgroup_reduction && op->src[2]->ne[0] % 32 == 0 && - op->src[8] == NULL; + op->src[8] == NULL && op->src[11] == NULL && + ggml_get_op_params_i32(op, 10) != 1; case GGML_OP_SOLVE_TRI: case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: diff --git a/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp b/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp index 463619fd5..7850baa62 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -4957,7 +4957,8 @@ static bool ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_GATED_LINEAR_ATTN: return true; case GGML_OP_GATED_DELTA_NET: - return op->src[8] == nullptr; + return op->src[8] == nullptr && op->src[11] == nullptr && + ggml_get_op_params_i32(op, 10) != 1; case GGML_OP_SSM_CONV: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && diff --git a/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 66c69692e..17db9166e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -15782,7 +15782,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm { // The Vulkan kernel addresses state by compact sequence row // and does not consume the physical-slot mapping in src[8]. - if (op->src[8] != nullptr) { + if (op->src[8] != nullptr || op->src[11] != nullptr || + ggml_get_op_params_i32(op, 10) == 1) { return false; } const uint32_t S_v = op->src[2]->ne[0]; diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index d6c889dad..87d86c6e2 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,16 @@ 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); + // 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); + 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 +5764,50 @@ 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); + 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 + // 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]); + } 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 +5817,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; } @@ -6811,6 +6865,31 @@ 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(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). [dt_bias | A] -> src[9], // op_params[10] = 1. (src[8] / op_params[2] belong to the compact-decode and // SpecLA variants.) 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..b21c17daa --- /dev/null +++ b/server/src/common/concurrency/chain_spec_shapes.h @@ -0,0 +1,103 @@ +#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_chain_verify_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 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.commit_rows = shape.accepted_rows + shape.ar_lanes; + return shape; +} + +// 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/ddtree.cpp b/server/src/common/ddtree.cpp index 9c083996c..8e633fadd 100644 --- a/server/src/common/ddtree.cpp +++ b/server/src/common/ddtree.cpp @@ -376,4 +376,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 026dc4ca4..f4f93c8de 100644 --- a/server/src/common/ddtree.h +++ b/server/src/common/ddtree.h @@ -8,6 +8,7 @@ #pragma once +#include #include #include #include @@ -88,4 +89,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/dflash2_batch.cpp b/server/src/common/dflash2_batch.cpp new file mode 100644 index 000000000..7a6d6b500 --- /dev/null +++ b/server/src/common/dflash2_batch.cpp @@ -0,0 +1,376 @@ +#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 + +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; +} + +} // 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) { + draft_tokens.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 + if (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)); + for (int lane = 0; lane < n_lanes; ++lane) { + draft_tokens[(size_t) lane][0] = last_tokens[(size_t) lane]; + 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; + 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; + 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]; + predecessor_row = n_lanes + selected_row; + } + } + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash2_head.cpp b/server/src/common/dflash2_head.cpp index 2cbcc2eb4..83273886a 100644 --- a/server/src/common/dflash2_head.cpp +++ b/server/src/common/dflash2_head.cpp @@ -1,5 +1,6 @@ #include "dflash2_head.h" +#include "dflash2_selector_validation.h" #include "ggml-alloc.h" #include @@ -62,6 +63,21 @@ bool dflash2_score_candidates(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_score_candidates: %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 diff --git a/server/src/common/dflash2_head.h b/server/src/common/dflash2_head.h index f054ac4fa..57048cd6c 100644 --- a/server/src/common/dflash2_head.h +++ b/server/src/common/dflash2_head.h @@ -26,6 +26,19 @@ bool dflash2_select_chain(const DraftWeights & dw, int32_t last_tok, std::vector & draft_tok); +// 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. +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); + // Selector-scored candidates for DDTree construction (DARTree-style): the // same per-position top-k + selector projections as the chain path, kept on // the host so the tree builder can ask for branch-conditioned scores. 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 05ec25aca..a9ac623b5 100644 --- a/server/src/common/dflash_draft_kv.cpp +++ b/server/src/common/dflash_draft_kv.cpp @@ -318,4 +318,144 @@ 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.hidden_by_lane.clear(); + batch.lane_states.clear(); + batch.meta_arena.clear(); + batch.n_lanes = 0; + batch.q_len = 0; + batch.built_for = nullptr; +} + +static bool draft_kv_batch_build( + DraftKvBatchGraph & batch, + const DraftWeights & dw, + ggml_backend_t backend, + const std::vector & lane_states) { + 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 != static_cast(&dw)) { + return false; + } + } + + draft_kv_batch_free(batch); + const int n_lanes = static_cast(lane_states.size()); + const size_t arena_size = + (32u + 16u * static_cast(n_lanes)) * 1024u * 1024u; + batch.meta_arena.resize(arena_size); + 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.hidden_by_lane.reserve(static_cast(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 output = build_draft_kv_step( + batch.g_ctx, batch.gf, dw, state->cache, step); + if (!output.hidden_states) { + draft_kv_batch_free(batch); + return false; + } + ggml_set_output(output.hidden_states); + ggml_build_forward_expand(batch.gf, output.hidden_states); + batch.hidden_by_lane.push_back(output.hidden_states); + } + + 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.built_for = &dw; + batch.lane_states = lane_states; + std::fprintf(stderr, + "[draft-kv-batch] packed backbone ready lanes=%d q_len=%d\n", + n_lanes, dw.block_size); + return true; +} + +bool draft_kv_batch_compute( + DraftKvBatchGraph & batch, + const DraftWeights & dw, + ggml_backend_t backend, + const std::vector & lane_states, + std::vector> & hidden_by_lane) { + hidden_by_lane.clear(); + if (lane_states.empty()) return false; + + const bool reusable = + batch.gf && batch.built_for == static_cast(&dw) && + batch.lane_states == lane_states; + if (!reusable && + !draft_kv_batch_build( + batch, dw, backend, lane_states)) { + return false; + } + 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 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)); + for (int lane = 0; lane < batch.n_lanes; ++lane) { + ggml_backend_tensor_get_async( + backend, batch.hidden_by_lane[static_cast(lane)], + hidden_by_lane[static_cast(lane)].data(), 0, + sizeof(float) * elements); + } + ggml_backend_synchronize(backend); + return true; +} + } // namespace dflash::common diff --git a/server/src/common/dflash_draft_kv.h b/server/src/common/dflash_draft_kv.h index 888f9b46b..6ef15e766 100644 --- a/server/src/common/dflash_draft_kv.h +++ b/server/src/common/dflash_draft_kv.h @@ -99,4 +99,33 @@ 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; + const void * built_for = nullptr; + std::vector lane_states; + + std::vector meta_arena; + ggml_context * g_ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + std::vector hidden_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 packed graph computes the shared backbone and +// returns one host-visible hidden-state block per lane. +bool draft_kv_batch_compute( + DraftKvBatchGraph & batch, + const DraftWeights & dw, + ggml_backend_t backend, + const std::vector & lane_states, + std::vector> & hidden_by_lane); + } // namespace dflash::common diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index df7b03f90..bfc7a80b3 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -174,6 +174,17 @@ std::string check_feature_compatibility( } } + const bool concurrent_local_chain = + arch == "qwen35" && args.paged_attention && + args.max_concurrency > 1 && args.draft_path != nullptr && + !args.ddtree_mode && !args.remote_draft.enabled() && + !args.device.is_layer_split() && + !args.device.is_tensor_parallel() && + !args.remote_target_shard.enabled() && + target_backend == draft_backend && + args.device.gpu == args.draft_device.gpu && + args.fa_window == 0; + // ── --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 @@ -191,10 +202,13 @@ 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) { + if ((args.draft_path != nullptr || args.remote_draft.enabled()) && + !concurrent_local_chain) { return "--paged-attention requires autoregressive decode without a " - "draft or DDTree"; + "draft, or concurrent local same-device DFlash2 chains"; + } + if (args.ddtree_mode) { + return "--paged-attention does not support DDTree"; } if (args.fa_window != 0) { return "--paged-attention requires full attention (--fa-window 0)"; @@ -243,10 +257,12 @@ 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(); + const int64_t chain_scratch = concurrent_local_chain + ? (int64_t)args.max_concurrency * paged_token_capacity(16) + : 0; + const int64_t max_pool_tokens = + ((int64_t)INT32_MAX - PAGED_BLOCK_SIZE - chain_scratch) / + 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/geometric_draft_topk_cuda.cu b/server/src/common/geometric_draft_topk_cuda.cu index ba287656d..e6cfaeb93 100644 --- a/server/src/common/geometric_draft_topk_cuda.cu +++ b/server/src/common/geometric_draft_topk_cuda.cu @@ -331,7 +331,10 @@ 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; + 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 +365,7 @@ 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); - // 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. + bool dispatched = false; #define DFLASH_TOPK_LAUNCH(KV, VEC) \ geometric_draft_topk_partial<<>>( \ lp_in, vocab, inv_t, split, \ @@ -374,6 +375,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 +389,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..774c822c2 100644 --- a/server/src/common/geometric_draft_topk_cuda.h +++ b/server/src/common/geometric_draft_topk_cuda.h @@ -25,6 +25,10 @@ namespace dflash::common { +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/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 900a7305e..affcbace7 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 @@ -36,6 +38,7 @@ struct StepGraph { 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 * tree_sizes = nullptr; // DDTree [n_tree_seqs], 0 = padding // SpecLA topology masks ([n_tokens, n_tokens] f32, host-filled; see // delta_net_specla.h). Created only when DFLASH_SPECLA capture is active. ggml_tensor * specla_m_strict = nullptr; @@ -61,11 +64,21 @@ 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; + // 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 +96,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; @@ -98,6 +117,7 @@ inline void step_graph_free(StepGraph & sg) { sg.built_view = false; sg.hidden_input = nullptr; sg.parent_ids = nullptr; + sg.tree_sizes = nullptr; sg.specla_m_strict = sg.specla_m_incl = sg.specla_m_eye = nullptr; sg.specla_hld = nullptr; sg.kv_write_rows = nullptr; @@ -105,6 +125,11 @@ inline void step_graph_free(StepGraph & sg) { sg.state_slot_ids = nullptr; 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 +141,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/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index 58c203195..6d1a3c2cd 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" @@ -549,6 +550,25 @@ bool load_draft_gguf(const std::string & path, 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/internal.h b/server/src/internal.h index f65fe7c58..07ab88877 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -498,13 +498,12 @@ struct TargetCache { ggml_tensor * specla_factor_ptrs = nullptr; // 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 concurrent + // tree 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; @@ -639,10 +638,11 @@ bool restore_target_cache_chain(const PrefixSnapshot * thick, // `n_seq_slots` (concurrent serving): number of sequence slots the cache // serves at once. > 1 requires paged_attention; it adds a trailing slot axis // to the recurrent state, widens the paged metadata to one block-table column -// per slot, and skips the spec-decode rollback tensors entirely (concurrent -// decode is AR-only). With +// per slot, and skips the legacy rollback tensors entirely. 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 owns fixed tree +// scratch and disjoint target-feature rings for GPU promotion. bool create_target_cache(const TargetWeights & w, int max_ctx, int max_verify_tokens, @@ -651,7 +651,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 @@ -668,7 +669,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); @@ -729,6 +731,10 @@ bool specla_commit_accepted(TargetCache & cache, 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; // SpecLA factor capture (DFLASH_SPECLA=1, docs/SPECLA.md). Persistent F32 // aliases into the bank written by this verify. In the HLD path the @@ -770,10 +776,12 @@ 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 - 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; @@ -799,6 +807,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. @@ -830,9 +842,19 @@ 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; + // 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. @@ -859,6 +881,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 0a54761f9..105fa93d7 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -11,8 +11,12 @@ #include "graph_builders.h" #include "attn_masks.h" #include "prefill_helpers.h" +#include "common/concurrency/chain_spec_shapes.h" +#include "common/ddtree.h" +#include "common/dflash2_head.h" #include "common/sampler.h" #include "internal.h" +#include "ggml-cuda.h" #include #include @@ -35,6 +39,300 @@ 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), 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(static_cast(n_slots)); + prepared_chain_drafts_.resize(static_cast(n_slots)); + + capture_features_ = tree_width_ > 1 && tree_width_ <= 16 && + tree_width_ == b_.dw_.block_size && b_.dw_.selector.enabled && + b_.cache_.target_feat && b_.cache_.target_feat_cap > 0 && + b_.cfg_.paged_attention && b_.cfg_.fa_window == 0 && + b_.cfg_.draft_path && !b_.cfg_.remote_draft.enabled() && + !b_.split_gpus_ && !b_.cfg_.device.is_tensor_parallel() && + b_.target_backend_ == b_.draft_backend_ && + b_.cfg_.draft_gpu == b_.cfg_.device.gpu && b_.w_.output; + if (!capture_features_) return; + + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * + static_cast(n_slots + 1); + params.no_alloc = true; + feature_view_ctx_ = ggml_init(params); + if (!feature_view_ctx_) { + capture_features_ = false; + return; + } + + const int cap = b_.cache_.target_feat_cap; + const int64_t feature_width = + static_cast(b_.w_.n_capture_layers) * b_.w_.n_embd; + slot_feature_mirrors_.resize(static_cast(n_slots)); + for (int slot = 0; slot < n_slots; ++slot) { + DraftFeatureMirror & mirror = + slot_feature_mirrors_[static_cast(slot)]; + mirror.target_feat = ggml_view_2d( + feature_view_ctx_, b_.cache_.target_feat, feature_width, cap, + b_.cache_.target_feat->nb[1], + static_cast(slot) * static_cast(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() { + draft_kv_batch_free(batch_draft_graph_); + for (std::unique_ptr & state : dummy_draft_kv_) { + if (state) draft_kv_free(*state); + } + for (std::unique_ptr & state : slot_draft_kv_) { + if (state) draft_kv_free(*state); + } + for (DraftFeatureMirror & mirror : slot_feature_mirrors_) { + draft_feature_mirror_free(mirror); + } + 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 >= static_cast(slot_feature_mirrors_.size())) { + return nullptr; + } + return &slot_feature_mirrors_[static_cast(slot)]; +} + +DraftKvState * Qwen35SeqEngine::ensure_slot_draft_kv(int slot) { + DraftFeatureMirror * mirror = slot_feature_mirror(slot); + if (!mirror || slot < 0 || + slot >= static_cast(slot_draft_kv_.size())) { + return nullptr; + } + std::unique_ptr & state = + slot_draft_kv_[static_cast(slot)]; + if (state && state->gf && + state->built_for == static_cast(&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::chain_spec_input_capable( + const StepInput & input) const { + if (!capture_features_ || !input.allow_speculation || + input.slot < 0 || input.slot >= slots_.slot_count()) { + return false; + } + const Qwen35Slot & slot = slots_.slot(input.slot); + return slot.decoding() && !slot.sampler.needs_logit_processing() && + slot.cur_pos >= 1 && + slot.cur_pos + tree_width_ <= slots_.max_context(); +} + +Qwen35SeqEngine::FixedServiceRound +Qwen35SeqEngine::make_fixed_service_round(const StepPlan & plan) const { + if (!plan.prefills.empty()) return PromptServiceRound{}; + SpeculativeDecodeRound round; + round.chain_lanes.resize(plan.decode.size(), 0); + for (size_t i = 0; i < plan.decode.size(); ++i) { + round.chain_lanes[i] = + chain_spec_input_capable(plan.decode[i]) ? 1 : 0; + } + return round; +} + +bool Qwen35SeqEngine::prepare_chain_drafts( + const std::vector & inputs, + const std::vector & selected) { + if (selected.size() != inputs.size() || !capture_features_ || + tree_width_ <= 1 || tree_width_ != b_.dw_.block_size) { + return false; + } + + struct Lane { + int slot = -1; + int32_t root = -1; + DraftKvState * state = nullptr; + DraftFeatureMirror * mirror = nullptr; + }; + std::vector lanes; + lanes.reserve(inputs.size()); + const int hidden = b_.w_.n_embd; + std::vector noise( + static_cast(tree_width_), b_.w_.mask_token_id); + std::vector noise_embed( + static_cast(hidden) * tree_width_); + + auto reset_lanes = [&]() { + for (const Lane & lane : lanes) { + if (lane.state) draft_kv_reset(*lane.state); + if (lane.slot >= 0 && + lane.slot < static_cast(prepared_chain_drafts_.size())) { + prepared_chain_drafts_[static_cast(lane.slot)] = {}; + } + } + for (const std::unique_ptr & state : dummy_draft_kv_) { + if (state) draft_kv_reset(*state); + } + }; + + for (size_t i = 0; i < inputs.size(); ++i) { + if (!selected[i]) continue; + const StepInput & input = inputs[i]; + if (!chain_spec_input_capable(input)) { + reset_lanes(); + return false; + } + prepared_chain_drafts_[static_cast(input.slot)] = {}; + DraftKvState * state = ensure_slot_draft_kv(input.slot); + DraftFeatureMirror * mirror = slot_feature_mirror(input.slot); + if (!state || !mirror) { + if (state) draft_kv_reset(*state); + reset_lanes(); + return false; + } + lanes.push_back({input.slot, input.token, state, mirror}); + if (!draft_kv_begin_step( + *state, b_.dw_, b_.draft_backend_, *mirror, + slots_.slot(input.slot).cur_pos)) { + reset_lanes(); + return false; + } + noise[0] = input.token; + std::fill( + noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), tree_width_, noise_embed.data())) { + reset_lanes(); + return false; + } + ggml_backend_tensor_set( + state->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + } + if (lanes.empty()) return true; + + const int bucket = chain_decode_bucket_width( + static_cast(lanes.size())); + std::vector batch_states; + std::vector roots; + batch_states.reserve(static_cast(bucket)); + roots.reserve(static_cast(bucket)); + for (const Lane & lane : lanes) { + batch_states.push_back(lane.state); + roots.push_back(lane.root); + } + + const int dummy_count = bucket - static_cast(lanes.size()); + const int cap = std::min( + lanes.front().mirror->cap, + std::max(1, b_.cfg_.draft_ctx_max)); + while (static_cast(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); + reset_lanes(); + return false; + } + dummy_draft_kv_.push_back(std::move(dummy)); + } + noise[0] = lanes.front().root; + std::fill(noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), tree_width_, noise_embed.data())) { + reset_lanes(); + return false; + } + for (int i = 0; i < dummy_count; ++i) { + DraftKvState * dummy = + dummy_draft_kv_[static_cast(i)].get(); + if (!draft_kv_begin_step( + *dummy, b_.dw_, b_.draft_backend_, + *lanes.front().mirror, 1)) { + reset_lanes(); + return false; + } + ggml_backend_tensor_set( + dummy->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + batch_states.push_back(dummy); + roots.push_back(lanes.front().root); + } + + std::vector> hidden_blocks; + std::vector> proposals; + if (!draft_kv_batch_compute( + batch_draft_graph_, b_.dw_, b_.draft_backend_, + batch_states, hidden_blocks) || + hidden_blocks.size() != batch_states.size()) { + reset_lanes(); + return false; + } + std::vector hidden_by_lane; + hidden_by_lane.reserve(hidden_blocks.size()); + for (const std::vector & block : hidden_blocks) { + hidden_by_lane.push_back(block.data()); + } + if (!dflash2_select_chains_batched( + b_.dw_, b_.draft_backend_, b_.w_.output, hidden_by_lane, + tree_width_, roots, proposals) || + proposals.size() != batch_states.size()) { + reset_lanes(); + return false; + } + + for (size_t lane_index = 0; lane_index < lanes.size(); ++lane_index) { + const Lane & lane = lanes[lane_index]; + std::vector & tokens = proposals[lane_index]; + if (tokens.size() != static_cast(tree_width_) || + tokens.front() != lane.root) { + reset_lanes(); + return false; + } + PreparedChainDraft & prepared = + prepared_chain_drafts_[static_cast(lane.slot)]; + prepared.valid = true; + prepared.generated = slots_.slot(lane.slot).generated_tokens(); + prepared.root = lane.root; + prepared.tokens = std::move(tokens); + } + return true; +} + bool Qwen35SeqEngine::token_is_eos(int32_t token) const { return b_.token_is_eos(token); } @@ -46,6 +344,15 @@ 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 (result.slot >= 0 && + result.slot < static_cast(slot_draft_kv_.size()) && + slot_draft_kv_[static_cast(result.slot)]) { + draft_kv_reset(*slot_draft_kv_[static_cast(result.slot)]); + } + if (result.slot >= 0 && result.slot < + static_cast(prepared_chain_drafts_.size())) { + prepared_chain_drafts_[static_cast(result.slot)] = {}; + } } return result; } @@ -157,6 +464,482 @@ Qwen35SeqEngine::PrefillStage Qwen35SeqEngine::stage_prefill_chunk( return stage; } +SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( + const StepPlan & plan, const std::vector & selected) { + StepResult result; + const std::vector & inputs = plan.decode; + if (!plan.prefills.empty() || selected.size() != inputs.size()) { + result.error = "invalid fixed chain round"; + return result; + } + + struct Proposal { + size_t input_index = 0; + int slot = -1; + int32_t root = -1; + DDTree tree; + std::vector tokens; + std::vector accepted; + std::vector path; + 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; + }; + + const int tree_width = 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(); + const int min_tokens = []() { + const char * value = std::getenv("DFLASH_MIN_TOKENS"); + return value ? std::max(0, std::atoi(value)) : 0; + }(); + + std::vector proposals; + std::vector ar_lanes; + std::vector proposal_for_input(inputs.size(), -1); + std::vector ar_for_input(inputs.size(), -1); + proposals.reserve(inputs.size()); + ar_lanes.reserve(inputs.size()); + + for (size_t i = 0; i < inputs.size(); ++i) { + const StepInput & input = inputs[i]; + if (selected[i]) { + PreparedChainDraft & prepared = + prepared_chain_drafts_[static_cast(input.slot)]; + if (!prepared.valid || prepared.root != input.token || + prepared.generated != + slots_.slot(input.slot).generated_tokens() || + prepared.tokens.size() != + static_cast(tree_width)) { + result.error = "prepared DFlash2 chain became stale"; + return result; + } + Proposal proposal; + proposal.input_index = i; + proposal.slot = input.slot; + proposal.root = input.token; + proposal.tokens = std::move(prepared.tokens); + prepared = {}; + proposal.tree = make_chain_verify_tree(proposal.tokens); + if (proposal.tree.n_nodes + 1 != tree_width) { + result.error = "prepared DFlash2 chain has invalid shape"; + return result; + } + proposal_for_input[i] = static_cast(proposals.size()); + proposals.push_back(std::move(proposal)); + } else { + ArLane lane; + lane.input_index = i; + lane.slot = input.slot; + lane.token = input.token; + ar_for_input[i] = static_cast(ar_lanes.size()); + ar_lanes.push_back(lane); + } + } + if (proposals.empty()) { + result.error = "fixed chain round has no speculative lanes"; + return result; + } + + // AR peers write their ordinary rows in the same target forward. Host + // history and positions remain staged until every promotion succeeds. + for (ArLane & lane : ar_lanes) { + const Qwen35SlotManager::StepAppend append = + slots_.append_token(lane.slot, lane.token); + if (!append.ok) { + result.error = append.busy + ? "paged KV pool exhausted during compact AR staging" + : "compact AR K/V staging failed"; + return result; + } + if (append.new_block >= 0 && + !upload_block_table_delta( + lane.slot, append.new_block_index, + &append.new_block, 1)) { + result.error = "compact AR block-table update failed"; + return result; + } + lane.position = append.position; + lane.physical_row = append.physical_row; + } + + const int spec_count = static_cast(proposals.size()); + const int ar_count = static_cast(ar_lanes.size()); + const int tree_bucket = chain_decode_bucket_width(spec_count); + const int tree_rows_count = tree_width * tree_bucket; + const int total_rows = ar_count + tree_rows_count; + + int max_prefix = 1; + for (const Proposal & proposal : proposals) { + max_prefix = std::max( + max_prefix, slots_.slot(proposal.slot).cur_pos); + } + for (const ArLane & lane : ar_lanes) { + max_prefix = std::max(max_prefix, lane.position + 1); + } + + StepGraph & graph = b_.sg_; + if (!build_target_step_paged_tree( + graph, b_.w_, b_.cache_, b_.target_backend_, + tree_width, tree_bucket, max_prefix, + tree_scratch_base_, tree_scratch_stride_, + b_.cfg_.kq_stride_pad, ar_count)) { + result.error = "fixed chain target graph build failed"; + return result; + } + + std::vector tokens(static_cast(total_rows), 0); + std::vector parents( + static_cast(tree_rows_count), -1); + std::vector tree_sizes( + static_cast(tree_bucket), 0); + std::vector mapped_slots( + static_cast(ar_count + tree_bucket), -1); + std::vector state_slots( + static_cast(ar_count + tree_bucket), 0); + std::vector query_slots( + static_cast(total_rows), -1); + std::vector query_positions( + static_cast(total_rows), -1); + std::vector write_rows( + static_cast(total_rows) * n_head_kv, scratch_row_); + std::vector positions( + static_cast(4) * total_rows, 0); + std::vector embeddings( + static_cast(hidden) * total_rows, 0.0f); + seq_lens_.assign(static_cast(n_slots), 0); + + for (int lane_index = 0; lane_index < ar_count; ++lane_index) { + const ArLane & lane = ar_lanes[static_cast(lane_index)]; + mapped_slots[static_cast(lane_index)] = lane.slot; + state_slots[static_cast(lane_index)] = lane.slot; + tokens[static_cast(lane_index)] = lane.token; + query_slots[static_cast(lane_index)] = lane.slot; + query_positions[static_cast(lane_index)] = lane.position; + seq_lens_[static_cast(lane.slot)] = lane.position + 1; + for (int axis = 0; axis < 3; ++axis) { + positions[static_cast(axis) * total_rows + lane_index] = + lane.position; + } + for (int head = 0; head < n_head_kv; ++head) { + write_rows[static_cast(head) * total_rows + lane_index] = + lane.physical_row; + } + } + + for (int lane_index = 0; lane_index < spec_count; ++lane_index) { + const Proposal & proposal = + proposals[static_cast(lane_index)]; + const int tree_base = lane_index * tree_width; + const int row_base = ar_count + tree_base; + const int mapped_lane = ar_count + lane_index; + tree_sizes[static_cast(lane_index)] = tree_width; + mapped_slots[static_cast(mapped_lane)] = proposal.slot; + 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 < tree_width; ++node) { + const int tree_row = tree_base + node; + const int row = row_base + node; + tokens[static_cast(row)] = + proposal.tokens[static_cast(node)]; + 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 + ? 0 + : proposal.tree.depths[static_cast(node) - 1]; + const int position = + slots_.slot(proposal.slot).cur_pos + depth; + for (int axis = 0; axis < 3; ++axis) { + positions[static_cast(axis) * total_rows + row] = + position; + } + for (int head = 0; head < n_head_kv; ++head) { + write_rows[static_cast(head) * total_rows + row] = + static_cast(tree_scratch_base_) + + static_cast(proposal.slot) * + tree_scratch_stride_ + node; + } + } + } + + if (!b_.w_.embedder.embed( + tokens.data(), total_rows, embeddings.data())) { + result.error = "fixed chain embedding failed"; + return result; + } + ggml_backend_tensor_set( + graph.inp_embed, embeddings.data(), 0, + sizeof(float) * embeddings.size()); + ggml_backend_tensor_set( + graph.positions, positions.data(), 0, + sizeof(int32_t) * positions.size()); + ggml_backend_tensor_set( + graph.parent_ids, parents.data(), 0, + sizeof(int32_t) * parents.size()); + ggml_backend_tensor_set( + graph.tree_sizes, tree_sizes.data(), 0, + sizeof(int32_t) * tree_sizes.size()); + if (detail::target_paged_tree_active_slots_need_upload(graph)) { + ggml_backend_tensor_set( + graph.active_slot_ids, mapped_slots.data(), 0, + sizeof(int32_t) * mapped_slots.size()); + } + ggml_backend_tensor_set( + graph.state_slot_ids, state_slots.data(), 0, + sizeof(int32_t) * state_slots.size()); + ggml_backend_tensor_set( + graph.paged_query_seq_ids, query_slots.data(), 0, + sizeof(int32_t) * query_slots.size()); + if (graph.paged_query_positions) { + ggml_backend_tensor_set( + graph.paged_query_positions, query_positions.data(), 0, + sizeof(int32_t) * query_positions.size()); + } + ggml_backend_tensor_set( + graph.kv_write_rows, write_rows.data(), 0, + sizeof(int64_t) * write_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_, graph.gf) != + GGML_STATUS_SUCCESS) { + result.error = "fixed chain target compute failed"; + return result; + } + + std::vector posterior( + static_cast(total_rows), -1); + ggml_backend_tensor_get( + graph.argmax_tokens, posterior.data(), 0, + sizeof(int32_t) * posterior.size()); + + for (int lane_index = 0; lane_index < spec_count; ++lane_index) { + Proposal & proposal = proposals[static_cast(lane_index)]; + const int row_base = ar_count + lane_index * tree_width; + const int32_t * lane_posterior = + posterior.data() + static_cast(row_base); + int32_t ignored_bonus = -1; + proposal.accepted = follow_verified_tree( + proposal.tree, lane_posterior, ignored_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, ignored_bonus); + if (proposal.accepted.empty()) { + result.error = "fixed chain accepted path is empty"; + return result; + } + proposal.path.reserve(proposal.accepted.size()); + for (int node : proposal.accepted) { + proposal.path.push_back(node == 0 + ? proposal.root + : proposal.tree.token_ids[ + static_cast(node) - 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); + if (proposal.path.empty()) { + result.error = "fixed chain min-token clamp removed the root"; + return result; + } + } + + std::vector accepted_prefixes( + static_cast(tree_bucket), 0); + std::vector commit_slots( + static_cast(tree_bucket), -1); + std::vector commit_rows( + static_cast(tree_rows_count), -1); + std::vector feature_commit_rows( + static_cast(total_rows), -1); + const int feature_cap = b_.cache_.target_feat_cap; + + for (int lane_index = 0; lane_index < ar_count; ++lane_index) { + const ArLane & lane = ar_lanes[static_cast(lane_index)]; + feature_commit_rows[static_cast(lane_index)] = + lane.slot * feature_cap + lane.position % feature_cap; + } + + for (int lane_index = 0; lane_index < spec_count; ++lane_index) { + Proposal & proposal = proposals[static_cast(lane_index)]; + const Qwen35SlotManager::StepAppend append = slots_.append_tokens( + proposal.slot, proposal.path.data(), + static_cast(proposal.path.size())); + if (!append.ok || + append.physical_rows.size() != proposal.path.size()) { + result.error = append.busy + ? "paged KV pool exhausted during chain staging" + : "accepted chain K/V staging failed"; + return result; + } + if (!upload_block_table_delta( + proposal.slot, append.first_new_block, + append.new_blocks.data(), append.new_blocks.size())) { + result.error = "accepted chain block-table update failed"; + return result; + } + + accepted_prefixes[static_cast(lane_index)] = + static_cast(proposal.path.size()); + commit_slots[static_cast(lane_index)] = proposal.slot; + for (size_t depth = 0; depth < proposal.accepted.size(); ++depth) { + const int node = proposal.accepted[depth]; + if (node != static_cast(depth)) { + result.error = "DFlash2 chain acceptance is not contiguous"; + return result; + } + const int flat = lane_index * tree_width + node; + const int graph_row = ar_count + flat; + commit_rows[static_cast(flat)] = + append.physical_rows[depth]; + feature_commit_rows[static_cast(graph_row)] = + proposal.slot * feature_cap + + (append.position + static_cast(depth)) % feature_cap; + } + const int seq_len = + append.position + static_cast(proposal.path.size()); + seq_lens_[static_cast(proposal.slot)] = seq_len; + } + + ggml_backend_tensor_set( + graph.accepted_prefixes, accepted_prefixes.data(), 0, + sizeof(int32_t) * accepted_prefixes.size()); + ggml_backend_tensor_set( + graph.commit_slot_ids, commit_slots.data(), 0, + sizeof(int32_t) * commit_slots.size()); + ggml_backend_tensor_set( + graph.commit_rows, commit_rows.data(), 0, + sizeof(int64_t) * commit_rows.size()); + ggml_backend_tensor_set( + graph.feature_commit_rows, feature_commit_rows.data(), 0, + sizeof(int32_t) * feature_commit_rows.size()); + + const size_t n_delta = b_.cache_.ssm_state.size(); + if (n_delta == 0 || graph.delta_captures.size() != n_delta || + b_.cache_.conv_state.size() != n_delta || + !graph.tree_features || !b_.cache_.target_feat) { + result.error = "fixed chain capture set is 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 = graph.delta_captures[layer]; + if (!capture.transition_journal || !capture.conv_input || + !b_.cache_.ssm_state[layer] || + !b_.cache_.conv_state[layer]) { + result.error = "fixed chain layer capture is 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 caches; + caches.reserve(b_.cache_.attn_k.size() + b_.cache_.attn_v.size()); + for (ggml_tensor * tensor : b_.cache_.attn_k) { + if (tensor) caches.push_back(tensor); + } + for (ggml_tensor * tensor : b_.cache_.attn_v) { + if (tensor) caches.push_back(tensor); + } + if (caches.empty()) { + result.error = "fixed chain K/V cache set is empty"; + return result; + } + + if (!ggml_backend_cuda_tree_cache_commit_many( + caches.data(), static_cast(caches.size()), + graph.commit_rows, graph.commit_slot_ids, + tree_scratch_base_, tree_scratch_stride_)) { + result.error = "fixed chain K/V promotion failed"; + return result; + } + if (!ggml_backend_cuda_tree_feature_commit( + graph.tree_features, b_.cache_.target_feat, + graph.feature_commit_rows)) { + result.error = "fixed chain 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), + graph.accepted_prefixes, graph.commit_slot_ids)) { + result.error = "fixed chain recurrent promotion failed"; + return result; + } + ggml_backend_tensor_set( + b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + + for (const StepInput & input : inputs) { + slots_.commit_step(input.slot); + } + for (int lane_index = 0; lane_index < spec_count; ++lane_index) { + Proposal & proposal = proposals[static_cast(lane_index)]; + const int graph_row = ar_count + lane_index * tree_width + + 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 = "fixed chain sampling failed"; + return result; + } + } + for (int lane_index = 0; lane_index < ar_count; ++lane_index) { + ArLane & lane = ar_lanes[static_cast(lane_index)]; + lane.pending = sample_graph_row( + lane.slot, lane_index, + &posterior[static_cast(lane_index)], &logits_buf_); + if (lane.pending < 0) { + result.error = "compact AR sampling failed"; + return result; + } + } + + result.decode.reserve(inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + DecodeOutput output; + output.slot = inputs[i].slot; + if (selected[i]) { + Proposal & proposal = + proposals[static_cast(proposal_for_input[i])]; + output.token = proposal.pending; + output.committed_tokens.assign( + proposal.path.begin() + 1, proposal.path.end()); + } else { + output.token = + ar_lanes[static_cast(ar_for_input[i])].pending; + } + result.decode.push_back(std::move(output)); + } + return result; +} + SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { StepResult result; std::vector & decode_outputs = result.decode; @@ -207,6 +990,19 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (inputs.empty() && plan.prefills.empty()) return result; + FixedServiceRound service_round = make_fixed_service_round(plan); + if (auto * decode_round = + std::get_if(&service_round)) { + const bool has_chain_lane = std::any_of( + decode_round->chain_lanes.begin(), + decode_round->chain_lanes.end(), + [](uint8_t selected) { return selected != 0; }); + if (has_chain_lane && + prepare_chain_drafts(inputs, decode_round->chain_lanes)) { + return step_chain_spec(plan, decode_round->chain_lanes); + } + } + const TargetWeights & w = b_.w_; StepGraph & sg = b_.sg_; const int hidden = w.n_embd; @@ -329,7 +1125,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 +1143,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 +1161,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))) { @@ -409,6 +1206,32 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { b_.target_backend_, sg.positions, pos_buf_.data(), 0, sizeof(int32_t) * pos_buf_.size()); + if (capture_features_) { + const int cap = b_.cache_.target_feat_cap; + const int dead_row = cap * n_slots; + feature_rows_.assign(static_cast(n_total), dead_row); + token_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_[static_cast(token_offset + row)] = + slot * cap + (prefill.kv_pos + row) % cap; + } + token_offset += prefill.chunk; + } + for (int row = 0; row < live_count; ++row) { + const int slot = live_slot_ids_[static_cast(row)]; + feature_rows_[static_cast(n_prefill + row)] = + slot * cap + + live_positions_[static_cast(row)] % cap; + } + ggml_backend_tensor_set_async( + b_.target_backend_, sg.target_feat_rows, + feature_rows_.data(), 0, + sizeof(int32_t) * feature_rows_.size()); + } + rows_buf_.assign((size_t)n_total * n_head_kv, scratch_row_); for (int h = 0; h < n_head_kv; ++h) { token_offset = 0; @@ -537,6 +1360,14 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { void Qwen35SeqEngine::retire(int slot) { if (!slots_.is_active(slot)) return; + 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)]); + } + if (slot >= 0 && + slot < static_cast(prepared_chain_drafts_.size())) { + prepared_chain_drafts_[static_cast(slot)] = {}; + } slots_.retire(slot); } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index d9391784c..c941dd242 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -22,10 +22,14 @@ #pragma once #include "common/concurrency/seq_engine.h" +#include "common/dflash_draft_kv.h" +#include "common/dflash_feature_ring.h" #include "qwen35_slot_manager.h" #include #include +#include +#include #include namespace dflash::common { @@ -34,25 +38,21 @@ class Qwen35Backend; class Qwen35SeqEngine final : public SeqEngine { public: - // `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. + // `pool` and `backend` must outlive the engine. `scratch_row` is outside + // the pool and any per-slot tree slabs; it is the K/V destination of + // graph-bucket padding rows. // `max_prefills` bounds scheduler-selected prompt slices per traversal. 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 = 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(); } @@ -95,6 +95,20 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector embeddings; }; + struct PromptServiceRound {}; + struct SpeculativeDecodeRound { + std::vector chain_lanes; + }; + using FixedServiceRound = + std::variant; + + struct PreparedChainDraft { + bool valid = false; + int generated = -1; + int32_t root = -1; + std::vector tokens; + }; + int max_prefills_; int mixed_prefill_tokens_; int long_mixed_prefill_tokens_; @@ -112,10 +126,30 @@ 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); + FixedServiceRound make_fixed_service_round( + const StepPlan & plan) const; + bool chain_spec_input_capable(const StepInput & input) const; + DraftFeatureMirror * slot_feature_mirror(int slot); + DraftKvState * ensure_slot_draft_kv(int slot); + bool prepare_chain_drafts( + const std::vector & inputs, + const std::vector & selected); + StepResult step_chain_spec( + const StepPlan & plan, const std::vector & selected); 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_; + std::vector> dummy_draft_kv_; + DraftKvBatchGraph batch_draft_graph_; + std::vector prepared_chain_drafts_; // Hoisted per-step buffers (reused across step() calls). std::vector output_rows_; @@ -131,6 +165,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/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 118963305..b859d09fd 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -8,11 +8,91 @@ #include #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 < 2 || tree_width > 16 || + 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, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride) { + 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 || + 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; +} + // ── build_layer_step ──────────────────────────────────────────── bool build_layer_step( @@ -336,7 +416,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 // every node property stable across AR decode steps -> captured graph @@ -494,10 +578,17 @@ 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); + } // 32k nodes: the chunked delta-net prefill graph (CS = 32) reaches ~17k // nodes at a 512-token ubatch. - sg.gf = ggml_new_graph_custom(sg.ctx, 32768, false); + graph_capacity = std::max(graph_capacity, 32768); + 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 @@ -558,6 +649,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; gi.specla_m_strict = sg.specla_m_strict; @@ -699,6 +791,163 @@ bool build_target_step_tree( return true; } +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, + int mapped_ar_seqs) { + (void)kq_stride_pad; + step_graph_free(sg); + + if (mapped_ar_seqs < 0 || mapped_ar_seqs > cache.n_seq_slots) { + 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)) { + 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 = 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; + 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 + + mapped_ar_seqs + n_tokens; ++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_mapped_seqs); + sg.state_slot_ids = + 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); + + 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.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); + } + + sg.gf = ggml_new_graph_custom(sg.ctx, graph_capacity, 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 = true; + gi.capture_delta_intermediate = false; + gi.capture_tree_commit = true; + 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.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; + 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; + sg.delta_captures = std::move(go.delta_captures); + sg.tree_features = go.tree_features; + if (!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"); + 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)); + } + if (!ggml_gallocr_alloc_graph(sg.alloc, sg.gf) || + !detail::target_paged_tree_uploads_ready(sg)) { + return false; + } + 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; +} + // ── build_lm_head_projection_step ─────────────────────────────── diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index a56a0f79b..4ab9f2266 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -24,6 +24,60 @@ 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 fixed-chain shape/capacity contract. DFlash2 uses widths 2..16, and +// concurrent serving supports 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 +// 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); + +// `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) && + (!sg.paged_query_positions || + allocated(sg.paged_query_positions)) && + 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. bool build_layer_step( StepGraph & sg, @@ -110,6 +164,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, @@ -148,6 +206,27 @@ bool build_target_step_tree( int kq_stride_pad = KQ_MASK_PAD, const SpecLAHLDSchedule * specla_hld = nullptr); +// Packed fixed-chain verify over a paged multi-slot cache. Tokens are +// 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; recurrent transitions and target +// features are exposed for post-verification promotion. +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, + int mapped_ar_seqs = 0); + // 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 239ba9280..c13c726cf 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -179,7 +179,8 @@ static FILE * open_dflash_floor_log() { // staging K/V or staging recurrent slab to reserve. static int64_t concurrent_fixed_cache_bytes( const TargetWeights & w, int max_ctx, int n_slots, - int64_t kv_bytes_per_token) { + int64_t kv_bytes_per_token, int64_t scratch_tokens, + bool fixed_chain) { const int64_t n_full_attn = w.n_layer / w.full_attention_interval; const int64_t n_delta = w.n_layer - n_full_attn; @@ -194,7 +195,10 @@ 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); + (fixed_chain + ? (int64_t)std::min(max_ctx, 4096) * n_slots + 1 + : (int64_t)std::min(max_ctx, 4096)) * + (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); @@ -202,7 +206,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; } @@ -448,6 +452,27 @@ bool Qwen35Backend::init() { set_last_error("--max-concurrency requires --paged-attention"); return false; } + const bool concurrent_local_chain = + n_slots > 1 && cfg_.paged_attention && cfg_.fa_window == 0 && + cfg_.draft_path && !cfg_.ddtree_mode && !use_remote_draft && + !tensor_parallel && !split_gpus_ && + target_backend_ == draft_backend_ && + cfg_.device.gpu == cfg_.draft_gpu && + dw_.selector.enabled && dw_.block_size > 1 && + dw_.block_size <= 16 && w_.output; + if (n_slots > 1 && cfg_.draft_path && !concurrent_local_chain) { + set_last_error( + "concurrent paged DFlash2 requires a selector-enabled local " + "same-device draft with block size in [2, 16] and a target " + "lm_head"); + return false; + } + const int tree_width = + concurrent_local_chain ? dw_.block_size : 0; + const int tree_stride = concurrent_local_chain + ? paged_token_capacity(tree_width) : 0; + const int64_t scratch_tokens = PAGED_BLOCK_SIZE + + (int64_t)n_slots * tree_stride; // 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 @@ -458,9 +483,12 @@ bool Qwen35Backend::init() { int64_t pool_tokens = 0; if (n_slots > 1) { if (cfg_.kv_pool_tokens > 0) { + const int64_t max_pool_tokens = + ((int64_t)INT32_MAX - scratch_tokens) / + PAGED_BLOCK_SIZE * PAGED_BLOCK_SIZE; pool_tokens = (int64_t)paged_token_capacity( (int)std::min( - cfg_.kv_pool_tokens, INT32_MAX - PAGED_BLOCK_SIZE)); + cfg_.kv_pool_tokens, max_pool_tokens)); } else { PagedKvAutoBudget budget; // TODO: Size tensor-parallel pools from each device's free memory @@ -469,7 +497,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, + scratch_tokens, concurrent_local_chain); pool_tokens = paged_kv_auto_pool_tokens( cfg_.device.max_ctx, n_slots, budget); const int64_t one_context = @@ -491,19 +520,20 @@ bool Qwen35Backend::init() { return false; } } - if (pool_tokens + PAGED_BLOCK_SIZE > INT32_MAX) { + if (pool_tokens + 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 + 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_chain)) { std::fprintf(stderr, "cache: %s\n", dflash27b_last_error()); return false; } @@ -535,12 +565,22 @@ bool Qwen35Backend::init() { 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); + if (concurrent_local_chain) { + std::fprintf(stderr, + "[parallel-chain] fixed DFlash2 width=%d, " + "same-device paged full-attention greedy lanes\n", + tree_width); + } std::printf("[parallel] %d decode slots, up to %d packed prefills " "(mixed short/long %d/%d at >=%d tokens, " "idle %d, quantum %d), " @@ -593,7 +633,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 && + !concurrent_local_chain) { 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_, @@ -1238,6 +1279,7 @@ DFlashTarget * Qwen35Backend::dflash_target() { void Qwen35Backend::shutdown() { const bool use_remote_draft = cfg_.remote_draft.enabled(); + seq_engine_.reset(); end_paged_sequence(); free_drafter(); step_graph_destroy(sg_); diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 66069e104..c1e44da4f 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -85,12 +85,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 @@ -108,7 +110,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) { @@ -120,6 +123,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; @@ -238,7 +246,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 = concurrent_tree + ? 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; @@ -277,9 +292,9 @@ 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. Fixed chain verification keeps + // speculative recurrent transitions in graph scratch and promotes only + // accepted prefixes through the compact GPU journal. if (!prefill_only && !multi_slot) { const int rb_tensors = 4 * n_delta; ggml_init_params ip{}; @@ -1109,7 +1124,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; @@ -1197,9 +1221,13 @@ 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 && 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. @@ -1265,12 +1293,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. Paged serving can + // map 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)); } @@ -1278,7 +1325,20 @@ 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. + // 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, paged_query_positions, + /*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 @@ -1286,6 +1346,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, @@ -1307,8 +1368,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, @@ -1402,6 +1463,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, // SpecLA topology masks (all three non-null together): route the // recurrence through the topology-masked factor-capture verify. @@ -1431,14 +1493,34 @@ 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(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 || - (!cap && !parent_ids && prefill_total + n_seqs == n_tokens)); + (mapped_tree + ? (!ragged && prefill_total == 0 && + 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 + // 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); + }; // Fully factorized SpecLA fallback: the current candidates do not mutate // durable state. The HLD route below may materialize the *previously* // accepted pending path while keeping current candidates speculative. @@ -1452,12 +1534,6 @@ static ggml_tensor * build_delta_net_block( GGML_ASSERT(!(use_specla_factorized || use_specla_hld) || (n_seqs == 1 && !ragged && !active_slot_ids)); - // Row-slices of a stacked projection are only contiguous for a single - // token; wider batches (verify/prefill) need a copy before reshape/unary. - auto contig = [&](ggml_tensor * t) { - return ggml_is_contiguous(t) ? t : ggml_cont(ctx, t); - }; - // ── Whole-batch projections ───────────────────────────────────── // qkv_mixed = wqkv @ cur [10240, n_tokens] // z = wqkv_gate @ cur [inner, n_tokens] @@ -1481,16 +1557,16 @@ static ggml_tensor * build_delta_net_block( // alpha = ssm_alpha @ cur [dt_rank, n_tokens] // One GEMV over the stacked (beta | alpha) alias when available. ggml_tensor * beta_2d = nullptr; - ggml_tensor * alpha = 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); // [2 * dt_rank, n_tokens] 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 = contig(ggml_view_2d(ctx, ba, num_v_heads, n_tokens, ba->nb[1], (size_t)num_v_heads * e)); + 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 = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_alpha, cur), L.ssm_alpha_s); + alpha_2d = apply_scale2(ctx, ggml_mul_mat(ctx, L.ssm_alpha, cur), L.ssm_alpha_s); } // Fused kernels (single-sequence chain path only): the conv step and the @@ -1512,25 +1588,6 @@ static ggml_tensor * build_delta_net_block( const char * s_env = std::getenv("DFLASH27B_CHUNKED"); return s_env && std::atoi(s_env) == 1; }(); - const bool chunked_call = chunked_env_on && can_skip_gdn_intermediate && !ragged && - !active_slot_ids && !use_specla_factorized && !use_specla_hld && n_tokens > 1; - const bool fused_plain = fused_kernels_env && !parent_ids && !ragged && !active_slot_ids && - !use_specla_factorized && !use_specla_hld; - const bool fused_conv = fused_plain; - const bool raw_gates = fused_plain && !chunked_call && L.ssm_gate_ba != nullptr; - - // beta = sigmoid(beta); g = softplus(alpha + ssm_dt_bias) * ssm_a - // (-A_log.exp() * softplus). In raw-gate mode the GDN kernel applies both - // itself (dt_bias / A are attached via ggml_gated_delta_net_set_raw_gates). - ggml_tensor * g_2d = nullptr; - if (raw_gates) { - g_2d = alpha; - } else { - beta_2d = ggml_sigmoid(ctx, beta_2d); - alpha = ggml_add(ctx, alpha, L.ssm_dt_bias); - alpha = ggml_softplus(ctx, alpha); - g_2d = ggml_mul(ctx, alpha, L.ssm_a); - } // ── Token-axis segments: prompt chunks first, then the decode batch ── struct DeltaSeg { @@ -1538,8 +1595,11 @@ 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; + ggml_tensor * active_ids; + ggml_tensor * state_ids; }; std::vector segs; segs.reserve((size_t)n_prefill_segments + 1); @@ -1555,14 +1615,35 @@ 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, nullptr, nullptr}); } if (active_slot_ids) { - segs.push_back({prefill_total, 1, n_seqs, true, - 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, conv_state, ssm_state}); + segs.push_back({0, n_tokens, n_seqs, false, false, + conv_state, ssm_state, nullptr, nullptr}); } const int n_segs = (int)segs.size(); @@ -1582,13 +1663,20 @@ 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; + 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 // 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 dense_chain = !ragged && !active_slot_ids && !seg_tree; + const bool inplace_state = (seg_active && !seg_tree) || (allow_inplace_state && can_skip_gdn_intermediate && !ragged && n_seq_tokens == 1); @@ -1601,12 +1689,34 @@ static ggml_tensor * build_delta_net_block( if (use_specla_hld || use_specla_factorized) { qkv_mixed = contig(qkv_mixed); // the SpecLA conv kernels raw-index x } + const bool use_chunked = chunked_env_on && can_skip_gdn_intermediate && + !ragged && !active_slot_ids && !seg_tree && + !use_specla_factorized && !use_specla_hld && n_seq_tokens > 1; + const bool fused_plain = fused_kernels_env && dense_chain && + !parent_ids && !seg_parent_ids && !use_specla_factorized && + !use_specla_hld; + const bool fused_conv = fused_plain; + const bool raw_gates = fused_plain && !use_chunked && L.ssm_gate_ba; + 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); + } ggml_tensor * conv_out = nullptr; if (use_specla_hld) { @@ -1633,7 +1743,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 { @@ -1641,74 +1751,77 @@ static ggml_tensor * build_delta_net_block( w.ssm_d_conv - 1, conv_channels, seg_seqs); } - if (fused_conv) { + if (fused_conv && !use_specla_factorized) { // 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 { - // 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) { - if (use_specla_factorized) { + // 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 (seg_cap && seg_cap->conv_input && use_specla_factorized) { // The consolidated SpecLA bank is [channels, layers, tokens]. // Capture only the raw current inputs; the compatibility commit // shifts the durable K-1 window and appends accepted tokens. - GGML_ASSERT(qkv_mixed->ne[0] == cap->conv_input->ne[0]); - GGML_ASSERT(n_seq_tokens <= cap->conv_input->ne[1]); - ggml_tensor * dst = ggml_view_3d(ctx, cap->conv_input, - cap->conv_input->ne[0], n_seq_tokens, 1, - cap->conv_input->nb[1], cap->conv_input->nb[2], 0); + GGML_ASSERT(qkv_mixed->ne[0] == seg_cap->conv_input->ne[0]); + GGML_ASSERT(n_seq_tokens <= seg_cap->conv_input->ne[1]); + ggml_tensor * dst = ggml_view_3d(ctx, seg_cap->conv_input, + seg_cap->conv_input->ne[0], n_seq_tokens, 1, + seg_cap->conv_input->nb[1], seg_cap->conv_input->nb[2], 0); GGML_ASSERT(ggml_nelements(qkv_mixed) == ggml_nelements(dst)); ggml_build_forward_expand(gf, ggml_cpy(ctx, qkv_mixed, dst)); - } else { + } else 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)); } - } - // ── Save the last (kernel-1) steps back to the conv state - // SpecLA: skipped — the window is speculative; the commit path - // shifts conv_state and appends the accepted raw inputs at commit time. - if (!use_specla_factorized) { + 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 + // SpecLA factorized: skipped — the window is speculative; the + // commit path shifts conv_state and appends accepted raw inputs. 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) { + if (!use_specla_factorized) { + 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( @@ -1717,23 +1830,24 @@ 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)); - } else { - ggml_build_forward_expand(gf, ggml_cpy(ctx, last_conv, seg.conv_st)); + 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)); + } } - } - // ── 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); - } + // ── 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 = 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); } + } // !use_specla_hld // conv_out: [conv_channels, n_tokens, n_seqs] const int64_t q_offset = 0; @@ -1786,21 +1900,36 @@ static ggml_tensor * build_delta_net_block( // produces); the chunked, compact-decode and SpecLA paths take the // materialized copies. if (num_k_heads != num_v_heads && - (chunked_call || seg_active || use_specla_factorized || use_specla_hld)) { + (use_chunked || seg_active || use_specla_factorized || use_specla_hld)) { 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); } // ── 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 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, seg.state_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 // 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. @@ -1816,10 +1945,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 @@ -1830,11 +1959,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. - // Chunked delta-net path (opt-in via DFLASH27B_CHUNKED, chain-only, no - // capture): decided whole-batch above; a segment only qualifies with - // more than one timestep. - const bool use_chunked = chunked_call && n_seq_tokens > 1; - ggml_tensor * output = nullptr; if (use_specla_hld) { @@ -1894,14 +2018,14 @@ 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) { + 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 @@ -1921,6 +2045,16 @@ static ggml_tensor * build_delta_net_block( ggml_gated_delta_net_set_raw_gates(result, L.ssm_gate_ba); } } + 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 (can_skip_gdn_intermediate) { ggml_gated_delta_net_set_skip_intermediate(result, true); } @@ -1935,7 +2069,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, @@ -1943,8 +2077,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)); } @@ -1959,10 +2093,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 @@ -1971,7 +2105,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); } } @@ -2139,7 +2273,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); @@ -2154,6 +2288,14 @@ 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; + const bool capture_tree_features = + in.capture_layers && in.capture_tree_commit && cache.target_feat; + std::vector capture_slices; + if (capture_with_rows || capture_tree_features) { + capture_slices.assign((size_t)N_CAPTURE, nullptr); + } for (int il = 0; il < w.n_layer; il++) { const TargetLayer & L = w.layers[il]; @@ -2183,7 +2325,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 @@ -2202,37 +2350,39 @@ 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. - cap_ptr->ssm_intermediate_states = cache.ssm_intermediate[dn_idx]; - cap_ptr->conv_input = cache.conv_input_cache[dn_idx]; - if (!cache.factor_k.empty()) { - const bool pending_alt = cache.specla_pending_bank != 0; - cap_ptr->pending_factor_k = pending_alt - ? cache.factor_k_alt[dn_idx] : cache.factor_k[dn_idx]; - cap_ptr->pending_factor_v_new = pending_alt - ? cache.factor_v_new_alt[dn_idx] : cache.factor_v_new[dn_idx]; - cap_ptr->pending_factor_g = pending_alt - ? cache.factor_g_ps_alt[dn_idx] : cache.factor_g_ps[dn_idx]; - cap_ptr->pending_conv_input = pending_alt - ? cache.conv_input_cache_alt[dn_idx] : cache.conv_input_cache[dn_idx]; - cap_ptr->factor_k = pending_alt - ? cache.factor_k[dn_idx] : cache.factor_k_alt[dn_idx]; - cap_ptr->factor_v_new = pending_alt - ? cache.factor_v_new[dn_idx] : cache.factor_v_new_alt[dn_idx]; - cap_ptr->factor_g_ps = pending_alt - ? cache.factor_g_ps[dn_idx] : cache.factor_g_ps_alt[dn_idx]; - cap_ptr->conv_input = pending_alt - ? cache.conv_input_cache[dn_idx] : cache.conv_input_cache_alt[dn_idx]; - cap_ptr->factor_ptrs = cache.specla_factor_ptrs; - cap_ptr->factor_n_layers = (int)cache.factor_k.size(); - cap_ptr->factor_layer = dn_idx; - cap_ptr->pending_bank = cache.specla_pending_bank; + 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]; + if (!cache.factor_k.empty()) { + const bool pending_alt = cache.specla_pending_bank != 0; + cap_ptr->pending_factor_k = pending_alt + ? cache.factor_k_alt[dn_idx] : cache.factor_k[dn_idx]; + cap_ptr->pending_factor_v_new = pending_alt + ? cache.factor_v_new_alt[dn_idx] : cache.factor_v_new[dn_idx]; + cap_ptr->pending_factor_g = pending_alt + ? cache.factor_g_ps_alt[dn_idx] : cache.factor_g_ps[dn_idx]; + cap_ptr->pending_conv_input = pending_alt + ? cache.conv_input_cache_alt[dn_idx] : cache.conv_input_cache[dn_idx]; + cap_ptr->factor_k = pending_alt + ? cache.factor_k[dn_idx] : cache.factor_k_alt[dn_idx]; + cap_ptr->factor_v_new = pending_alt + ? cache.factor_v_new[dn_idx] : cache.factor_v_new_alt[dn_idx]; + cap_ptr->factor_g_ps = pending_alt + ? cache.factor_g_ps[dn_idx] : cache.factor_g_ps_alt[dn_idx]; + cap_ptr->conv_input = pending_alt + ? cache.conv_input_cache[dn_idx] : cache.conv_input_cache_alt[dn_idx]; + cap_ptr->factor_ptrs = cache.specla_factor_ptrs; + cap_ptr->factor_n_layers = (int)cache.factor_k.size(); + cap_ptr->factor_layer = dn_idx; + cap_ptr->pending_bank = cache.specla_pending_bank; + } } } ggml_tensor * conv_st = cache.conv_state[dn_idx]; @@ -2264,6 +2414,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, in.specla_m_strict, in.specla_m_incl, @@ -2304,6 +2455,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_tree_features) { + 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; @@ -2311,8 +2469,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 = @@ -2342,6 +2498,28 @@ QwenGraphOutputs build_qwen35_graph( inpL = cur; } + if (capture_with_rows || capture_tree_features) { + 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); + 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 ggml_tensor * out = rms_norm_mul(ctx, inpL, w.out_norm, w.rms_eps); diff --git a/server/test/test_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp new file mode 100644 index 000000000..74dd401ab --- /dev/null +++ b/server/test/test_chain_spec_shapes.cpp @@ -0,0 +1,89 @@ +#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_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})); + 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}; + 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.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_lanes == 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.commit_rows == 2); + + 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_ddtree_path.cpp b/server/test/test_ddtree_path.cpp new file mode 100644 index 000000000..f408be5b4 --- /dev/null +++ b/server/test/test_ddtree_path.cpp @@ -0,0 +1,44 @@ +#include "common/ddtree.h" +#include "host_check.h" + +#include +#include +#include + +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; + 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); + CHECK((accepted == std::vector{0, 1, 2})); + CHECK(pending == 33); + + CHECK(truncate_verified_path(accepted, 2, posterior, pending)); + CHECK((accepted == std::vector{0, 1})); + CHECK(pending == 22); + + CHECK(!truncate_verified_path(accepted, 2, posterior, pending)); + CHECK(pending == 22); + + 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_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..4552048da 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,28 @@ 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); + } + + 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 +154,6 @@ 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. const Case cases[] = { // Realistic decode shape: Qwen3.5 vocab, small position batch. {15, 151936, 8, 1.0f}, @@ -145,6 +166,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 +177,25 @@ 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. { - 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] unsupported K 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_feature_gate.cpp b/server/test/test_feature_gate.cpp index e01f97938..e04cd4278 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -406,7 +406,7 @@ void test_feature_gate_paged_attention_requires_qwen35_monolithic() { } } -void test_feature_gate_paged_attention_requires_plain_ar_decode() { +void test_feature_gate_paged_attention_allows_fixed_local_chains() { BackendArgs base; base.model_path = "/nonexistent/model.gguf"; base.paged_attention = true; @@ -415,6 +415,18 @@ void test_feature_gate_paged_attention_requires_plain_ar_decode() { draft.draft_path = "/nonexistent/draft.gguf"; CHECK(!gate_result(draft, "qwen35", PlacementBackend::Cuda).empty()); + 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 remote_chain = concurrent_chain; + remote_chain.remote_draft.ipc_bin = "/usr/bin/draft-ipc"; + CHECK(!gate_result( + remote_chain, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs ddtree = base; ddtree.ddtree_mode = true; CHECK(!gate_result(ddtree, "qwen35", PlacementBackend::Cuda).empty()); @@ -519,6 +531,21 @@ 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 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. BackendArgs overflow = paged; @@ -698,7 +725,7 @@ TEST_CASE(FeatureGateFixture, feature_gate_suite) { test_feature_gate_remote_draft_requires_supported_arch(); test_feature_gate_layer_split_requires_supported_arch(); test_feature_gate_paged_attention_requires_qwen35_monolithic(); - test_feature_gate_paged_attention_requires_plain_ar_decode(); + test_feature_gate_paged_attention_allows_fixed_local_chains(); test_feature_gate_parallel_and_kv_pool_rules(); test_feature_warnings_silent_when_supported(); test_feature_warnings_report_inert_draft(); diff --git a/server/test/test_gdn_transition_journal.cpp b/server/test/test_gdn_transition_journal.cpp new file mode 100644 index 000000000..cc3e82e2d --- /dev/null +++ b/server/test/test_gdn_transition_journal.cpp @@ -0,0 +1,789 @@ +#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 int CONV_WINDOW = 3; +constexpr int CONV_CHANNELS = 7; +constexpr float FIELD_TOLERANCE = 5.0e-5f; +constexpr float STATE_TOLERANCE = 2.0e-4f; + +bool test_raw_gate_protocol() { + ggml_init_params params{}; + params.mem_size = 128*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, H, 1, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, 1, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, H, 1, 1); + ggml_tensor * g = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, 1, 1); + ggml_tensor * beta = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, H, 1, 1); + ggml_tensor * state = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, S, H, 1); + ggml_tensor * gate_ba = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2*H); + ggml_tensor * result = ggml_gated_delta_net( + ctx, q, k, v, g, beta, state); + + ggml_gated_delta_net_set_raw_gates(result, gate_ba); + const int32_t * op_params = + reinterpret_cast(result->op_params); + const bool ok = result->src[9] == gate_ba && result->src[10] == nullptr && + ggml_nelements(result->src[9]) == 2*H && op_params[2] == 0 && + op_params[10] == 1; + if (!ok) { + std::fprintf( + stderr, + "raw gate protocol: src9=%p src10=%p elements=%lld op2=%d op10=%d\n", + static_cast(result->src[9]), + static_cast(result->src[10]), + (long long) ggml_nelements(result->src[9]), + op_params[2], op_params[10]); + } + ggml_free(ctx); + return ok; +} + +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; +} + +size_t conv_input_index( + int sequence, int channel, int position) { + return ((size_t) sequence*CONV_CHANNELS + channel)* + (CONV_WINDOW + T) + position; +} + +size_t conv_state_index(int slot, int channel, int position) { + return ((size_t) slot*CONV_CHANNELS + channel)*CONV_WINDOW + position; +} + +std::vector make_conv_input() { + std::vector values( + (size_t)(CONV_WINDOW + T)*CONV_CHANNELS*B); + for (int sequence = 0; sequence < B; ++sequence) { + for (int channel = 0; channel < CONV_CHANNELS; ++channel) { + for (int position = 0; position < CONV_WINDOW + T; ++position) { + values[conv_input_index(sequence, channel, position)] = + 0.01f*sequence + 0.001f*channel + 0.0001f*position; + } + } + } + return values; +} + +std::vector conv_state_for_slots( + const std::vector & input, + const std::vector & slots, + const std::vector & prefixes, + int physical_slots) { + std::vector state( + (size_t)CONV_WINDOW*CONV_CHANNELS*physical_slots, -1.0f); + for (int sequence = 0; sequence < B; ++sequence) { + const int slot = slots[(size_t)sequence]; + if (slot < 0 || slot >= physical_slots) continue; + const int prefix = prefixes[(size_t)sequence]; + for (int channel = 0; channel < CONV_CHANNELS; ++channel) { + for (int k = 0; k < CONV_WINDOW; ++k) { + state[conv_state_index(slot, channel, k)] = + input[conv_input_index( + sequence, channel, prefix + k)]; + } + } + } + return state; +} + +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); + + 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 * conv_input = nullptr; + ggml_tensor * identity_conv_state = nullptr; + ggml_tensor * mapped_conv_state = nullptr; + ggml_tensor * accepted = nullptr; + ggml_tensor * slots = nullptr; +}; + +bool commit_many_one_layer( + const ggml_tensor * journal, + ggml_tensor * state, + const ggml_tensor * conv_input, + ggml_tensor * conv_state, + const ggml_tensor * accepted, + const ggml_tensor * slots) { + const ggml_tensor * journals[] = {journal}; + ggml_tensor * states[] = {state}; + const ggml_tensor * conv_inputs[] = {conv_input}; + ggml_tensor * conv_states[] = {conv_state}; + return ggml_backend_cuda_gdn_transition_journal_commit_many( + journals, states, conv_inputs, conv_states, 1, accepted, slots); +} + +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) { + 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.conv_input = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, CONV_WINDOW + T, + CONV_CHANNELS, B, 1); + tensors.identity_conv_state = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, CONV_WINDOW, + CONV_CHANNELS, B, 1); + tensors.mapped_conv_state = ggml_new_tensor_4d( + tensors.ctx, GGML_TYPE_F32, CONV_WINDOW, + CONV_CHANNELS, PHYSICAL_SLOTS, 1); + 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 * gate_ba = nullptr; + if (raw_gates) { + gate_ba = ggml_new_tensor_1d(tensors.ctx, GGML_TYPE_F32, 2*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, gate_ba); + } + 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); + const std::vector conv_input = make_conv_input(); + upload_f32(tensors.conv_input, conv_input); + if (raw_gates) { + std::vector packed_gates; + packed_gates.reserve(2*H); + packed_gates.insert( + packed_gates.end(), inputs.dt_bias.begin(), inputs.dt_bias.end()); + packed_gates.insert( + packed_gates.end(), inputs.gate_A.begin(), inputs.gate_A.end()); + upload_f32(gate_ba, packed_gates); + } + + 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()); + std::vector actual_conv( + (size_t)CONV_WINDOW*CONV_CHANNELS*B); + const std::vector zero_prefixes(B, 0); + const std::vector identity_conv_base = conv_state_for_slots( + conv_input, identity_slots, zero_prefixes, B); + for (int prefix = 0; ok && prefix <= T; ++prefix) { + std::fill(accepted.begin(), accepted.end(), prefix); + upload_f32(tensors.identity_state, inputs.state); + upload_f32(tensors.identity_conv_state, identity_conv_base); + 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 = commit_many_one_layer( + tensors.journal, tensors.identity_state, + tensors.conv_input, tensors.identity_conv_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); + ggml_backend_tensor_get( + tensors.identity_conv_state, actual_conv.data(), 0, + actual_conv.size()*sizeof(float)); + ok = ok && compare_vectors( + "convolution prefix", actual_conv, + conv_state_for_slots( + conv_input, identity_slots, accepted, B), + 0.0f); + } + } + + 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()); + const std::vector mapped_conv_base = conv_state_for_slots( + conv_input, mapped_slots, zero_prefixes, PHYSICAL_SLOTS); + const std::vector mapped_conv_expected = conv_state_for_slots( + conv_input, mapped_slots, mapped_prefixes, PHYSICAL_SLOTS); + std::vector mapped_conv_actual(mapped_conv_base.size()); + if (ok) { + upload_f32(tensors.mapped_state, mapped_base); + upload_f32(tensors.mapped_conv_state, mapped_conv_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 = commit_many_one_layer( + tensors.journal, tensors.mapped_state, + tensors.conv_input, tensors.mapped_conv_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); + ggml_backend_tensor_get( + tensors.mapped_conv_state, mapped_conv_actual.data(), 0, + mapped_conv_actual.size()*sizeof(float)); + ok = ok && compare_vectors( + "permuted/padded convolution", mapped_conv_actual, + mapped_conv_expected, 0.0f); + } + } + 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); + upload_f32(tensors.identity_conv_state, identity_conv_base); + 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 = !commit_many_one_layer( + tensors.journal, tensors.identity_state, + tensors.conv_input, tensors.identity_conv_state, + tensors.accepted, tensors.slots); + const std::vector out_of_range_slots{0, 99, 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, out_of_range_slots.data(), 0, + out_of_range_slots.size()*sizeof(out_of_range_slots[0])); + ok = ok && !commit_many_one_layer( + tensors.journal, tensors.identity_state, + tensors.conv_input, tensors.identity_conv_state, + tensors.accepted, tensors.slots); + const std::vector duplicate_slots{0, 0, 2, 3}; + ggml_backend_tensor_set(tensors.slots, duplicate_slots.data(), 0, + duplicate_slots.size()*sizeof(duplicate_slots[0])); + ok = ok && !commit_many_one_layer( + tensors.journal, tensors.identity_state, + tensors.conv_input, tensors.identity_conv_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); + ggml_backend_tensor_get( + tensors.identity_conv_state, actual_conv.data(), 0, + actual_conv.size()*sizeof(float)); + ok = ok && compare_vectors( + "transactional convolution validation", actual_conv, + identity_conv_base, 0.0f); + } + + std::printf("gdn transition journal %-10s: %s\n", name, + ok ? "PASS" : "FAIL"); + destroy(tensors); + 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 * conv_input = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, CONV_WINDOW + T, CONV_CHANNELS, B, 1); + ggml_tensor * conv_state = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, CONV_WINDOW, CONV_CHANNELS, B, 1); + + 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); + const std::vector conv_values = make_conv_input(); + upload_f32(conv_input, conv_values); + + 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}; + const std::vector zero_prefixes(B, 0); + upload_f32(committed_state, inputs.state); + upload_f32( + conv_state, + conv_state_for_slots( + conv_values, identity_slots, zero_prefixes, B)); + 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 && commit_many_one_layer( + journal, committed_state, conv_input, conv_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::vector actual_conv( + (size_t)CONV_WINDOW*CONV_CHANNELS*B); + ggml_backend_tensor_get( + conv_state, actual_conv.data(), 0, + actual_conv.size()*sizeof(float)); + ok = ok && compare_vectors( + "grouped tree convolution commit", actual_conv, + conv_state_for_slots( + conv_values, identity_slots, prefixes, B), + 0.0f); + } + + std::printf("gdn grouped tree journal : %s\n", + ok ? "PASS" : "FAIL"); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ok; +} +} // namespace + +int main() { + if (!test_raw_gate_protocol()) return 1; + 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"); + return 1; + } + bool ok = run_case(backend, false, false); + ok = run_case(backend, true, false) && ok; + ok = run_case(backend, false, true) && ok; + ok = run_grouped_tree_case(backend) && ok; + ggml_backend_free(backend); + return ok ? 0 : 1; +} diff --git a/server/test/test_paged_attention.cpp b/server/test/test_paged_attention.cpp index 78338f717..e283fc3dc 100644 --- a/server/test/test_paged_attention.cpp +++ b/server/test/test_paged_attention.cpp @@ -30,6 +30,14 @@ struct TestCase { bool corrupt_blocks; }; +struct TreeMetadata { + int width; + int scratch_stride; + std::vector parent_ids; + std::vector tree_sizes; + int ar_rows = 0; +}; + int clamped_seq_len(const TestCase & test_case, int seq) { return std::max( 0, std::min( @@ -51,6 +59,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 +165,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; @@ -141,56 +181,86 @@ std::vector reference_attention( // Mirrors the kernel: out-of-range slot ids and negative positions // are padding rows and leave zero output. if (physical_seq < 0 || physical_seq >= physical_n_seq) continue; + const bool tree_query = tree && seq >= tree->ar_rows; + const int tree_seq = tree_query + ? (seq - tree->ar_rows) / tree->width : 0; + const int query_node = tree_query + ? (seq - tree->ar_rows) % tree->width : -1; int kv_seq_len = clamped_seq_len(test_case, physical_seq); - if (query_positions && (*query_positions)[seq] < kv_seq_len) { + if (query_positions && !tree_query) { + if ((*query_positions)[seq] < 0) continue; // The inclusive causal clamp: row seq attends its sequence's // cached tokens [0, position]. - kv_seq_len = (*query_positions)[seq] + 1; + kv_seq_len = std::min( + kv_seq_len, (*query_positions)[seq] + 1); + } + const int tree_size = tree_query ? tree->tree_sizes[tree_seq] : 0; + if (tree_query && + (tree_size < 0 || tree_size > tree->width || + query_node >= tree_size)) { + continue; + } + + std::vector physical_rows; + physical_rows.reserve( + kv_seq_len + (tree_query ? 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_query) { + 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 +275,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 +285,25 @@ 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); + if (tree) { + GGML_ASSERT(tree->width > 0); + GGML_ASSERT(tree->scratch_stride >= tree->width); + GGML_ASSERT(tree->ar_rows >= 0); + GGML_ASSERT(tree->ar_rows == 0 || query_positions); + GGML_ASSERT( + tree->parent_ids.size() == + static_cast(tree->width) * tree->tree_sizes.size()); + GGML_ASSERT( + n_seq == tree->ar_rows + + 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 +338,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 +416,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 +436,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 +449,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 +483,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 +515,57 @@ 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{ + 16, 16, + { + -1, 0, 1, 2, 3, 4, 5, 6, + 7, 8, 9, 10, 11, 12, 13, 14, + -1, 0, 1, 2, 3, 4, 5, 6, + 7, -1, -1, -1, -1, -1, -1, -1, + }, + {16, 9}, + }; + const std::vector tree_slot_ids{ + 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, + }; + 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_mixed_tree_case() { + ggml_backend_t backend = ggml_backend_cuda_init(0); + REQUIRE_NOT_NULL(backend); + const TestCase mixed_case{"mixed-tree", 8, {33, 65, 17}, false}; + const TreeMetadata tree_metadata{ + 16, 16, + { + -1, 0, 1, 2, 3, 4, 5, 6, + 7, 8, 9, 10, 11, 12, 13, 14, + }, + {16}, + 2, + }; + std::vector query_slots{1, 0}; + query_slots.insert(query_slots.end(), 16, 2); + std::vector query_positions{7, 15}; + query_positions.insert(query_positions.end(), 16, -1); + CHECK(run_case(backend, mixed_case, GGML_TYPE_F16, GGML_TYPE_F16, + &query_slots, &query_positions, &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 +628,14 @@ TEST_CASE(PagedAttention, CompactThreeSlotBucketMatchesReference) { }, {2, 1, 0, -1}); } +TEST_CASE(PagedAttention, PackedTreesMatchReference) { + run_tree_case(); +} + +TEST_CASE(PagedAttention, CompactArAndFixedChainMatchReference) { + run_mixed_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_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index 6fb5de4d6..dcf6f4628 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" @@ -11,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; @@ -32,11 +34,151 @@ static std::vector get_tensor(const ggml_tensor * tensor) { return values; } +TEST_CASE(RecurrentSnapshotFixture, validates_paged_tree_capacity_and_uploads) { + 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( + 16, 16, graph_capacity) && graph_capacity == 32768); + CHECK(!dflash::common::detail::target_paged_tree_graph_capacity( + 17, 16, graph_capacity)); + + { + 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) { + { + 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); if (!backend) SKIP("CPU backend is unavailable"); + { + size_t graph_capacity = 0; + CHECK(dflash::common::detail::target_paged_tree_graph_capacity( + 16, 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; diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index 5539c3d45..535fda85e 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -85,7 +85,6 @@ int main() { mgr.commit_prefill(0); CHECK(mgr.slot(0).cur_pos == 20); - // Decode rows stage until the target compute succeeds. auto st = mgr.append_token(0, /*fed_token=*/42); CHECK(st.ok); CHECK(st.position == 20); @@ -336,7 +335,6 @@ int main() { CHECK(pool.free_block_count() == 0); } - // A multi-token stage owns its rows until one atomic commit. { PagedKvPool pool(8, 1, /*block_size=*/4); Qwen35SlotManager mgr(pool, /*max_ctx=*/32); From 9dbc13d0db5a9c7864de5b3965800aca4f1a067d Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 22 Aug 2026 06:44:42 +0000 Subject: [PATCH 03/11] fix(concurrency): correct speculative backend contracts Restore the paged-attention C ABI through a separate tree entry point. Reject unsupported GDN variants, initialize root-only chain trees, and size draft metadata from graph capacity. --- server/deps/llama.cpp/ggml/include/ggml.h | 52 ++++++++----------- .../llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp | 6 ++- .../ggml/src/ggml-sycl/ggml-sycl.cpp | 9 +++- .../ggml/src/ggml-vulkan/ggml-vulkan.cpp | 10 ++-- server/deps/llama.cpp/ggml/src/ggml.c | 46 +++++++++++++++- .../common/concurrency/chain_spec_shapes.h | 2 +- server/src/common/dflash_draft_kv.cpp | 43 ++++++++++++--- server/src/qwen35/qwen35_target_graph.cpp | 17 +++--- server/test/test_batched_gdn.cpp | 49 ++++++++++++++++- server/test/test_chain_spec_shapes.cpp | 11 ++++ server/test/test_paged_attention.cpp | 17 +++--- 11 files changed, 204 insertions(+), 58 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index ef3806dc3..f76f319d4 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2503,7 +2503,20 @@ extern "C" { // marks a padding row. NULL keeps the decode semantics (full cached // length per row). // - // parent_ids/tree_sizes optionally enable packed tree verification. + GGML_API struct ggml_tensor * ggml_paged_attn_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * block_table, + struct ggml_tensor * kv_seq_lens, + struct ggml_tensor * active_slot_ids, + struct ggml_tensor * query_positions, + float scale, + int block_size, + int max_kv_seq_len); + + // Packed tree verification over the same paged K/V pool. // 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 @@ -2512,9 +2525,11 @@ extern "C" { // 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( + // at or beyond tree_sizes[s] are excluded. query_positions may describe + // compact autoregressive rows in a mixed AR/tree batch; tree rows ignore + // it and read the full committed prefix. Pure tree batches pass NULL. + // tree_width is derived from parent_ids. + GGML_API struct ggml_tensor * ggml_paged_attn_ext_tree( struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, @@ -2526,31 +2541,10 @@ extern "C" { float scale, int block_size, int max_kv_seq_len, - struct ggml_tensor * parent_ids -#ifdef __cplusplus - = nullptr -#endif - , - struct ggml_tensor * tree_sizes -#ifdef __cplusplus - = nullptr -#endif - , - int tree_width -#ifdef __cplusplus - = 0 -#endif - , - int tree_scratch_base -#ifdef __cplusplus - = 0 -#endif - , - int tree_scratch_stride -#ifdef __cplusplus - = 0 -#endif - ); + struct ggml_tensor * parent_ids, + struct ggml_tensor * tree_sizes, + 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-cpu/ggml-cpu.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp index 00d05b3ee..e1db1f8e3 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -478,10 +478,12 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st // silently compute garbage. return ggml_get_op_params_i32(op, 0) != 1; case GGML_OP_GATED_DELTA_NET: - // The Specla GDN variant (op param 2 == 1) is stateful via HLD and is - // only supported by the CUDA kernel. + // The CPU kernel supports in-place and active-slot recurrence, + // but not tree parents, persistent intermediate storage, raw + // gates, transition journals, or SpecLA state. return ggml_get_op_params_i32(op, 2) != 1 && ggml_get_op_params_i32(op, 10) != 1 && + op->src[6] == nullptr && op->src[7] == nullptr && op->src[11] == nullptr; case GGML_OP_OUT_PROD: return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && diff --git a/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp b/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp index 7850baa62..3c7b37ec2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -4957,7 +4957,14 @@ static bool ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_GATED_LINEAR_ATTN: return true; case GGML_OP_GATED_DELTA_NET: - return op->src[8] == nullptr && op->src[11] == nullptr && + // The SYCL kernel consumes only src[0..5] and writes final state + // into the result tensor. + return op->src[6] == nullptr && + op->src[7] == nullptr && + op->src[8] == nullptr && + op->src[11] == nullptr && + ggml_get_op_params_i32(op, 1) == 0 && + ggml_get_op_params_i32(op, 2) != 1 && ggml_get_op_params_i32(op, 10) != 1; case GGML_OP_SSM_CONV: return op->type == GGML_TYPE_F32 && diff --git a/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 17db9166e..9bfd9eb39 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -15780,9 +15780,13 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm return true; // all inputs are contiguous, see ggml.c case GGML_OP_GATED_DELTA_NET: { - // The Vulkan kernel addresses state by compact sequence row - // and does not consume the physical-slot mapping in src[8]. - if (op->src[8] != nullptr || op->src[11] != nullptr || + // The Vulkan kernel consumes only src[0..5] and writes final + // state into the result tensor. Reject tree, persistent, + // active-slot, and in-place variants. + if (op->src[6] != nullptr || op->src[7] != nullptr || + op->src[8] != nullptr || op->src[11] != nullptr || + ggml_get_op_params_i32(op, 1) != 0 || + ggml_get_op_params_i32(op, 2) == 1 || ggml_get_op_params_i32(op, 10) == 1) { return false; } diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 87d86c6e2..99ea64307 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5697,7 +5697,7 @@ struct ggml_tensor * ggml_flash_attn_sparse( // ggml_paged_attn -struct ggml_tensor * ggml_paged_attn_ext( +static struct ggml_tensor * ggml_paged_attn_ext_impl( struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, @@ -5823,6 +5823,50 @@ struct ggml_tensor * ggml_paged_attn_ext( return result; } +struct ggml_tensor * ggml_paged_attn_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * block_table, + struct ggml_tensor * kv_seq_lens, + struct ggml_tensor * active_slot_ids, + struct ggml_tensor * query_positions, + float scale, + int block_size, + int max_kv_seq_len) { + return ggml_paged_attn_ext_impl( + ctx, q, k, v, block_table, kv_seq_lens, + active_slot_ids, query_positions, scale, block_size, + max_kv_seq_len, NULL, NULL, 0, 0, 0); +} + +struct ggml_tensor * ggml_paged_attn_ext_tree( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * block_table, + struct ggml_tensor * kv_seq_lens, + struct ggml_tensor * active_slot_ids, + struct ggml_tensor * query_positions, + float scale, + int block_size, + int max_kv_seq_len, + struct ggml_tensor * parent_ids, + struct ggml_tensor * tree_sizes, + int tree_scratch_base, + int tree_scratch_stride) { + GGML_ASSERT(parent_ids != NULL); + GGML_ASSERT(parent_ids->ne[0] > 0 && parent_ids->ne[0] <= INT_MAX); + const int tree_width = (int) parent_ids->ne[0]; + return ggml_paged_attn_ext_impl( + ctx, q, k, v, block_table, kv_seq_lens, + active_slot_ids, query_positions, scale, block_size, + max_kv_seq_len, parent_ids, tree_sizes, tree_width, + tree_scratch_base, tree_scratch_stride); +} + // ggml_flash_attn_back struct ggml_tensor * ggml_flash_attn_back( diff --git a/server/src/common/concurrency/chain_spec_shapes.h b/server/src/common/concurrency/chain_spec_shapes.h index b21c17daa..ef8640836 100644 --- a/server/src/common/concurrency/chain_spec_shapes.h +++ b/server/src/common/concurrency/chain_spec_shapes.h @@ -24,7 +24,7 @@ inline int chain_decode_bucket_width(int lanes) { inline DDTree make_chain_verify_tree( const std::vector & draft_tokens) { DDTree tree; - if (draft_tokens.size() <= 1) return tree; + if (draft_tokens.empty()) return tree; tree.n_nodes = static_cast(draft_tokens.size()) - 1; tree.token_ids.assign(draft_tokens.begin() + 1, draft_tokens.end()); diff --git a/server/src/common/dflash_draft_kv.cpp b/server/src/common/dflash_draft_kv.cpp index a9ac623b5..5cb92209f 100644 --- a/server/src/common/dflash_draft_kv.cpp +++ b/server/src/common/dflash_draft_kv.cpp @@ -4,6 +4,7 @@ #include #include #include +#include namespace dflash::common { @@ -352,9 +353,36 @@ 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 * static_cast(n_lanes)) * 1024u * 1024u; + constexpr size_t graph_nodes_per_lane = 4096; + constexpr size_t shared_graph_nodes = 2048; + const size_t n_lanes_size = lane_states.size(); + if (n_lanes_size > + (static_cast(std::numeric_limits::max()) - + shared_graph_nodes) / + graph_nodes_per_lane) { + return false; + } + const int n_lanes = static_cast(n_lanes_size); + const size_t graph_capacity = + graph_nodes_per_lane * n_lanes_size + shared_graph_nodes; + if (graph_capacity > std::numeric_limits::max() / 2) { + return false; + } + const size_t tensor_capacity = 2 * graph_capacity; + const size_t tensor_overhead = ggml_tensor_overhead(); + if (tensor_overhead != 0 && + tensor_capacity > + std::numeric_limits::max() / tensor_overhead) { + return false; + } + const size_t tensor_bytes = tensor_capacity * tensor_overhead; + const size_t graph_bytes = + ggml_graph_overhead_custom(graph_capacity, false); + if (tensor_bytes > + std::numeric_limits::max() - graph_bytes) { + return false; + } + const size_t arena_size = graph_bytes + tensor_bytes; batch.meta_arena.resize(arena_size); ggml_init_params params{}; params.mem_size = batch.meta_arena.size(); @@ -366,7 +394,7 @@ static bool draft_kv_batch_build( return false; } batch.gf = ggml_new_graph_custom( - batch.g_ctx, 4096 * n_lanes + 2048, false); + batch.g_ctx, graph_capacity, false); batch.hidden_by_lane.reserve(static_cast(n_lanes)); for (DraftKvState * state : lane_states) { @@ -413,8 +441,11 @@ static bool draft_kv_batch_build( batch.built_for = &dw; batch.lane_states = lane_states; std::fprintf(stderr, - "[draft-kv-batch] packed backbone ready lanes=%d q_len=%d\n", - n_lanes, dw.block_size); + "[draft-kv-batch] packed backbone ready lanes=%d q_len=%d " + "metadata=%.1f MiB\n", + n_lanes, dw.block_size, + static_cast(batch.meta_arena.size()) / + (1024.0 * 1024.0)); return true; } diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index c1e44da4f..608d20b0c 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -1312,12 +1312,17 @@ static ggml_tensor * build_full_attn_block( 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_tree_parent_ids, paged_tree_sizes, - tree_width, tree_scratch_base, tree_scratch_stride); + ggml_tensor * out = paged_tree + ? ggml_paged_attn_ext_tree( + ctx, q, cache_k, cache_v, paged_block_table, + paged_kv_seq_lens, row_seq_ids, row_positions, kq_scale, + PAGED_BLOCK_SIZE, launch_len, + paged_tree_parent_ids, paged_tree_sizes, + tree_scratch_base, tree_scratch_stride) + : 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); if (dense_token_layout) { out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); } diff --git a/server/test/test_batched_gdn.cpp b/server/test/test_batched_gdn.cpp index 1e639e049..5f845b9e3 100644 --- a/server/test/test_batched_gdn.cpp +++ b/server/test/test_batched_gdn.cpp @@ -48,6 +48,52 @@ constexpr size_t CONV_WEIGHT_ELEMS = static_cast(D_CONV) * CONV_CHANNELS bool check(const char * name, const float * batched, const float * reference, size_t count); +bool test_cpu_gdn_support_matrix(ggml_backend_t backend) { + ggml_init_params params{}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + + constexpr int n_tokens = 2; + ggml_tensor * q = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S_V, N_HEAD, n_tokens, 1); + ggml_tensor * k = ggml_dup_tensor(ctx, q); + ggml_tensor * v = ggml_dup_tensor(ctx, q); + ggml_tensor * g = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, 1, N_HEAD, n_tokens, 1); + ggml_tensor * beta = ggml_dup_tensor(ctx, g); + ggml_tensor * state = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S_V, S_V, N_HEAD, 1); + ggml_tensor * parents = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, n_tokens); + ggml_tensor * persistent = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, S_V, S_V, N_HEAD, n_tokens); + ggml_tensor * active_slots = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, 1); + + ggml_tensor * base = ggml_gated_delta_net( + ctx, q, k, v, g, beta, state); + ggml_tensor * inplace = ggml_gated_delta_net_inplace( + ctx, q, k, v, g, beta, state); + ggml_tensor * active = ggml_gated_delta_net_active_inplace( + ctx, q, k, v, g, beta, state, active_slots); + ggml_tensor * tree = ggml_gated_delta_net_tree( + ctx, q, k, v, g, beta, state, parents); + ggml_tensor * tree_persistent = ggml_gated_delta_net_tree_persist( + ctx, q, k, v, g, beta, state, parents, persistent); + + const bool ok = + ggml_backend_supports_op(backend, base) && + ggml_backend_supports_op(backend, inplace) && + ggml_backend_supports_op(backend, active) && + !ggml_backend_supports_op(backend, tree) && + !ggml_backend_supports_op(backend, tree_persistent); + std::printf("batched gdn CPU support matrix %s\n", ok ? "PASS" : "FAIL"); + ggml_free(ctx); + return ok; +} + void fill_uniform(std::mt19937 & rng, float lo, float hi, std::vector & values) { std::uniform_real_distribution dist(lo, hi); @@ -525,7 +571,8 @@ int main(int argc, char ** argv) { } std::mt19937 rng(20260728); - bool ok = test_gdn_sequential(backend, rng, /*inplace_batched=*/false); + bool ok = !cpu || test_cpu_gdn_support_matrix(backend); + ok = test_gdn_sequential(backend, rng, /*inplace_batched=*/false) && ok; ok = test_gdn_sequential(backend, rng, /*inplace_batched=*/true) && ok; ok = test_gdn_active_slots(backend, rng) && ok; ok = test_conv(backend, rng) && ok; diff --git a/server/test/test_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp index 74dd401ab..8ebf783b1 100644 --- a/server/test/test_chain_spec_shapes.cpp +++ b/server/test/test_chain_spec_shapes.cpp @@ -9,6 +9,17 @@ using namespace dflash::common; static int g_checks = 0; int main() { + const DDTree root_only = make_dspark_chain_tree({10}); + CHECK(root_only.n_nodes == 0); + CHECK((root_only.parents == std::vector{-1})); + CHECK(root_only.child_maps.size() == 1); + CHECK((root_only.visibility == std::vector{1})); + int root_pending = -1; + const int32_t root_posterior[] = {11}; + CHECK((follow_verified_tree( + root_only, root_posterior, root_pending) == std::vector{0})); + CHECK(root_pending == 11); + const std::vector draft = {10, 11, 12, 13}; const DDTree tree = make_chain_verify_tree(draft); CHECK(tree.n_nodes == 3); diff --git a/server/test/test_paged_attention.cpp b/server/test/test_paged_attention.cpp index e283fc3dc..620a8059c 100644 --- a/server/test/test_paged_attention.cpp +++ b/server/test/test_paged_attention.cpp @@ -352,12 +352,14 @@ bool run_case(ggml_backend_t backend, 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, parents, sizes, - tree ? tree->width : 0, - tree ? tree_scratch_base : 0, - tree ? tree->scratch_stride : 0); + ggml_tensor * output = tree + ? ggml_paged_attn_ext_tree( + ctx, q, k, v, table, kv_seq_lens, active, + positions, scale, BLOCK_SIZE, max_kv_seq_len, parents, sizes, + tree_scratch_base, tree->scratch_stride) + : ggml_paged_attn_ext( + ctx, q, k, v, table, kv_seq_lens, active, positions, + scale, BLOCK_SIZE, max_kv_seq_len); ggml_set_output(output); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, output); @@ -483,8 +485,7 @@ 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, - nullptr, nullptr, 0, 0, 0); + 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE, 1); const bool rejected = !ggml_backend_supports_op(backend, output); std::printf("paged attention unlaunchable GQA support %s\n", From 0e82772b037f3dd68e11c1ff01c251fc72c7e87b Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 22 Aug 2026 06:59:53 +0000 Subject: [PATCH 04/11] fix(concurrency): make tree commits fail closed Keep GDN journals inside the result buffer, restore the 10-source tensor layout, validate every tree destination before mutation, and allow selector widths that use the CPU top-k fallback. --- .../deps/llama.cpp/ggml/include/ggml-cuda.h | 20 +++ server/deps/llama.cpp/ggml/include/ggml.h | 12 +- .../llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp | 5 +- .../ggml/src/ggml-cuda/gated_delta_net.cu | 22 ++- .../src/ggml-cuda/gdn-transition-journal.cu | 123 +++++++++++++- .../ggml/src/ggml-metal/ggml-metal-device.m | 2 +- .../ggml/src/ggml-sycl/ggml-sycl.cpp | 1 - .../ggml/src/ggml-vulkan/ggml-vulkan.cpp | 2 +- server/deps/llama.cpp/ggml/src/ggml.c | 50 ++++-- .../src/common/dflash2_selector_validation.h | 8 +- server/src/common/step_graph.h | 2 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 20 +-- server/src/qwen35/qwen35_target_graph.cpp | 15 +- server/test/test_chain_spec_shapes.cpp | 2 +- .../test/test_dflash2_selector_validation.cpp | 20 +-- server/test/test_gdn_transition_journal.cpp | 150 ++++++++++++++++-- server/test/test_paged_attention.cpp | 8 +- server/test/test_seq_engine_contract.cpp | 5 - 18 files changed, 363 insertions(+), 104 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml-cuda.h b/server/deps/llama.cpp/ggml/include/ggml-cuda.h index 1281cd061..31b5e72a1 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -123,6 +123,26 @@ 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); +// Validate every packed-tree destination before changing any cache. Once the +// first kernel launches, a device failure is fatal because fallback cannot +// recover from a partially committed state. +GGML_BACKEND_API bool ggml_backend_cuda_tree_commit_transaction( + struct ggml_tensor * const * caches, + int n_caches, + const struct ggml_tensor * feature_source, + struct ggml_tensor * feature_destination, + const struct ggml_tensor * feature_destination_rows, + 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 * commit_rows, + const struct ggml_tensor * accepted_prefixes, + const struct ggml_tensor * active_slot_ids, + int tree_scratch_base, + int tree_scratch_stride); + // 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 f76f319d4..bda2c7fd5 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 12 +#define GGML_MAX_SRC 10 #define GGML_MAX_N_THREADS 512 #define GGML_MAX_OP_PARAMS 64 @@ -2959,11 +2959,11 @@ extern "C" { // CUDA/HIP fixed-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. A tree-form op may attach the journal when - // its parent table describes one root-inclusive chain. - GGML_API void ggml_gated_delta_net_set_transition_journal( - struct ggml_tensor * tensor, - struct ggml_tensor * journal); + // state-dependent reduction. The returned tensor is a view into the + // GDN result buffer, so no extra source slot is required. + GGML_API struct ggml_tensor * ggml_gated_delta_net_capture_transition_journal( + struct ggml_context * ctx, + struct ggml_tensor * tensor); // dflash extension: let the kernel derive the gates from the raw // projections instead of graph-side sigmoid/softplus ops: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp index e1db1f8e3..b3df79664 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -480,11 +480,10 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st case GGML_OP_GATED_DELTA_NET: // The CPU kernel supports in-place and active-slot recurrence, // but not tree parents, persistent intermediate storage, raw - // gates, transition journals, or SpecLA state. + // gates or SpecLA state. return ggml_get_op_params_i32(op, 2) != 1 && ggml_get_op_params_i32(op, 10) != 1 && - op->src[6] == nullptr && op->src[7] == nullptr && - op->src[11] == nullptr; + op->src[6] == nullptr && op->src[7] == nullptr; case GGML_OP_OUT_PROD: return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; 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 a43cf10fb..9f3e13d19 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 @@ -931,10 +931,6 @@ 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]. The packed tree caller - // uses it only with a root-inclusive linear parent chain. - 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); GGML_TENSOR_LOCALS(int64_t, nek, src_k, ne); @@ -976,8 +972,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 + const int journal_row_offset = ggml_get_op_params_i32(dst, 3); + float * transition_journal_d = journal_row_offset > 0 + ? dst_d + (int64_t) journal_row_offset*dst->ne[0] : nullptr; const bool persist_is_f16 = src_persist_inter && src_persist_inter->type == GGML_TYPE_F16; @@ -1007,14 +1004,13 @@ 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) { + if (transition_journal_d) { const int64_t journal_width = kda ? 3*S_v : 2*S_v + 1; - 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); + const int64_t journal_elements = + journal_width*H*n_tokens*n_seqs; + GGML_ASSERT( + (int64_t) journal_row_offset*dst->ne[0] + + journal_elements <= ggml_nelements(dst)); } // strides in floats (beta strides used for both g and beta offset computation) 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 d9a81a68d..7767e0ff0 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 @@ -169,14 +169,15 @@ bool same_device_pointer(const void * pointer, int expected_device) { } // namespace -extern "C" bool ggml_backend_cuda_gdn_transition_journal_commit_many( +static bool gdn_transition_journal_commit_many_impl( 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) { + const ggml_tensor * active_slot_ids, + bool commit) { if (!journals || !states || !conv_inputs || !conv_states || n_layers <= 0 || !accepted_prefixes || !active_slot_ids || accepted_prefixes->type != GGML_TYPE_I32 || @@ -250,6 +251,7 @@ extern "C" bool ggml_backend_cuda_gdn_transition_journal_commit_many( seen[(size_t) slot] = 1; } + if (!commit) return true; ggml_cuda_set_device(device); constexpr int threads = 256; (void) cudaGetLastError(); @@ -290,13 +292,35 @@ extern "C" bool ggml_backend_cuda_gdn_transition_journal_commit_many( return cudaDeviceSynchronize() == cudaSuccess; } -extern "C" bool ggml_backend_cuda_tree_cache_commit_many( +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 (!gdn_transition_journal_commit_many_impl( + journals, states, conv_inputs, conv_states, n_layers, + accepted_prefixes, active_slot_ids, false)) { + return false; + } + if (!gdn_transition_journal_commit_many_impl( + journals, states, conv_inputs, conv_states, n_layers, + accepted_prefixes, active_slot_ids, true)) { + GGML_ABORT("recurrent commit failed after mutation began"); + } + return true; +} + +static bool tree_cache_commit_many_impl( 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) { + int tree_scratch_stride, + bool commit) { if (!caches || n_caches <= 0 || !commit_rows || !active_slot_ids || commit_rows->type != GGML_TYPE_I64 || active_slot_ids->type != GGML_TYPE_I32 || @@ -342,6 +366,7 @@ extern "C" bool ggml_backend_cuda_tree_cache_commit_many( } } + if (!commit) return true; ggml_cuda_set_device(device); constexpr int threads = 256; (void) cudaGetLastError(); @@ -362,10 +387,31 @@ extern "C" bool ggml_backend_cuda_tree_cache_commit_many( return cudaDeviceSynchronize() == cudaSuccess; } -extern "C" bool ggml_backend_cuda_tree_feature_commit( +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 (!tree_cache_commit_many_impl( + caches, n_caches, commit_rows, active_slot_ids, + tree_scratch_base, tree_scratch_stride, false)) { + return false; + } + if (!tree_cache_commit_many_impl( + caches, n_caches, commit_rows, active_slot_ids, + tree_scratch_base, tree_scratch_stride, true)) { + GGML_ABORT("K/V commit failed after mutation began"); + } + return true; +} + +static bool tree_feature_commit_impl( const ggml_tensor * source, ggml_tensor * destination, - const ggml_tensor * destination_rows) { + const ggml_tensor * destination_rows, + bool commit) { if (!source || !destination || !destination_rows || source->type != destination->type || source->type != GGML_TYPE_BF16 || destination_rows->type != GGML_TYPE_I32 || @@ -387,6 +433,7 @@ extern "C" bool ggml_backend_cuda_tree_feature_commit( for (int row : rows) { if (row < -1 || row >= destination->ne[1]) return false; } + if (!commit) return true; ggml_cuda_set_device(device); constexpr int threads = 256; const dim3 grid( @@ -400,3 +447,67 @@ extern "C" bool ggml_backend_cuda_tree_feature_commit( 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 (!tree_feature_commit_impl( + source, destination, destination_rows, false)) { + return false; + } + if (!tree_feature_commit_impl( + source, destination, destination_rows, true)) { + GGML_ABORT("feature commit failed after mutation began"); + } + return true; +} + +extern "C" bool ggml_backend_cuda_tree_commit_transaction( + ggml_tensor * const * caches, + int n_caches, + const ggml_tensor * feature_source, + ggml_tensor * feature_destination, + const ggml_tensor * feature_destination_rows, + 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 * commit_rows, + const ggml_tensor * accepted_prefixes, + const ggml_tensor * active_slot_ids, + int tree_scratch_base, + int tree_scratch_stride) { + if (!tree_cache_commit_many_impl( + caches, n_caches, commit_rows, active_slot_ids, + tree_scratch_base, tree_scratch_stride, false) || + !tree_feature_commit_impl( + feature_source, feature_destination, + feature_destination_rows, false) || + !gdn_transition_journal_commit_many_impl( + journals, states, conv_inputs, conv_states, n_layers, + accepted_prefixes, active_slot_ids, false)) { + return false; + } + + // No recoverable error may escape after the first destination changes. + // Returning false here would let the caller run fallback from a mixed + // pre-commit/post-commit state. + if (!tree_cache_commit_many_impl( + caches, n_caches, commit_rows, active_slot_ids, + tree_scratch_base, tree_scratch_stride, true)) { + GGML_ABORT("tree commit failed after K/V mutation began"); + } + if (!tree_feature_commit_impl( + feature_source, feature_destination, + feature_destination_rows, true)) { + GGML_ABORT("tree commit failed after feature mutation began"); + } + if (!gdn_transition_journal_commit_many_impl( + journals, states, conv_inputs, conv_states, n_layers, + accepted_prefixes, active_slot_ids, true)) { + GGML_ABORT("tree commit failed after recurrent mutation began"); + } + return true; +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m b/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m index b8d710771..9c0b4cdb3 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m +++ b/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m @@ -1186,7 +1186,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te return true; case GGML_OP_GATED_DELTA_NET: return has_simdgroup_reduction && op->src[2]->ne[0] % 32 == 0 && - op->src[8] == NULL && op->src[11] == NULL && + op->src[8] == NULL && ggml_get_op_params_i32(op, 10) != 1; case GGML_OP_SOLVE_TRI: case GGML_OP_MUL_MAT: diff --git a/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp b/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp index 3c7b37ec2..50038792e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -4962,7 +4962,6 @@ static bool ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, const g return op->src[6] == nullptr && op->src[7] == nullptr && op->src[8] == nullptr && - op->src[11] == nullptr && ggml_get_op_params_i32(op, 1) == 0 && ggml_get_op_params_i32(op, 2) != 1 && ggml_get_op_params_i32(op, 10) != 1; diff --git a/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 9bfd9eb39..84c7f6064 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -15784,7 +15784,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm // state into the result tensor. Reject tree, persistent, // active-slot, and in-place variants. if (op->src[6] != nullptr || op->src[7] != nullptr || - op->src[8] != nullptr || op->src[11] != nullptr || + op->src[8] != nullptr || ggml_get_op_params_i32(op, 1) != 0 || ggml_get_op_params_i32(op, 2) == 1 || ggml_get_op_params_i32(op, 10) == 1) { diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 99ea64307..206cbe7aa 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -6882,6 +6882,7 @@ void ggml_gated_delta_net_set_skip_intermediate( bool skip_intermediate) { GGML_ASSERT(tensor != NULL); GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); + GGML_ASSERT(ggml_get_op_params_i32(tensor, 3) == 0); ggml_set_op_params_i32(tensor, 0, skip_intermediate ? 1 : 0); const struct ggml_tensor * v = tensor->src[2]; @@ -6909,13 +6910,15 @@ 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); +struct ggml_tensor * ggml_gated_delta_net_capture_transition_journal( + struct ggml_context * ctx, + struct ggml_tensor * tensor) { + GGML_ASSERT(ctx != NULL && tensor != NULL); GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); - GGML_ASSERT(journal->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(journal)); + GGML_ASSERT(tensor->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(tensor)); + GGML_ASSERT(ggml_get_op_params_i32(tensor, 3) == 0); + GGML_ASSERT(ggml_get_op_params_i32(tensor, 2) == 0); const struct ggml_tensor * v = tensor->src[2]; const struct ggml_tensor * g = tensor->src[3]; @@ -6926,12 +6929,35 @@ void ggml_gated_delta_net_set_transition_journal( 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; + GGML_ASSERT(journal_width > 0 && H > 0 && n_tokens > 0 && n_seqs > 0); + GGML_ASSERT(journal_width <= INT64_MAX/H); + const int64_t journal_head = journal_width*H; + GGML_ASSERT(journal_head <= INT64_MAX/n_tokens); + const int64_t journal_token = journal_head*n_tokens; + GGML_ASSERT(journal_token <= INT64_MAX/n_seqs); + const int64_t journal_elements = journal_token*n_seqs; + GGML_ASSERT(tensor->ne[0] > 0); + GGML_ASSERT(journal_elements <= INT64_MAX - (tensor->ne[0] - 1)); + const int64_t journal_rows = + (journal_elements + tensor->ne[0] - 1)/tensor->ne[0]; + GGML_ASSERT(tensor->ne[1] > 0 && tensor->ne[1] <= INT32_MAX); + GGML_ASSERT(journal_rows <= INT64_MAX - tensor->ne[1]); + const int32_t journal_row_offset = (int32_t) tensor->ne[1]; + GGML_ASSERT((size_t) journal_row_offset <= SIZE_MAX/tensor->nb[1]); + const size_t journal_byte_offset = + (size_t) journal_row_offset*tensor->nb[1]; + + tensor->ne[1] += journal_rows; + tensor->nb[2] = tensor->nb[1]*tensor->ne[1]; + tensor->nb[3] = tensor->nb[2]*tensor->ne[2]; + ggml_set_op_params_i32(tensor, 3, journal_row_offset); + + return ggml_view_4d( + ctx, tensor, journal_width, H, n_tokens, n_seqs, + (size_t) journal_width*sizeof(float), + (size_t) journal_head*sizeof(float), + (size_t) journal_token*sizeof(float), + journal_byte_offset); } // dflash: raw-gate mode (see ggml.h). [dt_bias | A] -> src[9], diff --git a/server/src/common/dflash2_selector_validation.h b/server/src/common/dflash2_selector_validation.h index 6bcd50403..d78a5fb39 100644 --- a/server/src/common/dflash2_selector_validation.h +++ b/server/src/common/dflash2_selector_validation.h @@ -1,7 +1,5 @@ #pragma once -#include "geometric_draft_topk_cuda.h" - #include #include @@ -33,9 +31,9 @@ inline bool validate_dflash2_selector_layout( 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"; + if (layout.top_k <= 0) { + error = "DFlash 2 selector top_k must be positive (got " + + std::to_string(layout.top_k) + ")"; return false; } if (layout.hproj_rank != layout.rank || diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index affcbace7..dbf342d60 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -78,7 +78,7 @@ struct StepGraph { 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 + ggml_tensor * feature_commit_rows = nullptr; // [n_tokens] 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; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 105fa93d7..9472ca53e 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -871,24 +871,16 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( return result; } - if (!ggml_backend_cuda_tree_cache_commit_many( + if (!ggml_backend_cuda_tree_commit_transaction( caches.data(), static_cast(caches.size()), - graph.commit_rows, graph.commit_slot_ids, - tree_scratch_base_, tree_scratch_stride_)) { - result.error = "fixed chain K/V promotion failed"; - return result; - } - if (!ggml_backend_cuda_tree_feature_commit( graph.tree_features, b_.cache_.target_feat, - graph.feature_commit_rows)) { - result.error = "fixed chain feature promotion failed"; - return result; - } - if (!ggml_backend_cuda_gdn_transition_journal_commit_many( + graph.feature_commit_rows, journals.data(), states.data(), conv_inputs.data(), conv_states.data(), static_cast(n_delta), - graph.accepted_prefixes, graph.commit_slot_ids)) { - result.error = "fixed chain recurrent promotion failed"; + graph.commit_rows, graph.accepted_prefixes, + graph.commit_slot_ids, + tree_scratch_base_, tree_scratch_stride_)) { + result.error = "fixed chain commit preflight failed"; return result; } ggml_backend_tensor_set( diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 608d20b0c..337b80004 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -2050,19 +2050,14 @@ static ggml_tensor * build_delta_net_block( ggml_gated_delta_net_set_raw_gates(result, L.ssm_gate_ba); } } - 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 (can_skip_gdn_intermediate) { ggml_gated_delta_net_set_skip_intermediate(result, true); } + if (seg_cap && seg_tree) { + seg_cap->transition_journal = + ggml_gated_delta_net_capture_transition_journal(ctx, result); + ggml_set_output(seg_cap->transition_journal); + } // Slice output and new_state out of the packed result const int64_t S_v = head_v_dim; diff --git a/server/test/test_chain_spec_shapes.cpp b/server/test/test_chain_spec_shapes.cpp index 8ebf783b1..0f2fa1aae 100644 --- a/server/test/test_chain_spec_shapes.cpp +++ b/server/test/test_chain_spec_shapes.cpp @@ -9,7 +9,7 @@ using namespace dflash::common; static int g_checks = 0; int main() { - const DDTree root_only = make_dspark_chain_tree({10}); + const DDTree root_only = make_chain_verify_tree({10}); CHECK(root_only.n_nodes == 0); CHECK((root_only.parents == std::vector{-1})); CHECK(root_only.child_maps.size() == 1); diff --git a/server/test/test_dflash2_selector_validation.cpp b/server/test/test_dflash2_selector_validation.cpp index bbc46acc8..450cb3bd5 100644 --- a/server/test/test_dflash2_selector_validation.cpp +++ b/server/test/test_dflash2_selector_validation.cpp @@ -28,23 +28,16 @@ int main() { CHECK(validate_dflash2_selector_layout(layout, error)); CHECK(error.empty()); - for (int K = 1; K <= 8; ++K) { + for (int K : {1, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17}) { 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.top_k = 0; + CHECK(!validate_dflash2_selector_layout(layout, error)); + CHECK(error.find("top_k must be positive") != std::string::npos); layout = valid_layout(); layout.succ_vocab--; @@ -73,6 +66,7 @@ int main() { layout.succ_vocab = 8; layout.target_output_vocab = 0; layout.target_declared_vocab = 0; + layout.top_k = 9; CHECK(!validate_dflash2_selector_layout(layout, error)); CHECK(error.find("exceeds codebook vocab") != std::string::npos); diff --git a/server/test/test_gdn_transition_journal.cpp b/server/test/test_gdn_transition_journal.cpp index cc3e82e2d..62a5df4a5 100644 --- a/server/test/test_gdn_transition_journal.cpp +++ b/server/test/test_gdn_transition_journal.cpp @@ -44,15 +44,14 @@ bool test_raw_gate_protocol() { ggml_gated_delta_net_set_raw_gates(result, gate_ba); const int32_t * op_params = reinterpret_cast(result->op_params); - const bool ok = result->src[9] == gate_ba && result->src[10] == nullptr && + const bool ok = result->src[9] == gate_ba && ggml_nelements(result->src[9]) == 2*H && op_params[2] == 0 && op_params[10] == 1; if (!ok) { std::fprintf( stderr, - "raw gate protocol: src9=%p src10=%p elements=%lld op2=%d op10=%d\n", + "raw gate protocol: src9=%p elements=%lld op2=%d op10=%d\n", static_cast(result->src[9]), - static_cast(result->src[10]), (long long) ggml_nelements(result->src[9]), op_params[2], op_params[10]); } @@ -417,8 +416,6 @@ bool run_case(ggml_backend_t backend, bool kda, bool raw_gates) { 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( @@ -447,7 +444,8 @@ bool run_case(ggml_backend_t backend, bool kda, bool raw_gates) { if (raw_gates) { ggml_gated_delta_net_set_raw_gates(result, gate_ba); } - ggml_gated_delta_net_set_transition_journal(result, tensors.journal); + tensors.journal = + ggml_gated_delta_net_capture_transition_journal(tensors.ctx, result); ggml_set_output(result); ggml_cgraph * graph = ggml_new_graph(tensors.ctx); ggml_build_forward_expand(graph, result); @@ -668,8 +666,6 @@ bool run_grouped_tree_case(ggml_backend_t backend) { 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( @@ -683,7 +679,8 @@ bool run_grouped_tree_case(ggml_backend_t backend) { 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_tensor * journal = + ggml_gated_delta_net_capture_transition_journal(ctx, result); ggml_set_output(result); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, result); @@ -770,6 +767,140 @@ bool run_grouped_tree_case(ggml_backend_t backend) { ggml_free(ctx); return ok; } + +bool test_tree_commit_preflight_is_non_mutating(ggml_backend_t backend) { + constexpr int state_size = 16; + constexpr int heads = 1; + constexpr int tokens = 2; + constexpr int sequences = 1; + constexpr int state_slots = 1; + constexpr int conv_window = 2; + constexpr int conv_channels = 3; + + ggml_init_params params{}; + params.mem_size = 256*1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + + ggml_tensor * cache = + ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 4, 8, 1, 1); + ggml_tensor * commit_rows = + ggml_new_tensor_2d(ctx, GGML_TYPE_I64, 1, sequences); + ggml_tensor * active_slots = + ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sequences); + ggml_tensor * feature_source = + ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 4, sequences); + ggml_tensor * feature_destination = + ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 4, 4); + ggml_tensor * feature_rows = + ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sequences); + ggml_tensor * journal = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, 2*state_size + 1, + heads, tokens, sequences); + ggml_tensor * state = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, state_size, state_size, + heads, state_slots); + ggml_tensor * conv_input = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, conv_window + tokens, + conv_channels, sequences, 1); + ggml_tensor * conv_state = ggml_new_tensor_4d( + ctx, GGML_TYPE_F32, conv_window, + conv_channels, state_slots, 1); + ggml_tensor * accepted = + ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sequences); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buffer) { + ggml_free(ctx); + return false; + } + + std::vector cache_before( + (size_t) ggml_nelements(cache), 1.25f); + std::vector feature_before( + (size_t) ggml_nelements(feature_destination), 0x3f80); + std::vector state_before( + (size_t) ggml_nelements(state), -2.5f); + std::vector conv_before( + (size_t) ggml_nelements(conv_state), 3.75f); + std::vector zeros( + (size_t) std::max( + ggml_nelements(journal), ggml_nelements(conv_input)), + 0.0f); + std::vector feature_source_values( + (size_t) ggml_nelements(feature_source), 0x4000); + const int64_t destination_row = 0; + const int32_t active_slot = 0; + const int32_t feature_row = 0; + const int32_t invalid_accepted = tokens + 1; + + ggml_backend_tensor_set( + cache, cache_before.data(), 0, + cache_before.size()*sizeof(float)); + ggml_backend_tensor_set( + feature_destination, feature_before.data(), 0, + feature_before.size()*sizeof(uint16_t)); + ggml_backend_tensor_set( + feature_source, feature_source_values.data(), 0, + feature_source_values.size()*sizeof(uint16_t)); + ggml_backend_tensor_set( + state, state_before.data(), 0, + state_before.size()*sizeof(float)); + ggml_backend_tensor_set( + conv_state, conv_before.data(), 0, + conv_before.size()*sizeof(float)); + ggml_backend_tensor_set( + journal, zeros.data(), 0, + (size_t) ggml_nelements(journal)*sizeof(float)); + ggml_backend_tensor_set( + conv_input, zeros.data(), 0, + (size_t) ggml_nelements(conv_input)*sizeof(float)); + ggml_backend_tensor_set( + commit_rows, &destination_row, 0, sizeof(destination_row)); + ggml_backend_tensor_set( + active_slots, &active_slot, 0, sizeof(active_slot)); + ggml_backend_tensor_set( + feature_rows, &feature_row, 0, sizeof(feature_row)); + ggml_backend_tensor_set( + accepted, &invalid_accepted, 0, sizeof(invalid_accepted)); + + ggml_tensor * caches[] = {cache}; + const ggml_tensor * journals[] = {journal}; + ggml_tensor * states[] = {state}; + const ggml_tensor * conv_inputs[] = {conv_input}; + ggml_tensor * conv_states[] = {conv_state}; + const bool rejected = !ggml_backend_cuda_tree_commit_transaction( + caches, 1, feature_source, feature_destination, feature_rows, + journals, states, conv_inputs, conv_states, 1, + commit_rows, accepted, active_slots, + /*tree_scratch_base=*/4, /*tree_scratch_stride=*/1); + + std::vector cache_after(cache_before.size()); + std::vector feature_after(feature_before.size()); + std::vector state_after(state_before.size()); + std::vector conv_after(conv_before.size()); + ggml_backend_tensor_get( + cache, cache_after.data(), 0, cache_after.size()*sizeof(float)); + ggml_backend_tensor_get( + feature_destination, feature_after.data(), 0, + feature_after.size()*sizeof(uint16_t)); + ggml_backend_tensor_get( + state, state_after.data(), 0, state_after.size()*sizeof(float)); + ggml_backend_tensor_get( + conv_state, conv_after.data(), 0, conv_after.size()*sizeof(float)); + + const bool ok = rejected && + cache_after == cache_before && + feature_after == feature_before && + state_after == state_before && + conv_after == conv_before; + std::printf("gdn tree commit preflight : %s\n", + ok ? "PASS" : "FAIL"); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ok; +} } // namespace int main() { @@ -784,6 +915,7 @@ int main() { ok = run_case(backend, true, false) && ok; ok = run_case(backend, false, true) && ok; ok = run_grouped_tree_case(backend) && ok; + ok = test_tree_commit_preflight_is_non_mutating(backend) && ok; ggml_backend_free(backend); return ok ? 0 : 1; } diff --git a/server/test/test_paged_attention.cpp b/server/test/test_paged_attention.cpp index 620a8059c..45e261f8a 100644 --- a/server/test/test_paged_attention.cpp +++ b/server/test/test_paged_attention.cpp @@ -188,11 +188,13 @@ std::vector reference_attention( ? (seq - tree->ar_rows) % tree->width : -1; int kv_seq_len = clamped_seq_len(test_case, physical_seq); if (query_positions && !tree_query) { - if ((*query_positions)[seq] < 0) continue; + const int32_t query_position = (*query_positions)[seq]; + if (query_position < 0) continue; // The inclusive causal clamp: row seq attends its sequence's // cached tokens [0, position]. - kv_seq_len = std::min( - kv_seq_len, (*query_positions)[seq] + 1); + if (query_position < kv_seq_len) { + kv_seq_len = query_position + 1; + } } const int tree_size = tree_query ? tree->tree_sizes[tree_seq] : 0; if (tree_query && diff --git a/server/test/test_seq_engine_contract.cpp b/server/test/test_seq_engine_contract.cpp index 56f270f41..3b42d9c0e 100644 --- a/server/test/test_seq_engine_contract.cpp +++ b/server/test/test_seq_engine_contract.cpp @@ -23,7 +23,6 @@ struct Faults { bool overconsume_prefill = false; bool drop_second_completion = false; bool retire_leaks = false; - bool burst_when_speculation_disabled = false; }; struct FakeCapabilities { @@ -112,10 +111,6 @@ 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; From accc84aac8910249e680b82e9ad2690d41837a0c Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 22 Aug 2026 07:42:23 +0000 Subject: [PATCH 05/11] concurrency: add service-round observability Add opt-in, bounded records for request lifecycle, scheduler rounds, prefill, KV pressure, and speculative decode. Serve the local dashboard from dflash_server. Generate Markdown and Perfetto files from the same model-neutral JSONL capture. --- .../benchmarks/concurrency/profile_report.py | 339 ++++++ .../concurrency/test_profile_report.py | 75 ++ server/CMakeLists.txt | 24 + server/docs/CONCURRENCY_OBSERVABILITY.md | 92 ++ server/docs/ENVIRONMENT.md | 7 +- server/share/observability.html | 988 ++++++++++++++++++ server/src/common/concurrency/seq_engine.h | 5 +- .../observability/inference_profile.cpp | 108 ++ .../common/observability/inference_profile.h | 158 +++ .../qwen35/concurrency/qwen35_seq_engine.cpp | 340 +++++- .../qwen35/concurrency/qwen35_seq_engine.h | 13 +- server/src/server/http_server.cpp | 68 +- server/src/server/http_server.h | 7 + server/src/server/observability.cpp | 437 ++++++++ server/src/server/observability.h | 135 +++ server/src/server/scheduler.cpp | 139 ++- server/test/test_inference_profile.cpp | 75 ++ server/test/test_observability.cpp | 94 ++ server/test/test_seq_engine_contract.cpp | 4 +- 19 files changed, 3029 insertions(+), 79 deletions(-) create mode 100644 harness/benchmarks/concurrency/profile_report.py create mode 100644 harness/benchmarks/concurrency/test_profile_report.py create mode 100644 server/docs/CONCURRENCY_OBSERVABILITY.md create mode 100644 server/share/observability.html create mode 100644 server/src/common/observability/inference_profile.cpp create mode 100644 server/src/common/observability/inference_profile.h create mode 100644 server/src/server/observability.cpp create mode 100644 server/src/server/observability.h create mode 100644 server/test/test_inference_profile.cpp create mode 100644 server/test/test_observability.cpp diff --git a/harness/benchmarks/concurrency/profile_report.py b/harness/benchmarks/concurrency/profile_report.py new file mode 100644 index 000000000..b06180f58 --- /dev/null +++ b/harness/benchmarks/concurrency/profile_report.py @@ -0,0 +1,339 @@ +#!/usr/bin/env python3 +"""Summarize a Lucebox concurrency profile and emit a Perfetto trace.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +from collections import Counter, defaultdict +from pathlib import Path +from typing import Any, Iterable + + +def load_records(path: Path) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + with path.open(encoding="utf-8") as source: + for line_number, line in enumerate(source, 1): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError(f"{path}:{line_number}: {error.msg}") from error + if not isinstance(record, dict) or "type" not in record: + raise ValueError(f"{path}:{line_number}: record needs a type") + records.append(record) + if not records or records[0].get("schema") != "lucebox.concurrency.v1": + raise ValueError(f"{path}: unsupported or missing profile schema") + return records + + +def ratio(numerator: float, denominator: float) -> float: + return numerator / denominator if denominator else math.nan + + +def percent(value: float) -> str: + return "n/a" if math.isnan(value) else f"{100.0 * value:.1f}%" + + +def percentile(values: Iterable[float], quantile: float) -> float: + ordered = sorted(values) + if not ordered: + return math.nan + index = (len(ordered) - 1) * quantile + low = math.floor(index) + high = math.ceil(index) + if low == high: + return ordered[low] + return ordered[low] * (high - index) + ordered[high] * (index - low) + + +def duration_ms(end: int, start: int) -> float: + return (end - start) / 1_000_000 if end and start and end >= start else math.nan + + +def fmt_ms(value: float) -> str: + return "n/a" if math.isnan(value) else f"{value:.2f} ms" + + +def sum_field(records: Iterable[dict[str, Any]], key: str) -> int: + return sum(int(record.get(key, 0)) for record in records) + + +def build_markdown(records: list[dict[str, Any]]) -> str: + steps = [record for record in records if record["type"] == "step"] + requests = [record for record in records if record["type"] == "request"] + bursts = [record for record in records if record["type"] == "token_burst"] + footer = next( + (record for record in reversed(records) if record["type"] == "footer"), + {}, + ) + failed_requests = sum( + not bool(request.get("ok")) for request in requests + ) + + phases: Counter[str] = Counter() + decisions: Counter[str] = Counter() + cohorts: dict[int, list[dict[str, Any]]] = defaultdict(list) + paths: Counter[str] = Counter() + for step in steps: + cohorts[int(step.get("live_slots", 0))].append(step) + paths[str(step.get("path", "unknown"))] += 1 + for span in step.get("phases", []): + phases[str(span.get("phase", "unknown"))] += int( + span.get("duration_ns", 0) + ) + for lane in step.get("lanes", []): + if lane.get("kind") == "decode": + decisions[str(lane.get("spec", "none"))] += 1 + + queue_ms = [ + duration_ms(int(request.get("admitted_ns", 0)), int(request.get("queued_ns", 0))) + for request in requests + ] + ttft_ms = [ + duration_ms(int(request.get("first_token_ns", 0)), int(request.get("queued_ns", 0))) + for request in requests + ] + e2e_ms = [ + duration_ms(int(request.get("completed_ns", 0)), int(request.get("queued_ns", 0))) + for request in requests + ] + queue_ms = [value for value in queue_ms if not math.isnan(value)] + ttft_ms = [value for value in ttft_ms if not math.isnan(value)] + e2e_ms = [value for value in e2e_ms if not math.isnan(value)] + + burst_times: dict[int, list[tuple[int, int]]] = defaultdict(list) + for burst in bursts: + burst_times[int(burst["request_id"])].append( + (int(burst["ready_ns"]), int(burst.get("token_count", 0))) + ) + inter_token_ms: list[float] = [] + for request_bursts in burst_times.values(): + request_bursts.sort() + for previous, current in zip(request_bursts, request_bursts[1:]): + token_count = max(1, current[1]) + inter_token_ms.append((current[0] - previous[0]) / 1_000_000 / token_count) + + eligible = sum_field(steps, "spec_eligible_lanes") + reserved = sum_field(steps, "spec_reserved_lanes") + attempted = sum_field(steps, "spec_attempted_lanes") + proposed = sum_field(steps, "spec_proposed_draft_tokens") + verified = sum_field(steps, "spec_verified_draft_tokens") + accepted = sum_field(steps, "spec_accepted_draft_tokens") + durable = sum_field(steps, "spec_durable_draft_tokens") + consumed = sum_field(steps, "spec_scheduler_consumed_tokens") + target_rows = sum_field(steps, "target_rows") + target_padding = sum_field(steps, "target_padding_rows") + draft_rows = sum_field(steps, "draft_rows") + draft_padding = sum_field(steps, "draft_padding_rows") + phase_total = sum(phases.values()) + + lines = [ + "# Lucebox concurrency profile", + "", + "## Run summary", + "", + "| Metric | Value |", + "| --- | ---: |", + f"| Captured rounds | {len(steps)} |", + f"| Requests | {len(requests)} |", + f"| Failed requests | {failed_requests} |", + f"| Queue delay p50 / p95 | {fmt_ms(percentile(queue_ms, 0.50))} / {fmt_ms(percentile(queue_ms, 0.95))} |", + f"| TTFT p50 / p95 | {fmt_ms(percentile(ttft_ms, 0.50))} / {fmt_ms(percentile(ttft_ms, 0.95))} |", + f"| End-to-end p50 / p95 | {fmt_ms(percentile(e2e_ms, 0.50))} / {fmt_ms(percentile(e2e_ms, 0.95))} |", + f"| Inter-burst token interval p50 / p95 | {fmt_ms(percentile(inter_token_ms, 0.50))} / {fmt_ms(percentile(inter_token_ms, 0.95))} |", + f"| Target padding | {target_padding} / {target_rows} ({percent(ratio(target_padding, target_rows))}) |", + f"| Draft padding | {draft_padding} / {draft_rows} ({percent(ratio(draft_padding, draft_rows))}) |", + "", + "The inter-burst interval divides each gap by the number of tokens made ready in the later burst. It is a scheduler-level estimate, not a per-token GPU timestamp.", + "", + "## Speculation funnel", + "", + "| Stage | Count | Conversion from previous |", + "| --- | ---: | ---: |", + ] + funnel = [ + ("Eligible lanes", eligible), + ("Reserved lanes", reserved), + ("Attempted lanes", attempted), + ("Proposed draft tokens", proposed), + ("Verified draft tokens", verified), + ("Accepted draft tokens", accepted), + ("Durable draft tokens", durable), + ("Scheduler-consumed draft tokens", consumed), + ] + previous = 0 + for index, (name, value) in enumerate(funnel): + conversion = ( + "n/a" if previous == 0 or index == 3 + else percent(ratio(value, previous)) + ) + lines.append(f"| {name} | {value} | {conversion} |") + previous = value + + lines.extend([ + "", + "### Suppression reasons", + "", + "| Decision | Decode lanes |", + "| --- | ---: |", + ]) + for decision, count in sorted(decisions.items()): + lines.append(f"| `{decision}` | {count} |") + + lines.extend([ + "", + "## Concurrency cohorts", + "", + "| Live slots | Rounds | Mean round | Target padding | Draft acceptance | Paths |", + "| ---: | ---: | ---: | ---: | ---: | --- |", + ]) + for live_slots, cohort in sorted(cohorts.items()): + mean_ms = statistics.fmean(int(step.get("duration_ns", 0)) for step in cohort) / 1_000_000 + cohort_target = sum_field(cohort, "target_rows") + cohort_target_padding = sum_field(cohort, "target_padding_rows") + cohort_proposed = sum_field(cohort, "spec_proposed_draft_tokens") + cohort_accepted = sum_field(cohort, "spec_accepted_draft_tokens") + cohort_paths = Counter(str(step.get("path", "unknown")) for step in cohort) + path_text = ", ".join(f"{name}={count}" for name, count in sorted(cohort_paths.items())) + lines.append( + f"| {live_slots} | {len(cohort)} | {mean_ms:.2f} ms | " + f"{percent(ratio(cohort_target_padding, cohort_target))} | " + f"{percent(ratio(cohort_accepted, cohort_proposed))} | {path_text} |" + ) + + lines.extend([ + "", + "## Phase time", + "", + "| Phase | Total | Share of measured phase time |", + "| --- | ---: | ---: |", + ]) + for phase, nanoseconds in phases.most_common(): + lines.append( + f"| `{phase}` | {nanoseconds / 1_000_000:.2f} ms | " + f"{percent(ratio(nanoseconds, phase_total))} |" + ) + + signals: list[str] = [] + if failed_requests: + signals.append( + f"{failed_requests}/{len(requests)} captured requests failed. " + "Inspect the first incomplete funnel or phase boundary." + ) + if accepted != durable: + signals.append( + "Accepted and durable draft token counts differ. Inspect state " + "promotion or commit before tuning proposal quality." + ) + if target_rows and ratio(target_padding, target_rows) > 0.20: + signals.append("Target graph padding exceeds 20%. Inspect cohort bucket shapes.") + if proposed and ratio(accepted, proposed) < 0.35: + signals.append("Draft acceptance is below 35%. Inspect proposal quality before increasing speculative width.") + if eligible and ratio(attempted, eligible) < 0.75: + signals.append("Fewer than 75% of eligible lanes reach an attempt. Inspect suppression reasons and prompt mixing.") + if requests and percentile(queue_ms, 0.95) > percentile(ttft_ms, 0.95) * 0.40: + signals.append("Queueing accounts for a large part of p95 TTFT. Inspect admission and KV pressure.") + if not signals: + signals.append("No default threshold fired. Use the cohort and phase tables to choose the next experiment.") + + lines.extend(["", "## Signals", ""]) + lines.extend(f"- {signal}" for signal in signals) + lines.extend([ + "", + "## Capture integrity", + "", + f"- Paths: {', '.join(f'{name}={count}' for name, count in sorted(paths.items())) or 'none'}", + f"- Dropped steps: {int(footer.get('dropped_steps', 0))}", + f"- Dropped requests: {int(footer.get('dropped_requests', 0))}", + f"- Dropped token bursts: {int(footer.get('dropped_token_bursts', 0))}", + "", + ]) + return "\n".join(lines) + + +def build_perfetto(records: list[dict[str, Any]]) -> dict[str, Any]: + events: list[dict[str, Any]] = [] + for record in records: + record_type = record["type"] + if record_type == "step": + started_ns = int(record.get("started_ns", 0)) + round_id = int(record.get("round_id", 0)) + for span in record.get("phases", []): + events.append({ + "name": str(span.get("phase", "unknown")), + "cat": "lucebox.round", + "ph": "X", + "pid": 1, + "tid": 1, + "ts": (started_ns + int(span.get("start_offset_ns", 0))) / 1000, + "dur": int(span.get("duration_ns", 0)) / 1000, + "args": { + "round_id": round_id, + "path": record.get("path", "unknown"), + "live_slots": int(record.get("live_slots", 0)), + }, + }) + elif record_type == "request": + request_id = int(record.get("request_id", 0)) + spans = [ + ("queue", int(record.get("queued_ns", 0)), int(record.get("admitted_ns", 0))), + ("prefill", int(record.get("admitted_ns", 0)), int(record.get("prefill_completed_ns", 0))), + ("decode", int(record.get("prefill_completed_ns", 0)), int(record.get("completed_ns", 0))), + ] + for name, start, end in spans: + if start and end >= start: + events.append({ + "name": name, + "cat": "lucebox.request", + "ph": "X", + "pid": 1, + "tid": 1000 + request_id, + "ts": start / 1000, + "dur": (end - start) / 1000, + "args": {"request_id": request_id}, + }) + elif record_type == "token_burst": + events.append({ + "name": "tokens_ready", + "cat": "lucebox.request", + "ph": "i", + "s": "t", + "pid": 1, + "tid": 1000 + int(record.get("request_id", 0)), + "ts": int(record.get("ready_ns", 0)) / 1000, + "args": { + "round_id": int(record.get("round_id", 0)), + "token_count": int(record.get("token_count", 0)), + }, + }) + events.sort(key=lambda event: (event.get("ts", 0), event.get("tid", 0))) + return {"displayTimeUnit": "ms", "traceEvents": events} + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("profile", type=Path) + parser.add_argument("--markdown", type=Path) + parser.add_argument("--perfetto", type=Path) + args = parser.parse_args() + + records = load_records(args.profile) + markdown = build_markdown(records) + if args.markdown: + args.markdown.write_text(markdown, encoding="utf-8") + else: + print(markdown) + if args.perfetto: + args.perfetto.write_text( + json.dumps(build_perfetto(records), indent=2) + "\n", + encoding="utf-8", + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/test_profile_report.py b/harness/benchmarks/concurrency/test_profile_report.py new file mode 100644 index 000000000..f5c989455 --- /dev/null +++ b/harness/benchmarks/concurrency/test_profile_report.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 + +import json +import tempfile +import unittest +from pathlib import Path + +import profile_report + + +class ProfileReportTest(unittest.TestCase): + def records(self): + return [ + {"type": "metadata", "schema": "lucebox.concurrency.v1"}, + { + "type": "step", "round_id": 1, "started_ns": 1_000_000, + "duration_ns": 2_000_000, "path": "speculative", + "live_slots": 4, "target_rows": 20, + "target_padding_rows": 4, "draft_rows": 16, + "draft_padding_rows": 0, "spec_eligible_lanes": 4, + "spec_reserved_lanes": 4, "spec_attempted_lanes": 4, + "spec_proposed_draft_tokens": 12, + "spec_verified_draft_tokens": 12, + "spec_accepted_draft_tokens": 8, + "spec_durable_draft_tokens": 8, + "spec_scheduler_consumed_tokens": 7, + "lanes": [{"kind": "decode", "spec": "selected"}], + "phases": [{"phase": "target_compute", + "start_offset_ns": 100, "duration_ns": 1000}], + }, + { + "type": "request", "request_id": 9, "ok": True, + "queued_ns": 100, "admitted_ns": 200, + "prefill_completed_ns": 500, "first_token_ns": 500, + "completed_ns": 1000, + }, + {"type": "token_burst", "request_id": 9, "round_id": 1, + "ready_ns": 500, "token_count": 1}, + {"type": "token_burst", "request_id": 9, "round_id": 2, + "ready_ns": 900, "token_count": 2}, + {"type": "footer", "dropped_steps": 0, + "dropped_requests": 0, "dropped_token_bursts": 0}, + ] + + def test_markdown_and_perfetto_share_the_records(self): + records = self.records() + markdown = profile_report.build_markdown(records) + self.assertIn("## Speculation funnel", markdown) + self.assertIn("| 4 | 1 | 2.00 ms | 20.0% | 66.7%", markdown) + trace = profile_report.build_perfetto(records) + names = [event["name"] for event in trace["traceEvents"]] + self.assertIn("target_compute", names) + self.assertIn("queue", names) + self.assertIn("tokens_ready", names) + + def test_loader_rejects_a_different_schema(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "profile.jsonl" + path.write_text(json.dumps({"type": "metadata", "schema": "bad"}) + "\n") + with self.assertRaisesRegex(ValueError, "unsupported"): + profile_report.load_records(path) + + def test_failure_and_durability_gap_are_actionable(self): + records = self.records() + records[1]["spec_durable_draft_tokens"] = 0 + records[2]["ok"] = False + + markdown = profile_report.build_markdown(records) + + self.assertIn("1/1 captured requests failed", markdown) + self.assertIn("Accepted and durable draft token counts differ", markdown) + + +if __name__ == "__main__": + unittest.main() diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 85dddf537..f4ec351e6 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -408,6 +408,7 @@ set(DFLASH27B_SRC_INCLUDE_DIRS ) add_library(dflash_common STATIC + src/common/observability/inference_profile.cpp src/errors.cpp src/qwen35/gguf_target_loader.cpp src/qwen35/qwen35_target_graph.cpp @@ -1436,6 +1437,23 @@ 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_inference_profile.cpp") + add_executable(test_inference_profile + test/test_inference_profile.cpp + src/common/observability/inference_profile.cpp) + target_include_directories(test_inference_profile PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_inference_profile) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_observability.cpp") + add_executable(test_observability + test/test_observability.cpp + src/server/observability.cpp + src/common/observability/inference_profile.cpp) + target_include_directories(test_observability PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_observability) + 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 @@ -1623,6 +1641,7 @@ if(DFLASH27B_TESTS) target_sources(test_server_unit PRIVATE src/server/http_server.cpp src/server/scheduler.cpp + src/server/observability.cpp src/server/model_card.cpp src/server/prompt_normalize.cpp src/qwen3/anchor_scan.cpp) @@ -2005,6 +2024,7 @@ if(DFLASH27B_SERVER) src/server/server_main.cpp src/server/http_server.cpp src/server/scheduler.cpp + src/server/observability.cpp src/server/model_card.cpp src/server/prompt_normalize.cpp ) @@ -2048,9 +2068,13 @@ if(DFLASH27B_SERVER) COMMAND ${CMAKE_COMMAND} -E copy_if_different "${CMAKE_CURRENT_SOURCE_DIR}/share/status.html" "$/share/status.html" + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${CMAKE_CURRENT_SOURCE_DIR}/share/observability.html" + "$/share/observability.html" COMMENT "Copying status.html to build/share/" ) install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/share/status.html" + "${CMAKE_CURRENT_SOURCE_DIR}/share/observability.html" DESTINATION share) endif() endif() diff --git a/server/docs/CONCURRENCY_OBSERVABILITY.md b/server/docs/CONCURRENCY_OBSERVABILITY.md new file mode 100644 index 000000000..379f52c08 --- /dev/null +++ b/server/docs/CONCURRENCY_OBSERVABILITY.md @@ -0,0 +1,92 @@ +# Concurrency observability + +Use the concurrency profiler to find scheduler, prefill, speculative decode, +padding, and KV pressure bottlenecks. It captures high-level serving phases. +It does not replace a kernel profiler. + +## Start a capture + +Set `DFLASH_PROF=concurrency` before starting `dflash_server`. + +```bash +DFLASH_PROF=concurrency \ +DFLASH_PROF_OUT=/tmp/lucebox-profile.jsonl \ +./dflash_server +``` + +Stop the server normally to write the JSONL file. Capture is bounded by the +`DFLASH_PROF_MAX_*` variables documented in +[ENVIRONMENT.md](ENVIRONMENT.md). The footer reports dropped records when a +bound is reached. + +## Inspect a live server + +The server exposes two read-only routes: + +- `/observability` serves the built-in Lucebox dashboard. +- `/observability/snapshot` returns one low-cardinality JSON snapshot. + +The dashboard derives ratios in the browser. The prefill ratio compares +executed prompt tokens with the scheduler's offered service budget. The server +publishes raw counts so the live and offline views use the same facts. + +When capture is disabled, the routes remain available and report that state. +The inference path passes a null profile pointer. Phase scopes do not read a +clock, allocate, lock, format, or write in this state. + +## Build an offline report + +Generate a Markdown summary and a Perfetto trace from the same JSONL capture. + +```bash +python3 harness/benchmarks/concurrency/profile_report.py \ + /tmp/lucebox-profile.jsonl \ + --markdown /tmp/lucebox-profile.md \ + --perfetto /tmp/lucebox-profile.perfetto.json +``` + +Open the Perfetto JSON at [ui.perfetto.dev](https://ui.perfetto.dev). It shows +round phase spans, request queue/prefill/decode spans, and token-ready bursts. + +## Read the speculation funnel + +The profiler keeps these stages separate: + +1. An eligible lane can use the configured speculative path. +2. A reserved lane is selected for this service round. +3. An attempted lane enters draft preparation. +4. Proposed draft children come from the drafter. The root token is excluded. +5. Verified draft children run through the target model. +6. Accepted draft children pass verification. +7. Durable draft children finish KV and recurrent-state promotion. +8. Scheduler-consumed draft children reach request generation state. + +The separately reported pending token is sampled after the accepted path. It +is not a draft child and does not inflate acceptance. + +Per-lane decisions explain why an eligible-looking decode did not speculate. +Examples include prompt work in the same round, unsupported sampling, caller +policy, insufficient context, unavailable features, and draft preparation +failure. + +## Understand phase timing + +Phase spans use the host steady clock around existing high-level operations. +Instrumentation does not add device synchronization. A target compute or +readback span therefore reflects the synchronization behavior already present +in that code path. + +Use the built-in profile to choose the next experiment. Use ROCTX and rocprof +afterward when the question becomes kernel scheduling, memory bandwidth, or a +specific device operation. + +## Extend another model + +The scheduler owns request IDs, lifecycle times, planned lanes, and the final +consumption count. A model adapter receives an optional `StepProfile *` and +fills only facts it owns, such as executed rows, padding, phase spans, KV +pressure, and speculation progress. + +The contract is model-neutral and fixed-capacity. A future non-batched C=1 +adapter can populate the same record without changing the report, metrics, or +dashboard. diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index da7d8bb09..6770cc6a3 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -18,7 +18,12 @@ consolidation of this list into CLI flags is tracked as follow-up work. |---|---|---| | `DFLASH_DRAFT_KV` | 1 | KILL SWITCH (remove after burn-in): =0 restores the legacy per-step drafter window recompute instead of the ring cache. | | `DFLASH_LAGUNA_SWA_RING` | 1 | KILL SWITCH (remove after burn-in): =0 keeps SWA layers on pool-sized caches under KVFlash. | -| `DFLASH_PROF` | unset | DEBUG: comma list of profilers (step,verify,prefill). Replaces DFLASH_LAGUNA_{STEP,VERIFY,PREFILL}_PROF. | +| `DFLASH_PROF` | unset | DEBUG: comma list of profilers (`step`, `verify`, `prefill`, `concurrency`). Replaces DFLASH_LAGUNA_{STEP,VERIFY,PREFILL}_PROF. | +| `DFLASH_PROF_OUT` | `concurrency-profile.jsonl` | Output path for the bounded concurrency profile when `DFLASH_PROF=concurrency`. | +| `DFLASH_PROF_WARMUP_ROUNDS` | `0` | Concurrency rounds to omit from the offline capture. Live totals still include them. | +| `DFLASH_PROF_MAX_ROUNDS` | `10000` | Maximum concurrency round records retained for shutdown export. | +| `DFLASH_PROF_MAX_REQUESTS` | `4096` | Maximum request lifecycle records retained for shutdown export. | +| `DFLASH_PROF_MAX_TOKEN_BURSTS` | `200000` | Maximum scheduler token-burst records retained for shutdown export. | | `GGML_CUDA_GRAPH_STATS` | unset | DEBUG: per-graph CUDA-graph replay/capture/eager counters. | | `GGML_CUDA_GRAPH_STATS_EVERY` | 200 | DEBUG: print period for the stats above (clamped to >=1). | | `DFLASH_ADAPTIVE_K_TAU` | 0 = off | Prefer the CLI: --adaptive-experts [tau]. Cumulative combine-weight threshold for per-token expert gating. | diff --git a/server/share/observability.html b/server/share/observability.html new file mode 100644 index 000000000..c6f42f50f --- /dev/null +++ b/server/share/observability.html @@ -0,0 +1,988 @@ + + + + + +Lucebox Observability + + + +
+
+
+
Lucebox telemetry
+

Concurrency observability

+

Waiting for the first scheduler snapshot.

+
+
+ connecting + no sample + schema unknown +
+
+ + + +
+
+
+

Live counters

+
round unknown
+
+
+
+
Rounds
+
0
+
+
+
Queue depth
+
0
+
+
+
Live slots
+
0
+
+
+
Decode lanes
+
0
+
+
+
+ +
+
+

Concurrency and KV

+
0 free
+
+
+
+
+
Latest round KV used
+
0%
+
+ +
+ Total blocks + 0 +
+
+ Free blocks + 0 +
+
+
+
+
Prefill budget used
+
0%
+
+ +
+ Budget tokens + 0 +
+
+ Executed tokens + 0 +
+
+
+
+ +
+
+

Pipeline phases

+
0 ms total
+
+
+
+
+
+ +
+
+

Speculation funnel

+
0% accepted
+
+
+
+
+
+ +
+
+

Utilization and padding

+
0 padded rows
+
+
+
+
Target padding
+
+ +
+
0%
+
+
+
Draft padding
+
+ +
+
0%
+
+
+
Speculative token share
+
+ +
+
0
+
+
+ Requests completed + 0 +
+
+ Requests failed + 0 +
+
+ Dropped steps + 0 +
+
+ Dropped requests + 0 +
+
+
+ +
+
+

Last step

+
duration unknown
+
+
+ + + + + + + + + + + + + + + + + + + +
RoundPathDurationLive slotsQueue
unknownunknownunknown00
+
+
+
+
+ + + + diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 0e22269b4..99395c461 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -55,6 +55,7 @@ #include #include "common/sampler.h" +#include "common/observability/inference_profile.h" namespace dflash::common { @@ -247,7 +248,9 @@ class SeqEngine { // 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; + virtual StepResult step( + const StepPlan & plan, + observability::StepProfile * profile = nullptr) = 0; // Release a slot's KV blocks and mark it free. Safe on failed slots. virtual void retire(int slot) = 0; diff --git a/server/src/common/observability/inference_profile.cpp b/server/src/common/observability/inference_profile.cpp new file mode 100644 index 000000000..ba6aa1cca --- /dev/null +++ b/server/src/common/observability/inference_profile.cpp @@ -0,0 +1,108 @@ +#include "common/observability/inference_profile.h" + +#include + +namespace dflash::common::observability { + +LaneProfile * StepProfile::add_lane(const LaneProfile & lane) noexcept { + if (lane_count >= lanes.size()) { + ++dropped_lanes; + return nullptr; + } + lanes[lane_count] = lane; + return &lanes[lane_count++]; +} + +LaneProfile * StepProfile::find_lane(int32_t slot, LaneKind kind) noexcept { + for (uint32_t i = 0; i < lane_count; ++i) { + if (lanes[i].slot == slot && lanes[i].kind == kind) { + return &lanes[i]; + } + } + return nullptr; +} + +void StepProfile::add_phase(PhaseSpan span) noexcept { + if (phase_count >= phases.size()) { + ++dropped_phases; + return; + } + phases[phase_count++] = span; +} + +const char * step_path_name(StepPath path) noexcept { + switch (path) { + case StepPath::Unknown: return "unknown"; + case StepPath::Packed: return "packed"; + case StepPath::Speculative: return "speculative"; + } + return "unknown"; +} + +const char * lane_kind_name(LaneKind kind) noexcept { + switch (kind) { + case LaneKind::Decode: return "decode"; + case LaneKind::Prefill: return "prefill"; + } + return "unknown"; +} + +const char * spec_decision_name(SpecDecision decision) noexcept { + switch (decision) { + case SpecDecision::None: return "none"; + case SpecDecision::Selected: return "selected"; + case SpecDecision::PromptWorkPresent: return "prompt_work_present"; + case SpecDecision::CallerDisallowed: return "caller_disallowed"; + case SpecDecision::FeatureUnavailable: return "feature_unavailable"; + case SpecDecision::SamplingUnsupported: return "sampling_unsupported"; + case SpecDecision::InsufficientContext: return "insufficient_context"; + case SpecDecision::DraftPrepareFailed: return "draft_prepare_failed"; + } + return "unknown"; +} + +const char * phase_name(Phase phase) noexcept { + switch (phase) { + case Phase::SchedulerPlan: return "scheduler_plan"; + case Phase::InputStaging: return "input_staging"; + case Phase::DraftPrepare: return "draft_prepare"; + case Phase::DraftCompute: return "draft_compute"; + case Phase::ProposalSelect: return "proposal_select"; + case Phase::TargetGraphBuild: return "target_graph_build"; + case Phase::MetadataUpload: return "metadata_upload"; + case Phase::TargetCompute: return "target_compute"; + case Phase::ReadbackSync: return "readback_sync"; + case Phase::Acceptance: return "acceptance"; + case Phase::StatePromotion: return "state_promotion"; + case Phase::SamplingCommit: return "sampling_commit"; + case Phase::OutputProcessing: return "output_processing"; + case Phase::ClientFlush: return "client_flush"; + } + return "unknown"; +} + +uint64_t steady_time_ns() noexcept { + const auto elapsed = std::chrono::steady_clock::now().time_since_epoch(); + return static_cast( + std::chrono::duration_cast(elapsed).count()); +} + +PhaseScope::PhaseScope( + StepProfile * profile, Phase phase, ProfileClock clock) noexcept + : profile_(profile), phase_(phase), clock_(profile ? clock : nullptr) { + if (clock_) started_ns_ = clock_(); +} + +PhaseScope::~PhaseScope() { + if (!profile_ || !clock_) return; + const uint64_t finished_ns = clock_(); + profile_->add_phase({ + phase_, + started_ns_ >= profile_->started_ns + ? started_ns_ - profile_->started_ns + : 0, + finished_ns >= started_ns_ ? finished_ns - started_ns_ : 0, + }); +} + +} diff --git a/server/src/common/observability/inference_profile.h b/server/src/common/observability/inference_profile.h new file mode 100644 index 000000000..17275bd73 --- /dev/null +++ b/server/src/common/observability/inference_profile.h @@ -0,0 +1,158 @@ +#pragma once + +#include +#include +#include + +namespace dflash::common::observability { + +inline constexpr uint32_t kProfileSchemaVersion = 1; +inline constexpr size_t kMaxProfileLanes = 64; +inline constexpr size_t kMaxProfilePhases = 32; +inline constexpr size_t kMaxSpecPositions = 64; + +enum class StepPath : uint8_t { + Unknown, + Packed, + Speculative, +}; + +enum class LaneKind : uint8_t { + Decode, + Prefill, +}; + +enum class SpecDecision : uint8_t { + None, + Selected, + PromptWorkPresent, + CallerDisallowed, + FeatureUnavailable, + SamplingUnsupported, + InsufficientContext, + DraftPrepareFailed, +}; + +enum class Phase : uint8_t { + SchedulerPlan, + InputStaging, + DraftPrepare, + DraftCompute, + ProposalSelect, + TargetGraphBuild, + MetadataUpload, + TargetCompute, + ReadbackSync, + Acceptance, + StatePromotion, + SamplingCommit, + OutputProcessing, + ClientFlush, +}; + +inline constexpr size_t kPhaseCount = + static_cast(Phase::ClientFlush) + 1; + +struct PhaseSpan { + Phase phase = Phase::SchedulerPlan; + uint64_t start_offset_ns = 0; + uint64_t duration_ns = 0; +}; + +struct LaneProfile { + uint64_t request_id = 0; + int32_t slot = -1; + LaneKind kind = LaneKind::Decode; + SpecDecision spec = SpecDecision::None; + uint32_t context_tokens = 0; + uint32_t requested_prefill_tokens = 0; + uint32_t executed_prefill_tokens = 0; + uint32_t proposed_draft_tokens = 0; + uint32_t verified_draft_tokens = 0; + uint32_t accepted_draft_tokens = 0; + uint32_t durable_draft_tokens = 0; + uint32_t scheduler_consumed_tokens = 0; + bool pending_token_sampled = false; + bool pending_token_consumed = false; +}; + +struct StepProfile { + uint32_t schema_version = kProfileSchemaVersion; + uint64_t round_id = 0; + uint64_t started_ns = 0; + uint64_t duration_ns = 0; + StepPath path = StepPath::Unknown; + bool ok = true; + + uint32_t queue_depth = 0; + uint32_t live_slots = 0; + uint32_t planned_decode_lanes = 0; + uint32_t planned_prefill_lanes = 0; + uint32_t planned_prefill_tokens = 0; + uint32_t executed_decode_lanes = 0; + uint32_t executed_prefill_lanes = 0; + uint32_t executed_prefill_tokens = 0; + + uint32_t spec_eligible_lanes = 0; + uint32_t spec_reserved_lanes = 0; + uint32_t spec_attempted_lanes = 0; + uint32_t spec_proposed_draft_tokens = 0; + uint32_t spec_verified_draft_tokens = 0; + uint32_t spec_accepted_draft_tokens = 0; + uint32_t spec_pending_tokens = 0; + uint32_t spec_durable_draft_tokens = 0; + uint32_t spec_scheduler_consumed_tokens = 0; + + uint32_t target_rows = 0; + uint32_t target_padding_rows = 0; + uint32_t draft_rows = 0; + uint32_t draft_padding_rows = 0; + uint32_t decode_bucket = 0; + uint32_t draft_bucket = 0; + uint32_t max_kv_len = 0; + uint32_t kv_blocks_total = 0; + uint32_t kv_blocks_free_before = 0; + uint32_t kv_blocks_free_after = 0; + uint32_t active_sequences = 0; + uint32_t target_forwards = 0; + uint32_t draft_forwards = 0; + + std::array proposed_by_position{}; + std::array accepted_by_position{}; + std::array lanes{}; + std::array phases{}; + uint32_t lane_count = 0; + uint32_t phase_count = 0; + uint32_t dropped_lanes = 0; + uint32_t dropped_phases = 0; + + LaneProfile * add_lane(const LaneProfile & lane) noexcept; + LaneProfile * find_lane(int32_t slot, LaneKind kind) noexcept; + void add_phase(PhaseSpan span) noexcept; +}; + +const char * step_path_name(StepPath path) noexcept; +const char * lane_kind_name(LaneKind kind) noexcept; +const char * spec_decision_name(SpecDecision decision) noexcept; +const char * phase_name(Phase phase) noexcept; + +uint64_t steady_time_ns() noexcept; +using ProfileClock = uint64_t (*)() noexcept; + +class PhaseScope final { +public: + PhaseScope(StepProfile * profile, Phase phase, + ProfileClock clock = steady_time_ns) noexcept; + ~PhaseScope(); + + PhaseScope(const PhaseScope &) = delete; + PhaseScope & operator=(const PhaseScope &) = delete; + +private: + StepProfile * profile_ = nullptr; + Phase phase_ = Phase::SchedulerPlan; + ProfileClock clock_ = nullptr; + uint64_t started_ns_ = 0; +}; + +} diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 9472ca53e..29d7b9c64 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -51,7 +51,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( 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), + prefill_quantum_(std::max(1, prefill_quantum)), pool_(pool), b_(backend), slots_(pool, max_ctx), scratch_row_(scratch_row), tree_width_(tree_width), tree_scratch_base_(tree_scratch_base), tree_scratch_stride_(tree_scratch_stride) { @@ -152,14 +152,29 @@ DraftKvState * Qwen35SeqEngine::ensure_slot_draft_kv(int slot) { bool Qwen35SeqEngine::chain_spec_input_capable( const StepInput & input) const { - if (!capture_features_ || !input.allow_speculation || - input.slot < 0 || input.slot >= slots_.slot_count()) { - return false; + return chain_spec_decision(input) == + observability::SpecDecision::Selected; +} + +observability::SpecDecision Qwen35SeqEngine::chain_spec_decision( + const StepInput & input) const { + using Decision = observability::SpecDecision; + if (!capture_features_ || tree_width_ <= 1) { + return Decision::FeatureUnavailable; + } + if (!input.allow_speculation) return Decision::CallerDisallowed; + if (input.slot < 0 || input.slot >= slots_.slot_count()) { + return Decision::InsufficientContext; } const Qwen35Slot & slot = slots_.slot(input.slot); - return slot.decoding() && !slot.sampler.needs_logit_processing() && - slot.cur_pos >= 1 && - slot.cur_pos + tree_width_ <= slots_.max_context(); + if (!slot.decoding() || slot.cur_pos < 1 || + slot.cur_pos + tree_width_ > slots_.max_context()) { + return Decision::InsufficientContext; + } + if (slot.sampler.needs_logit_processing()) { + return Decision::SamplingUnsupported; + } + return Decision::Selected; } Qwen35SeqEngine::FixedServiceRound @@ -176,7 +191,10 @@ Qwen35SeqEngine::make_fixed_service_round(const StepPlan & plan) const { bool Qwen35SeqEngine::prepare_chain_drafts( const std::vector & inputs, - const std::vector & selected) { + const std::vector & selected, + observability::StepProfile * profile) { + const uint64_t prepare_started_ns = + profile ? observability::steady_time_ns() : 0; if (selected.size() != inputs.size() || !capture_features_ || tree_width_ <= 1 || tree_width_ != b_.dw_.block_size) { return false; @@ -293,12 +311,33 @@ bool Qwen35SeqEngine::prepare_chain_drafts( roots.push_back(lanes.front().root); } + if (profile) { + const uint64_t prepare_finished_ns = observability::steady_time_ns(); + profile->add_phase({ + observability::Phase::DraftPrepare, + prepare_started_ns >= profile->started_ns + ? prepare_started_ns - profile->started_ns : 0, + prepare_finished_ns >= prepare_started_ns + ? prepare_finished_ns - prepare_started_ns : 0, + }); + profile->draft_bucket = static_cast(bucket); + profile->draft_rows = static_cast(tree_width_ * bucket); + profile->draft_padding_rows = static_cast( + tree_width_ * (bucket - static_cast(lanes.size()))); + profile->draft_forwards = 1; + } + std::vector> hidden_blocks; std::vector> proposals; - if (!draft_kv_batch_compute( + bool draft_ok = false; + { + observability::PhaseScope phase( + profile, observability::Phase::DraftCompute); + draft_ok = draft_kv_batch_compute( batch_draft_graph_, b_.dw_, b_.draft_backend_, - batch_states, hidden_blocks) || - hidden_blocks.size() != batch_states.size()) { + batch_states, hidden_blocks); + } + if (!draft_ok || hidden_blocks.size() != batch_states.size()) { reset_lanes(); return false; } @@ -307,10 +346,15 @@ bool Qwen35SeqEngine::prepare_chain_drafts( for (const std::vector & block : hidden_blocks) { hidden_by_lane.push_back(block.data()); } - if (!dflash2_select_chains_batched( + bool selected_ok = false; + { + observability::PhaseScope phase( + profile, observability::Phase::ProposalSelect); + selected_ok = dflash2_select_chains_batched( b_.dw_, b_.draft_backend_, b_.w_.output, hidden_by_lane, - tree_width_, roots, proposals) || - proposals.size() != batch_states.size()) { + tree_width_, roots, proposals); + } + if (!selected_ok || proposals.size() != batch_states.size()) { reset_lanes(); return false; } @@ -329,6 +373,22 @@ bool Qwen35SeqEngine::prepare_chain_drafts( prepared.generated = slots_.slot(lane.slot).generated_tokens(); prepared.root = lane.root; prepared.tokens = std::move(tokens); + if (profile) { + auto * lane_profile = profile->find_lane( + lane.slot, observability::LaneKind::Decode); + if (lane_profile) { + lane_profile->proposed_draft_tokens = + static_cast(tree_width_ - 1); + } + profile->spec_proposed_draft_tokens += + static_cast(tree_width_ - 1); + for (int position = 1; + position < tree_width_ && + position < static_cast(observability::kMaxSpecPositions); + ++position) { + ++profile->proposed_by_position[static_cast(position)]; + } + } } return true; } @@ -465,7 +525,8 @@ Qwen35SeqEngine::PrefillStage Qwen35SeqEngine::stage_prefill_chunk( } SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( - const StepPlan & plan, const std::vector & selected) { + const StepPlan & plan, const std::vector & selected, + observability::StepProfile * profile) { StepResult result; const std::vector & inputs = plan.decode; if (!plan.prefills.empty() || selected.size() != inputs.size()) { @@ -500,6 +561,11 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( const char * value = std::getenv("DFLASH_MIN_TOKENS"); return value ? std::max(0, std::atoi(value)) : 0; }(); + if (profile) { + profile->path = observability::StepPath::Speculative; + profile->executed_decode_lanes = + static_cast(inputs.size()); + } std::vector proposals; std::vector ar_lanes; @@ -575,6 +641,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( const int tree_bucket = chain_decode_bucket_width(spec_count); const int tree_rows_count = tree_width * tree_bucket; const int total_rows = ar_count + tree_rows_count; + if (profile) { + profile->target_rows = static_cast(total_rows); + profile->target_padding_rows = static_cast( + tree_width * (tree_bucket - spec_count)); + profile->decode_bucket = static_cast(tree_bucket); + profile->target_forwards = 1; + } int max_prefix = 1; for (const Proposal & proposal : proposals) { @@ -586,11 +659,17 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( } StepGraph & graph = b_.sg_; - if (!build_target_step_paged_tree( + bool graph_built = false; + { + observability::PhaseScope phase( + profile, observability::Phase::TargetGraphBuild); + graph_built = build_target_step_paged_tree( graph, b_.w_, b_.cache_, b_.target_backend_, tree_width, tree_bucket, max_prefix, tree_scratch_base_, tree_scratch_stride_, - b_.cfg_.kq_stride_pad, ar_count)) { + b_.cfg_.kq_stride_pad, ar_count); + } + if (!graph_built) { result.error = "fixed chain target graph build failed"; return result; } @@ -672,6 +751,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( } } + { + observability::PhaseScope phase( + profile, observability::Phase::MetadataUpload); if (!b_.w_.embedder.embed( tokens.data(), total_rows, embeddings.data())) { result.error = "fixed chain embedding failed"; @@ -711,18 +793,43 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( 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_, graph.gf) != - GGML_STATUS_SUCCESS) { + } + ggml_status target_status = GGML_STATUS_FAILED; + { + observability::PhaseScope phase( + profile, observability::Phase::TargetCompute); + target_status = ggml_backend_graph_compute( + b_.target_backend_, graph.gf); + } + if (target_status != GGML_STATUS_SUCCESS) { result.error = "fixed chain target compute failed"; return result; } std::vector posterior( static_cast(total_rows), -1); - ggml_backend_tensor_get( - graph.argmax_tokens, posterior.data(), 0, - sizeof(int32_t) * posterior.size()); + { + observability::PhaseScope phase( + profile, observability::Phase::ReadbackSync); + ggml_backend_tensor_get( + graph.argmax_tokens, posterior.data(), 0, + sizeof(int32_t) * posterior.size()); + } + if (profile) { + profile->spec_verified_draft_tokens = + static_cast(spec_count * (tree_width - 1)); + for (const Proposal & proposal : proposals) { + if (auto * lane = profile->find_lane( + proposal.slot, observability::LaneKind::Decode)) { + lane->verified_draft_tokens = + static_cast(tree_width - 1); + } + } + } + { + observability::PhaseScope phase( + profile, observability::Phase::Acceptance); for (int lane_index = 0; lane_index < spec_count; ++lane_index) { Proposal & proposal = proposals[static_cast(lane_index)]; const int row_base = ar_count + lane_index * tree_width; @@ -758,7 +865,28 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( return result; } } + } + if (profile) { + for (const Proposal & proposal : proposals) { + const uint32_t accepted = proposal.path.size() > 1 + ? static_cast(proposal.path.size() - 1) : 0; + profile->spec_accepted_draft_tokens += accepted; + if (auto * lane = profile->find_lane( + proposal.slot, observability::LaneKind::Decode)) { + lane->accepted_draft_tokens = accepted; + } + for (uint32_t position = 1; + position <= accepted && + position < observability::kMaxSpecPositions; + ++position) { + ++profile->accepted_by_position[position]; + } + } + } + { + observability::PhaseScope phase( + profile, observability::Phase::StatePromotion); std::vector accepted_prefixes( static_cast(tree_bucket), 0); std::vector commit_slots( @@ -890,6 +1018,21 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( for (const StepInput & input : inputs) { slots_.commit_step(input.slot); } + } + if (profile) { + for (const Proposal & proposal : proposals) { + const uint32_t durable = proposal.path.size() > 1 + ? static_cast(proposal.path.size() - 1) : 0; + profile->spec_durable_draft_tokens += durable; + if (auto * lane = profile->find_lane( + proposal.slot, observability::LaneKind::Decode)) { + lane->durable_draft_tokens = durable; + } + } + } + { + observability::PhaseScope phase( + profile, observability::Phase::SamplingCommit); for (int lane_index = 0; lane_index < spec_count; ++lane_index) { Proposal & proposal = proposals[static_cast(lane_index)]; const int graph_row = ar_count + lane_index * tree_width + @@ -902,6 +1045,18 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( return result; } } + } + if (profile) { + profile->spec_pending_tokens = static_cast(spec_count); + profile->max_kv_len = static_cast(max_prefix); + profile->kv_blocks_free_after = pool_.free_block_count(); + for (const Proposal & proposal : proposals) { + if (auto * lane = profile->find_lane( + proposal.slot, observability::LaneKind::Decode)) { + lane->pending_token_sampled = true; + } + } + } for (int lane_index = 0; lane_index < ar_count; ++lane_index) { ArLane & lane = ar_lanes[static_cast(lane_index)]; lane.pending = sample_graph_row( @@ -932,14 +1087,29 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( return result; } -SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { +SeqEngine::StepResult Qwen35SeqEngine::step( + const StepPlan & plan, observability::StepProfile * profile) { StepResult result; std::vector & decode_outputs = result.decode; std::vector & prefill_outputs = result.prefills; const std::vector & inputs = plan.decode; const int n_slots = slots_.slot_count(); + if (profile) { + profile->kv_blocks_total = pool_.physical_block_count(); + profile->kv_blocks_free_before = pool_.free_block_count(); + profile->kv_blocks_free_after = profile->kv_blocks_free_before; + profile->active_sequences = pool_.active_sequence_count(); + profile->planned_decode_lanes = + static_cast(plan.decode.size()); + profile->planned_prefill_lanes = + static_cast(plan.prefills.size()); + } auto fail_step = [&](const std::string & error) { + if (profile) { + profile->ok = false; + profile->kv_blocks_free_after = pool_.free_block_count(); + } result.decode.clear(); result.prefills.clear(); result.error = error; @@ -982,6 +1152,31 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (inputs.empty() && plan.prefills.empty()) return result; + if (profile) { + profile->planned_prefill_tokens = 0; + for (const PrefillSlice & slice : plan.prefills) { + profile->planned_prefill_tokens += + static_cast(slice.max_tokens); + } + for (const StepInput & input : inputs) { + const auto decision = chain_spec_decision(input); + auto * lane = profile->find_lane( + input.slot, observability::LaneKind::Decode); + if (lane) { + lane->context_tokens = static_cast( + slots_.slot(input.slot).cur_pos); + lane->spec = !plan.prefills.empty() && + decision == observability::SpecDecision::Selected + ? observability::SpecDecision::PromptWorkPresent + : decision; + } + if (decision == observability::SpecDecision::Selected) { + ++profile->spec_eligible_lanes; + if (plan.prefills.empty()) ++profile->spec_reserved_lanes; + } + } + } + FixedServiceRound service_round = make_fixed_service_round(plan); if (auto * decode_round = std::get_if(&service_round)) { @@ -989,17 +1184,40 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { decode_round->chain_lanes.begin(), decode_round->chain_lanes.end(), [](uint8_t selected) { return selected != 0; }); - if (has_chain_lane && - prepare_chain_drafts(inputs, decode_round->chain_lanes)) { - return step_chain_spec(plan, decode_round->chain_lanes); + if (has_chain_lane) { + if (profile) { + profile->spec_attempted_lanes = static_cast( + std::count(decode_round->chain_lanes.begin(), + decode_round->chain_lanes.end(), 1)); + } + if (prepare_chain_drafts( + inputs, decode_round->chain_lanes, profile)) { + return step_chain_spec( + plan, decode_round->chain_lanes, profile); + } + if (profile) { + for (size_t i = 0; i < inputs.size(); ++i) { + if (!decode_round->chain_lanes[i]) continue; + if (auto * lane = profile->find_lane( + inputs[i].slot, + observability::LaneKind::Decode)) { + lane->spec = + observability::SpecDecision::DraftPrepareFailed; + } + } + } } } + if (profile) profile->path = observability::StepPath::Packed; + const TargetWeights & w = b_.w_; StepGraph & sg = b_.sg_; const int hidden = w.n_embd; const int n_head_kv = w.n_head_kv; + const uint64_t input_started_ns = + profile ? observability::steady_time_ns() : 0; decode_outputs.reserve(inputs.size()); prefill_outputs.reserve(plan.prefills.size()); output_rows_.clear(); @@ -1044,6 +1262,12 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { live_physical_rows_.push_back(app.physical_row); live_slot_ids_.push_back(in.slot); max_kv_len = std::max(max_kv_len, app.position + 1); + if (profile) { + if (auto * lane = profile->find_lane( + in.slot, observability::LaneKind::Decode)) { + lane->context_tokens = static_cast(app.position + 1); + } + } out.failed = false; decode_outputs.push_back(std::move(out)); output_rows_.push_back(compact_row); @@ -1065,6 +1289,15 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { return fail_step("selected prefill work made no progress"); } prefills.push_back(std::move(prefill)); + if (profile) { + if (auto * lane = profile->find_lane( + slice.slot, observability::LaneKind::Prefill)) { + lane->context_tokens = static_cast( + prefills.back().kv_pos); + lane->executed_prefill_tokens = static_cast( + prefills.back().chunk); + } + } } const int live_count = (int)live_tokens_.size(); @@ -1103,6 +1336,27 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } const bool with_prefill = n_prefill > 0; const int n_total = n_prefill + decode_bucket; + if (profile) { + const uint64_t input_finished_ns = observability::steady_time_ns(); + profile->add_phase({ + observability::Phase::InputStaging, + input_started_ns >= profile->started_ns + ? input_started_ns - profile->started_ns : 0, + input_finished_ns >= input_started_ns + ? input_finished_ns - input_started_ns : 0, + }); + profile->executed_decode_lanes = + static_cast(live_count); + profile->executed_prefill_lanes = + static_cast(prefills.size()); + profile->executed_prefill_tokens = static_cast(n_prefill); + profile->target_rows = static_cast(n_total); + profile->target_padding_rows = static_cast( + decode_bucket - live_count); + profile->decode_bucket = static_cast(decode_bucket); + profile->max_kv_len = static_cast(max_kv_len); + profile->target_forwards = 1; + } const Qwen35RoctxMetadata roctx_metadata{ live_count, decode_bucket, n_prefill, (int)segments.size(), n_total, max_kv_len}; @@ -1113,6 +1367,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { : 0; bool built = false; + { + observability::PhaseScope phase( + profile, observability::Phase::TargetGraphBuild); if (with_prefill) { built = build_target_step( sg, w, b_.cache_, b_.target_backend_, @@ -1152,6 +1409,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { /*n_logits_rows=*/0, /*compact_slots=*/true); } + } if (!built || !sg.kv_write_rows || (capture_features_ && !sg.target_feat_rows) || (with_prefill && @@ -1160,6 +1418,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { return fail_step("packed prefill/decode graph build failed"); } + { + observability::PhaseScope phase( + profile, observability::Phase::MetadataUpload); embed_buf_.resize((size_t)hidden * n_total); int token_offset = 0; for (const PrefillStage & prefill : prefills) { @@ -1300,9 +1561,12 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { ggml_backend_tensor_set_async( b_.target_backend_, b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, sizeof(int32_t) * seq_lens_.size()); + } ggml_status st = GGML_STATUS_FAILED; { + observability::PhaseScope phase( + profile, observability::Phase::TargetCompute); const Qwen35RoctxRange roctx_compute( "qwen35.graph_compute", roctx_metadata); st = ggml_backend_graph_compute(b_.target_backend_, sg.gf); @@ -1313,16 +1577,21 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const int decode_row0 = with_prefill ? n_commits : 0; const int argmax_rows = with_prefill ? gather_rows : decode_bucket; - argmax_buf_.assign((size_t)argmax_rows, -1); - ggml_backend_tensor_get_async( - b_.target_backend_, sg.argmax_tokens, argmax_buf_.data(), 0, - sizeof(int32_t) * argmax_buf_.size()); { + observability::PhaseScope phase( + profile, observability::Phase::ReadbackSync); + argmax_buf_.assign((size_t)argmax_rows, -1); + ggml_backend_tensor_get_async( + b_.target_backend_, sg.argmax_tokens, argmax_buf_.data(), 0, + sizeof(int32_t) * argmax_buf_.size()); const Qwen35RoctxRange roctx_sync( "qwen35.argmax_readback", roctx_metadata); ggml_backend_synchronize(b_.target_backend_); } + { + observability::PhaseScope phase( + profile, observability::Phase::SamplingCommit); for (size_t oi = 0; oi < inputs.size(); ++oi) { DecodeOutput & out = decode_outputs[oi]; if (out.failed) continue; @@ -1347,6 +1616,17 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } prefill_outputs.push_back(std::move(out)); } + } + if (profile) { + profile->kv_blocks_free_after = pool_.free_block_count(); + for (const DecodeOutput & out : decode_outputs) { + if (out.failed) continue; + if (auto * lane = profile->find_lane( + out.slot, observability::LaneKind::Decode)) { + lane->pending_token_sampled = true; + } + } + } return result; } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index c941dd242..8379133c4 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -61,7 +61,9 @@ class Qwen35SeqEngine final : public SeqEngine { const std::vector & prompt, const SamplerCfg & sampler) override; - StepResult step(const StepPlan & plan) override; + StepResult step( + const StepPlan & plan, + observability::StepProfile * profile = nullptr) override; StepPlanLimits step_plan_limits(int decode_rows) const override { const bool mixed = decode_rows > 0; const int per_sequence = mixed ? 512 : 2048; @@ -128,15 +130,20 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector * logits_scratch = nullptr); FixedServiceRound make_fixed_service_round( const StepPlan & plan) const; + observability::SpecDecision chain_spec_decision( + const StepInput & input) const; bool chain_spec_input_capable(const StepInput & input) const; DraftFeatureMirror * slot_feature_mirror(int slot); DraftKvState * ensure_slot_draft_kv(int slot); bool prepare_chain_drafts( const std::vector & inputs, - const std::vector & selected); + const std::vector & selected, + observability::StepProfile * profile); StepResult step_chain_spec( - const StepPlan & plan, const std::vector & selected); + const StepPlan & plan, const std::vector & selected, + observability::StepProfile * profile); + PagedKvPool & pool_; Qwen35Backend & b_; Qwen35SlotManager slots_; int64_t scratch_row_ = 0; diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 813412f4b..663622df8 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -489,6 +489,8 @@ static const std::vector kApiEndpoints = { "GET /status", "GET /status/events", "GET /status/json", + "GET /observability", + "GET /observability/snapshot", "GET /v1/models", "POST /v1/chat/completions", "POST /v1/messages", @@ -1078,6 +1080,7 @@ HttpServer::HttpServer(ModelBackend & backend, config.disk_cache_min_tokens, config.disk_cache_continued_interval, config.disk_cache_cold_max_tokens}, backend) + , observability_(observability::ObservabilityConfig::from_env()) { #ifdef DFLASH_HAS_CURL curl_global_init(CURL_GLOBAL_DEFAULT); @@ -1092,6 +1095,7 @@ HttpServer::HttpServer(ModelBackend & backend, } disk_cache_.init(); status_html_path_ = resolve_status_html(); + observability_html_path_ = resolve_share_html("observability.html"); // PPP env overrides (operator-facing; no CLI flags required). auto env_truthy = [](const char * v) -> bool { @@ -1124,16 +1128,17 @@ HttpServer::HttpServer(ModelBackend & backend, config_.ppp_max_ephemeral_tokens); } -// Resolve path to share/status.html at startup. std::string HttpServer::resolve_status_html() { - // 1. DFLASH_SHARE_DIR env var + return resolve_share_html("status.html"); +} + +std::string HttpServer::resolve_share_html(const char * filename) { + if (!filename || !*filename) return {}; if (const char * dir = std::getenv("DFLASH_SHARE_DIR")) { - std::string path = std::string(dir) + "/status.html"; + std::string path = std::string(dir) + "/" + filename; struct stat st; if (::stat(path.c_str(), &st) == 0) return path; } - // 2. share/ relative to exe path (build dir or installed prefix) - { std::string exe_dir; #if defined(_WIN32) char exe_buf[MAX_PATH] = {}; @@ -1154,24 +1159,16 @@ std::string HttpServer::resolve_status_html() { } #endif if (!exe_dir.empty()) { - // 2a. /share/status.html (build directory layout) - { - std::string path = exe_dir + "/share/status.html"; - struct stat st; - if (::stat(path.c_str(), &st) == 0) return path; - } - // 2b. /../share/status.html (installed prefix layout) - { - std::string path = exe_dir + "/../share/status.html"; - struct stat st; - if (::stat(path.c_str(), &st) == 0) return path; - } - } + struct stat st; + std::string path = exe_dir + "/share/" + filename; + if (::stat(path.c_str(), &st) == 0) return path; + path = exe_dir + "/../share/" + filename; + if (::stat(path.c_str(), &st) == 0) return path; } - // 3. ./share/status.html (development) { struct stat st; - if (::stat("share/status.html", &st) == 0) return "share/status.html"; + std::string path = std::string("share/") + filename; + if (::stat(path.c_str(), &st) == 0) return path; } return {}; } @@ -1547,6 +1544,33 @@ void HttpServer::handle_client(SocketHandle fd) { return; } + if (hr.method == "GET" && hr.path == "/observability") { + if (observability_html_path_.empty()) { + send_error(fd, 404, + "observability.html not found. Set DFLASH_SHARE_DIR or place it in share/observability.html"); + socket_close(fd); + return; + } + std::ifstream ifs(observability_html_path_); + if (!ifs.is_open()) { + send_error(fd, 500, "failed to open observability.html"); + socket_close(fd); + return; + } + std::ostringstream oss; + oss << ifs.rdbuf(); + send_response(fd, 200, "text/html; charset=utf-8", oss.str()); + socket_close(fd); + return; + } + + if (hr.method == "GET" && hr.path == "/observability/snapshot") { + send_response(fd, 200, "application/json", + observability_.snapshot_json()); + socket_close(fd); + return; + } + // Status SSE stream: hold connection open and push updates. if (hr.method == "GET" && hr.path == "/status/events") { // Send SSE headers. @@ -3971,6 +3995,7 @@ void HttpServer::enqueue(ServerJob * job) { if (queue_tail_) queue_tail_->next = job; else queue_head_ = job; queue_tail_ = job; + job->profile_queued_ns = observability_.job_queued(); queue_cv_.notify_one(); } @@ -3990,6 +4015,7 @@ ServerJob * HttpServer::dequeue() { queue_head_ = j->next; if (!queue_head_) queue_tail_ = nullptr; j->next = nullptr; + observability_.job_dequeued(); return j; } @@ -4000,6 +4026,7 @@ ServerJob * HttpServer::try_dequeue() { queue_head_ = job->next; if (!queue_head_) queue_tail_ = nullptr; job->next = nullptr; + observability_.job_dequeued(); return job; } @@ -4017,6 +4044,7 @@ ServerJob * HttpServer::dequeue_for( queue_head_ = job->next; if (!queue_head_) queue_tail_ = nullptr; job->next = nullptr; + observability_.job_dequeued(); return job; } diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 9e14e8b1f..6659ceb27 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -30,6 +30,7 @@ #include "adaptive_keep_ratio.h" #include "server_status.h" #include "sse_emitter.h" +#include "server/observability.h" #include #include @@ -536,6 +537,8 @@ class HttpServer { // Resolve and cache path to share/status.html. std::string status_html_path_; std::string resolve_status_html(); + std::string observability_html_path_; + std::string resolve_share_html(const char * filename); // Track prompt tokens for each snapshot slot (for shutdown save). std::unordered_map> slot_tokens_; @@ -568,6 +571,7 @@ class HttpServer { std::condition_variable queue_cv_; ServerJob * queue_head_ = nullptr; ServerJob * queue_tail_ = nullptr; + observability::ObservabilityState observability_; std::atomic stopping_{false}; // Active client thread tracking. @@ -603,6 +607,9 @@ struct ServerJob { // First concurrent-scheduler attempt; retained across busy deferrals so // server-side prefill/elapsed telemetry does not erase queueing delay. std::chrono::steady_clock::time_point parallel_started_at{}; + uint64_t profile_request_id = 0; + uint64_t profile_queued_ns = 0; + uint64_t profile_admitted_ns = 0; std::unique_ptr emitter; }; diff --git a/server/src/server/observability.cpp b/server/src/server/observability.cpp new file mode 100644 index 000000000..42a04b42e --- /dev/null +++ b/server/src/server/observability.cpp @@ -0,0 +1,437 @@ +#include "server/observability.h" + +#include "common/prof_env.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common::observability { +namespace { + +uint64_t env_u64(const char * name, uint64_t fallback) { + const char * raw = std::getenv(name); + if (!raw || !*raw) return fallback; + char * end = nullptr; + const unsigned long long value = std::strtoull(raw, &end, 10); + return end && *end == '\0' ? static_cast(value) : fallback; +} + +size_t bounded_size_env(const char * name, size_t fallback) { + return static_cast(std::min( + env_u64(name, fallback), std::numeric_limits::max())); +} + +std::string json_escape(std::string_view value) { + std::ostringstream out; + for (const unsigned char c : value) { + switch (c) { + case '\\': out << "\\\\"; break; + case '"': out << "\\\""; break; + case '\b': out << "\\b"; break; + case '\f': out << "\\f"; break; + case '\n': out << "\\n"; break; + case '\r': out << "\\r"; break; + case '\t': out << "\\t"; break; + default: + if (c < 0x20) { + out << "\\u" << std::hex << std::setw(4) + << std::setfill('0') << static_cast(c) + << std::dec << std::setfill(' '); + } else { + out << static_cast(c); + } + } + } + return out.str(); +} + +void write_u32_array( + std::ostream & out, + const std::array & values) { + out << '['; + for (size_t i = 0; i < values.size(); ++i) { + if (i) out << ','; + out << values[i]; + } + out << ']'; +} + +void write_step_json(std::ostream & out, const StepProfile & step) { + out << "{\"type\":\"step\",\"schema_version\":" + << step.schema_version + << ",\"round_id\":" << step.round_id + << ",\"started_ns\":" << step.started_ns + << ",\"duration_ns\":" << step.duration_ns + << ",\"path\":\"" << step_path_name(step.path) << "\"" + << ",\"ok\":" << (step.ok ? "true" : "false") + << ",\"queue_depth\":" << step.queue_depth + << ",\"live_slots\":" << step.live_slots + << ",\"planned_decode_lanes\":" << step.planned_decode_lanes + << ",\"planned_prefill_lanes\":" << step.planned_prefill_lanes + << ",\"planned_prefill_tokens\":" << step.planned_prefill_tokens + << ",\"executed_decode_lanes\":" << step.executed_decode_lanes + << ",\"executed_prefill_lanes\":" << step.executed_prefill_lanes + << ",\"executed_prefill_tokens\":" << step.executed_prefill_tokens + << ",\"spec_eligible_lanes\":" << step.spec_eligible_lanes + << ",\"spec_reserved_lanes\":" << step.spec_reserved_lanes + << ",\"spec_attempted_lanes\":" << step.spec_attempted_lanes + << ",\"spec_proposed_draft_tokens\":" + << step.spec_proposed_draft_tokens + << ",\"spec_verified_draft_tokens\":" + << step.spec_verified_draft_tokens + << ",\"spec_accepted_draft_tokens\":" + << step.spec_accepted_draft_tokens + << ",\"spec_pending_tokens\":" << step.spec_pending_tokens + << ",\"spec_durable_draft_tokens\":" + << step.spec_durable_draft_tokens + << ",\"spec_scheduler_consumed_tokens\":" + << step.spec_scheduler_consumed_tokens + << ",\"target_rows\":" << step.target_rows + << ",\"target_padding_rows\":" << step.target_padding_rows + << ",\"draft_rows\":" << step.draft_rows + << ",\"draft_padding_rows\":" << step.draft_padding_rows + << ",\"decode_bucket\":" << step.decode_bucket + << ",\"draft_bucket\":" << step.draft_bucket + << ",\"max_kv_len\":" << step.max_kv_len + << ",\"kv_blocks_total\":" << step.kv_blocks_total + << ",\"kv_blocks_free_before\":" << step.kv_blocks_free_before + << ",\"kv_blocks_free_after\":" << step.kv_blocks_free_after + << ",\"active_sequences\":" << step.active_sequences + << ",\"target_forwards\":" << step.target_forwards + << ",\"draft_forwards\":" << step.draft_forwards + << ",\"dropped_lanes\":" << step.dropped_lanes + << ",\"dropped_phases\":" << step.dropped_phases + << ",\"proposed_by_position\":"; + write_u32_array(out, step.proposed_by_position); + out << ",\"accepted_by_position\":"; + write_u32_array(out, step.accepted_by_position); + out << ",\"lanes\":["; + for (uint32_t i = 0; i < step.lane_count; ++i) { + if (i) out << ','; + const LaneProfile & lane = step.lanes[i]; + out << "{\"request_id\":" << lane.request_id + << ",\"slot\":" << lane.slot + << ",\"kind\":\"" << lane_kind_name(lane.kind) << "\"" + << ",\"spec\":\"" << spec_decision_name(lane.spec) << "\"" + << ",\"context_tokens\":" << lane.context_tokens + << ",\"requested_prefill_tokens\":" + << lane.requested_prefill_tokens + << ",\"executed_prefill_tokens\":" + << lane.executed_prefill_tokens + << ",\"proposed_draft_tokens\":" + << lane.proposed_draft_tokens + << ",\"verified_draft_tokens\":" + << lane.verified_draft_tokens + << ",\"accepted_draft_tokens\":" + << lane.accepted_draft_tokens + << ",\"durable_draft_tokens\":" + << lane.durable_draft_tokens + << ",\"scheduler_consumed_tokens\":" + << lane.scheduler_consumed_tokens + << ",\"pending_token_sampled\":" + << (lane.pending_token_sampled ? "true" : "false") + << ",\"pending_token_consumed\":" + << (lane.pending_token_consumed ? "true" : "false") << '}'; + } + out << "],\"phases\":["; + for (uint32_t i = 0; i < step.phase_count; ++i) { + if (i) out << ','; + const PhaseSpan & span = step.phases[i]; + out << "{\"phase\":\"" << phase_name(span.phase) << "\"" + << ",\"start_offset_ns\":" << span.start_offset_ns + << ",\"duration_ns\":" << span.duration_ns << '}'; + } + out << "]}\n"; +} + +} + +ObservabilityConfig ObservabilityConfig::from_env() { + ObservabilityConfig config; + config.enabled = dflash_prof_enabled("concurrency"); + if (const char * path = std::getenv("DFLASH_PROF_OUT")) { + if (*path) config.output_path = path; + } + config.warmup_rounds = env_u64("DFLASH_PROF_WARMUP_ROUNDS", 0); + config.max_rounds = bounded_size_env("DFLASH_PROF_MAX_ROUNDS", 10000); + config.max_requests = bounded_size_env("DFLASH_PROF_MAX_REQUESTS", 4096); + config.max_token_bursts = bounded_size_env( + "DFLASH_PROF_MAX_TOKEN_BURSTS", 200000); + return config; +} + +ObservabilityState::ObservabilityState(ObservabilityConfig config) + : config_(std::move(config)) { + live_.enabled = config_.enabled; + if (!config_.enabled) return; + steps_.reserve(config_.max_rounds); + requests_.reserve(config_.max_requests); + token_bursts_.reserve(config_.max_token_bursts); + active_requests_.reserve(config_.max_requests); +} + +ObservabilityState::~ObservabilityState() { + flush(); +} + +uint64_t ObservabilityState::job_queued() noexcept { + if (!config_.enabled) return 0; + queue_depth_.fetch_add(1, std::memory_order_relaxed); + return steady_time_ns(); +} + +void ObservabilityState::job_dequeued() noexcept { + if (!config_.enabled) return; + uint32_t depth = queue_depth_.load(std::memory_order_relaxed); + while (depth != 0 && !queue_depth_.compare_exchange_weak( + depth, depth - 1, std::memory_order_relaxed)) {} +} + +uint32_t ObservabilityState::queue_depth() const noexcept { + return config_.enabled + ? queue_depth_.load(std::memory_order_relaxed) + : 0; +} + +void ObservabilityState::set_live_slots(uint32_t live_slots) { + if (!config_.enabled) return; + std::lock_guard lock(live_mu_); + live_.live_slots = live_slots; +} + +StepProfile * ObservabilityState::begin_step(uint32_t live_slots) noexcept { + if (!config_.enabled) return nullptr; + current_step_ = {}; + current_step_.round_id = next_round_id_++; + current_step_.started_ns = steady_time_ns(); + current_step_.queue_depth = queue_depth(); + current_step_.live_slots = live_slots; + return ¤t_step_; +} + +void ObservabilityState::commit_step(StepProfile * profile) { + if (!profile) return; + if (profile->duration_ns == 0) { + const uint64_t now = steady_time_ns(); + profile->duration_ns = now >= profile->started_ns + ? now - profile->started_ns : 0; + } + + uint64_t consumed = 0; + for (uint32_t i = 0; i < profile->lane_count; ++i) { + if (profile->lanes[i].kind == LaneKind::Decode) { + consumed += profile->lanes[i].scheduler_consumed_tokens; + } + } + + { + std::lock_guard lock(live_mu_); + ++live_.rounds; + live_.queue_depth = queue_depth(); + live_.live_slots = profile->live_slots; + live_.kv_blocks_total = profile->kv_blocks_total; + live_.kv_blocks_free = profile->kv_blocks_free_after; + live_.planned_prefill_tokens += profile->planned_prefill_tokens; + live_.executed_prefill_tokens += profile->executed_prefill_tokens; + live_.decode_lanes += profile->executed_decode_lanes; + live_.durable_decode_tokens += consumed; + live_.spec_eligible_lanes += profile->spec_eligible_lanes; + live_.spec_reserved_lanes += profile->spec_reserved_lanes; + live_.spec_attempted_lanes += profile->spec_attempted_lanes; + live_.spec_proposed_draft_tokens += + profile->spec_proposed_draft_tokens; + live_.spec_verified_draft_tokens += + profile->spec_verified_draft_tokens; + live_.spec_accepted_draft_tokens += + profile->spec_accepted_draft_tokens; + live_.spec_durable_draft_tokens += + profile->spec_durable_draft_tokens; + live_.spec_scheduler_consumed_tokens += + profile->spec_scheduler_consumed_tokens; + live_.target_rows += profile->target_rows; + live_.target_padding_rows += profile->target_padding_rows; + live_.draft_rows += profile->draft_rows; + live_.draft_padding_rows += profile->draft_padding_rows; + for (uint32_t i = 0; i < profile->phase_count; ++i) { + const size_t phase = static_cast(profile->phases[i].phase); + if (phase < live_.phase_ns.size()) { + live_.phase_ns[phase] += profile->phases[i].duration_ns; + } + } + live_.last_step = { + profile->round_id, profile->duration_ns, profile->path, + profile->queue_depth, profile->live_slots, + }; + } + + if (profile->round_id <= config_.warmup_rounds) return; + if (steps_.size() < config_.max_rounds) { + steps_.push_back(*profile); + } else { + std::lock_guard lock(live_mu_); + ++live_.dropped_steps; + } +} + +void ObservabilityState::record_request_admitted( + uint64_t request_id, std::string response_id, + uint32_t prompt_tokens, uint64_t queued_ns, + uint64_t admitted_ns) { + if (!config_.enabled) return; + if (requests_.size() >= config_.max_requests) { + std::lock_guard lock(live_mu_); + ++live_.dropped_requests; + return; + } + active_requests_[request_id] = requests_.size(); + requests_.push_back({request_id, std::move(response_id), false, + prompt_tokens, 0, queued_ns, admitted_ns}); +} + +void ObservabilityState::record_prefill_completed( + uint64_t request_id, uint64_t now_ns) { + if (!config_.enabled) return; + const auto it = active_requests_.find(request_id); + if (it != active_requests_.end()) { + requests_[it->second].prefill_completed_ns = now_ns; + } +} + +void ObservabilityState::record_token_burst( + uint64_t request_id, uint64_t round_id, + uint64_t ready_ns, uint32_t token_count) { + if (!config_.enabled || token_count == 0) return; + const auto it = active_requests_.find(request_id); + if (it != active_requests_.end() && + requests_[it->second].first_token_ns == 0) { + requests_[it->second].first_token_ns = ready_ns; + } + if (token_bursts_.size() < config_.max_token_bursts) { + token_bursts_.push_back({request_id, round_id, ready_ns, token_count}); + } else { + std::lock_guard lock(live_mu_); + ++live_.dropped_token_bursts; + } +} + +void ObservabilityState::record_request_finished( + uint64_t request_id, bool ok, uint32_t output_tokens, + uint64_t completed_ns) { + if (!config_.enabled) return; + const auto it = active_requests_.find(request_id); + if (it != active_requests_.end()) { + RequestRecord & request = requests_[it->second]; + request.ok = ok; + request.output_tokens = output_tokens; + request.completed_ns = completed_ns; + active_requests_.erase(it); + } + std::lock_guard lock(live_mu_); + ++live_.requests_completed; + if (!ok) ++live_.requests_failed; +} + +LiveMetricsSnapshot ObservabilityState::snapshot() const { + if (!config_.enabled) return live_; + std::lock_guard lock(live_mu_); + LiveMetricsSnapshot result = live_; + result.queue_depth = queue_depth(); + return result; +} + +std::string ObservabilityState::snapshot_json() const { + const LiveMetricsSnapshot s = snapshot(); + std::ostringstream out; + out << "{\"enabled\":" << (s.enabled ? "true" : "false") + << ",\"schema_version\":" << s.schema_version + << ",\"rounds\":" << s.rounds + << ",\"queue_depth\":" << s.queue_depth + << ",\"live_slots\":" << s.live_slots + << ",\"kv_blocks_total\":" << s.kv_blocks_total + << ",\"kv_blocks_free\":" << s.kv_blocks_free + << ",\"planned_prefill_tokens\":" << s.planned_prefill_tokens + << ",\"executed_prefill_tokens\":" << s.executed_prefill_tokens + << ",\"decode_lanes\":" << s.decode_lanes + << ",\"durable_decode_tokens\":" << s.durable_decode_tokens + << ",\"spec_eligible_lanes\":" << s.spec_eligible_lanes + << ",\"spec_reserved_lanes\":" << s.spec_reserved_lanes + << ",\"spec_attempted_lanes\":" << s.spec_attempted_lanes + << ",\"spec_proposed_draft_tokens\":" + << s.spec_proposed_draft_tokens + << ",\"spec_verified_draft_tokens\":" + << s.spec_verified_draft_tokens + << ",\"spec_accepted_draft_tokens\":" + << s.spec_accepted_draft_tokens + << ",\"spec_durable_draft_tokens\":" + << s.spec_durable_draft_tokens + << ",\"spec_scheduler_consumed_tokens\":" + << s.spec_scheduler_consumed_tokens + << ",\"target_rows\":" << s.target_rows + << ",\"target_padding_rows\":" << s.target_padding_rows + << ",\"draft_rows\":" << s.draft_rows + << ",\"draft_padding_rows\":" << s.draft_padding_rows + << ",\"requests_completed\":" << s.requests_completed + << ",\"requests_failed\":" << s.requests_failed + << ",\"dropped_steps\":" << s.dropped_steps + << ",\"dropped_requests\":" << s.dropped_requests + << ",\"dropped_token_bursts\":" << s.dropped_token_bursts + << ",\"phases\":{"; + for (size_t i = 0; i < s.phase_ns.size(); ++i) { + if (i) out << ','; + out << '"' << phase_name(static_cast(i)) << "\":" + << s.phase_ns[i]; + } + out << "},\"last_step\":{\"round_id\":" << s.last_step.round_id + << ",\"duration_ns\":" << s.last_step.duration_ns + << ",\"path\":\"" << step_path_name(s.last_step.path) << "\"" + << ",\"queue_depth\":" << s.last_step.queue_depth + << ",\"live_slots\":" << s.last_step.live_slots << "}}\n"; + return out.str(); +} + +void ObservabilityState::flush() { + if (!config_.enabled || config_.output_path.empty() || flushed_) return; + flushed_ = true; + std::ofstream out(config_.output_path); + if (!out) { + std::cerr << "[observability] failed to write " + << config_.output_path << '\n'; + return; + } + out << "{\"type\":\"metadata\",\"schema\":\"lucebox.concurrency.v1\"}\n"; + for (const StepProfile & step : steps_) write_step_json(out, step); + for (const RequestRecord & request : requests_) { + out << "{\"type\":\"request\",\"request_id\":" + << request.request_id + << ",\"response_id\":\"" << json_escape(request.response_id) + << "\",\"ok\":" << (request.ok ? "true" : "false") + << ",\"prompt_tokens\":" << request.prompt_tokens + << ",\"output_tokens\":" << request.output_tokens + << ",\"queued_ns\":" << request.queued_ns + << ",\"admitted_ns\":" << request.admitted_ns + << ",\"prefill_completed_ns\":" + << request.prefill_completed_ns + << ",\"first_token_ns\":" << request.first_token_ns + << ",\"completed_ns\":" << request.completed_ns << "}\n"; + } + for (const TokenBurst & burst : token_bursts_) { + out << "{\"type\":\"token_burst\",\"request_id\":" + << burst.request_id << ",\"round_id\":" << burst.round_id + << ",\"ready_ns\":" << burst.ready_ns + << ",\"token_count\":" << burst.token_count << "}\n"; + } + const LiveMetricsSnapshot s = snapshot(); + out << "{\"type\":\"footer\",\"dropped_steps\":" + << s.dropped_steps << ",\"dropped_requests\":" + << s.dropped_requests << ",\"dropped_token_bursts\":" + << s.dropped_token_bursts << "}\n"; +} + +} diff --git a/server/src/server/observability.h b/server/src/server/observability.h new file mode 100644 index 000000000..6fa91027d --- /dev/null +++ b/server/src/server/observability.h @@ -0,0 +1,135 @@ +#pragma once + +#include "common/observability/inference_profile.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common::observability { + +struct ObservabilityConfig { + bool enabled = false; + std::string output_path = "concurrency-profile.jsonl"; + uint64_t warmup_rounds = 0; + size_t max_rounds = 10000; + size_t max_requests = 4096; + size_t max_token_bursts = 200000; + + static ObservabilityConfig from_env(); +}; + +struct LastStepSnapshot { + uint64_t round_id = 0; + uint64_t duration_ns = 0; + StepPath path = StepPath::Unknown; + uint32_t queue_depth = 0; + uint32_t live_slots = 0; +}; + +struct LiveMetricsSnapshot { + bool enabled = false; + uint32_t schema_version = kProfileSchemaVersion; + uint64_t rounds = 0; + uint32_t queue_depth = 0; + uint32_t live_slots = 0; + uint32_t kv_blocks_total = 0; + uint32_t kv_blocks_free = 0; + uint64_t planned_prefill_tokens = 0; + uint64_t executed_prefill_tokens = 0; + uint64_t decode_lanes = 0; + uint64_t durable_decode_tokens = 0; + uint64_t spec_eligible_lanes = 0; + uint64_t spec_reserved_lanes = 0; + uint64_t spec_attempted_lanes = 0; + uint64_t spec_proposed_draft_tokens = 0; + uint64_t spec_verified_draft_tokens = 0; + uint64_t spec_accepted_draft_tokens = 0; + uint64_t spec_durable_draft_tokens = 0; + uint64_t spec_scheduler_consumed_tokens = 0; + uint64_t target_rows = 0; + uint64_t target_padding_rows = 0; + uint64_t draft_rows = 0; + uint64_t draft_padding_rows = 0; + uint64_t requests_completed = 0; + uint64_t requests_failed = 0; + uint64_t dropped_steps = 0; + uint64_t dropped_requests = 0; + uint64_t dropped_token_bursts = 0; + std::array phase_ns{}; + LastStepSnapshot last_step; +}; + +class ObservabilityState final { +public: + explicit ObservabilityState(ObservabilityConfig config); + ~ObservabilityState(); + + bool enabled() const noexcept { return config_.enabled; } + + uint64_t job_queued() noexcept; + void job_dequeued() noexcept; + uint32_t queue_depth() const noexcept; + void set_live_slots(uint32_t live_slots); + + StepProfile * begin_step(uint32_t live_slots) noexcept; + void commit_step(StepProfile * profile); + + void record_request_admitted( + uint64_t request_id, std::string response_id, + uint32_t prompt_tokens, uint64_t queued_ns, + uint64_t admitted_ns); + void record_prefill_completed(uint64_t request_id, uint64_t now_ns); + void record_token_burst( + uint64_t request_id, uint64_t round_id, + uint64_t ready_ns, uint32_t token_count); + void record_request_finished( + uint64_t request_id, bool ok, uint32_t output_tokens, + uint64_t completed_ns); + + LiveMetricsSnapshot snapshot() const; + std::string snapshot_json() const; + void flush(); + +private: + struct RequestRecord { + uint64_t request_id = 0; + std::string response_id; + bool ok = false; + uint32_t prompt_tokens = 0; + uint32_t output_tokens = 0; + uint64_t queued_ns = 0; + uint64_t admitted_ns = 0; + uint64_t prefill_completed_ns = 0; + uint64_t first_token_ns = 0; + uint64_t completed_ns = 0; + }; + + struct TokenBurst { + uint64_t request_id = 0; + uint64_t round_id = 0; + uint64_t ready_ns = 0; + uint32_t token_count = 0; + }; + + ObservabilityConfig config_; + std::atomic queue_depth_{0}; + uint64_t next_round_id_ = 1; + StepProfile current_step_; + + mutable std::mutex live_mu_; + LiveMetricsSnapshot live_; + + std::vector steps_; + std::vector requests_; + std::vector token_bursts_; + std::unordered_map active_requests_; + bool flushed_ = false; +}; + +} diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 4ee87f349..6a2fee6bd 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -31,6 +31,7 @@ struct SchedSlot { std::unique_ptr emitter; bool prefilling = false; uint64_t admission_order = 0; + uint64_t request_id = 0; std::chrono::steady_clock::time_point started_at{}; std::chrono::steady_clock::time_point decode_started_at{}; double prefill_s = 0.0; @@ -103,6 +104,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { prefilling == published_prefill_count) return; published_live_count = live_slots; published_prefill_count = prefilling; + observability_.set_live_slots(static_cast(live_slots)); if (live_slots > 0) { status_.set_concurrent_requests(live_slots, prefilling); } else status_.set_idle(); @@ -309,6 +311,13 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { idx, s.prefill_s, decode_s, decode_s > 0.0 ? out_tokens / decode_s : 0.0); + if (observability_.enabled()) { + observability_.record_request_finished( + s.request_id, !s.failed && backend_ok, + static_cast(out_tokens), + observability::steady_time_ns()); + } + engine.retire(idx); // A retirement may have released the blocks the head job needs. deferred_retry_at = {}; @@ -444,7 +453,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // Admission only claims the slot and queues the prompt. Prefill // advances one chunk per engine step alongside live decode. - auto ar = engine.admit(next_request_id, req.prompt_tokens, + const uint64_t request_id = next_request_id; + auto ar = engine.admit(request_id, req.prompt_tokens, req.sampler); if (ar.status == SeqEngine::AdmitResult::Status::busy) return AdmissionDisposition::Deferred; @@ -466,6 +476,10 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { return AdmissionDisposition::Retired; } next_request_id++; + job->profile_request_id = request_id; + if (observability_.enabled()) { + job->profile_admitted_ns = observability::steady_time_ns(); + } SchedSlot & s = slots[(size_t)ar.slot]; s = SchedSlot{}; @@ -473,6 +487,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.fd = job->fd; s.prefilling = true; s.admission_order = next_admission_order++; + s.request_id = request_id; s.started_at = started_at; s.decode_started_at = started_at; // sane on prefill failure s.n_gen_cap = std::min( @@ -486,6 +501,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.hook.hard_limit_remaining = eff_reply_for_n_gen; } live_slots++; + if (observability_.enabled()) { + observability_.record_request_admitted( + request_id, req.response_id, + static_cast(req.prompt_tokens.size()), + job->profile_queued_ns, job->profile_admitted_ns); + } publish_live_count(); return AdmissionDisposition::Admitted; }; @@ -613,27 +634,56 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // Phase 3 — Build one model-neutral batch plan: every decode row plus // a FIFO, engine-bounded subset of pending prompt work. The engine // lowers this plan into whatever graph/state representation it owns. - step_plan.decode.clear(); - prefill_candidates.clear(); - for (int i = 0; i < n_slots; i++) { - if (slots[(size_t)i].job && !slots[(size_t)i].prefilling) { - step_plan.decode.push_back( - {i, slots[(size_t)i].pending_tok, - slots[(size_t)i].hook.close_token_ids.empty()}); - } else if (slots[(size_t)i].job) { - prefill_candidates.push_back( - {i, slots[(size_t)i].admission_order}); + observability::StepProfile * profile = + observability_.begin_step(static_cast(live_slots)); + { + observability::PhaseScope phase( + profile, observability::Phase::SchedulerPlan); + step_plan.decode.clear(); + prefill_candidates.clear(); + for (int i = 0; i < n_slots; i++) { + if (slots[(size_t)i].job && !slots[(size_t)i].prefilling) { + step_plan.decode.push_back( + {i, slots[(size_t)i].pending_tok, + slots[(size_t)i].hook.close_token_ids.empty()}); + } else if (slots[(size_t)i].job) { + prefill_candidates.push_back( + {i, slots[(size_t)i].admission_order}); + } + } + const StepPlanLimits step_limits = + engine.step_plan_limits((int)step_plan.decode.size()); + step_plan.prefills = plan_prefill_slices( + prefill_candidates, step_limits, prefill_round_robin_start); + if (!prefill_candidates.empty()) ++prefill_round_robin_start; + } + if (profile) { + profile->planned_decode_lanes = + static_cast(step_plan.decode.size()); + profile->planned_prefill_lanes = + static_cast(step_plan.prefills.size()); + for (const auto & input : step_plan.decode) { + const SchedSlot & slot = slots[(size_t)input.slot]; + profile->add_lane({ + slot.request_id, input.slot, + observability::LaneKind::Decode, + }); + } + for (const auto & slice : step_plan.prefills) { + const SchedSlot & slot = slots[(size_t)slice.slot]; + observability::LaneProfile lane; + lane.request_id = slot.request_id; + lane.slot = slice.slot; + lane.kind = observability::LaneKind::Prefill; + lane.requested_prefill_tokens = + static_cast(slice.max_tokens); + profile->planned_prefill_tokens += + lane.requested_prefill_tokens; + profile->add_lane(lane); } - } - const StepPlanLimits step_limits = - engine.step_plan_limits((int)step_plan.decode.size()); - step_plan.prefills = plan_prefill_slices( - prefill_candidates, step_limits, prefill_round_robin_start); - if (!prefill_candidates.empty()) { - ++prefill_round_robin_start; } - SeqEngine::StepResult step_result = engine.step(step_plan); + SeqEngine::StepResult step_result = engine.step(step_plan, profile); const std::string protocol_error = validate_step_result(step_plan, step_result, n_slots); if (!protocol_error.empty()) { @@ -644,6 +694,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } if (!step_result.ok()) { + if (profile) profile->ok = false; + observability_.commit_step(profile); const std::string & error = step_result.error; std::fprintf(stderr, "[parallel] engine step failed: %s — " @@ -657,6 +709,9 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } continue; } + { + observability::PhaseScope output_phase( + profile, observability::Phase::OutputProcessing); for (const auto & out : step_result.decode) { if (out.slot < 0 || out.slot >= n_slots) continue; SchedSlot & s = slots[(size_t)out.slot]; @@ -667,11 +722,32 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.finished = true; continue; } - consume_decode_output_tokens(out, [&](int32_t token) { - if (s.finished) return false; + uint32_t consumed = 0; + uint32_t committed_consumed = 0; + for (int32_t token : out.committed_tokens) { + if (s.finished) break; advance_slot(s, token); - return !s.finished; - }); + ++consumed; + ++committed_consumed; + } + bool pending_consumed = false; + if (!s.finished) { + advance_slot(s, out.token); + ++consumed; + pending_consumed = true; + } + if (profile) { + if (auto * lane = profile->find_lane( + out.slot, observability::LaneKind::Decode)) { + lane->scheduler_consumed_tokens = consumed; + lane->pending_token_consumed = pending_consumed; + profile->spec_scheduler_consumed_tokens += std::min( + committed_consumed, lane->durable_draft_tokens); + } + observability_.record_token_burst( + s.request_id, profile->round_id, + observability::steady_time_ns(), consumed); + } } using PrefillStatus = SeqEngine::PrefillOutput::Status; for (const auto & out : step_result.prefills) { @@ -694,13 +770,28 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.prefill_s = std::chrono::duration( s.decode_started_at - s.started_at).count(); advance_slot(s, out.token); + if (profile) { + const uint64_t now_ns = observability::steady_time_ns(); + observability_.record_prefill_completed( + s.request_id, now_ns); + observability_.record_token_burst( + s.request_id, profile->round_id, now_ns, 1); + if (auto * lane = profile->find_lane( + out.slot, observability::LaneKind::Prefill)) { + lane->scheduler_consumed_tokens = 1; + lane->pending_token_consumed = true; + } + } continue; } } + } // Phase 4 — Non-blocking flush of every live slot's chunks. Progress // resets the stall clock; a reader that makes no progress for 30 s // or lets the buffer hit the cap is dropped (its slot retires). { + observability::PhaseScope flush_phase( + profile, observability::Phase::ClientFlush); const auto now = std::chrono::steady_clock::now(); for (int i = 0; i < n_slots; i++) { SchedSlot & s = slots[(size_t)i]; @@ -724,6 +815,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } } } + observability_.commit_step(profile); // Phase 5 — Reap: finish the drains, then hand back the blocks of // every slot that ended this iteration so the next admit can use them. service_drains(); @@ -768,6 +860,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { send_error(queued->fd, 503, "server shutting down"); finish_job(queued); } + observability_.flush(); } diff --git a/server/test/test_inference_profile.cpp b/server/test/test_inference_profile.cpp new file mode 100644 index 000000000..0271d1e02 --- /dev/null +++ b/server/test/test_inference_profile.cpp @@ -0,0 +1,75 @@ +#include "common/observability/inference_profile.h" + +#include +#include +#include + +using namespace dflash::common::observability; + +#define CHECK(condition) do { \ + if (!(condition)) { \ + std::fprintf(stderr, "CHECK failed at %s:%d: %s\n", \ + __FILE__, __LINE__, #condition); \ + return 1; \ + } \ +} while (false) + +namespace { + +uint64_t fake_now = 0; +int fake_clock_calls = 0; + +uint64_t fake_clock() noexcept { + ++fake_clock_calls; + fake_now += 10; + return fake_now; +} + +} + +int main() { + fake_now = 0; + fake_clock_calls = 0; + { + PhaseScope scope(nullptr, Phase::TargetCompute, fake_clock); + } + CHECK(fake_clock_calls == 0); + + StepProfile profile; + profile.started_ns = 5; + { + PhaseScope scope(&profile, Phase::TargetCompute, fake_clock); + } + CHECK(fake_clock_calls == 2); + CHECK(profile.phase_count == 1); + CHECK(profile.phases[0].phase == Phase::TargetCompute); + CHECK(profile.phases[0].start_offset_ns == 5); + CHECK(profile.phases[0].duration_ns == 10); + + LaneProfile lane; + lane.request_id = 42; + lane.slot = 3; + lane.kind = LaneKind::Decode; + CHECK(profile.add_lane(lane) != nullptr); + CHECK(profile.find_lane(3, LaneKind::Decode)->request_id == 42); + CHECK(profile.find_lane(3, LaneKind::Prefill) == nullptr); + + StepProfile full; + for (size_t i = 0; i < kMaxProfileLanes; ++i) { + lane.slot = static_cast(i); + CHECK(full.add_lane(lane) != nullptr); + } + CHECK(full.add_lane(lane) == nullptr); + CHECK(full.dropped_lanes == 1); + + CHECK(std::string_view(step_path_name(StepPath::Speculative)) == + "speculative"); + CHECK(std::string_view(spec_decision_name( + SpecDecision::PromptWorkPresent)) == + "prompt_work_present"); + CHECK(std::string_view(phase_name(Phase::ClientFlush)) == + "client_flush"); + + std::printf("test_inference_profile: passed\n"); + return 0; +} diff --git a/server/test/test_observability.cpp b/server/test/test_observability.cpp new file mode 100644 index 000000000..7f5f49c07 --- /dev/null +++ b/server/test/test_observability.cpp @@ -0,0 +1,94 @@ +#include "server/observability.h" + +#include +#include +#include +#include +#include + +using namespace dflash::common::observability; + +#define CHECK(condition) do { \ + if (!(condition)) { \ + std::fprintf(stderr, "CHECK failed at %s:%d: %s\n", \ + __FILE__, __LINE__, #condition); \ + return 1; \ + } \ +} while (false) + +int main() { + ObservabilityState disabled({}); + CHECK(disabled.job_queued() == 0); + CHECK(disabled.queue_depth() == 0); + CHECK(disabled.begin_step(4) == nullptr); + CHECK(disabled.snapshot_json().find("\"enabled\":false") != + std::string::npos); + + const auto output = std::filesystem::temp_directory_path() / + ("lucebox-observability-" + std::to_string(steady_time_ns()) + + ".jsonl"); + ObservabilityConfig config; + config.enabled = true; + config.output_path = output.string(); + config.max_rounds = 1; + config.max_requests = 1; + config.max_token_bursts = 1; + ObservabilityState state(config); + + const uint64_t queued_ns = state.job_queued(); + CHECK(queued_ns > 0); + CHECK(state.queue_depth() == 1); + state.job_dequeued(); + CHECK(state.queue_depth() == 0); + + state.set_live_slots(3); + CHECK(state.snapshot().live_slots == 3); + + state.record_request_admitted(7, "response", 10, queued_ns, + queued_ns + 10); + StepProfile * step = state.begin_step(4); + CHECK(step != nullptr); + step->path = StepPath::Speculative; + step->executed_decode_lanes = 4; + step->spec_eligible_lanes = 4; + step->spec_attempted_lanes = 4; + step->spec_proposed_draft_tokens = 12; + step->spec_accepted_draft_tokens = 8; + step->kv_blocks_total = 100; + step->kv_blocks_free_after = 80; + LaneProfile lane; + lane.request_id = 7; + lane.slot = 0; + lane.scheduler_consumed_tokens = 3; + step->add_lane(lane); + step->add_phase({Phase::TargetCompute, 1, 20}); + state.record_prefill_completed(7, queued_ns + 20); + state.record_token_burst(7, step->round_id, queued_ns + 30, 3); + state.record_request_finished(7, true, 3, queued_ns + 40); + state.commit_step(step); + + StepProfile * dropped = state.begin_step(4); + dropped->kv_blocks_total = 100; + dropped->kv_blocks_free_after = 80; + state.commit_step(dropped); + const LiveMetricsSnapshot snapshot = state.snapshot(); + CHECK(snapshot.rounds == 2); + CHECK(snapshot.kv_blocks_free == 80); + CHECK(snapshot.durable_decode_tokens == 3); + CHECK(snapshot.requests_completed == 1); + CHECK(snapshot.dropped_steps == 1); + state.flush(); + std::ifstream input(output); + const std::string jsonl{ + std::istreambuf_iterator(input), + std::istreambuf_iterator()}; + CHECK(jsonl.find("lucebox.concurrency.v1\"}\n{\"type\":\"step\"") != + std::string::npos); + CHECK(jsonl.find("\"type\":\"request\"") != std::string::npos); + CHECK(jsonl.find("\"type\":\"token_burst\"") != std::string::npos); + CHECK(jsonl.find("\"dropped_steps\":1") != std::string::npos); + std::filesystem::remove(output); + + std::printf("test_observability: passed\n"); + return 0; +} diff --git a/server/test/test_seq_engine_contract.cpp b/server/test/test_seq_engine_contract.cpp index 3b42d9c0e..d4b2ecb8a 100644 --- a/server/test/test_seq_engine_contract.cpp +++ b/server/test/test_seq_engine_contract.cpp @@ -82,7 +82,9 @@ class FakeSeqEngine final : public SeqEngine { return result; } - StepResult step(const StepPlan & plan) override { + StepResult step( + const StepPlan & plan, + observability::StepProfile * = nullptr) override { StepResult result; std::string validation_error; if (!valid_decode(plan, validation_error) && From 9ebde47c66783214576c3e7c2f2bace7583747f6 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 22 Aug 2026 08:26:22 +0000 Subject: [PATCH 06/11] concurrency: make profiles recoverable and comparable Add opt-in atomic checkpoints and run context to concurrency captures. Emit versioned JSON summaries for benchmark diffs. Remove redundant profile accounting and clarify adapter ownership. --- .../benchmarks/concurrency/profile_report.py | 333 +++++++++++++----- .../concurrency/test_profile_report.py | 18 +- server/CMakeLists.txt | 24 ++ server/docs/CONCURRENCY_OBSERVABILITY.md | 34 +- server/docs/ENVIRONMENT.md | 1 + server/src/common/concurrency/seq_engine.h | 12 +- .../observability/inference_profile.cpp | 1 + .../common/observability/inference_profile.h | 2 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 13 +- server/src/server/http_server.cpp | 26 +- server/src/server/http_server.h | 2 + server/src/server/observability.cpp | 110 +++++- server/src/server/observability.h | 17 + server/src/server/server_main.cpp | 2 + server/test/seq_engine_contract.h | 6 +- server/test/test_inference_profile.cpp | 2 + server/test/test_observability.cpp | 27 +- server/test/test_server_unit.cpp | 18 - 18 files changed, 514 insertions(+), 134 deletions(-) diff --git a/harness/benchmarks/concurrency/profile_report.py b/harness/benchmarks/concurrency/profile_report.py index b06180f58..3bf57efc6 100644 --- a/harness/benchmarks/concurrency/profile_report.py +++ b/harness/benchmarks/concurrency/profile_report.py @@ -54,32 +54,44 @@ def duration_ms(end: int, start: int) -> float: return (end - start) / 1_000_000 if end and start and end >= start else math.nan -def fmt_ms(value: float) -> str: - return "n/a" if math.isnan(value) else f"{value:.2f} ms" - - def sum_field(records: Iterable[dict[str, Any]], key: str) -> int: return sum(int(record.get(key, 0)) for record in records) -def build_markdown(records: list[dict[str, Any]]) -> str: - steps = [record for record in records if record["type"] == "step"] - requests = [record for record in records if record["type"] == "request"] - bursts = [record for record in records if record["type"] == "token_burst"] +def optional_number(value: float) -> float | None: + return None if math.isnan(value) else value + + +def percentile_pair(values: Iterable[float]) -> dict[str, float | None]: + finite = [value for value in values if not math.isnan(value)] + return { + "p50": optional_number(percentile(finite, 0.50)), + "p95": optional_number(percentile(finite, 0.95)), + } + + +def build_summary(records: list[dict[str, Any]]) -> dict[str, Any]: + metadata = next( + (record for record in records if record["type"] == "metadata"), + {}, + ) footer = next( (record for record in reversed(records) if record["type"] == "footer"), {}, ) - failed_requests = sum( - not bool(request.get("ok")) for request in requests - ) + steps = [record for record in records if record["type"] == "step"] + requests = [record for record in records if record["type"] == "request"] + bursts = [record for record in records if record["type"] == "token_burst"] phases: Counter[str] = Counter() decisions: Counter[str] = Counter() - cohorts: dict[int, list[dict[str, Any]]] = defaultdict(list) paths: Counter[str] = Counter() + cohorts: dict[int, list[dict[str, Any]]] = defaultdict(list) + proposed_by_position: list[int] = [] + accepted_by_position: list[int] = [] for step in steps: - cohorts[int(step.get("live_slots", 0))].append(step) + live_slots = int(step.get("live_slots", 0)) + cohorts[live_slots].append(step) paths[str(step.get("path", "unknown"))] += 1 for span in step.get("phases", []): phases[str(span.get("phase", "unknown"))] += int( @@ -88,48 +100,162 @@ def build_markdown(records: list[dict[str, Any]]) -> str: for lane in step.get("lanes", []): if lane.get("kind") == "decode": decisions[str(lane.get("spec", "none"))] += 1 + for source, target in ( + (step.get("proposed_by_position", []), proposed_by_position), + (step.get("accepted_by_position", []), accepted_by_position), + ): + if len(target) < len(source): + target.extend([0] * (len(source) - len(target))) + for index, value in enumerate(source): + target[index] += int(value) queue_ms = [ - duration_ms(int(request.get("admitted_ns", 0)), int(request.get("queued_ns", 0))) + duration_ms(int(request.get("admitted_ns", 0)), + int(request.get("queued_ns", 0))) for request in requests ] ttft_ms = [ - duration_ms(int(request.get("first_token_ns", 0)), int(request.get("queued_ns", 0))) + duration_ms(int(request.get("first_token_ns", 0)), + int(request.get("queued_ns", 0))) for request in requests ] e2e_ms = [ - duration_ms(int(request.get("completed_ns", 0)), int(request.get("queued_ns", 0))) + duration_ms(int(request.get("completed_ns", 0)), + int(request.get("queued_ns", 0))) for request in requests ] - queue_ms = [value for value in queue_ms if not math.isnan(value)] - ttft_ms = [value for value in ttft_ms if not math.isnan(value)] - e2e_ms = [value for value in e2e_ms if not math.isnan(value)] - burst_times: dict[int, list[tuple[int, int]]] = defaultdict(list) for burst in bursts: burst_times[int(burst["request_id"])].append( (int(burst["ready_ns"]), int(burst.get("token_count", 0))) ) - inter_token_ms: list[float] = [] + inter_burst_ms: list[float] = [] for request_bursts in burst_times.values(): request_bursts.sort() for previous, current in zip(request_bursts, request_bursts[1:]): - token_count = max(1, current[1]) - inter_token_ms.append((current[0] - previous[0]) / 1_000_000 / token_count) - - eligible = sum_field(steps, "spec_eligible_lanes") - reserved = sum_field(steps, "spec_reserved_lanes") - attempted = sum_field(steps, "spec_attempted_lanes") - proposed = sum_field(steps, "spec_proposed_draft_tokens") - verified = sum_field(steps, "spec_verified_draft_tokens") - accepted = sum_field(steps, "spec_accepted_draft_tokens") - durable = sum_field(steps, "spec_durable_draft_tokens") - consumed = sum_field(steps, "spec_scheduler_consumed_tokens") + inter_burst_ms.append( + (current[0] - previous[0]) / + 1_000_000 / max(1, current[1]) + ) + + funnel_keys = ( + "spec_eligible_lanes", + "spec_reserved_lanes", + "spec_attempted_lanes", + "spec_proposed_draft_tokens", + "spec_verified_draft_tokens", + "spec_accepted_draft_tokens", + "spec_durable_draft_tokens", + "spec_scheduler_consumed_tokens", + ) + funnel = {key: sum_field(steps, key) for key in funnel_keys} + acceptance_by_position = [ + optional_number(ratio(accepted, proposed)) + for proposed, accepted in zip( + proposed_by_position, accepted_by_position) + ] + + cohort_summary: dict[str, Any] = {} + for live_slots, cohort in sorted(cohorts.items()): + target_rows = sum_field(cohort, "target_rows") + target_padding = sum_field(cohort, "target_padding_rows") + draft_rows = sum_field(cohort, "draft_rows") + draft_padding = sum_field(cohort, "draft_padding_rows") + proposed = sum_field(cohort, "spec_proposed_draft_tokens") + accepted = sum_field(cohort, "spec_accepted_draft_tokens") + cohort_summary[str(live_slots)] = { + "rounds": len(cohort), + "mean_round_ms": statistics.fmean( + int(step.get("duration_ns", 0)) for step in cohort + ) / 1_000_000, + "target_padding_ratio": optional_number( + ratio(target_padding, target_rows) + ), + "draft_padding_ratio": optional_number( + ratio(draft_padding, draft_rows) + ), + "draft_acceptance_ratio": optional_number( + ratio(accepted, proposed) + ), + "paths": dict(sorted(Counter( + str(step.get("path", "unknown")) for step in cohort + ).items())), + } + target_rows = sum_field(steps, "target_rows") target_padding = sum_field(steps, "target_padding_rows") draft_rows = sum_field(steps, "draft_rows") draft_padding = sum_field(steps, "draft_padding_rows") - phase_total = sum(phases.values()) + return { + "schema": "lucebox.concurrency.summary.v1", + "run": { + key: value for key, value in metadata.items() + if key not in {"type", "schema"} + }, + "capture": { + "complete": bool(footer.get("complete", False)), + "rounds": len(steps), + "requests": len(requests), + "failed_requests": sum( + not bool(request.get("ok")) for request in requests + ), + "dropped_steps": int(footer.get("dropped_steps", 0)), + "dropped_requests": int(footer.get("dropped_requests", 0)), + "dropped_token_bursts": int( + footer.get("dropped_token_bursts", 0) + ), + "paths": dict(sorted(paths.items())), + }, + "latency_ms": { + "queue": percentile_pair(queue_ms), + "ttft": percentile_pair(ttft_ms), + "end_to_end": percentile_pair(e2e_ms), + "inter_burst_per_token": percentile_pair(inter_burst_ms), + }, + "speculation": { + **funnel, + "tree_widths": sorted({ + int(step.get("spec_tree_width", 0)) for step in steps + if int(step.get("spec_tree_width", 0)) > 0 + }), + "decisions": dict(sorted(decisions.items())), + "proposed_by_position": proposed_by_position, + "accepted_by_position": accepted_by_position, + "acceptance_by_position": acceptance_by_position, + }, + "padding": { + "target_rows": target_rows, + "target_padding_rows": target_padding, + "target_padding_ratio": optional_number( + ratio(target_padding, target_rows) + ), + "draft_rows": draft_rows, + "draft_padding_rows": draft_padding, + "draft_padding_ratio": optional_number( + ratio(draft_padding, draft_rows) + ), + }, + "cohorts": cohort_summary, + "phase_ns": dict(phases.most_common()), + } + + +def build_markdown(records: list[dict[str, Any]]) -> str: + summary = build_summary(records) + run = summary["run"] + capture = summary["capture"] + latency = summary["latency_ms"] + speculation = summary["speculation"] + padding = summary["padding"] + + def format_ms(value: float | None) -> str: + return "n/a" if value is None else f"{value:.2f} ms" + + def format_ratio(value: float | None) -> str: + return "n/a" if value is None else f"{100.0 * value:.1f}%" + + def format_pair(values: dict[str, float | None]) -> str: + return f"{format_ms(values['p50'])} / {format_ms(values['p95'])}" lines = [ "# Lucebox concurrency profile", @@ -138,17 +264,26 @@ def build_markdown(records: list[dict[str, Any]]) -> str: "", "| Metric | Value |", "| --- | ---: |", - f"| Captured rounds | {len(steps)} |", - f"| Requests | {len(requests)} |", - f"| Failed requests | {failed_requests} |", - f"| Queue delay p50 / p95 | {fmt_ms(percentile(queue_ms, 0.50))} / {fmt_ms(percentile(queue_ms, 0.95))} |", - f"| TTFT p50 / p95 | {fmt_ms(percentile(ttft_ms, 0.50))} / {fmt_ms(percentile(ttft_ms, 0.95))} |", - f"| End-to-end p50 / p95 | {fmt_ms(percentile(e2e_ms, 0.50))} / {fmt_ms(percentile(e2e_ms, 0.95))} |", - f"| Inter-burst token interval p50 / p95 | {fmt_ms(percentile(inter_token_ms, 0.50))} / {fmt_ms(percentile(inter_token_ms, 0.95))} |", - f"| Target padding | {target_padding} / {target_rows} ({percent(ratio(target_padding, target_rows))}) |", - f"| Draft padding | {draft_padding} / {draft_rows} ({percent(ratio(draft_padding, draft_rows))}) |", + f"| Git SHA | `{run.get('git_sha', 'unknown')}` |", + f"| Model | `{run.get('model_name', 'unknown')}` |", + f"| Configured concurrency | {run.get('max_concurrency', 'unknown')} |", + f"| Captured rounds | {capture['rounds']} |", + f"| Requests | {capture['requests']} |", + f"| Failed requests | {capture['failed_requests']} |", + f"| Queue delay p50 / p95 | {format_pair(latency['queue'])} |", + f"| TTFT p50 / p95 | {format_pair(latency['ttft'])} |", + f"| End-to-end p50 / p95 | {format_pair(latency['end_to_end'])} |", + f"| Inter-burst token interval p50 / p95 | {format_pair(latency['inter_burst_per_token'])} |", + f"| Target padding | {padding['target_padding_rows']} / " + f"{padding['target_rows']} " + f"({format_ratio(padding['target_padding_ratio'])}) |", + f"| Draft padding | {padding['draft_padding_rows']} / " + f"{padding['draft_rows']} " + f"({format_ratio(padding['draft_padding_ratio'])}) |", "", - "The inter-burst interval divides each gap by the number of tokens made ready in the later burst. It is a scheduler-level estimate, not a per-token GPU timestamp.", + "The inter-burst interval divides each gap by the number of tokens " + "made ready in the later burst. It is a scheduler-level estimate, " + "not a per-token GPU timestamp.", "", "## Speculation funnel", "", @@ -156,14 +291,19 @@ def build_markdown(records: list[dict[str, Any]]) -> str: "| --- | ---: | ---: |", ] funnel = [ - ("Eligible lanes", eligible), - ("Reserved lanes", reserved), - ("Attempted lanes", attempted), - ("Proposed draft tokens", proposed), - ("Verified draft tokens", verified), - ("Accepted draft tokens", accepted), - ("Durable draft tokens", durable), - ("Scheduler-consumed draft tokens", consumed), + ("Eligible lanes", speculation["spec_eligible_lanes"]), + ("Reserved lanes", speculation["spec_reserved_lanes"]), + ("Attempted lanes", speculation["spec_attempted_lanes"]), + ("Proposed draft tokens", + speculation["spec_proposed_draft_tokens"]), + ("Verified draft tokens", + speculation["spec_verified_draft_tokens"]), + ("Accepted draft tokens", + speculation["spec_accepted_draft_tokens"]), + ("Durable draft tokens", + speculation["spec_durable_draft_tokens"]), + ("Scheduler-consumed draft tokens", + speculation["spec_scheduler_consumed_tokens"]), ] previous = 0 for index, (name, value) in enumerate(funnel): @@ -181,30 +321,31 @@ def build_markdown(records: list[dict[str, Any]]) -> str: "| Decision | Decode lanes |", "| --- | ---: |", ]) - for decision, count in sorted(decisions.items()): + for decision, count in speculation["decisions"].items(): lines.append(f"| `{decision}` | {count} |") lines.extend([ "", "## Concurrency cohorts", "", - "| Live slots | Rounds | Mean round | Target padding | Draft acceptance | Paths |", + "| Live slots | Rounds | Mean round | Target padding | " + "Draft acceptance | Paths |", "| ---: | ---: | ---: | ---: | ---: | --- |", ]) - for live_slots, cohort in sorted(cohorts.items()): - mean_ms = statistics.fmean(int(step.get("duration_ns", 0)) for step in cohort) / 1_000_000 - cohort_target = sum_field(cohort, "target_rows") - cohort_target_padding = sum_field(cohort, "target_padding_rows") - cohort_proposed = sum_field(cohort, "spec_proposed_draft_tokens") - cohort_accepted = sum_field(cohort, "spec_accepted_draft_tokens") - cohort_paths = Counter(str(step.get("path", "unknown")) for step in cohort) - path_text = ", ".join(f"{name}={count}" for name, count in sorted(cohort_paths.items())) + for live_slots, cohort in sorted( + summary["cohorts"].items(), key=lambda item: int(item[0])): + path_text = ", ".join( + f"{name}={count}" for name, count in cohort["paths"].items() + ) lines.append( - f"| {live_slots} | {len(cohort)} | {mean_ms:.2f} ms | " - f"{percent(ratio(cohort_target_padding, cohort_target))} | " - f"{percent(ratio(cohort_accepted, cohort_proposed))} | {path_text} |" + f"| {live_slots} | {cohort['rounds']} | " + f"{cohort['mean_round_ms']:.2f} ms | " + f"{format_ratio(cohort['target_padding_ratio'])} | " + f"{format_ratio(cohort['draft_acceptance_ratio'])} | " + f"{path_text} |" ) + phase_total = sum(summary["phase_ns"].values()) lines.extend([ "", "## Phase time", @@ -212,44 +353,72 @@ def build_markdown(records: list[dict[str, Any]]) -> str: "| Phase | Total | Share of measured phase time |", "| --- | ---: | ---: |", ]) - for phase, nanoseconds in phases.most_common(): + for phase, nanoseconds in summary["phase_ns"].items(): lines.append( f"| `{phase}` | {nanoseconds / 1_000_000:.2f} ms | " f"{percent(ratio(nanoseconds, phase_total))} |" ) + eligible = speculation["spec_eligible_lanes"] + attempted = speculation["spec_attempted_lanes"] + proposed = speculation["spec_proposed_draft_tokens"] + accepted = speculation["spec_accepted_draft_tokens"] + durable = speculation["spec_durable_draft_tokens"] + queue_p95 = latency["queue"]["p95"] + ttft_p95 = latency["ttft"]["p95"] signals: list[str] = [] - if failed_requests: + if capture["failed_requests"]: signals.append( - f"{failed_requests}/{len(requests)} captured requests failed. " - "Inspect the first incomplete funnel or phase boundary." + f"{capture['failed_requests']}/{capture['requests']} captured " + "requests failed. Inspect the first incomplete funnel or phase " + "boundary." ) if accepted != durable: signals.append( "Accepted and durable draft token counts differ. Inspect state " "promotion or commit before tuning proposal quality." ) - if target_rows and ratio(target_padding, target_rows) > 0.20: - signals.append("Target graph padding exceeds 20%. Inspect cohort bucket shapes.") + if (padding["target_padding_ratio"] is not None and + padding["target_padding_ratio"] > 0.20): + signals.append( + "Target graph padding exceeds 20%. Inspect cohort bucket shapes." + ) if proposed and ratio(accepted, proposed) < 0.35: - signals.append("Draft acceptance is below 35%. Inspect proposal quality before increasing speculative width.") + signals.append( + "Draft acceptance is below 35%. Inspect proposal quality before " + "increasing speculative width." + ) if eligible and ratio(attempted, eligible) < 0.75: - signals.append("Fewer than 75% of eligible lanes reach an attempt. Inspect suppression reasons and prompt mixing.") - if requests and percentile(queue_ms, 0.95) > percentile(ttft_ms, 0.95) * 0.40: - signals.append("Queueing accounts for a large part of p95 TTFT. Inspect admission and KV pressure.") + signals.append( + "Fewer than 75% of eligible lanes reach an attempt. Inspect " + "suppression reasons and prompt mixing." + ) + if (capture["requests"] and queue_p95 is not None and + ttft_p95 is not None and queue_p95 > ttft_p95 * 0.40): + signals.append( + "Queueing accounts for a large part of p95 TTFT. Inspect " + "admission and KV pressure." + ) if not signals: - signals.append("No default threshold fired. Use the cohort and phase tables to choose the next experiment.") + signals.append( + "No default threshold fired. Use the cohort and phase tables to " + "choose the next experiment." + ) lines.extend(["", "## Signals", ""]) lines.extend(f"- {signal}" for signal in signals) + paths = ", ".join( + f"{name}={count}" for name, count in capture["paths"].items() + ) or "none" lines.extend([ "", "## Capture integrity", "", - f"- Paths: {', '.join(f'{name}={count}' for name, count in sorted(paths.items())) or 'none'}", - f"- Dropped steps: {int(footer.get('dropped_steps', 0))}", - f"- Dropped requests: {int(footer.get('dropped_requests', 0))}", - f"- Dropped token bursts: {int(footer.get('dropped_token_bursts', 0))}", + f"- Complete: {'yes' if capture['complete'] else 'no'}", + f"- Paths: {paths}", + f"- Dropped steps: {capture['dropped_steps']}", + f"- Dropped requests: {capture['dropped_requests']}", + f"- Dropped token bursts: {capture['dropped_token_bursts']}", "", ]) return "\n".join(lines) @@ -319,6 +488,7 @@ def main() -> int: parser.add_argument("profile", type=Path) parser.add_argument("--markdown", type=Path) parser.add_argument("--perfetto", type=Path) + parser.add_argument("--json-summary", type=Path) args = parser.parse_args() records = load_records(args.profile) @@ -332,6 +502,11 @@ def main() -> int: json.dumps(build_perfetto(records), indent=2) + "\n", encoding="utf-8", ) + if args.json_summary: + args.json_summary.write_text( + json.dumps(build_summary(records), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) return 0 diff --git a/harness/benchmarks/concurrency/test_profile_report.py b/harness/benchmarks/concurrency/test_profile_report.py index f5c989455..4219ed06d 100644 --- a/harness/benchmarks/concurrency/test_profile_report.py +++ b/harness/benchmarks/concurrency/test_profile_report.py @@ -11,11 +11,13 @@ class ProfileReportTest(unittest.TestCase): def records(self): return [ - {"type": "metadata", "schema": "lucebox.concurrency.v1"}, + {"type": "metadata", "schema": "lucebox.concurrency.v1", + "git_sha": "abc123", "max_concurrency": 4}, { "type": "step", "round_id": 1, "started_ns": 1_000_000, "duration_ns": 2_000_000, "path": "speculative", "live_slots": 4, "target_rows": 20, + "spec_tree_width": 8, "target_padding_rows": 4, "draft_rows": 16, "draft_padding_rows": 0, "spec_eligible_lanes": 4, "spec_reserved_lanes": 4, "spec_attempted_lanes": 4, @@ -38,7 +40,7 @@ def records(self): "ready_ns": 500, "token_count": 1}, {"type": "token_burst", "request_id": 9, "round_id": 2, "ready_ns": 900, "token_count": 2}, - {"type": "footer", "dropped_steps": 0, + {"type": "footer", "complete": True, "dropped_steps": 0, "dropped_requests": 0, "dropped_token_bursts": 0}, ] @@ -53,6 +55,18 @@ def test_markdown_and_perfetto_share_the_records(self): self.assertIn("queue", names) self.assertIn("tokens_ready", names) + summary = profile_report.build_summary(records) + self.assertEqual(summary["schema"], + "lucebox.concurrency.summary.v1") + self.assertEqual(summary["run"]["git_sha"], "abc123") + self.assertTrue(summary["capture"]["complete"]) + self.assertEqual(summary["cohorts"]["4"]["rounds"], 1) + self.assertEqual( + summary["speculation"]["spec_accepted_draft_tokens"], 8) + self.assertEqual(summary["speculation"]["tree_widths"], [8]) + self.assertAlmostEqual( + summary["padding"]["target_padding_ratio"], 0.2) + def test_loader_rejects_a_different_schema(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "profile.jsonl" diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index f4ec351e6..d316a86cb 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -2028,6 +2028,30 @@ if(DFLASH27B_SERVER) src/server/model_card.cpp src/server/prompt_normalize.cpp ) + execute_process( + COMMAND git rev-parse --verify HEAD + WORKING_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/.." + RESULT_VARIABLE _dflash_git_result + OUTPUT_VARIABLE _dflash_git_sha + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET) + if(NOT _dflash_git_result EQUAL 0 OR NOT _dflash_git_sha) + set(_dflash_git_sha "unknown") + else() + execute_process( + COMMAND git diff --quiet --ignore-submodules -- + WORKING_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/.." + RESULT_VARIABLE _dflash_git_dirty + ERROR_QUIET) + if(_dflash_git_dirty EQUAL 1) + string(APPEND _dflash_git_sha "-dirty") + endif() + endif() + target_compile_definitions(dflash_server PRIVATE + DFLASH_GIT_SHA="${_dflash_git_sha}") + unset(_dflash_git_result) + unset(_dflash_git_dirty) + unset(_dflash_git_sha) target_include_directories(dflash_server PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) if(DFLASH27B_GPU_BACKEND STREQUAL "hip") target_compile_definitions(dflash_server PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP) diff --git a/server/docs/CONCURRENCY_OBSERVABILITY.md b/server/docs/CONCURRENCY_OBSERVABILITY.md index 379f52c08..4480685a9 100644 --- a/server/docs/CONCURRENCY_OBSERVABILITY.md +++ b/server/docs/CONCURRENCY_OBSERVABILITY.md @@ -19,6 +19,28 @@ Stop the server normally to write the JSONL file. Capture is bounded by the [ENVIRONMENT.md](ENVIRONMENT.md). The footer reports dropped records when a bound is reached. +Set `DFLASH_PROF_CHECKPOINT_EVERY` to preserve an in-progress capture if the +server later crashes or hangs. Each checkpoint replaces the prior JSONL file +atomically and ends with `"complete": false`. Clean shutdown replaces it with +the final capture and `"complete": true`. Checkpoint writing runs on the +scheduler thread, so leave it disabled for overhead measurements. A signal +handler does not attempt to flush C++ streams because that is not +async-signal-safe. + +The metadata record includes the configured Git SHA, model and draft paths, +architecture, backend, maximum concurrency, DDTree budget, draft block-size +override, selected Qwen concurrency environment values, and both wall-clock +and steady-clock anchors. The anchors correlate service-round timestamps with +benchmark logs and external traces. Each Qwen step also records the effective +speculative tree width resolved from the drafter. + +Round retention is keep-first. Once `DFLASH_PROF_MAX_ROUNDS` is reached, later +rounds increment `dropped_steps`. The metadata record reports the retention +policy, `max_rounds`, and `step_record_bytes`, so the reserved memory and any +early-run bias are visible in every capture. Raise the limit for long benchmark +runs. Periodic checkpoints protect in-progress data but do not change the +retention policy. + ## Inspect a live server The server exposes two read-only routes: @@ -42,11 +64,15 @@ Generate a Markdown summary and a Perfetto trace from the same JSONL capture. python3 harness/benchmarks/concurrency/profile_report.py \ /tmp/lucebox-profile.jsonl \ --markdown /tmp/lucebox-profile.md \ - --perfetto /tmp/lucebox-profile.perfetto.json + --perfetto /tmp/lucebox-profile.perfetto.json \ + --json-summary /tmp/lucebox-profile.summary.json ``` Open the Perfetto JSON at [ui.perfetto.dev](https://ui.perfetto.dev). It shows round phase spans, request queue/prefill/decode spans, and token-ready bursts. +The JSON summary has a versioned schema for benchmark diffs. It includes run +context, latency percentiles, phase totals, padding ratios, concurrency +cohorts, suppression decisions, and acceptance by speculative position. ## Read the speculation funnel @@ -90,3 +116,9 @@ pressure, and speculation progress. The contract is model-neutral and fixed-capacity. A future non-batched C=1 adapter can populate the same record without changing the report, metrics, or dashboard. + +The detailed engine instrumentation currently lives in the Qwen35-family +adapter used by Qwen3.8. Another sequence engine can accept the optional +profile pointer and leave it untouched, but its capture will contain only the +scheduler-owned lifecycle and planning fields until that adapter records its +own execution facts. diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 6770cc6a3..12ae6ce8b 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -24,6 +24,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_PROF_MAX_ROUNDS` | `10000` | Maximum concurrency round records retained for shutdown export. | | `DFLASH_PROF_MAX_REQUESTS` | `4096` | Maximum request lifecycle records retained for shutdown export. | | `DFLASH_PROF_MAX_TOKEN_BURSTS` | `200000` | Maximum scheduler token-burst records retained for shutdown export. | +| `DFLASH_PROF_CHECKPOINT_EVERY` | `0` | Write an atomic in-progress JSONL snapshot after this many service rounds. `0` writes only on clean shutdown. | | `GGML_CUDA_GRAPH_STATS` | unset | DEBUG: per-graph CUDA-graph replay/capture/eager counters. | | `GGML_CUDA_GRAPH_STATS_EVERY` | 200 | DEBUG: print period for the stats above (clamped to >=1). | | `DFLASH_ADAPTIVE_K_TAU` | 0 = off | Prefer the CLI: --adaptive-experts [tau]. Cumulative combine-weight threshold for per-token expert gating. | diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 99395c461..322e25757 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -250,7 +250,7 @@ class SeqEngine { // backend mutation, but expose no consumable payload. virtual StepResult step( const StepPlan & plan, - observability::StepProfile * profile = nullptr) = 0; + observability::StepProfile * profile) = 0; // Release a slot's KV blocks and mark it free. Safe on failed slots. virtual void retire(int slot) = 0; @@ -259,16 +259,6 @@ class SeqEngine { virtual bool token_is_eos(int32_t token) const = 0; }; -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. diff --git a/server/src/common/observability/inference_profile.cpp b/server/src/common/observability/inference_profile.cpp index ba6aa1cca..f19be17af 100644 --- a/server/src/common/observability/inference_profile.cpp +++ b/server/src/common/observability/inference_profile.cpp @@ -51,6 +51,7 @@ const char * spec_decision_name(SpecDecision decision) noexcept { switch (decision) { case SpecDecision::None: return "none"; case SpecDecision::Selected: return "selected"; + case SpecDecision::InvalidSlot: return "invalid_slot"; case SpecDecision::PromptWorkPresent: return "prompt_work_present"; case SpecDecision::CallerDisallowed: return "caller_disallowed"; case SpecDecision::FeatureUnavailable: return "feature_unavailable"; diff --git a/server/src/common/observability/inference_profile.h b/server/src/common/observability/inference_profile.h index 17275bd73..57c265014 100644 --- a/server/src/common/observability/inference_profile.h +++ b/server/src/common/observability/inference_profile.h @@ -25,6 +25,7 @@ enum class LaneKind : uint8_t { enum class SpecDecision : uint8_t { None, Selected, + InvalidSlot, PromptWorkPresent, CallerDisallowed, FeatureUnavailable, @@ -109,6 +110,7 @@ struct StepProfile { uint32_t draft_padding_rows = 0; uint32_t decode_bucket = 0; uint32_t draft_bucket = 0; + uint32_t spec_tree_width = 0; uint32_t max_kv_len = 0; uint32_t kv_blocks_total = 0; uint32_t kv_blocks_free_before = 0; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 29d7b9c64..1a64c880d 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -164,7 +164,7 @@ observability::SpecDecision Qwen35SeqEngine::chain_spec_decision( } if (!input.allow_speculation) return Decision::CallerDisallowed; if (input.slot < 0 || input.slot >= slots_.slot_count()) { - return Decision::InsufficientContext; + return Decision::InvalidSlot; } const Qwen35Slot & slot = slots_.slot(input.slot); if (!slot.decoding() || slot.cur_pos < 1 || @@ -1099,10 +1099,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step( profile->kv_blocks_free_before = pool_.free_block_count(); profile->kv_blocks_free_after = profile->kv_blocks_free_before; profile->active_sequences = pool_.active_sequence_count(); - profile->planned_decode_lanes = - static_cast(plan.decode.size()); - profile->planned_prefill_lanes = - static_cast(plan.prefills.size()); + profile->spec_tree_width = static_cast( + std::max(0, tree_width_)); } auto fail_step = [&](const std::string & error) { @@ -1153,11 +1151,6 @@ SeqEngine::StepResult Qwen35SeqEngine::step( if (inputs.empty() && plan.prefills.empty()) return result; if (profile) { - profile->planned_prefill_tokens = 0; - for (const PrefillSlice & slice : plan.prefills) { - profile->planned_prefill_tokens += - static_cast(slice.max_tokens); - } for (const StepInput & input : inputs) { const auto decision = chain_spec_decision(input); auto * lane = profile->find_lane( diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 663622df8..cb982d611 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -46,6 +46,10 @@ using dflash::common::SocketHandle; +#ifndef DFLASH_GIT_SHA +#define DFLASH_GIT_SHA "unknown" +#endif + #if defined(_WIN32) #include #include @@ -1067,6 +1071,26 @@ static std::array compute_disk_cache_salt(const ServerConfig & cfg) // ─── HttpServer ───────────────────────────────────────────────────────── +namespace { + +observability::ObservabilityConfig make_observability_config( + const ServerConfig & server) { + observability::ObservabilityConfig config = + observability::ObservabilityConfig::from_env(); + config.git_sha = DFLASH_GIT_SHA; + config.model_name = server.model_name; + config.model_path = server.model_path; + config.draft_path = server.draft_path; + config.arch = server.arch; + config.runtime_backend = server.runtime_backend; + config.max_concurrency = server.max_concurrency; + config.ddtree_budget = server.ddtree_budget; + config.draft_block_size = server.draft_block_size; + return config; +} + +} + HttpServer::HttpServer(ModelBackend & backend, Tokenizer & tokenizer, const ServerConfig & config) @@ -1080,7 +1104,7 @@ HttpServer::HttpServer(ModelBackend & backend, config.disk_cache_min_tokens, config.disk_cache_continued_interval, config.disk_cache_cold_max_tokens}, backend) - , observability_(observability::ObservabilityConfig::from_env()) + , observability_(make_observability_config(config)) { #ifdef DFLASH_HAS_CURL curl_global_init(CURL_GLOBAL_DEFAULT); diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 6659ceb27..8a2935e11 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -176,6 +176,8 @@ struct ServerConfig { std::string runtime_backend; // "cuda" | "hip" | "cpu" int fa_window = 0; int ddtree_budget = 0; + int draft_block_size = 0; + int max_concurrency = 1; bool speculative_enabled = false; bool target_sharding = false; // Prefill chunk size (bargs.chunk). Exposed at /props.runtime.chunk so diff --git a/server/src/server/observability.cpp b/server/src/server/observability.cpp index 42a04b42e..45494ac93 100644 --- a/server/src/server/observability.cpp +++ b/server/src/server/observability.cpp @@ -3,7 +3,9 @@ #include "common/prof_env.h" #include +#include #include +#include #include #include #include @@ -28,6 +30,16 @@ size_t bounded_size_env(const char * name, size_t fallback) { env_u64(name, fallback), std::numeric_limits::max())); } +void capture_env( + ObservabilityConfig & config, + std::initializer_list names) { + for (const char * name : names) { + if (const char * value = std::getenv(name)) { + config.run_env.emplace_back(name, value); + } + } +} + std::string json_escape(std::string_view value) { std::ostringstream out; for (const unsigned char c : value) { @@ -99,6 +111,7 @@ void write_step_json(std::ostream & out, const StepProfile & step) { << ",\"draft_padding_rows\":" << step.draft_padding_rows << ",\"decode_bucket\":" << step.decode_bucket << ",\"draft_bucket\":" << step.draft_bucket + << ",\"spec_tree_width\":" << step.spec_tree_width << ",\"max_kv_len\":" << step.max_kv_len << ",\"kv_blocks_total\":" << step.kv_blocks_total << ",\"kv_blocks_free_before\":" << step.kv_blocks_free_before @@ -164,6 +177,18 @@ ObservabilityConfig ObservabilityConfig::from_env() { config.max_requests = bounded_size_env("DFLASH_PROF_MAX_REQUESTS", 4096); config.max_token_bursts = bounded_size_env( "DFLASH_PROF_MAX_TOKEN_BURSTS", 200000); + config.checkpoint_every_rounds = env_u64( + "DFLASH_PROF_CHECKPOINT_EVERY", 0); + capture_env(config, { + "DFLASH_QWEN35_DFLASH2_TREE", + "DFLASH_QWEN35_DSPARK_TREE", + "DFLASH_QWEN35_SPEC_STEP_RATIO", + "DFLASH_DRAFT_KV", + "DFLASH_DISABLE_DRAFT_SWA", + "DFLASH_QWEN35_NO_KVPAD", + "DFLASH27B_PREFILL_UBATCH", + "DFLASH_KVFLASH", + }); return config; } @@ -171,6 +196,10 @@ ObservabilityState::ObservabilityState(ObservabilityConfig config) : config_(std::move(config)) { live_.enabled = config_.enabled; if (!config_.enabled) return; + started_unix_ns_ = static_cast( + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()).count()); + started_steady_ns_ = steady_time_ns(); steps_.reserve(config_.max_rounds); requests_.reserve(config_.max_requests); token_bursts_.reserve(config_.max_token_bursts); @@ -278,6 +307,12 @@ void ObservabilityState::commit_step(StepProfile * profile) { std::lock_guard lock(live_mu_); ++live_.dropped_steps; } + if (config_.checkpoint_every_rounds != 0 && + profile->round_id - last_checkpoint_round_ >= + config_.checkpoint_every_rounds && + write_capture(false)) { + last_checkpoint_round_ = profile->round_id; + } } void ObservabilityState::record_request_admitted( @@ -396,16 +431,41 @@ std::string ObservabilityState::snapshot_json() const { return out.str(); } -void ObservabilityState::flush() { - if (!config_.enabled || config_.output_path.empty() || flushed_) return; - flushed_ = true; - std::ofstream out(config_.output_path); +bool ObservabilityState::write_capture(bool complete) { + if (!config_.enabled || config_.output_path.empty()) return false; + const std::string temporary_path = config_.output_path + ".tmp"; + std::ofstream out(temporary_path); if (!out) { std::cerr << "[observability] failed to write " - << config_.output_path << '\n'; - return; + << temporary_path << '\n'; + return false; + } + out << "{\"type\":\"metadata\",\"schema\":\"lucebox.concurrency.v1\"" + << ",\"schema_version\":" << kProfileSchemaVersion + << ",\"git_sha\":\"" << json_escape(config_.git_sha) << "\"" + << ",\"model_name\":\"" << json_escape(config_.model_name) << "\"" + << ",\"model_path\":\"" << json_escape(config_.model_path) << "\"" + << ",\"draft_path\":\"" << json_escape(config_.draft_path) << "\"" + << ",\"arch\":\"" << json_escape(config_.arch) << "\"" + << ",\"runtime_backend\":\"" + << json_escape(config_.runtime_backend) << "\"" + << ",\"max_concurrency\":" << config_.max_concurrency + << ",\"ddtree_budget\":" << config_.ddtree_budget + << ",\"draft_block_size\":" << config_.draft_block_size + << ",\"started_unix_ns\":" << started_unix_ns_ + << ",\"started_steady_ns\":" << started_steady_ns_ + << ",\"round_retention\":\"keep_first\"" + << ",\"max_rounds\":" << config_.max_rounds + << ",\"step_record_bytes\":" << sizeof(StepProfile) + << ",\"checkpoint_every_rounds\":" + << config_.checkpoint_every_rounds + << ",\"env\":{"; + for (size_t i = 0; i < config_.run_env.size(); ++i) { + if (i) out << ','; + out << '\"' << json_escape(config_.run_env[i].first) << "\":\"" + << json_escape(config_.run_env[i].second) << '\"'; } - out << "{\"type\":\"metadata\",\"schema\":\"lucebox.concurrency.v1\"}\n"; + out << "}}\n"; for (const StepProfile & step : steps_) write_step_json(out, step); for (const RequestRecord & request : requests_) { out << "{\"type\":\"request\",\"request_id\":" @@ -431,7 +491,41 @@ void ObservabilityState::flush() { out << "{\"type\":\"footer\",\"dropped_steps\":" << s.dropped_steps << ",\"dropped_requests\":" << s.dropped_requests << ",\"dropped_token_bursts\":" - << s.dropped_token_bursts << "}\n"; + << s.dropped_token_bursts << ",\"complete\":" + << (complete ? "true" : "false") << "}\n"; + out.close(); + if (!out) { + std::cerr << "[observability] failed to finish " + << temporary_path << '\n'; + std::error_code ignored; + std::filesystem::remove(temporary_path, ignored); + return false; + } + + std::error_code error; + std::filesystem::rename( + temporary_path, config_.output_path, error); +#if defined(_WIN32) + if (error) { + std::filesystem::remove(config_.output_path, error); + error.clear(); + std::filesystem::rename( + temporary_path, config_.output_path, error); + } +#endif + if (error) { + std::cerr << "[observability] failed to publish " + << config_.output_path << ": " << error.message() << '\n'; + std::error_code ignored; + std::filesystem::remove(temporary_path, ignored); + return false; + } + return true; +} + +void ObservabilityState::flush() { + if (!config_.enabled || config_.output_path.empty() || flushed_) return; + if (write_capture(true)) flushed_ = true; } } diff --git a/server/src/server/observability.h b/server/src/server/observability.h index 6fa91027d..f45729520 100644 --- a/server/src/server/observability.h +++ b/server/src/server/observability.h @@ -9,6 +9,7 @@ #include #include #include +#include #include namespace dflash::common::observability { @@ -20,6 +21,17 @@ struct ObservabilityConfig { size_t max_rounds = 10000; size_t max_requests = 4096; size_t max_token_bursts = 200000; + uint64_t checkpoint_every_rounds = 0; + std::string git_sha = "unknown"; + std::string model_name; + std::string model_path; + std::string draft_path; + std::string arch; + std::string runtime_backend; + int max_concurrency = 1; + int ddtree_budget = 0; + int draft_block_size = 0; + std::vector> run_env; static ObservabilityConfig from_env(); }; @@ -129,7 +141,12 @@ class ObservabilityState final { std::vector requests_; std::vector token_bursts_; std::unordered_map active_requests_; + uint64_t started_unix_ns_ = 0; + uint64_t started_steady_ns_ = 0; + uint64_t last_checkpoint_round_ = 0; bool flushed_ = false; + + bool write_capture(bool complete); }; } diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 176c66da8..1812f3480 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -1280,6 +1280,8 @@ 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.draft_block_size = bargs.draft_block_size; + sconfig.max_concurrency = bargs.max_concurrency; sconfig.speculative_enabled = bargs.ddtree_mode; sconfig.target_sharding = bargs.device.is_layer_split(); // KV type: report the operator's choice if set, else the family default diff --git a/server/test/seq_engine_contract.h b/server/test/seq_engine_contract.h index 63ca035ab..ead65b95d 100644 --- a/server/test/seq_engine_contract.h +++ b/server/test/seq_engine_contract.h @@ -179,7 +179,7 @@ inline std::vector check_seq_engine_contract(SeqEngine & engine) { }; auto execute = [&](const SeqEngine::StepPlan & plan) { - return apply_progress(plan, engine.step(plan)); + return apply_progress(plan, engine.step(plan, nullptr)); }; auto decode_inputs = [&]() { @@ -270,7 +270,7 @@ inline std::vector check_seq_engine_contract(SeqEngine & engine) { // terminal plan-validation failure and must not partially advance state. auto require_failed = [&](const SeqEngine::StepPlan & invalid, const char * message) { - const SeqEngine::StepResult result = engine.step(invalid); + const SeqEngine::StepResult result = engine.step(invalid, nullptr); require(!result.ok(), message); require(!result.error.empty(), "failed step must explain the validation error"); @@ -333,7 +333,7 @@ inline std::vector check_seq_engine_contract(SeqEngine & engine) { } retire_all(); - const SeqEngine::StepResult idle = engine.step({}); + const SeqEngine::StepResult idle = engine.step({}, nullptr); require(idle.ok(), "step() with no work must succeed"); require(idle.decode.empty() && idle.prefills.empty() && idle.error.empty(), diff --git a/server/test/test_inference_profile.cpp b/server/test/test_inference_profile.cpp index 0271d1e02..2b9667b76 100644 --- a/server/test/test_inference_profile.cpp +++ b/server/test/test_inference_profile.cpp @@ -67,6 +67,8 @@ int main() { CHECK(std::string_view(spec_decision_name( SpecDecision::PromptWorkPresent)) == "prompt_work_present"); + CHECK(std::string_view(spec_decision_name(SpecDecision::InvalidSlot)) == + "invalid_slot"); CHECK(std::string_view(phase_name(Phase::ClientFlush)) == "client_flush"); diff --git a/server/test/test_observability.cpp b/server/test/test_observability.cpp index 7f5f49c07..330bf6b49 100644 --- a/server/test/test_observability.cpp +++ b/server/test/test_observability.cpp @@ -33,6 +33,17 @@ int main() { config.max_rounds = 1; config.max_requests = 1; config.max_token_bursts = 1; + config.checkpoint_every_rounds = 1; + config.git_sha = "0123456789abcdef"; + config.model_name = "qwen38"; + config.model_path = "/models/target.gguf"; + config.draft_path = "/models/draft.gguf"; + config.arch = "qwen35"; + config.runtime_backend = "hip"; + config.max_concurrency = 4; + config.ddtree_budget = 8; + config.draft_block_size = 8; + config.run_env.emplace_back("DFLASH_DRAFT_KV", "1"); ObservabilityState state(config); const uint64_t queued_ns = state.job_queued(); @@ -66,6 +77,19 @@ int main() { state.record_token_burst(7, step->round_id, queued_ns + 30, 3); state.record_request_finished(7, true, 3, queued_ns + 40); state.commit_step(step); + { + std::ifstream checkpoint_input(output); + const std::string checkpoint{ + std::istreambuf_iterator(checkpoint_input), + std::istreambuf_iterator()}; + CHECK(checkpoint.find("\"complete\":false") != std::string::npos); + CHECK(checkpoint.find("\"git_sha\":\"0123456789abcdef\"") != + std::string::npos); + CHECK(checkpoint.find("\"max_concurrency\":4") != + std::string::npos); + CHECK(checkpoint.find("\"DFLASH_DRAFT_KV\":\"1\"") != + std::string::npos); + } StepProfile * dropped = state.begin_step(4); dropped->kv_blocks_total = 100; @@ -82,11 +106,12 @@ int main() { const std::string jsonl{ std::istreambuf_iterator(input), std::istreambuf_iterator()}; - CHECK(jsonl.find("lucebox.concurrency.v1\"}\n{\"type\":\"step\"") != + CHECK(jsonl.find("lucebox.concurrency.v1") != std::string::npos); CHECK(jsonl.find("\"type\":\"request\"") != std::string::npos); CHECK(jsonl.find("\"type\":\"token_burst\"") != std::string::npos); CHECK(jsonl.find("\"dropped_steps\":1") != std::string::npos); + CHECK(jsonl.find("\"complete\":true") != std::string::npos); std::filesystem::remove(output); std::printf("test_observability: passed\n"); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index c7907f4c9..94e7a9292 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -5943,21 +5943,3 @@ TEST_CASE(ServerUnitFixture, test_emitter_function_calls_param_with_literal_thin TEST_ASSERT(em.emit_token_count() == 3); TEST_ASSERT(em.emit_token_count() - em.first_content_token_index() == 1); } - -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; - const bool consumed_all = consume_decode_output_tokens( - burst, [&](int32_t token) { - emitted.push_back(token); - return token != 2; - }); - - TEST_ASSERT(!consumed_all); - TEST_ASSERT((emitted == std::vector{101, 2})); -} From bd6aeb3b6fa514490075a2d91435b0d36108117b Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 22 Aug 2026 08:40:24 +0000 Subject: [PATCH 07/11] concurrency: simplify profiling feedback loop Treat unfinished checkpoint requests as incomplete, reject invalid unsigned limits, and compact position histograms. Remove speculative report thresholds, dead job state, repeated summary work, and the duplicate dashboard panel. --- .../benchmarks/concurrency/profile_report.py | 66 +++++-------- .../concurrency/test_profile_report.py | 19 +++- server/share/observability.html | 96 ++----------------- .../qwen35/concurrency/qwen35_seq_engine.h | 2 +- server/src/server/http_server.h | 2 - server/src/server/observability.cpp | 34 +++++-- server/src/server/scheduler.cpp | 8 +- server/test/test_observability.cpp | 24 ++++- 8 files changed, 102 insertions(+), 149 deletions(-) diff --git a/harness/benchmarks/concurrency/profile_report.py b/harness/benchmarks/concurrency/profile_report.py index 3bf57efc6..978feffdd 100644 --- a/harness/benchmarks/concurrency/profile_report.py +++ b/harness/benchmarks/concurrency/profile_report.py @@ -197,7 +197,9 @@ def build_summary(records: list[dict[str, Any]]) -> dict[str, Any]: "rounds": len(steps), "requests": len(requests), "failed_requests": sum( - not bool(request.get("ok")) for request in requests + int(request.get("completed_ns", 0)) != 0 + and request.get("ok") is False + for request in requests ), "dropped_steps": int(footer.get("dropped_steps", 0)), "dropped_requests": int(footer.get("dropped_requests", 0)), @@ -240,8 +242,7 @@ def build_summary(records: list[dict[str, Any]]) -> dict[str, Any]: } -def build_markdown(records: list[dict[str, Any]]) -> str: - summary = build_summary(records) +def build_markdown(summary: dict[str, Any]) -> str: run = summary["run"] capture = summary["capture"] latency = summary["latency_ms"] @@ -359,54 +360,29 @@ def format_pair(values: dict[str, float | None]) -> str: f"{percent(ratio(nanoseconds, phase_total))} |" ) - eligible = speculation["spec_eligible_lanes"] - attempted = speculation["spec_attempted_lanes"] - proposed = speculation["spec_proposed_draft_tokens"] accepted = speculation["spec_accepted_draft_tokens"] durable = speculation["spec_durable_draft_tokens"] - queue_p95 = latency["queue"]["p95"] - ttft_p95 = latency["ttft"]["p95"] - signals: list[str] = [] + warnings: list[str] = [] if capture["failed_requests"]: - signals.append( + warnings.append( f"{capture['failed_requests']}/{capture['requests']} captured " "requests failed. Inspect the first incomplete funnel or phase " "boundary." ) if accepted != durable: - signals.append( + warnings.append( "Accepted and durable draft token counts differ. Inspect state " "promotion or commit before tuning proposal quality." ) - if (padding["target_padding_ratio"] is not None and - padding["target_padding_ratio"] > 0.20): - signals.append( - "Target graph padding exceeds 20%. Inspect cohort bucket shapes." - ) - if proposed and ratio(accepted, proposed) < 0.35: - signals.append( - "Draft acceptance is below 35%. Inspect proposal quality before " - "increasing speculative width." - ) - if eligible and ratio(attempted, eligible) < 0.75: - signals.append( - "Fewer than 75% of eligible lanes reach an attempt. Inspect " - "suppression reasons and prompt mixing." - ) - if (capture["requests"] and queue_p95 is not None and - ttft_p95 is not None and queue_p95 > ttft_p95 * 0.40): - signals.append( - "Queueing accounts for a large part of p95 TTFT. Inspect " - "admission and KV pressure." - ) - if not signals: - signals.append( - "No default threshold fired. Use the cohort and phase tables to " - "choose the next experiment." - ) - - lines.extend(["", "## Signals", ""]) - lines.extend(f"- {signal}" for signal in signals) + if not capture["complete"]: + warnings.append("Capture is incomplete.") + dropped = ( + capture["dropped_steps"] + + capture["dropped_requests"] + + capture["dropped_token_bursts"] + ) + if dropped: + warnings.append(f"Capture dropped {dropped} records.") paths = ", ".join( f"{name}={count}" for name, count in capture["paths"].items() ) or "none" @@ -419,8 +395,11 @@ def format_pair(values: dict[str, float | None]) -> str: f"- Dropped steps: {capture['dropped_steps']}", f"- Dropped requests: {capture['dropped_requests']}", f"- Dropped token bursts: {capture['dropped_token_bursts']}", - "", ]) + if warnings: + lines.extend(["", "### Warnings", ""]) + lines.extend(f"- {warning}" for warning in warnings) + lines.append("") return "\n".join(lines) @@ -492,7 +471,8 @@ def main() -> int: args = parser.parse_args() records = load_records(args.profile) - markdown = build_markdown(records) + summary = build_summary(records) + markdown = build_markdown(summary) if args.markdown: args.markdown.write_text(markdown, encoding="utf-8") else: @@ -504,7 +484,7 @@ def main() -> int: ) if args.json_summary: args.json_summary.write_text( - json.dumps(build_summary(records), indent=2, sort_keys=True) + "\n", + json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) return 0 diff --git a/harness/benchmarks/concurrency/test_profile_report.py b/harness/benchmarks/concurrency/test_profile_report.py index 4219ed06d..d6ea85a4f 100644 --- a/harness/benchmarks/concurrency/test_profile_report.py +++ b/harness/benchmarks/concurrency/test_profile_report.py @@ -46,7 +46,8 @@ def records(self): def test_markdown_and_perfetto_share_the_records(self): records = self.records() - markdown = profile_report.build_markdown(records) + markdown = profile_report.build_markdown( + profile_report.build_summary(records)) self.assertIn("## Speculation funnel", markdown) self.assertIn("| 4 | 1 | 2.00 ms | 20.0% | 66.7%", markdown) trace = profile_report.build_perfetto(records) @@ -79,11 +80,25 @@ def test_failure_and_durability_gap_are_actionable(self): records[1]["spec_durable_draft_tokens"] = 0 records[2]["ok"] = False - markdown = profile_report.build_markdown(records) + markdown = profile_report.build_markdown( + profile_report.build_summary(records)) self.assertIn("1/1 captured requests failed", markdown) self.assertIn("Accepted and durable draft token counts differ", markdown) + def test_incomplete_request_is_not_failed(self): + records = self.records() + records.insert(-1, { + "type": "request", "request_id": 10, "ok": None, + "queued_ns": 1000, "admitted_ns": 1100, + "completed_ns": 0, + }) + + summary = profile_report.build_summary(records) + + self.assertEqual(summary["capture"]["requests"], 2) + self.assertEqual(summary["capture"]["failed_requests"], 0) + if __name__ == "__main__": unittest.main() diff --git a/server/share/observability.html b/server/share/observability.html index c6f42f50f..68e503a89 100644 --- a/server/share/observability.html +++ b/server/share/observability.html @@ -199,37 +199,21 @@ padding: 14px; } -.span-3 { - grid-column: span 3; -} - -.span-4 { - grid-column: span 4; -} - .span-5 { grid-column: span 5; } -.span-6 { - grid-column: span 6; -} - .span-7 { grid-column: span 7; } -.span-8 { - grid-column: span 8; -} - .span-12 { grid-column: span 12; } .metric-grid { display: grid; - grid-template-columns: repeat(4, minmax(0, 1fr)); + grid-template-columns: repeat(6, minmax(0, 1fr)); gap: 10px; } @@ -313,10 +297,6 @@ background: linear-gradient(90deg, #ff4962, var(--danger)); } -.bar-fill.muted { - background: linear-gradient(90deg, #5d7293, #9db2d0); -} - .bar-value { color: var(--text); font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, "Liberation Mono", monospace; @@ -400,30 +380,6 @@ content: ""; } -.mono-table { - width: 100%; - border-collapse: collapse; -} - -.mono-table th, -.mono-table td { - padding: 8px 6px; - border-bottom: 1px solid rgba(157, 178, 208, 0.12); - font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, "Liberation Mono", monospace; - font-size: 0.78rem; - text-align: left; - vertical-align: top; -} - -.mono-table th { - color: var(--muted); - font-weight: 600; -} - -.mono-table td { - color: var(--text); -} - .notice { min-height: 42px; padding: 11px 12px; @@ -439,12 +395,8 @@ } @media (max-width: 1100px) { - .span-3, - .span-4, .span-5, - .span-6, - .span-7, - .span-8 { + .span-7 { grid-column: span 6; } @@ -467,12 +419,8 @@ justify-content: flex-start; } - .span-3, - .span-4, .span-5, - .span-6, .span-7, - .span-8, .span-12 { grid-column: span 12; } @@ -534,6 +482,14 @@

Live counters

Decode lanes
0
+
+
Last path
+
unknown
+
+
+
Last round time
+
unknown
+
@@ -642,34 +598,6 @@

Utilization and padding

-
-
-

Last step

-
duration unknown
-
-
- - - - - - - - - - - - - - - - - - - -
RoundPathDurationLive slotsQueue
unknownunknownunknown00
-
-
@@ -948,12 +876,8 @@

Last step

const path = lastStep.path || data.path || "unknown"; const duration = Number(lastStep.duration_ns || 0); setText("last-step-label", "round " + roundId); - setText("last-step-duration", fmtDuration(duration)); - setText("last-round", String(roundId)); setText("last-path", String(path)); setText("last-duration", fmtDuration(duration)); - setText("last-live-slots", fmt(n(lastStep, "live_slots", liveSlots))); - setText("last-queue-depth", fmt(n(lastStep, "queue_depth", queueDepth))); renderPhases(data); renderFunnel(data); diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 8379133c4..3f3dbc468 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -63,7 +63,7 @@ class Qwen35SeqEngine final : public SeqEngine { StepResult step( const StepPlan & plan, - observability::StepProfile * profile = nullptr) override; + observability::StepProfile * profile) override; StepPlanLimits step_plan_limits(int decode_rows) const override { const bool mixed = decode_rows > 0; const int per_sequence = mixed ? 512 : 2048; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 8a2935e11..441d88973 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -609,9 +609,7 @@ struct ServerJob { // First concurrent-scheduler attempt; retained across busy deferrals so // server-side prefill/elapsed telemetry does not erase queueing delay. std::chrono::steady_clock::time_point parallel_started_at{}; - uint64_t profile_request_id = 0; uint64_t profile_queued_ns = 0; - uint64_t profile_admitted_ns = 0; std::unique_ptr emitter; }; diff --git a/server/src/server/observability.cpp b/server/src/server/observability.cpp index 45494ac93..a6734b681 100644 --- a/server/src/server/observability.cpp +++ b/server/src/server/observability.cpp @@ -3,6 +3,7 @@ #include "common/prof_env.h" #include +#include #include #include #include @@ -20,9 +21,13 @@ namespace { uint64_t env_u64(const char * name, uint64_t fallback) { const char * raw = std::getenv(name); if (!raw || !*raw) return fallback; - char * end = nullptr; - const unsigned long long value = std::strtoull(raw, &end, 10); - return end && *end == '\0' ? static_cast(value) : fallback; + const std::string_view text(raw); + uint64_t value = 0; + const auto parsed = std::from_chars( + text.data(), text.data() + text.size(), value); + return parsed.ec == std::errc{} && + parsed.ptr == text.data() + text.size() + ? value : fallback; } size_t bounded_size_env(const char * name, size_t fallback) { @@ -66,9 +71,10 @@ std::string json_escape(std::string_view value) { void write_u32_array( std::ostream & out, - const std::array & values) { + const std::array & values, + size_t count) { out << '['; - for (size_t i = 0; i < values.size(); ++i) { + for (size_t i = 0; i < count; ++i) { if (i) out << ','; out << values[i]; } @@ -76,6 +82,14 @@ void write_u32_array( } void write_step_json(std::ostream & out, const StepProfile & step) { + size_t position_count = 0; + for (size_t i = kMaxSpecPositions; i > 0; --i) { + if (step.proposed_by_position[i - 1] != 0 || + step.accepted_by_position[i - 1] != 0) { + position_count = i; + break; + } + } out << "{\"type\":\"step\",\"schema_version\":" << step.schema_version << ",\"round_id\":" << step.round_id @@ -122,9 +136,9 @@ void write_step_json(std::ostream & out, const StepProfile & step) { << ",\"dropped_lanes\":" << step.dropped_lanes << ",\"dropped_phases\":" << step.dropped_phases << ",\"proposed_by_position\":"; - write_u32_array(out, step.proposed_by_position); + write_u32_array(out, step.proposed_by_position, position_count); out << ",\"accepted_by_position\":"; - write_u32_array(out, step.accepted_by_position); + write_u32_array(out, step.accepted_by_position, position_count); out << ",\"lanes\":["; for (uint32_t i = 0; i < step.lane_count; ++i) { if (i) out << ','; @@ -471,8 +485,10 @@ bool ObservabilityState::write_capture(bool complete) { out << "{\"type\":\"request\",\"request_id\":" << request.request_id << ",\"response_id\":\"" << json_escape(request.response_id) - << "\",\"ok\":" << (request.ok ? "true" : "false") - << ",\"prompt_tokens\":" << request.prompt_tokens + << "\",\"ok\":"; + if (request.completed_ns == 0) out << "null"; + else out << (request.ok ? "true" : "false"); + out << ",\"prompt_tokens\":" << request.prompt_tokens << ",\"output_tokens\":" << request.output_tokens << ",\"queued_ns\":" << request.queued_ns << ",\"admitted_ns\":" << request.admitted_ns diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 6a2fee6bd..074573333 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -476,10 +476,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { return AdmissionDisposition::Retired; } next_request_id++; - job->profile_request_id = request_id; - if (observability_.enabled()) { - job->profile_admitted_ns = observability::steady_time_ns(); - } + const uint64_t profile_admitted_ns = observability_.enabled() + ? observability::steady_time_ns() : 0; SchedSlot & s = slots[(size_t)ar.slot]; s = SchedSlot{}; @@ -505,7 +503,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { observability_.record_request_admitted( request_id, req.response_id, static_cast(req.prompt_tokens.size()), - job->profile_queued_ns, job->profile_admitted_ns); + job->profile_queued_ns, profile_admitted_ns); } publish_live_count(); return AdmissionDisposition::Admitted; diff --git a/server/test/test_observability.cpp b/server/test/test_observability.cpp index 330bf6b49..8edc743ef 100644 --- a/server/test/test_observability.cpp +++ b/server/test/test_observability.cpp @@ -1,6 +1,7 @@ #include "server/observability.h" #include +#include #include #include #include @@ -17,6 +18,18 @@ using namespace dflash::common::observability; } while (false) int main() { +#if defined(_WIN32) + _putenv_s("DFLASH_PROF_MAX_ROUNDS", "-1"); +#else + setenv("DFLASH_PROF_MAX_ROUNDS", "-1", 1); +#endif + CHECK(ObservabilityConfig::from_env().max_rounds == 10000); +#if defined(_WIN32) + _putenv_s("DFLASH_PROF_MAX_ROUNDS", ""); +#else + unsetenv("DFLASH_PROF_MAX_ROUNDS"); +#endif + ObservabilityState disabled({}); CHECK(disabled.job_queued() == 0); CHECK(disabled.queue_depth() == 0); @@ -65,6 +78,9 @@ int main() { step->spec_attempted_lanes = 4; step->spec_proposed_draft_tokens = 12; step->spec_accepted_draft_tokens = 8; + step->spec_tree_width = 3; + step->proposed_by_position[1] = 2; + step->accepted_by_position[1] = 1; step->kv_blocks_total = 100; step->kv_blocks_free_after = 80; LaneProfile lane; @@ -75,7 +91,6 @@ int main() { step->add_phase({Phase::TargetCompute, 1, 20}); state.record_prefill_completed(7, queued_ns + 20); state.record_token_burst(7, step->round_id, queued_ns + 30, 3); - state.record_request_finished(7, true, 3, queued_ns + 40); state.commit_step(step); { std::ifstream checkpoint_input(output); @@ -89,7 +104,13 @@ int main() { std::string::npos); CHECK(checkpoint.find("\"DFLASH_DRAFT_KV\":\"1\"") != std::string::npos); + CHECK(checkpoint.find("\"ok\":null") != std::string::npos); + CHECK(checkpoint.find( + "\"proposed_by_position\":[0,2]," + "\"accepted_by_position\":[0,1]") != + std::string::npos); } + state.record_request_finished(7, true, 3, queued_ns + 40); StepProfile * dropped = state.begin_step(4); dropped->kv_blocks_total = 100; @@ -110,6 +131,7 @@ int main() { std::string::npos); CHECK(jsonl.find("\"type\":\"request\"") != std::string::npos); CHECK(jsonl.find("\"type\":\"token_burst\"") != std::string::npos); + CHECK(jsonl.find("\"ok\":true") != std::string::npos); CHECK(jsonl.find("\"dropped_steps\":1") != std::string::npos); CHECK(jsonl.find("\"complete\":true") != std::string::npos); std::filesystem::remove(output); From 1b209ecee8977a5e2c758f73ccecf700d73a5025 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 22 Aug 2026 09:04:03 +0000 Subject: [PATCH 08/11] concurrency: correlate service rounds with rocprof --- harness/benchmarks/concurrency/README.md | 4 + .../concurrency/rocprof_server_wrapper.sh | 56 +++++++++++ .../concurrency/test_concurrency_tools.py | 91 ++++++++++++++++++ server/docs/CONCURRENCY_OBSERVABILITY.md | 96 ++++++++++++++++++- server/docs/ENVIRONMENT.md | 2 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 29 +++++- server/src/qwen35/qwen35_roctx.cpp | 26 ++++- server/src/qwen35/qwen35_roctx.h | 11 ++- server/test/test_qwen35_roctx.cpp | 20 +++- 9 files changed, 320 insertions(+), 15 deletions(-) create mode 100755 harness/benchmarks/concurrency/rocprof_server_wrapper.sh diff --git a/harness/benchmarks/concurrency/README.md b/harness/benchmarks/concurrency/README.md index af6bfd365..ad5ac6e1a 100644 --- a/harness/benchmarks/concurrency/README.md +++ b/harness/benchmarks/concurrency/README.md @@ -5,6 +5,10 @@ prefill and concurrent decode. It has two user-facing runners that share the streaming client, deterministic prompt generation, and summary tooling: the paired ragged workload runner and the canonical-suite runner. +To drill from a slow high-level service round into kernels, launches, copies, +graph replay/capture, or device idle gaps, use the local rocprof workflow in +[Concurrency observability](../../../server/docs/CONCURRENCY_OBSERVABILITY.md#drill-into-one-target-round). + ## Canonical and blog workloads The synthetic ragged profiles below isolate serving mechanics. To measure how diff --git a/harness/benchmarks/concurrency/rocprof_server_wrapper.sh b/harness/benchmarks/concurrency/rocprof_server_wrapper.sh new file mode 100755 index 000000000..31c5b7904 --- /dev/null +++ b/harness/benchmarks/concurrency/rocprof_server_wrapper.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Run a Lucebox server under a delayed rocprofv3 collection window while also +# keeping the high-level concurrency capture in the same artifact directory. + +PROFILED_SERVER_BIN="${PROFILED_SERVER_BIN:?set PROFILED_SERVER_BIN}" +ROCPROF_OUTPUT_DIR="${ROCPROF_OUTPUT_DIR:?set ROCPROF_OUTPUT_DIR}" +ROCPROF_START_SECONDS="${ROCPROF_START_SECONDS:-180}" +ROCPROF_DURATION_SECONDS="${ROCPROF_DURATION_SECONDS:-90}" + +if [[ ! "$ROCPROF_START_SECONDS" =~ ^[0-9]+$ ]] || + [[ ! "$ROCPROF_DURATION_SECONDS" =~ ^[1-9][0-9]*$ ]]; then + echo "ROCPROF_START_SECONDS must be non-negative and " \ + "ROCPROF_DURATION_SECONDS must be positive integer seconds" >&2 + exit 2 +fi + +if [[ ! -x "$PROFILED_SERVER_BIN" ]]; then + echo "PROFILED_SERVER_BIN is not executable: $PROFILED_SERVER_BIN" >&2 + exit 2 +fi + +if [[ -n "${ROCPROF_BIN:-}" ]]; then + rocprof_bin="$ROCPROF_BIN" +elif rocprof_bin="$(command -v rocprofv3 2>/dev/null)" && + [[ -n "$rocprof_bin" ]]; then + : +else + rocprof_bin="${ROCM_PATH:-/opt/rocm}/bin/rocprofv3" +fi +if [[ ! -x "$rocprof_bin" ]]; then + echo "rocprofv3 is not executable: $rocprof_bin" >&2 + exit 2 +fi + +mkdir -p "$ROCPROF_OUTPUT_DIR" +export DFLASH_PROF=concurrency +export DFLASH_PROF_OUT="${DFLASH_PROF_OUT:-$ROCPROF_OUTPUT_DIR/profile.jsonl}" +export DFLASH_QWEN35_ROCTX=1 + +exec "$rocprof_bin" \ + --marker-trace \ + --kernel-trace \ + --memory-copy-trace \ + --hip-runtime-trace \ + --group-by-queue true \ + --stats \ + --summary \ + --summary-output-file "$ROCPROF_OUTPUT_DIR/summary.txt" \ + --collection-period \ + "${ROCPROF_START_SECONDS}:${ROCPROF_DURATION_SECONDS}:1" \ + --output-format csv pftrace \ + --output-directory "$ROCPROF_OUTPUT_DIR" \ + --output-file trace \ + -- "$PROFILED_SERVER_BIN" "$@" diff --git a/harness/benchmarks/concurrency/test_concurrency_tools.py b/harness/benchmarks/concurrency/test_concurrency_tools.py index 565884a76..14c9ef425 100644 --- a/harness/benchmarks/concurrency/test_concurrency_tools.py +++ b/harness/benchmarks/concurrency/test_concurrency_tools.py @@ -88,6 +88,97 @@ def test_client_level_parser_rejects_reuse(self) -> None: class RunnerTests(unittest.TestCase): + def test_rocprof_wrapper_exports_capture_contract(self) -> None: + wrapper = HERE / "rocprof_server_wrapper.sh" + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + output = root / "capture" + invocation = root / "invocation.json" + fake_rocprof = root / "rocprofv3" + fake_server = root / "dflash_server" + fake_rocprof.write_text( + """#!/usr/bin/env python3 +import json +import os +import sys + +payload = { + "argv": sys.argv[1:], + "env": { + name: os.environ.get(name) + for name in ( + "DFLASH_PROF", + "DFLASH_PROF_OUT", + "DFLASH_QWEN35_ROCTX", + ) + }, +} +with open(os.environ["FAKE_ROCPROF_INVOCATION"], "w", encoding="utf-8") as out: + json.dump(payload, out) +""", + encoding="utf-8", + ) + fake_server.write_text("#!/usr/bin/env sh\nexit 0\n", encoding="utf-8") + fake_rocprof.chmod(0o755) + fake_server.chmod(0o755) + env = { + "PATH": os.environ.get("PATH", ""), + "PROFILED_SERVER_BIN": str(fake_server), + "ROCPROF_OUTPUT_DIR": str(output), + "ROCPROF_BIN": str(fake_rocprof), + "ROCPROF_START_SECONDS": "7", + "ROCPROF_DURATION_SECONDS": "11", + "FAKE_ROCPROF_INVOCATION": str(invocation), + } + + result = subprocess.run( + [str(wrapper), "--model", "qwen"], + env=env, + text=True, + capture_output=True, + check=False, + ) + + self.assertEqual(result.returncode, 0, result.stderr) + self.assertTrue(output.is_dir()) + payload = json.loads(invocation.read_text(encoding="utf-8")) + self.assertEqual( + payload["env"], + { + "DFLASH_PROF": "concurrency", + "DFLASH_PROF_OUT": str(output / "profile.jsonl"), + "DFLASH_QWEN35_ROCTX": "1", + }, + ) + self.assertEqual( + payload["argv"], + [ + "--marker-trace", + "--kernel-trace", + "--memory-copy-trace", + "--hip-runtime-trace", + "--group-by-queue", + "true", + "--stats", + "--summary", + "--summary-output-file", + str(output / "summary.txt"), + "--collection-period", + "7:11:1", + "--output-format", + "csv", + "pftrace", + "--output-directory", + str(output), + "--output-file", + "trace", + "--", + str(fake_server), + "--model", + "qwen", + ], + ) + def test_runners_isolate_and_record_selected_gpu(self) -> None: for script_name in ( "run_qwen36_concurrency.sh", diff --git a/server/docs/CONCURRENCY_OBSERVABILITY.md b/server/docs/CONCURRENCY_OBSERVABILITY.md index 4480685a9..357e36760 100644 --- a/server/docs/CONCURRENCY_OBSERVABILITY.md +++ b/server/docs/CONCURRENCY_OBSERVABILITY.md @@ -102,9 +102,99 @@ Instrumentation does not add device synchronization. A target compute or readback span therefore reflects the synchronization behavior already present in that code path. -Use the built-in profile to choose the next experiment. Use ROCTX and rocprof -afterward when the question becomes kernel scheduling, memory bandwidth, or a -specific device operation. +`target_compute` measures the host call that submits a graph. Device work can +continue after that scope and finish while the host is in `argmax_readback`. +Do not treat the host span as GPU busy time. The kernel, queue, and memory-copy +tracks in a native rocprof trace are authoritative for device timing and idle +gaps. + +## Drill into one target round + +Use the built-in profile to find an expensive service round, then use its +`round_id` to locate the same Qwen graph submission in a native rocprof trace. +For example, start a local C=4 server under the repository wrapper: + +```bash +PROFILED_SERVER_BIN=/absolute/path/to/dflash_server \ +ROCPROF_OUTPUT_DIR=/tmp/lucebox-c4-rocprof \ +ROCPROF_START_SECONDS=180 \ +ROCPROF_DURATION_SECONDS=60 \ +harness/benchmarks/concurrency/rocprof_server_wrapper.sh \ + +``` + +In another terminal, wait for the server to become healthy and run the normal +C=4 workload during the collection window. Adjust the start delay to cover +model loading and warmup on the target machine. The wrapper enables the +high-level concurrency capture and Qwen ROCTX markers, and writes local +artifacts under `ROCPROF_OUTPUT_DIR`. + +Build the high-level report after stopping the server normally: + +```bash +python3 harness/benchmarks/concurrency/profile_report.py \ + /tmp/lucebox-c4-rocprof/profile.jsonl \ + --markdown /tmp/lucebox-c4-rocprof/profile.md \ + --perfetto /tmp/lucebox-c4-rocprof/profile.perfetto.json \ + --json-summary /tmp/lucebox-c4-rocprof/profile.summary.json +``` + +Choose a `round_id` whose `target_compute` phase or cohort is interesting. +Open the generated `.pftrace` artifact at +[ui.perfetto.dev](https://ui.perfetto.dev), then search for +`qwen35.graph_compute round_id=`. The marker also carries the packed +or speculative path and the relevant live, bucket, row, prefill, and KV-length +shape. + +Use the native tracks to answer the device-level question: + +- Raw kernel names and durations provide evidence about which operations + dominate. Names are backend implementation details, not a stable + attention/GDN/expert classification. +- HIP runtime calls beside dispatches expose host launch overhead and launch + queues. +- Kernel tracks inside and between round markers show real device busy and + idle gaps. +- Memory-copy tracks show transfers on the critical path. +- Graph launch, capture, instantiate, and update APIs distinguish replay from + capture for that traced round. + +The built-in `profile.perfetto.json` and rocprof `.pftrace` artifact remain +separate by design. The former explains serving policy and request lifecycle; +the latter owns the device timeline. `round_id` joins them without translating +independent clocks or importing profiler-version-specific CSV schemas. + +`GGML_CUDA_GRAPH_STATS=1` is an optional aggregate cross-check. Set it before +the wrapper, and optionally set `GGML_CUDA_GRAPH_STATS_EVERY=1` for a short +diagnostic run. The counters are cumulative per graph key and cannot identify +an exact round; use the HIP graph APIs in the native trace for that. Stats are +emitted only by a graph-enabled build when the exercised path has a graph key. +The wrapper does not force graph stats on. + +Kernel tracing does not measure occupancy or memory bandwidth. Counter names +and compatible groups depend on the target GPU, so discover and validate them +there instead of hard-coding a repository default: + +```bash +rocprofv3-avail --device 0 list --pmc +rocprofv3-avail --device 0 pmc-check [ ...] +``` + +Run a separate pass with the validated counters and write it to a different +local output directory: + +```bash +rocprofv3 \ + --pmc \ + --output-format csv \ + --output-directory /tmp/lucebox-c4-pmc \ + --output-file counters \ + -- +``` + +Do not combine this counter pass with the timing trace. Hardware counters and +full tracing can each perturb execution; use these runs to diagnose a chosen +operation, not as throughput benchmark results. ## Extend another model diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 12ae6ce8b..88a88dcf0 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -49,7 +49,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_DS4_TP_FUSED_CACHE_SLOTS` | 2 | BURN-IN: number of heterogeneous verifier schedulers retained; higher values retain substantially more scratch on both GPUs. | | `DFLASH_DS4_VERIFY_FORCE_GRAPH_REPLAY` | unset | OPT-IN: bypass graph property scans only after warmup; scheduler-generation checks remain mandatory. | | `DFLASH_DS4_ROCTX` | unset | DEBUG: on HIP builds, dynamically load ROCTX and emit semantic DS4 prefill, speculative-decode, and layer-range markers for external rocprof traces. No events, timing, or device synchronization are added. | -| `DFLASH_QWEN35_ROCTX` | unset | DEBUG: on HIP builds, dynamically load ROCTX and mark Qwen concurrent steps, graph compute, and argmax readback with live, padded, and packed-prefill shape metadata. | +| `DFLASH_QWEN35_ROCTX` | unset | DEBUG: on HIP builds, dynamically load ROCTX and mark Qwen concurrent steps, graph compute, and argmax readback with live, padded, and packed-prefill shape metadata. With `DFLASH_PROF=concurrency`, markers also carry the matching high-level round ID and packed/speculative path. See [CONCURRENCY_OBSERVABILITY.md](CONCURRENCY_OBSERVABILITY.md#drill-into-one-target-round). | | `GGML_DS4_FA_SERIAL_INDEX_SCAN` | unset | DEBUG/A-B: restore the serial indexed-attention mask scan instead of the long-context HIP parallel scan. | | `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` | 1 for qualified long heterogeneous prefill | KILL SWITCH: =0 restores per-layer route/owner scratch allocation. | | `DFLASH_MOE_TP_*` / `DFLASH_MOE_HYBRID_PREFILL_EAGER` | unset | BURN-IN: model-neutral names for common heterogeneous-MoE scheduling and kernel policy. Existing `DFLASH_DS4_*` names remain compatibility aliases. | diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 1a64c880d..8f7e085dc 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -657,6 +657,19 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( for (const ArLane & lane : ar_lanes) { max_prefix = std::max(max_prefix, lane.position + 1); } + const Qwen35RoctxMetadata roctx_metadata{ + profile ? profile->round_id : 0, + "speculative", + tree_width, + static_cast(inputs.size()), + tree_bucket, + 0, + 0, + total_rows, + max_prefix, + }; + const Qwen35RoctxRange roctx_step( + "qwen35.concurrent_step", roctx_metadata); StepGraph & graph = b_.sg_; bool graph_built = false; @@ -798,6 +811,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( { observability::PhaseScope phase( profile, observability::Phase::TargetCompute); + const Qwen35RoctxRange roctx_compute( + "qwen35.graph_compute", roctx_metadata); target_status = ggml_backend_graph_compute( b_.target_backend_, graph.gf); } @@ -811,6 +826,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( { observability::PhaseScope phase( profile, observability::Phase::ReadbackSync); + const Qwen35RoctxRange roctx_sync( + "qwen35.argmax_readback", roctx_metadata); ggml_backend_tensor_get( graph.argmax_tokens, posterior.data(), 0, sizeof(int32_t) * posterior.size()); @@ -1351,8 +1368,16 @@ SeqEngine::StepResult Qwen35SeqEngine::step( profile->target_forwards = 1; } const Qwen35RoctxMetadata roctx_metadata{ - live_count, decode_bucket, n_prefill, (int)segments.size(), - n_total, max_kv_len}; + profile ? profile->round_id : 0, + "packed", + tree_width_, + live_count, + decode_bucket, + n_prefill, + static_cast(segments.size()), + n_total, + max_kv_len, + }; const Qwen35RoctxRange roctx_step("qwen35.concurrent_step", roctx_metadata); const int gather_rows = with_prefill ? (with_decode ? n_commits + decode_bucket diff --git a/server/src/qwen35/qwen35_roctx.cpp b/server/src/qwen35/qwen35_roctx.cpp index 016436723..5f426e879 100644 --- a/server/src/qwen35/qwen35_roctx.cpp +++ b/server/src/qwen35/qwen35_roctx.cpp @@ -92,6 +92,23 @@ void append(char * message, size_t capacity, size_t & used, const char * name, i const int n = std::snprintf(message + used, capacity - used, " %s=%d", name, value); if (n > 0) used += std::min((size_t)n, capacity - used - 1); } + +void append_u64(char * message, size_t capacity, size_t & used, + const char * name, uint64_t value) { + if (value == 0 || used >= capacity) return; + const int n = std::snprintf( + message + used, capacity - used, " %s=%llu", name, + static_cast(value)); + if (n > 0) used += std::min((size_t)n, capacity - used - 1); +} + +void append_text(char * message, size_t capacity, size_t & used, + const char * name, const char * value) { + if (!value || !value[0] || used >= capacity) return; + const int n = std::snprintf( + message + used, capacity - used, " %s=%s", name, value); + if (n > 0) used += std::min((size_t)n, capacity - used - 1); +} } // namespace bool qwen35_roctx_env_enabled(const char * value) { @@ -108,11 +125,14 @@ Qwen35RoctxRange::Qwen35RoctxRange(const char * scope, const Qwen35RoctxMetadata char message[256]; const int initial = std::snprintf(message, sizeof(message), "%s", scope); size_t used = initial > 0 ? std::min((size_t)initial, sizeof(message) - 1) : 0; - append(message, sizeof(message), used, "live", metadata.live); - append(message, sizeof(message), used, "bucket", metadata.bucket); + append_u64(message, sizeof(message), used, "round_id", metadata.round_id); + append_text(message, sizeof(message), used, "path", metadata.path); + append(message, sizeof(message), used, "spec_tree_width", metadata.spec_tree_width); + append(message, sizeof(message), used, "live_slots", metadata.live_slots); + append(message, sizeof(message), used, "decode_bucket", metadata.decode_bucket); append(message, sizeof(message), used, "prefill_tokens", metadata.prefill_tokens); append(message, sizeof(message), used, "prefill_segments", metadata.prefill_segments); - append(message, sizeof(message), used, "total_rows", metadata.total_rows); + append(message, sizeof(message), used, "target_rows", metadata.target_rows); append(message, sizeof(message), used, "max_kv_len", metadata.max_kv_len); if (callbacks.push(message) >= 0) { pop_ = callbacks.pop; pushed_ = true; } } diff --git a/server/src/qwen35/qwen35_roctx.h b/server/src/qwen35/qwen35_roctx.h index 5935b3c20..1127d4951 100644 --- a/server/src/qwen35/qwen35_roctx.h +++ b/server/src/qwen35/qwen35_roctx.h @@ -1,13 +1,18 @@ #pragma once +#include + namespace dflash::common { struct Qwen35RoctxMetadata { - int live = -1; - int bucket = -1; + uint64_t round_id = 0; + const char * path = nullptr; + int spec_tree_width = -1; + int live_slots = -1; + int decode_bucket = -1; int prefill_tokens = -1; int prefill_segments = -1; - int total_rows = -1; + int target_rows = -1; int max_kv_len = -1; }; diff --git a/server/test/test_qwen35_roctx.cpp b/server/test/test_qwen35_roctx.cpp index 98e856788..303fe0d9f 100644 --- a/server/test/test_qwen35_roctx.cpp +++ b/server/test/test_qwen35_roctx.cpp @@ -27,15 +27,29 @@ int main() { { Qwen35RoctxRange range( - "qwen35.graph_compute", {9, 12, 64, 2, 76, 511}, true, + "qwen35.graph_compute", + {42, "speculative", 8, 9, 12, 64, 2, 76, 511}, true, {push, pop}); CHECK(events.size() == 1); - CHECK(events[0] == "qwen35.graph_compute live=9 bucket=12 " - "prefill_tokens=64 prefill_segments=2 total_rows=76 " + CHECK(events[0] == "qwen35.graph_compute round_id=42 path=speculative " + "spec_tree_width=8 live_slots=9 decode_bucket=12 " + "prefill_tokens=64 prefill_segments=2 target_rows=76 " "max_kv_len=511"); } CHECK(events.size() == 2 && events[1] == "pop"); + events.clear(); + { + Qwen35RoctxRange range( + "qwen35.graph_compute", {0, nullptr, 8, 4, 4, 0, 0, 4, 128}, true, + {push, pop}); + CHECK(events.size() == 1); + CHECK(events[0] == "qwen35.graph_compute spec_tree_width=8 live_slots=4 " + "decode_bucket=4 prefill_tokens=0 prefill_segments=0 " + "target_rows=4 max_kv_len=128"); + } + CHECK(events.size() == 2 && events[1] == "pop"); + events.clear(); { Qwen35RoctxRange disabled("qwen35.graph_compute", {}, false, {push, pop}); } CHECK(events.empty()); From 22baa51ba6be33155cdbb19289a04528c44f9493 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 22 Aug 2026 09:40:07 +0000 Subject: [PATCH 09/11] concurrency: add folded profile views --- .../benchmarks/concurrency/profile_report.py | 141 +++++++++++++++++- .../concurrency/test_profile_report.py | 76 +++++++++- server/docs/CONCURRENCY_OBSERVABILITY.md | 23 ++- 3 files changed, 237 insertions(+), 3 deletions(-) diff --git a/harness/benchmarks/concurrency/profile_report.py b/harness/benchmarks/concurrency/profile_report.py index 978feffdd..2dc136f95 100644 --- a/harness/benchmarks/concurrency/profile_report.py +++ b/harness/benchmarks/concurrency/profile_report.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Summarize a Lucebox concurrency profile and emit a Perfetto trace.""" +"""Build reports from a Lucebox concurrency profile.""" from __future__ import annotations @@ -11,6 +11,9 @@ from pathlib import Path from typing import Any, Iterable +FOLDED_DIMENSIONS = frozenset(("path", "cohort", "phase")) +DEFAULT_FOLDED_STACK = ("path", "cohort", "phase") + def load_records(path: Path) -> list[dict[str, Any]]: records: list[dict[str, Any]] = [] @@ -462,12 +465,139 @@ def build_perfetto(records: list[dict[str, Any]]) -> dict[str, Any]: return {"displayTimeUnit": "ms", "traceEvents": events} +def parse_folded_stack(value: str) -> tuple[str, ...]: + dimensions = tuple(value.split(",")) + if len(dimensions) != 3 or set(dimensions) != FOLDED_DIMENSIONS: + raise argparse.ArgumentTypeError( + "stack must be a permutation of path,cohort,phase" + ) + return dimensions + + +def folded_atom(value: Any) -> str: + atom = str(value) + if not atom or ";" in atom or any(character.isspace() for character in atom): + raise ValueError(f"invalid folded-stack frame: {atom!r}") + return atom + + +def step_phase_buckets(step: dict[str, Any]) -> Counter[str]: + duration_ns = max(0, int(step.get("duration_ns", 0))) + events: dict[int, Counter[str]] = defaultdict(Counter) + for span in step.get("phases", []): + phase = folded_atom(span.get("phase", "unknown")) + raw_start = int(span.get("start_offset_ns", 0)) + raw_end = raw_start + max(0, int(span.get("duration_ns", 0))) + start = min(duration_ns, max(0, raw_start)) + end = min(duration_ns, max(0, raw_end)) + if end <= start: + continue + events[start][phase] += 1 + events[end][phase] -= 1 + + buckets: Counter[str] = Counter() + active: Counter[str] = Counter() + previous = 0 + for offset in sorted({0, duration_ns, *events}): + if offset > previous: + phases = sorted( + phase for phase, count in active.items() if count > 0 + ) + if not phases: + label = "unattributed" + elif len(phases) == 1: + label = phases[0] + else: + label = f"overlap({'+'.join(phases)})" + buckets[label] += offset - previous + active.update(events[offset]) + previous = offset + return buckets + + +def format_folded_weight(duration_ns: int, tokens: int | None) -> str: + if tokens is None or duration_ns % tokens == 0: + return str(duration_ns if tokens is None else duration_ns // tokens) + return f"{duration_ns / tokens:.6f}".rstrip("0").rstrip(".") + + +def build_folded( + records: list[dict[str, Any]], + *, + stack: tuple[str, ...] = DEFAULT_FOLDED_STACK, + per_token: bool = False, +) -> str: + if len(stack) != 3 or set(stack) != FOLDED_DIMENSIONS: + raise ValueError("stack must be a permutation of path,cohort,phase") + + durations: Counter[tuple[str, int, str]] = Counter() + tokens: Counter[tuple[str, int]] = Counter() + steps = [record for record in records if record["type"] == "step"] + for step in steps: + path = folded_atom(step.get("path", "unknown")) + cohort = max(0, int(step.get("live_slots", 0))) + group = (path, cohort) + tokens[group] += sum( + max(0, int(lane.get("scheduler_consumed_tokens", 0))) + for lane in step.get("lanes", []) + if lane.get("kind") == "decode" + ) + for phase, duration_ns in step_phase_buckets(step).items(): + durations[path, cohort, phase] += duration_ns + + lines: list[str] = [] + for (path, cohort, phase), duration_ns in durations.items(): + denominator = tokens[path, cohort] if per_token else None + if per_token and denominator == 0: + continue + frames = { + "path": path, + "cohort": f"C={cohort}", + "phase": phase, + } + folded_stack = ";".join(frames[dimension] for dimension in stack) + lines.append( + f"{folded_stack} {format_folded_weight(duration_ns, denominator)}" + ) + + if not per_token and len(steps) > 1: + inter_round_ns = 0 + ordered_steps = sorted( + steps, + key=lambda step: ( + int(step.get("started_ns", 0)), + int(step.get("round_id", 0)), + ), + ) + previous_end = ( + int(ordered_steps[0].get("started_ns", 0)) + + max(0, int(ordered_steps[0].get("duration_ns", 0))) + ) + for step in ordered_steps[1:]: + started_ns = int(step.get("started_ns", 0)) + inter_round_ns += max(0, started_ns - previous_end) + previous_end = max( + previous_end, + started_ns + max(0, int(step.get("duration_ns", 0))), + ) + if inter_round_ns: + lines.append(f"idle;inter_round {inter_round_ns}") + + return "" if not lines else "\n".join(sorted(lines)) + "\n" + + def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("profile", type=Path) parser.add_argument("--markdown", type=Path) parser.add_argument("--perfetto", type=Path) parser.add_argument("--json-summary", type=Path) + parser.add_argument("--folded", type=Path) + parser.add_argument("--folded-per-token", type=Path) + parser.add_argument( + "--stack", type=parse_folded_stack, default=DEFAULT_FOLDED_STACK, + metavar="path,cohort,phase", + ) args = parser.parse_args() records = load_records(args.profile) @@ -487,6 +617,15 @@ def main() -> int: json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) + if args.folded: + args.folded.write_text( + build_folded(records, stack=args.stack), encoding="utf-8" + ) + if args.folded_per_token: + args.folded_per_token.write_text( + build_folded(records, stack=args.stack, per_token=True), + encoding="utf-8", + ) return 0 diff --git a/harness/benchmarks/concurrency/test_profile_report.py b/harness/benchmarks/concurrency/test_profile_report.py index d6ea85a4f..641495fb9 100644 --- a/harness/benchmarks/concurrency/test_profile_report.py +++ b/harness/benchmarks/concurrency/test_profile_report.py @@ -26,7 +26,8 @@ def records(self): "spec_accepted_draft_tokens": 8, "spec_durable_draft_tokens": 8, "spec_scheduler_consumed_tokens": 7, - "lanes": [{"kind": "decode", "spec": "selected"}], + "lanes": [{"kind": "decode", "spec": "selected", + "scheduler_consumed_tokens": 2}], "phases": [{"phase": "target_compute", "start_offset_ns": 100, "duration_ns": 1000}], }, @@ -99,6 +100,79 @@ def test_incomplete_request_is_not_failed(self): self.assertEqual(summary["capture"]["requests"], 2) self.assertEqual(summary["capture"]["failed_requests"], 0) + def test_folded_views_merge_frames_and_normalize_by_group(self): + records = self.records() + records.insert(2, { + "type": "step", "round_id": 2, "started_ns": 4_000_000, + "duration_ns": 1_000, "path": "speculative", + "live_slots": 4, + "lanes": [{"kind": "decode", "scheduler_consumed_tokens": 3}], + "phases": [{"phase": "target_compute", + "start_offset_ns": 0, "duration_ns": 1_000}], + }) + + wall = profile_report.build_folded(records) + per_token = profile_report.build_folded(records, per_token=True) + + self.assertIn("speculative;C=4;target_compute 2000\n", wall) + self.assertIn("speculative;C=4;unattributed 1999000\n", wall) + self.assertIn("idle;inter_round 1000000\n", wall) + self.assertEqual( + per_token, + "speculative;C=4;target_compute 400\n" + "speculative;C=4;unattributed 399800\n", + ) + + def test_folded_coverage_partitions_overlap(self): + step = { + "type": "step", "duration_ns": 100, "path": "packed", + "live_slots": 2, + "phases": [ + {"phase": "a", "start_offset_ns": 10, "duration_ns": 50}, + {"phase": "b", "start_offset_ns": 40, "duration_ns": 40}, + ], + } + + buckets = profile_report.step_phase_buckets(step) + + self.assertEqual(buckets, { + "unattributed": 30, + "a": 30, + "overlap(a+b)": 20, + "b": 20, + }) + self.assertEqual(sum(buckets.values()), step["duration_ns"]) + + step["phases"] = [ + {"phase": "a", "start_offset_ns": -10, "duration_ns": 60}, + {"phase": "a", "start_offset_ns": 40, "duration_ns": 80}, + ] + self.assertEqual( + profile_report.step_phase_buckets(step), {"a": 100} + ) + + def test_folded_stack_order_is_validated(self): + stack = profile_report.parse_folded_stack("cohort,path,phase") + folded = profile_report.build_folded(self.records(), stack=stack) + self.assertIn("C=4;speculative;target_compute 1000\n", folded) + + with self.assertRaisesRegex(ValueError, "permutation"): + profile_report.build_folded( + self.records(), stack=("path", "path", "phase") + ) + with self.assertRaisesRegex(Exception, "permutation"): + profile_report.parse_folded_stack("path,cohort,kind") + + def test_folded_per_token_omits_groups_without_decode_tokens(self): + records = self.records() + records[1]["lanes"] = [ + {"kind": "prefill", "scheduler_consumed_tokens": 1} + ] + + self.assertEqual( + profile_report.build_folded(records, per_token=True), "" + ) + if __name__ == "__main__": unittest.main() diff --git a/server/docs/CONCURRENCY_OBSERVABILITY.md b/server/docs/CONCURRENCY_OBSERVABILITY.md index 357e36760..2b0c34954 100644 --- a/server/docs/CONCURRENCY_OBSERVABILITY.md +++ b/server/docs/CONCURRENCY_OBSERVABILITY.md @@ -65,7 +65,9 @@ python3 harness/benchmarks/concurrency/profile_report.py \ /tmp/lucebox-profile.jsonl \ --markdown /tmp/lucebox-profile.md \ --perfetto /tmp/lucebox-profile.perfetto.json \ - --json-summary /tmp/lucebox-profile.summary.json + --json-summary /tmp/lucebox-profile.summary.json \ + --folded /tmp/lucebox-profile.folded \ + --folded-per-token /tmp/lucebox-profile.per-token.folded ``` Open the Perfetto JSON at [ui.perfetto.dev](https://ui.perfetto.dev). It shows @@ -74,6 +76,25 @@ The JSON summary has a versioned schema for benchmark diffs. It includes run context, latency percentiles, phase totals, padding ratios, concurrency cohorts, suppression decisions, and acceptance by speculative position. +The folded files use `path;C=;phase` stacks. Identical stacks +merge across rounds. Pass `--stack cohort,path,phase` to put the concurrency +cohort first. The flag accepts any permutation of `path`, `cohort`, and +`phase`. + +`--folded` writes host wall nanoseconds. The `unattributed` phase covers time +inside a round that has no phase span. The `idle;inter_round` stack covers +positive host-clock gaps between retained rounds. Neither value proves that +the device was idle. + +`--folded-per-token` divides each path and cohort's phase totals by the durable +decode tokens that the scheduler consumed for that group. The report omits a +group when it has no durable decode tokens. Inter-round gaps have no honest +path or cohort owner, so the per-token file omits them. + +The v1 capture does not contain model FLOP counts, weight bytes, or device +machine balance. Folded stacks also have no portable color metadata. Use a +separate device profile for compute-bound or bandwidth-bound classification. + ## Read the speculation funnel The profiler keeps these stages separate: From f56e2e5997f441bd1ae8c9c163d16fd2bae60700 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 22 Aug 2026 10:57:42 +0000 Subject: [PATCH 10/11] harness: add c1-c5 prompt cohort matrix --- .../concurrency/concurrent_benchmark.py | 2 +- .../concurrency/generate_prompts.py | 15 ++++++++++++--- .../concurrency/run_qwen36_concurrency.sh | 4 ++-- .../concurrency/test_concurrency_tools.py | 19 +++++++++++-------- .../concurrency/test_concurrent_benchmark.py | 4 ++-- 5 files changed, 28 insertions(+), 16 deletions(-) diff --git a/harness/benchmarks/concurrency/concurrent_benchmark.py b/harness/benchmarks/concurrency/concurrent_benchmark.py index c05303869..ddd8587a0 100755 --- a/harness/benchmarks/concurrency/concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/concurrent_benchmark.py @@ -17,7 +17,7 @@ from typing import Any PromptInput = str | list[dict[str, str]] -DEFAULT_CLIENT_LEVELS = (2, 4, 8, 16) +DEFAULT_CLIENT_LEVELS = (1, 2, 3, 4, 5) def sha256_text(text: str) -> str: diff --git a/harness/benchmarks/concurrency/generate_prompts.py b/harness/benchmarks/concurrency/generate_prompts.py index abafadab6..9cf2092dc 100755 --- a/harness/benchmarks/concurrency/generate_prompts.py +++ b/harness/benchmarks/concurrency/generate_prompts.py @@ -18,7 +18,11 @@ "long": (2000, 2600, 3400, 4000), } -DEFAULT_CLIENT_LEVELS = (2, 4, 8, 16) +CLIENT_MATRICES = { + "profile-view": (1, 2, 3, 4, 5), + "legacy": (2, 4, 8, 16), +} +DEFAULT_CLIENT_LEVELS = CLIENT_MATRICES["profile-view"] WORD_BANK = ( "systems engineers compare latency throughput scheduling memory kernels queues " @@ -126,16 +130,21 @@ def main() -> int: "--profile", choices=["he-raw", *sorted(RAGGED_PROFILES)], required=True ) parser.add_argument( - "--clients", default=",".join(map(str, DEFAULT_CLIENT_LEVELS)), + "--clients", help="comma-separated, distinct concurrency levels for disjoint cohorts", ) + parser.add_argument( + "--matrix", choices=sorted(CLIENT_MATRICES), default="profile-view", + help="named cohort matrix used when --clients is omitted", + ) parser.add_argument("--out", type=Path, required=True) args = parser.parse_args() if args.out.exists(): parser.error(f"refusing to overwrite {args.out}") args.out.parent.mkdir(parents=True, exist_ok=True) try: - client_levels = parse_client_levels(args.clients) + raw_levels = args.clients or ",".join(map(str, CLIENT_MATRICES[args.matrix])) + client_levels = parse_client_levels(raw_levels) except ValueError as exc: parser.error(str(exc)) records = build_records(args.profile, client_levels) diff --git a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh index 09a4e42c7..7729f9f0c 100755 --- a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh +++ b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh @@ -15,7 +15,7 @@ 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:-2,4,8,16}" +CLIENTS="${CLIENTS:-1,2,3,4,5}" SLOTS="${SLOTS:-}" GPU_DEVICE="${GPU_DEVICE:-0}" EXPECTED_GPU_ARCH="${EXPECTED_GPU_ARCH:-}" @@ -32,7 +32,7 @@ usage() { 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=2/4/8/16. Defaults to one repeat for screening; use at +luce-k1, and llama at C=1/2/3/4/5. 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. GPU_DEVICE is the physical ROCr device exposed exclusively to both diff --git a/harness/benchmarks/concurrency/test_concurrency_tools.py b/harness/benchmarks/concurrency/test_concurrency_tools.py index 14c9ef425..8020e69ad 100644 --- a/harness/benchmarks/concurrency/test_concurrency_tools.py +++ b/harness/benchmarks/concurrency/test_concurrency_tools.py @@ -30,15 +30,15 @@ def load(name: str): class PromptGeneratorTests(unittest.TestCase): def test_cohorts_are_disjoint_ragged_and_mean_matched(self) -> None: records = generator.build_records("short") - self.assertEqual(len(records), 30) + self.assertEqual(len(records), 15) self.assertEqual( [row["cohort"] for row in records], - ["c2"] * 2 + ["c4"] * 4 + ["c8"] * 8 + ["c16"] * 16, + ["c1"] + ["c2"] * 2 + ["c3"] * 3 + ["c4"] * 4 + ["c5"] * 5, ) - self.assertEqual(len({row["prompt"] for row in records}), 30) + self.assertEqual(len({row["prompt"] for row in records}), 15) by_cohort = { cohort: [row for row in records if row["cohort"] == cohort] - for cohort in ("c2", "c4", "c8", "c16") + for cohort in ("c1", "c2", "c3", "c4", "c5") } means = { cohort: sum(row["target_words"] for row in rows) / len(rows) @@ -47,13 +47,15 @@ def test_cohorts_are_disjoint_ragged_and_mean_matched(self) -> None: self.assertEqual(len(set(means.values())), 1) self.assertEqual( {cohort: rows[0]["cohort_offset"] for cohort, rows in by_cohort.items()}, - {"c2": 0, "c4": 2, "c8": 6, "c16": 14}, + {"c1": 0, "c2": 1, "c3": 3, "c4": 6, "c5": 10}, ) self.assertEqual( {row["target_words"] for row in by_cohort["c2"]}, {250, 550} ) - for cohort in ("c4", "c8", "c16"): - self.assertEqual(len({row["target_words"] for row in by_cohort[cohort]}), 4) + self.assertEqual({row["target_words"] for row in by_cohort["c1"]}, {400}) + self.assertEqual(len({row["target_words"] for row in by_cohort["c3"]}), 3) + self.assertEqual(len({row["target_words"] for row in by_cohort["c4"]}), 4) + self.assertEqual(len({row["target_words"] for row in by_cohort["c5"]}), 5) for row in records: self.assertEqual(len(row["prompt"].split()), row["target_words"]) @@ -80,6 +82,7 @@ def test_extended_matrix_is_deterministic_without_changing_existing_cohorts(self offset += clients def test_client_level_parser_rejects_reuse(self) -> None: + self.assertEqual(generator.CLIENT_MATRICES["legacy"], (2, 4, 8, 16)) self.assertEqual(generator.parse_client_levels("2,4,8,16,32"), (2, 4, 8, 16, 32)) with self.assertRaisesRegex(ValueError, "distinct"): generator.parse_client_levels("2,4,2") @@ -191,7 +194,7 @@ def test_runners_isolate_and_record_selected_gpu(self) -> None: def test_ragged_runner_derives_offsets_from_requested_matrix(self) -> None: text = (HERE / "run_qwen36_concurrency.sh").read_text(encoding="utf-8") - self.assertIn('CLIENTS="${CLIENTS:-2,4,8,16}"', text) + self.assertIn('CLIENTS="${CLIENTS:-1,2,3,4,5}"', text) self.assertIn('prompt_offsets[$c]="$next_prompt_offset"', text) self.assertIn('--clients "$CLIENTS"', text) self.assertNotIn("prompt_offsets=([", text) diff --git a/harness/benchmarks/concurrency/test_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_concurrent_benchmark.py index 3d0e94906..c2deb66f7 100644 --- a/harness/benchmarks/concurrency/test_concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/test_concurrent_benchmark.py @@ -28,8 +28,8 @@ def test_sse_parser_handles_events_and_done(self) -> None: ['{"choices":[{"delta":{"content":"hi"}}]}', "[DONE]"], ) - def test_default_matrix_starts_at_c2(self) -> None: - self.assertEqual(benchmark.DEFAULT_CLIENT_LEVELS, (2, 4, 8, 16)) + def test_default_matrix_matches_profile_view_range(self) -> None: + self.assertEqual(benchmark.DEFAULT_CLIENT_LEVELS, (1, 2, 3, 4, 5)) def test_run_rejects_duplicate_client_levels(self) -> None: args = argparse.Namespace(client_levels=[2, 4, 2]) From f47b35c4f00794fc056f70b066b75aa37ad0db34 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 22 Aug 2026 10:57:56 +0000 Subject: [PATCH 11/11] harness: generate offline LuceGraph reports --- harness/.gitignore | 1 + .../benchmarks/concurrency/device_specs.json | 35 + .../concurrency/generate_prompts.py | 6 +- .../benchmarks/concurrency/profile_payload.py | 782 ++++++++++++++++++ .../benchmarks/concurrency/profile_report.py | 354 +------- .../benchmarks/concurrency/profile_view.py | 570 +++++++++++++ .../concurrency/test_concurrency_tools.py | 1 + .../concurrency/test_concurrent_benchmark.py | 2 +- .../concurrency/test_profile_view.py | 375 +++++++++ server/docs/CONCURRENCY_OBSERVABILITY.md | 44 + 10 files changed, 1847 insertions(+), 323 deletions(-) create mode 100644 harness/benchmarks/concurrency/device_specs.json create mode 100644 harness/benchmarks/concurrency/profile_payload.py create mode 100644 harness/benchmarks/concurrency/profile_view.py create mode 100644 harness/benchmarks/concurrency/test_profile_view.py diff --git a/harness/.gitignore b/harness/.gitignore index 2df70216b..7830f864e 100644 --- a/harness/.gitignore +++ b/harness/.gitignore @@ -1,4 +1,5 @@ .harness-work/ results/ *.json +!benchmarks/concurrency/device_specs.json __pycache__/ diff --git a/harness/benchmarks/concurrency/device_specs.json b/harness/benchmarks/concurrency/device_specs.json new file mode 100644 index 000000000..39775b563 --- /dev/null +++ b/harness/benchmarks/concurrency/device_specs.json @@ -0,0 +1,35 @@ +{ + "devices": { + "gfx1151": { + "name": "Strix Halo", + "mem_bw_gbps": 256, + "fp16_tflops": 60, + "note": "Published peak FP16 vector throughput and shared-memory bandwidth. Bandwidth is shared with the CPU.", + "source_urls": [ + "https://www.amd.com/en/products/processors/desktops/ryzen/ryzen-ai-halo/ryzen-ai-max-plus-395.html", + "https://ir.amd.com/news-events/press-releases/detail/1270/amd-expands-ai-leadership-across-client-graphics-and-software-with-new-ryzen-ryzen-ai-and-amd-rocm-announcements-at-ces-2026" + ] + }, + "gfx1201": { + "name": "R9700", + "mem_bw_gbps": 640, + "fp16_tflops": 95.7, + "note": "Published peak FP16 vector throughput and memory bandwidth.", + "source_urls": [ + "https://www.amd.com/content/dam/amd/en/documents/partner-hub/radeon-pro/radeon-ai-pro-r9700-datasheet.pdf" + ] + } + }, + "models": { + "Qwen3.6-27B-Q4_K_M.gguf": { + "weight_bytes": 16800000000, + "active_params": 27000000000, + "kv_bytes_per_token_per_seq": 18432, + "note": "VERIFY weight_bytes against the deployed GGUF file. KV bytes are derived from the Q4_0 cache layout and must be rechecked if the cache type changes.", + "source_urls": [ + "https://huggingface.co/Qwen/Qwen3.6-27B/blob/main/README.md", + "https://huggingface.co/unsloth/Qwen3.6-27B-GGUF/blob/main/Qwen3.6-27B-Q4_K_M.gguf" + ] + } + } +} diff --git a/harness/benchmarks/concurrency/generate_prompts.py b/harness/benchmarks/concurrency/generate_prompts.py index 9cf2092dc..76a7aaac8 100755 --- a/harness/benchmarks/concurrency/generate_prompts.py +++ b/harness/benchmarks/concurrency/generate_prompts.py @@ -19,10 +19,10 @@ } CLIENT_MATRICES = { - "profile-view": (1, 2, 3, 4, 5), + "lucegraph": (1, 2, 3, 4, 5), "legacy": (2, 4, 8, 16), } -DEFAULT_CLIENT_LEVELS = CLIENT_MATRICES["profile-view"] +DEFAULT_CLIENT_LEVELS = CLIENT_MATRICES["lucegraph"] WORD_BANK = ( "systems engineers compare latency throughput scheduling memory kernels queues " @@ -134,7 +134,7 @@ def main() -> int: help="comma-separated, distinct concurrency levels for disjoint cohorts", ) parser.add_argument( - "--matrix", choices=sorted(CLIENT_MATRICES), default="profile-view", + "--matrix", choices=sorted(CLIENT_MATRICES), default="lucegraph", help="named cohort matrix used when --clients is omitted", ) parser.add_argument("--out", type=Path, required=True) diff --git a/harness/benchmarks/concurrency/profile_payload.py b/harness/benchmarks/concurrency/profile_payload.py new file mode 100644 index 000000000..afb9e25bb --- /dev/null +++ b/harness/benchmarks/concurrency/profile_payload.py @@ -0,0 +1,782 @@ +from __future__ import annotations + +import math +import statistics +from collections import Counter, defaultdict +from pathlib import Path +from typing import Any, Iterable, TypeAlias + +ReportPayload: TypeAlias = dict[str, Any] + +PHASE_ORDER = ( + "scheduler_plan", + "input_staging", + "draft_prepare", + "draft_compute", + "proposal_select", + "target_graph_build", + "metadata_upload", + "target_compute", + "readback_sync", + "acceptance", + "state_promotion", + "sampling_commit", + "output_processing", + "client_flush", +) + +FUNNEL_KEYS = ( + "spec_eligible_lanes", + "spec_reserved_lanes", + "spec_attempted_lanes", + "spec_proposed_draft_tokens", + "spec_verified_draft_tokens", + "spec_accepted_draft_tokens", + "spec_durable_draft_tokens", + "spec_scheduler_consumed_tokens", +) + + +def ratio(numerator: float, denominator: float) -> float: + return numerator / denominator if denominator else math.nan + + +def percentile(values: Iterable[float], quantile: float) -> float: + ordered = sorted(values) + if not ordered: + return math.nan + index = (len(ordered) - 1) * quantile + low = math.floor(index) + high = math.ceil(index) + if low == high: + return ordered[low] + return ordered[low] * (high - index) + ordered[high] * (index - low) + + +def optional_number(value: float) -> float | None: + return None if math.isnan(value) else value + + +def percentile_pair(values: Iterable[float]) -> dict[str, float | None]: + finite = [value for value in values if not math.isnan(value)] + return { + "p50": optional_number(percentile(finite, 0.50)), + "p95": optional_number(percentile(finite, 0.95)), + } + + +def duration_ms(end: int, start: int) -> float: + return (end - start) / 1_000_000 if end and start and end >= start else math.nan + + +def sum_field(records: Iterable[dict[str, Any]], key: str) -> int: + return sum(int(record.get(key, 0)) for record in records) + + +def folded_atom(value: Any) -> str: + atom = str(value) + if not atom or ";" in atom or any(character.isspace() for character in atom): + raise ValueError(f"invalid folded-stack frame: {atom!r}") + return atom + + +def step_phase_buckets(step: dict[str, Any]) -> Counter[str]: + duration_ns = max(0, int(step.get("duration_ns", 0))) + events: dict[int, Counter[str]] = defaultdict(Counter) + for span in step.get("phases", []): + phase = folded_atom(span.get("phase", "unknown")) + raw_start = int(span.get("start_offset_ns", 0)) + raw_end = raw_start + max(0, int(span.get("duration_ns", 0))) + start = min(duration_ns, max(0, raw_start)) + end = min(duration_ns, max(0, raw_end)) + if end <= start: + continue + events[start][phase] += 1 + events[end][phase] -= 1 + + buckets: Counter[str] = Counter() + active: Counter[str] = Counter() + previous = 0 + for offset in sorted({0, duration_ns, *events}): + if offset > previous: + phases = sorted( + phase for phase, count in active.items() if count > 0 + ) + if not phases: + label = "unattributed" + elif len(phases) == 1: + label = phases[0] + else: + label = f"overlap({'+'.join(phases)})" + buckets[label] += offset - previous + active.update(events[offset]) + previous = offset + return buckets + + +def phase_sort_key(phase: str) -> tuple[int, str]: + if phase in PHASE_ORDER: + return PHASE_ORDER.index(phase), phase + if phase == "unattributed": + return len(PHASE_ORDER) + 1, phase + if phase.startswith("overlap("): + return len(PHASE_ORDER) + 2, phase + return len(PHASE_ORDER), phase + + +def inter_round_idle(steps: list[dict[str, Any]]) -> dict[str, Any]: + ordered = sorted( + steps, + key=lambda step: ( + int(step.get("started_ns", 0)), + int(step.get("round_id", 0)), + ), + ) + gaps: list[int] = [] + if ordered: + previous_end = int(ordered[0].get("started_ns", 0)) + max( + 0, int(ordered[0].get("duration_ns", 0)) + ) + for step in ordered[1:]: + started_ns = int(step.get("started_ns", 0)) + gaps.append(max(0, started_ns - previous_end)) + previous_end = max( + previous_end, + started_ns + max(0, int(step.get("duration_ns", 0))), + ) + return { + "phase": "idle;inter_round", + "total_ns": sum(gaps), + "gaps": len(gaps), + "p50_ns": optional_number(percentile(gaps, 0.50)), + "p95_ns": optional_number(percentile(gaps, 0.95)), + "note": "Positive host-clock gaps between retained rounds. No cohort or token denominator applies.", + } + + +def _decode_tokens(step: dict[str, Any]) -> int: + return sum( + max(0, int(lane.get("scheduler_consumed_tokens", 0))) + for lane in step.get("lanes", []) + if lane.get("kind") == "decode" + ) + + +def _resolve_model_spec( + metadata: dict[str, Any], + catalog: dict[str, Any], + *, + draft: bool = False, +) -> tuple[str | None, dict[str, Any] | None]: + models = catalog.get("models", {}) + if not isinstance(models, dict): + return None, None + candidates: list[str] = [] + if draft: + draft_path = str(metadata.get("draft_path", "")) + if draft_path: + candidates.append(Path(draft_path).name) + else: + model_name = str(metadata.get("model_name", "")) + model_path = str(metadata.get("model_path", "")) + if model_name: + candidates.append(model_name) + if model_path: + candidates.append(Path(model_path).name) + for candidate in candidates: + spec = models.get(candidate) + if isinstance(spec, dict): + return candidate, spec + return None, None + + +def _classification( + phase: str, + steps: list[dict[str, Any]], + metadata: dict[str, Any], + catalog: dict[str, Any], + device_key: str | None, +) -> dict[str, Any]: + devices = catalog.get("devices", {}) + device = devices.get(device_key) if isinstance(devices, dict) and device_key else None + if not isinstance(device, dict): + return {"class": "neutral", "reason": "no device spec"} + if phase == "unattributed" or phase.startswith("idle;"): + return {"class": "idle", "reason": "idle or unattributed host time"} + if phase not in {"target_compute", "draft_compute"}: + return {"class": "overhead", "reason": "non-compute phase"} + + draft = phase == "draft_compute" + model_key, model = _resolve_model_spec(metadata, catalog, draft=draft) + if not isinstance(model, dict): + return { + "class": "overhead" if draft else "neutral", + "reason": "no draft model spec" if draft else "no target model spec", + } + + rounds = len(steps) + rows_key = "draft_rows" if draft else "target_rows" + padding_key = "draft_padding_rows" if draft else "target_padding_rows" + rows = sum_field(steps, rows_key) / rounds if rounds else 0.0 + mean_max_kv_len = sum_field(steps, "max_kv_len") / rounds if rounds else 0.0 + live_slots = int(steps[0].get("live_slots", 0)) if steps else 0 + weight_bytes = float(model["weight_bytes"]) + active_params = float(model["active_params"]) + kv_bytes = float(model["kv_bytes_per_token_per_seq"]) + bytes_moved = weight_bytes + kv_bytes * mean_max_kv_len * live_slots + flops = 2.0 * active_params * rows + arithmetic_intensity = flops / bytes_moved if bytes_moved else 0.0 + machine_balance = ( + float(device["fp16_tflops"]) * 1_000.0 / float(device["mem_bw_gbps"]) + ) + headroom = arithmetic_intensity / machine_balance if machine_balance else None + boundness = "bandwidth" if arithmetic_intensity < machine_balance else "compute" + total_rows = sum_field(steps, rows_key) + padding_fraction = ( + sum_field(steps, padding_key) / total_rows if total_rows else 0.0 + ) + padding_note = None + if padding_fraction: + if boundness == "bandwidth": + padding_note = "padding ≈ free when bandwidth-bound" + else: + padding_note = f"padding is {padding_fraction:.1%} of executed rows" + return { + "class": boundness, + "reason": "analytic roofline", + "device": device_key, + "device_name": device.get("name", device_key), + "model": model_key, + "arithmetic_intensity_flops_per_byte": arithmetic_intensity, + "machine_balance_flops_per_byte": machine_balance, + "headroom": headroom, + "estimated_flops": flops, + "estimated_bytes": bytes_moved, + "mean_rows": rows, + "mean_max_kv_len": mean_max_kv_len, + "padding_fraction": padding_fraction, + "padding_note": padding_note, + } + + +def _group_payload( + steps: list[dict[str, Any]], + path: str, + cohort: int, + metadata: dict[str, Any], + catalog: dict[str, Any], + device_key: str | None, +) -> dict[str, Any]: + per_round = [step_phase_buckets(step) for step in steps] + phase_names = sorted( + {phase for buckets in per_round for phase in buckets}, + key=phase_sort_key, + ) + round_count = len(steps) + token_count = sum(_decode_tokens(step) for step in steps) + total_ns = sum(max(0, int(step.get("duration_ns", 0))) for step in steps) + phases = [] + for phase in phase_names: + values = [int(buckets.get(phase, 0)) for buckets in per_round] + phase_total = sum(values) + phases.append({ + "phase": phase, + "total_ns": phase_total, + "ns_per_round": phase_total / round_count if round_count else None, + "ns_per_token": phase_total / token_count if token_count else None, + "wall_share": phase_total / total_ns if total_ns else None, + "p50_ns": optional_number(percentile(values, 0.50)), + "p95_ns": optional_number(percentile(values, 0.95)), + "classification": _classification( + phase, steps, metadata, catalog, device_key + ), + }) + return { + "path": path, + "cohort": cohort, + "label": f"C={cohort}", + "rounds": round_count, + "durable_tokens": token_count, + "total_ns": total_ns, + "mean_target_rows": sum_field(steps, "target_rows") / round_count, + "mean_draft_rows": sum_field(steps, "draft_rows") / round_count, + "mean_max_kv_len": sum_field(steps, "max_kv_len") / round_count, + "phases": phases, + } + + +def build_phase_groups( + records: list[dict[str, Any]], + *, + device_specs: dict[str, Any] | None = None, + device_key: str | None = None, +) -> dict[str, Any]: + metadata = next( + (record for record in records if record.get("type") == "metadata"), {} + ) + steps = [record for record in records if record.get("type") == "step"] + catalog = device_specs or {} + all_groups: dict[int, list[dict[str, Any]]] = defaultdict(list) + path_groups: dict[tuple[str, int], list[dict[str, Any]]] = defaultdict(list) + for step in steps: + path = folded_atom(step.get("path", "unknown")) + cohort = max(0, int(step.get("live_slots", 0))) + all_groups[cohort].append(step) + path_groups[path, cohort].append(step) + return { + "all": [ + _group_payload(group, "all", cohort, metadata, catalog, device_key) + for cohort, group in sorted(all_groups.items()) + ], + "paths": { + path: [ + _group_payload(group, path, cohort, metadata, catalog, device_key) + for (group_path, cohort), group in sorted(path_groups.items()) + if group_path == path + ] + for path in sorted({path for path, _ in path_groups}) + }, + "inter_round_idle": inter_round_idle(steps), + } + + +def format_folded_weight(duration_ns: int, tokens: int | None) -> str: + if tokens is None or duration_ns % tokens == 0: + return str(duration_ns if tokens is None else duration_ns // tokens) + return f"{duration_ns / tokens:.6f}".rstrip("0").rstrip(".") + + +def build_folded( + records: list[dict[str, Any]], + *, + stack: tuple[str, ...] = ("path", "cohort", "phase"), + per_token: bool = False, +) -> str: + if len(stack) != 3 or set(stack) != {"path", "cohort", "phase"}: + raise ValueError("stack must be a permutation of path,cohort,phase") + phase_groups = build_phase_groups(records) + lines: list[str] = [] + for path_groups in phase_groups["paths"].values(): + for group in path_groups: + if per_token and not group["durable_tokens"]: + continue + for phase in group["phases"]: + frames = { + "path": group["path"], + "cohort": group["label"], + "phase": phase["phase"], + } + folded_stack = ";".join(frames[dimension] for dimension in stack) + lines.append( + f"{folded_stack} " + f"{format_folded_weight(phase['total_ns'], group['durable_tokens'] if per_token else None)}" + ) + idle_ns = phase_groups["inter_round_idle"]["total_ns"] + if not per_token and idle_ns: + lines.append(f"idle;inter_round {idle_ns}") + return "" if not lines else "\n".join(sorted(lines)) + "\n" + + +def build_summary(records: list[dict[str, Any]]) -> dict[str, Any]: + metadata = next( + (record for record in records if record["type"] == "metadata"), {} + ) + footer = next( + (record for record in reversed(records) if record["type"] == "footer"), {} + ) + steps = [record for record in records if record["type"] == "step"] + requests = [record for record in records if record["type"] == "request"] + bursts = [record for record in records if record["type"] == "token_burst"] + + phases: Counter[str] = Counter() + decisions: Counter[str] = Counter() + paths: Counter[str] = Counter() + cohorts: dict[int, list[dict[str, Any]]] = defaultdict(list) + proposed_by_position: list[int] = [] + accepted_by_position: list[int] = [] + for step in steps: + live_slots = int(step.get("live_slots", 0)) + cohorts[live_slots].append(step) + paths[str(step.get("path", "unknown"))] += 1 + for span in step.get("phases", []): + phases[str(span.get("phase", "unknown"))] += int( + span.get("duration_ns", 0) + ) + for lane in step.get("lanes", []): + if lane.get("kind") == "decode": + decisions[str(lane.get("spec", "none"))] += 1 + for source, target in ( + (step.get("proposed_by_position", []), proposed_by_position), + (step.get("accepted_by_position", []), accepted_by_position), + ): + if len(target) < len(source): + target.extend([0] * (len(source) - len(target))) + for index, value in enumerate(source): + target[index] += int(value) + + queue_ms = [ + duration_ms(int(request.get("admitted_ns", 0)), int(request.get("queued_ns", 0))) + for request in requests + ] + ttft_ms = [ + duration_ms(int(request.get("first_token_ns", 0)), int(request.get("queued_ns", 0))) + for request in requests + ] + e2e_ms = [ + duration_ms(int(request.get("completed_ns", 0)), int(request.get("queued_ns", 0))) + for request in requests + ] + burst_times: dict[int, list[tuple[int, int]]] = defaultdict(list) + for burst in bursts: + burst_times[int(burst["request_id"])].append( + (int(burst["ready_ns"]), int(burst.get("token_count", 0))) + ) + inter_burst_ms: list[float] = [] + for request_bursts in burst_times.values(): + request_bursts.sort() + for previous, current in zip(request_bursts, request_bursts[1:]): + inter_burst_ms.append( + (current[0] - previous[0]) / 1_000_000 / max(1, current[1]) + ) + + funnel = {key: sum_field(steps, key) for key in FUNNEL_KEYS} + acceptance_by_position = [ + optional_number(ratio(accepted, proposed)) + for proposed, accepted in zip(proposed_by_position, accepted_by_position) + ] + cohort_summary: dict[str, Any] = {} + for live_slots, cohort in sorted(cohorts.items()): + target_rows = sum_field(cohort, "target_rows") + target_padding = sum_field(cohort, "target_padding_rows") + draft_rows = sum_field(cohort, "draft_rows") + draft_padding = sum_field(cohort, "draft_padding_rows") + proposed = sum_field(cohort, "spec_proposed_draft_tokens") + accepted = sum_field(cohort, "spec_accepted_draft_tokens") + cohort_summary[str(live_slots)] = { + "rounds": len(cohort), + "mean_round_ms": statistics.fmean( + int(step.get("duration_ns", 0)) for step in cohort + ) / 1_000_000, + "target_padding_ratio": optional_number(ratio(target_padding, target_rows)), + "draft_padding_ratio": optional_number(ratio(draft_padding, draft_rows)), + "draft_acceptance_ratio": optional_number(ratio(accepted, proposed)), + "paths": dict(sorted(Counter( + str(step.get("path", "unknown")) for step in cohort + ).items())), + } + + target_rows = sum_field(steps, "target_rows") + target_padding = sum_field(steps, "target_padding_rows") + draft_rows = sum_field(steps, "draft_rows") + draft_padding = sum_field(steps, "draft_padding_rows") + return { + "schema": "lucebox.concurrency.summary.v1", + "run": { + key: value for key, value in metadata.items() + if key not in {"type", "schema"} + }, + "capture": { + "complete": bool(footer.get("complete", False)), + "rounds": len(steps), + "requests": len(requests), + "failed_requests": sum( + int(request.get("completed_ns", 0)) != 0 and request.get("ok") is False + for request in requests + ), + "dropped_steps": int(footer.get("dropped_steps", 0)), + "dropped_requests": int(footer.get("dropped_requests", 0)), + "dropped_token_bursts": int(footer.get("dropped_token_bursts", 0)), + "paths": dict(sorted(paths.items())), + }, + "latency_ms": { + "queue": percentile_pair(queue_ms), + "ttft": percentile_pair(ttft_ms), + "end_to_end": percentile_pair(e2e_ms), + "inter_burst_per_token": percentile_pair(inter_burst_ms), + }, + "speculation": { + **funnel, + "tree_widths": sorted({ + int(step.get("spec_tree_width", 0)) for step in steps + if int(step.get("spec_tree_width", 0)) > 0 + }), + "decisions": dict(sorted(decisions.items())), + "proposed_by_position": proposed_by_position, + "accepted_by_position": accepted_by_position, + "acceptance_by_position": acceptance_by_position, + }, + "padding": { + "target_rows": target_rows, + "target_padding_rows": target_padding, + "target_padding_ratio": optional_number(ratio(target_padding, target_rows)), + "draft_rows": draft_rows, + "draft_padding_rows": draft_padding, + "draft_padding_ratio": optional_number(ratio(draft_padding, draft_rows)), + }, + "cohorts": cohort_summary, + "phase_ns": dict(phases.most_common()), + } + + +def _capture_bounds(records: list[dict[str, Any]]) -> dict[str, int | None]: + timestamps: list[int] = [] + request_keys = ( + "queued_ns", + "admitted_ns", + "prefill_completed_ns", + "first_token_ns", + "completed_ns", + ) + for record in records: + record_type = record.get("type") + if record_type == "step": + started_ns = int(record.get("started_ns", 0)) + if started_ns > 0: + timestamps.append(started_ns) + timestamps.append( + started_ns + max(0, int(record.get("duration_ns", 0))) + ) + elif record_type == "request": + timestamps.extend( + timestamp + for key in request_keys + if (timestamp := int(record.get(key, 0))) > 0 + ) + elif record_type == "token_burst": + ready_ns = int(record.get("ready_ns", 0)) + if ready_ns > 0: + timestamps.append(ready_ns) + earliest_ns = min(timestamps) if timestamps else None + latest_ns = max(timestamps) if timestamps else None + + return { + "earliest_steady_ns": earliest_ns, + "latest_steady_ns": latest_ns, + "duration_ns": ( + latest_ns - earliest_ns if earliest_ns is not None else None + ), + } + + +def _request_payload(records: list[dict[str, Any]]) -> dict[str, Any]: + requests = sorted( + (record for record in records if record.get("type") == "request"), + key=lambda request: ( + int(request.get("queued_ns", 0)), + int(request.get("request_id", 0)), + ), + ) + embedded = requests[:1000] + boundaries = [ + int(request.get(key, 0)) + for request in embedded + for key in ( + "queued_ns", "admitted_ns", "prefill_completed_ns", + "first_token_ns", "completed_ns", + ) + if int(request.get(key, 0)) > 0 + ] + origin = min(boundaries) if boundaries else 0 + end = max(boundaries) if boundaries else origin + + def span(request: dict[str, Any], start_key: str, end_key: str) -> int | None: + start = int(request.get(start_key, 0)) + finish = int(request.get(end_key, 0)) + return finish - start if start and finish >= start else None + + return { + "total": len(requests), + "embedded": len(embedded), + "display_limit": 200, + "origin_ns": origin, + "end_ns": end, + "rows": [ + { + "request_id": int(request.get("request_id", 0)), + "ok": request.get("ok"), + "open_ended": request.get("ok") is None or not int(request.get("completed_ns", 0)), + "prompt_tokens": int(request.get("prompt_tokens", 0)), + "output_tokens": int(request.get("output_tokens", 0)), + "start_offset_ns": max(0, int(request.get("queued_ns", 0)) - origin), + "queue_ns": span(request, "queued_ns", "admitted_ns"), + "prefill_ns": span(request, "admitted_ns", "prefill_completed_ns"), + "first_decode_ns": span(request, "prefill_completed_ns", "first_token_ns"), + "decode_ns": span(request, "first_token_ns", "completed_ns"), + } + for request in embedded + ], + } + + +def _run_notices( + records: list[dict[str, Any]], + catalog: dict[str, Any], + device_key: str | None, +) -> list[dict[str, str]]: + metadata = next( + (record for record in records if record.get("type") == "metadata"), {} + ) + notices: list[dict[str, str]] = [] + devices = catalog.get("devices", {}) + if not device_key: + notices.append({ + "kind": "device", + "level": "warn", + "message": "No device selected. Boundness coloring is neutral. Pass --device with a checked-in device key.", + }) + elif not isinstance(devices, dict) or device_key not in devices: + notices.append({ + "kind": "device", + "level": "warn", + "message": f"No device spec for {device_key}. Boundness coloring is neutral.", + }) + else: + target_key, target_spec = _resolve_model_spec(metadata, catalog) + if target_spec is None: + notices.append({ + "kind": "model", + "level": "warn", + "message": "No target model spec matched model_name or the model_path basename. Target compute is neutral.", + }) + draft_compute = any( + span.get("phase") == "draft_compute" + for record in records if record.get("type") == "step" + for span in record.get("phases", []) + ) + draft_key, draft_spec = _resolve_model_spec(metadata, catalog, draft=True) + if draft_compute and draft_spec is None: + notices.append({ + "kind": "model", + "level": "warn", + "message": "No draft model spec matched draft_path. Draft compute remains overhead.", + }) + del target_key, draft_key + return notices + + +def _device_payload(catalog: dict[str, Any], device_key: str | None) -> dict[str, Any]: + devices = catalog.get("devices", {}) + device = devices.get(device_key) if isinstance(devices, dict) and device_key else None + if not isinstance(device, dict): + return {"key": device_key, "known": False, "name": device_key or "not selected"} + return { + "key": device_key, + "known": True, + "name": device.get("name", device_key), + "mem_bw_gbps": device.get("mem_bw_gbps"), + "fp16_tflops": device.get("fp16_tflops"), + "note": device.get("note", ""), + } + + +def build_run_payload( + records: list[dict[str, Any]], + *, + device_specs: dict[str, Any] | None = None, + device_key: str | None = None, +) -> ReportPayload: + catalog = device_specs or {} + summary = build_summary(records) + cohorts = sorted(int(cohort) for cohort in summary["cohorts"]) + return { + "run": summary["run"], + "capture": summary["capture"], + "capture_bounds": _capture_bounds(records), + "latency_ms": summary["latency_ms"], + "phase_groups": build_phase_groups( + records, device_specs=catalog, device_key=device_key + ), + "requests": _request_payload(records), + "speculation": summary["speculation"], + "padding": summary["padding"], + "device": _device_payload(catalog, device_key), + "notices": _run_notices(records, catalog, device_key), + "mixed_run_cohorts": len(cohorts) > 1, + "cohorts": cohorts, + } + + +def _group_index(run: ReportPayload) -> dict[tuple[str, int, str], dict[str, Any]]: + groups = list(run["phase_groups"]["all"]) + for path_groups in run["phase_groups"]["paths"].values(): + groups.extend(path_groups) + return { + (group["path"], group["cohort"], phase["phase"]): phase + for group in groups + for phase in group["phases"] + } + + +def _diff_payload(current: ReportPayload, baseline: ReportPayload) -> dict[str, Any]: + warnings: list[str] = [] + for key, label in (("model_name", "model"), ("arch", "model architecture")): + current_value = current["run"].get(key) + baseline_value = baseline["run"].get(key) + if current_value != baseline_value: + warnings.append( + f"Baseline {label} mismatch: {baseline_value or 'unknown'} vs {current_value or 'unknown'}." + ) + current_device = current["device"].get("key") + baseline_device = baseline["device"].get("key") + if current_device != baseline_device: + warnings.append( + f"Baseline device mismatch: {baseline_device or 'not selected'} " + f"vs {current_device or 'not selected'}." + ) + current_index = _group_index(current) + baseline_index = _group_index(baseline) + rows = [] + for path, cohort, phase_name in sorted( + current_index.keys() | baseline_index.keys(), + key=lambda key: (key[0], key[1], phase_sort_key(key[2])), + ): + current_phase = current_index.get((path, cohort, phase_name), {}) + baseline_phase = baseline_index.get((path, cohort, phase_name), {}) + current_value = current_phase.get("ns_per_token") + baseline_value = baseline_phase.get("ns_per_token") + delta = ( + current_value - baseline_value + if current_value is not None and baseline_value is not None else None + ) + delta_percent = ( + delta / baseline_value + if delta is not None and baseline_value else None + ) + rows.append({ + "path": path, + "cohort": cohort, + "phase": phase_name, + "baseline_ns_per_token": baseline_value, + "current_ns_per_token": current_value, + "delta_ns_per_token": delta, + "delta_percent": delta_percent, + }) + return {"warnings": warnings, "rows": rows} + + +def build_report_payload( + records: list[dict[str, Any]], + baseline_records: list[dict[str, Any]] | None = None, + *, + device_specs: dict[str, Any] | None = None, + device_key: str | None = None, + baseline_device_key: str | None = None, +) -> ReportPayload: + current = build_run_payload( + records, device_specs=device_specs, device_key=device_key + ) + baseline = ( + build_run_payload( + baseline_records, + device_specs=device_specs, + device_key=baseline_device_key or device_key, + ) + if baseline_records is not None else None + ) + return { + "schema": "lucebox.concurrency.report.v1", + "phase_order": list(PHASE_ORDER), + "current": current, + "baseline": baseline, + "diff": _diff_payload(current, baseline) if baseline is not None else None, + } diff --git a/harness/benchmarks/concurrency/profile_report.py b/harness/benchmarks/concurrency/profile_report.py index 2dc136f95..6c8b0d91e 100644 --- a/harness/benchmarks/concurrency/profile_report.py +++ b/harness/benchmarks/concurrency/profile_report.py @@ -6,10 +6,16 @@ import argparse import json import math -import statistics -from collections import Counter, defaultdict from pathlib import Path -from typing import Any, Iterable +from typing import Any + +from profile_payload import ( + build_folded, + build_summary, + format_folded_weight, + step_phase_buckets, +) +from profile_view import build_html, load_device_specs FOLDED_DIMENSIONS = frozenset(("path", "cohort", "phase")) DEFAULT_FOLDED_STACK = ("path", "cohort", "phase") @@ -41,210 +47,6 @@ def percent(value: float) -> str: return "n/a" if math.isnan(value) else f"{100.0 * value:.1f}%" -def percentile(values: Iterable[float], quantile: float) -> float: - ordered = sorted(values) - if not ordered: - return math.nan - index = (len(ordered) - 1) * quantile - low = math.floor(index) - high = math.ceil(index) - if low == high: - return ordered[low] - return ordered[low] * (high - index) + ordered[high] * (index - low) - - -def duration_ms(end: int, start: int) -> float: - return (end - start) / 1_000_000 if end and start and end >= start else math.nan - - -def sum_field(records: Iterable[dict[str, Any]], key: str) -> int: - return sum(int(record.get(key, 0)) for record in records) - - -def optional_number(value: float) -> float | None: - return None if math.isnan(value) else value - - -def percentile_pair(values: Iterable[float]) -> dict[str, float | None]: - finite = [value for value in values if not math.isnan(value)] - return { - "p50": optional_number(percentile(finite, 0.50)), - "p95": optional_number(percentile(finite, 0.95)), - } - - -def build_summary(records: list[dict[str, Any]]) -> dict[str, Any]: - metadata = next( - (record for record in records if record["type"] == "metadata"), - {}, - ) - footer = next( - (record for record in reversed(records) if record["type"] == "footer"), - {}, - ) - steps = [record for record in records if record["type"] == "step"] - requests = [record for record in records if record["type"] == "request"] - bursts = [record for record in records if record["type"] == "token_burst"] - - phases: Counter[str] = Counter() - decisions: Counter[str] = Counter() - paths: Counter[str] = Counter() - cohorts: dict[int, list[dict[str, Any]]] = defaultdict(list) - proposed_by_position: list[int] = [] - accepted_by_position: list[int] = [] - for step in steps: - live_slots = int(step.get("live_slots", 0)) - cohorts[live_slots].append(step) - paths[str(step.get("path", "unknown"))] += 1 - for span in step.get("phases", []): - phases[str(span.get("phase", "unknown"))] += int( - span.get("duration_ns", 0) - ) - for lane in step.get("lanes", []): - if lane.get("kind") == "decode": - decisions[str(lane.get("spec", "none"))] += 1 - for source, target in ( - (step.get("proposed_by_position", []), proposed_by_position), - (step.get("accepted_by_position", []), accepted_by_position), - ): - if len(target) < len(source): - target.extend([0] * (len(source) - len(target))) - for index, value in enumerate(source): - target[index] += int(value) - - queue_ms = [ - duration_ms(int(request.get("admitted_ns", 0)), - int(request.get("queued_ns", 0))) - for request in requests - ] - ttft_ms = [ - duration_ms(int(request.get("first_token_ns", 0)), - int(request.get("queued_ns", 0))) - for request in requests - ] - e2e_ms = [ - duration_ms(int(request.get("completed_ns", 0)), - int(request.get("queued_ns", 0))) - for request in requests - ] - burst_times: dict[int, list[tuple[int, int]]] = defaultdict(list) - for burst in bursts: - burst_times[int(burst["request_id"])].append( - (int(burst["ready_ns"]), int(burst.get("token_count", 0))) - ) - inter_burst_ms: list[float] = [] - for request_bursts in burst_times.values(): - request_bursts.sort() - for previous, current in zip(request_bursts, request_bursts[1:]): - inter_burst_ms.append( - (current[0] - previous[0]) / - 1_000_000 / max(1, current[1]) - ) - - funnel_keys = ( - "spec_eligible_lanes", - "spec_reserved_lanes", - "spec_attempted_lanes", - "spec_proposed_draft_tokens", - "spec_verified_draft_tokens", - "spec_accepted_draft_tokens", - "spec_durable_draft_tokens", - "spec_scheduler_consumed_tokens", - ) - funnel = {key: sum_field(steps, key) for key in funnel_keys} - acceptance_by_position = [ - optional_number(ratio(accepted, proposed)) - for proposed, accepted in zip( - proposed_by_position, accepted_by_position) - ] - - cohort_summary: dict[str, Any] = {} - for live_slots, cohort in sorted(cohorts.items()): - target_rows = sum_field(cohort, "target_rows") - target_padding = sum_field(cohort, "target_padding_rows") - draft_rows = sum_field(cohort, "draft_rows") - draft_padding = sum_field(cohort, "draft_padding_rows") - proposed = sum_field(cohort, "spec_proposed_draft_tokens") - accepted = sum_field(cohort, "spec_accepted_draft_tokens") - cohort_summary[str(live_slots)] = { - "rounds": len(cohort), - "mean_round_ms": statistics.fmean( - int(step.get("duration_ns", 0)) for step in cohort - ) / 1_000_000, - "target_padding_ratio": optional_number( - ratio(target_padding, target_rows) - ), - "draft_padding_ratio": optional_number( - ratio(draft_padding, draft_rows) - ), - "draft_acceptance_ratio": optional_number( - ratio(accepted, proposed) - ), - "paths": dict(sorted(Counter( - str(step.get("path", "unknown")) for step in cohort - ).items())), - } - - target_rows = sum_field(steps, "target_rows") - target_padding = sum_field(steps, "target_padding_rows") - draft_rows = sum_field(steps, "draft_rows") - draft_padding = sum_field(steps, "draft_padding_rows") - return { - "schema": "lucebox.concurrency.summary.v1", - "run": { - key: value for key, value in metadata.items() - if key not in {"type", "schema"} - }, - "capture": { - "complete": bool(footer.get("complete", False)), - "rounds": len(steps), - "requests": len(requests), - "failed_requests": sum( - int(request.get("completed_ns", 0)) != 0 - and request.get("ok") is False - for request in requests - ), - "dropped_steps": int(footer.get("dropped_steps", 0)), - "dropped_requests": int(footer.get("dropped_requests", 0)), - "dropped_token_bursts": int( - footer.get("dropped_token_bursts", 0) - ), - "paths": dict(sorted(paths.items())), - }, - "latency_ms": { - "queue": percentile_pair(queue_ms), - "ttft": percentile_pair(ttft_ms), - "end_to_end": percentile_pair(e2e_ms), - "inter_burst_per_token": percentile_pair(inter_burst_ms), - }, - "speculation": { - **funnel, - "tree_widths": sorted({ - int(step.get("spec_tree_width", 0)) for step in steps - if int(step.get("spec_tree_width", 0)) > 0 - }), - "decisions": dict(sorted(decisions.items())), - "proposed_by_position": proposed_by_position, - "accepted_by_position": accepted_by_position, - "acceptance_by_position": acceptance_by_position, - }, - "padding": { - "target_rows": target_rows, - "target_padding_rows": target_padding, - "target_padding_ratio": optional_number( - ratio(target_padding, target_rows) - ), - "draft_rows": draft_rows, - "draft_padding_rows": draft_padding, - "draft_padding_ratio": optional_number( - ratio(draft_padding, draft_rows) - ), - }, - "cohorts": cohort_summary, - "phase_ns": dict(phases.most_common()), - } - - def build_markdown(summary: dict[str, Any]) -> str: run = summary["run"] capture = summary["capture"] @@ -474,118 +276,6 @@ def parse_folded_stack(value: str) -> tuple[str, ...]: return dimensions -def folded_atom(value: Any) -> str: - atom = str(value) - if not atom or ";" in atom or any(character.isspace() for character in atom): - raise ValueError(f"invalid folded-stack frame: {atom!r}") - return atom - - -def step_phase_buckets(step: dict[str, Any]) -> Counter[str]: - duration_ns = max(0, int(step.get("duration_ns", 0))) - events: dict[int, Counter[str]] = defaultdict(Counter) - for span in step.get("phases", []): - phase = folded_atom(span.get("phase", "unknown")) - raw_start = int(span.get("start_offset_ns", 0)) - raw_end = raw_start + max(0, int(span.get("duration_ns", 0))) - start = min(duration_ns, max(0, raw_start)) - end = min(duration_ns, max(0, raw_end)) - if end <= start: - continue - events[start][phase] += 1 - events[end][phase] -= 1 - - buckets: Counter[str] = Counter() - active: Counter[str] = Counter() - previous = 0 - for offset in sorted({0, duration_ns, *events}): - if offset > previous: - phases = sorted( - phase for phase, count in active.items() if count > 0 - ) - if not phases: - label = "unattributed" - elif len(phases) == 1: - label = phases[0] - else: - label = f"overlap({'+'.join(phases)})" - buckets[label] += offset - previous - active.update(events[offset]) - previous = offset - return buckets - - -def format_folded_weight(duration_ns: int, tokens: int | None) -> str: - if tokens is None or duration_ns % tokens == 0: - return str(duration_ns if tokens is None else duration_ns // tokens) - return f"{duration_ns / tokens:.6f}".rstrip("0").rstrip(".") - - -def build_folded( - records: list[dict[str, Any]], - *, - stack: tuple[str, ...] = DEFAULT_FOLDED_STACK, - per_token: bool = False, -) -> str: - if len(stack) != 3 or set(stack) != FOLDED_DIMENSIONS: - raise ValueError("stack must be a permutation of path,cohort,phase") - - durations: Counter[tuple[str, int, str]] = Counter() - tokens: Counter[tuple[str, int]] = Counter() - steps = [record for record in records if record["type"] == "step"] - for step in steps: - path = folded_atom(step.get("path", "unknown")) - cohort = max(0, int(step.get("live_slots", 0))) - group = (path, cohort) - tokens[group] += sum( - max(0, int(lane.get("scheduler_consumed_tokens", 0))) - for lane in step.get("lanes", []) - if lane.get("kind") == "decode" - ) - for phase, duration_ns in step_phase_buckets(step).items(): - durations[path, cohort, phase] += duration_ns - - lines: list[str] = [] - for (path, cohort, phase), duration_ns in durations.items(): - denominator = tokens[path, cohort] if per_token else None - if per_token and denominator == 0: - continue - frames = { - "path": path, - "cohort": f"C={cohort}", - "phase": phase, - } - folded_stack = ";".join(frames[dimension] for dimension in stack) - lines.append( - f"{folded_stack} {format_folded_weight(duration_ns, denominator)}" - ) - - if not per_token and len(steps) > 1: - inter_round_ns = 0 - ordered_steps = sorted( - steps, - key=lambda step: ( - int(step.get("started_ns", 0)), - int(step.get("round_id", 0)), - ), - ) - previous_end = ( - int(ordered_steps[0].get("started_ns", 0)) - + max(0, int(ordered_steps[0].get("duration_ns", 0))) - ) - for step in ordered_steps[1:]: - started_ns = int(step.get("started_ns", 0)) - inter_round_ns += max(0, started_ns - previous_end) - previous_end = max( - previous_end, - started_ns + max(0, int(step.get("duration_ns", 0))), - ) - if inter_round_ns: - lines.append(f"idle;inter_round {inter_round_ns}") - - return "" if not lines else "\n".join(sorted(lines)) + "\n" - - def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("profile", type=Path) @@ -594,12 +284,27 @@ def main() -> int: parser.add_argument("--json-summary", type=Path) parser.add_argument("--folded", type=Path) parser.add_argument("--folded-per-token", type=Path) + parser.add_argument("--html", type=Path) + parser.add_argument("--baseline", type=Path) + parser.add_argument( + "--device", + help="device_specs.json key used for analytic boundness coloring", + ) + parser.add_argument( + "--baseline-device", + help="device key for the baseline capture; defaults to --device", + ) parser.add_argument( "--stack", type=parse_folded_stack, default=DEFAULT_FOLDED_STACK, metavar="path,cohort,phase", ) args = parser.parse_args() + if args.baseline and not args.html: + parser.error("--baseline requires --html") + if args.baseline_device and not args.baseline: + parser.error("--baseline-device requires --baseline") + records = load_records(args.profile) summary = build_summary(records) markdown = build_markdown(summary) @@ -626,6 +331,17 @@ def main() -> int: build_folded(records, stack=args.stack, per_token=True), encoding="utf-8", ) + if args.html: + device_specs = load_device_specs() + if args.device: + device_specs["selected_device"] = args.device + if args.baseline_device: + device_specs["baseline_device"] = args.baseline_device + baseline_records = load_records(args.baseline) if args.baseline else None + args.html.write_text( + build_html(records, baseline_records, device_specs), + encoding="utf-8", + ) return 0 diff --git a/harness/benchmarks/concurrency/profile_view.py b/harness/benchmarks/concurrency/profile_view.py new file mode 100644 index 000000000..b2fd9a58b --- /dev/null +++ b/harness/benchmarks/concurrency/profile_view.py @@ -0,0 +1,570 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from profile_payload import ReportPayload, build_report_payload + +DEVICE_SPECS_PATH = Path(__file__).with_name("device_specs.json") + + +def load_device_specs(path: Path = DEVICE_SPECS_PATH) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict) or not isinstance(value.get("devices"), dict): + raise ValueError(f"{path}: device specs need a devices object") + if not isinstance(value.get("models"), dict): + raise ValueError(f"{path}: device specs need a models object") + return value + + +def _payload_from_selection( + records: list[dict[str, Any]], + baseline_records: list[dict[str, Any]] | None, + device_specs: dict[str, Any] | None, +) -> ReportPayload: + selection = device_specs or load_device_specs() + return build_report_payload( + records, + baseline_records, + device_specs=selection, + device_key=selection.get("selected_device"), + baseline_device_key=selection.get("baseline_device"), + ) + + +def build_html( + records: list[dict[str, Any]], + baseline_records: list[dict[str, Any]] | None = None, + device_specs: dict[str, Any] | None = None, +) -> str: + payload = _payload_from_selection(records, baseline_records, device_specs) + encoded = json.dumps( + payload, + ensure_ascii=False, + allow_nan=False, + separators=(",", ":"), + ).replace("<", "\\u003c") + return _HTML.replace("__REPORT_PAYLOAD__", encoded) + + +_HTML = r''' + + + + +LuceGraph + + + +
+
+
+
Lucebox / offline profiler
+

LuceGraph

+

+
+
+
+ +
+
+

Phase-budget wall

+
+ + + +
+
+
+
+
+
+
Focus or select a segment to inspect its percentiles and roofline facts.
+
+
+ +
+

Request waterfall

+
+
+ +
+

Speculation funnel

+
+
+ +
+
+ + + +
+

Capture contract

+
+
+
+ + + + + +''' diff --git a/harness/benchmarks/concurrency/test_concurrency_tools.py b/harness/benchmarks/concurrency/test_concurrency_tools.py index 8020e69ad..e827d888d 100644 --- a/harness/benchmarks/concurrency/test_concurrency_tools.py +++ b/harness/benchmarks/concurrency/test_concurrency_tools.py @@ -83,6 +83,7 @@ def test_extended_matrix_is_deterministic_without_changing_existing_cohorts(self def test_client_level_parser_rejects_reuse(self) -> None: self.assertEqual(generator.CLIENT_MATRICES["legacy"], (2, 4, 8, 16)) + self.assertEqual(generator.CLIENT_MATRICES["lucegraph"], (1, 2, 3, 4, 5)) self.assertEqual(generator.parse_client_levels("2,4,8,16,32"), (2, 4, 8, 16, 32)) with self.assertRaisesRegex(ValueError, "distinct"): generator.parse_client_levels("2,4,2") diff --git a/harness/benchmarks/concurrency/test_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_concurrent_benchmark.py index c2deb66f7..ca1e6acaa 100644 --- a/harness/benchmarks/concurrency/test_concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/test_concurrent_benchmark.py @@ -28,7 +28,7 @@ def test_sse_parser_handles_events_and_done(self) -> None: ['{"choices":[{"delta":{"content":"hi"}}]}', "[DONE]"], ) - def test_default_matrix_matches_profile_view_range(self) -> None: + def test_default_matrix_matches_lucegraph_range(self) -> None: self.assertEqual(benchmark.DEFAULT_CLIENT_LEVELS, (1, 2, 3, 4, 5)) def test_run_rejects_duplicate_client_levels(self) -> None: diff --git a/harness/benchmarks/concurrency/test_profile_view.py b/harness/benchmarks/concurrency/test_profile_view.py new file mode 100644 index 000000000..62bcee308 --- /dev/null +++ b/harness/benchmarks/concurrency/test_profile_view.py @@ -0,0 +1,375 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import copy +import json +import re +import shutil +import subprocess +import tempfile +import unittest +from pathlib import Path + +import profile_payload +import profile_view + +HERE = Path(__file__).parent + + +class ProfileViewTest(unittest.TestCase): + def records(self) -> list[dict]: + return [ + { + "type": "metadata", + "schema": "lucebox.concurrency.v1", + "git_sha": "current123", + "model_name": "Qwen3.6-27B-Q4_K_M.gguf", + "model_path": "/models/Qwen3.6-27B-Q4_K_M.gguf", + "draft_path": "/models/unlisted-draft.gguf", + "arch": "qwen35", + "runtime_backend": "hip", + "max_concurrency": 2, + "started_unix_ns": 1_700_000_000_000_000_000, + "retention_policy": "keep_first", + "label": "safe ", + }, + { + "type": "step", + "round_id": 1, + "started_ns": 1_000, + "duration_ns": 100, + "path": "packed", + "live_slots": 1, + "target_rows": 10, + "target_padding_rows": 2, + "max_kv_len": 20, + "lanes": [ + {"kind": "decode", "scheduler_consumed_tokens": 2} + ], + "phases": [ + { + "phase": "target_compute", + "start_offset_ns": 10, + "duration_ns": 50, + }, + { + "phase": "readback_sync", + "start_offset_ns": 40, + "duration_ns": 40, + }, + ], + "proposed_by_position": [], + "accepted_by_position": [], + }, + { + "type": "step", + "round_id": 2, + "started_ns": 1_200, + "duration_ns": 200, + "path": "speculative", + "live_slots": 2, + "target_rows": 20, + "target_padding_rows": 0, + "draft_rows": 8, + "draft_padding_rows": 2, + "max_kv_len": 30, + "spec_eligible_lanes": 2, + "spec_reserved_lanes": 2, + "spec_attempted_lanes": 2, + "spec_proposed_draft_tokens": 4, + "spec_verified_draft_tokens": 4, + "spec_accepted_draft_tokens": 3, + "spec_durable_draft_tokens": 3, + "spec_scheduler_consumed_tokens": 3, + "lanes": [ + {"kind": "decode", "scheduler_consumed_tokens": 4} + ], + "phases": [ + { + "phase": "target_compute", + "start_offset_ns": 0, + "duration_ns": 100, + }, + { + "phase": "draft_compute", + "start_offset_ns": 100, + "duration_ns": 50, + }, + ], + "proposed_by_position": [2, 2], + "accepted_by_position": [2, 1], + }, + { + "type": "request", + "request_id": 1, + "ok": True, + "queued_ns": 900, + "admitted_ns": 920, + "prefill_completed_ns": 950, + "first_token_ns": 1_000, + "completed_ns": 1_350, + "prompt_tokens": 12, + "output_tokens": 6, + }, + { + "type": "request", + "request_id": 2, + "ok": None, + "queued_ns": 1_100, + "admitted_ns": 1_150, + "prefill_completed_ns": 0, + "first_token_ns": 0, + "completed_ns": 0, + }, + { + "type": "footer", + "complete": True, + "dropped_steps": 0, + "dropped_requests": 0, + "dropped_token_bursts": 0, + }, + ] + + def specs(self, device: str = "bandwidth") -> dict: + return { + "selected_device": device, + "devices": { + "bandwidth": { + "name": "Bandwidth fixture", + "mem_bw_gbps": 1000, + "fp16_tflops": 100, + }, + "compute": { + "name": "Compute fixture", + "mem_bw_gbps": 1000, + "fp16_tflops": 0.01, + }, + }, + "models": { + "Qwen3.6-27B-Q4_K_M.gguf": { + "weight_bytes": 100, + "active_params": 100, + "kv_bytes_per_token_per_seq": 0, + } + }, + } + + def test_checked_in_device_and_model_facts(self) -> None: + specs = profile_view.load_device_specs() + + self.assertEqual(specs["devices"]["gfx1201"]["mem_bw_gbps"], 640) + self.assertEqual(specs["devices"]["gfx1201"]["fp16_tflops"], 95.7) + self.assertEqual(specs["devices"]["gfx1151"]["mem_bw_gbps"], 256) + self.assertEqual(specs["devices"]["gfx1151"]["fp16_tflops"], 60) + model = specs["models"]["Qwen3.6-27B-Q4_K_M.gguf"] + self.assertEqual( + ( + model["weight_bytes"], + model["active_params"], + model["kv_bytes_per_token_per_seq"], + ), + (16_800_000_000, 27_000_000_000, 18_432), + ) + + def embedded_payload(self, html: str) -> dict: + match = re.search( + r'', + html, + re.DOTALL, + ) + self.assertIsNotNone(match) + return json.loads(match.group(1)) + + def test_embedded_aggregates_round_trip(self) -> None: + html = profile_view.build_html(self.records(), device_specs=self.specs()) + payload = self.embedded_payload(html) + + self.assertEqual(html.count('type="application/json" id="data"'), 1) + self.assertIn("LuceGraph", html) + self.assertIn("

LuceGraph

", html) + self.assertNotIn("", html, re.DOTALL) + javascript = [script for script in scripts if not script.lstrip().startswith("{")] + self.assertEqual(len(javascript), 1) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "lucegraph.js" + path.write_text(javascript[0], encoding="utf-8") + result = subprocess.run( + ["node", "--check", str(path)], + text=True, + capture_output=True, + check=False, + ) + self.assertEqual(result.returncode, 0, result.stderr) + + def test_cli_writes_one_offline_file(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + capture = root / "capture.jsonl" + report = root / "report.html" + capture.write_text( + "".join(json.dumps(record) + "\n" for record in self.records()), + encoding="utf-8", + ) + result = subprocess.run( + [ + "python3", + str(HERE / "profile_report.py"), + str(capture), + "--html", + str(report), + "--device", + "gfx1201", + ], + text=True, + capture_output=True, + check=False, + ) + html = report.read_text(encoding="utf-8") if report.exists() else "" + + self.assertEqual(result.returncode, 0, result.stderr) + self.assertTrue(html.startswith("")) + self.assertIsNone(re.search(r"https?://", html)) + self.assertEqual(self.embedded_payload(html)["current"]["device"]["key"], "gfx1201") + + def test_c1_only_without_speculation_stays_renderable(self) -> None: + records = [ + record for record in self.records() + if record.get("type") not in {"step", "request"} + or record.get("round_id") == 1 + or record.get("request_id") == 1 + ] + html = profile_view.build_html(records) + payload = self.embedded_payload(html) + + self.assertEqual(payload["current"]["cohorts"], [1]) + self.assertFalse(payload["current"]["mixed_run_cohorts"]) + self.assertEqual(payload["current"]["speculation"]["spec_attempted_lanes"], 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/server/docs/CONCURRENCY_OBSERVABILITY.md b/server/docs/CONCURRENCY_OBSERVABILITY.md index 2b0c34954..0260cffe9 100644 --- a/server/docs/CONCURRENCY_OBSERVABILITY.md +++ b/server/docs/CONCURRENCY_OBSERVABILITY.md @@ -91,6 +91,50 @@ decode tokens that the scheduler consumed for that group. The report omits a group when it has no durable decode tokens. Inter-round gaps have no honest path or cohort owner, so the per-token file omits them. +### Build a LuceGraph report + +Write one self-contained HTML file that opens from `file://`. + +```bash +python3 harness/benchmarks/concurrency/profile_report.py \ + /tmp/lucebox-profile.jsonl \ + --html /tmp/lucegraph.html \ + --device gfx1201 +``` + +The LuceGraph report contains three linked views. The phase-budget wall +compares +concurrency cohorts with per-token, per-round, and wall-share normalization. +The request waterfall separates queue, prefill, first-decode, and decode time. +The speculation strip shows the funnel and acceptance by draft position. + +The wall uses the exclusive buckets from `step_phase_buckets`. It includes +unattributed time, overlapping spans, and inter-round host gaps. Select a wall +segment to inspect its zero-inclusive per-round p50 and p95 values. The report +embeds aggregates instead of raw step records. + +A capture with several observed `live_slots` values shows a mixed-run badge. +Admission and tail drain can create these cohorts even when the configured +concurrency is fixed. Use separate fixed-C captures for cohort comparisons +that exclude this bias. + +Pass `--baseline` to compare two captures. + +```bash +python3 harness/benchmarks/concurrency/profile_report.py \ + current.jsonl --html lucegraph-diff.html --device gfx1201 \ + --baseline baseline.jsonl --baseline-device gfx1151 +``` + +`device_specs.json` is the source of device and model facts. The capture's +`arch` field names the model adapter, not the GPU architecture, so the +report never infers a device from capture metadata. Omit `--device` or pass +an unknown key to render neutral segments with a visible notice. + +The analytic roofline classifier applies only to `target_compute` and +`draft_compute`. Segment details show arithmetic intensity, machine +balance, and headroom. Recheck any fact whose note starts with `VERIFY`. + The v1 capture does not contain model FLOP counts, weight bytes, or device machine balance. Folded stacks also have no portable color metadata. Use a separate device profile for compute-bound or bandwidth-bound classification.