server : keep the tokens a recompute restore replays out of the request's prompt stats

A recompute restore re-enters prompt processing to put back the cells the park dropped. Those tokens are submitted with is_prompt set, so every one of them was added to n_prompt_processed and pushed t_prompt_last to the end of the replay, while n_gen deliberately carries across the park. The reported prompt length grew with every park, and the generation time covered only the tokens after the last re-prefill, so predicted_per_second was inflated by the ratio of the two.

The replay is still counted in the server-wide prompt metrics, where it is real compute; it just no longer moves the slot's prompt count or the prompt/generation boundary.

Author: danielhanchen

tools/server/server-context.cpp

@@ -5177,8 +5177,12 @@ struct server_context_impl {
                         }
 
                         metrics_queue_prompt(n_tokens_out);
-                        slot.stats.n_prompt_processed += n_tokens_out;
-                        slot.stats.update_prompt_last();
+
+                        // [TAG_PREEMPT] a re-prefill puts back what a park dropped: real compute, counted above, but it is not the request's prompt
+                        if (!slot.preempt_reprefill) {
+                            slot.stats.n_prompt_processed += n_tokens_out;
+                            slot.stats.update_prompt_last();
+                        }
 
                         // add the mtmd chunk to cache
                         {
@@ -5918,7 +5922,8 @@ struct server_context_impl {
             n_prompt_tokens++;
 
             auto & slot = slots[t.id_slot];
-            if (slot.stats.is_set()) {
+            // [TAG_PREEMPT] replayed tokens stay out of the slot's prompt count, they were counted when the request first processed its prompt
+            if (slot.stats.is_set() && !slot.preempt_reprefill) {
                 slot.stats.n_prompt_processed++;
             }
         }
@@ -5936,7 +5941,8 @@ struct server_context_impl {
         for (int i = off; i < off + n_tokens; ++i) {
             const auto & t = batch.tokens[i];
             auto & slot = slots[t.id_slot];
-            if (t.is_prompt && slot.stats.is_set()) {
+            // [TAG_PREEMPT] a re-prefill must not move the prompt/generation boundary: n_gen carries across the park, so the generation time would then cover only the tokens after it
+            if (t.is_prompt && slot.stats.is_set() && !slot.preempt_reprefill) {
                 slot.stats.set_prompt_last(t_now);
             }
         }

tools/server/tests/unit/test_preempt.py

@@ -878,6 +878,26 @@ def test_a_recompute_park_is_reported_to_the_client_and_to_metrics():
     assert plain.body["preempt"] == {"parks": 0, "recomputes": 0}
 
 
+def test_a_recompute_restore_does_not_count_its_replay_as_prompt():
+    # the re-prefill puts back what the park dropped: counted as prompt it would move the prompt/generation boundary, and since n_gen carries across the park the generation time would then cover only the tokens after the last re-prefill
+    os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1"
+    os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8"
+    _start(n_ctx=2048, n_batch=2048)
+
+    n_prompt = 1950
+    n_predict = 24
+    res = _complete(n_predict, _prompt_of(n_prompt, _PROMPT_C))
+
+    assert res.status_code == 200, res.body
+    assert res.body["preempt"]["recomputes"] >= 1, res.body["preempt"]
+
+    timings = res.body["timings"]
+    assert timings["prompt_n"] == n_prompt, timings
+    assert timings["predicted_n"] == n_predict, timings
+    # every re-prefill happens inside the generation, so the generation holds the longer time of the two
+    assert timings["predicted_ms"] > timings["prompt_ms"], timings
+
+
 def test_slots_reports_the_recomputes_of_the_current_task():
     # a reader watching the slots sees the same count the request is given at the end
     os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1"