server: stabilize DFlash np2 scheduling

Anbeeld · Anbeeld · commit d33c4c01bde4 · 2026-06-03T04:39:13.000+02:00
Prioritize pending DFlash prompts and initial profit baseline cycles before target-generation rows. Restrict multi-seq target verification to one row per slot because multi-row capture only preserves the final internal ubatch, and rotate multi-row verifier cycles across slots to keep np2 runs fair.
diff --git a/tests/test-dflash-plumbing.cpp b/tests/test-dflash-plumbing.cpp
@@ -1766,6 +1766,27 @@ int main(int argc, char ** argv) {
     ok &= expect(server_context.find("dflash_multiseq_rows") != std::string::npos &&
                  server_context.find("uneven-rows") != std::string::npos,
         "server must not multi-seq batch DFlash target views with uneven per-slot rows because GPU hidden capture is per-ubatch");
+    ok &= expect(server_context.find("dflash_batch_has_tg") != std::string::npos &&
+                 server_context.find("!dflash_batch_has_tg && (params_base.cont_batching || batch.n_tokens == 0)") != std::string::npos,
+        "server must defer prompt batching when the current DFlash batch already contains TG verifier rows");
+    ok &= expect(server_context.find("dflash_try_reserve_tg_rows") != std::string::npos &&
+                 server_context.find("dflash_skip_tg_slot_for_next_cycle") != std::string::npos &&
+                 server_context.find("dflash uneven TG rows deferred") != std::string::npos,
+        "server must defer uneven DFlash TG slots instead of decoding a combined view with partial hidden capture");
+    ok &= expect(server_context.find("dflash_multi_seq_tg_rows_safe") != std::string::npos &&
+                 server_context.find("rows == 1") != std::string::npos,
+        "server must only multi-seq batch DFlash target verification when each slot contributes one row");
+    ok &= expect(server_context.find("dflash_next_tg_slot") != std::string::npos &&
+                 server_context.find("dflash_tg_slot_start") != std::string::npos,
+        "server must fairly rotate one-slot DFlash multi-row verifier cycles instead of always starting at slot 0");
+    ok &= expect(server_context.find("dflash_has_pending_prompt") != std::string::npos &&
+                 server_context.find("dflash_has_pending_prompt && slot.can_speculate()") != std::string::npos &&
+                 server_context.find("params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH &&\n            needs_reeval &&\n            std::any_of(slots.begin(), slots.end(), [](const server_slot & slot) {\n                return slot.state == SLOT_STATE_STARTED || slot.state == SLOT_STATE_PROCESSING_PROMPT;\n            })") == std::string::npos,
+        "DFlash must prioritize unfinished prompt prefill before TG rows on all target types, not only recurrent/hybrid targets");
+    ok &= expect(server_context.find("dflash_has_profit_baseline_slot") != std::string::npos &&
+                 server_context.find("dflash_force_profit_baseline_cycle") != std::string::npos &&
+                 server_context.find("dflash_force_profit_baseline_cycle && n_draft_max > 0") != std::string::npos,
+        "DFlash adaptive profit baseline slots must get one-row target cycles instead of starving behind multi-row verifier batches");
     ok &= expect(common_h.find("int32_t dflash_max_slots = 0") != std::string::npos &&
                  server_context.find("params_base.speculative.dflash_max_slots > 0") != std::string::npos &&
                  server_context.find("DFlash enabled for all %d slots") != std::string::npos,
@@ -1786,7 +1807,7 @@ int main(int argc, char ** argv) {
     ok &= expect(server_context.find("dflash_multislot_flat_accept_barrier") != std::string::npos &&
                  server_context.find("struct dflash_flat_accept_prefetch") != std::string::npos &&
                  server_context.find("common_speculative_update_logits_deferred_dflash_kv(slot.get_spec(), ctx_tgt, batch_tokens, prefetched.n_hidden_keep);") != std::string::npos &&
-                 server_context.find("dflash_recurrent_has_pending_prompt") != std::string::npos &&
+                 server_context.find("dflash_has_pending_prompt") != std::string::npos &&
                  server_context.find("dflash_recurrent_draft_rr") == std::string::npos &&
                  server_context.find("dflash_recurrent_cycle_slot_id") == std::string::npos,
         "DFlash recurrent/hybrid multi-slot scheduling must isolate prompt prefill, pre-sample every flat speculative slot before rollback, and remove the old one-slot recurrent round-robin");
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
@@ -1879,6 +1879,7 @@ struct server_context_impl {
     int trace = 0;
     int slots_debug = 0;
     int n_empty_consecutive = 0;
+    int dflash_next_tg_slot = 0;
 
     std::unique_ptr<server_prompt_cache> prompt_cache;
 
@@ -4421,16 +4422,64 @@ struct server_context_impl {
         std::vector<llama_tokens> batched_drafts(slots.size());
         std::vector<std::vector<float>> batched_log_probs(slots.size());
         const bool use_rejection_sampling = params_base.speculative.sample_temp > 0.0f && params_base.sampling.temp > 0.0f;
-        const bool dflash_recurrent_has_pending_prompt =
+        int dflash_tg_batch_rows = -1;
+        auto dflash_multi_seq_tg_rows_safe = [](int rows) {
+            return rows == 1;
+        };
+        auto dflash_try_reserve_tg_rows = [&](const server_slot & slot, int rows) {
+            if (params_base.speculative.type() != COMMON_SPECULATIVE_TYPE_DFLASH ||
+                    params_base.speculative.branch_budget != 0 ||
+                    rows <= 0) {
+                return true;
+            }
+
+            if (batch.n_tokens == 0 || dflash_tg_batch_rows < 0) {
+                dflash_tg_batch_rows = rows;
+                return true;
+            }
+
+            if (dflash_tg_batch_rows == rows && dflash_multi_seq_tg_rows_safe(rows)) {
+                return true;
+            }
+
+            if (dflash_server_crash_trace_enabled()) {
+                SLT_INF(slot, "dflash uneven TG rows deferred: reason=%s batch_rows=%d slot_rows=%d batch_tokens=%d\n",
+                        dflash_tg_batch_rows == rows ? "multi-row" : "uneven",
+                        dflash_tg_batch_rows, rows, batch.n_tokens);
+            }
+            return false;
+        };
+        auto dflash_skip_tg_slot_for_next_cycle = [](server_slot & slot) {
+            slot.spec_i_batch.clear();
+            slot.spec_pad_i_batch.clear();
+            slot.spec_draft.clear();
+            slot.draft_log_probs.clear();
+        };
+        const bool dflash_has_pending_prompt =
             params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH &&
-            needs_reeval &&
             std::any_of(slots.begin(), slots.end(), [](const server_slot & slot) {
-                return slot.state == SLOT_STATE_STARTED || slot.state == SLOT_STATE_PROCESSING_PROMPT;
+                return slot.can_speculate() &&
+                    (slot.state == SLOT_STATE_STARTED || slot.state == SLOT_STATE_PROCESSING_PROMPT);
             });
+        const bool dflash_has_profit_baseline_slot =
+            params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH &&
+            params_base.speculative.branch_budget == 0 &&
+            std::any_of(slots.begin(), slots.end(), [](const server_slot & slot) {
+                return slot.state == SLOT_STATE_GENERATING &&
+                    slot.can_speculate() &&
+                    slot.dm_adaptive &&
+                    server_adaptive_dm_uses_profit_controller(slot.dm_controller) &&
+                    (slot.profit_baseline_probe_pending ||
+                        !slot.profit_baseline_ready() ||
+                        slot.adaptive_n_max < 0);
+            });
+        const bool dflash_force_profit_baseline_cycle =
+            dflash_has_profit_baseline_slot && !dflash_has_pending_prompt;
         std::vector<int> dflash_cohort_n_draft_max(slots.size(), -1);
         if (params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH &&
                 params_base.speculative.branch_budget == 0 &&
-                !dflash_recurrent_has_pending_prompt) {
+                !dflash_has_pending_prompt &&
+                !dflash_force_profit_baseline_cycle) {
             std::vector<int> normal_n_max(slots.size(), 0);
             std::vector<int> cohort_cap_n_max(slots.size(), 0);
 
@@ -4474,7 +4523,8 @@ struct server_context_impl {
                 if (slot.state == SLOT_STATE_GENERATING &&
                         slot.can_speculate() &&
                         dflash_scheduler_n_draft_max(slot, false) > 0 &&
-                        !dflash_recurrent_has_pending_prompt) {
+                        !dflash_has_pending_prompt &&
+                        !dflash_force_profit_baseline_cycle) {
                     n_drafting++;
                 }
             }
@@ -4491,7 +4541,8 @@ struct server_context_impl {
                     if (slot.state == SLOT_STATE_GENERATING &&
                             slot.can_speculate() &&
                             dflash_scheduler_n_draft_max(slot, false) > 0 &&
-                            !dflash_recurrent_has_pending_prompt) {
+                            !dflash_has_pending_prompt &&
+                            !dflash_force_profit_baseline_cycle) {
                         batch_specs.push_back(slot.get_spec());
                         batch_id_lasts.push_back(slot.sampled);
                         batch_slot_ids.push_back(slot.id);
@@ -4531,14 +4582,26 @@ struct server_context_impl {
         bool ddtree_batch_active = false;
         std::vector<server_slot *> mtp_drafting;
         std::vector<server_slot *> mtp_batching; // slots with existing draft that need batch building but not re-drafting
+        const bool dflash_rotate_tg_slots =
+            params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH &&
+            params_base.speculative.branch_budget == 0 &&
+            !slots.empty();
+        const int dflash_tg_slot_start = dflash_rotate_tg_slots
+            ? std::clamp(dflash_next_tg_slot, 0, (int) slots.size() - 1)
+            : 0;
 
         // first, add sampled tokens from any ongoing sequences
-        for (auto & slot : slots) {
+        for (int slot_offset = 0; slot_offset < (int) slots.size(); ++slot_offset) {
+            const int slot_index = dflash_rotate_tg_slots
+                ? (dflash_tg_slot_start + slot_offset) % (int) slots.size()
+                : slot_offset;
+            auto & slot = slots[slot_index];
+
             if (slot.state != SLOT_STATE_GENERATING) {
                 continue;
             }
 
-            if (dflash_recurrent_has_pending_prompt && slot.can_speculate()) {
+            if (dflash_has_pending_prompt && slot.can_speculate()) {
                 continue;
             }
 
@@ -4554,6 +4617,9 @@ struct server_context_impl {
                     params_base.speculative.branch_budget == 0) {
                 n_draft_max = dflash_flat_effective_draft_max(ctx_dft_shared.get(), n_draft_max);
             }
+            if (dflash_force_profit_baseline_cycle && n_draft_max > 0) {
+                n_draft_max = 0;
+            }
             const bool dflash_active_grammar =
                 params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH &&
                 slot.smpl && common_sampler_has_active_grammar(slot.smpl.get());
@@ -4730,11 +4796,20 @@ struct server_context_impl {
 
                 // MTP batch building is deferred to the post-loop phase below
                 if (!use_mtp_spec) {
+                    const bool small_draft =
+                        slot.spec_draft.empty() ||
+                        slot.task->params.speculative.n_min > (int) slot.spec_draft.size();
+                    const int dflash_slot_tg_rows = small_draft ? 1 : 1 + (int) slot.spec_draft.size();
+                    if (!dflash_try_reserve_tg_rows(slot, dflash_slot_tg_rows)) {
+                        dflash_skip_tg_slot_for_next_cycle(slot);
+                        continue;
+                    }
+
                     slot.spec_i_batch.push_back(batch.n_tokens);
                     common_batch_add(batch, slot.sampled, slot.prompt.tokens.pos_next(), { slot.id }, true);
                     slot.prompt.tokens.push_back(slot.sampled);
 
-                    if (slot.spec_draft.empty() || slot.task->params.speculative.n_min > (int) slot.spec_draft.size()) {
+                    if (small_draft) {
                         if (!slot.spec_draft.empty()) {
                             SLT_DBG(slot, "ignoring small draft: %d < %d\n", (int) slot.spec_draft.size(), slot.task->params.speculative.n_min);
                         }
@@ -4828,6 +4903,11 @@ struct server_context_impl {
                 t_draft_total += ggml_time_us() - t_draft_slot_start;
                 n_slots_drafted++;
             } else {
+                if (!dflash_try_reserve_tg_rows(slot, 1)) {
+                    dflash_skip_tg_slot_for_next_cycle(slot);
+                    continue;
+                }
+
                 if (slot.can_speculate() && slot.dm_adaptive &&
                         server_adaptive_dm_uses_profit_controller(slot.dm_controller) &&
                         slot.profit_expects_baseline_sample()) {
@@ -4920,12 +5000,15 @@ struct server_context_impl {
         // track how many TG tokens are in the batch vs total, to detect
         // pure-verify batches where multi-seq batching is safe.
         const int32_t n_tg_tokens = batch.n_tokens;
+        const bool dflash_batch_has_tg =
+            params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH &&
+            n_tg_tokens > 0;
 
         float  alora_scale       = -1.0f;
         size_t alora_disabled_id = 0;
 
         // next, batch any pending prompts without exceeding n_batch
-        if (!ddtree_batch_active && (params_base.cont_batching || batch.n_tokens == 0)) {
+        if (!ddtree_batch_active && !dflash_batch_has_tg && (params_base.cont_batching || batch.n_tokens == 0)) {
             for (auto & slot : slots) {
                 if (!slot.is_processing()) {
                     continue;
@@ -7142,6 +7225,13 @@ struct server_context_impl {
         if (n_slots_drafted > 0) {
             const int64_t t_cycle_total = ggml_time_us() - t_cycle_start;
             const int64_t t_other = t_cycle_total - t_draft_total - t_verify_total - t_accept_total;
+            if (params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH &&
+                    params_base.speculative.branch_budget == 0 &&
+                    dflash_cycle_slot &&
+                    dflash_cycle_slot->id >= 0 &&
+                    !slots.empty()) {
+                dflash_next_tg_slot = (dflash_cycle_slot->id + 1) % (int) slots.size();
+            }
             if (profile_dflash_cycle && n_slots_drafted == 1 && dflash_cycle_slot &&
                     params_base.speculative.type() == COMMON_SPECULATIVE_TYPE_DFLASH) {
                 const common_dflash_ring_stats ring_stats =
