perf(qwen35): add opt-in F32 rollback checkpoints for single-GPU inference (#515)

* perf(qwen35): add opt-in F32 rollback checkpoints

* fix(qwen35): address rollback policy review findings

* refactor(qwen35): centralize rollback diagnostics

* fix(qwen35): guard null stream in rollback diagnostics

Treat a null diagnostics stream as a no-op before calling fprintf, and cover the diagnostics-enabled null-stream path in the rollback policy test. Addresses the latest Cubic P3 on PR #515.

---------

Co-authored-by: Drew Schuyler <[email protected]>

Author: Amidwestnoob

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})

server/src/common/chain_rollback_policy.h

@@ -0,0 +1,77 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+
+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;
+}
+
+// Rollback diagnostics shared by the two spec-decode loops
+// (Qwen35Backend::do_spec_decode and run_dflash_spec_decode). Keeping the
+// counters and the print format in one place prevents the loops from
+// drifting when the diagnostics change.
+struct RollbackDiag {
+    int accept_hist[17] = {};
+    int fast_low = 0;         // fast rollbacks below the F16 breakeven (accept_n < 5)
+    int fast_high = 0;        // fast rollbacks at accept_n >= 5
+    int legacy_replay = 0;
+    int failed_fallback = 0;
+
+    void record_accept(int accept_n) {
+        accept_hist[std::min(accept_n, 16)]++;
+    }
+    void record_fast_rollback(int accept_n) {
+        if (accept_n < 5) fast_low++;
+        else fast_high++;
+    }
+    void record_failed_fallback() { failed_fallback++; }
+    void record_legacy_replay()   { legacy_replay++; }
+
+    void print(const ChainRollbackPolicy & policy, std::FILE * out) const {
+        if (!out || !policy.diagnostics) return;
+        std::fprintf(out,
+            "[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",
+            policy.checkpoint_f32 ? "F32" : "default",
+            policy.fast_rollback_threshold,
+            fast_low, fast_high, legacy_replay, failed_fallback,
+            accept_hist[1], accept_hist[2], accept_hist[3],
+            accept_hist[4], accept_hist[5], accept_hist[6],
+            accept_hist[7], accept_hist[8], accept_hist[9],
+            accept_hist[10], accept_hist[11], accept_hist[12],
+            accept_hist[13], accept_hist[14], accept_hist[15],
+            accept_hist[16]);
+    }
+};
+
+}  // namespace dflash::common

server/src/common/dflash_spec_decode.cpp

@@ -12,6 +12,8 @@
 #include <cstdio>
 #include <vector>
 
+#include "chain_rollback_policy.h"
+
 namespace dflash::common {
 
 namespace {
@@ -84,6 +86,8 @@ 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();
+    RollbackDiag rollback_diag;
 
     auto t_dec0 = std::chrono::steady_clock::now();
     while (n_generated < n_gen) {
@@ -213,9 +217,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_diag.record_accept(accept_n);
         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<int32_t> replay_tok((size_t)commit_n);
         for (int i = 0; i < commit_n; i++) {
@@ -236,15 +241,18 @@ bool run_dflash_spec_decode(
             if (target.rollback_to(committed, commit_n)) {
                 last_tok = target_tok[commit_n - 1];
                 fast_rolled_back = true;
+                rollback_diag.record_fast_rollback(accept_n);
             } 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_diag.record_failed_fallback();
             }
         }
         if (!fast_rolled_back) {
+            rollback_diag.record_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 +304,7 @@ 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);
+    rollback_diag.print(rollback_policy, stdout);
     if (n_hint_proposed > 0) {
         std::printf("[target-split-dflash] hint tokens: %d/%d accepted (%.1f%%)\n",
                     n_hint_accepted, n_hint_proposed,

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,

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,14 +75,20 @@ 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(
         sg.ctx, sg.gf, w, cache, layer_idx,
         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,

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

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,15 @@ 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();
+    RollbackDiag rollback_diag;
 
     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);
+        rollback_diag.print(rollback_policy, stderr);
     };
 
     // kvflash: an in-pool prompt prefills contiguously without registering
@@ -2488,11 +2492,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_diag.record_accept(accept_n);
         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 +2513,18 @@ 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;
+                rollback_diag.record_fast_rollback(accept_n);
             } 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_diag.record_failed_fallback();
             }
         }
         if (!fast_rolled_back) {
+            rollback_diag.record_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.)

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].