diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index dacd851bb..8c9138718 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -741,6 +741,12 @@ if(DFLASH27B_TESTS) target_include_directories(test_derived_scalars PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) add_test(NAME derived_scalars COMMAND test_derived_scalars) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_chain_rollback_policy.cpp") + add_executable(test_chain_rollback_policy test/test_chain_rollback_policy.cpp) + target_include_directories(test_chain_rollback_policy PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src/common) + add_test(NAME chain_rollback_policy COMMAND test_chain_rollback_policy) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_kvflash_placement.cpp") add_executable(test_kvflash_placement test/test_kvflash_placement.cpp) target_include_directories(test_kvflash_placement PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) diff --git a/server/src/common/chain_rollback_policy.h b/server/src/common/chain_rollback_policy.h new file mode 100644 index 000000000..c3fd5be37 --- /dev/null +++ b/server/src/common/chain_rollback_policy.h @@ -0,0 +1,38 @@ +#pragma once + +#include +#include + +namespace dflash::common { + +struct ChainRollbackPolicy { + bool checkpoint_f32 = false; + int fast_rollback_threshold = 5; + bool diagnostics = false; +}; + +inline bool env_flag_enabled(const char * name) { + const char * value = std::getenv(name); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + +inline ChainRollbackPolicy resolve_chain_rollback_policy() { + ChainRollbackPolicy policy; + policy.checkpoint_f32 = env_flag_enabled("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32"); + policy.diagnostics = env_flag_enabled("DFLASH_SINGLE_CHAIN_ROLLBACK_DIAG"); + + // Lower thresholds are valid only with exact F32 checkpoints. This keeps + // the established F16 behavior unchanged when no opt-in flags are set. + if (policy.checkpoint_f32) { + const char * value = std::getenv("DFLASH_FAST_ROLLBACK_THRESHOLD"); + if (value != nullptr) { + const int requested = std::atoi(value); + if (requested >= 1 && requested <= 5) { + policy.fast_rollback_threshold = requested; + } + } + } + return policy; +} + +} // namespace dflash::common diff --git a/server/src/common/dflash_spec_decode.cpp b/server/src/common/dflash_spec_decode.cpp index 74fc1807b..6c7ea3cdc 100644 --- a/server/src/common/dflash_spec_decode.cpp +++ b/server/src/common/dflash_spec_decode.cpp @@ -12,6 +12,8 @@ #include #include +#include "chain_rollback_policy.h" + namespace dflash::common { namespace { @@ -84,6 +86,13 @@ bool run_dflash_spec_decode( int n_accept_sum = 0; int n_hint_proposed = 0; int n_hint_accepted = 0; + const ChainRollbackPolicy rollback_policy = resolve_chain_rollback_policy(); + int rollback_accept_hist[17] = {}; + int rollback_fast_low = 0; + int rollback_fast_high = 0; + int rollback_legacy_replay = 0; + int rollback_failed_fallback = 0; + auto t_dec0 = std::chrono::steady_clock::now(); while (n_generated < n_gen) { @@ -213,9 +222,10 @@ bool run_dflash_spec_decode( // ── Commit accepted tokens to KV state ────────────────────────── // Adaptive: use fast-rollback when acceptance is high enough to benefit. - constexpr int kFastRollbackThreshold = 5; + rollback_accept_hist[std::min(accept_n, 16)]++; const bool use_fast_rollback = - target.supports_fast_rollback() && (accept_n >= kFastRollbackThreshold); + target.supports_fast_rollback() && + (accept_n >= rollback_policy.fast_rollback_threshold); std::vector replay_tok((size_t)commit_n); for (int i = 0; i < commit_n; i++) { @@ -236,15 +246,19 @@ bool run_dflash_spec_decode( if (target.rollback_to(committed, commit_n)) { last_tok = target_tok[commit_n - 1]; fast_rolled_back = true; + if (accept_n < 5) rollback_fast_low++; + else rollback_fast_high++; } else { // Rollback failed (e.g. CUDA error / unsupported state type). // The pre-verify snapshot is still valid, so degrade to the // legacy restore+replay path below instead of aborting. std::fprintf(stderr, "dflash-spec rollback_to failed; " "falling back to restore+replay\n"); + rollback_failed_fallback++; } } if (!fast_rolled_back) { + rollback_legacy_replay++; // Legacy path: restore SSM snapshot and replay accepted + bonus tokens. // (When falling back from fast-rollback, bonus_tok is already -1 and // replay_tok/commit_n reflect the budget-clamped accepted set.) @@ -296,6 +310,19 @@ bool run_dflash_spec_decode( std::printf("[target-split-dflash] %d draft steps, accepted=%d/%d (%.1f%%), avg commit/step=%.2f\n", n_draft_steps, n_accept_sum, total_draft_pos, accept_pct, n_draft_steps > 0 ? (double)n_generated / (double)n_draft_steps : 0.0); + if (rollback_policy.diagnostics) { + std::printf("[chain-rollback-policy] checkpoint=%s threshold=%d fast_low=%d fast_high=%d legacy_replay=%d failed_fallback=%d accept_hist=1:%d,2:%d,3:%d,4:%d,5:%d,6:%d,7:%d,8:%d,9:%d,10:%d,11:%d,12:%d,13:%d,14:%d,15:%d,16+:%d\n", + rollback_policy.checkpoint_f32 ? "F32" : "default", + rollback_policy.fast_rollback_threshold, + rollback_fast_low, rollback_fast_high, rollback_legacy_replay, + rollback_failed_fallback, + rollback_accept_hist[1], rollback_accept_hist[2], rollback_accept_hist[3], + rollback_accept_hist[4], rollback_accept_hist[5], rollback_accept_hist[6], + rollback_accept_hist[7], rollback_accept_hist[8], rollback_accept_hist[9], + rollback_accept_hist[10], rollback_accept_hist[11], rollback_accept_hist[12], + rollback_accept_hist[13], rollback_accept_hist[14], rollback_accept_hist[15], + rollback_accept_hist[16]); + } if (n_hint_proposed > 0) { std::printf("[target-split-dflash] hint tokens: %d/%d accepted (%.1f%%)\n", n_hint_accepted, n_hint_proposed, diff --git a/server/src/internal.h b/server/src/internal.h index ea43482fc..6f5a37405 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -654,7 +654,8 @@ ggml_tensor * build_qwen35_layer( int fa_window = 0, ggml_tensor * q_tail_capture = nullptr, int q_tail_start = 0, - ggml_tensor * kv_write_rows = nullptr); + ggml_tensor * kv_write_rows = nullptr, + ggml_tensor * parent_ids = nullptr); // Overload that also exposes the MoE router selection tensor (if MoE layer). ggml_tensor * build_qwen35_layer( @@ -673,7 +674,8 @@ ggml_tensor * build_qwen35_layer( ggml_tensor * q_tail_capture, int q_tail_start, ggml_tensor ** moe_selected_out, - ggml_tensor * kv_write_rows = nullptr); + ggml_tensor * kv_write_rows = nullptr, + ggml_tensor * parent_ids = nullptr); QwenLayerPrefnOutputs build_qwen35_layer_prefn( ggml_context * ctx, diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 842314b0d..4e4fdf912 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -24,7 +24,8 @@ bool build_layer_step( bool capture, int fa_window, int kq_stride_pad, - bool kvflash) { + bool kvflash, + bool tree_mode) { if (kvflash) with_mask = true; step_graph_free(sg); @@ -74,6 +75,12 @@ bool build_layer_step( } } + if (tree_mode && !is_attn) { + sg.parent_ids = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + ggml_set_name(sg.parent_ids, "parent_ids"); + ggml_set_input(sg.parent_ids); + } + sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); ggml_tensor * layer_out = build_qwen35_layer( @@ -81,7 +88,7 @@ bool build_layer_step( sg.inp_embed, sg.positions, sg.attn_mask, kv_start, n_tokens, capture, fa_window, /*q_tail_capture=*/nullptr, /*q_tail_start=*/0, - sg.kv_write_rows); + sg.kv_write_rows, sg.parent_ids); if (!layer_out) return false; ggml_tensor * out_view = ggml_view_2d(sg.ctx, act_out, diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index 7557a1f7e..2cff7804a 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -39,7 +39,8 @@ bool build_layer_step( bool capture, int fa_window = 0, int kq_stride_pad = KQ_MASK_PAD, - bool kvflash = false); + bool kvflash = false, + bool tree_mode = false); // `kvflash`: pooled mode — KV rows go through a set_rows input // (sg.kv_write_rows, [n_tokens, n_head_kv] ne0-major slots) and the mask diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 24ce524ae..28cc43e71 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1,4 +1,5 @@ #include "qwen35_backend.h" +#include "common/chain_rollback_policy.h" #include "placement/skip_park_guard.h" #include "qwen35_dflash_target.h" #include "graph_builders.h" @@ -1927,12 +1928,33 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int n_hint_proposed = 0; int n_hint_accepted = 0; int target_forwards = 0; + const ChainRollbackPolicy rollback_policy = resolve_chain_rollback_policy(); + int rollback_accept_hist[17] = {}; + int rollback_fast_low = 0; + int rollback_fast_high = 0; + int rollback_legacy_replay = 0; + int rollback_failed_fallback = 0; + auto log_target_forward_stats = [&]() { std::fprintf(stderr, "[spec-decode] target_forwards=%d forwards_per_token=%.6f forwards_per_step=%.3f\n", target_forwards, n_generated > 0 ? (double)target_forwards / n_generated : 0.0, n_draft_steps > 0 ? (double)target_forwards / n_draft_steps : 0.0); + if (rollback_policy.diagnostics) { + std::fprintf(stderr, + "[chain-rollback-policy] checkpoint=%s threshold=%d fast_low=%d fast_high=%d legacy_replay=%d failed_fallback=%d accept_hist=1:%d,2:%d,3:%d,4:%d,5:%d,6:%d,7:%d,8:%d,9:%d,10:%d,11:%d,12:%d,13:%d,14:%d,15:%d,16+:%d\n", + rollback_policy.checkpoint_f32 ? "F32" : "default", + rollback_policy.fast_rollback_threshold, + rollback_fast_low, rollback_fast_high, rollback_legacy_replay, + rollback_failed_fallback, + rollback_accept_hist[1], rollback_accept_hist[2], rollback_accept_hist[3], + rollback_accept_hist[4], rollback_accept_hist[5], rollback_accept_hist[6], + rollback_accept_hist[7], rollback_accept_hist[8], rollback_accept_hist[9], + rollback_accept_hist[10], rollback_accept_hist[11], rollback_accept_hist[12], + rollback_accept_hist[13], rollback_accept_hist[14], rollback_accept_hist[15], + rollback_accept_hist[16]); + } }; // kvflash: an in-pool prompt prefills contiguously without registering @@ -2488,11 +2510,12 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // 6. Fix state: adaptive fast-rollback vs legacy replay. // Fast-rollback (implicit bonus, skip replay) is profitable when // accept_n is large enough that skipping the replay saves more compute - // than the cost of deferring the bonus to the next step. Breakeven - // is around accept_n ≈ 5. Below that, legacy replay is cheaper. - constexpr int kFastRollbackThreshold = 5; + // than the cost of deferring the bonus to the next step. The default + // threshold is 5; exact F32 checkpoints may opt in to a lower value. + rollback_accept_hist[std::min(accept_n, 16)]++; const bool use_fast_rollback = - target->supports_fast_rollback() && (accept_n >= kFastRollbackThreshold); + target->supports_fast_rollback() && + (accept_n >= rollback_policy.fast_rollback_threshold); int replay_last_tok = -1; bool fast_rolled_back = false; @@ -2508,15 +2531,19 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (target->rollback_to(committed, commit_n)) { replay_last_tok = target_tok[commit_n - 1]; fast_rolled_back = true; + if (accept_n < 5) rollback_fast_low++; + else rollback_fast_high++; } else { // Rollback failed (CUDA error / unsupported state type). The // pre-verify snapshot is still valid, so degrade to the legacy // restore+replay path below instead of aborting the request. std::fprintf(stderr, "spec-decode: rollback_to failed; " "falling back to restore+replay\n"); + rollback_failed_fallback++; } } if (!fast_rolled_back) { + rollback_legacy_replay++; // Legacy replay: restore SSM snapshot, replay accepted + bonus tokens. // (When falling back from fast-rollback, bonus_tok is -1 and commit_n // is the budget-clamped accepted count.) diff --git a/server/src/qwen35/qwen35_dflash_target.cpp b/server/src/qwen35/qwen35_dflash_target.cpp index dc30ff7c4..19be74052 100644 --- a/server/src/qwen35/qwen35_dflash_target.cpp +++ b/server/src/qwen35/qwen35_dflash_target.cpp @@ -376,13 +376,24 @@ bool Qwen35DFlashTarget::rollback_to_tree( (size_t)rollback_dfs * cap.ssm_intermediate_states->nb[3]; const void * ssm_src = (const char *)cap.ssm_intermediate_states->data + ssm_src_offset; - const auto to_fp32 = ggml_get_to_fp32_cuda(cap.ssm_intermediate_states->type); - if (!to_fp32) { - std::fprintf(stderr, "rollback_to_tree: no fp32 converter for type %d (layer %d)\n", - (int)cap.ssm_intermediate_states->type, il); - return false; + if (cap.ssm_intermediate_states->type == GGML_TYPE_F32) { + const cudaError_t ce = cudaMemcpyAsync(cache_.ssm_state[il]->data, ssm_src, + ssm_elems * sizeof(float), + cudaMemcpyDeviceToDevice, stream); + if (ce != cudaSuccess) { + std::fprintf(stderr, "rollback_to_tree: F32 SSM copy failed at layer %d: %s\n", + il, cudaGetErrorString(ce)); + return false; + } + } else { + const auto to_fp32 = ggml_get_to_fp32_cuda(cap.ssm_intermediate_states->type); + if (!to_fp32) { + std::fprintf(stderr, "rollback_to_tree: no fp32 converter for type %d (layer %d)\n", + (int)cap.ssm_intermediate_states->type, il); + return false; + } + to_fp32(ssm_src, (float *)cache_.ssm_state[il]->data, (int64_t)ssm_elems, stream); } - to_fp32(ssm_src, (float *)cache_.ssm_state[il]->data, (int64_t)ssm_elems, stream); // Conv state ← the K-1 most recent inputs along rollback_dfs's ancestry. const int K_conv = 4; @@ -551,7 +562,6 @@ bool Qwen35DFlashTarget::rollback_to(int base_pos, int commit_n) { cache_.cur_pos = base_pos + commit_n; return true; } - const int rollback_idx = commit_n - 1; // index into per-step intermediates cudaStream_t stream = nullptr; @@ -589,16 +599,30 @@ bool Qwen35DFlashTarget::rollback_to(int base_pos, int commit_n) { (size_t)rollback_idx * cap.ssm_intermediate_states->nb[3]; const void * ssm_src = (const char *)cap.ssm_intermediate_states->data + ssm_src_offset; - const auto to_fp32 = ggml_get_to_fp32_cuda(cap.ssm_intermediate_states->type); - if (!to_fp32) { - if (kFastRollbackDiag) { - std::fprintf(stderr, "rollback_to: no fp32 converter type=%d layer=%d\n", - (int)cap.ssm_intermediate_states->type, il); + if (cap.ssm_intermediate_states->type == GGML_TYPE_F32) { + const size_t ssm_bytes = ssm_elems * sizeof(float); + const cudaError_t ce = cudaMemcpyAsync(cache_.ssm_state[il]->data, ssm_src, + ssm_bytes, + cudaMemcpyDeviceToDevice, stream); + if (ce != cudaSuccess) { + if (kFastRollbackDiag) { + std::fprintf(stderr, "rollback_to: F32 SSM copy failed layer=%d: %s\n", + il, cudaGetErrorString(ce)); + } + return false; } - return false; + } else { + const auto to_fp32 = ggml_get_to_fp32_cuda(cap.ssm_intermediate_states->type); + if (!to_fp32) { + if (kFastRollbackDiag) { + std::fprintf(stderr, "rollback_to: no fp32 converter type=%d layer=%d\n", + (int)cap.ssm_intermediate_states->type, il); + } + return false; + } + to_fp32(ssm_src, (float *)cache_.ssm_state[il]->data, + (int64_t)ssm_elems, stream); } - to_fp32(ssm_src, (float *)cache_.ssm_state[il]->data, - (int64_t)ssm_elems, stream); // Conv rollback: copy conv_input[commit_n..commit_n+K-2, :, :] // into cache.conv_state[il]. diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index ea4684e09..24808177a 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -35,6 +35,7 @@ #include "kv_quant.h" #include "qwen35_ops.h" #include "qwen35moe_ffn.h" +#include "common/chain_rollback_policy.h" #include #include @@ -377,6 +378,12 @@ bool migrate_prefill_cache(const TargetWeights & w, cache.rollback_ctx = ggml_init(ip); if (!cache.rollback_ctx) { set_last_error("rollback cache ggml_init failed"); return false; } + // Preserve the established F16 default. Opt in to PR #506's F32 checkpoint + // representation for single-GPU validation with an explicit environment flag. + const ChainRollbackPolicy rollback_policy = resolve_chain_rollback_policy(); + const ggml_type checkpoint_type = rollback_policy.checkpoint_f32 + ? GGML_TYPE_F32 : GGML_TYPE_F16; + int dn_idx = 0; for (int il = 0; il < w.n_layer; il++) { if (((il + 1) % w.full_attention_interval) != 0) { @@ -384,7 +391,7 @@ bool migrate_prefill_cache(const TargetWeights & w, head_v_dim, head_v_dim, w.ssm_dt_rank); ggml_tensor * Cn = ggml_new_tensor_2d(cache.rollback_ctx, GGML_TYPE_F32, w.ssm_d_conv - 1, conv_ch); - ggml_tensor * Si = ggml_new_tensor_4d(cache.rollback_ctx, GGML_TYPE_F16, + ggml_tensor * Si = ggml_new_tensor_4d(cache.rollback_ctx, checkpoint_type, head_v_dim, head_v_dim, w.ssm_dt_rank, max_verify_tokens); ggml_tensor * Ci = ggml_new_tensor_3d(cache.rollback_ctx, GGML_TYPE_F32, @@ -404,6 +411,18 @@ bool migrate_prefill_cache(const TargetWeights & w, } cache.rollback_buf = ggml_backend_alloc_ctx_tensors(cache.rollback_ctx, backend); + if (rollback_policy.diagnostics) { + size_t checkpoint_bytes = 0; + for (ggml_tensor * t : cache.ssm_intermediate) { + if (t) checkpoint_bytes += ggml_nbytes(t); + } + const ggml_type allocated_checkpoint_type = cache.ssm_intermediate.empty() || !cache.ssm_intermediate[0] + ? GGML_TYPE_COUNT : cache.ssm_intermediate[0]->type; + std::fprintf(stderr, + "[target-single][chain-rollback] checkpoint_dtype=%s delta_layers=%d max_verify_tokens=%d checkpoint_bytes=%zu\n", + allocated_checkpoint_type == GGML_TYPE_COUNT ? "missing" : ggml_type_name(allocated_checkpoint_type), + n_delta, max_verify_tokens, checkpoint_bytes); + } if (!cache.rollback_buf) { set_last_error("ggml_backend_alloc_ctx_tensors failed for rollback cache"); ggml_free(cache.rollback_ctx); @@ -792,6 +811,9 @@ static ggml_tensor * build_delta_net_block( 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 @@ -1044,7 +1066,8 @@ static ggml_tensor * build_single_layer( ggml_tensor * q_tail_capture = nullptr, int q_tail_start = 0, ggml_tensor ** moe_selected_out = nullptr, - ggml_tensor * kv_write_rows = nullptr) + ggml_tensor * kv_write_rows = nullptr, + ggml_tensor * parent_ids = nullptr) { const int hidden = w.n_embd; const float eps = w.rms_eps; @@ -1076,9 +1099,16 @@ static ggml_tensor * build_single_layer( for (int il = 0; il < layer_idx; il++) { if (((il + 1) % w.full_attention_interval) != 0) dn_idx++; } + DeltaNetCapture cap{}; + DeltaNetCapture * cap_ptr = nullptr; + if (capture) { + cap_ptr = ∩ + cap_ptr->ssm_intermediate_states = cache.ssm_intermediate[dn_idx]; + cap_ptr->conv_input = cache.conv_input_cache[dn_idx]; + } cur = build_delta_net_block(ctx, gf, w, L, cur, cache.conv_state[dn_idx], cache.ssm_state[dn_idx], - n_tokens, nullptr, nullptr, + n_tokens, cap_ptr, parent_ids, /*skip_gdn_intermediate=*/true); } @@ -1334,12 +1364,13 @@ ggml_tensor * build_qwen35_layer( int fa_window, ggml_tensor * q_tail_capture, int q_tail_start, - ggml_tensor * kv_write_rows) + ggml_tensor * kv_write_rows, + ggml_tensor * parent_ids) { return build_single_layer(ctx, gf, w, cache, layer_idx, inp, positions, attn_mask, kv_start, n_tokens, capture, fa_window, q_tail_capture, q_tail_start, nullptr, - kv_write_rows); + kv_write_rows, parent_ids); } ggml_tensor * build_qwen35_layer( @@ -1358,12 +1389,13 @@ ggml_tensor * build_qwen35_layer( ggml_tensor * q_tail_capture, int q_tail_start, ggml_tensor ** moe_selected_out, - ggml_tensor * kv_write_rows) + ggml_tensor * kv_write_rows, + ggml_tensor * parent_ids) { return build_single_layer(ctx, gf, w, cache, layer_idx, inp, positions, attn_mask, kv_start, n_tokens, capture, fa_window, q_tail_capture, q_tail_start, moe_selected_out, - kv_write_rows); + kv_write_rows, parent_ids); } QwenLayerPrefnOutputs build_qwen35_layer_prefn( diff --git a/server/test/test_chain_rollback_policy.cpp b/server/test/test_chain_rollback_policy.cpp new file mode 100644 index 000000000..7ae1e2345 --- /dev/null +++ b/server/test/test_chain_rollback_policy.cpp @@ -0,0 +1,67 @@ +#include "chain_rollback_policy.h" + +#include +#include + +using dflash::common::resolve_chain_rollback_policy; + +static int failures = 0; + +#define CHECK(expr) do { \ + if (!(expr)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #expr); \ + ++failures; \ + } \ +} while (0) + +static void clear_policy_env() { + unsetenv("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32"); + unsetenv("DFLASH_FAST_ROLLBACK_THRESHOLD"); + unsetenv("DFLASH_SINGLE_CHAIN_ROLLBACK_DIAG"); +} + +int main() { + clear_policy_env(); + auto policy = resolve_chain_rollback_policy(); + CHECK(!policy.checkpoint_f32); + CHECK(policy.fast_rollback_threshold == 5); + CHECK(!policy.diagnostics); + + // A threshold flag alone must not alter the established F16 policy. + setenv("DFLASH_FAST_ROLLBACK_THRESHOLD", "2", 1); + policy = resolve_chain_rollback_policy(); + CHECK(!policy.checkpoint_f32); + CHECK(policy.fast_rollback_threshold == 5); + + setenv("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32", "1", 1); + policy = resolve_chain_rollback_policy(); + CHECK(policy.checkpoint_f32); + CHECK(policy.fast_rollback_threshold == 2); + + // Boolean flags follow the project's non-empty, non-"0" convention. + setenv("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32", "true", 1); + CHECK(resolve_chain_rollback_policy().checkpoint_f32); + setenv("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32", "yes", 1); + CHECK(resolve_chain_rollback_policy().checkpoint_f32); + setenv("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32", "on", 1); + CHECK(resolve_chain_rollback_policy().checkpoint_f32); + setenv("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32", "0", 1); + CHECK(!resolve_chain_rollback_policy().checkpoint_f32); + setenv("DFLASH_SINGLE_CHAIN_CHECKPOINT_F32", "1", 1); + + // Invalid values degrade safely to the default threshold. + setenv("DFLASH_FAST_ROLLBACK_THRESHOLD", "0", 1); + CHECK(resolve_chain_rollback_policy().fast_rollback_threshold == 5); + setenv("DFLASH_FAST_ROLLBACK_THRESHOLD", "6", 1); + CHECK(resolve_chain_rollback_policy().fast_rollback_threshold == 5); + setenv("DFLASH_FAST_ROLLBACK_THRESHOLD", "garbage", 1); + CHECK(resolve_chain_rollback_policy().fast_rollback_threshold == 5); + + setenv("DFLASH_SINGLE_CHAIN_ROLLBACK_DIAG", "1", 1); + CHECK(resolve_chain_rollback_policy().diagnostics); + + clear_policy_env(); + if (failures != 0) return 1; + std::printf("chain rollback policy tests passed\n"); + return 0; +} diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 7f5473995..fc6a341f4 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -3617,10 +3617,21 @@ int main(int argc, char ** argv) { (size_t)rollback_dfs * cap.ssm_intermediate_states->nb[3]; const void * ssm_src = (const char *)cap.ssm_intermediate_states->data + ssm_src_offset; - ggml_get_to_fp32_cuda(cap.ssm_intermediate_states->type)( - ssm_src, (float *)cache.ssm_state[il]->data, - (int64_t)ssm_elems, stream); - cudaError_t ce = cudaSuccess; // launch error checked in the conv block below + cudaError_t ce = cudaSuccess; + if (cap.ssm_intermediate_states->type == GGML_TYPE_F32) { + ce = cudaMemcpyAsync(cache.ssm_state[il]->data, ssm_src, + ssm_elems * sizeof(float), + cudaMemcpyDeviceToDevice, stream); + } else { + ggml_get_to_fp32_cuda(cap.ssm_intermediate_states->type)( + ssm_src, (float *)cache.ssm_state[il]->data, + (int64_t)ssm_elems, stream); + } + if (ce != cudaSuccess) { + std::fprintf(stderr, "ddtree F32 SSM rollback il=%d: %s\n", + il, cudaGetErrorString(ce)); + return 1; + } // Conv rollback: copy the K-1 most recent inputs along // the rolled-back token's ANCESTRY (not DFS order). Two @@ -3952,10 +3963,21 @@ int main(int argc, char ** argv) { (size_t)rollback_idx * cap.ssm_intermediate_states->nb[3]; const void * ssm_src = (const char *)cap.ssm_intermediate_states->data + ssm_src_offset; - ggml_get_to_fp32_cuda(cap.ssm_intermediate_states->type)( - ssm_src, (float *)cache.ssm_state[il]->data, - (int64_t)ssm_elems, stream); cudaError_t ce = cudaSuccess; + if (cap.ssm_intermediate_states->type == GGML_TYPE_F32) { + ce = cudaMemcpyAsync(cache.ssm_state[il]->data, ssm_src, + ssm_elems * sizeof(float), + cudaMemcpyDeviceToDevice, stream); + } else { + ggml_get_to_fp32_cuda(cap.ssm_intermediate_states->type)( + ssm_src, (float *)cache.ssm_state[il]->data, + (int64_t)ssm_elems, stream); + } + if (ce != cudaSuccess) { + std::fprintf(stderr, "F32 SSM rollback il=%d: %s\n", + il, cudaGetErrorString(ce)); + return 1; + } // ── Conv rollback: copy conv_input[commit_n..commit_n+K-2, :, :] // into cache.conv_state[il].