From 65f8f5464fdc03d7f1b2ad5a3f3aff023f1f48dc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 12:18:50 +0000 Subject: [PATCH 01/81] server: tell a streaming client when its slot is parked and restored A slot parked by the preemption path produces nothing until its cells come back, and to a client that is indistinguishable from a hung server: the stream goes silent, read timeouts fire, and a chat that was merely waiting for room is torn down as broken. Push a small out-of-band result to the task's response queue when a streaming slot is parked and when it is restored. The HTTP layer writes it as an SSE comment, ": preempted" and ": resumed", which is legal SSE that every existing client ignores, so the body of the response is unchanged by preemption. While parked the ping runs every 2 s as ": preempt-keepalive" regardless of --sse-ping, so proxies and client read timeouts survive a wait that is long by design. Notices that arrive before the first real result (a slot parked while it was still processing its prompt) are sent in front of it. Non-streaming requests see nothing. Harness test: forced parks every 8 tokens on /completion and /v1/chat/completions carry the comments in park/resume order and generate the same tokens as the unparked run; a non-streaming request is untouched; two streams that overflow the pool together both finish and the parked one says so. --- tools/server/server-context.cpp | 60 ++++++- tools/server/server-task.cpp | 7 + tools/server/server-task.h | 14 ++ .../server/tests/unit/test_preempt_notify.py | 168 ++++++++++++++++++ 4 files changed, 241 insertions(+), 8 deletions(-) create mode 100644 tools/server/tests/unit/test_preempt_notify.py diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 6723c51397e..74614d13f6a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -76,6 +76,7 @@ enum slot_state { // room the copy goes back and the slot carries on with the same sampler, the same generated // text and the same open stream. A streaming client sees a pause, not an error. constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation +constexpr int64_t PREEMPT_KEEPALIVE_MS = 2000; // SSE keepalive period while a streaming slot is parked constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked @@ -2106,6 +2107,25 @@ struct server_context_impl { queue_results.send(std::move(res)); } + // [TAG_PREEMPT] tell a streaming client that its slot was parked or restored. The + // HTTP layer turns this into an SSE comment, so a client that does not know about + // preemption sees nothing, and one that does can show a pause instead of a stall. + void send_preempt_notice(server_slot & slot, bool parked) { + if (!slot.task || !slot.task->params.stream) { + return; + } + + auto res = std::make_unique(); + + res->id = slot.task->id; + res->index = slot.task->index; + res->id_slot = slot.id; + res->parked = parked; + res->n_preempt = slot.n_preempt; + + queue_results.send(std::move(res)); + } + void send_partial_response(server_slot & slot, const completion_token_output & tkn, bool is_progress, bool is_begin = false) { auto res = std::make_unique(); @@ -3092,6 +3112,8 @@ struct server_context_impl { metrics.n_resume++; + send_preempt_notice(*best, false); + SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", (ggml_time_us() - best->t_preempt_us) / 1e6, best->prompt.n_tokens(), @@ -3109,6 +3131,8 @@ struct server_context_impl { slot.preempt_save()) { metrics.n_preempt++; + send_preempt_notice(slot, true); + SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); } @@ -3145,6 +3169,8 @@ struct server_context_impl { metrics.n_preempt++; + send_preempt_notice(*victim, true); + SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, (ggml_time_us() - t_start) / 1e3, @@ -4739,7 +4765,16 @@ std::unique_ptr server_routes::handle_completions_impl( // in streaming mode, the first error must be treated as non-stream response // this is to match the OAI API behavior // ref: https://github.com/ggml-org/llama.cpp/pull/16486#discussion_r2419657309 + // [TAG_PREEMPT] a slot can be parked while still processing its prompt, before any + // token exists. Those notices arrive ahead of the first real result; keep them and + // send them in front of it, so the client learns about the wait it just had. + std::string preempt_prefix; auto first_result = rd.next(req.should_stop); + while (first_result != nullptr && dynamic_cast(first_result.get()) != nullptr) { + const auto * notice = static_cast(first_result.get()); + preempt_prefix += notice->parked ? ": preempted\n\n" : ": resumed\n\n"; + first_result = rd.next(req.should_stop); + } if (first_result == nullptr) { GGML_ASSERT(req.should_stop()); return res; // connection is closed @@ -4759,17 +4794,17 @@ std::unique_ptr server_routes::handle_completions_impl( // to be sent immediately json first_result_json = first_result->to_json(); if (first_result_json == nullptr) { - res->data = ""; // simply send HTTP headers and status code + res->data = preempt_prefix; // simply send HTTP headers and status code } else if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { - res->data = format_anthropic_sse(first_result_json); + res->data = preempt_prefix + format_anthropic_sse(first_result_json); } else if (res_type == TASK_RESPONSE_TYPE_OAI_RESP) { - res->data = format_oai_resp_sse(first_result_json); + res->data = preempt_prefix + format_oai_resp_sse(first_result_json); } else { - res->data = format_oai_sse(first_result_json); + res->data = preempt_prefix + format_oai_sse(first_result_json); } res->status = 200; res->content_type = "text/event-stream"; - res->set_next([res_this = res.get(), res_type, sse_ping_interval](std::string & output) -> bool { + res->set_next([res_this = res.get(), res_type, sse_ping_interval, parked = false](std::string & output) mutable -> bool { static auto format_error = [](task_response_type res_type, const json & res_json) { if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { return format_anthropic_sse({ @@ -4820,10 +4855,14 @@ std::unique_ptr server_routes::handle_completions_impl( // receive subsequent results bool timeout = false; int64_t start_time = ggml_time_ms(); - auto result = rd.next([&timeout, &start_time, sse_ping_interval, &effective_should_stop]() { + // [TAG_PREEMPT] a parked slot produces nothing for as long as the pool is + // full, so while parked the ping runs every 2 s regardless of --sse-ping and + // is named, so a client can tell "waiting for cells" from "slow". + const int64_t ping_ms = parked ? PREEMPT_KEEPALIVE_MS : (sse_ping_interval > 0 ? (int64_t) sse_ping_interval * 1000 : -1); + auto result = rd.next([&timeout, &start_time, ping_ms, &effective_should_stop]() { if (effective_should_stop()) { return true; // should_stop condition met - } else if (sse_ping_interval > 0 && ggml_time_ms() - start_time > (int64_t)sse_ping_interval * 1000) { + } else if (ping_ms > 0 && ggml_time_ms() - start_time > ping_ms) { timeout = true; return true; // timeout } @@ -4833,7 +4872,7 @@ std::unique_ptr server_routes::handle_completions_impl( if (timeout) { // some clients may time out (e.g. undici) will time out if no data is received for a while, so we need to send a ping to keep the connection alive SRV_DBG("%s", "sending SSE ping\n"); - output = ":\n\n"; + output = parked ? ": preempt-keepalive\n\n" : ":\n\n"; return true; } @@ -4849,6 +4888,11 @@ std::unique_ptr server_routes::handle_completions_impl( output = format_error(res_type, res_json); SRV_DBG("%s", "error received during streaming, terminating stream\n"); return false; // terminate on error + } else if (const auto * notice = dynamic_cast(result.get())) { + // [TAG_PREEMPT] an SSE comment: invisible to clients that do not know + // about preemption, a pause indicator for the ones that do + parked = notice->parked; + output = parked ? ": preempted\n\n" : ": resumed\n\n"; } else { GGML_ASSERT( dynamic_cast(result.get()) != nullptr diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 9afe3c7f06a..48fc77c960f 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1023,6 +1023,13 @@ void server_task_result_cmpl_partial::update(task_result_state & state) { } } +json server_task_result_preempt_notice::to_json() { + return json { + {"preempted", parked}, + {"n_preempt", n_preempt}, + }; +} + json server_task_result_cmpl_partial::to_json() { GGML_ASSERT(is_updated && "update() must be called before to_json()"); if (is_begin) { diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 00734924bc6..f20f891655f 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -392,6 +392,20 @@ struct server_task_result_cmpl_final : server_task_result { json to_json_anthropic_stream(); }; +// [TAG_PREEMPT] out-of-band notice for a streaming task whose slot was parked or restored. +// Serialised as an SSE comment (": preempted", ": resumed"), which every existing client +// ignores, so the body of the response is unchanged by preemption. Never sent to a +// non-streaming task. +struct server_task_result_preempt_notice : server_task_result { + bool parked = false; // true when the slot was just parked, false when restored + int32_t n_preempt = 0; // how many times this task has been parked so far + + virtual bool is_stop() override { + return false; + } + virtual json to_json() override; +}; + struct server_task_result_cmpl_partial : server_task_result { std::string content; llama_tokens tokens; diff --git a/tools/server/tests/unit/test_preempt_notify.py b/tools/server/tests/unit/test_preempt_notify.py new file mode 100644 index 00000000000..86a4e66a674 --- /dev/null +++ b/tools/server/tests/unit/test_preempt_notify.py @@ -0,0 +1,168 @@ +import os +import tempfile +import pytest +import requests +from utils import * + +# [TAG_PREEMPT] A streaming client is told when its slot is parked and when it is +# restored, as SSE comments, and the body is byte for byte what it is without any park. +# Comments are legal SSE that every existing client ignores; a client that knows about +# preemption can show "paused" instead of a dead stream, and a keepalive every 2 s while +# parked keeps proxies and read timeouts from giving up on a wait that is by design long. + +server = ServerPreset.tinyllama2() + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.n_slots = 2 + server.kv_unified = True + server.temperature = 0.0 + server.seed = 42 + # A build without libcurl cannot fetch the model itself; point it at a local copy. + local = os.environ.get("LLAMA_SERVER_TEST_MODEL") + if local: + server.model_hf_repo = None + server.model_hf_file = None + server.model_file = local + fd, server.log_path = tempfile.mkstemp(suffix=".log") + os.close(fd) + yield + os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + + +def _stream_raw(path: str, data: dict) -> tuple[list[str], list[str]]: + """The SSE lines of one streaming request: (comment lines, data lines).""" + url = f"http://{server.server_host}:{server.server_port}{path}" + res = requests.post(url, json=data, stream=True) + assert res.status_code == 200 + comments, datas = [], [] + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if line.startswith(":"): + comments.append(line) + elif line.startswith("data: "): + datas.append(line[6:]) + return comments, datas + + +def _content(datas: list[str]) -> str: + out = "" + for d in datas: + if d == "[DONE]": + break + j = json.loads(d) + if "content" in j: + out += j["content"] + for ch in j.get("choices", []) or []: + delta = ch.get("delta") or {} + out += delta.get("content") or "" + return out + + +def _completion_payload(n_predict: int) -> dict: + return { + "n_predict": n_predict, + "prompt": "Hi how are you", + "ignore_eos": True, + "temperature": 0.0, + "seed": 42, + "stream": True, + } + + +def _chat_payload(n_predict: int) -> dict: + return { + "max_tokens": n_predict, + "messages": [{"role": "user", "content": "Hi how are you"}], + "temperature": 0.0, + "seed": 42, + "stream": True, + } + + +def test_a_stream_announces_its_parks_and_the_body_is_unchanged(): + global server + server.n_ctx = 512 + server.start() + ref_comments, ref_datas = _stream_raw("/completion", _completion_payload(64)) + assert not any(c.startswith(": preempted") or c.startswith(": resumed") for c in ref_comments) + assert _content(ref_datas) + server.stop() + + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + comments, datas = _stream_raw("/completion", _completion_payload(64)) + + parked = [c for c in comments if c == ": preempted"] + resumed = [c for c in comments if c == ": resumed"] + assert len(parked) >= 6, comments + assert len(resumed) == len(parked), comments + # Every park is followed by its resume before the next park. + seq = [c for c in comments if c in (": preempted", ": resumed")] + assert seq == [": preempted", ": resumed"] * len(parked), seq + # The generated text is byte for byte the unparked text, token by token. Only the + # final chunk's wall-clock timings differ between the two runs. + def _pieces(ds): + return [json.loads(d).get("content") for d in ds if d != "[DONE]"] + + assert _pieces(datas) == _pieces(ref_datas) + assert _content(datas) == _content(ref_datas) + final, ref_final = json.loads(datas[-1]), json.loads(ref_datas[-1]) + assert final["tokens_predicted"] == ref_final["tokens_predicted"] == 64 + + +def test_the_oai_chat_stream_carries_the_same_comments(): + global server + server.n_ctx = 512 + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + comments, datas = _stream_raw("/v1/chat/completions", _chat_payload(48)) + assert ": preempted" in comments and ": resumed" in comments + assert datas[-1] == "[DONE]" + assert _content(datas) + + +def test_non_streaming_requests_see_nothing(): + global server + server.n_ctx = 512 + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + res = server.make_request("POST", "/completion", data={ + "n_predict": 32, + "prompt": "Hi how are you", + "ignore_eos": True, + "temperature": 0.0, + "seed": 42, + }) + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == 32 + assert "preempted" not in res.body + + +def test_two_overflowing_streams_both_finish_and_the_parked_one_says_so(): + # The pair from test_preempt: each alone fits, together they do not, so one is + # parked until the other finishes. The parked stream must carry the comments and + # finish with its full output. + global server + server.n_ctx = 256 + server.start() + + n_predict = 160 + p1 = _completion_payload(n_predict) | {"prompt": "Once upon a time there was a brave knight who"} + p2 = _completion_payload(n_predict) | {"prompt": "The quick brown fox jumps over the lazy dog and"} + results = parallel_function_calls([ + (_stream_raw, ("/completion", p1)), + (_stream_raw, ("/completion", p2)), + ]) + announced = 0 + for comments, datas in results: + final = json.loads([d for d in datas if d != "[DONE]"][-1]) + assert final["timings"]["predicted_n"] == n_predict + assert final["truncated"] is False + if ": preempted" in comments: + announced += 1 + assert ": resumed" in comments + assert announced >= 1, [r[0] for r in results] From 601cfe44334819a0a9ae4c1b8b32df8a81dc40a0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 19:49:31 +0000 Subject: [PATCH 02/81] cuda: add GGML_CUDA_BATCH_INVARIANT so a row does not depend on its batch The number of tokens in a batch selects the matmul implementation, the flash attention kernel, and inside several of them how the K loop or the KV cache is divided between threads and blocks. All of those change the order in which the partial products of one destination element are summed, so a request decoding next to three others produces different bits than the same request decoding alone, even at temperature 0. GGML_CUDA_BATCH_INVARIANT=1 computes every destination column, and attends every query row, with the configuration a batch of one would use. =2 does the same but only where the batch-of-one configuration actually differs, which leaves the quantized projections batched because MMVQ already uses the same nwarps for one to four columns. Flash attention additionally pins the vector kernel, pins the split over the KV cache to one block per tile, and scans the mask for the sequence's own extent, so neither the query count nor the length of a shared KV cache selects the algorithm. Measured on a B200 with Qwen3.5-4B-UD-Q4_K_XL: with the KV cache state held equal, the 1831 node decode graph goes from 1190 nodes whose sequence-0 row differs between a one token and a four token ubatch to 0, and 256 greedy decode steps that diverged at step 125 become identical. --- ggml/src/ggml-cuda/common.cuh | 3 + ggml/src/ggml-cuda/fattn-common.cuh | 14 ++- ggml/src/ggml-cuda/fattn.cu | 37 ++++++ ggml/src/ggml-cuda/ggml-cuda.cu | 173 +++++++++++++++++++++++----- ggml/src/ggml-cuda/mmvq.cu | 15 +++ ggml/src/ggml-cuda/mmvq.cuh | 5 + 6 files changed, 217 insertions(+), 30 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 14dd1098c97..51ad1d6aa4d 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -49,6 +49,9 @@ #define GGML_CUDA_CC_PASCAL 600 #define GGML_CUDA_CC_DP4A 610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products +// [TAG_BATCH_INVARIANT] 0 = off, 1 = split every batched matmul, 2 = split only where it changes bits +int ggml_cuda_batch_invariant(); + #define GGML_CUDA_CC_VOLTA 700 #define GGML_CUDA_CC_TURING 750 #define GGML_CUDA_CC_AMPERE 800 diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index e67cc7fdf78..123e5071911 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1091,7 +1091,10 @@ void launch_fattn( // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped. // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or // multiple sequences of possibly different lengths. - if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) { + // [TAG_BATCH_INVARIANT] Without this scan the KV loop runs to K->ne[1], which grows with the + // other sequences sharing the cache. Scanning the mask bounds it by the sequence's own extent. + const bool batch_invariant_KV_max = ggml_cuda_batch_invariant() != 0; + if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) { const int64_t s31 = mask->nb[1] / sizeof(half2); const int64_t s33 = mask->nb[3] / sizeof(half2); @@ -1148,6 +1151,15 @@ void launch_fattn( if (ntiles_dst % blocks_num.x != 0) { // Fixup is only needed if the SMs work on fractional tiles. dst_tmp_meta.alloc((size_t(blocks_num.x) * ncols * (2 + DV/2))); } + } else if (ggml_cuda_batch_invariant()) { + // [TAG_BATCH_INVARIANT] How the KV cache is split between blocks, and therefore the order + // in which the partial attention results are combined, follows K->ne[1]. That length grows + // with the other sequences sharing the cache, so pin the split to a single block per tile. + parallel_blocks = 1; + + blocks_num.x = ntiles_x; + blocks_num.y = parallel_blocks; + blocks_num.z = ntiles_z_gqa*K->ne[2]*Q->ne[3]; } else { // parallel_blocks must not be larger than what the tensor size allows: parallel_blocks = std::min(parallel_blocks, ntiles_KV); diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ab7a3b297c0..b2f6f0660ee 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -457,6 +457,13 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path. const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0; + // [TAG_BATCH_INVARIANT] Every choice below switches on Q->ne[1] or on K->ne[1], and both + // grow with the other sequences in the batch and in the shared KV cache. Pin the kernel a + // batch of one would use so a request is never moved onto a different algorithm by its neighbours. + if (ggml_cuda_batch_invariant() && can_use_vector_kernel && Q->ne[1] == 1) { + return BEST_FATTN_KERNEL_VEC; + } + // If Turing tensor cores are available, use them: if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { if (can_use_vector_kernel) { @@ -569,6 +576,36 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); + + // [TAG_BATCH_INVARIANT] Attend one query row at a time, as a batch of one would. + if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1) { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * mask = dst->src[3]; + + for (int64_t i = 0; i < Q->ne[1]; ++i) { + ggml_tensor Q_row = *Q; + Q_row.ne[1] = 1; + Q_row.data = (char *) Q->data + i*Q->nb[1]; + + ggml_tensor mask_row; + ggml_tensor dst_row = *dst; + // ne[2] keeps running to the end of dst so that the scratch space for F16 copies of + // K and V, which is placed right behind dst, is still put in the same place. + dst_row.ne[2] = dst->ne[2] - i; + dst_row.data = (char *) dst->data + i*dst->nb[2]; + dst_row.src[0] = &Q_row; + if (mask) { + mask_row = *mask; + mask_row.ne[1] = 1; + mask_row.data = (char *) mask->data + i*mask->nb[1]; + dst_row.src[3] = &mask_row; + } + + ggml_cuda_flash_attn_ext(ctx, &dst_row); + } + return; + } + switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2456f7dcc62..f9aa99ad003 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1758,6 +1758,12 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, } static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { + // [TAG_BATCH_INVARIANT] mul_mat+GLU is only fused for a single destination column, so + // leaving it on would give a solo request a different code path from a batched one. + if (ggml_cuda_batch_invariant()) { + return false; + } + ggml_tensor * src0 = tensor->src[0]; ggml_tensor * src1 = tensor->src[1]; const ggml_tensor * dst = tensor; @@ -1785,6 +1791,12 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { } static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { + // [TAG_BATCH_INVARIANT] mul_mat+GLU is only fused for a single destination column, so + // leaving it on would give a solo request a different code path from a batched one. + if (ggml_cuda_batch_invariant()) { + return false; + } + ggml_tensor * src0 = tensor->src[0]; ggml_tensor * src1 = tensor->src[1]; const ggml_tensor * dst = tensor; @@ -1813,60 +1825,163 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } -static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - GGML_TENSOR_BINARY_OP_LOCALS +// [TAG_BATCH_INVARIANT] +// The number of tokens in a batch picks both the matmul implementation below and, inside +// several of them, how the K loop is divided between threads. Both change the order in +// which the partial products of one destination element are summed, so the same request +// produces different bits depending on how many other requests decode alongside it. +// +// GGML_CUDA_BATCH_INVARIANT removes that dependency: +// 1 - compute every destination column on its own, exactly as a batch of one would. +// 2 - split off only the columns whose batch-of-one configuration differs from the +// batched one, leaving the already invariant matmuls batched. +int ggml_cuda_batch_invariant() { + static const int mode = []() { + const char * val = getenv("GGML_CUDA_BATCH_INVARIANT"); + return val ? atoi(val) : 0; + }(); + return mode; +} - const int32_t hint = ggml_get_op_params_i32(dst, 1); - if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { - return; - } +enum ggml_cuda_mm_path { + GGML_CUDA_MM_CUBLAS_UNSUPPORTED, + GGML_CUDA_MM_MMVF, + GGML_CUDA_MM_MMVF_TRANSPOSED, + GGML_CUDA_MM_MMF, + GGML_CUDA_MM_MMVQ, + GGML_CUDA_MM_MMQ, + GGML_CUDA_MM_CUBLAS, +}; +// The implementation ggml_cuda_mul_mat would pick for a batch of ne11 columns. +static ggml_cuda_mm_path ggml_cuda_mul_mat_path( + int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, int64_t ne11) { // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. // Therefore, in such cases use cuBLAS. const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; if (bad_padding_clear || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); - return; + return GGML_CUDA_MM_CUBLAS_UNSUPPORTED; } - - const int cc = ggml_cuda_info().devices[ctx.device].cc; - const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; - if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) { // The custom F16 vector kernel can be used over batched cuBLAS GEMM. // But this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) - ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMVF; } // A transposed vector can still use MMVQ (i.e. ne01 == 1) - if (ne01 == 1 && ne11 > MMVF_MAX_BATCH_SIZE && ne2 == 1 && ne3 == 1 + if (src0->ne[1] == 1 && ne11 > MMVF_MAX_BATCH_SIZE && dst->ne[2] == 1 && dst->ne[3] == 1 && src0->type == GGML_TYPE_F32 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && ggml_cuda_should_use_mmvf(src1->type, cc, src1->ne, src1->nb, /*ne11 =*/ 1)) { - ggml_tensor dst_vec = *dst; - dst_vec.ne[0] = ne11; - dst_vec.ne[1] = 1; - dst_vec.nb[1] = dst_vec.nb[0]*ne11; - dst_vec.nb[2] = dst_vec.nb[1]; - dst_vec.nb[3] = dst_vec.nb[1]; - ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); - return; + return GGML_CUDA_MM_MMVF_TRANSPOSED; } if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { - ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMF; } if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMVQ; } if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) { - ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return GGML_CUDA_MM_MMQ; + } + return GGML_CUDA_MM_CUBLAS; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); + +// Recompute dst one column at a time so that each column sees the batch-of-one configuration. +// Returns false when the batched launch already gives every column that same value. +static bool ggml_cuda_mul_mat_split_columns( + ggml_backend_cuda_context & ctx, int cc, int warp_size, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + const int64_t ncols_dst = dst->ne[1]; + if (ncols_dst <= 1 || src1->ne[1] != ncols_dst) { + return false; + } + // Only the token dimension is split, batched matmuls (attention) keep their shape. + if (src1->ne[2] != 1 || src1->ne[3] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) { + return false; + } + + if (ggml_cuda_batch_invariant() >= 2) { + const ggml_cuda_mm_path path_one = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1); + const ggml_cuda_mm_path path_batched = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ncols_dst); + if (path_one == path_batched) { + // Same implementation, but it still has to sum each destination element in the same order. + if (path_batched == GGML_CUDA_MM_MMVF) { + return false; // the block size follows K alone + } + if (path_batched == GGML_CUDA_MM_MMVQ && + ggml_cuda_mmvq_matches_single_column(src0->type, cc, ncols_dst)) { + return false; + } + } + } + + for (int64_t i = 0; i < ncols_dst; ++i) { + ggml_tensor src1_col = *src1; + ggml_tensor dst_col = *dst; + + src1_col.ne[1] = 1; + src1_col.nb[2] = src1_col.nb[1]; + src1_col.nb[3] = src1_col.nb[1]; + src1_col.data = (char *) src1->data + i*src1->nb[1]; + + dst_col.ne[1] = 1; + dst_col.nb[2] = dst_col.nb[1]; + dst_col.nb[3] = dst_col.nb[1]; + dst_col.data = (char *) dst->data + i*dst->nb[1]; + + ggml_cuda_mul_mat(ctx, src0, &src1_col, &dst_col); + } + return true; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + GGML_TENSOR_BINARY_OP_LOCALS + + const int32_t hint = ggml_get_op_params_i32(dst, 1); + if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { + return; + } + + const int cc = ggml_cuda_info().devices[ctx.device].cc; + const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; + + if (ggml_cuda_batch_invariant() && ggml_cuda_mul_mat_split_columns(ctx, cc, warp_size, src0, src1, dst)) { return; } - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + + switch (ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ne11)) { + case GGML_CUDA_MM_CUBLAS_UNSUPPORTED: + case GGML_CUDA_MM_CUBLAS: + ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + return; + case GGML_CUDA_MM_MMVF: + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MM_MMVF_TRANSPOSED: { + ggml_tensor dst_vec = *dst; + dst_vec.ne[0] = ne11; + dst_vec.ne[1] = 1; + dst_vec.nb[1] = dst_vec.nb[0]*ne11; + dst_vec.nb[2] = dst_vec.nb[1]; + dst_vec.nb[3] = dst_vec.nb[1]; + ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); + return; + } + case GGML_CUDA_MM_MMF: + ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MM_MMVQ: + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MM_MMQ: + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return; + } + GGML_ABORT("fatal error"); } // returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 97053480980..b14ef9681c5 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -541,6 +541,21 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } +// [TAG_BATCH_INVARIANT] +bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst) { + if (ncols_dst < 1 || ncols_dst > MMVQ_MAX_BATCH_SIZE) { + return false; + } + const mmvq_parameter_table_id table_id = get_device_table_id(cc); + if (table_id == MMVQ_PARAMETERS_GB10) { + // There nwarps also depends on the K loop trip count, which the caller does not pass in. + return ncols_dst == 1; + } + // blocks_per_iter, which is what assigns K blocks to threads, is proportional to nwarps. + // rows_per_cuda_block only changes which rows a block owns, not the order within a row. + return calc_nwarps(type, 1, table_id) == calc_nwarps(type, (int) ncols_dst, table_id); +} + template __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 5605bf7a4e6..61a88b851ec 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -4,6 +4,11 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11); +// [TAG_BATCH_INVARIANT] +// True when an MMVQ launch of ncols_dst columns sums each destination element in the same +// order as a launch of a single column, i.e. when the column count leaves nwarps unchanged. +bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst); + // Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID, // based on the quantization type and GPU architecture (compute capability). int get_mmvq_mmid_max_batch(ggml_type type, int cc); From b5c10293842a3e51dafce13b95829ea53af7146f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 20:52:21 +0000 Subject: [PATCH 03/81] cuda: bound the batch-invariant split with GGML_CUDA_BATCH_INVARIANT_MAX_COLS Splitting a prompt-sized batch costs far more than splitting a decode-sized one: a 273 token prefill becomes 273 single column matmuls and 273 single row attention launches per layer, which took prompt processing from 2731 to 225 tok/s on a B200 while four-chat decode only lost 7 percent. GGML_CUDA_BATCH_INVARIANT_MAX_COLS caps the width the split applies to, 0 keeps the previous unbounded behaviour. At 8 it covers every decode batch the server can form and leaves prefill alone, which restores prompt processing to 2696 tok/s and four-chat wall throughput to 127.2 against 136.4 unpatched. The bound gives up invariance for the prompt phase, so it is opt-in rather than the default. --- ggml/src/ggml-cuda/common.cuh | 3 +++ ggml/src/ggml-cuda/fattn.cu | 4 +++- ggml/src/ggml-cuda/ggml-cuda.cu | 12 ++++++++++++ 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 51ad1d6aa4d..728aa08dcb5 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -51,6 +51,9 @@ #define GGML_CUDA_CC_DP4A 610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products // [TAG_BATCH_INVARIANT] 0 = off, 1 = split every batched matmul, 2 = split only where it changes bits int ggml_cuda_batch_invariant(); +// Widest batch the split is applied to, 0 = no bound. Prompt-sized batches cost far more to +// split than decode-sized ones, and only prompt-phase invariance is given up by bounding it. +int ggml_cuda_batch_invariant_max_cols(); #define GGML_CUDA_CC_VOLTA 700 #define GGML_CUDA_CC_TURING 750 diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index b2f6f0660ee..6d20a756016 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -578,7 +578,9 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst ggml_cuda_set_device(ctx.device); // [TAG_BATCH_INVARIANT] Attend one query row at a time, as a batch of one would. - if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1) { + const int fattn_max_cols = ggml_cuda_batch_invariant_max_cols(); + if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1 && + (fattn_max_cols <= 0 || dst->src[0]->ne[1] <= fattn_max_cols)) { const ggml_tensor * Q = dst->src[0]; const ggml_tensor * mask = dst->src[3]; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index f9aa99ad003..a007bad9f44 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1843,6 +1843,14 @@ int ggml_cuda_batch_invariant() { return mode; } +int ggml_cuda_batch_invariant_max_cols() { + static const int max_cols = []() { + const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); + return val ? atoi(val) : 0; + }(); + return max_cols; +} + enum ggml_cuda_mm_path { GGML_CUDA_MM_CUBLAS_UNSUPPORTED, GGML_CUDA_MM_MMVF, @@ -1903,6 +1911,10 @@ static bool ggml_cuda_mul_mat_split_columns( if (src1->ne[2] != 1 || src1->ne[3] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) { return false; } + const int max_cols = ggml_cuda_batch_invariant_max_cols(); + if (max_cols > 0 && ncols_dst > max_cols) { + return false; + } if (ggml_cuda_batch_invariant() >= 2) { const ggml_cuda_mm_path path_one = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1); From f3ce9725e6ef793aa6b4a4d0fda89ff1896c2678 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 21:39:00 +0000 Subject: [PATCH 04/81] cuda: add exact concurrency with canonical paged attention --- ggml/src/ggml-cuda/fattn-common.cuh | 6 +- ggml/src/ggml-cuda/fattn-vec.cuh | 18 +- ggml/src/ggml-cuda/fattn.cu | 15 ++ ggml/src/ggml-cuda/ggml-cuda.cu | 30 +++ scripts/batchinv/README.md | 52 ++++ scripts/batchinv/bench.py | 53 ++++ scripts/batchinv/divergence.py | 164 +++++++++++++ scripts/batchinv/probe.cpp | 368 ++++++++++++++++++++++++++++ scripts/batchinv/prompts.py | 53 ++++ src/llama-graph.cpp | 10 +- src/llama-graph.h | 4 +- src/llama-kv-cache.cpp | 97 ++++++++ src/llama-kv-cache.h | 7 + tests/test-backend-ops.cpp | 46 ++++ 14 files changed, 914 insertions(+), 9 deletions(-) create mode 100644 scripts/batchinv/README.md create mode 100644 scripts/batchinv/bench.py create mode 100644 scripts/batchinv/divergence.py create mode 100644 scripts/batchinv/probe.cpp create mode 100644 scripts/batchinv/prompts.py diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index 123e5071911..f6aa5b03ec1 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1094,7 +1094,7 @@ void launch_fattn( // [TAG_BATCH_INVARIANT] Without this scan the KV loop runs to K->ne[1], which grows with the // other sequences sharing the cache. Scanning the mask bounds it by the sequence's own extent. const bool batch_invariant_KV_max = ggml_cuda_batch_invariant() != 0; - if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) { + if (!dst->src[5] && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) { const int64_t s31 = mask->nb[1] / sizeof(half2); const int64_t s33 = mask->nb[3] / sizeof(half2); @@ -1151,7 +1151,7 @@ void launch_fattn( if (ntiles_dst % blocks_num.x != 0) { // Fixup is only needed if the SMs work on fractional tiles. dst_tmp_meta.alloc((size_t(blocks_num.x) * ncols * (2 + DV/2))); } - } else if (ggml_cuda_batch_invariant()) { + } else if (dst->src[5] || ggml_cuda_batch_invariant()) { // [TAG_BATCH_INVARIANT] How the KV cache is split between blocks, and therefore the order // in which the partial attention results are combined, follows K->ne[1]. That length grows // with the other sequences sharing the cache, so pin the split to a single block per tile. @@ -1226,7 +1226,7 @@ void launch_fattn( V_data, mask ? ((const char *) mask->data) : nullptr, sinks ? ((const char *) sinks->data) : nullptr, - KV_max.ptr, + dst->src[5] ? (const int *) dst->src[5]->data : KV_max.ptr, !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr, scale, max_bias, m0, m1, n_head_log2, logit_softcap, Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3], diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 69dd9368624..f402795942c 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -16,7 +16,7 @@ static constexpr __device__ int ggml_cuda_fattn_vec_get_nthreads_device() { #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wpass-failed" #endif // __clang__ -template // D == head size +template // D == head size __launch_bounds__(ggml_cuda_fattn_vec_get_nthreads_device(), 1) static __global__ void flash_attn_ext_vec( const char * Q_ptr, @@ -247,13 +247,25 @@ static __global__ void flash_attn_ext_vec( #endif // V_DOT2_F32_F16_AVAILABLE } - const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11; + // In the paged specialization KV_max carries [count, physical page IDs...] per query. + // The loop and each warp's recurrence follow logical positions, never physical addresses. + static_assert(!paged || ncols == 1, "paged attention has one query per block"); + const int * pages = paged ? KV_max + (sequence*int(ne01.z) + ic0)*(1 + ne11/FATTN_KQ_STRIDE) : nullptr; + const int k_VKQ_max = paged ? pages[0]*FATTN_KQ_STRIDE : (KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11); + const char * K_base = K; + const char * V_base = V; + const half * mask_base = maskh; K += blockIdx.y*nthreads * nb11; V += blockIdx.y*nthreads * nb21; maskh += blockIdx.y*nthreads; for (int k_VKQ_0 = blockIdx.y*nthreads; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nthreads, - // Increment pointers after each loop: K += gridDim.y*nthreads*nb11, V += gridDim.y*nthreads*nb21, maskh += gridDim.y*nthreads) { + if constexpr (paged) { + const int physical = pages[1 + k_VKQ_0/FATTN_KQ_STRIDE]*FATTN_KQ_STRIDE + k_VKQ_0%FATTN_KQ_STRIDE; + K = K_base + int64_t(physical)*nb11; + V = V_base + int64_t(physical)*nb21; + maskh = mask_base + physical; + } // Calculate KQ tile and keep track of new maximum KQ values: float KQ_reg[ncols]; // KQ in registers. diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 6d20a756016..eff18212272 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -577,6 +577,21 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); + if (dst->src[5]) { + GGML_ASSERT(dst->src[0]->ne[0] == 256 && dst->src[2]->ne[0] == 256); + GGML_ASSERT(dst->src[1]->type == GGML_TYPE_F16 && dst->src[2]->type == GGML_TYPE_F16); + GGML_ASSERT(dst->src[3] && dst->src[0]->ne[3] == 1); + GGML_ASSERT(dst->src[5]->type == GGML_TYPE_I32 && ggml_is_contiguous(dst->src[5])); + GGML_ASSERT(dst->src[5]->ne[0] == 1 + dst->src[1]->ne[1]/FATTN_KQ_STRIDE); + GGML_ASSERT(dst->src[5]->ne[1] == dst->src[0]->ne[1]); + float softcap; + memcpy(&softcap, (const float *) dst->op_params + 2, sizeof(softcap)); + GGML_ASSERT(softcap == 0.0f); + fattn_kernel_t kernel = flash_attn_ext_vec<256, 1, GGML_TYPE_F16, GGML_TYPE_F16, false, true>; + launch_fattn<256, 1, 1>(ctx, dst, kernel, 4, 0, 128, false, false, false); + return; + } + // [TAG_BATCH_INVARIANT] Attend one query row at a time, as a batch of one would. const int fattn_max_cols = ggml_cuda_batch_invariant_max_cols(); if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1 && diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a007bad9f44..251758c60ee 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1835,8 +1835,17 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { // 1 - compute every destination column on its own, exactly as a batch of one would. // 2 - split off only the columns whose batch-of-one configuration differs from the // batched one, leaving the already invariant matmuls batched. +static bool ggml_cuda_exact_concurrency() { + static const bool exact = []() { + const char * value = getenv("LLAMA_EXACT_CONCURRENCY"); + return value && atoi(value) != 0; + }(); + return exact; +} + int ggml_cuda_batch_invariant() { static const int mode = []() { + if (ggml_cuda_exact_concurrency()) { return 2; } const char * val = getenv("GGML_CUDA_BATCH_INVARIANT"); return val ? atoi(val) : 0; }(); @@ -1845,6 +1854,7 @@ int ggml_cuda_batch_invariant() { int ggml_cuda_batch_invariant_max_cols() { static const int max_cols = []() { + if (ggml_cuda_exact_concurrency()) { return 0; } const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); return val ? atoi(val) : 0; }(); @@ -1903,6 +1913,26 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor static bool ggml_cuda_mul_mat_split_columns( ggml_backend_cuda_context & ctx, int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + // Recurrent-model output projections broadcast one weight matrix over sequence + // planes. These are token projections too, even though ne[2] or ne[3] is > 1. + // Normalize each plane before applying the existing selective column policy. + if (ggml_cuda_exact_concurrency() && src0->ne[2] == 1 && src0->ne[3] == 1 && + (dst->ne[2] > 1 || dst->ne[3] > 1) && + src1->ne[2] == dst->ne[2] && src1->ne[3] == dst->ne[3]) { + for (int64_t i3 = 0; i3 < dst->ne[3]; ++i3) { + for (int64_t i2 = 0; i2 < dst->ne[2]; ++i2) { + ggml_tensor src_plane = *src1; + ggml_tensor dst_plane = *dst; + src_plane.ne[2] = src_plane.ne[3] = 1; + dst_plane.ne[2] = dst_plane.ne[3] = 1; + src_plane.data = (char *) src1->data + i2*src1->nb[2] + i3*src1->nb[3]; + dst_plane.data = (char *) dst->data + i2*dst->nb[2] + i3*dst->nb[3]; + ggml_cuda_mul_mat(ctx, src0, &src_plane, &dst_plane); + } + } + return true; + } + const int64_t ncols_dst = dst->ne[1]; if (ncols_dst <= 1 || src1->ne[1] != ncols_dst) { return false; diff --git a/scripts/batchinv/README.md b/scripts/batchinv/README.md new file mode 100644 index 00000000000..62a51487047 --- /dev/null +++ b/scripts/batchinv/README.md @@ -0,0 +1,52 @@ +# Exact concurrency experiment p + +Opt in before loading the model with `LLAMA_EXACT_CONCURRENCY=1`. This also forces +`GGML_CUDA_BATCH_INVARIANT=2` with no column limit, including during prefill. + +The experimental policy supports unified, offloaded F16 K/V, causal flash attention, +256-dimensional K and V heads, no attention soft cap, and no sliding window. +Shared-weight matmuls over multiple sequence planes are normalized to one plane +before the inherited selective column dispatcher. Without this, the recurrent +output projection bypasses batch invariance during concurrent prefill. +It is measured on text prompts with Qwen3.5-4B on one B200. Context shifting, +position division, cross-sequence prefix copies, shared-prefix input tokens, and +whole-context state loading are unsupported. Per-sequence state save and restore +is supported. Unsupported cache transformations assert instead of silently +violating the page invariant. + +The allocator owns pages of 256 cells on behalf of one (sequence, position/256). +Position modulo 256 fixes the cell offset. Empty pages remain in the unified pool +and can be allocated by any sequence. The metadata is derived from live cells so +allocation rollback, tail removal, and sequence removal do not need another +transaction log. Restoring a sequence allocates free pages from this same pool. +The cost is up to 255 reserved cells per active sequence tail, plus holes introduced +by partial range removal. + +Attention receives an I32 page table in source 5: `[count, physical page IDs...]` +for each query, sorted by logical position. The physical K/V view and mask span +the pool, but the attention loop only visits the query's logical pages. The +final page is padded to 256 cells using the existing causal mask. Wholly future +pages in a prefill ubatch are excluded from the query's table. + +The vector attention specialization runs one query per block, four warps, and +`parallel_blocks=1`. It reads K/V directly from the physical pages, with two +128-cell softmax iterations per page, in logical page order. There is no K/V +gather and no split-K combine. The default vector specialization has no page +lookup. The ordinary path allocates no page metadata and launches no extra kernels. + +`FATTN_KQ_STRIDE=256` is a mask-scan stride, not the actual MMA rescaling tile. +For K/V head size 256, the Ampere-or-newer MMA configurations use 64 KV rows at +8 query/head columns and 32 KV rows at 16/32/64 columns. The retained vector path +has a 128-cell iteration. Both divide the 256-cell placement page. + +The probe is adapted from the existing batch-invariant harness and rejects +nonfinite logits and attention. `PROBE_B_REVERSE=1` fills neighbours before P0; +`PROBE_RESTORE=1` parks P0, releases a neighbour, restores P0, then rebuilds the +neighbour. Compute rows can be compared across this relocation; physical cache +views and index tensors must not be mistaken for sequence-0 compute outputs. + +`divergence.py --reference FILE` compares against an existing unparked solo token +reference. `bench.py --modes 0,1 --pairs 3` measures default off against exact mode +on, with 256 predicted tokens. Set `UNSLOTH_WORKSPACE` to the model parent workspace +and `LD_LIBRARY_PATH` to this build's bin directory. The harness uses GPU 3; +select a port in 9601-9610 explicitly. diff --git a/scripts/batchinv/bench.py b/scripts/batchinv/bench.py new file mode 100644 index 00000000000..64d8b32244b --- /dev/null +++ b/scripts/batchinv/bench.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""Cost of the knob: solo tok/s and four-chat aggregate tok/s, knob off and on, back to back.""" +import argparse, json, os, sys, time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from divergence import Server, completion, run_concurrent +from prompts import PROMPTS + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--binary", required=True) + ap.add_argument("--spec", default="none") + ap.add_argument("--n-predict", type=int, default=256) + ap.add_argument("--pairs", type=int, default=3) + ap.add_argument("--port", type=int, default=9602) + ap.add_argument("--modes", default="0,1") + ap.add_argument("--out", required=True) + a = ap.parse_args() + + modes = a.modes.split(",") + rows = [] + for pair in range(a.pairs): + for mode in modes: + env = {"LLAMA_EXACT_CONCURRENCY": mode, "GGML_CUDA_BATCH_INVARIANT": "0" if mode == "0" else "2"} + with Server(a.port, a.binary, [], env, a.out + ".server.log", a.spec) as s: + completion(a.port, PROMPTS["P0"], 32) # warm + solo = completion(a.port, PROMPTS["P0"], a.n_predict) + outs, wall = run_concurrent(a.port, ["P0", "P1", "P2", "P3"], a.n_predict) + row = { + "pair": pair, "mode": mode, "spec": a.spec, + "solo_tok_per_s": solo["timings"]["predicted_per_second"], + "solo_prompt_tok_per_s": solo["timings"]["prompt_per_second"], + "four_aggregate_tok_per_s": sum(o["timings"]["predicted_per_second"] for o in outs.values()), + "four_wall_s": wall, + "four_total_tokens": sum(len(o["tokens"]) for o in outs.values()), + } + row["four_wall_tok_per_s"] = row["four_total_tokens"] / wall + rows.append(row) + print(json.dumps(row), flush=True) + with open(a.out, "w") as f: + json.dump(rows, f, indent=2) + + print("\n=== summary ===", flush=True) + for mode in modes: + rs = [r for r in rows if r["mode"] == mode] + for k in ("solo_tok_per_s", "four_aggregate_tok_per_s", "four_wall_tok_per_s", "solo_prompt_tok_per_s"): + vals = sorted(r[k] for r in rs) + print(f"mode={mode} {k}: median {vals[len(vals)//2]:.1f} values {[round(v,1) for v in vals]}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py new file mode 100644 index 00000000000..e451afca0e1 --- /dev/null +++ b/scripts/batchinv/divergence.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Baseline / patched divergence harness: solo P0 vs P0 sharing batches with P1..P3.""" +import argparse, json, os, signal, subprocess, sys, threading, time, urllib.request, urllib.error + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from prompts import PROMPTS + +WS = os.environ["UNSLOTH_WORKSPACE"] +MODEL = f"{WS}/models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf" + + +def post(port, path, payload, timeout=1800): + req = urllib.request.Request(f"http://127.0.0.1:{port}{path}", + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.loads(r.read().decode()) + + +def get(port, path, timeout=10): + with urllib.request.urlopen(f"http://127.0.0.1:{port}{path}", timeout=timeout) as r: + return json.loads(r.read().decode()) + + +def completion(port, prompt, n_predict): + return post(port, "/completion", { + "prompt": prompt, "n_predict": n_predict, "temperature": 0.0, "top_k": 1, + "top_p": 1.0, "min_p": 0.0, "typical_p": 1.0, "seed": 0, + "repeat_penalty": 1.0, "presence_penalty": 0.0, "frequency_penalty": 0.0, + "cache_prompt": False, "return_tokens": True, "samplers": ["top_k", "temperature"], + }) + + +class Server: + def __init__(self, port, binary, extra, env_extra, log_path, spec, kv_unified=True): + self.port, self.log_path = port, log_path + self.args = [binary, "-m", MODEL, "--port", str(port), "--host", "127.0.0.1", + "--parallel", "4", "-c", "8192", + "--flash-attn", "on", "--metrics", "-ngl", "99", "--no-warmup", + "--seed", "0", "--spec-type", spec] + if kv_unified: + self.args += ["--kv-unified"] + if spec == "draft-mtp": + self.args += ["--spec-draft-n-max", "2"] + self.args += extra + self.env = dict(os.environ) + self.env["CUDA_VISIBLE_DEVICES"] = "3" + self.env.update(env_extra) + + def __enter__(self): + self.fh = open(self.log_path, "ab") + self.fh.write(("\n=== " + " ".join(self.args) + "\n=== env " + + json.dumps({k: v for k, v in self.env.items() + if k.startswith("GGML") or k == "CUDA_VISIBLE_DEVICES"}) + "\n").encode()) + self.fh.flush() + self.p = subprocess.Popen(self.args, stdout=self.fh, stderr=subprocess.STDOUT, + env=self.env, start_new_session=True) + print(f"[server] pid={self.p.pid} port={self.port} log={self.log_path}", flush=True) + deadline = time.time() + 600 + while time.time() < deadline: + if self.p.poll() is not None: + raise RuntimeError(f"server died rc={self.p.returncode}, see {self.log_path}") + try: + if get(self.port, "/health").get("status") == "ok": + print("[server] ready", flush=True) + return self + except Exception: + time.sleep(1.0) + raise RuntimeError("server did not become healthy") + + def __exit__(self, *a): + print(f"[server] stopping pid={self.p.pid}", flush=True) + try: + os.killpg(os.getpgid(self.p.pid), signal.SIGTERM) + self.p.wait(timeout=60) + except Exception: + try: + os.killpg(os.getpgid(self.p.pid), signal.SIGKILL) + except Exception: + pass + self.fh.close() + + +def run_concurrent(port, names, n_predict): + barrier = threading.Barrier(len(names)) + out = {} + + def work(name): + barrier.wait() + out[name] = completion(port, PROMPTS[name], n_predict) + + ts = [threading.Thread(target=work, args=(n,)) for n in names] + t0 = time.time() + for t in ts: + t.start() + for t in ts: + t.join() + return out, time.time() - t0 + + +def first_diff(a, b): + for i, (x, y) in enumerate(zip(a, b)): + if x != y: + return i + return None if len(a) == len(b) else min(len(a), len(b)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--reference") + ap.add_argument("--label", required=True) + ap.add_argument("--port", type=int, default=9601) + ap.add_argument("--binary", required=True) + ap.add_argument("--spec", default="none") + ap.add_argument("--n-predict", type=int, default=512) + ap.add_argument("--repeats", type=int, default=3) + ap.add_argument("--env", action="append", default=[]) + ap.add_argument("--extra", action="append", default=[]) + ap.add_argument("--out", required=True) + ap.add_argument("--no-kv-unified", action="store_true") + a = ap.parse_args() + + env_extra = dict(kv.split("=", 1) for kv in a.env) + res = {"label": a.label, "spec": a.spec, "n_predict": a.n_predict, + "env": env_extra, "extra": a.extra, "binary": a.binary, + "kv_unified": not a.no_kv_unified} + + with Server(a.port, a.binary, a.extra, env_extra, a.out + ".server.log", a.spec, + kv_unified=not a.no_kv_unified) as s: + solo = completion(a.port, PROMPTS["P0"], a.n_predict) + ref = json.load(open(a.reference))["tokens"] if a.reference else solo["tokens"] + res["solo_first_diff"] = first_diff(ref, solo["tokens"]) + res["reference"] = a.reference + res["solo"] = {"n_tokens": len(ref), "tok_per_s": solo["timings"]["predicted_per_second"], + "text_sha": None} + # solo repeat, to prove solo itself is stable + solo2 = completion(a.port, PROMPTS["P0"], a.n_predict) + res["solo_repeat_first_diff"] = first_diff(ref, solo2["tokens"]) + res["rounds"] = [] + for r in range(a.repeats): + outs, wall = run_concurrent(a.port, ["P0", "P1", "P2", "P3"], a.n_predict) + p0 = outs["P0"]["tokens"] + fd = first_diff(ref, p0) + agg = sum(outs[n]["timings"]["predicted_per_second"] for n in outs) + row = {"round": r, "first_diff": fd, "n_tokens": len(p0), + "identical": fd is None, "wall_s": wall, + "p0_tok_per_s": outs["P0"]["timings"]["predicted_per_second"], + "aggregate_tok_per_s": agg, + "per_req_n": {n: len(outs[n]["tokens"]) for n in outs}, "p0_tokens": p0} + res["rounds"].append(row) + print(f"[round {r}] first_diff={fd} identical={fd is None} wall={wall:.1f}s agg={agg:.1f} tok/s", flush=True) + with urllib.request.urlopen(f"http://127.0.0.1:{a.port}/metrics") as response: + res["metrics"] = response.read().decode() + with open(a.out + ".p0_solo.json", "w") as f: + json.dump({"tokens": ref, "content": solo["content"]}, f) + + with open(a.out, "w") as f: + json.dump(res, f, indent=2) + print(json.dumps({k: v for k, v in res.items() if k != "rounds"}, indent=2), flush=True) + print(json.dumps(res["rounds"], indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/batchinv/probe.cpp b/scripts/batchinv/probe.cpp new file mode 100644 index 00000000000..151463cd768 --- /dev/null +++ b/scripts/batchinv/probe.cpp @@ -0,0 +1,368 @@ +// Locate the first graph op whose sequence-0 output changes when the decode batch +// holds four sequences instead of one. Prompt KV for seq 0 is built identically in +// both phases, so the only difference is the width of the final decode ubatch. +#include "llama.h" +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct node_rec { + std::string name; + std::string op; + std::string tname; + int64_t ne[4]; + int64_t gdn_tokens = 0, gdn_seqs = 0; + size_t esize = 0; // bytes per element, 0 = not byte comparable + std::vector data; // empty when skipped + bool contiguous = false; + uint64_t hash = 0; +}; + +static bool g_record = false; +static std::vector * g_sink = nullptr; + +static uint64_t fnv1a(const uint8_t * p, size_t n) { + uint64_t h = 1469598103934665603ULL; + for (size_t i = 0; i < n; ++i) { h ^= p[i]; h *= 1099511628211ULL; } + return h; +} + +static bool eval_cb(struct ggml_tensor * t, bool ask, void * /*ud*/) { + if (!g_record) return false; + if (ask) return true; + + node_rec r; + r.name = ggml_get_name(t); + r.tname = ggml_type_name(t->type); + r.op = t->op == GGML_OP_NONE ? "LEAF" : ggml_op_name(t->op); + if (t->op == GGML_OP_UNARY) r.op = std::string("UNARY_") + ggml_unary_op_name(ggml_get_unary_op(t)); + if (t->op == GGML_OP_GLU) r.op = std::string("GLU_") + ggml_glu_op_name(ggml_get_glu_op(t)); + for (int i = 0; i < 4; ++i) r.ne[i] = t->ne[i]; + r.contiguous = ggml_is_contiguous(t); + if (r.name == "linear_attn_out-0") { + fprintf(stderr, "linear_attn_out-0: weight=%s input=[%lld,%lld,%lld,%lld]\n", + ggml_type_name(t->src[0]->type), (long long)t->src[1]->ne[0], + (long long)t->src[1]->ne[1], (long long)t->src[1]->ne[2], (long long)t->src[1]->ne[3]); + } + if (t->op == GGML_OP_GATED_DELTA_NET) { + r.gdn_tokens = t->src[2]->ne[2]; + r.gdn_seqs = t->src[2]->ne[3]; + } + + const size_t nbytes = ggml_nbytes(t); + if (r.contiguous && ggml_blck_size(t->type) == 1 && nbytes <= (256u << 20)) { + r.esize = ggml_type_size(t->type); + r.data.resize(nbytes); + ggml_backend_tensor_get(t, r.data.data(), 0, nbytes); + if (t->op == GGML_OP_FLASH_ATTN_EXT) { + for (size_t i = 0; i < nbytes/sizeof(float); ++i) { + float v; memcpy(&v, r.data.data() + i*sizeof(float), sizeof(float)); + if (!std::isfinite(v)) { fprintf(stderr, "nonfinite attention: %s\n", t->name); exit(5); } + } + } + r.hash = fnv1a(r.data.data(), nbytes); + if (nbytes > (64u << 20)) { r.data.clear(); } // keep the hash only for the big ones + } + g_sink->push_back(std::move(r)); + return true; +} + +static std::string slurp(const char * path) { + std::ifstream f(path); + std::stringstream ss; ss << f.rdbuf(); + std::string s = ss.str(); + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back(); + return s; +} + +static std::vector tokenize(const llama_vocab * v, const std::string & s) { + std::vector out(s.size() + 16); + int n = llama_tokenize(v, s.c_str(), (int) s.size(), out.data(), (int) out.size(), true, false); + if (n < 0) { out.resize(-n); n = llama_tokenize(v, s.c_str(), (int) s.size(), out.data(), (int) out.size(), true, false); } + out.resize(n); + return out; +} + +struct batch_holder { + std::vector tok; + std::vector pos; + std::vector nsid; + std::vector sid; + std::vector sidp; + std::vector out; + llama_batch get() { + sidp.resize(tok.size()); + for (size_t i = 0; i < tok.size(); ++i) sidp[i] = &sid[i]; + llama_batch b{}; + b.n_tokens = (int32_t) tok.size(); + b.token = tok.data(); b.pos = pos.data(); b.n_seq_id = nsid.data(); + b.seq_id = sidp.data(); b.logits = out.data(); + return b; + } +}; + +static llama_token greedy(llama_context * ctx, int32_t i, int n_vocab) { + const float * l = llama_get_logits_ith(ctx, i); + for (int k = 0; k < n_vocab; ++k) { + if (!std::isfinite(l[k])) { fprintf(stderr, "nonfinite logits at %d\n", k); exit(4); } + } + int best = 0; + for (int k = 1; k < n_vocab; ++k) if (l[k] > l[best]) best = k; + return best; +} + +// Feed a prompt as one decode call for one sequence, return the greedy next token. +static llama_token feed(llama_context * ctx, const std::vector & p, llama_seq_id seq, int n_vocab) { + batch_holder h; + for (size_t i = 0; i < p.size(); ++i) { + h.tok.push_back(p[i]); h.pos.push_back((llama_pos) i); + h.nsid.push_back(1); h.sid.push_back(seq); + h.out.push_back(i + 1 == p.size()); + } + llama_batch b = h.get(); + if (llama_decode(ctx, b) != 0) { fprintf(stderr, "decode failed\n"); exit(1); } + return greedy(ctx, (int32_t) p.size() - 1, n_vocab); +} + +int main(int argc, char ** argv) { + const bool prefill = getenv("PROBE_PREFILL") != nullptr; + const char * model_path = argv[1]; + const int n_seqs = argc > 2 ? atoi(argv[2]) : 4; // width of the probed decode batch + const char * out_path = argc > 3 ? argv[3] : nullptr; + std::vector prompts; + for (int i = 4; i < argc; ++i) prompts.push_back(slurp(argv[i])); + + llama_backend_init(); + llama_model_params mp = llama_model_default_params(); + mp.n_gpu_layers = 99; + llama_model * model = llama_model_load_from_file(model_path, mp); + if (!model) { fprintf(stderr, "model load failed\n"); return 1; } + const llama_vocab * vocab = llama_model_get_vocab(model); + const int n_vocab = llama_vocab_n_tokens(vocab); + + std::vector> ptok; + for (auto & s : prompts) ptok.push_back(tokenize(vocab, s)); + for (size_t i = 0; i < ptok.size(); ++i) fprintf(stderr, "prompt %zu: %zu tokens\n", i, ptok[i].size()); + + auto make_ctx = [&]() { + llama_context_params cp = llama_context_default_params(); + cp.n_ctx = 8192; cp.n_batch = 2048; cp.n_ubatch = 512; + if (prefill) { cp.n_ubatch = 2048; } + cp.n_seq_max = 4; cp.kv_unified = true; + cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + cp.cb_eval = eval_cb; cp.cb_eval_user_data = nullptr; + cp.no_perf = true; + return llama_init_from_model(model, cp); + }; + + std::vector rec_a, rec_b; + llama_token first_tok[4] = {0, 0, 0, 0}; + + // Phase A: decode ubatch width 1. PROBE_A_FILL controls how many sequences are + // already in the shared KV cache, which is what sets K->ne[1] for attention. + const int a_fill = getenv("PROBE_A_FILL") ? atoi(getenv("PROBE_A_FILL")) : 1; + // PROBE_A_PERM reorders which prompt goes into which sequence in phase A. With the same + // multiset of prompts the cache keeps its length but the masked cells hold different data. + int a_perm[4] = {0, 1, 2, 3}; + if (const char * perm = getenv("PROBE_A_PERM")) { + for (int k = 0; k < 4 && perm[2*k]; ++k) a_perm[k] = perm[2*k] - '0'; + } + { + llama_context * ctx = make_ctx(); + g_sink = &rec_a; g_record = prefill; + for (int s = 0; s < a_fill; ++s) { + const llama_token t = feed(ctx, ptok[a_perm[s]], s, n_vocab); + if (a_perm[s] == 0) first_tok[0] = t; + } + batch_holder h; + h.tok = {first_tok[0]}; h.pos = {(llama_pos) ptok[0].size()}; + h.nsid = {1}; h.sid = {0}; h.out = {1}; + llama_batch b = h.get(); + g_sink = &rec_a; g_record = !prefill; + if (llama_decode(ctx, b) != 0) { fprintf(stderr, "A decode failed\n"); return 1; } + g_record = false; + llama_free(ctx); + } + + // Phase B: same seq-0 prompt KV, then a decode ubatch holding n_seqs tokens. + { + llama_context * ctx = make_ctx(); + if (prefill) { + batch_holder h; + for (int seq = 0; seq < n_seqs; ++seq) { + for (size_t i = 0; i < ptok[0].size(); ++i) { + h.tok.push_back(ptok[seq][i%ptok[seq].size()]); h.pos.push_back(i); + h.nsid.push_back(1); h.sid.push_back(seq); h.out.push_back(i+1 == ptok[0].size()); + } + } + auto b = h.get(); + g_sink = &rec_b; g_record = true; + if (llama_decode(ctx, b) != 0) { return 6; } + g_record = false; + for (int seq = 0; seq < n_seqs; ++seq) { + first_tok[seq] = greedy(ctx, (seq+1)*ptok[0].size()-1, n_vocab); + } + } else for (int k = 0; k < n_seqs; ++k) { + const int s = getenv("PROBE_B_REVERSE") ? n_seqs - 1 - k : k; + first_tok[s] = feed(ctx, ptok[s], s, n_vocab); + } + if (getenv("PROBE_RESTORE")) { + std::vector state(llama_state_seq_get_size(ctx, 0)); + if (llama_state_seq_get_data(ctx, state.data(), state.size(), 0) != state.size()) { return 2; } + llama_memory_seq_rm(llama_get_memory(ctx), 0, -1, -1); + llama_memory_seq_rm(llama_get_memory(ctx), 1, -1, -1); + if (llama_state_seq_set_data(ctx, state.data(), state.size(), 0) != state.size()) { return 3; } + first_tok[1] = feed(ctx, ptok[1], 1, n_vocab); + } + if (first_tok[0] != 0 && rec_a.size()) {} + batch_holder h; + for (int s = 0; s < n_seqs; ++s) { + h.tok.push_back(first_tok[s]); h.pos.push_back((llama_pos) ptok[s].size()); + h.nsid.push_back(1); h.sid.push_back(s); h.out.push_back(1); + } + llama_batch b = h.get(); + g_sink = &rec_b; g_record = !prefill; + if (!prefill && llama_decode(ctx, b) != 0) { fprintf(stderr, "B decode failed\n"); return 1; } + g_record = false; + llama_free(ctx); + } + + // Optional: keep decoding and report the first step at which seq 0's token differs. + const int n_steps = getenv("PROBE_STEPS") ? atoi(getenv("PROBE_STEPS")) : 0; + int first_bad_step = -1; + if (n_steps > 0) { + std::vector tok_a, tok_b; + for (int phase = 0; phase < 2; ++phase) { + const int fill = phase == 0 ? a_fill : n_seqs; + const int width = phase == 0 ? 1 : n_seqs; + std::vector & out = phase == 0 ? tok_a : tok_b; + llama_context * ctx = make_ctx(); + std::vector next(4, 0); + std::vector pos(4, 0); + for (int s = 0; s < fill; ++s) { + const int p = phase == 0 ? a_perm[s] : s; + next[s] = feed(ctx, ptok[p], s, n_vocab); + pos[s] = (llama_pos) ptok[p].size(); + } + for (int step = 0; step < n_steps; ++step) { + batch_holder h; + for (int s = 0; s < width; ++s) { + h.tok.push_back(next[s]); h.pos.push_back(pos[s]); + h.nsid.push_back(1); h.sid.push_back(s); h.out.push_back(1); + } + llama_batch b = h.get(); + if (llama_decode(ctx, b) != 0) { fprintf(stderr, "step decode failed\n"); exit(1); } + out.push_back(next[0]); + for (int s = 0; s < width; ++s) { next[s] = greedy(ctx, s, n_vocab); pos[s] += 1; } + } + llama_free(ctx); + } + for (int i = 0; i < n_steps; ++i) { + if (tok_a[i] != tok_b[i]) { first_bad_step = i; break; } + } + fprintf(stderr, "steps: %d first differing step: %d\n", n_steps, first_bad_step); + } + + fprintf(stderr, "nodes: A=%zu B=%zu first tokens: %d %d %d %d\n", + rec_a.size(), rec_b.size(), first_tok[0], first_tok[1], first_tok[2], first_tok[3]); + + // Walk both node lists in order and compare seq 0's slice. + FILE * out = out_path ? fopen(out_path, "w") : stdout; + fprintf(out, "{\"n_seqs\":%d,\"first_bad_step\":%d,\"nodes_a\":%zu,\"nodes_b\":%zu,\"diffs\":[", n_seqs, first_bad_step, rec_a.size(), rec_b.size()); + size_t n = rec_a.size() < rec_b.size() ? rec_a.size() : rec_b.size(); + int emitted = 0; + for (size_t i = 0; i < n; ++i) { + const node_rec & A = rec_a[i]; + const node_rec & B = rec_b[i]; + const char * verdict = nullptr; + double max_abs = 0.0; + size_t ndiff = 0, ncmp = 0; + + if (A.name != B.name || A.op != B.op) { + verdict = "misaligned"; + } else if (A.op == "GATED_DELTA_NET" && A.gdn_tokens == B.gdn_tokens && + !A.data.empty() && !B.data.empty()) { + // Packed GDN outputs put ALL token outputs before ALL sequence states. + // Sequence 0's state therefore moves when the number of sequences changes. + const size_t output = A.ne[0]*A.gdn_tokens; + const size_t state = A.ne[0]*A.ne[1]/A.gdn_seqs - output; + for (size_t k = 0; k < output + state; ++k) { + const size_t ia = k < output ? k : A.gdn_seqs*output + k-output; + const size_t ib = k < output ? k : B.gdn_seqs*output + k-output; + float va, vb; + memcpy(&va, A.data.data()+ia*4, 4); memcpy(&vb, B.data.data()+ib*4, 4); + ++ncmp; + if (memcmp(&va, &vb, 4)) { + ++ndiff; + if (std::abs(double(va)-vb) > max_abs) { max_abs = std::abs(double(va)-vb); } + } + } + verdict = ndiff ? "row-differs" : nullptr; + } else if (A.esize == 0 || B.esize == 0 || A.esize != B.esize) { + verdict = "skipped"; + } else { + int tdim = -1; bool same = true; + for (int d = 0; d < 4; ++d) { + if (A.ne[d] == B.ne[d]) continue; + same = false; + if (B.ne[d] == n_seqs*A.ne[d] && tdim < 0) tdim = d; else { tdim = -2; break; } + } + if (tdim == -2) { + verdict = "shape-incomparable"; + } else if (same) { + verdict = (A.hash == B.hash) ? nullptr : "whole-tensor-differs"; + } else if (A.data.empty() || B.data.empty()) { + verdict = "too-large"; + } else { + // compare element (.., i_tdim = 0, ..) across all other indices + int64_t st[4] = {1, A.ne[0], A.ne[0]*A.ne[1], A.ne[0]*A.ne[1]*A.ne[2]}; + int64_t stb[4] = {1, B.ne[0], B.ne[0]*B.ne[1], B.ne[0]*B.ne[1]*B.ne[2]}; + for (int64_t i3 = 0; i3 < A.ne[3]; ++i3) + for (int64_t i2 = 0; i2 < A.ne[2]; ++i2) + for (int64_t i1 = 0; i1 < A.ne[1]; ++i1) + for (int64_t i0 = 0; i0 < A.ne[0]; ++i0) { + int64_t idx[4] = {i0, i1, i2, i3}; + + size_t oa = 0, ob = 0; + for (int d = 0; d < 4; ++d) { oa += idx[d]*st[d]; ob += idx[d]*stb[d]; } + ncmp++; + const uint8_t * pa = A.data.data() + oa*A.esize; + const uint8_t * pb = B.data.data() + ob*B.esize; + if (memcmp(pa, pb, A.esize) != 0) { + ndiff++; + if (A.esize == 4) { + float fa, fb; memcpy(&fa, pa, 4); memcpy(&fb, pb, 4); + double d2 = fa - fb; if (d2 < 0) d2 = -d2; + if (d2 > max_abs) max_abs = d2; + } + } + } + verdict = ndiff ? "row-differs" : nullptr; + } + } + { + if (emitted++) fprintf(out, ","); + fprintf(out, "\n{\"i\":%zu,\"name\":\"%s\",\"op\":\"%s\",\"ne_a\":[%lld,%lld,%lld,%lld]," + "\"ne_b\":[%lld,%lld,%lld,%lld],\"type\":\"%s\",\"verdict\":\"%s\",\"ndiff\":%zu,\"ncmp\":%zu,\"max_abs\":%.6g}", + i, A.name.c_str(), A.op.c_str(), + (long long)A.ne[0],(long long)A.ne[1],(long long)A.ne[2],(long long)A.ne[3], + (long long)B.ne[0],(long long)B.ne[1],(long long)B.ne[2],(long long)B.ne[3], + A.tname.c_str(), verdict ? verdict : "same", ndiff, ncmp, max_abs); + } + } + fprintf(out, "\n]}\n"); + if (out_path) fclose(out); + + llama_model_free(model); + llama_backend_free(); + return 0; +} diff --git a/scripts/batchinv/prompts.py b/scripts/batchinv/prompts.py new file mode 100644 index 00000000000..860b65fe08d --- /dev/null +++ b/scripts/batchinv/prompts.py @@ -0,0 +1,53 @@ +# Four distinct prompts, each about 300 tokens of raw text (no chat template). +_BODIES = { +"P0": """The history of numerical computing is a history of compromises between speed and exactness. +Early machines used fixed point arithmetic because it was cheap, and programmers carried scaling +factors in their heads. Floating point hardware moved the bookkeeping into silicon, but it did not +remove the compromise, it only hid it. Addition of floating point numbers is commutative but it is +not associative, so the order in which a long sum is accumulated changes the last few bits of the +result. On a single processor that order is fixed by the program text and nobody notices. On a +parallel processor the order is fixed by how the work was divided, and the division is chosen for +speed, not for reproducibility. A reduction split across two warps sums a different set of partial +products than the same reduction split across four warps, and the two answers differ in the low +bits. Nothing is wrong with either answer. Both are within a fraction of an ulp of the exact value. +The trouble begins when a downstream decision is discrete. A comparison, a rounding to an integer, +or the selection of the largest element of a vector turns a difference of one bit into a difference +of one branch, and from there the two computations walk away from each other and never come back. +Explain, carefully and at length, why this matters for a system that serves many users at once, +and what an engineer would have to give up to make the answer depend only on the request and not +on what else the machine happened to be doing at the time. Discuss the cost.""", +"P1": """Consider a public library that lends physical books and must decide how many copies of a +popular title to buy. The librarian has a fixed budget, a waiting list that grows and shrinks, and +a shelf that is already full. Every copy purchased shortens the queue for that title and lengthens +the queue for every other title, because the money and the shelf space are shared. The obvious +policy, buy copies of whatever has the longest queue, is unstable, because a title that briefly +becomes fashionable will absorb the whole budget and then sit unread for a decade. A better policy +has to weigh how long the demand is likely to last against how long the book will remain useful, +and it has to do this with almost no information. Describe in detail how you would design such a +policy, what data you would collect, how you would test it without harming readers, and how you +would know whether it was working. Consider what happens when the budget is cut in half without +warning, when a title is suddenly assigned as required reading by a local school, and when the +shelf itself must shrink because the building is being renovated. Explain the tradeoffs plainly.""", +"P2": """A small coastal town has one bridge to the mainland and it is failing. The engineers say it +has perhaps eight years left. Replacing it costs more than the town has ever spent on anything. +Repairing it buys maybe four years and costs a third as much, and the repair work closes the bridge +for two months in the summer, which is when the town earns most of its money. Doing nothing is +free until the day it is not. The town council is split, the ferry operator has opinions, and the +regional government will match funds only for a replacement, only if construction begins within +three years, and only if the town covers the first quarter of the cost itself. Write a long and +careful analysis of the options available to the council. Identify the assumptions that matter +most, the ones where being wrong changes the recommendation, and say how the council could cheaply +find out whether those assumptions hold. Then give a recommendation and state honestly what would +have to be true for the recommendation to be wrong. Do not hedge. Commit to an answer at the end.""", +"P3": """Describe the process by which a large body of water freezes over in winter, beginning with +the surface layer and working downward, and explain why the ice floats rather than sinking, why a +deep lake takes much longer to freeze than a shallow one of the same surface area, and why the +temperature at the bottom of a frozen lake settles near four degrees Celsius rather than at zero. +Then explain what this means for the animals that live there, how fish survive a winter under a +solid lid, why a heavy snowfall on top of the ice can be more dangerous to them than the cold +itself, and what happens in the spring when the whole column overturns. Use plain language and +avoid equations. Where a common explanation is wrong or incomplete, say so and give the better one. +Be thorough. Assume the reader is curious and patient but has no training in physics or biology, +and would rather understand one thing properly than be told five things quickly.""", +} +PROMPTS = {k: " ".join(v.split()) for k, v in _BODIES.items()} diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0e..71fce85bca7 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -468,6 +468,7 @@ void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) { } void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { + if (self_pages && self_pages->buffer) { mctx->set_input_pages(self_pages, ubatch); } mctx->set_input_k_idxs(self_k_idxs, ubatch); mctx->set_input_v_idxs(self_v_idxs, ubatch); @@ -1084,6 +1085,7 @@ void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) { } void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { + if (inp_attn->self_pages) { mctx->get_attn()->set_input_pages(inp_attn->self_pages, ubatch); } mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); @@ -2547,7 +2549,8 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, - int il) const { + int il, + ggml_tensor * pages) const { const bool v_trans = v->nb[1] > v->nb[2]; // split the batch into streams if needed @@ -2580,6 +2583,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); + cur->src[5] = pages; res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ggml_flash_attn_ext_add_sinks(cur, sinks); @@ -2769,6 +2773,8 @@ static std::unique_ptr build_attn_inp_kv_impl( inp->self_kq_mask_cnv = inp->self_kq_mask; } + inp->self_pages = mctx_cur->build_input_pages(ctx0, ubatch); + GGML_ASSERT(!inp->self_pages || (cparams.flash_attn && cparams.causal_attn)); inp->self_k_rot = mctx_cur->build_input_k_rot(ctx0); inp->self_v_rot = mctx_cur->build_input_v_rot(ctx0); @@ -2831,7 +2837,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il, inp->self_pages); cb(cur, "kqv_out", il); if (inp->self_v_rot) { diff --git a/src/llama-graph.h b/src/llama-graph.h index b388e028cb5..26f2169532d 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -319,6 +319,7 @@ class llm_graph_input_attn_no_cache : public llm_graph_input_i { class llm_graph_input_attn_kv : public llm_graph_input_i { public: + ggml_tensor * self_pages = nullptr; // I32 [1 + physical pages, n_tokens] llm_graph_input_attn_kv( const llama_hparams & hparams, const llama_cparams & cparams, @@ -1172,7 +1173,8 @@ struct llm_graph_context { ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] float kq_scale, - int il) const; + int il, + ggml_tensor * pages = nullptr) const; llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ec0f5a75314..df643a047a9 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -84,6 +84,14 @@ llama_kv_cache::llama_kv_cache( v_cells_impl(other ? other->v_cells_impl : std::make_shared()), v_cells(*v_cells_impl) { + const char * exact_env = getenv("LLAMA_EXACT_CONCURRENCY"); + exact_pages = exact_env && atoi(exact_env) != 0; + if (exact_pages) { + GGML_ASSERT(unified && offload && !v_trans && n_swa == 0); + GGML_ASSERT(type_k == GGML_TYPE_F16 && type_v == GGML_TYPE_F16); + GGML_ASSERT(kv_size % exact_page_size == 0); + } + // shared cells view the source cache's K/V tensors, so the cell count // follows the source allocation: a fitted target can be smaller than the // draft default and oversized views would overflow the source tensors @@ -447,6 +455,7 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { } void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + GGML_ASSERT(!exact_pages || seq_id_src == seq_id_dst); // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -566,6 +575,7 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) { } void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + GGML_ASSERT(!exact_pages || shift == 0); // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -616,6 +626,7 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll } void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + GGML_ASSERT(!exact_pages || d == 1); // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -961,6 +972,48 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, } } + if (exact_pages) { + // Reconstruct page ownership from live cells. Empty pages are immediately reusable; + // prepare() can roll back its speculative allocations without a second metadata log. + const auto & cells = v_cells[0]; + using page_key = std::pair; + std::map pages; + std::vector occupied(cells.size()/exact_page_size, false); + std::vector assigned(cells.size(), false); + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.is_empty(i)) { continue; } + GGML_ASSERT(cells.seq_count(i) == 1); + const auto pos = cells.pos_get(i); + GGML_ASSERT(pos >= 0 && uint32_t(pos)%exact_page_size == i%exact_page_size); + const page_key key {cells.seq_get(i), pos/exact_page_size}; + auto ins = pages.emplace(key, i/exact_page_size); + GGML_ASSERT(ins.first->second == i/exact_page_size); + occupied[i/exact_page_size] = true; + } + slot_info res {0, 0, {0}, {{}}}; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + GGML_ASSERT(ubatch.n_seq_id[i] == 1 && ubatch.pos[i] >= 0); + const page_key key {ubatch.seq_id[i][0], ubatch.pos[i]/exact_page_size}; + auto it = pages.find(key); + if (it == pages.end()) { + // Round-robin free-page search deliberately permits nonmonotonic physical order. + uint32_t page = v_heads[0]/exact_page_size; + uint32_t tested = 0; + while (tested < occupied.size() && occupied[page%occupied.size()]) { ++page; ++tested; } + if (tested == occupied.size()) { return {}; } + page %= occupied.size(); + occupied[page] = true; + it = pages.emplace(key, page).first; + } + const uint32_t idx = it->second*exact_page_size + ubatch.pos[i]%exact_page_size; + if (!cells.is_empty(idx) || assigned[idx]) { return {}; } + assigned[idx] = true; + res.idxs[0].push_back(idx); + } + if (cont && !res.is_contiguous()) { return {}; } + return res; + } + uint32_t n_tokens = ubatch.n_tokens; uint32_t n_seqs = 1; @@ -1232,7 +1285,50 @@ const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const { return v_cells[seq_to_stream[seq_id]]; } +ggml_tensor * llama_kv_cache::build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const { + if (!exact_pages) { return nullptr; } + auto * pages = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1 + get_size()/exact_page_size, ubatch.n_tokens); + ggml_set_input(pages); + ggml_set_name(pages, "attn_logical_pages"); + return pages; +} + +void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const { + GGML_ASSERT(exact_pages && dst->ne[1] == ubatch->n_tokens); + std::map> pages; + const auto & cells = v_cells[0]; + for (uint32_t i = 0; i < cells.size(); ++i) { + if (!cells.is_empty(i)) { + GGML_ASSERT(cells.seq_count(i) == 1); + pages[cells.seq_get(i)][cells.pos_get(i)/exact_page_size] = i/exact_page_size; + } + } + std::vector data(ggml_nelements(dst), -1); + for (uint32_t i = 0; i < ubatch->n_tokens; ++i) { + GGML_ASSERT(ubatch->n_seq_id[i] == 1); + auto * row = data.data() + i*dst->ne[0]; + row[0] = 0; + for (const auto & page : pages[ubatch->seq_id[i][0]]) { + // Exclude wholly future pages even when prefill includes later query rows. + if (page.first*exact_page_size > uint32_t(ubatch->pos[i])) { break; } + row[++row[0]] = page.second; + } + } + ggml_backend_tensor_set(dst, data.data(), 0, data.size()*sizeof(int32_t)); +} + +ggml_tensor * llama_kv_cache_context::build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const { + return kv->build_input_pages(ctx, ubatch); +} + +void llama_kv_cache_context::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const { + kv->set_input_pages(dst, ubatch); +} + uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { + // The physical view spans the pool. The page map, independently padded per query, + // is the only loop bound for exact attention; neighbours cannot extend that loop. + if (exact_pages) { return get_size(); } uint32_t result = 0; // pad the n_kv value so that the graph remains constant across batches and can be reused @@ -2283,6 +2379,7 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 GGML_ASSERT(cells.seq_has(idx, dest_seq_id)); } } else { + GGML_ASSERT(!exact_pages && "exact mode supports per-sequence restore only"); // whole KV cache restore if (cell_count > cells.size()) { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 6cb6dbd2f98..fa257422f02 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -171,6 +171,8 @@ class llama_kv_cache : public llama_memory_i { // uint32_t get_n_kv(const slot_info & sinfo) const; + ggml_tensor * build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const; + void set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const; // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; @@ -235,6 +237,9 @@ class llama_kv_cache : public llama_memory_i { std::vector v_stream; }; + static constexpr uint32_t exact_page_size = 256; + bool exact_pages = false; + bool v_trans = true; // the value tensor is transposed const uint32_t n_seq_max = 1; @@ -365,6 +370,8 @@ class llama_kv_cache_context : public llama_memory_context_i { // uint32_t get_n_kv() const; + ggml_tensor * build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const; + void set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const; ggml_type type_k() const; ggml_type type_v() const; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 53e93a1448d..a7511396791 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7193,6 +7193,41 @@ struct test_flash_attn_ext : public test_case { } }; +// Same mathematical attention as the CPU mask reference, but visit nonadjacent pages +// in a different order. Covers a partial tail and different page counts per query. +struct test_flash_attn_ext_pages : public test_flash_attn_ext { + test_flash_attn_ext_pages(int64_t batch) : + test_flash_attn_ext(256, 256, 2, {8, 1}, 1024, batch) {} + + std::string vars() override { return test_flash_attn_ext::vars() + ",exact_pages=1"; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + auto * out = test_flash_attn_ext::build_graph(ctx); + out->src[5] = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 5, nb); + ggml_set_name(out->src[5], "pages"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + test_flash_attn_ext::initialize_tensors(ctx); + auto * pages = ggml_get_tensor(ctx, "pages"); + auto * mask = ggml_get_tensor(ctx, "m"); + std::vector ids(5*nb, -1); + std::vector values(1024*nb, ggml_fp32_to_fp16(-INFINITY)); + for (int64_t q = 0; q < nb; ++q) { + ids[5*q] = q%2 ? 1 : 2; + ids[5*q + 1] = 2; + ids[5*q + 2] = 0; + for (int j = 0; j < 256; ++j) { values[1024*q + 512 + j] = ggml_fp32_to_fp16(0.0f); } + if (q%2 == 0) { + for (int j = 0; j < 17; ++j) { values[1024*q + j] = ggml_fp32_to_fp16(0.0f); } + } + } + ggml_backend_tensor_set(pages, ids.data(), 0, ids.size()*sizeof(int32_t)); + ggml_backend_tensor_set(mask, values.data(), 0, values.size()*sizeof(ggml_fp16_t)); + } +}; + // GGML_OP_CROSS_ENTROPY_LOSS struct test_cross_entropy_loss : public test_case { const ggml_type type; @@ -9170,6 +9205,14 @@ static std::vector> make_test_cases_eval() { } } + // Shared weights over sequence planes, as in a recurrent-model output projection. + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0}) { + for (int n : {1, 17, 307}) { + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {4, 1})); + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {1, 4})); + } + } + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_MXFP4, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); @@ -9937,6 +9980,9 @@ static std::vector> make_test_cases_eval() { } // mixed quant and Q1_0 test cases + for (int64_t batch : {1, 4, 12}) { + test_cases.emplace_back(new test_flash_attn_ext_pages(batch)); + } test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(72, 72, 4, {1, 1}, 96, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0)); From a94f76feaa9e7862ec2768e17d6a1b75a21c71e3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 23:20:47 +0000 Subject: [PATCH 05/81] batch: keep a prompt ubatch to one sequence under LLAMA_EXACT_CONCURRENCY The recurrent half of a hybrid model is not invariant to the shape of the ubatch. With the attention half made exact, a prompt processed in ubatches it shares with other sequences' prompt tokens still leaves a different gated delta net state than the same prompt processed alone: the first node to show it is the layer 1 recurrent state, 2.2e5 of 5.2e5 elements, max 7.3e-4, and over a 512 token generation it flips a token at step 79. split_equal grows an optional cap on the number of sequence sets per ubatch. The hybrid memory passes 1 when the mode is on and some sequence contributes more than one token to the batch, which is the prompt phase. A plain decode step, one token per sequence, is already exact under the gather and stays batched, so the cost falls on prompt processing only. --- src/llama-batch.cpp | 21 ++++++++++++++++++++- src/llama-batch.h | 8 +++++++- src/llama-impl.cpp | 11 +++++++++++ src/llama-impl.h | 6 ++++++ src/llama-memory-hybrid.cpp | 9 ++++++++- 5 files changed, 52 insertions(+), 3 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f48..50f70bb0ff3 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -507,7 +507,21 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) { return ubatch_add(idxs, idxs.size(), false); } -llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail) { +bool llama_batch_allocr::has_multi_token_seq() const { + std::vector n_per_seq(n_seq_max, 0); + + for (int32_t i = 0; i < batch.n_tokens; ++i) { + for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { + if (++n_per_seq[batch.seq_id[i][s]] > 1) { + return true; + } + } + } + + return false; +} + +llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t n_seqs_max) { if (sequential && has_cpl) { LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__); @@ -547,6 +561,11 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, if (cur_seq_set.size() > n_ubatch) { break; } + + // [TAG_EXACT_CONCURRENCY] + if (n_seqs_max > 0 && cur_seq_set.size() >= n_seqs_max) { + break; + } } } diff --git a/src/llama-batch.h b/src/llama-batch.h index a3d1889d4a0..d354c442d03 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -105,7 +105,13 @@ class llama_batch_allocr { // make ubatches of equal-length sequences sets // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch - llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail); + // n_seqs_max = maximum sequence sets per ubatch, 0 = no limit + // [TAG_EXACT_CONCURRENCY] passing 1 keeps a ubatch to a single sequence + llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t n_seqs_max = 0); + + // [TAG_EXACT_CONCURRENCY] true if some sequence contributes more than one token to the batch, + // i.e. this is not a plain one-token-per-sequence decode step + bool has_multi_token_seq() const; // sequence-set-wise split - each ubatch contains a single sequence-set llama_ubatch split_seq(uint32_t n_ubatch); diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index b3a94b946d2..bad0e55237a 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -169,3 +170,13 @@ std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i) { return gguf_data_to_str(type, gguf_get_val_data(ctx_gguf, i), 0); } } + +// [TAG_EXACT_CONCURRENCY] +bool llama_exact_concurrency() { + static const bool enabled = []() { + const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); + return val && atoi(val) != 0; + }(); + + return enabled; +} diff --git a/src/llama-impl.h b/src/llama-impl.h index 4988b06d2ca..9b64431fedd 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -103,3 +103,9 @@ std::string llama_format_tensor_shape(const std::vector & ne); std::string llama_format_tensor_shape(const struct ggml_tensor * t); std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i); + +// [TAG_EXACT_CONCURRENCY] +// opt-in mode under which a sequence's attention depends only on its own cells, in position order, +// so that its output does not change when other sequences share the KV cache. Off by default. +// Reads the same LLAMA_EXACT_CONCURRENCY variable as the paged KV cache and the CUDA backend. +bool llama_exact_concurrency(); diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6..ba54ab12923 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -86,7 +86,14 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // so that the rollback snapshots remain valid const uint32_t n_rs_seq = mem_recr->n_rs_seq; - ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0); + // [TAG_EXACT_CONCURRENCY] the recurrent half of a hybrid model is not invariant to + // the shape of the ubatch: a prompt processed next to other sequences' prompt tokens + // leaves a different gated delta net state than the same prompt processed alone. + // Keeping such a ubatch to a single sequence removes that. A plain decode step, one + // token per sequence, is already exact and stays batched. + const uint32_t n_seqs_max = llama_exact_concurrency() && balloc.has_multi_token_seq() ? 1 : 0; + + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, n_seqs_max); } if (ubatch.n_tokens == 0) { From 07d82f022d44d8339bd0270f1aedcf821c1b7e79 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 23:31:58 +0000 Subject: [PATCH 06/81] cuda: let exact mode bound the column policy when prompt ubatches are per sequence --- ggml/src/ggml-cuda/ggml-cuda.cu | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 251758c60ee..2b1754baea6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1854,9 +1854,13 @@ int ggml_cuda_batch_invariant() { int ggml_cuda_batch_invariant_max_cols() { static const int max_cols = []() { - if (ggml_cuda_exact_concurrency()) { return 0; } + // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill + // matmul shapes match its solo run, so exact mode no longer needs the column policy to be + // unbounded there. Honour an explicit bound when one is set; default to unbounded. const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); - return val ? atoi(val) : 0; + if (val) { return atoi(val); } + if (ggml_cuda_exact_concurrency()) { return 0; } + return 0; }(); return max_cols; } From 5c6d79e1a8587c168eafb3c8d40cdc87eceb7fad Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 23:34:40 +0000 Subject: [PATCH 07/81] ggml: add a non-blocking query for backend events ggml_backend_event_synchronize() is the only way to find out whether the work recorded before an event has finished, and it answers by waiting for it. A caller that issued an asynchronous copy so that it could get on with something else has no way to ask "is it done yet" without giving that up again. ggml_backend_event_query() is that question. It is optional, and it is the last field of ggml_backend_device_i so that a backend which does not implement it needs no change: a missing entry is NULL and the generic implementation falls back to a blocking synchronize and returns true, which is correct, just no better than what a caller could do already. CUDA implements it with cudaEventQuery, treating cudaErrorNotReady as the answer "not yet" rather than as a failure, and clearing it so it is not reported against the next call. The other sixteen device interfaces get an explicit NULL. Trailing initializers could have been left off, since these are positional aggregate initializers and the new member would be value-initialized, but -Wmissing-field-initializers is part of -Wextra and becomes an error under LLAMA_FATAL_WARNINGS. --- ggml/include/ggml-backend.h | 3 +++ ggml/src/ggml-backend-impl.h | 5 +++++ ggml/src/ggml-backend-meta.cpp | 1 + ggml/src/ggml-backend.cpp | 12 ++++++++++++ ggml/src/ggml-blas/ggml-blas.cpp | 1 + ggml/src/ggml-cann/ggml-cann.cpp | 1 + ggml/src/ggml-cpu/ggml-cpu.cpp | 1 + ggml/src/ggml-cuda/ggml-cuda.cu | 17 +++++++++++++++++ ggml/src/ggml-et/ggml-et.cpp | 1 + ggml/src/ggml-hexagon/ggml-hexagon.cpp | 1 + ggml/src/ggml-metal/ggml-metal.cpp | 1 + ggml/src/ggml-opencl/ggml-opencl.cpp | 1 + ggml/src/ggml-openvino/ggml-openvino.cpp | 1 + ggml/src/ggml-rpc/ggml-rpc.cpp | 1 + ggml/src/ggml-sycl/ggml-sycl.cpp | 1 + ggml/src/ggml-virtgpu/ggml-backend-device.cpp | 1 + ggml/src/ggml-vulkan/ggml-vulkan.cpp | 1 + ggml/src/ggml-webgpu/ggml-webgpu.cpp | 1 + ggml/src/ggml-zdnn/ggml-zdnn.cpp | 1 + ggml/src/ggml-zendnn/ggml-zendnn.cpp | 1 + 20 files changed, 53 insertions(+) diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36e3..30d8304d492 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -125,6 +125,9 @@ extern "C" { GGML_API void ggml_backend_event_free(ggml_backend_event_t event); GGML_API void ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backend); GGML_API void ggml_backend_event_synchronize(ggml_backend_event_t event); + // non-blocking: true once everything recorded before the event has completed. + // backends without a query implementation fall back to a blocking synchronize and return true. + GGML_API bool ggml_backend_event_query(ggml_backend_event_t event); GGML_API void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event); // diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 40cea024c3d..902b0963afa 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -200,6 +200,11 @@ extern "C" { ggml_backend_event_t (*event_new) (ggml_backend_dev_t dev); void (*event_free) (ggml_backend_dev_t dev, ggml_backend_event_t event); void (*event_synchronize) (ggml_backend_dev_t dev, ggml_backend_event_t event); + + // (optional) non-blocking completion test for an event. + // kept last so that backends that do not implement it need no change: a missing entry + // is NULL, and ggml_backend_event_query() then falls back to a blocking synchronize. + bool (*event_query) (ggml_backend_dev_t dev, ggml_backend_event_t event); }; struct ggml_backend_device { diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 3ec40fb1af7..d531ae4b5fa 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -193,6 +193,7 @@ static const ggml_backend_device_i ggml_backend_meta_device_iface = { /* .event_new = */ nullptr, /* .event_free = */ nullptr, /* .event_synchronize = */ nullptr, + /* .event_query = */ NULL, }; static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev) { diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index e519bdf50a1..a56ab30862a 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -551,6 +551,18 @@ void ggml_backend_event_synchronize(ggml_backend_event_t event) { event->device->iface.event_synchronize(event->device, event); } +bool ggml_backend_event_query(ggml_backend_event_t event) { + GGML_ASSERT(event); + + if (event->device->iface.event_query == NULL) { + // no way to ask: the honest answer is to wait for it and then say yes + ggml_backend_event_synchronize(event); + return true; + } + + return event->device->iface.event_query(event->device, event); +} + void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { GGML_ASSERT(backend); GGML_ASSERT(backend->iface.event_wait != NULL); diff --git a/ggml/src/ggml-blas/ggml-blas.cpp b/ggml/src/ggml-blas/ggml-blas.cpp index e4b5bd25474..7271b6b632b 100644 --- a/ggml/src/ggml-blas/ggml-blas.cpp +++ b/ggml/src/ggml-blas/ggml-blas.cpp @@ -469,6 +469,7 @@ static const struct ggml_backend_device_i ggml_backend_blas_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // backend reg interface diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index 5e5541aac94..f0bc9205bfa 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2948,6 +2948,7 @@ static const ggml_backend_device_i ggml_backend_cann_device_interface = { /* .event_new = */ ggml_backend_cann_device_event_new, /* .event_free = */ ggml_backend_cann_device_event_free, /* .event_synchronize = */ ggml_backend_cann_device_event_synchronize, + /* .event_query = */ NULL, }; // backend reg diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 8cece71f186..a11d707aa14 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -500,6 +500,7 @@ static const struct ggml_backend_device_i ggml_backend_cpu_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // CPU backend - backend (reg) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2456f7dcc62..91902099011 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5375,6 +5375,22 @@ static void ggml_backend_cuda_device_event_synchronize(ggml_backend_dev_t dev, g CUDA_CHECK(cudaEventSynchronize((cudaEvent_t)event->context)); } +static bool ggml_backend_cuda_device_event_query(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + + const cudaError_t err = cudaEventQuery((cudaEvent_t)event->context); + + if (err == cudaErrorNotReady) { + // not an error: clear it so it is not reported against the next call + (void) cudaGetLastError(); + return false; + } + + CUDA_CHECK(err); + + return true; +} + static const ggml_backend_device_i ggml_backend_cuda_device_interface = { /* .get_name = */ ggml_backend_cuda_device_get_name, /* .get_description = */ ggml_backend_cuda_device_get_description, @@ -5391,6 +5407,7 @@ static const ggml_backend_device_i ggml_backend_cuda_device_interface = { /* .event_new = */ ggml_backend_cuda_device_event_new, /* .event_free = */ ggml_backend_cuda_device_event_free, /* .event_synchronize = */ ggml_backend_cuda_device_event_synchronize, + /* .event_query = */ ggml_backend_cuda_device_event_query, }; // backend reg diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp index b87b189a57a..ace1cdedad1 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1684,6 +1684,7 @@ static const struct ggml_backend_device_i ggml_backend_et_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; /* diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index e8a5009b381..497ae043a4e 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -4274,6 +4274,7 @@ static const struct ggml_backend_device_i ggml_backend_hexagon_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; //** backend registry diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 9756d47050c..8962e2f7f00 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -818,6 +818,7 @@ static ggml_backend_device_i ggml_backend_metal_device_i = { /* .event_new = */ ggml_backend_metal_device_event_new, /* .event_free = */ ggml_backend_metal_device_event_free, /* .event_synchronize = */ ggml_backend_metal_device_event_synchronize, + /* .event_query = */ NULL, }; // backend registry diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 64f3325b2a5..fc8dbbac80a 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -11332,6 +11332,7 @@ struct ggml_backend_device_i ggml_backend_opencl_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; } diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index e299e16c778..995f1328442 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -1454,6 +1454,7 @@ static const struct ggml_backend_device_i ggml_backend_openvino_device_interface /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; struct ggml_backend_openvino_reg_context { diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 69a8a08ae17..4a814d3f534 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1945,6 +1945,7 @@ static const struct ggml_backend_device_i ggml_backend_rpc_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // backend reg interface diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 0573643d834..6c0cc073117 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6448,6 +6448,7 @@ static const ggml_backend_device_i ggml_backend_sycl_device_interface = { /* .event_new = */ ggml_backend_sycl_device_event_new, /* .event_free = */ ggml_backend_sycl_device_event_free, /* .event_synchronize = */ ggml_backend_sycl_device_event_synchronize, + /* .event_query = */ NULL, }; // backend reg diff --git a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp index 987ce9dd110..13a70c594df 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp @@ -157,4 +157,5 @@ const ggml_backend_device_i ggml_backend_remoting_device_interface = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c1d86aaac5c..b94f520185a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -18798,6 +18798,7 @@ static const struct ggml_backend_device_i ggml_backend_vk_device_i = { /* .event_new = */ ggml_backend_vk_device_event_new, /* .event_free = */ ggml_backend_vk_device_event_free, /* .event_synchronize = */ ggml_backend_vk_device_event_synchronize, + /* .event_query = */ NULL, }; static const char * ggml_backend_vk_reg_get_name(ggml_backend_reg_t reg) { diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 2434848a55a..4de48d7fa69 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4661,6 +4661,7 @@ static struct ggml_backend_device_i ggml_backend_webgpu_device_i = { /* .event_new = */ ggml_backend_webgpu_device_event_new, /* .event_free = */ ggml_backend_webgpu_device_event_free, /* .event_synchronize = */ ggml_backend_webgpu_device_event_synchronize, + /* .event_query = */ NULL, }; /* End GGML Backend Device Interface */ diff --git a/ggml/src/ggml-zdnn/ggml-zdnn.cpp b/ggml/src/ggml-zdnn/ggml-zdnn.cpp index 4007ac9dfc7..bbd74fb9d5a 100644 --- a/ggml/src/ggml-zdnn/ggml-zdnn.cpp +++ b/ggml/src/ggml-zdnn/ggml-zdnn.cpp @@ -547,6 +547,7 @@ static ggml_backend_device_i ggml_backend_zdnn_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // diff --git a/ggml/src/ggml-zendnn/ggml-zendnn.cpp b/ggml/src/ggml-zendnn/ggml-zendnn.cpp index ec7ce233145..89c6c36a0f1 100644 --- a/ggml/src/ggml-zendnn/ggml-zendnn.cpp +++ b/ggml/src/ggml-zendnn/ggml-zendnn.cpp @@ -781,6 +781,7 @@ static const struct ggml_backend_device_i ggml_backend_zendnn_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // backend reg interface From b0487839371a7e395596d40c9dbc9c2a1eea19d5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 23:34:40 +0000 Subject: [PATCH 08/81] llama: coalesce sequence state transfers, and issue them asynchronously Two changes to how a sequence's state is copied out of and back into the cache, the first of which the second one needs. Coalescing. The save side works out which cells belong to the sequence, merges them into ranges and emits one write per range per tensor. The restore side does not: it emits one read per cell, thousands of them, even when the cells it was given are a handful of long runs. Merging fragments that are adjacent in both the tensor and the buffer fixes both sides at once, and covers the transposed V layout where the same runs are emitted once per embedding row. Sequences sharing a unified cache take their cells in turn, so what is left after merging is a regular comb rather than one block; a comb is what a strided copy describes, so runs of one length at a constant stride become a single 2d transfer. Measured on a 4B at -c 8192 with four chats, a 1989-cell sequence goes from 1989 transfers per tensor to about 160, and a sequence that has the cache to itself to one. Asynchronous transfers. llama_state_seq_copy is a transfer that can be issued and left running: it owns the host buffer, a backend per device holding part of the cache so the copies get a stream of their own rather than queueing behind the graphs, and an event per device to say when its half is done. The buffer is pinned where the backend offers pinned memory, which is what makes the copies overlap at all, and grow-only, because page-locking a hundred MiB costs about as long as the copy it is for and a caller parking the same sequence repeatedly asks for a slightly different size each time. The restore side of the asynchronous path deliberately does not use the whole-tensor staging the synchronous one does. Staging reads a tensor, patches the sequence's bytes into the host copy and writes the tensor back, which keeps the neighbours only while nothing else is touching the cache. These copies exist so that decoding can carry on beside them, so the write-back would undo whatever the sequences sharing the tensor wrote to their own cells in the meantime. Writing only this sequence's runs cannot, and coalescing is what makes that affordable. llama_state_seq_copy_init() returns NULL when no backend can copy asynchronously, so a caller keeps the synchronous calls on those. --- include/llama.h | 58 ++++ src/llama-context.cpp | 644 +++++++++++++++++++++++++++++++++++++++++- src/llama-context.h | 9 + 3 files changed, 702 insertions(+), 9 deletions(-) diff --git a/include/llama.h b/include/llama.h index a04177f9f7d..2c61b3d2a71 100644 --- a/include/llama.h +++ b/include/llama.h @@ -927,6 +927,64 @@ extern "C" { llama_seq_id dest_seq_id, llama_state_seq_flags flags); + // [TAG_STATE_ASYNC] asynchronous per-sequence state transfer + // + // llama_state_seq_get_data_ext / set_data_ext do not return until every byte has moved, + // so a caller that copies a sequence out of the cache to make room stops doing anything + // else for as long as the copy takes. A transfer object issues the same copies on a + // stream of its own and hands back control immediately; the caller polls + // llama_state_seq_copy_done() and gets on with its other work in between. + // + // The transfer owns the host buffer it reads from or writes into. That buffer is pinned + // when the backend offers pinned memory, which is what makes the copy fast, and it + // cannot be freed while a copy is still using it. + // + // Between issuing and completion the caller must not touch the buffer, must not free or + // reuse the cells of a sequence being read, and must not decode a sequence being + // written. llama_state_seq_copy_free() waits for an outstanding copy first. + struct llama_state_seq_copy; + + // NULL if the context's backends cannot copy asynchronously; the caller then uses the + // synchronous llama_state_seq_*_data_ext calls + LLAMA_API struct llama_state_seq_copy * llama_state_seq_copy_init(struct llama_context * ctx); + LLAMA_API void llama_state_seq_copy_free(struct llama_state_seq_copy * cpy); + + // Size the transfer's host buffer, keeping no contents; NULL on failure. Grow-only: + // page-locking host memory is far too slow to do once per transfer, so the memory is + // kept between them and only given back by llama_state_seq_copy_buf_free(). + LLAMA_API uint8_t * llama_state_seq_copy_buf_resize (struct llama_state_seq_copy * cpy, size_t size); + LLAMA_API uint8_t * llama_state_seq_copy_buf (struct llama_state_seq_copy * cpy); + LLAMA_API size_t llama_state_seq_copy_buf_size (struct llama_state_seq_copy * cpy); + // host memory actually held, which is what a caller budgeting host RAM has to count + LLAMA_API size_t llama_state_seq_copy_buf_capacity(struct llama_state_seq_copy * cpy); + LLAMA_API void llama_state_seq_copy_buf_free (struct llama_state_seq_copy * cpy); + + // true when the buffer is page-locked, i.e. when the copies can really overlap + LLAMA_API bool llama_state_seq_copy_buf_is_pinned(struct llama_state_seq_copy * cpy); + + // issue the copies; return the number of bytes covered, 0 on failure + LLAMA_API size_t llama_state_seq_copy_get( + struct llama_state_seq_copy * cpy, + size_t size, + llama_seq_id seq_id, + llama_state_seq_flags flags); + + LLAMA_API size_t llama_state_seq_copy_set( + struct llama_state_seq_copy * cpy, + size_t size, + llama_seq_id dest_seq_id, + llama_state_seq_flags flags); + + // transfers the last issue posted: one per run of adjacent cells, per tensor + LLAMA_API size_t llama_state_seq_copy_n_copies(struct llama_state_seq_copy * cpy); + + // microseconds the last issue spent waiting for the compute streams before it could start + LLAMA_API int64_t llama_state_seq_copy_sync_us(struct llama_state_seq_copy * cpy); + + // non-blocking completion test, and the blocking wait behind it + LLAMA_API bool llama_state_seq_copy_done(struct llama_state_seq_copy * cpy); + LLAMA_API void llama_state_seq_copy_wait(struct llama_state_seq_copy * cpy); + // // Decoding // diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 66940d4fc61..0747d1c6366 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2558,16 +2558,145 @@ class llama_io_write_dummy : public llama_io_write_i { size_t size_written = 0; }; +// [TAG_STATE_COALESCE] one transfer per run of cells, not one per cell +// +// A sequence's state is emitted in cell order, so a run of cells that is contiguous in the +// cache is contiguous both in the tensor and in the host buffer, and the fragments covering +// it are one transfer. The save side already coalesces its cells into ranges before it emits +// them; the restore side does not, and asks for one transfer per cell even when the cells it +// was given are a handful of long runs. Merging here fixes both sides at once, and covers +// the transposed V layout, where the same runs are emitted once per embedding row. +template +static size_t llama_io_run_end(const std::vector & infos, size_t i) { + size_t end = i + 1; + + while (end < infos.size() && + infos[end].tensor == infos[end - 1].tensor && + infos[end].offset == infos[end - 1].offset + infos[end - 1].size && + infos[end].ptr == infos[end - 1].ptr + infos[end - 1].size) { + end++; + } + + return end; +} + +template +static size_t llama_io_run_size(const std::vector & infos, size_t i, size_t end) { + size_t size = 0; + + for (size_t j = i; j < end; ++j) { + size += infos[j].size; + } + + return size; +} + +// [TAG_STATE_COALESCE] runs of one length at a constant stride are a single strided copy +// +// Sequences sharing a unified cache take their cells in turn, so a sequence's cells are not +// one block but a regular comb: a few cells, a gap, a few cells, for as long as the sequence +// is. Merging adjacent cells still leaves hundreds of runs per tensor, and at a few +// microseconds to post each one that is tens of milliseconds spent issuing copies. A comb is +// exactly what a strided copy describes, so one call replaces a whole group of runs. +// +// emit(tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); n_copies == 1 means +// an ordinary contiguous transfer and the strides are not meaningful. +template +static void llama_io_emit(const std::vector & infos, size_t first, size_t last, emit_t emit) { + // the runs of adjacent cells, as index ranges into infos + std::vector> runs; + + for (size_t i = first; i < last; ) { + const size_t end = llama_io_run_end(infos, i); + + runs.emplace_back(i, end); + + i = end; + } + + for (size_t r = 0; r < runs.size(); ) { + const auto & head = infos[runs[r].first]; + + const size_t size = llama_io_run_size(infos, runs[r].first, runs[r].second); + + size_t n_copies = 1; + size_t stride_tensor = 0; + size_t stride_data = 0; + + if (r + 1 < runs.size()) { + const auto & next = infos[runs[r + 1].first]; + + if (next.tensor == head.tensor && next.offset > head.offset && next.ptr > head.ptr && + llama_io_run_size(infos, runs[r + 1].first, runs[r + 1].second) == size) { + stride_tensor = next.offset - head.offset; + stride_data = (size_t) (next.ptr - head.ptr); + + // a strided copy may not have its rows overlap, on either side + if (stride_tensor >= size && stride_data >= size) { + while (r + n_copies < runs.size()) { + const auto & cur = infos[runs[r + n_copies].first]; + + if (cur.tensor != head.tensor || + cur.offset != head.offset + n_copies * stride_tensor || + cur.ptr != head.ptr + n_copies * stride_data || + llama_io_run_size(infos, runs[r + n_copies].first, runs[r + n_copies].second) != size) { + break; + } + + n_copies++; + } + } + } + } + + emit(head.tensor, head.ptr, head.offset, size, n_copies, stride_tensor, stride_data); + + r += n_copies; + } +} + +// a null backend means the caller wants the copy to have happened by the time this returns +static void llama_io_get(ggml_backend_t backend, ggml_tensor * tensor, void * ptr, + size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { + if (n_copies > 1) { + if (backend) { + ggml_backend_tensor_get_2d_async(backend, tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + } else { + ggml_backend_tensor_get_2d(tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + } + } else if (backend) { + ggml_backend_tensor_get_async(backend, tensor, ptr, offset, size); + } else { + ggml_backend_tensor_get(tensor, ptr, offset, size); + } +} + +static void llama_io_set(ggml_backend_t backend, ggml_tensor * tensor, const void * ptr, + size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { + if (n_copies > 1) { + if (backend) { + ggml_backend_tensor_set_2d_async(backend, tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + } else { + ggml_backend_tensor_set_2d(tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + } + } else if (backend) { + ggml_backend_tensor_set_async(backend, tensor, ptr, offset, size); + } else { + ggml_backend_tensor_set(tensor, ptr, offset, size); + } +} + class llama_io_write_host : public llama_io_write_i { public: llama_io_write_host( uint8_t * p, size_t len) : ptr(p), buf_size(len) {} ~llama_io_write_host() { - // TODO: add backend support to batch tensor_get? or some other way to speed this up - for (const auto & winfo : winfos) { - ggml_backend_tensor_get(winfo.tensor, winfo.ptr, winfo.offset, winfo.size); - } + llama_io_emit(winfos, 0, winfos.size(), + [](ggml_tensor * tensor, uint8_t * ptr, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + llama_io_get(nullptr, tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + }); } void write(const void * src, size_t size) override { @@ -2623,13 +2752,23 @@ class llama_io_read_host : public llama_io_read_i { while (end < rinfos.size() && rinfos[end].tensor == tensor) { end++; } + // [TAG_STATE_COALESCE] the fragments the restore emits are one per cell; what + // matters is how many runs of adjacent cells they form, because that is how many + // transfers they actually cost. Count the runs first, and only fall back to + // staging the whole tensor when even the runs are too many. + size_t n_runs = 0; + llama_io_emit(rinfos, i, end, + [&n_runs](ggml_tensor *, const uint8_t *, size_t, size_t, size_t, size_t, size_t) { + n_runs++; + }); + const size_t tensor_bytes = ggml_nbytes(tensor); auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; // A fragmented sequence can require thousands of synchronous device // transfers per layer. For bounded tensors, stage the tensor once and // preserve every byte belonging to other sequences. Bound scratch RAM // and leave ordinary contiguous transfers on their original fast path. - if (end - i >= 64 && tensor_bytes <= 64 * 1024 * 1024 && + if (n_runs >= 64 && tensor_bytes <= 64 * 1024 * 1024 && !ggml_backend_buffer_is_host(buffer)) { std::vector staging; try { @@ -2649,10 +2788,13 @@ class llama_io_read_host : public llama_io_read_i { continue; } } - for (; i < end; ++i) { - const auto & rinfo = rinfos[i]; - ggml_backend_tensor_set(rinfo.tensor, rinfo.ptr, rinfo.offset, rinfo.size); - } + llama_io_emit(rinfos, i, end, + [](ggml_tensor * tensor, const uint8_t * ptr, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + llama_io_set(nullptr, tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + }); + + i = end; } } @@ -2998,6 +3140,321 @@ size_t llama_context::state_set_data(const uint8_t * src, size_t size) { } } +// [TAG_STATE_ASYNC] a sequence state transfer that runs beside the decode instead of in it +// +// Everything the transfer needs to outlive the call that issued it lives here: the host +// buffer the bytes land in or come from, one backend per device holding part of the cache +// (each with a stream of its own, so the copies never queue behind the graphs), and one +// event per device to tell the caller when its half is finished. +struct llama_state_seq_copy { + llama_context * ctx = nullptr; + + struct dev_copy { + ggml_backend_ptr backend; + ggml_backend_event_t event = nullptr; + bool pending = false; + }; + + std::map devs; + + ggml_backend_buffer_ptr host_buf; + + uint8_t * data = nullptr; + size_t size = 0; // bytes the current transfer covers + size_t capacity = 0; // bytes actually held, kept across transfers + bool pinned = false; + bool can_pin = false; + + // transfers the last issue actually posted, i.e. runs of adjacent cells over all tensors + size_t n_copies = 0; + // microseconds the last issue spent draining the compute streams before it could start + int64_t t_sync_us = 0; + + ~llama_state_seq_copy() { + wait(); + + for (auto & it : devs) { + if (it.second.event) { + ggml_backend_event_free(it.second.event); + } + } + } + + // The stream this tensor is copied on, or null when it needs no stream: tensors already + // in host memory are a memcpy, and a tensor in a split or otherwise non-default buffer + // fails the buffer check every backend's async copy asserts, so both take the plain + // synchronous path. Handing a backend out marks it, so record() knows which ones ran. + ggml_backend_t backend_for(const ggml_tensor * t) { + ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; + + if (!buf || ggml_backend_buffer_is_host(buf)) { + return nullptr; + } + + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf); + + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); + + if (!dev || buft != ggml_backend_dev_buffer_type(dev)) { + return nullptr; + } + + auto it = devs.find(dev); + + if (it == devs.end()) { + return nullptr; + } + + it->second.pending = true; + + return it->second.backend.get(); + } + + // close every stream the transfer just used + void record() { + + for (auto & it : devs) { + if (it.second.pending) { + ggml_backend_event_record(it.second.event, it.second.backend.get()); + } + } + } + + bool done() { + bool res = true; + + for (auto & it : devs) { + if (!it.second.pending) { + continue; + } + + if (ggml_backend_event_query(it.second.event)) { + it.second.pending = false; + } else { + res = false; + } + } + + return res; + } + + void wait() { + for (auto & it : devs) { + if (!it.second.pending) { + continue; + } + + ggml_backend_event_synchronize(it.second.event); + + it.second.pending = false; + } + } + + // Grow-only. Pinning host memory is expensive -- a hundred MiB of it costs about as long + // as the copy it is for -- and a caller that parks the same sequence over and over asks + // for a slightly different size every time, so freeing between transfers would put that + // cost back on the very loop this is keeping clear. The memory is given back by + // buf_free() when the caller is finished with the slot, not between two of its parks. + uint8_t * buf_resize(size_t size_new) { + if (size_new <= capacity) { + size = size_new; + + return size_new == 0 ? nullptr : data; + } + + // never move memory a copy could still be reading or writing + wait(); + + host_buf.reset(); + + data = nullptr; + size = 0; + capacity = 0; + pinned = false; + + ggml_backend_buffer_type_t host_buft = host_buffer_type(); + + ggml_backend_buffer_t buf = ggml_backend_buft_alloc_buffer(host_buft, size_new); + + if (!buf) { + return nullptr; + } + + uint8_t * base = (uint8_t *) ggml_backend_buffer_get_base(buf); + + if (!base) { + ggml_backend_buffer_free(buf); + return nullptr; + } + + host_buf.reset(buf); + + data = base; + size = size_new; + capacity = size_new; + // a host buffer type may quietly hand back ordinary memory when pinning is turned + // off, so believe the buffer that came back rather than the type that was asked + pinned = can_pin && ggml_backend_buffer_get_type(buf) == host_buft; + + return data; + } + + void buf_free() { + wait(); + + host_buf.reset(); + + data = nullptr; + size = 0; + capacity = 0; + pinned = false; + } + + // Pinned host memory is the point of allocating through the backend at all: a copy in or + // out of pageable memory is staged through a pinned bounce buffer by the driver and + // blocks, which is exactly the stall being removed here. + ggml_backend_buffer_type_t host_buffer_type() { + for (auto & it : devs) { + ggml_backend_buffer_type_t buft = ggml_backend_dev_host_buffer_type(it.first); + + if (buft) { + return buft; + } + } + + return ggml_backend_cpu_buffer_type(); + } +}; + +class llama_io_write_host_async : public llama_io_write_i { +public: + llama_io_write_host_async(uint8_t * p, size_t len, llama_state_seq_copy & cpy) : + ptr(p), buf_size(len), cpy(cpy) {} + + ~llama_io_write_host_async() { + llama_io_emit(winfos, 0, winfos.size(), + [this](ggml_tensor * tensor, uint8_t * ptr, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + llama_io_get(cpy.backend_for(tensor), tensor, ptr, offset, size, + n_copies, stride_tensor, stride_data); + + cpy.n_copies++; + }); + + cpy.record(); + } + + void write(const void * src, size_t size) override { + if (size > buf_size) { + throw std::runtime_error("unexpectedly reached end of buffer"); + } + memcpy(ptr, src, size); + ptr += size; + size_written += size; + buf_size -= size; + } + + void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { + if (size > buf_size) { + throw std::runtime_error("unexpectedly reached end of buffer"); + } + + winfos.push_back({tensor, ptr, size, offset}); + + ptr += size; + size_written += size; + buf_size -= size; + } + + size_t n_bytes() override { + return size_written; + } + +private: + uint8_t * ptr; + size_t buf_size = 0; + size_t size_written = 0; + + struct write_info { + ggml_tensor * tensor; + uint8_t * ptr; + size_t size; + size_t offset; + }; + std::vector winfos; + + llama_state_seq_copy & cpy; +}; + +class llama_io_read_host_async : public llama_io_read_i { +public: + llama_io_read_host_async(const uint8_t * p, size_t len, llama_state_seq_copy & cpy) : + ptr(p), buf_size(len), cpy(cpy) {} + + ~llama_io_read_host_async() { + // No whole-tensor staging here, unlike the synchronous path above. Staging reads a + // tensor, patches this sequence's bytes into the host copy and writes the whole + // tensor back, which preserves the neighbours only while nothing else is touching + // the cache. These copies are issued precisely so that decoding can carry on beside + // them, so a write-back would undo whatever the sequences sharing the tensor wrote + // to their own cells in the meantime. Writing only this sequence's runs cannot: + // every byte in them belongs to the sequence being restored. That is affordable + // because the runs have been coalesced -- one transfer per run of adjacent cells, + // which is what staging was working around in the first place. + llama_io_emit(rinfos, 0, rinfos.size(), + [this](ggml_tensor * tensor, const uint8_t * ptr, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + llama_io_set(cpy.backend_for(tensor), tensor, ptr, offset, size, + n_copies, stride_tensor, stride_data); + + cpy.n_copies++; + }); + + cpy.record(); + } + + void read(void * dst, size_t size) override { + if (size > buf_size) { + throw std::runtime_error("unexpectedly reached end of buffer"); + } + memcpy(dst, ptr, size); + ptr += size; + size_read += size; + buf_size -= size; + } + + void read_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { + if (size > buf_size) { + throw std::runtime_error("unexpectedly reached end of buffer"); + } + + rinfos.push_back({tensor, ptr, size, offset}); + + ptr += size; + size_read += size; + buf_size -= size; + } + + size_t n_bytes() override { + return size_read; + } + +private: + const uint8_t * ptr; + size_t buf_size = 0; + size_t size_read = 0; + + struct read_info { + ggml_tensor * tensor; + const uint8_t * ptr; + size_t size; + size_t offset; + }; + std::vector rinfos; + + llama_state_seq_copy & cpy; +}; + static constexpr uint32_t io_magic = 0xaf143cd8; size_t llama_context::state_seq_get_size(llama_seq_id seq_id, llama_state_seq_flags flags) { @@ -3071,6 +3528,117 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr } } +// [TAG_STATE_ASYNC] + +llama_state_seq_copy * llama_context::state_seq_copy_init() { + std::unique_ptr cpy(new llama_state_seq_copy()); + + cpy->ctx = this; + + for (auto & backend : backends) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend.get()); + + if (!dev || cpy->devs.find(dev) != cpy->devs.end()) { + continue; + } + + ggml_backend_dev_props props; + ggml_backend_dev_get_props(dev, &props); + + if (!props.caps.async || !props.caps.events) { + continue; + } + + // a backend of its own, not the one the graphs are computed on: that one moves its + // copies to whichever stream it is currently using, so a transfer posted to it could + // end up ordered behind a graph -- which is the stall this exists to avoid + ggml_backend_t backend_cpy = ggml_backend_dev_init(dev, nullptr); + + if (!backend_cpy) { + continue; + } + + ggml_backend_event_t event = ggml_backend_event_new(dev); + + if (!event) { + ggml_backend_free(backend_cpy); + continue; + } + + auto & dc = cpy->devs[dev]; + + dc.backend.reset(backend_cpy); + dc.event = event; + } + + if (cpy->devs.empty()) { + return nullptr; + } + + cpy->can_pin = cpy->host_buffer_type() != ggml_backend_cpu_buffer_type(); + + return cpy.release(); +} + +size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { + if (!cpy.data) { + return 0; + } + + // The copies run on their own stream and are ordered against nothing, so the decode that + // produced these cells has to be finished before they are read. This is the one part of + // the transfer that stays on the caller's thread, and it costs nothing where it is used: + // a caller preempting a sequence does it between two decodes, with the previous one + // already drained by the sampling that followed it. + const int64_t t_sync = ggml_time_us(); + synchronize(); + cpy.t_sync_us = ggml_time_us() - t_sync; + + cpy.n_copies = 0; + + llama_io_write_host_async io(cpy.data, size, cpy); + + try { + io.write(&io_magic, sizeof(io_magic)); + io.write(&seq_id, sizeof(seq_id)); + + return state_seq_write_data(io, seq_id, flags); + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("%s: error saving state: %s\n", __func__, err.what()); + return 0; + } +} + +size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { + if (!cpy.data) { + return 0; + } + + const int64_t t_sync = ggml_time_us(); + synchronize(); + cpy.t_sync_us = ggml_time_us() - t_sync; + + cpy.n_copies = 0; + + llama_io_read_host_async io(cpy.data, size, cpy); + + try { + uint32_t magic_read; + io.read(&magic_read, sizeof(magic_read)); + if (io_magic != magic_read) { + throw std::runtime_error("wrong sequence state magic"); + } + + llama_seq_id seq_id_read; + io.read(&seq_id_read, sizeof(seq_id_read)); + + return state_seq_read_data(io, seq_id, flags); + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); + return 0; + } +} + bool llama_context::state_load_file(const char * filepath, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { llama_file file(filepath, "rb"); @@ -4125,6 +4693,64 @@ size_t llama_state_seq_set_data_ext(llama_context * ctx, const uint8_t * src, si return ctx->state_seq_set_data(seq_id, src, size, flags); } +// [TAG_STATE_ASYNC] + +llama_state_seq_copy * llama_state_seq_copy_init(llama_context * ctx) { + return ctx->state_seq_copy_init(); +} + +void llama_state_seq_copy_free(llama_state_seq_copy * cpy) { + delete cpy; // waits for anything still in flight +} + +uint8_t * llama_state_seq_copy_buf_resize(llama_state_seq_copy * cpy, size_t size) { + return cpy->buf_resize(size); +} + +uint8_t * llama_state_seq_copy_buf(llama_state_seq_copy * cpy) { + return cpy->data; +} + +size_t llama_state_seq_copy_buf_size(llama_state_seq_copy * cpy) { + return cpy->size; +} + +size_t llama_state_seq_copy_buf_capacity(llama_state_seq_copy * cpy) { + return cpy->capacity; +} + +size_t llama_state_seq_copy_n_copies(llama_state_seq_copy * cpy) { + return cpy->n_copies; +} + +int64_t llama_state_seq_copy_sync_us(llama_state_seq_copy * cpy) { + return cpy->t_sync_us; +} + +void llama_state_seq_copy_buf_free(llama_state_seq_copy * cpy) { + cpy->buf_free(); +} + +bool llama_state_seq_copy_buf_is_pinned(llama_state_seq_copy * cpy) { + return cpy->can_pin; +} + +size_t llama_state_seq_copy_get(llama_state_seq_copy * cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { + return cpy->ctx->state_seq_copy_get(*cpy, size, seq_id, flags); +} + +size_t llama_state_seq_copy_set(llama_state_seq_copy * cpy, size_t size, llama_seq_id dest_seq_id, llama_state_seq_flags flags) { + return cpy->ctx->state_seq_copy_set(*cpy, size, dest_seq_id, flags); +} + +bool llama_state_seq_copy_done(llama_state_seq_copy * cpy) { + return cpy->done(); +} + +void llama_state_seq_copy_wait(llama_state_seq_copy * cpy) { + cpy->wait(); +} + size_t llama_state_seq_save_file(llama_context * ctx, const char * filepath, llama_seq_id seq_id, const llama_token * tokens, size_t n_token_count) { ctx->synchronize(); diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b56..3d66f2a948a 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -39,6 +39,9 @@ struct llama_memory_buffer { using llama_memory_buffers = std::map; +// [TAG_STATE_ASYNC] defined in llama-context.cpp +struct llama_state_seq_copy; + struct llama_context { // init scheduler and compute buffers, reserve worst-case graphs llama_context( @@ -156,6 +159,12 @@ struct llama_context { size_t state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, size_t size, llama_state_seq_flags flags); size_t state_seq_set_data(llama_seq_id seq_id, const uint8_t * src, size_t size, llama_state_seq_flags flags); + // [TAG_STATE_ASYNC] the same two transfers, issued on a stream of their own and left running + llama_state_seq_copy * state_seq_copy_init(); + + size_t state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags); + size_t state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id dest_seq_id, llama_state_seq_flags flags); + bool state_load_file( const char * filepath, llama_token * tokens_out, From 1d98f93cbfe4d4ed97cdea5b2aa9ed1c41302d1a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 23:35:19 +0000 Subject: [PATCH 09/81] server: take the park and restore copies off the decode loop preempt_save() and preempt_restore() run inside update_slots(), so while one sequence is copied out of or back into the KV pool every other slot stops. On a 4B at -c 8192 with four chats that is a 250 ms freeze at a park and 149 ms at a restore, seen by chats that had nothing to do with either, against a p99 inter-token gap of 19 ms when nothing is being parked. The park was cheap for the slot it saved; it was the other three that paid for it. A park now has two halves. preempt_save() issues the copy and leaves the slot PREEMPTING: the cells are still its own, because the copy is still reading them, and nobody may take them. update_slots() polls the event each iteration and only then releases the cells and marks the slot PREEMPTED. A restore is the mirror, RESTORING: the cells are allocated and owned by the sequence, so nobody else can take them, but they do not hold its state until the copy lands, which is why the slot is not scheduled and its drafter not rearmed until it does. An asynchronous park does not hand its cells back before update_slots() carries on, so it has to fire earlier than a synchronous one, or the slots that keep decoding have nowhere to put their tokens and end up waiting for the copy after all. preempt_n_margin() keeps eight decode steps of every running slot clear ahead of the pool filling, which is about the tenth of a second a copy of one sequence takes. The same figure gates a resume, so that a slot is not put back into a pool it would immediately have to be taken out of again. When that lookahead is not enough the loop waits for the outstanding park rather than let the KV-full path end every request, which is no worse than the synchronous path and is the last thing tried before giving up. Everything that reads a slot's state had to learn the two new ones. is_processing() is deliberately left as "not idle", because it is what keeps NEXT_RESPONSE posted and the loop polling; narrowing it would deadlock a server whose only slots are mid-copy. preempt_kv_used() deliberately still counts them, since a slot on its way out has not released its cells and one on its way back in has already been given them. release() waits for an outstanding copy before freeing the buffer and handing the cells on, which is the path a cancelled request and every error path take, and where a transfer would otherwise outlive the memory on both ends. --preempt-async (LLAMA_ARG_PREEMPT_ASYNC) is on by default and falls back to the synchronous path on a backend that cannot copy asynchronously, saying so once at load. --no-preempt-async keeps the old behaviour, so both can be compared on one binary. The pinned buffers are held for as long as the task that parked owns the slot rather than freed between two of its parks, so --preempt-ram now bounds the host memory actually held; it still reads zero once the slots are released. Measured on the same four chats, survivors now see 38 to 43 ms at a park and 19 ms at a restore under LLAMA_SERVER_PREEMPT_EVERY=64, against 127 to 158 ms and 112 to 115 ms before, and four-chat throughput goes from 179-184 to 243-267 tok/s. What is left is issue cost: about 11500 transfers at 4 us each, because four chats interleaving in one pool leave a sequence in roughly 160 runs per tensor. A sequence that has the pool to itself is 66 transfers and 0.26 ms. Tests: the asynchronous path is byte-identical to an uninterrupted run and to the synchronous path, two slots that overflow the pool together finish with the tokens they produce alone, cancelling while a copy is in flight leaves no slot stuck and no parked memory held, and --no-preempt-async really does switch it off. --- common/arg.cpp | 10 + common/common.h | 1 + tools/server/README.md | 1 + tools/server/server-context.cpp | 489 ++++++++++++++++++++++-- tools/server/tests/unit/test_preempt.py | 180 +++++++++ 5 files changed, 657 insertions(+), 24 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 5bfa4adcdf0..30a208e707f 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1717,6 +1717,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.preempt_ram_mib = value; } ).set_env("LLAMA_ARG_PREEMPT_RAM").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--preempt-async"}, + {"--no-preempt-async"}, + "copy a parked sequence out of and back into the KV cache on a stream of its own, so the " + "slots that keep running do not wait for it (default: enabled, needs a backend that can " + "copy asynchronously, otherwise the copies are synchronous as before)", + [](common_params & params, bool value) { + params.preempt_async = value; + } + ).set_env("LLAMA_ARG_PREEMPT_ASYNC").set_examples({LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-kvu", "--kv-unified"}, {"-no-kvu", "--no-kv-unified"}, diff --git a/common/common.h b/common/common.h index c99269f9a96..1c220e6b785 100644 --- a/common/common.h +++ b/common/common.h @@ -615,6 +615,7 @@ struct common_params { int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. int32_t preempt_ram_mib = 8192; // host RAM for parked (preempted) sequences: -1 = no limit, 0 = disable preemption + bool preempt_async = true; // park and restore on a stream of their own, off the decode loop std::string hostname = "127.0.0.1"; std::string public_path = ""; // NOLINT diff --git a/tools/server/README.md b/tools/server/README.md index 7b4a0330340..bec3f09ddff 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -165,6 +165,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) | | `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | | `--preempt-ram N` | with a unified KV cache, park a slot in host RAM instead of failing every slot when the cache fills; N is the maximum host RAM for parked sequences in MiB (default: 8192, -1 - no limit, 0 - disable)
(env: LLAMA_ARG_PREEMPT_RAM) | +| `--preempt-async`, `--no-preempt-async` | copy a parked sequence out of and back into the KV cache on a stream of its own, so the slots that keep running do not wait for it (default: enabled, needs a backend that can copy asynchronously, otherwise the copies are synchronous as before)
(env: LLAMA_ARG_PREEMPT_ASYNC) | | `-kvu, --kv-unified, -no-kvu, --no-kv-unified` | use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)
(env: LLAMA_ARG_KV_UNIFIED) | | `--cache-idle-slots, --no-cache-idle-slots` | save idle slots to the prompt cache on new task, and clear them when using unified KV (default: enabled, requires cache-ram)
(env: LLAMA_ARG_CACHE_IDLE_SLOTS) | | `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)
(env: LLAMA_ARG_CONTEXT_SHIFT) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 6723c51397e..f15f78f8201 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -60,6 +60,8 @@ enum slot_state { SLOT_STATE_DONE_PROMPT, SLOT_STATE_GENERATING, SLOT_STATE_PREEMPTED, // [TAG_PREEMPT] cells released, everything needed to resume is in host RAM + SLOT_STATE_PREEMPTING, // [TAG_PREEMPT_ASYNC] the copy out is running; the cells are still this slot's + SLOT_STATE_RESTORING, // [TAG_PREEMPT_ASYNC] the copy back in is running; the cells are allocated but not yet filled }; // [TAG_PREEMPT] server-side request preemption @@ -80,6 +82,34 @@ constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked +// [TAG_PREEMPT_ASYNC] how far ahead of the pool filling an asynchronous park is triggered +// +// A synchronous park hands the cells back before update_slots() goes on, so it only has to +// fire once the next step would not fit. An asynchronous one does not: the copy is still +// reading the cells, and they are only released when it lands. The slots that keep decoding +// in the meantime need somewhere to put their tokens, so the park has to be triggered this +// many decode steps before the pool would actually have run out. Too small and the pool +// fills while the copy is still running, and the decode ends up waiting for it after all, +// which is no worse than not doing this at all but no better either. Too large and slots +// are parked, and so parked again, earlier and more often than they need to be, which costs +// more in total than the one late park it avoided. +// +// Eight steps of every running slot is about a tenth of a second of runway at the speeds a +// handful of parallel chats decode at, which is the order a copy of one sequence takes. +constexpr int32_t PREEMPT_N_ASYNC_STEPS = 8; + +struct llama_state_seq_copy_deleter { + void operator()(llama_state_seq_copy * cpy) const { llama_state_seq_copy_free(cpy); } +}; + +using llama_state_seq_copy_ptr = std::shared_ptr; + +static llama_state_seq_copy_ptr llama_state_seq_copy_make(llama_context * ctx) { + llama_state_seq_copy * cpy = ctx ? llama_state_seq_copy_init(ctx) : nullptr; + + return cpy ? llama_state_seq_copy_ptr(cpy, llama_state_seq_copy_deleter{}) : llama_state_seq_copy_ptr(); +} + struct server_slot; // forward declaration struct server_batch { @@ -320,32 +350,210 @@ struct server_slot { slot_state state_before_preempt = SLOT_STATE_IDLE; std::vector preempt_state_tgt; std::vector preempt_state_dft; + + // [TAG_PREEMPT_ASYNC] the two transfers this slot parks and resumes through + // + // They own the pinned host buffers the sequence lives in while it is parked, so when + // they exist the std::vectors above stay empty and the state is in the transfers. Held + // by shared_ptr only because the slots are built with emplace_back into a vector that + // reallocates as it grows, and a slot must survive being moved. + llama_state_seq_copy_ptr preempt_cpy_tgt; + llama_state_seq_copy_ptr preempt_cpy_dft; + + bool preempt_is_async() const { + return (bool) preempt_cpy_tgt; + } + + // microseconds the last park or resume spent draining the compute streams + int64_t preempt_sync_us() const { + if (!preempt_is_async()) { + return 0; + } + + return llama_state_seq_copy_sync_us(preempt_cpy_tgt.get()) + + (preempt_cpy_dft ? llama_state_seq_copy_sync_us(preempt_cpy_dft.get()) : 0); + } + + // transfers the last park or resume posted, which is what its issue cost is made of + size_t preempt_n_copies() const { + if (!preempt_is_async()) { + return 0; + } + + return llama_state_seq_copy_n_copies(preempt_cpy_tgt.get()) + + (preempt_cpy_dft ? llama_state_seq_copy_n_copies(preempt_cpy_dft.get()) : 0); + } + + // [TAG_PREEMPT_ASYNC] a copy is running for this slot: it is not decoding and must not be + // scheduled, but it still owns cells, so it is neither running nor parked + bool preempt_in_flight() const { + return state == SLOT_STATE_PREEMPTING || state == SLOT_STATE_RESTORING; + } + + // parked, or on its way out or back in: in none of these does the slot take part in a decode + bool preempt_is_out() const { + return state == SLOT_STATE_PREEMPTED || preempt_in_flight(); + } int32_t n_preempt = 0; // times the CURRENT task has been preempted int32_t n_preempt_fail = 0; // consecutive failed restores int64_t t_preempt_us = 0; // when it was parked + int64_t t_preempt_copy_us = 0; // [TAG_PREEMPT_ASYNC] when the current copy was issued size_t preempt_state_size() const { + if (preempt_is_async()) { + // the capacity, not the live size: the pinned buffers are kept between two parks + // of the same task because page-locking them again would cost as much as the + // copy, so what --preempt-ram has to bound is what is held, not what is in use + return llama_state_seq_copy_buf_capacity(preempt_cpy_tgt.get()) + + (preempt_cpy_dft ? llama_state_seq_copy_buf_capacity(preempt_cpy_dft.get()) : 0); + } + return preempt_state_tgt.size() + preempt_state_dft.size(); } void preempt_state_free() { + // resizing waits for anything still in flight first: this is called from release(), + // which a cancelled request reaches while its copy may still be reading or writing + // the buffer, and freeing it underneath a running transfer would be a use-after-free + if (preempt_cpy_tgt) { + llama_state_seq_copy_buf_free(preempt_cpy_tgt.get()); + } + + if (preempt_cpy_dft) { + llama_state_seq_copy_buf_free(preempt_cpy_dft.get()); + } + preempt_state_tgt.clear(); preempt_state_tgt.shrink_to_fit(); preempt_state_dft.clear(); preempt_state_dft.shrink_to_fit(); } + // [TAG_PREEMPT_ASYNC] give up on an outstanding copy without using its result + void preempt_copy_wait() { + if (preempt_cpy_tgt) { + llama_state_seq_copy_wait(preempt_cpy_tgt.get()); + } + + if (preempt_cpy_dft) { + llama_state_seq_copy_wait(preempt_cpy_dft.get()); + } + } + // bytes preempt_save() would need for this slot right now size_t preempt_state_required() const { return llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) + (ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0); } + // take the slot out of the step that is about to be built + // + // The draft is a prediction, not a result, so it goes with the cells. Preemption runs + // before the batch is built, so spec_i_batch is empty and prompt.tokens already holds + // exactly the tokens the state covers -- including the rollback done by the checkpoint + // path when a draft was only partially accepted. + void preempt_detach() { + spec_draft.clear(); + spec_i_batch.clear(); + spec_ckpt.clear(); + spec_is_replay = false; + + i_batch = -1; + } + + // [TAG_PREEMPT_ASYNC] has the copy out finished? if so, the cells can finally go + bool preempt_save_poll() { + if (!llama_state_seq_copy_done(preempt_cpy_tgt.get())) { + return false; + } + + if (preempt_cpy_dft && !llama_state_seq_copy_done(preempt_cpy_dft.get())) { + return false; + } + + mem.seq_rm(id, -1, -1); + + state = SLOT_STATE_PREEMPTED; + + return true; + } + + // [TAG_PREEMPT_ASYNC] has the copy back in finished? if so, the slot can decode again + bool preempt_restore_poll() { + if (!llama_state_seq_copy_done(preempt_cpy_tgt.get())) { + return false; + } + + if (preempt_cpy_dft && !llama_state_seq_copy_done(preempt_cpy_dft.get())) { + return false; + } + + // the state is back in the cache, so the buffers hold nothing that matters; the + // memory itself is kept for the next park of this task and handed back by release() + llama_state_seq_copy_buf_resize(preempt_cpy_tgt.get(), 0); + + if (preempt_cpy_dft) { + llama_state_seq_copy_buf_resize(preempt_cpy_dft.get(), 0); + } + + n_preempt_fail = 0; + + state = state_before_preempt; + + // same call the DONE_PROMPT -> GENERATING transition makes; it reads the restored + // sequence, so it has to wait for the copy like everything else + if (state == SLOT_STATE_GENERATING && can_speculate()) { + common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); + } + + return true; + } + // copy the sequence out of the cache and release its cells + // + // [TAG_PREEMPT_ASYNC] With a transfer this returns as soon as the copy has been issued, + // leaving the slot PREEMPTING: the cells are still its own, because the copy is still + // reading them, and nobody may take them until preempt_save_poll() says the copy landed. bool preempt_save() { const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); const size_t size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + if (preempt_is_async()) { + if (!llama_state_seq_copy_buf_resize(preempt_cpy_tgt.get(), size_tgt) || + (size_dft > 0 && (!preempt_cpy_dft || + !llama_state_seq_copy_buf_resize(preempt_cpy_dft.get(), size_dft)))) { + SLT_ERR(*this, "failed to allocate %.3f MiB of pinned host memory for the preemption state\n", + (size_tgt + size_dft) / (1024.0 * 1024.0)); + preempt_state_free(); + return false; + } + + if (llama_state_seq_copy_get(preempt_cpy_tgt.get(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { + SLT_ERR(*this, "%s", "failed to issue the copy of the target sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + if (size_dft > 0 && + llama_state_seq_copy_get(preempt_cpy_dft.get(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + SLT_ERR(*this, "%s", "failed to issue the copy of the draft sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + preempt_detach(); + + // note: no mem.seq_rm() here. The copy is still reading these cells, so they are + // released in preempt_save_poll() once it has finished with them. + state_before_preempt = state; + state = SLOT_STATE_PREEMPTING; + t_preempt_us = ggml_time_us(); + + n_preempt++; + + return true; + } + try { preempt_state_tgt.resize(size_tgt); preempt_state_dft.resize(size_dft); @@ -369,16 +577,7 @@ struct server_slot { return false; } - // The draft is a prediction, not a result, so it goes with the cells. Preemption - // runs before the batch is built, so spec_i_batch is empty and prompt.tokens already - // holds exactly the tokens the state above covers -- including the rollback done by - // the checkpoint path when a draft was only partially accepted. - spec_draft.clear(); - spec_i_batch.clear(); - spec_ckpt.clear(); - spec_is_replay = false; - - i_batch = -1; + preempt_detach(); // note: prompt.tokens is deliberately kept. It is the mirror of the state just // copied out, and the resume needs it to know how many cells to ask for. @@ -394,7 +593,33 @@ struct server_slot { } // put the sequence back; the slot then continues from the token it was about to decode + // + // [TAG_PREEMPT_ASYNC] With a transfer this returns as soon as the copy has been issued, + // leaving the slot RESTORING: the cells are allocated and owned by this sequence, so + // nobody else can take them, but they do not hold its state until the copy lands, which + // is why the slot is not scheduled until preempt_restore_poll() says so. bool preempt_restore() { + if (preempt_is_async()) { + const size_t size_tgt = llama_state_seq_copy_buf_size(preempt_cpy_tgt.get()); + const size_t size_dft = preempt_cpy_dft ? llama_state_seq_copy_buf_size(preempt_cpy_dft.get()) : 0; + + if (llama_state_seq_copy_set(preempt_cpy_tgt.get(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt || + (size_dft > 0 && + llama_state_seq_copy_set(preempt_cpy_dft.get(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft)) { + // no room after all: let whatever was already issued finish before the + // half-written sequence is dropped, or the cells would go while a copy is + // still writing into them + preempt_copy_wait(); + mem.seq_rm(id, -1, -1); + n_preempt_fail++; + return false; + } + + state = SLOT_STATE_RESTORING; + + return true; + } + const size_t size_tgt = preempt_state_tgt.size(); const size_t size_dft = preempt_state_dft.size(); @@ -649,7 +874,13 @@ struct server_slot { // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be // released while parked), so the mirror of them must not outlive them: the next // task on this slot would otherwise take a prefix match against an empty cache - if (state == SLOT_STATE_PREEMPTED) { + // + // [TAG_PREEMPT_ASYNC] a slot can also be released with a copy still running, by + // a cancelled request or by the error paths. Wait for it before anything else: + // the host buffer is about to be freed and the cells about to be handed to the + // next task, and a transfer still reading or writing either would outlive both. + if (preempt_is_out()) { + preempt_copy_wait(); preempt_state_free(); prompt_clear(); } @@ -796,7 +1027,7 @@ struct server_slot { {"n_ctx", n_ctx}, {"speculative", can_speculate()}, {"is_processing", is_processing()}, - {"is_preempted", state == SLOT_STATE_PREEMPTED}, + {"is_preempted", preempt_is_out()}, {"n_preempt", n_preempt}, }; @@ -1381,6 +1612,22 @@ struct server_context_impl { } }; + // [TAG_PREEMPT_ASYNC] one transfer per context, made once and reused for every + // park and resume this slot ever does, because each owns a backend and a stream + if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { + slot.preempt_cpy_tgt = llama_state_seq_copy_make(ctx_tgt); + + if (slot.preempt_cpy_tgt && ctx_dft) { + slot.preempt_cpy_dft = llama_state_seq_copy_make(ctx_dft); + + if (!slot.preempt_cpy_dft) { + // a draft that cannot go asynchronously would have to be waited for + // in the middle of the park, so the whole slot stays synchronous + slot.preempt_cpy_tgt.reset(); + } + } + } + slot.reset(); } @@ -1402,6 +1649,31 @@ struct server_context_impl { } } + // [TAG_PREEMPT_ASYNC] the slots either all park through a transfer or none do + { + preempt_async_ok = !slots.empty(); + + for (const auto & slot : slots) { + preempt_async_ok = preempt_async_ok && slot.preempt_is_async(); + } + + if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { + if (preempt_async_ok) { + SRV_INF("preemption: parking and resuming asynchronously, %s host memory\n", + llama_state_seq_copy_buf_is_pinned(slots[0].preempt_cpy_tgt.get()) ? "pinned" : "pageable"); + } else { + SRV_WRN("%s", "preemption: this backend cannot copy asynchronously, parking and resuming synchronously\n"); + } + } + + if (!preempt_async_ok) { + for (auto & slot : slots) { + slot.preempt_cpy_tgt.reset(); + slot.preempt_cpy_dft.reset(); + } + } + } + { const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; @@ -2554,7 +2826,7 @@ struct server_context_impl { if (slot.is_processing()) { n_processing_slots++; } - if (slot.state == SLOT_STATE_PREEMPTED) { + if (slot.preempt_is_out()) { n_preempted_slots++; } } @@ -2855,6 +3127,11 @@ struct server_context_impl { // uninterrupted one is the preemption's fault and nothing else's. int32_t preempt_test_every = 0; + // [TAG_PREEMPT_ASYNC] whether the slots park and resume through a transfer. False when + // --no-preempt-async was given, or when the backend cannot copy asynchronously, in which + // case every park and resume is the synchronous one it always was. + bool preempt_async_ok = false; + int32_t preempt_n_spec_max() const { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } @@ -2878,7 +3155,13 @@ struct server_context_impl { const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; - return preempt_ram_used() + slot.preempt_state_required() <= budget; + // whatever this slot already holds is counted by preempt_ram_used() and will be + // reused, so parking it again only costs what it does not have yet + const size_t held = slot.preempt_state_size(); + const size_t need = slot.preempt_state_required(); + const size_t extra = need > held ? need - held : 0; + + return preempt_ram_used() + extra <= budget; } // cells the slot will ask for on its next step once it is back in the pool @@ -2908,12 +3191,40 @@ struct server_context_impl { continue; // parked: its cells are in host RAM, not in the pool } + // [TAG_PREEMPT_ASYNC] deliberately not skipped: a slot whose copy is still + // running holds cells either way. One on its way out has not released them yet + // because the copy is still reading them, and one on its way back in has already + // been given them. Skipping either would hand the same cells out twice. + res += slot.prompt.n_tokens(); } return res; } + // [TAG_PREEMPT_ASYNC] the room the pool is kept clear of + // + // A synchronous park releases the cells before update_slots() carries on, so it only has + // to fire once the next step would not fit. An asynchronous one leaves them held until + // its copy lands, so it has to fire early enough that everything still decoding has + // somewhere to put its tokens until then. The same figure gates a resume, so that a slot + // is not put back into a pool it would immediately have to be taken out of again. + int32_t preempt_n_margin() const { + if (!preempt_async_active()) { + return PREEMPT_N_MARGIN; + } + + int32_t n_running = 0; + + for (const auto & slot : slots) { + if (slot.is_processing() && !slot.preempt_is_out()) { + n_running++; + } + } + + return PREEMPT_N_MARGIN + n_running * (1 + preempt_n_spec_max()) * PREEMPT_N_ASYNC_STEPS; + } + // cells those slots are about to ask for on the next decode int32_t preempt_kv_reserve() const { const int32_t n_spec = preempt_n_spec_max(); @@ -2936,6 +3247,21 @@ struct server_context_impl { res_pmt += std::max(1, std::min(n_batch, n_left)); } break; + case SLOT_STATE_RESTORING: + { + // [TAG_PREEMPT_ASYNC] its cells are already counted by preempt_kv_used(), + // but it starts decoding as soon as its copy lands, so the step it will + // take has to be reserved now -- otherwise the pool is handed out from + // under it and its first step preempts somebody else straight away + if (slot.state_before_preempt == SLOT_STATE_GENERATING) { + res += 1 + n_spec; + } else { + const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + + res_pmt += std::max(1, std::min(n_batch, n_left)); + } + } break; + // a preempting slot is on its way out and will not decode: nothing to reserve default: break; } @@ -2953,7 +3279,7 @@ struct server_context_impl { int32_t n_running = 0; for (auto & slot : slots) { - if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + if (slot.is_processing() && !slot.preempt_is_out()) { n_running++; if (!leader || slot.prompt.n_tokens() > leader->prompt.n_tokens()) { @@ -3011,11 +3337,75 @@ struct server_context_impl { // at a token boundary, prompt.tokens is exactly what the cache holds for it, and no // draft is in flight, so a slot can be removed from the picture without unpicking a // half-decoded batch + // [TAG_PREEMPT_ASYNC] is any slot parking or resuming through a transfer right now + bool preempt_async_active() const { + return preempt_async_ok; + } + + // Pick up the copies that have landed since the last iteration. This runs before + // anything reads preempt_kv_used(), so a park whose cells came back is seen as free + // room straight away and a resume that landed can be scheduled in the same iteration. + void update_preempt_copies() { + for (auto & slot : slots) { + if (slot.state == SLOT_STATE_PREEMPTING) { + if (slot.preempt_save_poll()) { + metrics.n_preempt++; + + SLT_WRN(slot, "park completed after %.2f ms: %d cells released, %.1f MiB parked, kv %d/%d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, + slot.prompt.n_tokens(), + slot.preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_ctx); + } + } else if (slot.state == SLOT_STATE_RESTORING) { + if (slot.preempt_restore_poll()) { + metrics.n_resume++; + + SLT_WRN(slot, "restore completed after %.2f ms: %d tokens back in the cache, kv %d/%d, preemptions %d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, + slot.prompt.n_tokens(), + preempt_kv_used(), n_ctx, + slot.n_preempt); + } + } + } + } + + // Wait for one outstanding park, the last thing tried before giving up on finding room. + // It is what keeps a pool that fills faster than the copies drain no worse than the + // synchronous path: the decode waits for the copy exactly as it used to. + bool preempt_wait_in_flight() { + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_PREEMPTING) { + continue; + } + + slot.preempt_copy_wait(); + + if (!slot.preempt_save_poll()) { + continue; + } + + metrics.n_preempt++; + + SLT_WRN(slot, "park completed after %.2f ms (waited for): %d cells released, kv %d/%d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, + slot.prompt.n_tokens(), + preempt_kv_used(), n_ctx); + + return true; + } + + return false; + } + void update_preemption() { if (!params_base.kv_unified || slots.size() < 2) { return; // with a cache per slot, no slot can take another one's cells } + update_preempt_copies(); + if (params_base.preempt_ram_mib == 0) { return; // --preempt-ram 0: the KV-full retry ladder, as before } @@ -3055,7 +3445,7 @@ struct server_context_impl { // continue, so give those cells up first - same call the KV-full path makes. for (;;) { for (auto * slot : parked) { - if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + PREEMPT_N_MARGIN <= n_cells) { + if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + preempt_n_margin() <= n_cells) { best = slot; break; } @@ -3072,6 +3462,8 @@ struct server_context_impl { const int64_t t_start = ggml_time_us(); + best->t_preempt_copy_us = t_start; + if (!best->preempt_restore()) { // update_slots() runs in a tight loop while tasks are pending, so a counter // alone burns its whole budget in a couple of milliseconds. Give up only on @@ -3090,6 +3482,22 @@ struct server_context_impl { break; } + // [TAG_PREEMPT_ASYNC] with a transfer the copy has only been issued; the slot is + // RESTORING and update_preempt_copies() counts it and logs it when it lands + if (best->state == SLOT_STATE_RESTORING) { + SLT_WRN(*best, "resumed after %.2f s: %d tokens, restore issued in %.2f ms (%zu transfers, %.2f ms sync), kv %d/%d, preemptions %d\n", + (ggml_time_us() - best->t_preempt_us) / 1e6, + best->prompt.n_tokens(), + (ggml_time_us() - t_start) / 1e3, + best->preempt_n_copies(), best->preempt_sync_us() / 1e3, + preempt_kv_used(), n_cells, + best->n_preempt); + + // it holds cells now but is not decoding yet, so there is nothing more to + // decide about it this iteration + continue; + } + metrics.n_resume++; SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", @@ -3105,12 +3513,19 @@ struct server_context_impl { for (auto & slot : slots) { if (slot.state == SLOT_STATE_GENERATING && (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && - preempt_fits_budget(slot) && - slot.preempt_save()) { - metrics.n_preempt++; + preempt_fits_budget(slot)) { + slot.t_preempt_copy_us = ggml_time_us(); + + if (slot.preempt_save()) { + // [TAG_PREEMPT_ASYNC] a slot left PREEMPTING is counted by + // update_preempt_copies() when its copy lands, not here + if (slot.state == SLOT_STATE_PREEMPTED) { + metrics.n_preempt++; + } - SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", - (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); + SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", + (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); + } } } } @@ -3119,7 +3534,7 @@ struct server_context_impl { for (;;) { const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); - if (n_used + PREEMPT_N_MARGIN <= n_cells) { + if (n_used + preempt_n_margin() <= n_cells) { break; } @@ -3128,6 +3543,15 @@ struct server_context_impl { continue; } + // [TAG_PREEMPT_ASYNC] Out of room for the step about to be built, rather than + // merely short of the lookahead the asynchronous path keeps. A park that has + // been issued but not landed is holding cells that are already spoken for, and + // waiting for it is both quicker and more useful than parking somebody else, + // whose cells would not come back this iteration either. + if (n_used + PREEMPT_N_MARGIN > n_cells && preempt_wait_in_flight()) { + continue; + } + server_slot * victim = preempt_pick_victim(); if (!victim) { @@ -3139,10 +3563,27 @@ struct server_context_impl { const int32_t n_tokens = victim->prompt.n_tokens(); const int64_t t_start = ggml_time_us(); + victim->t_preempt_copy_us = t_start; + if (!victim->preempt_save()) { break; // could not park it; the existing retry ladder is still behind us } + // [TAG_PREEMPT_ASYNC] the copy has only been issued; the cells are still the + // victim's until it lands, so nothing further can be decided about the pool this + // iteration. update_preempt_copies() picks it up on the next one, and the step + // that wanted the room is built from whatever is free right now. + if (victim->state == SLOT_STATE_PREEMPTING) { + SLT_WRN(*victim, "preempted: %d cells, park issued in %.2f ms (%zu transfers, %.2f ms sync), %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", + n_tokens, + (ggml_time_us() - t_start) / 1e3, + victim->preempt_n_copies(), victim->preempt_sync_us() / 1e3, + victim->preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_cells, n_used, + victim->n_preempt); + break; + } + metrics.n_preempt++; SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", @@ -3485,7 +3926,7 @@ struct server_context_impl { // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to // batch; it takes no part in this pass until it is restored - if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { + if (!slot.is_processing() || slot.preempt_is_out()) { return; } @@ -4429,7 +4870,7 @@ struct server_context_impl { metrics.n_decode++; for (const auto & slot : slots) { // [TAG_PREEMPT] a parked slot is processing but took no part in this decode - if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + if (slot.is_processing() && !slot.preempt_is_out()) { metrics.n_busy_slots++; } metrics.n_tokens_max = std::max(metrics.n_tokens_max, (uint64_t) slot.prompt.n_tokens()); diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 0da885bcafd..1aba01751bd 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -39,6 +39,7 @@ def create_server(): yield os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) + os.environ.pop("LLAMA_ARG_PREEMPT_ASYNC", None) def _complete(n_predict: int, prompt: str = "Hi how are you"): @@ -267,3 +268,182 @@ def test_metrics_and_slots_report_the_parked_state(): res = server.make_request("GET", "/slots") assert res.status_code == 200 assert sum(slot["n_preempt"] for slot in res.body) == 0, "n_preempt is per task and resets with the slot" + + +# [TAG_PREEMPT_ASYNC] parking and resuming on a stream of their own +# +# The copies are only asynchronous on a backend that can copy asynchronously and signal an +# event, which today means a GPU one. On a CPU-only build the server says so and falls back +# to the synchronous path, and the tests below that need the asynchronous one skip. + +_ASYNC_BANNER = "parking and resuming asynchronously" + + +def _start_async(**kwargs) -> str: + """Start the server with the asynchronous path asked for, and return its log so far.""" + os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "1" + for key, value in kwargs.items(): + setattr(server, key, value) + server.start() + with open(server.log_path) as f: + return f.read() + + +def _require_async(text: str): + if _ASYNC_BANNER not in text: + pytest.skip("this backend cannot copy asynchronously, the async park path is not exercised") + + +def test_async_preemption_does_not_change_the_output(): + # The same question the synchronous determinism test asks, of the asynchronous path: + # with one request the batch has the same shape at every step, so a continuation that + # was parked and resumed through a transfer and is not byte-identical to an + # uninterrupted one is the transfer's fault and nothing else's. + global server + server.n_ctx = 512 + server.n_gpu_layer = 99 + text = _start_async() + _require_async(text) + + res_plain = _complete(64) + assert res_plain.status_code == 200 + + server.stop() + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + log = LogReader(server.log_path) + + res_preempted = _complete(64) + assert res_preempted.status_code == 200 + + text = log.drain() + _require_async(text) + assert text.count("preempted on request") >= 6 + assert text.count("resumed after") >= 6 + # the asynchronous path is the one that ran, not the synchronous fallback: only it + # splits a park and a resume into an issue and a completion + assert "park completed after" in text + assert "restore issued in" in text + assert "restore completed after" in text + + assert res_preempted.body["content"] == res_plain.body["content"] + assert res_preempted.body["tokens"] == res_plain.body["tokens"] + + +def test_async_preemption_under_load_keeps_every_slot_and_its_output(): + # Two requests that do not fit the pool together, parked and resumed asynchronously + # while the other one keeps decoding. Every slot must finish, and finish with exactly + # the tokens it produces when it has the pool to itself. + global server + server.n_ctx = 256 + server.n_gpu_layer = 99 + text = _start_async() + _require_async(text) + + n_predict = 160 + prompts = [ + "Once upon a time there was a brave knight who", + "The quick brown fox jumps over the lazy dog and", + ] + + # each one alone, for the reference tokens + alone = [_complete(n_predict, prompt) for prompt in prompts] + for res in alone: + assert res.status_code == 200 + + server.stop() + server.start() + log = LogReader(server.log_path) + + together = parallel_function_calls([(_complete, (n_predict, prompt)) for prompt in prompts]) + + text = log.drain() + _require_async(text) + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + for res, ref in zip(together, alone): + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False + assert res.body["tokens"] == ref.body["tokens"] + + +def _cancel_soon(n_predict: int, prompt: str, timeout: float): + try: + server.make_request("POST", "/completion", data={ + "n_predict": n_predict, + "prompt": prompt, + "ignore_eos": True, + "temperature": 0.0, + "seed": 42, + }, timeout=timeout) + except Exception: + pass # the point is the drop, not the response + + +def test_cancel_while_a_copy_is_in_flight_frees_the_slot(): + # A cancelled request can reach release() with a park or a resume still running, which + # is where the host buffer is freed and the cells are handed on. Both have to wait for + # the copy first. LLAMA_SERVER_PREEMPT_EVERY keeps every slot cycling between the two + # states, so cancelling at a spread of moments lands in both; what is asserted is that + # the server survives it, the slots come back, and it still answers correctly. + global server + server.n_ctx = 512 + server.n_gpu_layer = 99 + # every 8 tokens, so a slot spends most of its life in one of the two copy states, but + # not so often that the abandoned requests take minutes to drain + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + text = _start_async() + _require_async(text) + + for i in range(4): + _cancel_soon(96, "Once upon a time there was a brave knight who", 0.05 + 0.1 * i) + + # every slot back, and none of them still holding a parked sequence + deadline = time.time() + 120 + while time.time() < deadline: + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + if all(not slot["is_processing"] for slot in res.body): + break + time.sleep(0.2) + else: + pytest.fail("a slot never came back after a cancel during a copy") + + for slot in res.body: + assert slot["is_preempted"] is False + + if server.server_metrics: + res = server.make_request("GET", "/metrics") + for line in res.body.splitlines(): + if line.startswith("llamacpp:preempt_ram_bytes"): + assert float(line.split(" ", 1)[1]) == 0, "a cancelled slot kept its parked memory" + + # and the server still works + res = _complete(16) + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == 16 + + +def test_no_preempt_async_falls_back_to_the_synchronous_path(): + # The flag has to really switch it off, so that the two can be compared on one binary. + global server + server.n_ctx = 512 + server.n_gpu_layer = 99 + os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "0" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + log = LogReader(server.log_path) + + res = _complete(64) + assert res.status_code == 200 + + text = log.drain() + assert _ASYNC_BANNER not in text + assert "park issued in" not in text + assert "restore issued in" not in text + # the synchronous path still parks and resumes + assert text.count("preempted on request") >= 6 + assert text.count("resumed after") >= 6 From 4814a26852444bc8200038f6fa18827b874bffcb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 02:10:23 +0000 Subject: [PATCH 10/81] cuda: give every MUL_MAT_ID token the single-token configuration A mixture-of-experts decode groups the ubatch's tokens by the expert they routed to, so the column count of an expert's matmul, the rows the per-expert copy gathers and the width the activations are quantized at all depend on what the other tokens in the ubatch picked. The quantized path makes it visible: at one token per ubatch MUL_MAT_ID runs mul_mat_vec_q at ncols_dst 1, four warps dividing the K loop and a shared memory reduction across them, and at more than one token it runs the dedicated MoE kernel, one warp per token with a warp only reduction. Whether those two agree bit for bit depends on the quantization type and on K. Measured on a B200 they agree for Q4_K and Q5_K up to K 2048 and disagree for Q6_K and Q8_0 from K 512, which is why on Qwen3.6-35B-A3B-UD-Q4_K_XL the Q4_K gate and up projections of every layer matched and the three Q6_K down projections, layers 34, 38 and 39, did not. Under GGML_CUDA_BATCH_INVARIANT the node is now computed one token at a time. Each call then sees the shapes a batch of one has, whatever the neighbours routed to, which covers the ids variants of MMVQ, MMQ and MMF and the sorted per expert fallback with one change. GGML_CUDA_BATCH_INVARIANT_MAX_COLS bounds it the way it bounds the MUL_MAT column split. The CUDA graph fallback check is evaluated against the single-token path as well, since that is what a split node runs. On the 35B the sequence-0 slice of one decode step goes from 195 of 3727 nodes differing between a one token and a four token ubatch to 0, and a standalone MUL_MAT_ID probe over Q4_K, Q5_K, Q6_K, Q8_0, Q4_0, Q3_K, Q2_K, MXFP4, F16 and BF16 at K 512, 2048 and 4096 goes to 0 at every token count up to 8. --- ggml/src/ggml-cuda/ggml-cuda.cu | 78 +++++++++++++++++++++++++++++++-- tests/test-backend-ops.cpp | 11 +++++ 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2b1754baea6..0995acdfd6a 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2030,6 +2030,25 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor GGML_ABORT("fatal error"); } +// [TAG_BATCH_INVARIANT] +// True when the batch-invariant policy computes this MUL_MAT_ID one token at a time. +// Every expert product then reduces the way it would in a batch of one, whatever the +// rest of the ubatch routed to. +static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) { + if (!ggml_cuda_batch_invariant()) { + return false; + } + const int64_t ntokens = dst->ne[2]; + if (ntokens <= 1) { + return false; + } + const int max_cols = ggml_cuda_batch_invariant_max_cols(); + if (max_cols > 0 && ntokens > max_cols) { + return false; + } + return true; +} + // returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization // [TAG_MUL_MAT_ID_CUDA_GRAPHS] static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int cc) { @@ -2040,9 +2059,13 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c return true; } - if (dst->ne[2] <= MMVQ_MAX_BATCH_SIZE) { + // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path + // that decides whether the stream is synchronized is the single-token one. + const int64_t ntokens = ggml_cuda_mul_mat_id_splits_tokens(dst) ? 1 : dst->ne[2]; + + if (ntokens <= MMVQ_MAX_BATCH_SIZE) { if (ggml_is_quantized(src0->type)) { - if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)) { + if (ntokens <= get_mmvq_mmid_max_batch(src0->type, cc)) { return false; } } else if (GGML_CUDA_CC_IS_AMD(cc)) { @@ -2050,17 +2073,57 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c } } - if (ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[2], /*n_experts=*/src0->ne[2])) { + if (ggml_cuda_should_use_mmq(src0->type, cc, ntokens, /*n_experts=*/src0->ne[2])) { return false; } - if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, ntokens, /*mul_mat_id=*/true)) { return false; } return true; } +static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// [TAG_BATCH_INVARIANT] +// Recompute dst one token at a time. Every implementation below groups the ubatch's tokens +// by the expert they routed to, so the column count of an expert's matmul, the tokens the +// per-expert copy gathers and the width the activations are quantized at all depend on what +// the other tokens in the ubatch picked. Handing each token its own call removes that: the +// callee sees the shapes a batch of one has, whatever the neighbours did. +static void ggml_cuda_mul_mat_id_split_tokens(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * ids = dst->src[2]; + + const int64_t ntokens = dst->ne[2]; + + for (int64_t i = 0; i < ntokens; ++i) { + ggml_tensor src1_token = *src1; + ggml_tensor ids_token = *ids; + ggml_tensor dst_token = *dst; + + // src1 is [ne10, ne11, ntokens], one expert list per token in ids [n_expert_used, ntokens] + src1_token.ne[2] = 1; + src1_token.nb[3] = src1_token.nb[2]; + src1_token.data = (char *) src1->data + i*src1->nb[2]; + + ids_token.ne[1] = 1; + ids_token.nb[2] = ids_token.nb[1]; + ids_token.nb[3] = ids_token.nb[1]; + ids_token.data = (char *) ids->data + i*ids->nb[1]; + + dst_token.ne[2] = 1; + dst_token.nb[3] = dst_token.nb[2]; + dst_token.data = (char *) dst->data + i*dst->nb[2]; + + dst_token.src[1] = &src1_token; + dst_token.src[2] = &ids_token; + + ggml_cuda_mul_mat_id(ctx, &dst_token); + } +} + static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -2073,6 +2136,13 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + // [TAG_BATCH_INVARIANT] + if (ggml_cuda_mul_mat_id_splits_tokens(dst)) { + GGML_ASSERT(ne3 == 1 && src1->ne[3] == 1 && ids->ne[2] == 1 && ids->ne[3] == 1); + ggml_cuda_mul_mat_id_split_tokens(ctx, dst); + return; + } + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { static_assert(MMVQ_MAX_BATCH_SIZE == MMVF_MAX_BATCH_SIZE); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a7511396791..685193be77d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9213,6 +9213,17 @@ static std::vector> make_test_cases_eval() { } } + // Mixture-of-experts projections at the token counts a decode ubatch forms. The gate and up + // projections broadcast one activation row over the expert list, the down projection carries + // one row per expert, and the mixed quantization of a real MoE gguf puts different types on + // the two. 17 tokens is past the width the exact-concurrency policy pins. + for (ggml_type type_a : {GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16}) { + for (int n : {1, 2, 4, 8, 17}) { + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, true, 512, n, 2048)); + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, false, 2048, n, 512)); + } + } + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_MXFP4, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); From 7b0d7edb9eccadde314cffdd7445d5c9aefb8e10 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 02:10:34 +0000 Subject: [PATCH 11/81] cuda: leave the top-k routing unfused under GGML_CUDA_BATCH_INVARIANT ggml_cuda_check_fusion_memory_ranges accepts the top-k MoE subgraph through an explicit ggml_nrows(node) == 1 exception, which skips the aliasing test on the grounds that each row is read entirely before it is written. With more than one token the generic overlap test runs instead and refuses. So a request decoding alone computes its routing weights with the fused warp local top-k kernel and the same request decoding next to three others computes them with the softmax, argsort, get_rows, sum, clamp and divide chain. Two algorithms for one set of routing weights is exactly the batch dependence this knob removes, and the same reason mul_mat plus GLU fusion is already off here. A node probe cannot see this, which is worth recording: registering an eval callback disables fusion, so both sides of the comparison take the unfused chain and agree. It only appears when the two are compared without one. On Qwen3.6-35B-A3B, with every node of one decode step already byte-identical, sequence 0's logits differed in 248319 of 248320 entries from the first decode step, by up to 2.6e-1, and the greedy stream flipped a token at step 47. The dense 4B is unaffected either way, and disabling every CUDA fusion removes it, which is what named the fused op. With the routing left unfused under the knob, 512 decode steps of sequence 0 alone against sequence 0 next to three neighbours are byte-identical on the 35B, and every cell of the server matrix reads identical. --- ggml/src/ggml-cuda/ggml-cuda.cu | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 0995acdfd6a..30c33b23eaf 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3528,9 +3528,14 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } - //topk-moe - if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || - cgraph->nodes[i]->op == GGML_OP_ARGSORT) { + // topk-moe + // [TAG_BATCH_INVARIANT] The routing fusion passes its memory-range check only when the ubatch + // holds one token, so a request decoding alone picks the fused warp-local top-k kernel and the + // same request decoding next to neighbours picks the softmax, argsort and normalize chain. + // Two algorithms for one set of routing weights is the batch dependence this mode removes. + if (!ggml_cuda_batch_invariant() && + (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || + cgraph->nodes[i]->op == GGML_OP_ARGSORT)) { ggml_cuda_topk_moe_args args; const bool can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args); std::vector ops; From 65860ea386d10a8bb7c662fdc72ce0d7b689b757 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 05:43:07 +0000 Subject: [PATCH 12/81] cuda: default exact mode to a bound that covers the speculative verify batch Exact mode still defaults the column policy to unbounded, so every measurement that recovered the prefill cost had to set GGML_CUDA_BATCH_INVARIANT_MAX_COLS by hand, and the value everything was measured at, 8, does not cover a speculative verify ubatch by construction: with --parallel 4 and --spec-type draft-mtp --spec-draft-n-max 2 a verify ubatch holds one accepted token plus two drafts per slot, up to 12 columns, and above the bound neither the MUL_MAT column split nor the MUL_MAT_ID per token split fires. Default the bound to 16 in exact mode instead, which covers four slots at up to three tokens each. An explicitly set GGML_CUDA_BATCH_INVARIANT_MAX_COLS still wins, in either mode, so a deployment with more slots or a wider draft can raise it. Measured on one B200, Qwen3.6-35B-A3B-UD-Q4_K_XL and Qwen3.5-4B-UD-Q4_K_XL, --parallel 4 --kv-unified -c 8192 --flash-attn on -ngl 99 --seed 0, greedy sampling, 512 predicted tokens, P0 solo twice then P0 concurrent with P1 to P3, three rounds per cell. Every P0 completion was byte identical to its solo reference and every solo repeat matched: 35B, draft-mtp n-max 2, MAX_COLS=16 identical x3 35B, draft-mtp n-max 2, MAX_COLS=16, PREEMPT_EVERY=64 identical x3, 98/98 parks 35B, draft-mtp n-max 4, MAX_COLS=8 (verify up to 20) identical x3 35B, draft-mtp n-max 4, MAX_COLS=32 identical x3 4B, draft-mtp n-max 2, MAX_COLS=16, PREEMPT_EVERY=64 identical x3, 98/98 parks 35B, draft-mtp n-max 2, new default, no env var identical x3 35B, spec off, new default, no env var identical x3 The 35B MTP cell at 16 reproduces the acceptance counters of the same cell at 8 exactly, 4091 of 6103 draft tokens, so raising the bound does not perturb the generation. Cost, three interleaved MAX_COLS=8 against MAX_COLS=16 pairs on the 35B with speculation on, medians: solo decode 25.42 against 25.37 tok/s, four chat aggregate decode 54.55 against 54.23 tok/s, four chat wall 19.98 against 20.15 s. Four interleaved pairs with speculation off, where a decode ubatch is four columns wide and both bounds must behave identically, medians 31.39 against 31.40 tok/s solo and 111.41 against 111.41 tok/s aggregate. All within the run to run spread. test-backend-ops test -b CUDA0 -o MUL_MAT_ID,MUL_MAT under LLAMA_EXACT_CONCURRENCY=1 GGML_CUDA_BATCH_INVARIANT=2 MAX_COLS=16: 2136/2136 passed, 2/2 backends. --- ggml/src/ggml-cuda/ggml-cuda.cu | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 30c33b23eaf..4adee44ca61 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1856,10 +1856,15 @@ int ggml_cuda_batch_invariant_max_cols() { static const int max_cols = []() { // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill // matmul shapes match its solo run, so exact mode no longer needs the column policy to be - // unbounded there. Honour an explicit bound when one is set; default to unbounded. + // unbounded there. An explicit bound always wins, in either mode. const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); if (val) { return atoi(val); } - if (ggml_cuda_exact_concurrency()) { return 0; } + // Exact mode then only has to cover the widest ubatch a decode step can build: one column + // per slot, times one plus the number of speculative draft tokens carried with it. 16 + // covers the default four slots at up to three tokens each, which is what + // --spec-type draft-mtp --spec-draft-n-max 2 produces. More slots, or a wider draft, need + // the bound set explicitly; above it the column split does not fire. + if (ggml_cuda_exact_concurrency()) { return 16; } return 0; }(); return max_cols; From d00cb7f91182a75e294a74ea0baf3eaf94fcb097 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:12:58 +0000 Subject: [PATCH 13/81] server: refuse n > 1 under exact concurrency instead of aborting in the cache A review of #194 pointed at the new GGML_ASSERT in llama_kv_cache::seq_cp, and it is right. Reproduced on this branch with the 4B on one B200: LLAMA_EXACT_CONCURRENCY=1, POST /completion {"n": 2} -> llama-kv-cache.cpp:458: GGML_ASSERT(!exact_pages || seq_id_src == seq_id_dst) failed, through server_context_impl::decode -> common_memory::seq_cp -> process aborted, the next request gets connection refused With the mode off the same request is served normally, so this is reachable by any client of an exact-mode server and takes every other request on the machine with it. Two changes: The server refuses the request. n_cmpl > 1 works by copying the parent's cells to a second sequence id, and exact mode gives a page to one sequence, so there is nowhere for that copy to land. Rejecting it where the task is built turns it into a 400 with a reason. The check reads LLAMA_EXACT_CONCURRENCY from the environment, the same way the KV cache, the batch splitter and the CUDA backend each do, because the answer is needed before a context exists and the mode has no other representation. The cache stops aborting. seq_cp, seq_add and seq_div log an error and return instead of asserting, so a caller this branch does not know about degrades to a refused operation rather than killing the server. The guards also move below the shared-cells early return, which the asserts sat above: a draft cache forwards these calls and copies nothing of its own, and it should not be judged by a rule about cells it does not own. After: the n=2 request returns 400 on both /completion and /v1/completions, the server stays up, and a following ordinary request returns 200. --- src/llama-kv-cache.cpp | 35 ++++++++++++++++++++++++++++++--- tools/server/server-context.cpp | 25 +++++++++++++++++++++++ 2 files changed, 57 insertions(+), 3 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index df643a047a9..297786fb889 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -455,12 +455,24 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { } void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { - GGML_ASSERT(!exact_pages || seq_id_src == seq_id_dst); // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; } + // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so cells cannot be shared + // between two of them. Refuse the operation rather than abort the process: a server + // rejects the request that would reach here (n_cmpl > 1), and any caller this does not + // cover degrades to a failed copy it can report instead of killing every other request + // on the machine. Placed after the shared-cells return so a draft cache, which copies + // nothing of its own, is unaffected. + if (exact_pages && seq_id_src != seq_id_dst) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between " + "sequences (%d -> %d); ignoring the copy\n", + __func__, seq_id_src, seq_id_dst); + return; + } + GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size()); GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size()); @@ -575,12 +587,21 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) { } void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { - GGML_ASSERT(!exact_pages || shift == 0); // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; } + // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo the + // page size, so shifting positions would put every cell of the sequence in the wrong + // place. Context shift is unsupported in exact mode; say so rather than abort. + if (exact_pages && shift != 0) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions " + "(seq %d, shift %d); ignoring the shift\n", + __func__, seq_id, shift); + return; + } + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1"); @@ -626,12 +647,20 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll } void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { - GGML_ASSERT(!exact_pages || d == 1); // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; } + // [TAG_EXACT_CONCURRENCY] same reason as seq_add: dividing positions breaks the + // identity between a cell's position and its offset inside its page. + if (exact_pages && d != 1) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions " + "(seq %d, d %d); ignoring the division\n", + __func__, seq_id, d); + return; + } + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1"); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e1cea00b29b..999a6e723a9 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -37,6 +37,19 @@ constexpr int HTTP_POLLING_SECONDS = 1; +// [TAG_EXACT_CONCURRENCY] the knob is read from the environment by the KV cache, the batch +// splitter and the CUDA backend independently, because it has to be answered before a +// context exists. The server needs the same answer to refuse the one request shape the mode +// cannot serve, so it reads it the same way rather than growing a public API for it. +static bool server_exact_concurrency() { + static const bool enabled = []() { + const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); + return val && atoi(val) != 0; + }(); + + return enabled; +} + static common_speculative_output_limits server_output_limits(const common_params & params) { if (params.embedding || (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) { @@ -5177,6 +5190,18 @@ std::unique_ptr server_routes::handle_completions_impl( task.params.oaicompat_cmpl_id = completion_id; task.params.oaicompat_model = meta->model_name; + // [TAG_EXACT_CONCURRENCY] the children of an n_cmpl > 1 task are served by + // copying the parent's cells to another sequence id, and exact mode gives a KV + // page to one sequence, so there is nothing for that copy to land in. Refuse + // the request here, where it becomes a 400 the client can read, rather than + // letting it reach seq_cp with nothing to do. + if (task.params.n_cmpl > 1 && server_exact_concurrency()) { + throw std::runtime_error( + "n > 1 is not supported while LLAMA_EXACT_CONCURRENCY is set: each " + "completion needs its own sequence, and in exact mode a KV page belongs " + "to a single sequence. Send n separate requests, or unset the variable."); + } + // prepare child tasks if (task.params.n_cmpl > 1) { int n_children = task.params.n_cmpl - 1; From 35097d84c1897a5869610aa8ab14a13ef9c6c163 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:13:57 +0000 Subject: [PATCH 14/81] server: plan the kv pool in cells rather than tokens A review of #194 predicted that exact mode and #184's preemption planner cannot both be right about how full the pool is, and on this branch they are not. Reproduced with the 4B on one B200, LLAMA_EXACT_CONCURRENCY=1, four chats with 1000-token prompts generating 2048 tokens each at --parallel 4 --kv-unified -c 8192 --spec-type draft-mtp, no forced-park knob, three rounds: 0 of 4 completions, every round, every chat ending in "Context size has been exceeded", with 0 parks and 0 restores. Nothing was ever parked. preempt_kv_used(), preempt_n_need() and preempt_kv_reserve() count tokens, and exact mode's allocator hands out 256-cell pages, one page to one (sequence, position / 256) pair. Four sequences can therefore be holding up to 1020 cells that no other sequence can be given, and the planner, seeing room in tokens that find_slot cannot find in pages, never reaches the threshold that would park anybody. The retry ladder then halves n_batch to 1 and ends every request, which is the pre-#184 behaviour that preemption exists to remove. Ask the memory how it allocates instead of assuming. llama_memory_i gains a non-pure alloc_granularity() defaulting to 1, so no module that allocates a cell per token needs a change; llama_kv_cache returns its page size under exact mode and 1 otherwise, the hybrid memory forwards to its attention half, and llama_memory_alloc_granularity() exposes it. The server reads it once at load and rounds: preempt_kv_used() charges every slot's tail page in full preempt_n_need() rounds what a resume must be given, since a restore takes fresh pages preempt_kv_reserve() reserves the cells the next step ADDS rather than its tokens, because on a rounded used figure a step is free until it crosses a page boundary and costs a whole page when it does, and that crossing is the only moment the pool can run out preempt_n_margin() rounds the asynchronous lookahead up to a page, since any of the steps it covers can cost one With a granularity of 1 every one of these is the arithmetic it was, so nothing changes with the mode off. After, same configuration and three rounds: 4 of 4 completions each round, 2 parks and 2 restores each round, no context errors, and P0 byte-identical to its solo run in all three. The same run with exact mode off is also 4 of 4 with 2 parks, unchanged from the parent. --- include/llama.h | 8 ++++ src/llama-context.cpp | 8 ++++ src/llama-kv-cache.cpp | 7 ++++ src/llama-kv-cache.h | 3 ++ src/llama-memory-hybrid.cpp | 6 +++ src/llama-memory-hybrid.h | 2 + src/llama-memory.h | 9 ++++ tools/server/server-context.cpp | 73 ++++++++++++++++++++++++++++----- 8 files changed, 105 insertions(+), 11 deletions(-) diff --git a/include/llama.h b/include/llama.h index 2c61b3d2a71..b297739d1e4 100644 --- a/include/llama.h +++ b/include/llama.h @@ -795,6 +795,14 @@ extern "C" { // Check if the memory supports shifting LLAMA_API bool llama_memory_can_shift(llama_memory_t mem); + // [TAG_EXACT_CONCURRENCY] Cells the memory allocates in one indivisible unit. + // + // 1 in every ordinary configuration. Larger where a mode places cells in blocks, and + // then a sequence of n tokens occupies round_up(n, granularity) cells. A caller that + // decides whether the pool has room by counting tokens has to round the same way, or it + // will believe there is space that cannot be handed out. + LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem); + // // State / sessions // diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0747d1c6366..ccaff40745d 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4598,6 +4598,14 @@ bool llama_memory_can_shift(llama_memory_t mem) { return mem->get_can_shift(); } +uint32_t llama_memory_alloc_granularity(llama_memory_t mem) { + if (!mem) { + return 1; + } + + return mem->alloc_granularity(); +} + // llama state API // deprecated diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 297786fb889..4df8d3b4ff1 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1252,6 +1252,13 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & } } +uint32_t llama_kv_cache::alloc_granularity() const { + // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so a + // sequence holding n tokens holds round_up(n, exact_page_size) cells: its tail page is + // charged in full whether or not it is full. + return exact_pages ? exact_page_size : 1; +} + bool llama_kv_cache::get_can_shift() const { // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot. if (model.arch == LLM_ARCH_STEP35) { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index fa257422f02..e55c5a23003 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -131,6 +131,9 @@ class llama_kv_cache : public llama_memory_i { bool get_can_shift() const override; + // [TAG_EXACT_CONCURRENCY] the page size under exact mode, 1 otherwise + uint32_t alloc_granularity() const override; + void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index ba54ab12923..5851f72d875 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -142,6 +142,12 @@ bool llama_memory_hybrid::get_can_shift() const { return mem_attn->get_can_shift(); } +uint32_t llama_memory_hybrid::alloc_granularity() const { + // the recurrent half holds one state per sequence rather than per token, so the + // attention half is the one whose cells a caller is planning capacity for + return mem_attn->alloc_granularity(); +} + void llama_memory_hybrid::clear(bool data) { mem_attn->clear(data); mem_recr->clear(data); diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb7499..70ba19ca323 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -58,6 +58,8 @@ class llama_memory_hybrid : public llama_memory_i { bool get_can_shift() const override; + uint32_t alloc_granularity() const override; + void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645..51539a03919 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -100,6 +100,15 @@ struct llama_memory_i { // getters virtual bool get_can_shift() const = 0; + // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit. + // + // 1 for every module that allocates a cell per token, which is all of them unless a mode + // is on that allocates in larger blocks. Where it is larger, a sequence of n tokens + // occupies round_up(n, granularity) cells, and a caller that plans pool capacity by + // counting tokens will believe there is room that does not exist. Not pure, so a module + // that has never heard of this inherits the answer that has always been true of it. + virtual uint32_t alloc_granularity() const { return 1; } + // // ops // diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 999a6e723a9..1a5fdc26934 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1688,6 +1688,18 @@ struct server_context_impl { } } + // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates, rather than assume a cell + // per token. 1 in every ordinary configuration, so this changes nothing unless a + // mode that places cells in blocks is on. + { + preempt_alloc_granularity = (int32_t) std::max(1u, llama_memory_alloc_granularity(llama_get_memory(ctx_tgt))); + + if (preempt_alloc_granularity > 1) { + SRV_INF("preemption: the kv pool allocates %d cells at a time, planning in pages\n", + preempt_alloc_granularity); + } + } + { const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; @@ -3165,6 +3177,28 @@ struct server_context_impl { // case every park and resume is the synchronous one it always was. bool preempt_async_ok = false; + // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time: 1 normally, the page size + // under exact concurrency. Everything below plans in cells rather than tokens because of + // it -- with a page size of 256, four sequences can hold 1020 cells the pool cannot hand + // to anybody else, and a planner counting tokens sees room that find_slot cannot find. + int32_t preempt_alloc_granularity = 1; + + // cells a sequence of n tokens actually occupies + int32_t preempt_n_cells(int32_t n_tokens) const { + const int32_t g = preempt_alloc_granularity; + + if (g <= 1 || n_tokens <= 0) { + return n_tokens; + } + + return ((n_tokens + g - 1) / g) * g; + } + + // cells a slot holding n_tokens has to be given for a step of n_step more + int32_t preempt_n_cells_step(int32_t n_tokens, int32_t n_step) const { + return preempt_n_cells(n_tokens + n_step) - preempt_n_cells(n_tokens); + } + int32_t preempt_n_spec_max() const { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } @@ -3209,7 +3243,10 @@ struct server_context_impl { res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); } - return res; + // [TAG_EXACT_CONCURRENCY] a restore takes fresh pages and its tail page is charged in + // full, so what the pool has to have free for this slot is the rounded figure. Under + // counting here is what admits a resume that find_slot then cannot satisfy. + return preempt_n_cells(res); } // Cells the pool is holding right now. A released slot keeps its prompt in the cache @@ -3229,7 +3266,7 @@ struct server_context_impl { // because the copy is still reading them, and one on its way back in has already // been given them. Skipping either would hand the same cells out twice. - res += slot.prompt.n_tokens(); + res += preempt_n_cells(slot.prompt.n_tokens()); } return res; @@ -3255,7 +3292,12 @@ struct server_context_impl { } } - return PREEMPT_N_MARGIN + n_running * (1 + preempt_n_spec_max()) * PREEMPT_N_ASYNC_STEPS; + // [TAG_EXACT_CONCURRENCY] the lookahead is a number of decode steps, and under a + // page allocator any one of them can cost a whole page rather than a cell. Round the + // runway up to a page so the park is issued with at least one page of real room + // behind it; with a page size of 1 this is the figure it always was. + return preempt_n_cells( + PREEMPT_N_MARGIN + n_running * (1 + preempt_n_spec_max()) * PREEMPT_N_ASYNC_STEPS); } // cells those slots are about to ask for on the next decode @@ -3266,19 +3308,27 @@ struct server_context_impl { int32_t res = 0; int32_t res_pmt = 0; + // [TAG_EXACT_CONCURRENCY] each slot reserves the cells its next step ADDS, not the + // tokens it adds. preempt_kv_used() already charges every slot's tail page in full, + // so with a page size of 1 these are the same number and nothing changes; with a + // larger one the step is free until it crosses a page boundary and costs a whole + // page when it does. Reserving tokens on top of a rounded used figure would miss + // exactly that crossing, which is the only moment the pool can actually run out. for (const auto & slot : slots) { + const int32_t n_cur = slot.prompt.n_tokens(); + switch (slot.state) { case SLOT_STATE_GENERATING: case SLOT_STATE_DONE_PROMPT: { - res += 1 + n_spec; + res += preempt_n_cells_step(n_cur, 1 + n_spec); } break; case SLOT_STATE_STARTED: case SLOT_STATE_PROCESSING_PROMPT: { - const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + const int32_t n_left = slot.task ? slot.task->n_tokens() - n_cur : 0; - res_pmt += std::max(1, std::min(n_batch, n_left)); + res_pmt += preempt_n_cells_step(n_cur, std::max(1, std::min(n_batch, n_left))); } break; case SLOT_STATE_RESTORING: { @@ -3287,11 +3337,11 @@ struct server_context_impl { // take has to be reserved now -- otherwise the pool is handed out from // under it and its first step preempts somebody else straight away if (slot.state_before_preempt == SLOT_STATE_GENERATING) { - res += 1 + n_spec; + res += preempt_n_cells_step(n_cur, 1 + n_spec); } else { - const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + const int32_t n_left = slot.task ? slot.task->n_tokens() - n_cur : 0; - res_pmt += std::max(1, std::min(n_batch, n_left)); + res_pmt += preempt_n_cells_step(n_cur, std::max(1, std::min(n_batch, n_left))); } } break; // a preempting slot is on its way out and will not decode: nothing to reserve @@ -3300,8 +3350,9 @@ struct server_context_impl { } } - // one batch is all the prompt slots get between them, however many are waiting - return res + std::min(res_pmt, n_batch); + // one batch is all the prompt slots get between them, however many are waiting; in + // cells that batch can straddle one boundary more than it has tokens for + return res + std::min(res_pmt, preempt_n_cells(n_batch)); } // Keep the slot that is furthest along -- it is the closest to finishing and to giving From e7e88e9de8bbb7d559bddc8969795480f5674edd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:30:36 +0000 Subject: [PATCH 15/81] ggml: map the event query onto HIP and MUSA ggml-cuda.cu is compiled for ROCm and for MUSA through the vendor headers, which rename every cuda* name it uses. The non-blocking event query added cudaEventQuery and cudaErrorNotReady, and neither header maps them, so both builds stop at an undeclared identifier while the adjacent cudaEventSynchronize has been mapped all along. hipEventQuery and musaEventQuery have the same signature and the same convention: success when everything recorded before the event has finished, hipErrorNotReady or musaErrorNotReady while it has not, which is exactly what the query reads them as. --- ggml/src/ggml-cuda/vendors/hip.h | 2 ++ ggml/src/ggml-cuda/vendors/musa.h | 2 ++ 2 files changed, 4 insertions(+) diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h index 9aa558f3f4c..83d37aeeb2a 100644 --- a/ggml/src/ggml-cuda/vendors/hip.h +++ b/ggml/src/ggml-cuda/vendors/hip.h @@ -58,10 +58,12 @@ #define cudaDeviceSynchronize hipDeviceSynchronize #define cudaError_t hipError_t #define cudaErrorMemoryAllocation hipErrorOutOfMemory +#define cudaErrorNotReady hipErrorNotReady #define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled #define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled #define cudaEventCreateWithFlags hipEventCreateWithFlags #define cudaEventDisableTiming hipEventDisableTiming +#define cudaEventQuery hipEventQuery #define cudaEventRecord hipEventRecord #define cudaEventSynchronize hipEventSynchronize #define cudaEvent_t hipEvent_t diff --git a/ggml/src/ggml-cuda/vendors/musa.h b/ggml/src/ggml-cuda/vendors/musa.h index 6d725c7ec19..ebecf679950 100644 --- a/ggml/src/ggml-cuda/vendors/musa.h +++ b/ggml/src/ggml-cuda/vendors/musa.h @@ -46,10 +46,12 @@ #define cudaDeviceSynchronize musaDeviceSynchronize #define cudaError_t musaError_t #define cudaErrorMemoryAllocation musaErrorMemoryAllocation +#define cudaErrorNotReady musaErrorNotReady #define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled #define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled #define cudaEventCreateWithFlags musaEventCreateWithFlags #define cudaEventDisableTiming musaEventDisableTiming +#define cudaEventQuery musaEventQuery #define cudaEventRecord musaEventRecord #define cudaEventSynchronize musaEventSynchronize #define cudaEvent_t musaEvent_t From d33f04fd639b8209a0d6b94d97803b4efc391591 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:33:35 +0000 Subject: [PATCH 16/81] llama: require a real event query before copying a sequence asynchronously state_seq_copy_init() took any device advertising async and events, but ggml_backend_event_query() is optional: a device that does not implement it gets the generic fallback, which answers "is it done" by waiting for it. Metal, Vulkan and SYCL all advertise both capabilities and all leave event_query null, so they were handed a transfer object, told the caller the copies were asynchronous, and then blocked it for the whole copy on its first poll. That is the stall the transfer exists to remove, made worse by the caller no longer expecting it. ggml_backend_dev_supports_event_query() is the question the fallback hides, and state_seq_copy_init() now asks it. A device without a query is left out, so those backends get NULL and keep the synchronous llama_state_seq_*_data_ext calls they always used, which is the documented behaviour and is what the server already falls back to. The reason is logged once. --- ggml/include/ggml-backend.h | 6 +++++- ggml/src/ggml-backend.cpp | 5 +++++ include/llama.h | 5 +++-- src/llama-context.cpp | 19 +++++++++++++++++++ 4 files changed, 32 insertions(+), 3 deletions(-) diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 30d8304d492..d21bf40dd58 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -126,7 +126,8 @@ extern "C" { GGML_API void ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backend); GGML_API void ggml_backend_event_synchronize(ggml_backend_event_t event); // non-blocking: true once everything recorded before the event has completed. - // backends without a query implementation fall back to a blocking synchronize and return true. + // backends without a query implementation fall back to a blocking synchronize and return true, + // which ggml_backend_dev_supports_event_query() tells apart from a real non-blocking query. GGML_API bool ggml_backend_event_query(ggml_backend_event_t event); GGML_API void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event); @@ -193,6 +194,9 @@ extern "C" { GGML_API ggml_backend_buffer_t ggml_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void * ptr, size_t size, size_t max_tensor_size); GGML_API bool ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_tensor * op); + // whether ggml_backend_event_query() on this device really is non-blocking, i.e. whether + // the device implements it rather than falling back to a blocking synchronize + GGML_API bool ggml_backend_dev_supports_event_query(ggml_backend_dev_t device); GGML_API bool ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft); GGML_API bool ggml_backend_dev_offload_op(ggml_backend_dev_t device, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index a56ab30862a..b13d9c811c4 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -639,6 +639,11 @@ bool ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_t return device->iface.supports_op(device, op); } +bool ggml_backend_dev_supports_event_query(ggml_backend_dev_t device) { + GGML_ASSERT(device); + return device->iface.event_query != NULL; +} + bool ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) { GGML_ASSERT(device); return device->iface.supports_buft(device, buft); diff --git a/include/llama.h b/include/llama.h index 2c61b3d2a71..26db483e6c2 100644 --- a/include/llama.h +++ b/include/llama.h @@ -944,8 +944,9 @@ extern "C" { // written. llama_state_seq_copy_free() waits for an outstanding copy first. struct llama_state_seq_copy; - // NULL if the context's backends cannot copy asynchronously; the caller then uses the - // synchronous llama_state_seq_*_data_ext calls + // NULL if the context's backends cannot copy asynchronously, or cannot say whether a + // copy has finished without waiting for it, which would put the stall straight back; the + // caller then uses the synchronous llama_state_seq_*_data_ext calls LLAMA_API struct llama_state_seq_copy * llama_state_seq_copy_init(struct llama_context * ctx); LLAMA_API void llama_state_seq_copy_free(struct llama_state_seq_copy * cpy); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0747d1c6366..d2e929a1429 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -13,6 +13,7 @@ #include "llama-sampler.h" #include "llama.h" +#include #include #include #include @@ -3549,6 +3550,24 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } + // A device that advertises events but does not implement event_query is no use + // here. ggml_backend_event_query() then answers the only way it can, by waiting for + // the event, so the first poll of a transfer blocks the caller for the whole copy -- + // the very stall this exists to remove, except that the caller has been told the + // copy is asynchronous and has stopped looking for it. Such a device is left out, so + // that state_seq_copy_init() returns NULL and the caller keeps the synchronous calls + // it already had. + if (!ggml_backend_dev_supports_event_query(dev)) { + static std::atomic warned(false); + + if (!warned.exchange(true)) { + LLAMA_LOG_INFO("%s: %s cannot test an event without waiting for it, so sequence " + "states are copied synchronously\n", __func__, ggml_backend_dev_name(dev)); + } + + continue; + } + // a backend of its own, not the one the graphs are computed on: that one moves its // copies to whichever stream it is currently using, so a transfer posted to it could // end up ordered behind a graph -- which is the stall this exists to avoid From 0160ea46748249c54be27bc6c0d4cd4a2421d4cc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:36:42 +0000 Subject: [PATCH 17/81] server: wait for the parks still in flight before tearing the contexts down destroy() resets llama_init and nulls ctx_tgt and ctx_dft, but the slots are declared after llama_init and are still alive at that point, and one of them can be holding a park or a resume that is still reading or writing KV tensors of the context being freed. release() already makes that wait for a single slot, on the path a cancelled request takes; nothing made it for all of them. The sleeping-state path is where it shows: /sleep calls destroy() and the server carries on running, so a copy issued an iteration earlier is left pointing at freed tensors and load_model() then clears the slots, running the transfer destructor's own wait against the same memory. Shutdown has the same hole with less time to notice it. destroy() now waits for every slot's outstanding copy and lets go of the transfers before anything is freed, which also means the next context does not inherit a backend and a host buffer belonging to the previous one. --- tools/server/server-context.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f15f78f8201..e31a96f427f 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1272,6 +1272,23 @@ struct server_context_impl { int64_t t_last_load_progress_ms = 0; void destroy() { + // [TAG_PREEMPT_ASYNC] the slots outlive this call -- they are declared after + // llama_init, so they are still there when it is reset here, and load_model() clears + // them only after the next context exists -- and any one of them may be holding a + // park or a resume that is still reading or writing KV tensors of the contexts about + // to be freed. release() makes the same wait for a single slot; this is the one that + // covers all of them, and on the sleeping-state path it is the only one there is, + // because the server carries on running afterwards. + for (auto & slot : slots) { + slot.preempt_copy_wait(); + + // the transfer holds a backend and an event of its own, and its host buffer is + // no use to the context that comes back: let go of both before that context's + // successor makes new ones + slot.preempt_cpy_tgt.reset(); + slot.preempt_cpy_dft.reset(); + } + spec.reset(); spec_init.reset(); From c0d92970dd208439ccd97fad226e20250c5c04c8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:36:57 +0000 Subject: [PATCH 18/81] server: charge a resume candidate its own lookahead before admitting it preempt_n_margin() keeps eight decode steps of every running slot clear ahead of the pool filling, and the resume gate uses the same figure so that a slot is not put back into a pool it would immediately have to leave. It was not doing that for the slot being resumed. The candidate is still PREEMPTED while it is being considered, so the loop that counts running slots skips it, and the runway it needs appears only after it has been let in, at which point the pool is short by exactly that much and somebody gets parked. At -c 256 with a 1 + n_spec step and the eight-step runway, totals from 233 to 240 cells admit a restore that then cannot take its first step, and under load the same slot was seen restored and parked again five times over. preempt_n_margin() takes the number of slots that are about to be running as well as those that already are, and the resume gate passes one for the candidate. Everything else keeps the count it had. preempt_kv_reserve() already reserves a restoring slot's next step; this is the eight-step runway behind it. --- tools/server/server-context.cpp | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e31a96f427f..1effa101417 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3226,12 +3226,19 @@ struct server_context_impl { // its copy lands, so it has to fire early enough that everything still decoding has // somewhere to put its tokens until then. The same figure gates a resume, so that a slot // is not put back into a pool it would immediately have to be taken out of again. - int32_t preempt_n_margin() const { + // + // n_additional_running is for slots that are not running yet but are about to be: a + // resume candidate is still PREEMPTED while it is being considered, so it is not counted + // by the loop below, yet the moment it is admitted it starts decoding and needs the same + // runway as everybody else. Admitting it without charging it that runway is what the + // margin exists to prevent, and it showed up as a slot restored and parked again a few + // iterations later, over and over. + int32_t preempt_n_margin(int32_t n_additional_running = 0) const { if (!preempt_async_active()) { return PREEMPT_N_MARGIN; } - int32_t n_running = 0; + int32_t n_running = n_additional_running; for (const auto & slot : slots) { if (slot.is_processing() && !slot.preempt_is_out()) { @@ -3457,12 +3464,14 @@ struct server_context_impl { server_slot * best = nullptr; // Room for the sequence AND for the next step of everything already running, - // so that a resume cannot immediately trigger the preemption of someone else. + // and for the lookahead of the candidate itself, which is about to become one of + // them: a resume must not immediately trigger the preemption of someone else, or + // of itself. // A cached prompt on an idle slot is worth less than a conversation waiting to // continue, so give those cells up first - same call the KV-full path makes. for (;;) { for (auto * slot : parked) { - if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + preempt_n_margin() <= n_cells) { + if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + preempt_n_margin(1) <= n_cells) { best = slot; break; } From cd54f6089240237cce65c030d3555dc7848d91c4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:37:18 +0000 Subject: [PATCH 19/81] llama: report the host memory a transfer holds, not the kind it asked for llama_state_seq_copy_buf_is_pinned() returned can_pin, which is worked out from the buffer type the backend offers and is fixed for the life of the transfer. The header promises the buffer is page-locked. Those are different questions: the CUDA host buffer type is handed out whether or not pinning is available, and under GGML_CUDA_NO_PINNED its allocation falls back to an ordinary CPU buffer, so the server logged "pinned host memory" while every park ran through pageable memory. It was also true before any buffer existed and after buf_free(). buf_resize() already records which it got, by comparing the buffer that came back against the type that was asked for, so is_pinned() now returns that. llama_state_seq_copy_buf_can_pin() is the capability question, for a caller that wants to know before allocating anything. The load banner asked the capability question at a point where no buffer exists and printed the answer as though one did. It now says what the backend offers, in those words, and the first park reports what the buffer it allocated actually turned out to be. --- include/llama.h | 8 +++++++- src/llama-context.cpp | 7 +++++++ tools/server/server-context.cpp | 36 +++++++++++++++++++++++++++++++-- 3 files changed, 48 insertions(+), 3 deletions(-) diff --git a/include/llama.h b/include/llama.h index 26db483e6c2..a5be59cb8f8 100644 --- a/include/llama.h +++ b/include/llama.h @@ -960,9 +960,15 @@ extern "C" { LLAMA_API size_t llama_state_seq_copy_buf_capacity(struct llama_state_seq_copy * cpy); LLAMA_API void llama_state_seq_copy_buf_free (struct llama_state_seq_copy * cpy); - // true when the buffer is page-locked, i.e. when the copies can really overlap + // true when the buffer that is held right now is page-locked, i.e. when the copies can + // really overlap. False while no buffer is held, since none is page-locked then: a + // caller asking before the first resize wants llama_state_seq_copy_buf_can_pin(). LLAMA_API bool llama_state_seq_copy_buf_is_pinned(struct llama_state_seq_copy * cpy); + // true when the backend offers pinned host memory at all. It is what the next resize + // will ask for, not what any buffer is: an allocation can still come back pageable. + LLAMA_API bool llama_state_seq_copy_buf_can_pin(struct llama_state_seq_copy * cpy); + // issue the copies; return the number of bytes covered, 0 on failure LLAMA_API size_t llama_state_seq_copy_get( struct llama_state_seq_copy * cpy, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index d2e929a1429..a807944b4b5 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4751,6 +4751,13 @@ void llama_state_seq_copy_buf_free(llama_state_seq_copy * cpy) { } bool llama_state_seq_copy_buf_is_pinned(llama_state_seq_copy * cpy) { + // what was allocated, not what could be: a host buffer type is free to hand back + // ordinary memory, which is what CUDA does under GGML_CUDA_NO_PINNED, and there is + // nothing page-locked before the first resize or after buf_free() + return cpy->pinned; +} + +bool llama_state_seq_copy_buf_can_pin(llama_state_seq_copy * cpy) { return cpy->can_pin; } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1effa101417..007f3298cb0 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1289,6 +1289,10 @@ struct server_context_impl { slot.preempt_cpy_dft.reset(); } + // the next context allocates its own host buffers, so say again what they turn out + // to be + preempt_ram_kind_logged = false; + spec.reset(); spec_init.reset(); @@ -1676,8 +1680,11 @@ struct server_context_impl { if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { if (preempt_async_ok) { - SRV_INF("preemption: parking and resuming asynchronously, %s host memory\n", - llama_state_seq_copy_buf_is_pinned(slots[0].preempt_cpy_tgt.get()) ? "pinned" : "pageable"); + // no buffer has been allocated yet, so this is what the backend offers, + // not what is held. What was actually got is reported by the first park, + // because a host buffer type may still hand back ordinary memory. + SRV_INF("preemption: parking and resuming asynchronously, backend offers %s host memory\n", + llama_state_seq_copy_buf_can_pin(slots[0].preempt_cpy_tgt.get()) ? "pinned" : "pageable"); } else { SRV_WRN("%s", "preemption: this backend cannot copy asynchronously, parking and resuming synchronously\n"); } @@ -3153,6 +3160,27 @@ struct server_context_impl { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } + // [TAG_PREEMPT_ASYNC] whether the kind of host memory the parks actually got has been + // reported. It is only knowable once a buffer exists, and it is worth knowing: pinned + // memory is what lets the copies overlap, and a host buffer type is free to hand back + // ordinary memory instead of failing. + bool preempt_ram_kind_logged = false; + + void preempt_log_ram_kind(const server_slot & slot) { + if (preempt_ram_kind_logged || !slot.preempt_is_async()) { + return; + } + + if (llama_state_seq_copy_buf_capacity(slot.preempt_cpy_tgt.get()) == 0) { + return; // nothing held, so nothing to report yet + } + + preempt_ram_kind_logged = true; + + SRV_INF("preemption: parking into %s host memory\n", + llama_state_seq_copy_buf_is_pinned(slot.preempt_cpy_tgt.get()) ? "pinned" : "pageable"); + } + // host RAM the parked sequences hold right now size_t preempt_ram_used() const { size_t res = 0; @@ -3543,6 +3571,8 @@ struct server_context_impl { slot.t_preempt_copy_us = ggml_time_us(); if (slot.preempt_save()) { + preempt_log_ram_kind(slot); + // [TAG_PREEMPT_ASYNC] a slot left PREEMPTING is counted by // update_preempt_copies() when its copy lands, not here if (slot.state == SLOT_STATE_PREEMPTED) { @@ -3595,6 +3625,8 @@ struct server_context_impl { break; // could not park it; the existing retry ladder is still behind us } + preempt_log_ram_kind(*victim); + // [TAG_PREEMPT_ASYNC] the copy has only been issued; the cells are still the // victim's until it lands, so nothing further can be decided about the pool this // iteration. update_preempt_copies() picks it up on the next one, and the step From 22c90bd9ffcebf70f66bb9f865711e330fa77a8c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:40:27 +0000 Subject: [PATCH 20/81] llama: check the size and the flags a sequence transfer is issued with Both issue functions validated only that a buffer existed. The caller's size was handed straight to the io object, which then validated every fragment against that number rather than against the allocation, so a save issued with a size larger than the buffer wrote past the end of it and a restore read whatever was next on the heap and sent it to the device. Unlike the legacy API the library owns this buffer, so it can simply check: a size of zero, or one beyond llama_state_seq_copy_buf_size(), is refused with a log line. The flags word had the same problem from the other end. LLAMA_STATE_SEQ_FLAGS_ON_DEVICE asks for the tensor data to stay in device buffers, and both functions built the host serializers regardless, while llama_state_seq_get_size_ext() with that flag reports a state without the tensor bytes in it. A caller pairing the documented size call with these ones sized a buffer for the metadata and then tried to fill it with the whole sequence. The flag is refused here and the restriction is written down in llama.h; the synchronous calls still serve it. tests/test-state-seq-copy.cpp covers both refusals in both directions, checks that a refused call posts nothing, that the same call at the buffer's own size still round-trips the sequence byte-for-byte, and that a transfer reports itself as pinned only while it holds memory that is. It skips itself where no backend can copy asynchronously. --- include/llama.h | 6 +- src/llama-context.cpp | 34 +++++++- tests/CMakeLists.txt | 5 ++ tests/test-state-seq-copy.cpp | 146 ++++++++++++++++++++++++++++++++++ 4 files changed, 188 insertions(+), 3 deletions(-) create mode 100644 tests/test-state-seq-copy.cpp diff --git a/include/llama.h b/include/llama.h index a5be59cb8f8..3c64888d25f 100644 --- a/include/llama.h +++ b/include/llama.h @@ -969,7 +969,11 @@ extern "C" { // will ask for, not what any buffer is: an allocation can still come back pageable. LLAMA_API bool llama_state_seq_copy_buf_can_pin(struct llama_state_seq_copy * cpy); - // issue the copies; return the number of bytes covered, 0 on failure + // Issue the copies; return the number of bytes covered, 0 on failure. size must be + // between 1 and llama_state_seq_copy_buf_size(): the buffer belongs to the transfer, and + // a size beyond it is refused rather than believed. LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is + // refused too, since these copies serialise through host memory; use + // llama_state_seq_get_data_ext / set_data_ext for that flag. LLAMA_API size_t llama_state_seq_copy_get( struct llama_state_seq_copy * cpy, size_t size, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a807944b4b5..09e53123da6 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3600,7 +3600,22 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { } size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { - if (!cpy.data) { + // Unlike the legacy API the library owns this buffer, so the extent the io object is + // built with can be checked instead of believed. Every bounds check inside that object + // validates against the extent it was given, so a size larger than the allocation makes + // all of them agree with the caller and the copy runs past the buffer. + if (!cpy.data || size == 0 || size > cpy.size) { + LLAMA_LOG_ERROR("%s: cannot cover %zu bytes, the transfer's buffer holds %zu\n", __func__, size, cpy.size); + return 0; + } + + // LLAMA_STATE_SEQ_FLAGS_ON_DEVICE asks for the tensor data to be left in device buffers, + // and this path has nowhere to leave it: it serialises through the host buffer it owns, + // which is the whole point of it. llama_state_seq_get_size_ext() with that flag reports + // a metadata-sized state, so a caller pairing the two would size a buffer for one thing + // and fill it with another; the synchronous calls serve that flag. + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + LLAMA_LOG_ERROR("%s: LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is not supported here, the copies go through host memory\n", __func__); return 0; } @@ -3629,7 +3644,22 @@ size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size } size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { - if (!cpy.data) { + // Unlike the legacy API the library owns this buffer, so the extent the io object is + // built with can be checked instead of believed. Every bounds check inside that object + // validates against the extent it was given, so a size larger than the allocation makes + // all of them agree with the caller and the copy runs past the buffer. + if (!cpy.data || size == 0 || size > cpy.size) { + LLAMA_LOG_ERROR("%s: cannot cover %zu bytes, the transfer's buffer holds %zu\n", __func__, size, cpy.size); + return 0; + } + + // LLAMA_STATE_SEQ_FLAGS_ON_DEVICE asks for the tensor data to be left in device buffers, + // and this path has nowhere to leave it: it serialises through the host buffer it owns, + // which is the whole point of it. llama_state_seq_get_size_ext() with that flag reports + // a metadata-sized state, so a caller pairing the two would size a buffer for one thing + // and fill it with another; the synchronous calls serve that flag. + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + LLAMA_LOG_ERROR("%s: LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is not supported here, the copies go through host memory\n", __func__); return 0; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b9f9d4b78af..55f46d7c6b1 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -299,6 +299,11 @@ llama_build_and_test(test-backend-sampler.cpp LABEL "model") llama_build_and_test(test-state-restore-fragmented.cpp LABEL "model" ARGS -m "${MODEL_DEST}") set_tests_properties(test-state-restore-fragmented PROPERTIES FIXTURES_REQUIRED test-download-model) +# Guards on the asynchronous per-sequence state transfer +# Skips itself on a backend that cannot copy asynchronously +llama_build_and_test(test-state-seq-copy.cpp LABEL "model" ARGS -m "${MODEL_DEST}") +set_tests_properties(test-state-seq-copy PROPERTIES FIXTURES_REQUIRED test-download-model) + # Test state save/load functionality llama_build_and_test(test-save-load-state.cpp LABEL "model" ARGS -m "${MODEL_DEST}") set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED test-download-model) diff --git a/tests/test-state-seq-copy.cpp b/tests/test-state-seq-copy.cpp new file mode 100644 index 00000000000..a4633954cee --- /dev/null +++ b/tests/test-state-seq-copy.cpp @@ -0,0 +1,146 @@ +// [TAG_STATE_ASYNC] guards on the asynchronous per-sequence state transfer +// +// llama_state_seq_copy_get / _set take a size and a flags word from the caller and hand both +// to an io object that validates everything else against them. The buffer belongs to the +// transfer, so a size larger than it is refused rather than believed, and +// LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is refused because these copies serialise through host +// memory. This also checks that the buffer reports itself as page-locked only while it holds +// memory that is. +// +// Skipped, not failed, on a backend that cannot copy asynchronously: there is no transfer to +// make and the synchronous calls are what a caller uses there. + +#include "arg.h" +#include "common.h" +#include "llama.h" + +#include +#include +#include + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + fprintf(stderr, "%s : FAILED at line %d: %s\n", __func__, \ + __LINE__, #cond); \ + return 1; \ + } \ + } while (0) + +int main(int argc, char ** argv) { + common_params params; + + params.sampling.seed = 1234; + params.kv_unified = true; + params.n_parallel = 2; + params.n_ctx = 256; + + common_init(); + + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + ggml_backend_load_all(); + + common_init_result_ptr llama_init = common_init_from_params(params); + + llama_context * ctx = llama_init->context(); + + if (llama_init->model() == nullptr || ctx == nullptr) { + fprintf(stderr, "%s : failed to init\n", __func__); + return 1; + } + + // put something in the cache to copy: two sequences interleaved, so the cells of each + // are a comb rather than one block, which is what the transfer is built for + std::vector tokens(60, 1); + + llama_batch batch = llama_batch_init(params.n_parallel*tokens.size(), 0, 1); + for (size_t i = 0; i < tokens.size(); i++) { + for (int s = 0; s < params.n_parallel; ++s) { + common_batch_add(batch, tokens[i], i, {s}, false); + } + } + batch.logits[batch.n_tokens - 1] = true; + + if (llama_decode(ctx, batch)) { + fprintf(stderr, "%s : failed to decode\n", __func__); + llama_batch_free(batch); + return 1; + } + + llama_batch_free(batch); + + llama_state_seq_copy * cpy = llama_state_seq_copy_init(ctx); + + if (cpy == nullptr) { + fprintf(stderr, "%s : this backend cannot copy sequence states asynchronously, skipping\n", __func__); + return 0; + } + + const int seq_id = 1; + const size_t size = llama_state_seq_get_size_ext(ctx, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE); + + CHECK(size > 0); + + // nothing is allocated yet, so nothing is page-locked yet, whatever the backend offers + CHECK(llama_state_seq_copy_buf_is_pinned(cpy) == false); + CHECK(llama_state_seq_copy_buf(cpy) == nullptr); + + CHECK(llama_state_seq_copy_buf_resize(cpy, size) != nullptr); + CHECK(llama_state_seq_copy_buf_size(cpy) == size); + + fprintf(stderr, "%s : seq %d state is %zu bytes, %s host memory (backend offers %s)\n", + __func__, seq_id, size, + llama_state_seq_copy_buf_is_pinned(cpy) ? "pinned" : "pageable", + llama_state_seq_copy_buf_can_pin(cpy) ? "pinned" : "pageable"); + + // a size beyond the buffer the transfer owns is refused, on both directions + CHECK(llama_state_seq_copy_get(cpy, size + 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + CHECK(llama_state_seq_copy_set(cpy, size + 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + + // and so is an empty one, which cannot even hold the header + CHECK(llama_state_seq_copy_get(cpy, 0, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + CHECK(llama_state_seq_copy_set(cpy, 0, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + + // ON_DEVICE keeps the tensor data off the host, which is where these copies go + CHECK(llama_state_seq_copy_get(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) == 0); + CHECK(llama_state_seq_copy_set(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) == 0); + + // none of that may have posted anything + CHECK(llama_state_seq_copy_done(cpy)); + + fprintf(stderr, "%s : oversized, empty and ON_DEVICE transfers are all refused\n", __func__); + + // the same call at the size the transfer does own still works, and round-trips + std::vector before(llama_state_seq_get_size(ctx, seq_id)); + CHECK(llama_state_seq_get_data(ctx, before.data(), before.size(), seq_id) == before.size()); + + CHECK(llama_state_seq_copy_get(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == size); + llama_state_seq_copy_wait(cpy); + + llama_memory_seq_rm(llama_get_memory(ctx), seq_id, -1, -1); + + CHECK(llama_state_seq_copy_set(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == size); + llama_state_seq_copy_wait(cpy); + + std::vector after(llama_state_seq_get_size(ctx, seq_id)); + CHECK(after.size() == before.size()); + CHECK(llama_state_seq_get_data(ctx, after.data(), after.size(), seq_id) == after.size()); + CHECK(before == after); + + fprintf(stderr, "%s : a transfer at the buffer's own size round-trips seq %d byte-for-byte\n", + __func__, seq_id); + + // giving the memory back leaves nothing page-locked to report + llama_state_seq_copy_buf_free(cpy); + CHECK(llama_state_seq_copy_buf_is_pinned(cpy) == false); + CHECK(llama_state_seq_copy_buf_capacity(cpy) == 0); + + llama_state_seq_copy_free(cpy); + + fprintf(stderr, "%s : SUCCESS\n", __func__); + + return 0; +} From e8f8b2fcf1c85ba16777c3ccd1a851f0fffcee1e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:41:15 +0000 Subject: [PATCH 21/81] ggml: stop collecting an error the CUDA event query never sets The cudaErrorNotReady branch of the CUDA event query called cudaGetLastError() on the belief that the result had to be cleared. It does not: cudaEventQuery() returns cudaErrorNotReady as its return value without recording it in the thread's last-error state, so the only thing that call can collect is an error somebody else planted and has not looked at yet. Checked on a B200 with CUDA 13.1. An unrelated cudaSetDevice(99) leaves 101 pending; cudaEventQuery() on an outstanding event returns 600 and cudaPeekAtLastError() still reads 101 afterwards, so the cudaGetLastError() returned 101 and left the state clean. A real launch failure would have been thrown away the same way, and its owner would never have seen it. --- ggml/src/ggml-cuda/ggml-cuda.cu | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 91902099011..1fca6352403 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5380,9 +5380,10 @@ static bool ggml_backend_cuda_device_event_query(ggml_backend_dev_t dev, ggml_ba const cudaError_t err = cudaEventQuery((cudaEvent_t)event->context); + // not an error, and nothing to clear: cudaEventQuery() returns cudaErrorNotReady + // without recording it as the thread's last error, so collecting one here would only + // consume somebody else's, and a real launch failure would be swallowed if (err == cudaErrorNotReady) { - // not an error: clear it so it is not reported against the next call - (void) cudaGetLastError(); return false; } From 86315eac6bc5a110a2df4e7f927b39dc7a9a6c25 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:41:15 +0000 Subject: [PATCH 22/81] ggml: bump the backend API version for the new device interface member ggml_backend_device_i gained event_query, so a device interface built against the previous header is one member shorter than the one ggml now reads. Every in-tree initializer was updated, but a backend loaded from a shared library is not: ggml_backend_reg_load_backend() accepts it on api_version alone, and a prebuilt .so still reporting 2 would have been let in and its iface.event_query read past the end of the object. Rejecting it is what the version is for. --- ggml/src/ggml-backend-impl.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 902b0963afa..bb3be31217b 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -8,7 +8,7 @@ extern "C" { #endif - #define GGML_BACKEND_API_VERSION 2 + #define GGML_BACKEND_API_VERSION 3 // // Backend buffer type From 3cc003acf55f8d8195e50f805dede52fbb808051 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:41:33 +0000 Subject: [PATCH 23/81] kv-cache: refuse exact mode when a KV layer is not on the CUDA backend The mode was gated on unified && offload && !v_trans && n_swa == 0 and never on where the attention layers actually run. Only the CUDA FLASH_ATTN_EXT reads src[5]; the CPU, Metal, Vulkan, SYCL, OpenCL and CANN kernels ignore it. So with -ngl 0, a partial -ngl, or a non-CUDA GPU the pool was still paged and the page table was still attached, but attention traversed physical cell order: the output stayed correct and neighbour and relocation independence were silently lost while the mode reported itself as on. Check the placement where the cache is built instead. Every KV layer must be offloaded and its device must belong to the CUDA family backend, which is also built as ROCm and MUSA and carries the same paged kernel; anything else fails the load naming the layer and the backend it landed on. As a second line, the FLASH_ATTN_EXT of every backend that would ignore the page table now refuses an op with src[5] set, so a scheduler decision made after the load cannot route it somewhere that walks the pool in physical order. The CPU is deliberately left accepting it and says so in place: it is the reference test-backend-ops compares the paged CUDA kernel against, and that test builds a mask which selects exactly the listed cells. The four remaining preconditions were one bare GGML_ASSERT each, so a quantized KV cache, an SWA model, a transposed V cache or a -c that is not a multiple of 256 aborted at model load without naming which one failed. Each now logs what it needs and which flag sets it, and the load returns an error the way every other KV cache failure does. The context size check also moved after the shared-source override, so it tests the size the cache is actually built at. -ngl 10 with LLAMA_EXACT_CONCURRENCY=1: llama_kv_cache: LLAMA_EXACT_CONCURRENCY is set but layer 0 keeps its KV cache on CPU, which has no paged attention: every layer must be offloaded to the CUDA backend (pass -ngl to offload all layers and do not pass --no-kv-offload) --- ggml/src/ggml-cann/ggml-cann.cpp | 5 ++ ggml/src/ggml-cpu/ggml-cpu.cpp | 7 +++ ggml/src/ggml-metal/ggml-metal-device.m | 5 ++ ggml/src/ggml-opencl/ggml-opencl.cpp | 5 ++ ggml/src/ggml-sycl/ggml-sycl.cpp | 4 +- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 5 ++ src/llama-kv-cache.cpp | 62 +++++++++++++++++++++++-- 7 files changed, 87 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index 5e5541aac94..0e901ed0160 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2656,6 +2656,11 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten return true; case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only + // the CUDA backend reads + if (op->src[5]) { + return false; + } #ifdef ASCEND_310P // FA not support on 310p device return false; diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 8cece71f186..a548b33bd71 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -474,6 +474,13 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return ggml_is_contiguous(op->src[0]); case GGML_OP_SSM_SCAN: return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1; + // [TAG_EXACT_CONCURRENCY] note: GGML_OP_FLASH_ATTN_EXT with src[5] set, the + // exact-concurrency page table, is deliberately still accepted here. The CPU ignores the + // page table and attends in physical cell order, which is why every other backend refuses + // it, but the CPU is also the reference that test-backend-ops compares the paged CUDA + // kernel against, and that test builds a mask which selects exactly the listed cells. A KV + // cache layer cannot reach the CPU under the mode anyway: llama_kv_cache refuses to + // construct unless every KV layer is on the CUDA backend. default: return true; } diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 19c57820e85..90873f2fab0 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1592,6 +1592,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_ROLL: return ggml_is_contiguous(op->src[0]); case GGML_OP_FLASH_ATTN_EXT: + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the + // CUDA backend reads; walking the pool in physical order here would be silently wrong + if (op->src[5] != NULL) { + return false; + } // for new head sizes, add checks here if (op->src[0]->ne[0] != 32 && op->src[0]->ne[0] != 40 && diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 64f3325b2a5..effd11714f9 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7842,6 +7842,11 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te case GGML_OP_MEAN: return op->src[0]->type == GGML_TYPE_F32; case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the + // CUDA backend reads + if (op->src[5]) { + return false; + } // The E17 compilers segfault while building FA kernels, skip E17 for now if (adreno_e17_compiler_quirks(backend_ctx)) { return false; diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 0573643d834..69a344ab790 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6342,7 +6342,9 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_SOLVE_TRI: return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K; case GGML_OP_FLASH_ATTN_EXT: - return ggml_sycl_flash_attn_ext_supported(device, op); + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the + // CUDA backend reads + return op->src[5] == nullptr && ggml_sycl_flash_attn_ext_supported(device, op); default: return false; } diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c1d86aaac5c..4a2347219b3 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -18192,6 +18192,11 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only + // the CUDA backend reads + if (op->src[5]) { + return false; + } bool coopmat2 = device->coopmat2; uint32_t HSK = op->src[1]->ne[0]; uint32_t HSV = op->src[2]->ne[0]; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index df643a047a9..dcaf7bf687c 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -61,6 +61,29 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { // llama_kv_cache // +// [TAG_EXACT_CONCURRENCY] +// The paged attention specialization that reads the logical page table lives in the CUDA backend +// sources, which are also built as the ROCm and MUSA backends. Every other backend ignores src[5] +// and walks the pool in physical cell order, so a KV layer placed there would silently lose the +// guarantee the mode exists to provide. +static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) { + if (!dev) { + return false; + } + + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + if (!reg) { + return false; + } + + const char * name = ggml_backend_reg_name(reg); + if (!name) { + return false; + } + + return strcmp(name, "CUDA") == 0 || strcmp(name, "ROCm") == 0 || strcmp(name, "MUSA") == 0; +} + llama_kv_cache::llama_kv_cache( const llama_model & model, const llama_hparams & hparams, @@ -86,11 +109,6 @@ llama_kv_cache::llama_kv_cache( const char * exact_env = getenv("LLAMA_EXACT_CONCURRENCY"); exact_pages = exact_env && atoi(exact_env) != 0; - if (exact_pages) { - GGML_ASSERT(unified && offload && !v_trans && n_swa == 0); - GGML_ASSERT(type_k == GGML_TYPE_F16 && type_v == GGML_TYPE_F16); - GGML_ASSERT(kv_size % exact_page_size == 0); - } // shared cells view the source cache's K/V tensors, so the cell count // follows the source allocation: a fitted target can be smaller than the @@ -105,6 +123,30 @@ llama_kv_cache::llama_kv_cache( GGML_ASSERT(kv_size % n_pad == 0); + // [TAG_EXACT_CONCURRENCY] + // Every one of these is reachable from the command line, so report which one failed by name + // instead of aborting on a bare assert that only prints a file and a line. + if (exact_pages) { + const char * unsupported = nullptr; + + if (!unified) { + unsupported = "it needs a unified KV cache (pass --kv-unified)"; + } else if (v_trans) { + unsupported = "it needs a non-transposed V cache (pass --flash-attn on)"; + } else if (n_swa != 0) { + unsupported = "the paged pool does not support sliding window attention"; + } else if (type_k != GGML_TYPE_F16 || type_v != GGML_TYPE_F16) { + unsupported = "it needs an F16 KV cache (do not pass --cache-type-k or --cache-type-v)"; + } else if (kv_size % exact_page_size != 0) { + unsupported = "the context size must be a multiple of 256 (pass -c as a multiple of 256)"; + } + + if (unsupported) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but %s\n", __func__, unsupported); + throw std::runtime_error("exact concurrency: unsupported KV cache configuration"); + } + } + const uint32_t n_layer = hparams.n_layer_all; // define a comparator for the buft -> ctx map to ensure that the order is well-defined: @@ -228,6 +270,16 @@ llama_kv_cache::llama_kv_cache( LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); + // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order while + // the mode still reports itself as on, so refuse the load instead + if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, " + "which has no paged attention: every layer must be offloaded to the CUDA backend " + "(pass -ngl to offload all layers and do not pass --no-kv-offload)\n", + __func__, il, dev_name); + throw std::runtime_error("exact concurrency: KV cache layer is not on the CUDA backend"); + } + ggml_context * ctx = ctx_for_buft(buft); if (!ctx) { throw std::runtime_error("failed to create ggml context for kv cache"); From 66583931d928a472dcbbdc6a6c94a6e8fe5b52e2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:43:01 +0000 Subject: [PATCH 24/81] kv-cache: report the cache transformations exact mode cannot do get_can_shift() still returned true under exact_pages, so --context-shift and --cache-reuse N passed every startup capability gate and then aborted the whole process on the first seq_add with a nonzero shift. Both are opt-in flags, so the default was safe, but llama-server accepted them silently and died on the first request that needed them. The server already disables both at load for a cache that cannot shift, with a warning, so returning false there reuses that path: srv load_model: ctx_shift is not supported by this context, it will be disabled The remaining aborts are reachable the same way, from one request parameter or one API call, and abort() is not an acceptable answer to either in a network server. seq_cp between two different sequences, a nonzero seq_add and a seq_div now log which transformation was refused and on which sequence and return without touching the cells, and the whole-context branch of state_read_meta() logs and returns false the way the two failure paths next to it already do, so llama_state_load_file() reports a recoverable error through an API that is designed for one instead of killing the process. The page invariant is protected exactly as before: none of these paths can now run and leave a cell outside the page its position belongs to. --- src/llama-kv-cache.cpp | 42 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 38 insertions(+), 4 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index dcaf7bf687c..b39018c3ac4 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -507,7 +507,14 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { } void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { - GGML_ASSERT(!exact_pages || seq_id_src == seq_id_dst); + // [TAG_EXACT_CONCURRENCY] a page belongs to one (sequence, position/256) pair, so two sequences + // cannot share physical cells. Refuse the copy rather than abort the process: this is reachable + // from a request parameter. + if (exact_pages && seq_id_src != seq_id_dst) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so cells cannot be shared between " + "sequences: ignoring the copy from seq %d to seq %d\n", __func__, seq_id_src, seq_id_dst); + return; + } // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -627,7 +634,14 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) { } void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { - GGML_ASSERT(!exact_pages || shift == 0); + // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so a + // shift would have to move the cells too. get_can_shift() reports this so that --context-shift + // and --cache-reuse are turned off at load; this is the guard for the library API. + if (exact_pages && shift != 0) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so positions cannot be shifted: " + "ignoring the shift of %d on seq %d\n", __func__, shift, seq_id); + return; + } // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -678,7 +692,13 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll } void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { - GGML_ASSERT(!exact_pages || d == 1); + // [TAG_EXACT_CONCURRENCY] same reason as seq_add: the offset inside a page is derived from the + // position, so dividing the positions would leave every cell in the wrong slot. + if (exact_pages && d != 1) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so positions cannot be divided: " + "ignoring the division by %d on seq %d\n", __func__, d, seq_id); + return; + } // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -1276,6 +1296,13 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & } bool llama_kv_cache::get_can_shift() const { + // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so the + // paged pool cannot shift positions. Reporting it here is what makes the server disable + // --context-shift and --cache-reuse at load, with a warning, instead of accepting both and + // failing on the first request that needs them. + if (exact_pages) { + return false; + } // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot. if (model.arch == LLM_ARCH_STEP35) { return false; @@ -2431,9 +2458,16 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 GGML_ASSERT(cells.seq_has(idx, dest_seq_id)); } } else { - GGML_ASSERT(!exact_pages && "exact mode supports per-sequence restore only"); // whole KV cache restore + // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical + // index, which the paged pool owns. Report it like every other failure in this function. + if (exact_pages) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state " + "restore only\n", __func__); + return false; + } + if (cell_count > cells.size()) { LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__); return false; From b81d8e7018bbc45f1a7c4ef93e92db69e3a2d62a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:44:24 +0000 Subject: [PATCH 25/81] batch: isolate only the prompt sequences, in every memory type Two things were wrong with the one-sequence prompt ubatch rule. It only existed in the hybrid memory. A dense transformer with a unified cache takes split_simple, which packs every sequence's prompt tokens into one ubatch, so its prefill matmuls ran at a width the solo run never sees and, with the column policy bounded, produced K and V the solo run never produces. Exact mode was therefore not exact by construction on dense models, which is most of what it will be pointed at. A pure recurrent model still called the three-argument split_equal for the same reason the hybrid one no longer does. Both now take the rule: llama_kv_cache::init_batch keeps split_simple for a plain decode step and switches to the sequence-set split when a prompt is present, and llama_memory_recurrent::init_batch passes the flag through. The rule itself then serialized more than it had to. has_multi_token_seq() scanned the whole original batch and ignored used[], so it stayed true after the prompt had been consumed, and the n_seqs_max cap it fed capped every sequence set including one-token decode sets. One prompt chunk plus three decodes therefore became four single-sequence ubatches and the three chats decoded one at a time for the whole prefill, which contradicts the comment saying a plain decode step stays batched. The predicate now skips used[] tokens, and the cap became an isolate_multi_token_seqs flag: a sequence set with more than one token left to place takes a ubatch of its own, sets with one token left keep grouping. One prompt next to three decodes now costs one extra ubatch, not three. The KV cache constructor also read getenv("LLAMA_EXACT_CONCURRENCY") directly while llama_exact_concurrency() and ggml_cuda_exact_concurrency() each cache the first value they see, so a process that created one context with the knob unset and then set it got a paged cache on top of a dispatcher still in default mode. It now reads the same cached value as the other two. --- src/llama-batch.cpp | 40 +++++++++++++++++++++++++++++----- src/llama-batch.h | 14 +++++++----- src/llama-kv-cache.cpp | 16 +++++++++++--- src/llama-memory-hybrid.cpp | 6 ++--- src/llama-memory-recurrent.cpp | 6 ++++- 5 files changed, 63 insertions(+), 19 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 50f70bb0ff3..4b73ab2478b 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -511,6 +511,11 @@ bool llama_batch_allocr::has_multi_token_seq() const { std::vector n_per_seq(n_seq_max, 0); for (int32_t i = 0; i < batch.n_tokens; ++i) { + // tokens already placed in an earlier ubatch do not make the rest of the batch a prompt + if (used[i]) { + continue; + } + for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { if (++n_per_seq[batch.seq_id[i][s]] > 1) { return true; @@ -521,7 +526,7 @@ bool llama_batch_allocr::has_multi_token_seq() const { return false; } -llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t n_seqs_max) { +llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, bool isolate_multi_token_seqs) { if (sequential && has_cpl) { LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__); @@ -554,6 +559,34 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } if (add) { + // [TAG_EXACT_CONCURRENCY] a sequence set that still has more than one token to place is + // a prompt, and a prompt shares its arithmetic with whatever else is in the ubatch, so + // give it a ubatch of its own. Sets with one token left are a plain decode step, which + // is already exact, so keep grouping those: isolating them too would make one prompt + // serialize every concurrent decode for the whole of the prefill. + if (isolate_multi_token_seqs) { + uint32_t n_left = 0; + + for (const auto idx : seq_set_map[seq_set[i]]) { + if (!used[idx]) { + ++n_left; + } + } + + if (n_left > 1) { + if (!cur_seq_set.empty()) { + // let the sets already taken have this ubatch; the prompt gets the next one + break; + } + + cur_seq_set.push_back(seq_set[i]); + + last_seq_id = batch.seq_id[i][0]; + + break; + } + } + cur_seq_set.push_back(seq_set[i]); last_seq_id = batch.seq_id[i][0]; @@ -561,11 +594,6 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, if (cur_seq_set.size() > n_ubatch) { break; } - - // [TAG_EXACT_CONCURRENCY] - if (n_seqs_max > 0 && cur_seq_set.size() >= n_seqs_max) { - break; - } } } diff --git a/src/llama-batch.h b/src/llama-batch.h index d354c442d03..4bd2aa98f9f 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -105,12 +105,14 @@ class llama_batch_allocr { // make ubatches of equal-length sequences sets // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch - // n_seqs_max = maximum sequence sets per ubatch, 0 = no limit - // [TAG_EXACT_CONCURRENCY] passing 1 keeps a ubatch to a single sequence - llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t n_seqs_max = 0); - - // [TAG_EXACT_CONCURRENCY] true if some sequence contributes more than one token to the batch, - // i.e. this is not a plain one-token-per-sequence decode step + // isolate_multi_token_seqs = [TAG_EXACT_CONCURRENCY] a sequence set with more than one token + // left to place is given a ubatch of its own; sets with a single token left are + // still grouped together, so a prompt next to three decodes costs one extra + // ubatch and does not serialize the three decodes + llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, bool isolate_multi_token_seqs = false); + + // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than one token left to place, + // i.e. what remains of the batch is not a plain one-token-per-sequence decode step bool has_multi_token_seq() const; // sequence-set-wise split - each ubatch contains a single sequence-set diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index b39018c3ac4..16c04fb6e3b 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -107,8 +107,10 @@ llama_kv_cache::llama_kv_cache( v_cells_impl(other ? other->v_cells_impl : std::make_shared()), v_cells(*v_cells_impl) { - const char * exact_env = getenv("LLAMA_EXACT_CONCURRENCY"); - exact_pages = exact_env && atoi(exact_env) != 0; + // [TAG_EXACT_CONCURRENCY] read the knob through the one cached reader that the graph and the + // CUDA dispatcher also use, so a process that sets it between two context creations cannot end + // up with a paged cache on top of a dispatcher that is still in default mode + exact_pages = llama_exact_concurrency(); // shared cells view the source cache's K/V tensors, so the cell count // follows the source allocation: a fitted target can be smaller than the @@ -791,7 +793,15 @@ llama_memory_context_ptr llama_kv_cache::init_batch( std::vector ubatches; while (true) { - auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0); + // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt tokens into one + // ubatch, so a sequence's prefill would run at a width its solo run never sees. Take + // the sequence-set split instead, which can give each prompt a ubatch of its own; a + // plain decode step has nothing to isolate and keeps taking split_simple. + const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq(); + + auto ubatch = n_stream == 1 && !isolate + ? balloc.split_simple(n_ubatch) + : balloc.split_equal(n_ubatch, n_stream > 1, 0, isolate); if (ubatch.n_tokens == 0) { break; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index ba54ab12923..4ebd476aa8b 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -89,11 +89,11 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // [TAG_EXACT_CONCURRENCY] the recurrent half of a hybrid model is not invariant to // the shape of the ubatch: a prompt processed next to other sequences' prompt tokens // leaves a different gated delta net state than the same prompt processed alone. - // Keeping such a ubatch to a single sequence removes that. A plain decode step, one + // Giving such a sequence a ubatch of its own removes that. A plain decode step, one // token per sequence, is already exact and stays batched. - const uint32_t n_seqs_max = llama_exact_concurrency() && balloc.has_multi_token_seq() ? 1 : 0; + const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq(); - ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, n_seqs_max); + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } if (ubatch.n_tokens == 0) { diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7..f639a25c5df 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -431,7 +431,11 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & // [TAG_RECURRENT_ROLLBACK_SPLITS] // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch // so that the rollback snapshots remain valid - ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0); + // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: a recurrent state that a + // prompt leaves behind depends on what shared its ubatch, so isolate the prompts + const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq(); + + ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } if (ubatch.n_tokens == 0) { From c7027d60dd8c48a0b72d11a9ba07c27e7b8c9504 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:49:37 +0000 Subject: [PATCH 26/81] cuda: derive the exact mode column bound from the decode width The fixed default of 16 silently turns exact mode off above 16 columns: MUL_MAT, MUL_MAT_ID and the per-row attention split all fall back to the neighbour dependent batched path. --parallel 6 --spec-draft-n-max 2 gives 18 columns and --parallel 8 gives 24; both are ordinary server configurations and neither said anything. The source comment documented the cliff, nothing at runtime did. The bound only ever had to cover the widest ubatch a decode step can build, since a prompt ubatch holds one sequence and gets its exactness from that. So let the caller report that width. ggml_backend_cuda_set_exact_decode_width(), reachable directly or through ggml_backend_reg_get_proc_address(), takes one column per slot times one plus the draft length, and exact mode defaults the bound to it. common computes it from n_parallel and the speculative type and reports it before the warmup, which is the first graph any of these tools computes, and refuses at startup an explicitly set GGML_CUDA_BATCH_INVARIANT_MAX_COLS that is smaller: GGML_CUDA_BATCH_INVARIANT_MAX_COLS is 8 but LLAMA_EXACT_CONCURRENCY needs at least 12 to cover a decode step of 4 slots, above which a matmul is left batched and its rows depend on the other rows in the ubatch. Raise it to 12, set it to 0 for no bound, or unset it to let it default to 12. --parallel 4 with speculation off now defaults to 4 rather than 16 and with --spec-type draft-mtp --spec-draft-n-max 2 to 12 rather than 16, which is the same guarantee over a narrower range of shapes: a decode ubatch of that model cannot be wider than that, and everything above it is a prompt. When nothing reported a width, the default stays 16 and the dispatcher warns once per process the first time a MUL_MAT or MUL_MAT_ID above the bound is left unsplit, naming both numbers. It deliberately stays quiet once a width is known, because then the only batches above the bound are prompt ubatches and warning on those would fire on every prefill for a case that is working as intended. --- common/arg.cpp | 7 +++ common/common.cpp | 82 +++++++++++++++++++++++++++++++++ common/common.h | 13 ++++++ ggml/include/ggml-cuda.h | 9 ++++ ggml/src/ggml-cuda/ggml-cuda.cu | 77 +++++++++++++++++++++++++------ 5 files changed, 175 insertions(+), 13 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 5bfa4adcdf0..e4f2e8f2517 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1304,6 +1304,13 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e exit(0); } params.lr.init(); + + // [TAG_EXACT_CONCURRENCY] refuse a column bound that cannot cover a decode step before + // anything is loaded, rather than running with the guarantee quietly switched off + if (!common_exact_concurrency_init(ctx_arg.params)) { + ctx_arg.params = params_org; + return false; + } } catch (const std::invalid_argument & ex) { fprintf(stderr, "%s\n", ex.what()); ctx_arg.params = params_org; diff --git a/common/common.cpp b/common/common.cpp index 3d54bd6002d..d2beebc8e2f 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1,4 +1,5 @@ #include "ggml.h" +#include "ggml-backend.h" #include "gguf.h" #include "build-info.h" @@ -1433,6 +1434,82 @@ std::vector & common_init_result::lora() { return pimpl->lora; } +// [TAG_EXACT_CONCURRENCY] +bool common_exact_concurrency() { + static const bool enabled = []() { + const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); + return val && atoi(val) != 0; + }(); + + return enabled; +} + +// [TAG_EXACT_CONCURRENCY] +int common_exact_decode_width(const common_params & params) { + const int n_slots = std::max(1, params.n_parallel); + + // the draft tokens a slot carries into the verify ubatch alongside its accepted token + int n_draft = 0; + + for (const auto type : params.speculative.types) { + switch (type) { + case COMMON_SPECULATIVE_TYPE_NONE: + break; + case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: + n_draft = std::max(n_draft, params.speculative.ngram_mod.n_max); + break; + case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: + n_draft = std::max(n_draft, (int) params.speculative.ngram_simple.size_m); + break; + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: + n_draft = std::max(n_draft, (int) params.speculative.ngram_map_k.size_m); + break; + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: + n_draft = std::max(n_draft, (int) params.speculative.ngram_map_k4v.size_m); + break; + default: + n_draft = std::max(n_draft, params.speculative.draft.n_max); + break; + } + } + + return n_slots*(1 + std::max(0, n_draft)); +} + +// [TAG_EXACT_CONCURRENCY] +bool common_exact_concurrency_init(const common_params & params) { + if (!common_exact_concurrency()) { + return true; + } + + const int n_cols = common_exact_decode_width(params); + + const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); + if (bound) { + const int max_cols = atoi(bound); + if (max_cols > 0 && max_cols < n_cols) { + COM_ERR("GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at " + "least %d to cover a decode step of %d slots, above which a matmul is left " + "batched and its rows depend on the other rows in the ubatch. Raise it to %d, " + "set it to 0 for no bound, or unset it to let it default to %d.\n", + max_cols, n_cols, std::max(1, params.n_parallel), n_cols, n_cols); + return false; + } + } + + // the CUDA backend may not be present or may be loaded dynamically, so go through the registry + for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { + ggml_backend_reg_t reg = ggml_backend_reg_get(i); + + auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width"); + if (fn) { + fn(n_cols); + } + } + + return true; +} + common_init_result_ptr common_init_from_params(common_params & params, bool model_only) { common_init_result_ptr res(new common_init_result(params, model_only)); @@ -1454,6 +1531,11 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode const llama_vocab * vocab = llama_model_get_vocab(model); + // [TAG_EXACT_CONCURRENCY] before the warmup, which is the first graph this process computes + if (!common_exact_concurrency_init(params)) { + return res; + } + if (params.ctx_shift && !llama_memory_can_shift(llama_get_memory(lctx))) { COM_WRN("%s", "KV cache shifting is not supported for this context, disabling KV cache shifting\n"); params.ctx_shift = false; diff --git a/common/common.h b/common/common.h index c99269f9a96..2be2fab6a8b 100644 --- a/common/common.h +++ b/common/common.h @@ -931,6 +931,19 @@ using common_init_result_ptr = std::unique_ptr; common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false); +// [TAG_EXACT_CONCURRENCY] +// true when LLAMA_EXACT_CONCURRENCY is set for this process +bool common_exact_concurrency(); + +// the widest ubatch a decode step can build with these parameters: one column per slot, times one +// plus the number of speculative draft tokens carried with it. Under exact mode this is what the +// CUDA column policy has to cover, and what its default bound is derived from. +int common_exact_decode_width(const common_params & params); + +// report that width to the CUDA backend, and refuse an explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS +// that is smaller than it. Returns false if the configuration must not run. +bool common_exact_concurrency_init(const common_params & params); + struct llama_model_params common_model_params_to_llama ( common_params & params); struct llama_context_params common_context_params_to_llama(const common_params & params); diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h index 1cd81eeaebc..c3dd87c97b7 100644 --- a/ggml/include/ggml-cuda.h +++ b/ggml/include/ggml-cuda.h @@ -38,6 +38,15 @@ GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total); GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); + +// [TAG_EXACT_CONCURRENCY] +// Report the widest ubatch a decode step of this process can build: one column per slot, times one +// plus the number of speculative draft tokens carried with it. Under LLAMA_EXACT_CONCURRENCY the +// column policy then defaults to that width instead of a fixed number, so --parallel or a wider +// draft cannot silently push a decode above the bound and leave it batched. An explicitly set +// GGML_CUDA_BATCH_INVARIANT_MAX_COLS still wins. Call before the first graph is computed. Also +// available through ggml_backend_reg_get_proc_address(). +GGML_BACKEND_API void ggml_backend_cuda_set_exact_decode_width(int n_cols); GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 4adee44ca61..c8f43bae382 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1852,22 +1852,67 @@ int ggml_cuda_batch_invariant() { return mode; } +// [TAG_EXACT_CONCURRENCY] the widest decode ubatch the caller says it can build, 0 if it never said +static std::atomic g_exact_decode_width{0}; + +void ggml_backend_cuda_set_exact_decode_width(int n_cols) { + g_exact_decode_width.store(n_cols > 0 ? n_cols : 0, std::memory_order_relaxed); +} + int ggml_cuda_batch_invariant_max_cols() { - static const int max_cols = []() { - // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill - // matmul shapes match its solo run, so exact mode no longer needs the column policy to be - // unbounded there. An explicit bound always wins, in either mode. + // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill + // matmul shapes match its solo run, so exact mode no longer needs the column policy to be + // unbounded there. An explicit bound always wins, in either mode. + static const int explicit_cols = []() { const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); - if (val) { return atoi(val); } - // Exact mode then only has to cover the widest ubatch a decode step can build: one column - // per slot, times one plus the number of speculative draft tokens carried with it. 16 - // covers the default four slots at up to three tokens each, which is what - // --spec-type draft-mtp --spec-draft-n-max 2 produces. More slots, or a wider draft, need - // the bound set explicitly; above it the column split does not fire. - if (ggml_cuda_exact_concurrency()) { return 16; } - return 0; + return val ? atoi(val) : -1; }(); - return max_cols; + + if (explicit_cols >= 0) { + return explicit_cols; + } + + if (!ggml_cuda_exact_concurrency()) { + return 0; + } + + // Exact mode only has to cover the widest ubatch a decode step can build: one column per slot, + // times one plus the number of speculative draft tokens carried with it. Use that width when + // the caller reported it through ggml_backend_cuda_set_exact_decode_width(). Nothing reported + // it, so fall back to 16, which covers four slots at up to three tokens each, which is what + // --parallel 4 --spec-type draft-mtp --spec-draft-n-max 2 produces. Above the bound the column + // split does not fire, and ggml_cuda_warn_above_exact_bound() says so once. + const int width = g_exact_decode_width.load(std::memory_order_relaxed); + + return width > 0 ? width : 16; +} + +// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so its rows depend on the +// other rows in it and the mode does not hold for that op. Say so once, rather than never. +// +// Only when nothing reported a decode width. When one was reported the bound is derived from it, so +// the only batches above the bound are prompt ubatches, and those hold a single sequence under this +// mode: their exactness comes from that, not from the column policy, and leaving them batched is +// the whole point of having a bound at all. Warning on those would be crying wolf on every prefill. +static void ggml_cuda_warn_above_exact_bound(const char * op, int64_t ncols, int max_cols) { + if (!ggml_cuda_exact_concurrency()) { + return; + } + + if (g_exact_decode_width.load(std::memory_order_relaxed) > 0) { + return; + } + + static std::atomic_flag warned = ATOMIC_FLAG_INIT; + if (warned.test_and_set(std::memory_order_relaxed)) { + return; + } + + GGML_LOG_WARN("%s: LLAMA_EXACT_CONCURRENCY is set, but this %s is %d columns wide while " + "GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d, so it is left batched and its result depends " + "on the other columns in the ubatch. Raise the bound, set it to 0 for no bound, or call " + "ggml_backend_cuda_set_exact_decode_width() with the widest decode this process builds. " + "Reported once.\n", __func__, op, (int) ncols, max_cols); } enum ggml_cuda_mm_path { @@ -1952,6 +1997,7 @@ static bool ggml_cuda_mul_mat_split_columns( } const int max_cols = ggml_cuda_batch_invariant_max_cols(); if (max_cols > 0 && ncols_dst > max_cols) { + ggml_cuda_warn_above_exact_bound("MUL_MAT", ncols_dst, max_cols); return false; } @@ -2049,6 +2095,7 @@ static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) { } const int max_cols = ggml_cuda_batch_invariant_max_cols(); if (max_cols > 0 && ntokens > max_cols) { + ggml_cuda_warn_above_exact_bound("MUL_MAT_ID", ntokens, max_cols); return false; } return true; @@ -5733,6 +5780,10 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } + // [TAG_EXACT_CONCURRENCY] + if (strcmp(name, "ggml_backend_cuda_set_exact_decode_width") == 0) { + return (void *)ggml_backend_cuda_set_exact_decode_width; + } return nullptr; } From 9cd1222437a91a4750651801b5ef146819791c77 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:50:19 +0000 Subject: [PATCH 27/81] graph: refuse exact mode on the V-less attention layouts Page tables are wired into llm_graph_input_attn_kv only. llm_graph_input_attn_k has no self_pages member and its build_attn calls build_attn_mha without a pages argument, as do the DeepSeek sparse and sliding window variants. A model on one of those layouts still got its cells placed in pages by the allocator and then attended in physical order after a park and a restore, so the mode reported itself as on and lost the one invariant it exists for. That is the same silent failure the CUDA placement gate was added to stop, so answer it the same way: log which layout it is and fail the context. Rejecting is the smaller correct change of the two. Wiring self_pages into llm_graph_input_attn_k is four lines and looks tempting, but it fixes one of the four V-less input classes and DeepSeek 3.2 uses two of them: its sparse layers build their mask from a top-k selection and would stay unpaged, leaving the model half paged, which is worse than refused. None of these architectures was measured here, and the paged kernel also requires 256-dimensional K and V heads, which none of them was checked against. Reaches the user through the path llama_init_from_model already has for a context that cannot be built: llm_graph_reject_exact_concurrency: LLAMA_EXACT_CONCURRENCY is set, but this model uses the V-less KV (attn_k) attention layout, which carries no page table and would attend in physical cell order llama_init_from_model: failed to initialize the context: exact concurrency: unsupported attention layout --- src/llama-graph.cpp | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 71fce85bca7..f8f5c83a2b6 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -21,11 +21,36 @@ #include #include #include +#include #include #include // dedup helpers +// [TAG_EXACT_CONCURRENCY] +// The page table is wired into llm_graph_input_attn_kv only. The V-less layouts build their +// attention without one, so a model on one of those would get its cells placed in pages by the +// allocator and then attend in physical cell order anyway: the mode would report itself as on and +// lose the one invariant it exists for, which is the same silent failure the CUDA placement gate +// was added to stop. Refuse the context instead. +// +// Rejecting is the smaller correct change here. Wiring self_pages into llm_graph_input_attn_k alone +// is four lines, but it fixes only one of the four V-less input classes, and DeepSeek 3.2 uses two +// of them: its sparse layers rewrite the mask from a top-k selection and would still be unpaged, so +// the model would end up half paged, which is worse than refused. None of these architectures was +// measured, and the paged kernel additionally requires 256-dimensional K and V heads. +static void llm_graph_reject_exact_concurrency(const char * layout) { + if (!llama_exact_concurrency()) { + return; + } + + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, but this model uses the %s attention " + "layout, which carries no page table and would attend in physical cell order\n", + __func__, layout); + + throw std::runtime_error("exact concurrency: unsupported attention layout"); +} + static ggml_tensor * build_attn_inp_kq_mask( ggml_context * ctx, const llama_kv_cache_context * mctx, @@ -2871,6 +2896,8 @@ static std::unique_ptr build_attn_inp_k_impl( const llama_cparams & cparams, const llama_kv_cache_context * mctx_cur) { + llm_graph_reject_exact_concurrency("V-less KV (attn_k)"); + auto inp = std::make_unique(hparams, cparams, mctx_cur); { @@ -3247,6 +3274,8 @@ static std::unique_ptr build_attn_inp_k_dsa_impl( const llama_cparams & cparams, const llama_kv_cache_dsa_context * mctx_cur) { + llm_graph_reject_exact_concurrency("sparse V-less KV (attn_k_dsa)"); + auto inp = std::make_unique(hparams, cparams, mctx_cur); { @@ -3364,6 +3393,8 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const { const auto * mctx_cur = static_cast(mctx); + llm_graph_reject_exact_concurrency("V-less sliding window KV (attn_k_iswa)"); + auto inp = std::make_unique(hparams, cparams, mctx_cur); { From b5e9ebdf2c6b42df8a241d321f876ddb3ebddb43 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:50:35 +0000 Subject: [PATCH 28/81] batchinv: stop the harness from certifying a reduced run run_concurrent lost worker exceptions. A Python thread exception only prints a traceback, join() returns, and the partial dict was returned, so a round where P1 to P3 failed and P0 succeeded still classified P0 as identical and aggregated throughput over whichever requests happened to survive. The evidence harness could certify a solo run as a clean four-way concurrency result. Exceptions are now collected under a lock and raised after join, the barrier is aborted so the other workers do not block on one that will never fill, and every expected name has to be present before the result is returned. bench.py imports the same helper, so its throughput number is covered too. Server.__enter__ raised after Popen had already started the server, and Python does not call __exit__ when __enter__ raises, so a server that started but never reported healthy kept the GPU, the port and the log handle. The health wait is now wrapped and tears the server down before re-raising. __exit__ also waits after the SIGKILL path instead of leaving a zombie, and says in place that it is POSIX only. The run record and the server log header captured only variables starting with GGML plus CUDA_VISIBLE_DEVICES, so an inherited LLAMA_EXACT_CONCURRENCY was invisible in both and a run intended as the mode-off reference could silently have been an exact-mode run while the JSON said "env": {}. That is the baseline the whole divergence claim rests on. Both now record the environment the server actually inherited, from an explicit allowlist that includes LLAMA_EXACT_CONCURRENCY, GGML_CUDA_BATCH_INVARIANT, GGML_CUDA_BATCH_INVARIANT_MAX_COLS, LLAMA_SERVER_PREEMPT_EVERY and CUDA_VISIBLE_DEVICES, along with the resolved model path and the full server command line. What the run asked for is kept separately as env_requested. UNSLOTH_WORKSPACE was read at import, so both tools raised KeyError before argparse ran and even --help failed. The model path is resolved when the server arguments are built and raises a named RuntimeError. The README still said the mode forces GGML_CUDA_BATCH_INVARIANT=2 with no column limit including during prefill, which stopped being true two commits before this branch. It now states the bound, where its default comes from, that an explicit value below the decode width is refused at startup, and the load-time refusals for placement and the V-less layouts. --- scripts/batchinv/README.md | 31 +++++++-- scripts/batchinv/divergence.py | 118 ++++++++++++++++++++++++--------- 2 files changed, 113 insertions(+), 36 deletions(-) diff --git a/scripts/batchinv/README.md b/scripts/batchinv/README.md index 62a51487047..82b769b6d04 100644 --- a/scripts/batchinv/README.md +++ b/scripts/batchinv/README.md @@ -1,18 +1,31 @@ # Exact concurrency experiment p Opt in before loading the model with `LLAMA_EXACT_CONCURRENCY=1`. This also forces -`GGML_CUDA_BATCH_INVARIANT=2` with no column limit, including during prefill. +`GGML_CUDA_BATCH_INVARIANT=2` and gives `GGML_CUDA_BATCH_INVARIANT_MAX_COLS` a +default. The column policy only has to cover the widest ubatch a decode step can +build, one column per slot times one plus the draft length, because a prompt +ubatch is kept to one sequence and gets its exactness from that instead. Tools +built on `common` report that width, so the default is `--parallel` times one plus +`--spec-draft-n-max`, and an explicitly set `GGML_CUDA_BATCH_INVARIANT_MAX_COLS` +smaller than it is refused at startup. Nothing reported a width, the default is 16 +and the dispatcher warns once the first time a `MUL_MAT` or `MUL_MAT_ID` above the +bound is left unsplit. Set the variable to `0` for no bound; above the bound the +column policy does not fire, including during prefill. The experimental policy supports unified, offloaded F16 K/V, causal flash attention, 256-dimensional K and V heads, no attention soft cap, and no sliding window. Shared-weight matmuls over multiple sequence planes are normalized to one plane before the inherited selective column dispatcher. Without this, the recurrent output projection bypasses batch invariance during concurrent prefill. -It is measured on text prompts with Qwen3.5-4B on one B200. Context shifting, -position division, cross-sequence prefix copies, shared-prefix input tokens, and -whole-context state loading are unsupported. Per-sequence state save and restore -is supported. Unsupported cache transformations assert instead of silently -violating the page invariant. +It is measured on text prompts with Qwen3.5-4B on one B200. Every KV layer has to +be on the CUDA backend, since no other backend reads the page table; a partial or +absent offload fails the load naming the layer. The V-less attention layouts have +no page table either, and a model on one of those is refused at context creation. +Context shifting, position division, cross-sequence prefix copies, shared-prefix +input tokens, and whole-context state loading are unsupported. Per-sequence state +save and restore is supported. Unsupported cache transformations are refused with +a logged error and leave the cells untouched; `--context-shift` and +`--cache-reuse` are reported as unsupported at load and disabled there. The allocator owns pages of 256 cells on behalf of one (sequence, position/256). Position modulo 256 fixes the cell offset. Empty pages remain in the unified pool @@ -50,3 +63,9 @@ reference. `bench.py --modes 0,1 --pairs 3` measures default off against exact m on, with 256 predicted tokens. Set `UNSLOTH_WORKSPACE` to the model parent workspace and `LD_LIBRARY_PATH` to this build's bin directory. The harness uses GPU 3; select a port in 9601-9610 explicitly. + +A concurrent request that fails now fails the run instead of being dropped from +the result, and every run records the environment the server actually inherited, +including `LLAMA_EXACT_CONCURRENCY`, under `env` in the JSON and in the server log +header, so a run labelled as the mode-off reference can be checked rather than +trusted. Teardown is POSIX only. diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py index e451afca0e1..3e554f465b7 100644 --- a/scripts/batchinv/divergence.py +++ b/scripts/batchinv/divergence.py @@ -5,8 +5,22 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from prompts import PROMPTS -WS = os.environ["UNSLOTH_WORKSPACE"] -MODEL = f"{WS}/models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf" +# Environment recorded with every run. LLAMA_EXACT_CONCURRENCY inherited from the shell is what +# decides whether a run labelled as the mode-off reference actually was one, so it is not optional. +RECORDED_ENV = ("LLAMA_EXACT_CONCURRENCY", "GGML_CUDA_BATCH_INVARIANT", + "GGML_CUDA_BATCH_INVARIANT_MAX_COLS", "LLAMA_SERVER_PREEMPT_EVERY", + "CUDA_VISIBLE_DEVICES") + +MODEL_REL = "models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf" + + +def model_path(): + """Resolved when the server args are built, so --help works without the variable set.""" + ws = os.environ.get("UNSLOTH_WORKSPACE") + if not ws: + raise RuntimeError("UNSLOTH_WORKSPACE is not set; it must point at the workspace holding " + + MODEL_REL) + return os.path.join(ws, MODEL_REL) def post(port, path, payload, timeout=1800): @@ -34,7 +48,7 @@ def completion(port, prompt, n_predict): class Server: def __init__(self, port, binary, extra, env_extra, log_path, spec, kv_unified=True): self.port, self.log_path = port, log_path - self.args = [binary, "-m", MODEL, "--port", str(port), "--host", "127.0.0.1", + self.args = [binary, "-m", model_path(), "--port", str(port), "--host", "127.0.0.1", "--parallel", "4", "-c", "8192", "--flash-attn", "on", "--metrics", "-ngl", "99", "--no-warmup", "--seed", "0", "--spec-type", spec] @@ -46,48 +60,78 @@ def __init__(self, port, binary, extra, env_extra, log_path, spec, kv_unified=Tr self.env = dict(os.environ) self.env["CUDA_VISIBLE_DEVICES"] = "3" self.env.update(env_extra) + # what the server will actually see, not what this run meant to set + self.env_resolved = {k: self.env[k] for k in RECORDED_ENV if k in self.env} + self.p = None + self.fh = None def __enter__(self): self.fh = open(self.log_path, "ab") self.fh.write(("\n=== " + " ".join(self.args) + "\n=== env " + - json.dumps({k: v for k, v in self.env.items() - if k.startswith("GGML") or k == "CUDA_VISIBLE_DEVICES"}) + "\n").encode()) + json.dumps(self.env_resolved) + "\n").encode()) self.fh.flush() self.p = subprocess.Popen(self.args, stdout=self.fh, stderr=subprocess.STDOUT, env=self.env, start_new_session=True) print(f"[server] pid={self.p.pid} port={self.port} log={self.log_path}", flush=True) - deadline = time.time() + 600 - while time.time() < deadline: - if self.p.poll() is not None: - raise RuntimeError(f"server died rc={self.p.returncode}, see {self.log_path}") - try: - if get(self.port, "/health").get("status") == "ok": - print("[server] ready", flush=True) - return self - except Exception: - time.sleep(1.0) - raise RuntimeError("server did not become healthy") + try: + deadline = time.time() + 600 + while time.time() < deadline: + if self.p.poll() is not None: + raise RuntimeError(f"server died rc={self.p.returncode}, see {self.log_path}") + try: + if get(self.port, "/health").get("status") == "ok": + print("[server] ready", flush=True) + return self + except Exception: + time.sleep(1.0) + raise RuntimeError("server did not become healthy") + except BaseException: + # __exit__ is not called when __enter__ raises, so a server that started but never + # reported healthy would keep the GPU, the port and the log handle + self.__exit__(None, None, None) + raise def __exit__(self, *a): - print(f"[server] stopping pid={self.p.pid}", flush=True) - try: - os.killpg(os.getpgid(self.p.pid), signal.SIGTERM) - self.p.wait(timeout=60) - except Exception: + # note: POSIX only. On Windows this needs CREATE_NEW_PROCESS_GROUP at Popen and + # terminate()/kill() here; the runs this harness backs are Linux only. + if self.p is not None: + print(f"[server] stopping pid={self.p.pid}", flush=True) try: - os.killpg(os.getpgid(self.p.pid), signal.SIGKILL) + os.killpg(os.getpgid(self.p.pid), signal.SIGTERM) + self.p.wait(timeout=60) except Exception: - pass - self.fh.close() + try: + os.killpg(os.getpgid(self.p.pid), signal.SIGKILL) + except Exception: + pass + try: + self.p.wait(timeout=60) + except Exception: + pass + self.p = None + if self.fh is not None: + self.fh.close() + self.fh = None def run_concurrent(port, names, n_predict): barrier = threading.Barrier(len(names)) + lock = threading.Lock() out = {} + errors = [] def work(name): - barrier.wait() - out[name] = completion(port, PROMPTS[name], n_predict) + try: + barrier.wait() + res = completion(port, PROMPTS[name], n_predict) + except BaseException as e: + with lock: + errors.append((name, e)) + # release the others rather than let them block on a barrier that will never fill + barrier.abort() + return + with lock: + out[name] = res ts = [threading.Thread(target=work, args=(n,)) for n in names] t0 = time.time() @@ -95,7 +139,18 @@ def work(name): t.start() for t in ts: t.join() - return out, time.time() - t0 + wall = time.time() - t0 + + # a thread exception used to only print a traceback, so a run where P1..P3 failed and P0 + # succeeded was still reported as a clean four-way concurrency result + if errors: + raise RuntimeError("concurrent requests failed: " + + "; ".join(f"{n}: {type(e).__name__}: {e}" for n, e in errors)) + missing = set(names) - set(out) + if missing: + raise RuntimeError(f"concurrent requests produced no result for {sorted(missing)}") + + return out, wall def first_diff(a, b): @@ -121,12 +176,15 @@ def main(): a = ap.parse_args() env_extra = dict(kv.split("=", 1) for kv in a.env) + server = Server(a.port, a.binary, a.extra, env_extra, a.out + ".server.log", a.spec, + kv_unified=not a.no_kv_unified) res = {"label": a.label, "spec": a.spec, "n_predict": a.n_predict, - "env": env_extra, "extra": a.extra, "binary": a.binary, + "env_requested": env_extra, "env": server.env_resolved, + "model": server.args[2], "args": server.args, + "extra": a.extra, "binary": a.binary, "kv_unified": not a.no_kv_unified} - with Server(a.port, a.binary, a.extra, env_extra, a.out + ".server.log", a.spec, - kv_unified=not a.no_kv_unified) as s: + with server as s: solo = completion(a.port, PROMPTS["P0"], a.n_predict) ref = json.load(open(a.reference))["tokens"] if a.reference else solo["tokens"] res["solo_first_diff"] = first_diff(ref, solo["tokens"]) From bba1668572a4995c4210d9ccd791399d9aea6af8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:21:36 +0000 Subject: [PATCH 29/81] server: do not leave an issued async park holding the room the decode needs A review of #192 pointed at the victim loop in update_preemption(). When the pool has no room for the step about to be built and no park is in flight, the loop issues the victim's asynchronous park and breaks. The cells are held until the copy lands, so the batch is built into a pool that has not got smaller, llama_decode returns 1, and the retry ladder halves n_batch to 1 in microseconds without ever polling the copy, ending in "Context size has been exceeded" for every slot. The synchronous path freed the cells before returning, so it could not do this. I could not reproduce it. Six live rounds on the 4B at -c 8192, four chats with 1000-token prompts and 2048 tokens each, with the fourth chat's prompt held back 20 s so it arrives into a pool the other three have filled, exact mode on and off, on a binary without this change: 4 of 4 every round, no context errors, and the retry ladder was not entered once ("failed to find free space" appears zero times in both server logs). The reason is that preempt_kv_reserve() counts an incoming prompt chunk before it is allocated, so the planner crosses the lookahead threshold an iteration before the pool actually fills, and every park in those runs was issued with the 80 cells of asynchronous runway still ahead of it, never at the hard threshold this is about. Committing it anyway, because the described state is real even if these workloads do not reach it, and the change is inert unless it is reached: * The victim loop goes round again instead of leaving, but only when n_used + PREEMPT_N_MARGIN > n_cells, that is when there is no room for the step itself rather than merely less than the asynchronous lookahead wants. The next pass reaches preempt_wait_in_flight() and waits for the park just issued, which is what that function was written for and no worse than the synchronous path. Short of the lookahead only, it still breaks, because parking early and letting the copy run beside the decode is the entire point of #192. * On llama_decode returning 1, an outstanding park is waited for before any batch width is given up. Halving n_batch returns no cells, so without this the ladder can walk to n_batch == 1 and end every request while the room it needed was one event query away. Safe at that point because the slot was detached before the batch was built, so completing its park cannot change what is about to be retried; that is also why update_preemption() itself is not called from here. New test, test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_ everything: three slots generating near the ceiling and a fourth request whose prompt does not fit in what is left, which is the shape the existing tests miss because their victim holds almost no cells. The three are sized to oversubscribe the pool between them so the pressure does not depend on when the fourth arrives. It is kept for the shape it covers rather than as an attribution: it passes either way, and the attribution above was done at live scale. --- tools/server/server-context.cpp | 33 ++++++++++++++ tools/server/tests/unit/test_preempt.py | 57 +++++++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1a5fdc26934..5a8d440c25b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3695,6 +3695,27 @@ struct server_context_impl { victim->preempt_state_size() / (1024.0 * 1024.0), preempt_kv_used(), n_cells, n_used, victim->n_preempt); + + // [TAG_PREEMPT_ASYNC] Whether we may leave now depends on which of the two + // thresholds we are under. + // + // Short of the lookahead only: there is still room for the step about to be + // built, the park is early by design, and leaving is the whole point -- the + // copy runs beside the decode and update_preempt_copies() collects it next + // iteration. + // + // Out of room for the step itself: the cells are held until the copy lands, + // so leaving now builds a batch into a pool that has not got smaller. The + // decode fails, and the retry ladder halves n_batch to 1 without ever + // polling the copy, ending in "Context size has been exceeded" for every + // slot. The synchronous path did not have this problem because it returned + // the cells before it returned. Go round instead: the next pass reaches + // preempt_wait_in_flight() and waits for the park just issued, which is no + // worse than the synchronous path and is what it was written for. + if (n_used + PREEMPT_N_MARGIN > n_cells) { + continue; + } + break; } @@ -4649,6 +4670,18 @@ struct server_context_impl { } } + // [TAG_PREEMPT_ASYNC] Before giving up any batch width: a park that has been + // issued and not yet landed is holding cells that are already spoken for, and + // waiting for it returns them. Halving the batch returns nothing, so without + // this the ladder can run all the way down to n_batch == 1 and end every + // request while the room it needed was moments from arriving. Safe from here + // because the slot was detached before this batch was built, so completing its + // park cannot change what the batch about to be retried contains. + if (ret == 1 && preempt_wait_in_flight()) { + SRV_WRN("%s", "waited for an in-flight park before retrying the decode\n"); + return false; // retry at the same batch size, with the cells it freed + } + // retry with half the batch size to try to find a free slot in the KV cache if (!try_clear_idle_slots()) { n_batch /= 2; diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 1aba01751bd..4c3e8c465d8 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -447,3 +447,60 @@ def test_no_preempt_async_falls_back_to_the_synchronous_path(): # the synchronous path still parks and resumes assert text.count("preempted on request") >= 6 assert text.count("resumed after") >= 6 + + +def test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_everything(): + # [TAG_PREEMPT_ASYNC] The case the async path made worse than the synchronous one, and + # that the existing tests miss because their victim holds almost no cells. + # + # Three slots are well into generating when a fourth request arrives whose prompt does + # not fit in what is left. update_preemption() picks a victim and issues its park, but + # an asynchronous park does not return the cells before update_slots() carries on. If + # the loop leaves at that point, the batch is built into a pool that has not got any + # smaller, llama_decode returns 1, and the retry ladder halves n_batch to 1 in + # microseconds without ever polling the copy -- ending every request with "Context size + # has been exceeded" while the room it wanted was one event query away. + # + # Pass is what the synchronous path gave: a park, and all four requests finish. + global server + server.n_ctx = 512 + server.n_gpu_layer = 99 + server.n_slots = 4 + server.start() + log = LogReader(server.log_path) + + prompt_a, n_a = _prompt_of_about(100, "Alpha") + prompt_b, n_b = _prompt_of_about(100, "Bravo") + prompt_c, n_c = _prompt_of_about(100, "Charlie") + prompt_d, n_d = _prompt_of_about(150, "Delta") + + # A, B and C oversubscribe the pool between them, so the pressure does not depend on + # when D arrives, and every occupant is holding real cells rather than the handful the + # other tests park. Each of the four still fits on its own. + n_predict_abc = 130 + n_predict_d = 40 + assert max(n_a, n_b, n_c) + n_predict_abc < 512 and n_d + n_predict_d < 512 + assert n_a + n_b + n_c + 3 * n_predict_abc > 512 + + def _late(n_predict, prompt): + # D's prompt arrives into a pool the other three have already grown into; this + # model decodes about 120 tokens a second, so they are all still running + time.sleep(0.25) + return _complete(n_predict, prompt) + + results = parallel_function_calls([ + (_complete, (n_predict_abc, prompt_a)), + (_complete, (n_predict_abc, prompt_b)), + (_complete, (n_predict_abc, prompt_c)), + (_late, (n_predict_d, prompt_d)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted" in text + + for i, res in enumerate(results): + assert res.status_code == 200, (i, res.body) + for i in range(3): + assert results[i].body["timings"]["predicted_n"] == n_predict_abc + assert results[3].body["timings"]["predicted_n"] == n_predict_d From 83af5be139e319a2fab89421d615d7f34f04cf36 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 07:05:06 +0000 Subject: [PATCH 30/81] server: make the result queue timeout a deadline, so a parked stream is kept alive Found by putting #190 and #192 together and then looking for the keepalive that #190 promises. It never arrives. Live, the 4B on one B200, two streaming completions that do not fit together so one is parked until the other finishes, every SSE line timestamped as it arrives: 6.19s B : preempted 16.82s B : resumed A 10.63 s silence on a stream whose whole point is that it says ": preempt- keepalive" every 2 s. Four-chat runs at -c 8192 and -c 4096 show the same: parks of up to 14.59 s by the server's own "resumed after" line, and not one keepalive on any stream in any run. server_response::send() notify_all()s a single condition variable for every result of every task, and server_response::recv_with_timeout() waited with wait_for(), which restarts on every wakeup. A reader waiting on a task that is producing nothing is therefore woken by every token every other task produces, and its wait_for() never elapses. On a server with any traffic at all the timeout is not a timeout: whoever waits for a quiet task waits indefinitely. A parked slot is the worst possible case for this, because a slot is only ever parked while the others are busy, so the keepalive was unreachable by construction. The same applies to the ordinary --sse-ping, which likewise only fired on an otherwise idle server, and to the should_stop polling in server_response_reader::next(), whose own comment says it happens every polling_interval_seconds and did not. Compute the deadline once and wait_until() it. Spurious wakeups then re-check the queue and go back to waiting for the same instant, which is what every caller already reads the argument as meaning. After, the same two streams: 6.26s B : preempted 9.26s B : preempt-keepalive 12.26s B : preempt-keepalive 14.26s B : preempt-keepalive 16.75s B : resumed Three keepalives across a 10.49 s park, at the 2 s period plus the reader's 1 s polling granularity. The probe is scripts/integ_keepalive_probe.py. No harness test: stories260K generates several hundred tokens a second, and at the context sizes the harness uses a park lasts two or three seconds, which is the keepalive period itself. Every sizing I tried either parked for milliseconds at a time as the pool oscillated around full, or did not park at all. A test that straddles the period it is testing would be worse than none, so the regression is pinned by the live probe above. --- tools/server/server-queue.cpp | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d8..bd74db89408 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -448,6 +448,18 @@ server_task_result_ptr server_response::recv(const std::unordered_set & id_ } server_task_result_ptr server_response::recv_with_timeout(const std::unordered_set & id_tasks, int timeout) { + // [TAG_PREEMPT] The timeout is a deadline, not a per-wait duration. + // + // send() notify_all()s on one condition variable for every result of every task, so a + // reader waiting on a task that is producing nothing is woken by every token every other + // task produces. With wait_for() each of those wakeups restarted the wait, and on a busy + // server the timeout was never reached at all: whoever was waiting for a quiet task + // waited forever, however small the timeout they asked for. That is exactly the + // situation of a parked slot, which by definition exists because the others are busy, so + // neither its 2 s keepalive nor the ordinary --sse-ping could ever fire for it. Waiting + // until a fixed point instead makes the timeout mean what every caller reads it as. + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout); + while (true) { std::unique_lock lock(mutex_results); @@ -459,7 +471,7 @@ server_task_result_ptr server_response::recv_with_timeout(const std::unordered_s } } - std::cv_status cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout)); + std::cv_status cr_res = condition_results.wait_until(lock, deadline); if (!running) { RES_DBG("%s : queue result stop\n", __func__); std::terminate(); // we cannot return here since the caller is HTTP code From 888603d03578749219b190fd1dd0ed9554496fdb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:21:36 +0000 Subject: [PATCH 31/81] server: do not leave an issued async park holding the room the decode needs A review of #192 pointed at the victim loop in update_preemption(). When the pool has no room for the step about to be built and no park is in flight, the loop issues the victim's asynchronous park and breaks. The cells are held until the copy lands, so the batch is built into a pool that has not got smaller, llama_decode returns 1, and the retry ladder halves n_batch to 1 in microseconds without ever polling the copy, ending in "Context size has been exceeded" for every slot. The synchronous path freed the cells before returning, so it could not do this. I could not reproduce it. Six live rounds on the 4B at -c 8192, four chats with 1000-token prompts and 2048 tokens each, with the fourth chat's prompt held back 20 s so it arrives into a pool the other three have filled, exact mode on and off, on a binary without this change: 4 of 4 every round, no context errors, and the retry ladder was not entered once ("failed to find free space" appears zero times in both server logs). The reason is that preempt_kv_reserve() counts an incoming prompt chunk before it is allocated, so the planner crosses the lookahead threshold an iteration before the pool actually fills, and every park in those runs was issued with the 80 cells of asynchronous runway still ahead of it, never at the hard threshold this is about. Committing it anyway, because the described state is real even if these workloads do not reach it, and the change is inert unless it is reached: * The victim loop goes round again instead of leaving, but only when n_used + PREEMPT_N_MARGIN > n_cells, that is when there is no room for the step itself rather than merely less than the asynchronous lookahead wants. The next pass reaches preempt_wait_in_flight() and waits for the park just issued, which is what that function was written for and no worse than the synchronous path. Short of the lookahead only, it still breaks, because parking early and letting the copy run beside the decode is the entire point of #192. * On llama_decode returning 1, an outstanding park is waited for before any batch width is given up. Halving n_batch returns no cells, so without this the ladder can walk to n_batch == 1 and end every request while the room it needed was one event query away. Safe at that point because the slot was detached before the batch was built, so completing its park cannot change what is about to be retried; that is also why update_preemption() itself is not called from here. New test, test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_ everything: three slots generating near the ceiling and a fourth request whose prompt does not fit in what is left, which is the shape the existing tests miss because their victim holds almost no cells. The three are sized to oversubscribe the pool between them so the pressure does not depend on when the fourth arrives. It is kept for the shape it covers rather than as an attribution: it passes either way, and the attribution above was done at live scale. --- tools/server/server-context.cpp | 33 ++++++++++++++ tools/server/tests/unit/test_preempt.py | 57 +++++++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 007f3298cb0..5509d4ea06c 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3639,6 +3639,27 @@ struct server_context_impl { victim->preempt_state_size() / (1024.0 * 1024.0), preempt_kv_used(), n_cells, n_used, victim->n_preempt); + + // [TAG_PREEMPT_ASYNC] Whether we may leave now depends on which of the two + // thresholds we are under. + // + // Short of the lookahead only: there is still room for the step about to be + // built, the park is early by design, and leaving is the whole point -- the + // copy runs beside the decode and update_preempt_copies() collects it next + // iteration. + // + // Out of room for the step itself: the cells are held until the copy lands, + // so leaving now builds a batch into a pool that has not got smaller. The + // decode fails, and the retry ladder halves n_batch to 1 without ever + // polling the copy, ending in "Context size has been exceeded" for every + // slot. The synchronous path did not have this problem because it returned + // the cells before it returned. Go round instead: the next pass reaches + // preempt_wait_in_flight() and waits for the park just issued, which is no + // worse than the synchronous path and is what it was written for. + if (n_used + PREEMPT_N_MARGIN > n_cells) { + continue; + } + break; } @@ -4593,6 +4614,18 @@ struct server_context_impl { } } + // [TAG_PREEMPT_ASYNC] Before giving up any batch width: a park that has been + // issued and not yet landed is holding cells that are already spoken for, and + // waiting for it returns them. Halving the batch returns nothing, so without + // this the ladder can run all the way down to n_batch == 1 and end every + // request while the room it needed was moments from arriving. Safe from here + // because the slot was detached before this batch was built, so completing its + // park cannot change what the batch about to be retried contains. + if (ret == 1 && preempt_wait_in_flight()) { + SRV_WRN("%s", "waited for an in-flight park before retrying the decode\n"); + return false; // retry at the same batch size, with the cells it freed + } + // retry with half the batch size to try to find a free slot in the KV cache if (!try_clear_idle_slots()) { n_batch /= 2; diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 1aba01751bd..4c3e8c465d8 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -447,3 +447,60 @@ def test_no_preempt_async_falls_back_to_the_synchronous_path(): # the synchronous path still parks and resumes assert text.count("preempted on request") >= 6 assert text.count("resumed after") >= 6 + + +def test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_everything(): + # [TAG_PREEMPT_ASYNC] The case the async path made worse than the synchronous one, and + # that the existing tests miss because their victim holds almost no cells. + # + # Three slots are well into generating when a fourth request arrives whose prompt does + # not fit in what is left. update_preemption() picks a victim and issues its park, but + # an asynchronous park does not return the cells before update_slots() carries on. If + # the loop leaves at that point, the batch is built into a pool that has not got any + # smaller, llama_decode returns 1, and the retry ladder halves n_batch to 1 in + # microseconds without ever polling the copy -- ending every request with "Context size + # has been exceeded" while the room it wanted was one event query away. + # + # Pass is what the synchronous path gave: a park, and all four requests finish. + global server + server.n_ctx = 512 + server.n_gpu_layer = 99 + server.n_slots = 4 + server.start() + log = LogReader(server.log_path) + + prompt_a, n_a = _prompt_of_about(100, "Alpha") + prompt_b, n_b = _prompt_of_about(100, "Bravo") + prompt_c, n_c = _prompt_of_about(100, "Charlie") + prompt_d, n_d = _prompt_of_about(150, "Delta") + + # A, B and C oversubscribe the pool between them, so the pressure does not depend on + # when D arrives, and every occupant is holding real cells rather than the handful the + # other tests park. Each of the four still fits on its own. + n_predict_abc = 130 + n_predict_d = 40 + assert max(n_a, n_b, n_c) + n_predict_abc < 512 and n_d + n_predict_d < 512 + assert n_a + n_b + n_c + 3 * n_predict_abc > 512 + + def _late(n_predict, prompt): + # D's prompt arrives into a pool the other three have already grown into; this + # model decodes about 120 tokens a second, so they are all still running + time.sleep(0.25) + return _complete(n_predict, prompt) + + results = parallel_function_calls([ + (_complete, (n_predict_abc, prompt_a)), + (_complete, (n_predict_abc, prompt_b)), + (_complete, (n_predict_abc, prompt_c)), + (_late, (n_predict_d, prompt_d)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted" in text + + for i, res in enumerate(results): + assert res.status_code == 200, (i, res.body) + for i in range(3): + assert results[i].body["timings"]["predicted_n"] == n_predict_abc + assert results[3].body["timings"]["predicted_n"] == n_predict_d From 5a791e03c4b28052475b04b1bb15b29fd244528c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 07:05:06 +0000 Subject: [PATCH 32/81] server: make the result queue timeout a deadline, so a parked stream is kept alive Found by putting #190 and #192 together and then looking for the keepalive that #190 promises. It never arrives. Live, the 4B on one B200, two streaming completions that do not fit together so one is parked until the other finishes, every SSE line timestamped as it arrives: 6.19s B : preempted 16.82s B : resumed A 10.63 s silence on a stream whose whole point is that it says ": preempt- keepalive" every 2 s. Four-chat runs at -c 8192 and -c 4096 show the same: parks of up to 14.59 s by the server's own "resumed after" line, and not one keepalive on any stream in any run. server_response::send() notify_all()s a single condition variable for every result of every task, and server_response::recv_with_timeout() waited with wait_for(), which restarts on every wakeup. A reader waiting on a task that is producing nothing is therefore woken by every token every other task produces, and its wait_for() never elapses. On a server with any traffic at all the timeout is not a timeout: whoever waits for a quiet task waits indefinitely. A parked slot is the worst possible case for this, because a slot is only ever parked while the others are busy, so the keepalive was unreachable by construction. The same applies to the ordinary --sse-ping, which likewise only fired on an otherwise idle server, and to the should_stop polling in server_response_reader::next(), whose own comment says it happens every polling_interval_seconds and did not. Compute the deadline once and wait_until() it. Spurious wakeups then re-check the queue and go back to waiting for the same instant, which is what every caller already reads the argument as meaning. After, the same two streams: 6.26s B : preempted 9.26s B : preempt-keepalive 12.26s B : preempt-keepalive 14.26s B : preempt-keepalive 16.75s B : resumed Three keepalives across a 10.49 s park, at the 2 s period plus the reader's 1 s polling granularity. The probe is scripts/integ_keepalive_probe.py. No harness test: stories260K generates several hundred tokens a second, and at the context sizes the harness uses a park lasts two or three seconds, which is the keepalive period itself. Every sizing I tried either parked for milliseconds at a time as the pool oscillated around full, or did not park at all. A test that straddles the period it is testing would be worse than none, so the regression is pinned by the live probe above. --- tools/server/server-queue.cpp | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d8..bd74db89408 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -448,6 +448,18 @@ server_task_result_ptr server_response::recv(const std::unordered_set & id_ } server_task_result_ptr server_response::recv_with_timeout(const std::unordered_set & id_tasks, int timeout) { + // [TAG_PREEMPT] The timeout is a deadline, not a per-wait duration. + // + // send() notify_all()s on one condition variable for every result of every task, so a + // reader waiting on a task that is producing nothing is woken by every token every other + // task produces. With wait_for() each of those wakeups restarted the wait, and on a busy + // server the timeout was never reached at all: whoever was waiting for a quiet task + // waited forever, however small the timeout they asked for. That is exactly the + // situation of a parked slot, which by definition exists because the others are busy, so + // neither its 2 s keepalive nor the ordinary --sse-ping could ever fire for it. Waiting + // until a fixed point instead makes the timeout mean what every caller reads it as. + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout); + while (true) { std::unique_lock lock(mutex_results); @@ -459,7 +471,7 @@ server_task_result_ptr server_response::recv_with_timeout(const std::unordered_s } } - std::cv_status cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout)); + std::cv_status cr_res = condition_results.wait_until(lock, deadline); if (!running) { RES_DBG("%s : queue result stop\n", __func__); std::terminate(); // we cannot return here since the caller is HTTP code From 712bee75c69120f869b9ae68e45ffe53978522e6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 08:34:48 +0000 Subject: [PATCH 33/81] kv-cache: maintain page ownership instead of rebuilding it per ubatch find_slot() and set_input_pages() each rebuilt the (sequence, logical page) to physical page map from a scan of every live cell into an std::map, so the paged allocator did O(cells log pages) work twice per ubatch, twice per decode step, and the CPU cost grew with the size of the pool rather than with the number of pages in it. Keep the ownership in a flat vector with one entry per physical page, claimed in apply_ubatch() as cells are placed and marked dirty by the paths that remove them, seq_rm(), seq_keep() and clear(). prepare() snapshots it alongside the cells so that undoing a speculative placement puts back what the allocator knew rather than forcing a rebuild. Both readers now build their lookup from one entry per page: 32 entries at -c 8192 and 256 at -c 65536, against 8192 and 65536 cells. find_slot also stops allocating a cells-sized bitmap per call. Nothing about the placement policy changes, and the derived-from-live-cells rebuild is still there and still authoritative: LLAMA_KV_CACHE_DEBUG=1 runs it on every read and asserts that the maintained ownership says exactly what the cells say. Four chats, 937 token prompts, 1536 tokens each, ignore_eos, speculation off, --parallel 4 --kv-unified --flash-attn on -ngl 99, three interleaved pairs on one B200, medians of four-chat aggregate decode tok/s: -c 8192 off 151.03 exact 136.70 0.905 before -c 8192 off 150.83 exact 137.94 0.915 after -c 65536 off 158.34 exact 136.29 0.861 before -c 65536 off 159.48 exact 143.62 0.901 after So at 32 pages it is worth about a point, and at 256 pages it is worth four, which is what a cost that followed the cell count and now follows the page count should look like. With LLAMA_KV_CACHE_DEBUG=1 and LLAMA_SERVER_PREEMPT_EVERY=32, which parks and restores every slot every 32 tokens and so exercises every path that marks the ownership dirty, no assert fires and P0 stays byte identical to its solo reference. --- scripts/batchinv/divergence.py | 2 +- src/llama-kv-cache.cpp | 149 +++++++++++++++++++++++++++------ src/llama-kv-cache.h | 22 +++++ 3 files changed, 147 insertions(+), 26 deletions(-) diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py index 3e554f465b7..ec2c3a00467 100644 --- a/scripts/batchinv/divergence.py +++ b/scripts/batchinv/divergence.py @@ -9,7 +9,7 @@ # decides whether a run labelled as the mode-off reference actually was one, so it is not optional. RECORDED_ENV = ("LLAMA_EXACT_CONCURRENCY", "GGML_CUDA_BATCH_INVARIANT", "GGML_CUDA_BATCH_INVARIANT_MAX_COLS", "LLAMA_SERVER_PREEMPT_EVERY", - "CUDA_VISIBLE_DEVICES") + "LLAMA_KV_CACHE_DEBUG", "LLAMA_BATCH_DEBUG", "CUDA_VISIBLE_DEVICES") MODEL_REL = "models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf" diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 16c04fb6e3b..06729158133 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include static bool ggml_is_power_of_2(int n) { @@ -426,7 +427,76 @@ llama_kv_cache::llama_kv_cache( debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0; } +// [TAG_EXACT_CONCURRENCY] +void llama_kv_cache::exact_pages_rebuild() const { + const auto & cells = v_cells[0]; + + exact_page_owner.assign(cells.size()/exact_page_size, exact_page{}); + + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.is_empty(i)) { + continue; + } + + GGML_ASSERT(cells.seq_count(i) == 1); + + const auto pos = cells.pos_get(i); + + GGML_ASSERT(pos >= 0 && uint32_t(pos)%exact_page_size == i%exact_page_size); + + const exact_page cur { cells.seq_get(i), llama_pos(pos/(llama_pos) exact_page_size) }; + + auto & owner = exact_page_owner[i/exact_page_size]; + + GGML_ASSERT(owner.seq < 0 || (owner.seq == cur.seq && owner.lpg == cur.lpg)); + + owner = cur; + } + + exact_page_owner_dirty = false; +} + +// [TAG_EXACT_CONCURRENCY] +void llama_kv_cache::exact_pages_sync() const { + if (exact_page_owner_dirty) { + exact_pages_rebuild(); + + return; + } + + if (debug > 0) { + // the incrementally maintained ownership has to say what the cells say + const auto kept = exact_page_owner; + + exact_pages_rebuild(); + + GGML_ASSERT(kept.size() == exact_page_owner.size()); + + for (size_t p = 0; p < kept.size(); ++p) { + GGML_ASSERT(kept[p].seq == exact_page_owner[p].seq && kept[p].lpg == exact_page_owner[p].lpg); + } + } +} + +// [TAG_EXACT_CONCURRENCY] +void llama_kv_cache::exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos) { + if (exact_page_owner_dirty || exact_page_owner.empty()) { + // the next sync rebuilds from the cells anyway + return; + } + + const exact_page cur { seq, llama_pos(pos/(llama_pos) exact_page_size) }; + + auto & owner = exact_page_owner[idx/exact_page_size]; + + GGML_ASSERT(owner.seq < 0 || (owner.seq == cur.seq && owner.lpg == cur.lpg)); + + owner = cur; +} + void llama_kv_cache::clear(bool data) { + exact_page_owner_dirty = true; + for (uint32_t s = 0; s < n_stream; ++s) { v_cells[s].reset(); v_heads[s] = 0; @@ -445,6 +515,9 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { return true; } + // [TAG_EXACT_CONCURRENCY] a removal can empty a page, which only the cells know + exact_page_owner_dirty = true; + // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG] GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); @@ -614,6 +687,9 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) { return; } + // [TAG_EXACT_CONCURRENCY] as in seq_rm, this can empty pages + exact_page_owner_dirty = true; + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); auto & cells = v_cells[seq_to_stream[seq_id]]; @@ -848,6 +924,10 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector v_heads_old; // old positions of the heads, before placing the ubatch std::vector v_cells; // copy of the old cells, before placing the ubatch + + // [TAG_EXACT_CONCURRENCY] page ownership before placing the ubatch, so that undoing the + // speculative placement does not force a rebuild from every cell on the next ubatch + std::vector exact_page_owner_old; }; // remember the old state of the cells so we can restore it in the end @@ -868,7 +948,7 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vectorv_cells[s]); head = it->v_heads_old[s]; } + + // [TAG_EXACT_CONCURRENCY] the speculative placements are being undone behind the + // allocator's back. Put back what it knew before, unless something during the placement + // removed cells as well, in which case only the cells can say what is left. + if (!exact_page_owner_dirty) { + exact_page_owner = it->exact_page_owner_old; + } } if (!success) { @@ -1055,23 +1142,26 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, } if (exact_pages) { - // Reconstruct page ownership from live cells. Empty pages are immediately reusable; - // prepare() can roll back its speculative allocations without a second metadata log. + // Page ownership is maintained as cells are placed and invalidated when they are removed, + // so the allocator reads one entry per physical page rather than scanning every cell. The + // claims this call makes are local: prepare() can still roll back its speculative + // placements, and empty pages stay immediately reusable. const auto & cells = v_cells[0]; + + exact_pages_sync(); + using page_key = std::pair; + + std::vector owner = exact_page_owner; std::map pages; - std::vector occupied(cells.size()/exact_page_size, false); - std::vector assigned(cells.size(), false); - for (uint32_t i = 0; i < cells.size(); ++i) { - if (cells.is_empty(i)) { continue; } - GGML_ASSERT(cells.seq_count(i) == 1); - const auto pos = cells.pos_get(i); - GGML_ASSERT(pos >= 0 && uint32_t(pos)%exact_page_size == i%exact_page_size); - const page_key key {cells.seq_get(i), pos/exact_page_size}; - auto ins = pages.emplace(key, i/exact_page_size); - GGML_ASSERT(ins.first->second == i/exact_page_size); - occupied[i/exact_page_size] = true; + + for (uint32_t p = 0; p < owner.size(); ++p) { + if (owner[p].seq >= 0) { + pages.emplace(page_key {owner[p].seq, owner[p].lpg}, p); + } } + + std::set assigned; slot_info res {0, 0, {0}, {{}}}; for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { GGML_ASSERT(ubatch.n_seq_id[i] == 1 && ubatch.pos[i] >= 0); @@ -1081,15 +1171,14 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, // Round-robin free-page search deliberately permits nonmonotonic physical order. uint32_t page = v_heads[0]/exact_page_size; uint32_t tested = 0; - while (tested < occupied.size() && occupied[page%occupied.size()]) { ++page; ++tested; } - if (tested == occupied.size()) { return {}; } - page %= occupied.size(); - occupied[page] = true; + while (tested < owner.size() && owner[page%owner.size()].seq >= 0) { ++page; ++tested; } + if (tested == owner.size()) { return {}; } + page %= owner.size(); + owner[page] = exact_page {key.first, key.second}; it = pages.emplace(key, page).first; } const uint32_t idx = it->second*exact_page_size + ubatch.pos[i]%exact_page_size; - if (!cells.is_empty(idx) || assigned[idx]) { return {}; } - assigned[idx] = true; + if (!cells.is_empty(idx) || !assigned.insert(idx).second) { return {}; } res.idxs[0].push_back(idx); } if (cont && !res.is_contiguous()) { return {}; } @@ -1274,6 +1363,13 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & for (int32_t s = 0; s < ubatch.n_seq_id[i]; s++) { cells.seq_add(idx, ubatch.seq_id[i][s]); } + + // [TAG_EXACT_CONCURRENCY] the page this cell belongs to is now owned by its sequence + if (exact_pages) { + GGML_ASSERT(ubatch.n_seq_id[i] == 1); + + exact_pages_claim(idx, ubatch.seq_id[i][0], ubatch.pos[i]); + } } } @@ -1384,12 +1480,15 @@ ggml_tensor * llama_kv_cache::build_input_pages(ggml_context * ctx, const llama_ void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const { GGML_ASSERT(exact_pages && dst->ne[1] == ubatch->n_tokens); + + // [TAG_EXACT_CONCURRENCY] one entry per physical page, not one per cell + exact_pages_sync(); + std::map> pages; - const auto & cells = v_cells[0]; - for (uint32_t i = 0; i < cells.size(); ++i) { - if (!cells.is_empty(i)) { - GGML_ASSERT(cells.seq_count(i) == 1); - pages[cells.seq_get(i)][cells.pos_get(i)/exact_page_size] = i/exact_page_size; + for (uint32_t p = 0; p < exact_page_owner.size(); ++p) { + const auto & owner = exact_page_owner[p]; + if (owner.seq >= 0) { + pages[owner.seq][owner.lpg] = p; } } std::vector data(ggml_nelements(dst), -1); diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index fa257422f02..8228752d38b 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -240,6 +240,28 @@ class llama_kv_cache : public llama_memory_i { static constexpr uint32_t exact_page_size = 256; bool exact_pages = false; + // [TAG_EXACT_CONCURRENCY] + // Which (sequence, logical page) owns each physical page of the pool; seq < 0 means the page is + // free. Kept current as cells are placed, and marked dirty by the paths that remove cells, so + // that find_slot() and set_input_pages() read one entry per page instead of rebuilding the map + // from every live cell twice per ubatch. Mutable because set_input_pages() is const. + struct exact_page { + llama_seq_id seq = -1; + llama_pos lpg = -1; + }; + + mutable std::vector exact_page_owner; + mutable bool exact_page_owner_dirty = true; + + // bring exact_page_owner up to date; rebuilds only when a removal marked it dirty + void exact_pages_sync() const; + + // recompute it from the live cells + void exact_pages_rebuild() const; + + // record that a cell of (seq, pos) now lives at physical cell idx + void exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos); + bool v_trans = true; // the value tensor is transposed const uint32_t n_seq_max = 1; From a2f9c081f4a9fe3fa5bc368d6a0646fb6a083e5a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:12:58 +0000 Subject: [PATCH 34/81] server: refuse n > 1 under exact concurrency instead of aborting in the cache A review of #194 pointed at the new GGML_ASSERT in llama_kv_cache::seq_cp, and it is right. Reproduced on this branch with the 4B on one B200: LLAMA_EXACT_CONCURRENCY=1, POST /completion {"n": 2} -> llama-kv-cache.cpp:458: GGML_ASSERT(!exact_pages || seq_id_src == seq_id_dst) failed, through server_context_impl::decode -> common_memory::seq_cp -> process aborted, the next request gets connection refused With the mode off the same request is served normally, so this is reachable by any client of an exact-mode server and takes every other request on the machine with it. Two changes: The server refuses the request. n_cmpl > 1 works by copying the parent's cells to a second sequence id, and exact mode gives a page to one sequence, so there is nowhere for that copy to land. Rejecting it where the task is built turns it into a 400 with a reason. The check reads LLAMA_EXACT_CONCURRENCY from the environment, the same way the KV cache, the batch splitter and the CUDA backend each do, because the answer is needed before a context exists and the mode has no other representation. The cache stops aborting. seq_cp, seq_add and seq_div log an error and return instead of asserting, so a caller this branch does not know about degrades to a refused operation rather than killing the server. The guards also move below the shared-cells early return, which the asserts sat above: a draft cache forwards these calls and copies nothing of its own, and it should not be judged by a rule about cells it does not own. After: the n=2 request returns 400 on both /completion and /v1/completions, the server stays up, and a following ordinary request returns 200. --- src/llama-kv-cache.cpp | 55 +++++++++++++++++++-------------- tools/server/server-context.cpp | 25 +++++++++++++++ 2 files changed, 57 insertions(+), 23 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 06729158133..8ecf701a557 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -582,19 +582,24 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { } void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { - // [TAG_EXACT_CONCURRENCY] a page belongs to one (sequence, position/256) pair, so two sequences - // cannot share physical cells. Refuse the copy rather than abort the process: this is reachable - // from a request parameter. - if (exact_pages && seq_id_src != seq_id_dst) { - LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so cells cannot be shared between " - "sequences: ignoring the copy from seq %d to seq %d\n", __func__, seq_id_src, seq_id_dst); - return; - } // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; } + // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so cells cannot be shared + // between two of them. Refuse the operation rather than abort the process: a server + // rejects the request that would reach here (n_cmpl > 1), and any caller this does not + // cover degrades to a failed copy it can report instead of killing every other request + // on the machine. Placed after the shared-cells return so a draft cache, which copies + // nothing of its own, is unaffected. + if (exact_pages && seq_id_src != seq_id_dst) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between " + "sequences (%d -> %d); ignoring the copy\n", + __func__, seq_id_src, seq_id_dst); + return; + } + GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size()); GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size()); @@ -712,19 +717,21 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) { } void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { - // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so a - // shift would have to move the cells too. get_can_shift() reports this so that --context-shift - // and --cache-reuse are turned off at load; this is the guard for the library API. - if (exact_pages && shift != 0) { - LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so positions cannot be shifted: " - "ignoring the shift of %d on seq %d\n", __func__, shift, seq_id); - return; - } // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; } + // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo the + // page size, so shifting positions would put every cell of the sequence in the wrong + // place. Context shift is unsupported in exact mode; say so rather than abort. + if (exact_pages && shift != 0) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions " + "(seq %d, shift %d); ignoring the shift\n", + __func__, seq_id, shift); + return; + } + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1"); @@ -770,18 +777,20 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll } void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { - // [TAG_EXACT_CONCURRENCY] same reason as seq_add: the offset inside a page is derived from the - // position, so dividing the positions would leave every cell in the wrong slot. - if (exact_pages && d != 1) { - LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so positions cannot be divided: " - "ignoring the division by %d on seq %d\n", __func__, d, seq_id); - return; - } // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; } + // [TAG_EXACT_CONCURRENCY] same reason as seq_add: dividing positions breaks the + // identity between a cell's position and its offset inside its page. + if (exact_pages && d != 1) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions " + "(seq %d, d %d); ignoring the division\n", + __func__, seq_id, d); + return; + } + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1"); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 6723c51397e..9f3d1c53eb5 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -37,6 +37,19 @@ constexpr int HTTP_POLLING_SECONDS = 1; +// [TAG_EXACT_CONCURRENCY] the knob is read from the environment by the KV cache, the batch +// splitter and the CUDA backend independently, because it has to be answered before a +// context exists. The server needs the same answer to refuse the one request shape the mode +// cannot serve, so it reads it the same way rather than growing a public API for it. +static bool server_exact_concurrency() { + static const bool enabled = []() { + const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); + return val && atoi(val) != 0; + }(); + + return enabled; +} + static common_speculative_output_limits server_output_limits(const common_params & params) { if (params.embedding || (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) { @@ -4686,6 +4699,18 @@ std::unique_ptr server_routes::handle_completions_impl( task.params.oaicompat_cmpl_id = completion_id; task.params.oaicompat_model = meta->model_name; + // [TAG_EXACT_CONCURRENCY] the children of an n_cmpl > 1 task are served by + // copying the parent's cells to another sequence id, and exact mode gives a KV + // page to one sequence, so there is nothing for that copy to land in. Refuse + // the request here, where it becomes a 400 the client can read, rather than + // letting it reach seq_cp with nothing to do. + if (task.params.n_cmpl > 1 && server_exact_concurrency()) { + throw std::runtime_error( + "n > 1 is not supported while LLAMA_EXACT_CONCURRENCY is set: each " + "completion needs its own sequence, and in exact mode a KV page belongs " + "to a single sequence. Send n separate requests, or unset the variable."); + } + // prepare child tasks if (task.params.n_cmpl > 1) { int n_children = task.params.n_cmpl - 1; From 72aca44b5e6f7b8bb3a5bf33bcbad16790435971 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 08:47:59 +0000 Subject: [PATCH 35/81] memory: let the server ask how many cells an allocation takes Under exact concurrency the KV cache hands out 256-cell pages, one to each (sequence, position / 256) pair, so a sequence can hold up to 255 cells that nobody else can be given. A preemption planner that counts tokens does not see those cells: it believes there is room, never parks anybody, and the pool fills until every request ends in the old context error. llama_memory_i gains alloc_granularity(), defaulting to 1 so no module that allocates a cell per token changes; the KV cache returns its page size under exact mode and the hybrid memory forwards to its attention half. llama_memory_alloc_granularity() exposes it. The server side of this, rounding its planner figures by that value, follows separately. --- include/llama.h | 8 ++++++++ src/llama-context.cpp | 8 ++++++++ src/llama-kv-cache.cpp | 7 +++++++ src/llama-kv-cache.h | 3 +++ src/llama-memory-hybrid.cpp | 6 ++++++ src/llama-memory-hybrid.h | 2 ++ src/llama-memory.h | 9 +++++++++ 7 files changed, 43 insertions(+) diff --git a/include/llama.h b/include/llama.h index a04177f9f7d..43d79b40a2a 100644 --- a/include/llama.h +++ b/include/llama.h @@ -795,6 +795,14 @@ extern "C" { // Check if the memory supports shifting LLAMA_API bool llama_memory_can_shift(llama_memory_t mem); + // [TAG_EXACT_CONCURRENCY] Cells the memory allocates in one indivisible unit. + // + // 1 in every ordinary configuration. Larger where a mode places cells in blocks, and + // then a sequence of n tokens occupies round_up(n, granularity) cells. A caller that + // decides whether the pool has room by counting tokens has to round the same way, or it + // will believe there is space that cannot be handed out. + LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem); + // // State / sessions // diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 66940d4fc61..3ab84de780c 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4030,6 +4030,14 @@ bool llama_memory_can_shift(llama_memory_t mem) { return mem->get_can_shift(); } +uint32_t llama_memory_alloc_granularity(llama_memory_t mem) { + if (!mem) { + return 1; + } + + return mem->alloc_granularity(); +} + // llama state API // deprecated diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 8ecf701a557..dced34f6a14 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1410,6 +1410,13 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & } } +uint32_t llama_kv_cache::alloc_granularity() const { + // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so a + // sequence holding n tokens holds round_up(n, exact_page_size) cells: its tail page is + // charged in full whether or not it is full. + return exact_pages ? exact_page_size : 1; +} + bool llama_kv_cache::get_can_shift() const { // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so the // paged pool cannot shift positions. Reporting it here is what makes the server disable diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 8228752d38b..af3a04be39d 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -131,6 +131,9 @@ class llama_kv_cache : public llama_memory_i { bool get_can_shift() const override; + // [TAG_EXACT_CONCURRENCY] the page size under exact mode, 1 otherwise + uint32_t alloc_granularity() const override; + void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 4ebd476aa8b..48fdea8b49c 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -142,6 +142,12 @@ bool llama_memory_hybrid::get_can_shift() const { return mem_attn->get_can_shift(); } +uint32_t llama_memory_hybrid::alloc_granularity() const { + // the recurrent half holds one state per sequence rather than per token, so the + // attention half is the one whose cells a caller is planning capacity for + return mem_attn->alloc_granularity(); +} + void llama_memory_hybrid::clear(bool data) { mem_attn->clear(data); mem_recr->clear(data); diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb7499..70ba19ca323 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -58,6 +58,8 @@ class llama_memory_hybrid : public llama_memory_i { bool get_can_shift() const override; + uint32_t alloc_granularity() const override; + void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645..51539a03919 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -100,6 +100,15 @@ struct llama_memory_i { // getters virtual bool get_can_shift() const = 0; + // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit. + // + // 1 for every module that allocates a cell per token, which is all of them unless a mode + // is on that allocates in larger blocks. Where it is larger, a sequence of n tokens + // occupies round_up(n, granularity) cells, and a caller that plans pool capacity by + // counting tokens will believe there is room that does not exist. Not pure, so a module + // that has never heard of this inherits the answer that has always been true of it. + virtual uint32_t alloc_granularity() const { return 1; } + // // ops // From da8556d316d2f7bb66b690d87d089019e5a720ce Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 09:23:52 +0000 Subject: [PATCH 36/81] server: plan the kv pool in cells rather than tokens Exact mode and #184's preemption planner cannot both be right about how full the pool is, and on this branch they are not. Reproduced with the 4B on one B200, LLAMA_EXACT_CONCURRENCY=1, four chats with 1000-token prompts generating 2048 tokens each at --parallel 4 --kv-unified -c 8192 --spec-type draft-mtp --spec-draft-n-max 2, no forced-park knob, three rounds: 0 of 4, 0 of 4 and 1 of 4 completions, 12 "Context size has been exceeded", 0 parks and 0 restores. Nothing was ever parked. preempt_kv_used(), preempt_n_need() and preempt_kv_reserve() count tokens, and exact mode's allocator hands out 256-cell pages, one page to one (sequence, position / 256) pair. Four sequences can therefore be holding up to 1020 cells that no other sequence can be given, and the planner, seeing room in tokens that find_slot cannot find in pages, never reaches the threshold that would park anybody. The retry ladder then halves n_batch to 1 and ends every request, which is the pre-#184 behaviour that preemption exists to remove. Ask the memory how it allocates instead of assuming. The server reads llama_memory_alloc_granularity() once at load and rounds: preempt_kv_used() charges every slot's tail page in full, because a page belongs to one sequence however little of it is used preempt_n_need() rounds what a resume must be given, since a restore takes fresh pages preempt_kv_reserve() reserves the cells the next step ADDS rather than its tokens, because on a rounded used figure a step is free until it crosses a page boundary and costs a whole page when it does, and that crossing is the only moment the pool can run out preempt_n_margin() rounds the spare cells up to a page, since a margin of eight is no margin at all where a step can cost 256 With a granularity of 1 every one of these is the arithmetic it was, which is pinned by static assertions on the two rounding helpers rather than left to be read: at 1 they are the identity, so nothing changes with the mode off. After, same configuration and three rounds: 4 of 4 completions each round, 3, 4 and 4 parks and the same number of restores, no context errors, and P0 byte-identical to its solo run in all three. The same run with exact mode off is also 4 of 4 with 3, 2 and 2 parks, and its planner figures are still the token counts they always were. LLAMA_SERVER_PREEMPT_GRANULARITY overrides the figure the memory reports. It is a test knob, next to LLAMA_SERVER_PREEMPT_EVERY: the paged attention kernel needs a head size of 256, which the harness model does not have, so this is the only way to reach the paged arithmetic from tools/server/tests. The new test drives two slots over a 256-cell pool at a granularity of 64 and asserts every figure the planner logs is a whole number of blocks. Counting tokens the same run logs kv 119/256 and wanted 249; counting cells it logs kv 64/256 and wanted 256. --- tools/server/server-context.cpp | 117 ++++++++++++++++++++++-- tools/server/tests/unit/test_preempt.py | 44 +++++++++ 2 files changed, 152 insertions(+), 9 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 9f3d1c53eb5..947ec514772 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -93,6 +93,45 @@ constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked +// [TAG_EXACT_CONCURRENCY] The planner above counts cells, not tokens, because the two are not +// the same number under every mode. llama_memory_alloc_granularity() reports how many cells the +// pool hands out at a time: 1 in every ordinary configuration, and the exact concurrency page +// size when that mode is on, where one page belongs to one (sequence, position / page) pair and +// a sequence of n tokens therefore occupies round_up(n, page) cells. Four sequences can be +// holding up to 4 * (page - 1) cells that nobody else can be given, and a planner counting +// tokens sees room in the pool that find_slot cannot find in pages: it never reaches the +// threshold that would park anybody, the retry ladder halves n_batch to 1, and every request +// ends in the context error that preemption exists to remove. + +// cells a run of n_tokens occupies when the pool allocates g at a time +static constexpr int32_t preempt_n_cells_g(int32_t n_tokens, int32_t g) { + return (g <= 1 || n_tokens <= 0) ? n_tokens : ((n_tokens + g - 1) / g) * g; +} + +// cells a run of n_tokens has to be given for a step of n_step more: nothing until the step +// crosses a page boundary, a whole page when it does +static constexpr int32_t preempt_n_cells_step_g(int32_t n_tokens, int32_t n_step, int32_t g) { + return preempt_n_cells_g(n_tokens + n_step, g) - preempt_n_cells_g(n_tokens, g); +} + +// At a granularity of 1 both are the identity, so every figure the planner computes is exactly +// the arithmetic it did before it started asking the memory how it allocates, and nothing +// changes in any configuration that does not page. +static_assert(preempt_n_cells_g(0, 1) == 0 && preempt_n_cells_g(1, 1) == 1 && + preempt_n_cells_g(8191, 1) == 8191 && preempt_n_cells_g(-3, 1) == -3, + "at a granularity of 1 a run of n tokens has to cost exactly n cells"); +static_assert(preempt_n_cells_step_g(0, 1, 1) == 1 && preempt_n_cells_step_g(8191, 1, 1) == 1 && + preempt_n_cells_step_g(1000, 512, 1) == 512, + "at a granularity of 1 a step of n tokens has to cost exactly n cells"); + +// and the page arithmetic itself, so the rounding cannot be changed by accident +static_assert(preempt_n_cells_g(1, 256) == 256 && preempt_n_cells_g(256, 256) == 256 && + preempt_n_cells_g(257, 256) == 512, + "a tail page is charged in full"); +static_assert(preempt_n_cells_step_g(255, 1, 256) == 0 && preempt_n_cells_step_g(256, 1, 256) == 256 && + preempt_n_cells_step_g(256, 257, 256) == 512, + "a step is free until it crosses a page boundary and costs whole pages when it does"); + struct server_slot; // forward declaration struct server_batch { @@ -1415,6 +1454,28 @@ struct server_context_impl { } } + // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates, rather than assume a cell per + // token. 1 in every ordinary configuration, so this changes nothing unless a mode that + // places cells in blocks is on. + { + preempt_alloc_granularity = (int32_t) std::max(1u, llama_memory_alloc_granularity(llama_get_memory(ctx_tgt))); + + // a test knob: the paged attention kernel only supports a head size of 256, so a + // harness model cannot turn exact concurrency on, and this is the only way to reach + // the paged arithmetic of the planner from the server tests + const char * LLAMA_SERVER_PREEMPT_GRANULARITY = getenv("LLAMA_SERVER_PREEMPT_GRANULARITY"); + + if (LLAMA_SERVER_PREEMPT_GRANULARITY) { + preempt_alloc_granularity = std::max(1, atoi(LLAMA_SERVER_PREEMPT_GRANULARITY)); + + SRV_WRN("LLAMA_SERVER_PREEMPT_GRANULARITY = %d (test knob: planning the kv pool in blocks of %d cells)\n", + preempt_alloc_granularity, preempt_alloc_granularity); + } else if (preempt_alloc_granularity > 1) { + SRV_INF("preemption: the kv pool allocates %d cells at a time, planning in pages\n", + preempt_alloc_granularity); + } + } + { const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; @@ -2868,6 +2929,30 @@ struct server_context_impl { // uninterrupted one is the preemption's fault and nothing else's. int32_t preempt_test_every = 0; + // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load from the + // memory itself: 1 in every ordinary configuration, the page size under exact concurrency. + // Everything below plans in cells because of it. LLAMA_SERVER_PREEMPT_GRANULARITY overrides + // it, which is how the harness reaches the paged arithmetic on a model whose head size the + // paged attention kernel does not support. + int32_t preempt_alloc_granularity = 1; + + // cells a slot holding n_tokens actually occupies + int32_t preempt_n_cells(int32_t n_tokens) const { + return preempt_n_cells_g(n_tokens, preempt_alloc_granularity); + } + + // cells a slot holding n_tokens has to be given for a step of n_step more + int32_t preempt_n_cells_step(int32_t n_tokens, int32_t n_step) const { + return preempt_n_cells_step_g(n_tokens, n_step, preempt_alloc_granularity); + } + + // Cells kept spare on top of the reservation. A step that crosses a page boundary costs a + // whole page rather than a cell, so a margin of a few cells is no margin at all under a page + // allocator: round it up to one page. With a granularity of 1 this is PREEMPT_N_MARGIN. + int32_t preempt_n_margin() const { + return preempt_n_cells(PREEMPT_N_MARGIN); + } + int32_t preempt_n_spec_max() const { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } @@ -2906,7 +2991,10 @@ struct server_context_impl { res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); } - return res; + // [TAG_EXACT_CONCURRENCY] a restore takes fresh pages and its tail page is charged in + // full, so what the pool has to have free for this slot is the rounded figure. Under + // counting here is what admits a resume that find_slot then cannot satisfy. + return preempt_n_cells(res); } // Cells the pool is holding right now. A released slot keeps its prompt in the cache @@ -2921,7 +3009,9 @@ struct server_context_impl { continue; // parked: its cells are in host RAM, not in the pool } - res += slot.prompt.n_tokens(); + // [TAG_EXACT_CONCURRENCY] the slot's tail page is charged in full: it belongs to + // this sequence and cannot be given to anybody else, however little of it is used + res += preempt_n_cells(slot.prompt.n_tokens()); } return res; @@ -2935,27 +3025,36 @@ struct server_context_impl { int32_t res = 0; int32_t res_pmt = 0; + // [TAG_EXACT_CONCURRENCY] each slot reserves the cells its next step ADDS, not the + // tokens it adds. preempt_kv_used() already charges every slot's tail page in full, so + // with a granularity of 1 these are the same number and nothing changes; with a larger + // one the step is free until it crosses a page boundary and costs a whole page when it + // does. Reserving tokens on top of a rounded used figure would miss exactly that + // crossing, which is the only moment the pool can actually run out. for (const auto & slot : slots) { + const int32_t n_cur = slot.prompt.n_tokens(); + switch (slot.state) { case SLOT_STATE_GENERATING: case SLOT_STATE_DONE_PROMPT: { - res += 1 + n_spec; + res += preempt_n_cells_step(n_cur, 1 + n_spec); } break; case SLOT_STATE_STARTED: case SLOT_STATE_PROCESSING_PROMPT: { - const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + const int32_t n_left = slot.task ? slot.task->n_tokens() - n_cur : 0; - res_pmt += std::max(1, std::min(n_batch, n_left)); + res_pmt += preempt_n_cells_step(n_cur, std::max(1, std::min(n_batch, n_left))); } break; default: break; } } - // one batch is all the prompt slots get between them, however many are waiting - return res + std::min(res_pmt, n_batch); + // one batch is all the prompt slots get between them, however many are waiting; in + // cells that batch can straddle one boundary more than it has tokens for + return res + std::min(res_pmt, preempt_n_cells(n_batch)); } // Keep the slot that is furthest along -- it is the closest to finishing and to giving @@ -3068,7 +3167,7 @@ struct server_context_impl { // continue, so give those cells up first - same call the KV-full path makes. for (;;) { for (auto * slot : parked) { - if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + PREEMPT_N_MARGIN <= n_cells) { + if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + preempt_n_margin() <= n_cells) { best = slot; break; } @@ -3132,7 +3231,7 @@ struct server_context_impl { for (;;) { const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); - if (n_used + PREEMPT_N_MARGIN <= n_cells) { + if (n_used + preempt_n_margin() <= n_cells) { break; } diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 0da885bcafd..1e22d0a6ddb 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -1,4 +1,5 @@ import os +import re import time import tempfile import pytest @@ -38,6 +39,7 @@ def create_server(): os.close(fd) yield os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + os.environ.pop("LLAMA_SERVER_PREEMPT_GRANULARITY", None) os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) @@ -111,6 +113,48 @@ def test_two_slots_that_overflow_the_pool_together_both_finish(): assert len(res.body["tokens"]) == n_predict +def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages(): + # A pool that hands out cells in blocks gives a whole block to one sequence, so a sequence + # of n tokens occupies round_up(n, block) cells and holds the rest of its tail block against + # everybody else. The planner has to count those cells: counting tokens, it sees room the + # allocator cannot find, never parks anybody, and the retry ladder ends every request. + # + # llama_memory_alloc_granularity() reports the block size, and the only mode that returns + # more than 1 today is exact concurrency, whose paged attention kernel needs a head size this + # model does not have. LLAMA_SERVER_PREEMPT_GRANULARITY injects the figure instead: what is + # under test is the server's arithmetic, which is the same at 64 as at 256. + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_GRANULARITY"] = "64" + server.start() + log = LogReader(server.log_path) + assert "LLAMA_SERVER_PREEMPT_GRANULARITY = 64" in log.drain() + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + # every figure the planner logs is a whole number of blocks: "kv N/256" is what the pool is + # holding and "(wanted N)" is that plus what the next decode reserves. Counting tokens, both + # land wherever the sequences happen to be. + held = [int(n) for n in re.findall(r"kv (\d+)/256", text)] + wanted = [int(n) for n in re.findall(r"\(wanted (\d+)\)", text)] + assert held and wanted, f"the planner logged no figures:\n{text}" + assert all(n % 64 == 0 for n in held + wanted), f"not whole blocks: {held} {wanted}" + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_predict + _WORDS = ( "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor " From c6c3cb671ded0994d682003b534f8d6c6593ce3b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 10:50:22 +0000 Subject: [PATCH 37/81] exact: group speculative verify batches and slice the column split Under LLAMA_EXACT_CONCURRENCY a decode step with speculative drafts cost about twice what the same step cost with the mode off, and the mode itself was not where the time went. The batch splitter isolated any sequence set with more than one token left to place, on the reasoning that such a set is a prompt whose prefill would otherwise share a ubatch with other sequences. A speculative verify batch is such a set too: with two MTP drafts every slot brings three tokens, so each slot's verify step became a ubatch of its own and every decode step ran the whole graph once per slot. The splitter now isolates by width. llama_set_exact_decode_tokens() tells the library how many tokens one sequence contributes to a decode step, one plus the draft length, and only a set with more tokens than that left to place is a prompt. The server already derives the CUDA column bound from the same figure, so a grouped verify batch is at most that wide and the column policy keeps every column at its batch-of-one arithmetic. The default of 1 is the previous behaviour. On the CUDA side the column policy recomputed a batch above the bound one column at a time, reading the weights once per column. A column's result depends on the implementation and, for MMVQ, on the warp count of the launch, and neither depends on the other columns in the launch, so the split now runs in the widest slices whose configuration matches a batch of one: a twelve-column verify batch is three MMVQ launches on a table whose configuration holds up to four columns, rather than twelve. Mode 1 of GGML_CUDA_BATCH_INVARIANT still recomputes one column at a time. Qwen3.5-4B, four chats with two MTP drafts each, exact mode on: byte identical to the solo run in every round, with and without a forced park every 64 tokens, at about twice the previous aggregate rate. Qwen3.6-35B-A3B, the same cell: identical in every round at 92 to 96 tok/s against 55 before. MUL_MAT op tests 1217 of 1217 under the mode; server preemption tests 7 of 7. --- common/common.cpp | 3 ++ ggml/src/ggml-cuda/ggml-cuda.cu | 53 ++++++++++++++++++++++++++++----- include/llama.h | 8 +++++ src/llama-batch.cpp | 22 +++++++------- src/llama-batch.h | 20 +++++++------ src/llama-impl.cpp | 12 ++++++++ src/llama-kv-cache.cpp | 2 +- src/llama-memory-hybrid.cpp | 2 +- src/llama-memory-recurrent.cpp | 2 +- 9 files changed, 94 insertions(+), 30 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index d2beebc8e2f..8a1c76793cc 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1484,6 +1484,9 @@ bool common_exact_concurrency_init(const common_params & params) { const int n_cols = common_exact_decode_width(params); + // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step is + llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel))); + const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); if (bound) { const int max_cols = atoi(bound); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index c8f43bae382..2d3fc641ae6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1962,7 +1962,35 @@ static ggml_cuda_mm_path ggml_cuda_mul_mat_path( static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); -// Recompute dst one column at a time so that each column sees the batch-of-one configuration. +// [TAG_BATCH_INVARIANT] +// The widest slice of columns that can be recomputed in one launch while every column in it still +// sums the way a batch of one would. A column's result depends on the implementation and, for +// MMVQ, on the warp count of the launch, and neither depends on the values of the other columns +// in the launch, so a slice as wide as the batch-of-one configuration reaches gives each of its +// columns the batch-of-one value while reading the weights once for all of them instead of once +// per column. A twelve-column speculative decode over a table whose configuration holds up to four +// columns then costs three weight reads rather than twelve. Always below ncols_dst, so the +// recursive call cannot land back here with the same shape. +static int64_t ggml_cuda_mul_mat_invariant_width( + int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, + ggml_cuda_mm_path path_one, int64_t ncols_dst) { + if (path_one != GGML_CUDA_MM_MMVF && path_one != GGML_CUDA_MM_MMVQ) { + return 1; + } + const int64_t widest = path_one == GGML_CUDA_MM_MMVF ? MMVF_MAX_BATCH_SIZE : MMVQ_MAX_BATCH_SIZE; + for (int64_t w = std::min(ncols_dst - 1, widest); w > 1; --w) { + if (ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, w) != path_one) { + continue; + } + if (path_one == GGML_CUDA_MM_MMVQ && !ggml_cuda_mmvq_matches_single_column(src0->type, cc, w)) { + continue; + } + return w; + } + return 1; +} + +// Recompute dst in slices of columns so that each column sees the batch-of-one configuration. // Returns false when the batched launch already gives every column that same value. static bool ggml_cuda_mul_mat_split_columns( ggml_backend_cuda_context & ctx, int cc, int warp_size, @@ -2001,6 +2029,9 @@ static bool ggml_cuda_mul_mat_split_columns( return false; } + // Mode 1 recomputes one column at a time. Mode 2 recomputes in the widest slices that keep the + // batch-of-one arithmetic, which is what the exact concurrency mode runs under. + int64_t width = 1; if (ggml_cuda_batch_invariant() >= 2) { const ggml_cuda_mm_path path_one = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1); const ggml_cuda_mm_path path_batched = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ncols_dst); @@ -2014,20 +2045,26 @@ static bool ggml_cuda_mul_mat_split_columns( return false; } } + width = ggml_cuda_mul_mat_invariant_width(cc, warp_size, src0, src1, dst, path_one, ncols_dst); + if (width >= ncols_dst) { + width = 1; + } } - for (int64_t i = 0; i < ncols_dst; ++i) { + for (int64_t i = 0; i < ncols_dst; i += width) { + const int64_t n = std::min(width, ncols_dst - i); + ggml_tensor src1_col = *src1; ggml_tensor dst_col = *dst; - src1_col.ne[1] = 1; - src1_col.nb[2] = src1_col.nb[1]; - src1_col.nb[3] = src1_col.nb[1]; + src1_col.ne[1] = n; + src1_col.nb[2] = n*src1_col.nb[1]; + src1_col.nb[3] = n*src1_col.nb[1]; src1_col.data = (char *) src1->data + i*src1->nb[1]; - dst_col.ne[1] = 1; - dst_col.nb[2] = dst_col.nb[1]; - dst_col.nb[3] = dst_col.nb[1]; + dst_col.ne[1] = n; + dst_col.nb[2] = n*dst_col.nb[1]; + dst_col.nb[3] = n*dst_col.nb[1]; dst_col.data = (char *) dst->data + i*dst->nb[1]; ggml_cuda_mul_mat(ctx, src0, &src1_col, &dst_col); diff --git a/include/llama.h b/include/llama.h index 43d79b40a2a..0e44ff052c2 100644 --- a/include/llama.h +++ b/include/llama.h @@ -803,6 +803,14 @@ extern "C" { // will believe there is space that cannot be handed out. LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem); + // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or + // 1 plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence + // set with more tokens than this left to place is a prompt and is prefilled in a ubatch of its + // own; a set at or below it is a decode step and stays grouped with the other decodes, so a + // speculative verify batch is not run once per sequence. Process-wide, default 1. + LLAMA_API void llama_set_exact_decode_tokens(uint32_t n_tokens); + LLAMA_API uint32_t llama_exact_decode_tokens(void); + // // State / sessions // diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 4b73ab2478b..4080ecd11d0 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -507,7 +507,7 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) { return ubatch_add(idxs, idxs.size(), false); } -bool llama_batch_allocr::has_multi_token_seq() const { +bool llama_batch_allocr::has_seq_wider_than(uint32_t n_tokens) const { std::vector n_per_seq(n_seq_max, 0); for (int32_t i = 0; i < batch.n_tokens; ++i) { @@ -517,7 +517,7 @@ bool llama_batch_allocr::has_multi_token_seq() const { } for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { - if (++n_per_seq[batch.seq_id[i][s]] > 1) { + if (++n_per_seq[batch.seq_id[i][s]] > n_tokens) { return true; } } @@ -526,7 +526,7 @@ bool llama_batch_allocr::has_multi_token_seq() const { return false; } -llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, bool isolate_multi_token_seqs) { +llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above) { if (sequential && has_cpl) { LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__); @@ -559,12 +559,14 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } if (add) { - // [TAG_EXACT_CONCURRENCY] a sequence set that still has more than one token to place is - // a prompt, and a prompt shares its arithmetic with whatever else is in the ubatch, so - // give it a ubatch of its own. Sets with one token left are a plain decode step, which - // is already exact, so keep grouping those: isolating them too would make one prompt - // serialize every concurrent decode for the whole of the prefill. - if (isolate_multi_token_seqs) { + // [TAG_EXACT_CONCURRENCY] a sequence set that still has more tokens to place than a + // decode step carries is a prompt, and a prompt shares its arithmetic with whatever + // else is in the ubatch, so give it a ubatch of its own. Sets at or below that width + // are decode steps, plain or speculative, whose columns the backend's column policy + // already keeps exact, so keep grouping those: isolating them too would make one + // prompt serialize every concurrent decode for the whole of the prefill, and would run + // a speculative verify step once per sequence. + if (isolate_seqs_above > 0) { uint32_t n_left = 0; for (const auto idx : seq_set_map[seq_set[i]]) { @@ -573,7 +575,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } } - if (n_left > 1) { + if (n_left > isolate_seqs_above) { if (!cur_seq_set.empty()) { // let the sets already taken have this ubatch; the prompt gets the next one break; diff --git a/src/llama-batch.h b/src/llama-batch.h index 4bd2aa98f9f..7b638b20d30 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -105,15 +105,17 @@ class llama_batch_allocr { // make ubatches of equal-length sequences sets // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch - // isolate_multi_token_seqs = [TAG_EXACT_CONCURRENCY] a sequence set with more than one token - // left to place is given a ubatch of its own; sets with a single token left are - // still grouped together, so a prompt next to three decodes costs one extra - // ubatch and does not serialize the three decodes - llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, bool isolate_multi_token_seqs = false); - - // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than one token left to place, - // i.e. what remains of the batch is not a plain one-token-per-sequence decode step - bool has_multi_token_seq() const; + // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this many + // tokens left to place is a prompt and is given a ubatch of its own; sets at or + // below it are decode steps (one token, or one plus the speculative drafts) and + // stay grouped together, so a prompt next to three decodes costs one extra ubatch + // and does not serialize the three decodes, and a speculative verify batch is not + // run once per sequence + llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above = 0); + + // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than n_tokens left to place, + // i.e. what remains of the batch holds a prompt rather than decode steps only + bool has_seq_wider_than(uint32_t n_tokens) const; // sequence-set-wise split - each ubatch contains a single sequence-set llama_ubatch split_seq(uint32_t n_ubatch); diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index bad0e55237a..8c95e842f99 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -180,3 +181,14 @@ bool llama_exact_concurrency() { return enabled; } + +// [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h +static std::atomic g_exact_decode_tokens{1}; + +void llama_set_exact_decode_tokens(uint32_t n_tokens) { + g_exact_decode_tokens.store(n_tokens > 0 ? n_tokens : 1, std::memory_order_relaxed); +} + +uint32_t llama_exact_decode_tokens(void) { + return g_exact_decode_tokens.load(std::memory_order_relaxed); +} diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index dced34f6a14..bf37cdd291f 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -882,7 +882,7 @@ llama_memory_context_ptr llama_kv_cache::init_batch( // ubatch, so a sequence's prefill would run at a width its solo run never sees. Take // the sequence-set split instead, which can give each prompt a ubatch of its own; a // plain decode step has nothing to isolate and keeps taking split_simple. - const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq(); + const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0; auto ubatch = n_stream == 1 && !isolate ? balloc.split_simple(n_ubatch) diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 48fdea8b49c..7f502f3fa0c 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -91,7 +91,7 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // leaves a different gated delta net state than the same prompt processed alone. // Giving such a sequence a ubatch of its own removes that. A plain decode step, one // token per sequence, is already exact and stays batched. - const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq(); + const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0; ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index f639a25c5df..8943faf316c 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -433,7 +433,7 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & // so that the rollback snapshots remain valid // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: a recurrent state that a // prompt leaves behind depends on what shared its ubatch, so isolate the prompts - const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq(); + const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0; ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } From 379ca5d42a187c05dd4b663be749b088422fa033 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 11:27:32 +0000 Subject: [PATCH 38/81] cuda: run the single-token MUL_MAT_ID configuration over every token in one launch Under the batch-invariant policy a MUL_MAT_ID with several tokens was recomputed one token at a time, re-entering the op once per token, so a decode step of four slots with two MTP drafts each cost twelve serial expert launches per projection. The single-column MMVQ kernel already carries a sample axis. With ids, a sample is now a token: the wrapper launches the ncols_dst = 1 configuration once with the tokens on the z axis, y advancing by a token per sample and the expert index read per sample, so every (token, expert slot) block runs exactly the instructions the token alone would run, on the same data, with the same warp count, row split and K loop. Only the block indices differ. The stock single-token launch has one sample and is unchanged, as is every launch without the knob. A quantized expert matrix takes this path for any token count under the knob; other types still go one token at a time. Single-op probe, token 0's slice against its own single-token run, knob on, Q4_K, Q5_K, Q6_K and Q8_0 at K 2048 in both the gate-up and the down layout and Q6_K at K 512 (the shape whose multi-token kernel differed), 2 to 12 tokens: 0 differing elements in every one of 99 comparisons. --- ggml/src/ggml-cuda/ggml-cuda.cu | 7 +++++++ ggml/src/ggml-cuda/mmvq.cu | 28 +++++++++++++++++++++++++--- 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2d3fc641ae6..e5e66f95265 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2228,6 +2228,13 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * // [TAG_BATCH_INVARIANT] if (ggml_cuda_mul_mat_id_splits_tokens(dst)) { GGML_ASSERT(ne3 == 1 && src1->ne[3] == 1 && ids->ne[2] == 1 && ids->ne[3] == 1); + // A quantized expert matrix takes the single-token MMVQ path for every token count, and + // that path can put the tokens on its sample axis in one launch rather than being + // re-entered once per token. Anything else is still recomputed one token at a time. + if (ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); + return; + } ggml_cuda_mul_mat_id_split_tokens(ctx, dst); return; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index b14ef9681c5..b4e5196f172 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -592,9 +592,13 @@ static __global__ void mul_mat_vec_q( uint32_t sample_dst; ggml_cuda_pdl_sync(); - channel_x = ncols_dst == 1 && ids ? ids[channel_dst] : fastdiv(channel_dst, channel_ratio); - channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; sample_dst = blockIdx.z; + // [TAG_BATCH_INVARIANT] with ids, a sample is a token: the batch-invariant MUL_MAT_ID launch + // puts every token of the batch on the z axis of one single-column launch, so each (token, + // expert slot) block runs the exact single-token configuration. The stock single-token launch + // has one sample, where this indexing is ids[channel_dst] as before. + channel_x = ncols_dst == 1 && ids ? ids[sample_dst*ids_stride + channel_dst] : fastdiv(channel_dst, channel_ratio); + channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; const uint32_t sample_x = fastdiv(sample_dst, sample_ratio); const uint32_t sample_y = sample_dst; @@ -1281,7 +1285,14 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT( nb0 == ts_dst); GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); - GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE); + // [TAG_BATCH_INVARIANT] under the knob a MUL_MAT_ID with several tokens is computed as one + // launch of the single-token configuration with the tokens on the sample axis, so every + // (token, expert slot) block reduces exactly as the token alone would. The token count is + // then not bounded by the column templates. + const bool tokens_as_samples = ids && ne2 > 1 && ggml_cuda_batch_invariant(); + + GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE || tokens_as_samples); + GGML_ASSERT(!tokens_as_samples || !fusion); const float * src1_d = (const float *) src1->data; const int32_t * ids_d = ids ? (const int32_t *) ids->data : nullptr; @@ -1369,6 +1380,17 @@ void ggml_cuda_mul_mat_vec_q( const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; + if (tokens_as_samples) { + GGML_ASSERT(ne03 == 1 && ne13 == 1 && ne3 == 1); + // one column, one sample per token: y advances by s12 per token, dst by s2, x not at all + mul_mat_vec_q_switch_type( + src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, + ne01, 1, s01, stride_col_y, stride_col_dst, + ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, + 1, ne2, s03, s12, s2, ids_stride, stream); + return; + } + mul_mat_vec_q_switch_type( src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, ne01, ncols_dst, s01, stride_col_y, stride_col_dst, From f4e45646d6c76992d9aa3730eb660fc9e97caba0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 12:55:09 +0000 Subject: [PATCH 39/81] batch: group only sets with the same number of tokens left under exact mode With prompts isolated by width, sets at or below the decode width were grouped whatever their token counts, and the equal-length expansion then placed a three-token verify step beside a two-token one as two tokens now and one later. Attention and the gated delta net do not care, and the 4B reads identical either way, but a memory that reduces over a chunk of tokens, such as a chunked state space scan, would sum in a different order than the solo run's single three-token ubatch. A set now joins the ubatch only if it has as many tokens left as the first set taken, so every set in a ubatch finishes in that ubatch and each sequence's step has the shape it has alone. Sets with a different count wait for a later ubatch. Speculative verify steps of equal width, the common case, still share one ubatch. 4B, two MTP drafts, three rounds each: identical to the solo run with and without a forced park every 64 tokens; speculation off with forced parks identical; server preemption tests 7 of 7. --- src/llama-batch.cpp | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 4080ecd11d0..0c28cec32ba 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -537,6 +537,10 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, llama_seq_id last_seq_id = -1; + // [TAG_EXACT_CONCURRENCY] tokens left in the first set taken, when isolating: only sets with + // the same count join it, so that every set in the ubatch finishes in this ubatch + uint32_t n_left_first = 0; + // determine the non-overlapping sequence sets participating in this ubatch for (int32_t i = 0; i < batch.n_tokens; ++i) { if (used[i]) { @@ -565,7 +569,12 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, // are decode steps, plain or speculative, whose columns the backend's column policy // already keeps exact, so keep grouping those: isolating them too would make one // prompt serialize every concurrent decode for the whole of the prefill, and would run - // a speculative verify step once per sequence. + // a speculative verify step once per sequence. Grouped sets must have the same number + // of tokens left: the equal-length expansion below would otherwise place a three-token + // verify step beside a two-token one as two tokens now and one later, and a memory + // that reduces over a chunk of tokens (a chunked state space scan) would then sum in a + // different order than the solo run's single three-token ubatch. A set with a + // different count waits for a later ubatch. if (isolate_seqs_above > 0) { uint32_t n_left = 0; @@ -587,6 +596,12 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, break; } + + if (cur_seq_set.empty()) { + n_left_first = n_left; + } else if (n_left != n_left_first) { + continue; + } } cur_seq_set.push_back(seq_set[i]); From 1a6f7da42092a7ec7192831e7ea8deb00b8f7dce Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 15:36:01 +0000 Subject: [PATCH 40/81] exact: refuse by name what the mode cannot run, before it runs Three gaps found in review, all of the same kind: exact mode accepted an input it could not implement and either ran something else or asserted. A model whose K or V heads are not 256 wide passed every load-time check while the paged attention kernel refuses such heads, so attention ran unpaged on the CPU with the mode reporting itself as on. The load now fails with the head widths in the message, as the other preconditions do. MLA layouts are refused the same way. A token carrying several sequence ids reached the page placement and hit an assertion there. init_batch now refuses the batch with a logged error and the failed-prepare status, for the unified cache and the hybrid one. On a hybrid memory the attention half already refused a cross-sequence copy, a position shift and a position division under the mode, but the recurrent half went ahead, leaving the two halves describing different states. The hybrid memory now refuses all three before either half is touched. The harness model (8-wide heads) is now refused at load under the mode instead of loading unpaged; the 4B is unaffected (two MTP rounds identical); server preemption tests 7 of 7. --- src/llama-batch.cpp | 10 ++++++++++ src/llama-batch.h | 3 +++ src/llama-kv-cache.cpp | 17 +++++++++++++++++ src/llama-memory-hybrid.cpp | 27 +++++++++++++++++++++++++++ 4 files changed, 57 insertions(+) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 0c28cec32ba..5db683db281 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -507,6 +507,16 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) { return ubatch_add(idxs, idxs.size(), false); } +bool llama_batch_allocr::has_shared_tokens() const { + for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.n_seq_id[i] > 1) { + return true; + } + } + + return false; +} + bool llama_batch_allocr::has_seq_wider_than(uint32_t n_tokens) const { std::vector n_per_seq(n_seq_max, 0); diff --git a/src/llama-batch.h b/src/llama-batch.h index 7b638b20d30..52a375ad49e 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -117,6 +117,9 @@ class llama_batch_allocr { // i.e. what remains of the batch holds a prompt rather than decode steps only bool has_seq_wider_than(uint32_t n_tokens) const; + // [TAG_EXACT_CONCURRENCY] true if some token carries more than one sequence id + bool has_shared_tokens() const; + // sequence-set-wise split - each ubatch contains a single sequence-set llama_ubatch split_seq(uint32_t n_ubatch); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index bf37cdd291f..1c31e57a761 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -275,6 +275,15 @@ llama_kv_cache::llama_kv_cache( // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order while // the mode still reports itself as on, so refuse the load instead + // [TAG_EXACT_CONCURRENCY] the paged attention kernel handles 256-wide K and V heads only; + // any other width would run unpaged on the CPU while the mode reports itself as on + if (exact_pages && (hparams.n_embd_head_k(il) != 256 || (!is_mla && hparams.n_embd_head_v(il) != 256) || is_mla)) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d has %u-wide K heads and %u-wide V heads%s, " + "and the paged attention kernel supports 256-wide K and V heads only\n", + __func__, il, hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), is_mla ? " (MLA)" : ""); + throw std::runtime_error("exact concurrency: unsupported attention head size"); + } + if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, " "which has no paged attention: every layer must be offloaded to the CUDA backend " @@ -874,6 +883,14 @@ llama_memory_context_ptr llama_kv_cache::init_batch( GGML_UNUSED(embd_all); do { + // [TAG_EXACT_CONCURRENCY] a token shared by several sequences would be one cell in a page + // that belongs to one sequence; the placement asserts on it later, so refuse it here + if (exact_pages && balloc.has_shared_tokens()) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support tokens shared by several sequence ids; " + "give every token exactly one sequence id\n", __func__); + break; + } + balloc.split_reset(); std::vector ubatches; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 7f502f3fa0c..3fea4bb0ba3 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -66,6 +66,13 @@ llama_memory_hybrid::llama_memory_hybrid( llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { do { + // [TAG_EXACT_CONCURRENCY] refused before the attention half asserts on it, see llama_kv_cache::init_batch + if (llama_exact_concurrency() && balloc.has_shared_tokens()) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support tokens shared by several sequence ids; " + "give every token exactly one sequence id\n", __func__); + break; + } + balloc.split_reset(); // follow the recurrent pattern for creating the ubatch splits @@ -163,6 +170,14 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 } void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + // [TAG_EXACT_CONCURRENCY] the attention half refuses this under exact mode; refuse it here + // before either half is touched, so the two halves cannot end up describing different states + if (llama_exact_concurrency() && seq_id_src != seq_id_dst) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between sequences (%d -> %d); ignoring the copy\n", + __func__, seq_id_src, seq_id_dst); + return; + } + mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1); mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1); } @@ -173,11 +188,23 @@ void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + if (llama_exact_concurrency() && shift != 0) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions (seq %d, shift %d); ignoring the shift\n", + __func__, seq_id, shift); + return; + } + mem_attn->seq_add(seq_id, p0, p1, shift); mem_recr->seq_add(seq_id, p0, p1, shift); } void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + if (llama_exact_concurrency() && d != 1) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions (seq %d, d %d); ignoring the division\n", + __func__, seq_id, d); + return; + } + mem_attn->seq_div(seq_id, p0, p1, d); mem_recr->seq_div(seq_id, p0, p1, d); } From 5864dae95a3ef4173b68ffa29842db33275272c1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 16:25:21 +0000 Subject: [PATCH 41/81] server: tell the stream about a park made as a last resort The last resort parks through the same call as the planner but did not send the stream comment the planner sends, so a client that shows the pause from that comment showed nothing for a slot parked this way. --- tools/server/server-context.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 2c89d7368dc..092ffd301b4 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4219,6 +4219,8 @@ struct server_context_impl { metrics.n_preempt++; n_parked++; + send_preempt_notice(*victim, true); + SLT_WRN(*victim, "preempted as a last resort: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, (ggml_time_us() - t_start) / 1e3, From cd1cd4e6d37474b9ee6157179e44e351f2aeed57 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 16:36:25 +0000 Subject: [PATCH 42/81] server: count a slot being restored in the asynchronous lookahead margin A slot on its way back in already holds its cells and starts decoding the moment its copy lands, so it needs the same runway as the slots already running. Leaving it out let two back-to-back restores land into a pool that then had to park someone again at once. --- tools/server/server-context.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index cea3665774d..0395a3e22f6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3337,7 +3337,9 @@ struct server_context_impl { int32_t n_running = n_additional_running; for (const auto & slot : slots) { - if (slot.is_processing() && !slot.preempt_is_out()) { + // A slot on its way back in already holds its cells and starts decoding the + // moment its copy lands, so it needs the runway now; one on its way out does not. + if (slot.is_processing() && (!slot.preempt_is_out() || slot.state == SLOT_STATE_RESTORING)) { n_running++; } } From 0c9fc0ed8b03c7845c16feee12df7ed3c0db765e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 16:48:25 +0000 Subject: [PATCH 43/81] exact: every context reports the widest decode step it can build The column bound the CUDA backend splits at was reported by the server only. A program using the library directly got a fallback of 16 columns and, above it, a once-only warning that the mode did not hold for that op; with 32 sequences in a step that is silent inexactness after one line of log. llama_set_exact_decode_width() is the one place the width is reported, and it never lowers what was reported. Every llama_context reports its own at creation, n_seq_max times the per-sequence width, so a decode of any context stays within the bound its kernels split at whether or not the caller knew there was one. The server's common path goes through the same call, so the figure it reports is the one its contexts would. Probe: 32 sequences on the 4B, greedy, each with its own copy of the prompt. Before, the warning fires; after, it does not, and all 32 match the solo run either way at this length. --- common/common.cpp | 12 +++--------- include/llama.h | 9 +++++++++ src/llama-context.cpp | 8 ++++++++ src/llama-impl.cpp | 29 +++++++++++++++++++++++++++++ 4 files changed, 49 insertions(+), 9 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 8a1c76793cc..abca3630e13 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1500,15 +1500,9 @@ bool common_exact_concurrency_init(const common_params & params) { } } - // the CUDA backend may not be present or may be loaded dynamically, so go through the registry - for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { - ggml_backend_reg_t reg = ggml_backend_reg_get(i); - - auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width"); - if (fn) { - fn(n_cols); - } - } + // a context created later reports n_seq_max times the per-sequence width, which is this + // figure again; reporting it here as well covers a caller that decodes before that + llama_set_exact_decode_width((uint32_t) n_cols); return true; } diff --git a/include/llama.h b/include/llama.h index 0e44ff052c2..299249c5daf 100644 --- a/include/llama.h +++ b/include/llama.h @@ -811,6 +811,15 @@ extern "C" { LLAMA_API void llama_set_exact_decode_tokens(uint32_t n_tokens); LLAMA_API uint32_t llama_exact_decode_tokens(void); + // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns: the + // sequences a context can hold times the tokens each contributes to a step. Every context + // reports its own at creation and a backend keeps the widest it has heard, so a decode of any + // context stays within the bound its kernels split at. A caller that builds wider steps than + // the contexts imply (a draft of its own, say) reports the width itself, before creating the + // context or before the first decode. Never lowers what was reported. + LLAMA_API void llama_set_exact_decode_width(uint32_t n_cols); + LLAMA_API uint32_t llama_exact_decode_width(void); + // // State / sessions // diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 3ab84de780c..e1840ac57d3 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -101,6 +101,14 @@ llama_context::llama_context( throw std::runtime_error("n_seq_max must be <= " + std::to_string(LLAMA_MAX_SEQ)); } + // [TAG_EXACT_CONCURRENCY] the widest decode step this context can build: one column per + // sequence, times the tokens a sequence contributes to a step. Reported so that a backend + // splitting columns for exactness covers it without the caller having to know the bound; a + // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width. + if (llama_exact_concurrency()) { + llama_set_exact_decode_width(cparams.n_seq_max * llama_exact_decode_tokens()); + } + cparams.n_rs_seq = params.n_rs_seq; if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) { LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n", diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index 8c95e842f99..ad63fc5ea3b 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -1,5 +1,6 @@ #include "llama-impl.h" +#include "ggml-backend.h" #include "gguf.h" #include "llama.h" @@ -192,3 +193,31 @@ void llama_set_exact_decode_tokens(uint32_t n_tokens) { uint32_t llama_exact_decode_tokens(void) { return g_exact_decode_tokens.load(std::memory_order_relaxed); } + +// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h. A backend that +// splits columns to make a decode exact reads it through ggml_backend_cuda_set_exact_decode_width, +// reached through the registry so that a backend that is absent or loaded late costs nothing. +static std::atomic g_exact_decode_width{0}; + +void llama_set_exact_decode_width(uint32_t n_cols) { + uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed); + + while (n_cols > cur) { + if (g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { + for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { + ggml_backend_reg_t reg = ggml_backend_reg_get(i); + + auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width"); + if (fn) { + fn((int) n_cols); + } + } + + return; + } + } +} + +uint32_t llama_exact_decode_width(void) { + return g_exact_decode_width.load(std::memory_order_relaxed); +} From 46e7fa742fa7c78ad0bfb8c7247f5146b2a0c2e2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 17:43:09 +0000 Subject: [PATCH 44/81] exact: four refusals and one width from review The decode width per slot came from a switch over the speculation types that did not know the ngram cache verifies eight tokens, so under that type the per-sequence width was 4 rather than 9 and every verify group was taken for a prompt and isolated. The width now comes from common_speculative_n_max(), the same place the speculation code takes it. DFlash drafting turns causal attention off on its draft context, and the paged attention the mode runs on needs it; the next graph asserted. The server refuses the combination at setup by name, and a context with a cache under the mode refuses to turn causal attention off. A library user setting GGML_CUDA_BATCH_INVARIANT_MAX_COLS below a context's decode width was not caught: the explicit bound wins in the backend and its warning is silent once a width was reported. The context refuses to be created, the way the server's setup refuses. A width reported before a backend was loaded never reached it, since the setter only forwarded on a raise. The widest figure now goes to every backend on every call, and every context reports at creation. --- common/common.cpp | 38 +++++++++++++------------------------- src/llama-context.cpp | 24 +++++++++++++++++++++++- src/llama-impl.cpp | 24 +++++++++++++----------- 3 files changed, 49 insertions(+), 37 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index abca3630e13..7b6fc8ad2e2 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1448,32 +1448,11 @@ bool common_exact_concurrency() { int common_exact_decode_width(const common_params & params) { const int n_slots = std::max(1, params.n_parallel); - // the draft tokens a slot carries into the verify ubatch alongside its accepted token - int n_draft = 0; + // the draft tokens a slot carries into the verify ubatch alongside its accepted token, per + // speculation type, from the same place the speculation code takes its own width + const int n_draft = std::max(0, (int) common_speculative_n_max(¶ms.speculative)); - for (const auto type : params.speculative.types) { - switch (type) { - case COMMON_SPECULATIVE_TYPE_NONE: - break; - case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: - n_draft = std::max(n_draft, params.speculative.ngram_mod.n_max); - break; - case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: - n_draft = std::max(n_draft, (int) params.speculative.ngram_simple.size_m); - break; - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: - n_draft = std::max(n_draft, (int) params.speculative.ngram_map_k.size_m); - break; - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: - n_draft = std::max(n_draft, (int) params.speculative.ngram_map_k4v.size_m); - break; - default: - n_draft = std::max(n_draft, params.speculative.draft.n_max); - break; - } - } - - return n_slots*(1 + std::max(0, n_draft)); + return n_slots*(1 + n_draft); } // [TAG_EXACT_CONCURRENCY] @@ -1482,6 +1461,15 @@ bool common_exact_concurrency_init(const common_params & params) { return true; } + // DFlash drafting turns causal attention off on its draft context, and the paged + // attention the mode runs on needs it; say so instead of asserting in the graph + for (const auto type : params.speculative.types) { + if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash: it disables causal attention, which the paged attention needs\n"); + return false; + } + } + const int n_cols = common_exact_decode_width(params); // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step is diff --git a/src/llama-context.cpp b/src/llama-context.cpp index e1840ac57d3..c95493e06e6 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -106,7 +106,22 @@ llama_context::llama_context( // splitting columns for exactness covers it without the caller having to know the bound; a // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width. if (llama_exact_concurrency()) { - llama_set_exact_decode_width(cparams.n_seq_max * llama_exact_decode_tokens()); + const uint32_t n_cols = cparams.n_seq_max * llama_exact_decode_tokens(); + + // an explicit column bound wins over the reported width in the backend, so one below + // this context's width would leave its decodes batched above the bound with the mode + // still reporting itself on; refuse it here, the way the server's setup does + if (const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS")) { + const int max_cols = atoi(bound); + + if (max_cols > 0 && (uint32_t) max_cols < n_cols) { + LLAMA_LOG_ERROR("%s: GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at least %u to cover a decode step of %u sequences; raise it, set it to 0 for no bound, or unset it\n", + __func__, max_cols, n_cols, cparams.n_seq_max); + throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); + } + } + + llama_set_exact_decode_width(n_cols); } cparams.n_rs_seq = params.n_rs_seq; @@ -1196,6 +1211,13 @@ void llama_context::set_causal_attn(bool value) { return; } + // [TAG_EXACT_CONCURRENCY] the paged attention the mode runs on is causal; a context with a + // cache under the mode keeps causal attention rather than asserting in the next graph + if (!value && memory && llama_exact_concurrency()) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so causal attention cannot be turned off; the change is refused\n", __func__); + return; + } + cparams.causal_attn = value; sched_need_reserve = true; diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index ad63fc5ea3b..f8ce14014f1 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -202,18 +202,20 @@ static std::atomic g_exact_decode_width{0}; void llama_set_exact_decode_width(uint32_t n_cols) { uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed); - while (n_cols > cur) { - if (g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { - for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { - ggml_backend_reg_t reg = ggml_backend_reg_get(i); - - auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width"); - if (fn) { - fn((int) n_cols); - } - } + while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { + } + + // The widest figure so far goes to every backend on every call, not only when it grew: a + // width reported before a backend was loaded would otherwise never reach it, and every + // context reports at creation, by which time the backends are there. + const uint32_t widest = g_exact_decode_width.load(std::memory_order_relaxed); + + for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { + ggml_backend_reg_t reg = ggml_backend_reg_get(i); - return; + auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width"); + if (fn) { + fn((int) widest); } } } From 7fb42b582d1b29f997fb0d8b4edc97ca142f862d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 17:44:10 +0000 Subject: [PATCH 45/81] server: keep the park notices of a multi-prompt stream apart A request with several prompts streams them through one reader, and one flag stood for all of them: a prompt resuming cleared it while another was still parked, which switched the parked keepalive off for the one still waiting. The reader now keeps the set of parked prompts and runs the keepalive while any of them is parked. The comment names the prompt it is about for prompts after the first; prompt 0 keeps the bare form a single-prompt client matches on. --- tools/server/server-context.cpp | 36 +++++++++++++++++++++++++++++---- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1a8aa8ad62d..94e46a9289e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -18,6 +18,7 @@ #include "mtmd-helper.h" #include +#include #include #include #include @@ -90,6 +91,19 @@ constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is // LLAMA_SERVER_PREEMPT_RESUME=head (the default) or pass; read once in load_model() and logged. static bool g_preempt_resume_head_of_line = true; +// [TAG_PREEMPT] the SSE comment for a park or a resume. A request with several prompts +// streams them through one reader, so the comment names the prompt it is about, except for +// prompt 0, whose comment stays the bare form a single-prompt client matches on. +static std::string preempt_notice_comment(const server_task_result_preempt_notice & notice) { + std::string res = notice.parked ? ": preempted" : ": resumed"; + + if (notice.index > 0) { + res += " " + std::to_string(notice.index); + } + + return res + "\n\n"; +} + static bool preempt_resume_head_of_line() { return g_preempt_resume_head_of_line; } @@ -5055,10 +5069,16 @@ std::unique_ptr server_routes::handle_completions_impl( // token exists. Those notices arrive ahead of the first real result; keep them and // send them in front of it, so the client learns about the wait it just had. std::string preempt_prefix; + std::set parked_idx; // prompts of this request that are parked right now auto first_result = rd.next(req.should_stop); while (first_result != nullptr && dynamic_cast(first_result.get()) != nullptr) { const auto * notice = static_cast(first_result.get()); - preempt_prefix += notice->parked ? ": preempted\n\n" : ": resumed\n\n"; + preempt_prefix += preempt_notice_comment(*notice); + if (notice->parked) { + parked_idx.insert(notice->index); + } else { + parked_idx.erase(notice->index); + } first_result = rd.next(req.should_stop); } if (first_result == nullptr) { @@ -5090,7 +5110,11 @@ std::unique_ptr server_routes::handle_completions_impl( } res->status = 200; res->content_type = "text/event-stream"; - res->set_next([res_this = res.get(), res_type, sse_ping_interval, parked = false](std::string & output) mutable -> bool { + res->set_next([res_this = res.get(), res_type, sse_ping_interval, parked_idx](std::string & output) mutable -> bool { + // [TAG_PREEMPT] the keepalive runs while ANY prompt of the request is parked: with + // several prompts in one stream, one resuming does not mean the others did + const bool parked = !parked_idx.empty(); + static auto format_error = [](task_response_type res_type, const json & res_json) { if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { return format_anthropic_sse({ @@ -5177,8 +5201,12 @@ std::unique_ptr server_routes::handle_completions_impl( } else if (const auto * notice = dynamic_cast(result.get())) { // [TAG_PREEMPT] an SSE comment: invisible to clients that do not know // about preemption, a pause indicator for the ones that do - parked = notice->parked; - output = parked ? ": preempted\n\n" : ": resumed\n\n"; + if (notice->parked) { + parked_idx.insert(notice->index); + } else { + parked_idx.erase(notice->index); + } + output = preempt_notice_comment(*notice); } else { GGML_ASSERT( dynamic_cast(result.get()) != nullptr From eae462431761ee3ac5774873eae2c99e43bda6df Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 18:13:27 +0000 Subject: [PATCH 46/81] exact concurrency: refuse a non-causal context with a cache at creation, keep a page boundary per prompt slot in the reserve cap A context created with LLAMA_ATTENTION_TYPE_NON_CAUSAL under the mode used to pass creation and assert on its first graph; llama_context now refuses it right after the memory is created, so callers get an error from llama_init_from_model instead of an abort. Contexts without a cache are unaffected. The reserve cap gave every waiting prompt slot one batch between them, rounded to a page once; under page allocation each prompt slot can cross a boundary of its own within that batch, so the cap now keeps one boundary per prompt slot. --- src/llama-context.cpp | 7 +++++++ tools/server/server-context.cpp | 13 ++++++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c95493e06e6..f3b1df04c48 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -416,6 +416,13 @@ llama_context::llama_context( }; memory.reset(model.create_memory(params_mem, cparams)); + + // [TAG_EXACT_CONCURRENCY] the paged attention the mode runs on is causal; a context + // created non-causal with a cache would assert on its first graph, so it is refused here + if (llama_exact_concurrency() && memory && !cparams.causal_attn) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so it cannot be created with non-causal attention\n", __func__); + throw std::runtime_error("exact concurrency: non-causal attention is not supported with a KV cache"); + } } // init backends diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 56b48cca60c..f61cb2406ad 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3166,7 +3166,18 @@ struct server_context_impl { // one batch is all the prompt slots get between them, however many are waiting; in // cells that batch can straddle one boundary more than it has tokens for - return res + std::min(res_pmt, preempt_n_cells(n_batch)); + // one batch is all the prompt slots get between them, however many are waiting; under + // page allocation each of them can still cross a page boundary of its own within that + // batch, so the cap keeps one boundary per prompt slot on top of the batch + int32_t n_pmt = 0; + + for (const auto & slot : slots) { + if (slot.state == SLOT_STATE_STARTED || slot.state == SLOT_STATE_PROCESSING_PROMPT) { + n_pmt++; + } + } + + return res + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1)); } // Keep the slot that is furthest along -- it is the closest to finishing and to giving From dbd82ca5ee65c21a2006b10a847875056957a787 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 18:24:14 +0000 Subject: [PATCH 47/81] exact concurrency: the decode width of every context follows the token figure A context reported n_seq_max times the tokens per sequence once, at creation. Raising the process-wide token figure afterwards widened the decode step of every earlier context while the width they reported stayed put, so a context created under a narrower figure could batch above the bound it reported. Each context now reports its sequence count; the widest count seen times the current token figure is re-reported whenever the figure changes. --- src/llama-context.cpp | 4 +++- src/llama-impl.cpp | 21 +++++++++++++++++++++ src/llama-impl.h | 4 ++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index f3b1df04c48..d14123779a5 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -105,6 +105,8 @@ llama_context::llama_context( // sequence, times the tokens a sequence contributes to a step. Reported so that a backend // splitting columns for exactness covers it without the caller having to know the bound; a // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width. + // The sequence count is what is reported: the tokens figure can be raised later for the + // whole process, and the width then follows it for this context too. if (llama_exact_concurrency()) { const uint32_t n_cols = cparams.n_seq_max * llama_exact_decode_tokens(); @@ -121,7 +123,7 @@ llama_context::llama_context( } } - llama_set_exact_decode_width(n_cols); + llama_exact_report_n_seq(cparams.n_seq_max); } cparams.n_rs_seq = params.n_rs_seq; diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index f8ce14014f1..51c50898bdb 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -186,8 +186,29 @@ bool llama_exact_concurrency() { // [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h static std::atomic g_exact_decode_tokens{1}; +// the most sequences any context so far was created with. The tokens figure is process +// wide, so raising it widens the decode step of every context that already exists; the +// width those contexts reported at creation is re-reported here with the new figure, or a +// context created under a narrower figure would batch above the bound it reported. +static std::atomic g_exact_max_n_seq{0}; + +void llama_exact_report_n_seq(uint32_t n_seq) { + uint32_t cur = g_exact_max_n_seq.load(std::memory_order_relaxed); + + while (n_seq > cur && !g_exact_max_n_seq.compare_exchange_weak(cur, n_seq, std::memory_order_relaxed)) { + } + + llama_set_exact_decode_width(g_exact_max_n_seq.load(std::memory_order_relaxed) * llama_exact_decode_tokens()); +} + void llama_set_exact_decode_tokens(uint32_t n_tokens) { g_exact_decode_tokens.store(n_tokens > 0 ? n_tokens : 1, std::memory_order_relaxed); + + const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed); + + if (n_seq > 0) { + llama_set_exact_decode_width(n_seq * llama_exact_decode_tokens()); + } } uint32_t llama_exact_decode_tokens(void) { diff --git a/src/llama-impl.h b/src/llama-impl.h index 9b64431fedd..3c720cec150 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -109,3 +109,7 @@ std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i); // so that its output does not change when other sequences share the KV cache. Off by default. // Reads the same LLAMA_EXACT_CONCURRENCY variable as the paged KV cache and the CUDA backend. bool llama_exact_concurrency(); + +// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so that the +// decode width every context needs is known to the backend and follows llama_set_exact_decode_tokens +void llama_exact_report_n_seq(uint32_t n_seq); From 42e536970c7d2348f0e7cb8a7fe67bbefe362e11 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 18:33:13 +0000 Subject: [PATCH 48/81] server: a round with a context shift waits for every park and restore copy first A context shift is recorded by pre_decode_shift() and applied inside the next llama_decode as one graph over the whole K cache, in place. A restore still writing its cells on its own stream could be read half done and written back stale, and a park still reading its cells would read through the rewrite. The round that recorded a shift now waits for every copy in flight before it decodes; shifts are rare, so the wait is too. The rotation park also stamps its issue time, so its completion line measures the copy rather than the previous park. --- tools/server/server-context.cpp | 48 +++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 8c77bac7adc..13ca47798e9 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3227,6 +3227,10 @@ struct server_context_impl { // set by preempt_last_resort(): the batch being decoded was given up, stop the chunk loop bool preempt_batch_abandoned = false; + // [TAG_PREEMPT_ASYNC] a context shift was recorded this round; it is applied inside the + // next llama_decode as one graph over the whole K cache, in place + bool preempt_shift_pending = false; + int32_t preempt_n_spec_max() const { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } @@ -3537,6 +3541,42 @@ struct server_context_impl { } } + // [TAG_PREEMPT_ASYNC] wait for every copy in flight, parks and restores alike. The + // context shift a slot recorded this round is applied inside the next llama_decode as one + // graph over the whole K cache, in place: a restore still writing its cells on its own + // stream could be read half done and written back stale, and a park still reading its + // cells would read through the rewrite. Shifts are rare, so this round waits. + void preempt_wait_for_shift() { + if (!preempt_shift_pending) { + return; + } + + preempt_shift_pending = false; + + while (preempt_wait_in_flight()) { + } + + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_RESTORING) { + continue; + } + + slot.preempt_copy_wait(); + + if (!slot.preempt_restore_poll()) { + continue; + } + + metrics.n_resume++; + + SLT_WRN(slot, "restore completed after %.2f ms (waited for, a context shift is due): %d tokens back in the cache, kv %d/%d, preemptions %d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, + slot.prompt.n_tokens(), + preempt_kv_used(), n_ctx, + slot.n_preempt); + } + } + // Wait for one outstanding park, the last thing tried before giving up on finding room. // It is what keeps a pool that fills faster than the copies drain no worse than the // synchronous path: the decode waits for the copy exactly as it used to. @@ -3682,10 +3722,14 @@ struct server_context_impl { continue; } + const int64_t t_start = ggml_time_us(); + if (!preempt_fits_budget(slot) || !slot.preempt_save()) { continue; } + slot.t_preempt_copy_us = t_start; + // [TAG_PREEMPT_ASYNC] an asynchronous park is counted when its copy // lands, and the head is re-examined on the pass that sees the room if (slot.state != SLOT_STATE_PREEMPTING) { @@ -3956,6 +4000,9 @@ struct server_context_impl { llama_batch batch_view; int32_t off_next = 0; int32_t n_batch = llama_n_batch(ctx_tgt); + + preempt_wait_for_shift(); + for (int32_t off = 0; off < batch.size(); off = off_next) { const int32_t n_tokens = std::min(n_batch, batch.size() - off); try { @@ -4047,6 +4094,7 @@ struct server_context_impl { SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); slot.n_ctx_shift++; + preempt_shift_pending = true; slot.mem.seq_rm (slot.id, n_keep , n_keep + n_discard); slot.mem.seq_add(slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); From ab40f16149fac3bc59d69d8a7264e8436ad1eab3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 18:35:59 +0000 Subject: [PATCH 49/81] server: a stream parked before its first token starts with the notice; the rotation park is announced A request parked while still processing its prompt has no token to send. The route used to hold the response until the first ordinary result, so the client saw nothing, not even the headers, until the slot resumed, and the parked keepalive never ran. The stream now starts on the first notice; the data-less signal a prompt sends before its first token is skipped once the stream is open, since it has nothing to add. The head-of-line rotation parks a resident through the same preempt_save() as any other park but never announced it, so that stream stayed silent while parked and later carried an unmatched resume. It is announced now. Two tests: the notice is the first thing on the wire and arrives while the other stream still runs; both streams of a rotation carry paired notices. --- tools/server/server-context.cpp | 50 +++++++----- .../server/tests/unit/test_preempt_notify.py | 81 +++++++++++++++++++ 2 files changed, 112 insertions(+), 19 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index de920419f6c..9bf5f25c3f6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3277,6 +3277,8 @@ struct server_context_impl { metrics.n_preempt++; + send_preempt_notice(slot, true); + SLT_WRN(slot, "rotated out after %d context shifts: %d cells released, %.1f MiB parked, a head parked %.1f s takes its turn, preemptions %d\n", slot.n_ctx_shift, slot.prompt.n_tokens(), slot.preempt_state_size() / (1024.0 * 1024.0), @@ -5115,34 +5117,37 @@ std::unique_ptr server_routes::handle_completions_impl( std::string preempt_prefix; std::set parked_idx; // prompts of this request that are parked right now auto first_result = rd.next(req.should_stop); - while (first_result != nullptr && dynamic_cast(first_result.get()) != nullptr) { + if (first_result != nullptr && dynamic_cast(first_result.get()) != nullptr) { + // [TAG_PREEMPT] parked before any token exists. The stream starts now, with the + // notice, so the parked keepalive runs through the wait instead of the client + // seeing nothing until the slot resumes; the first ordinary result follows in + // the stream, an error included, since the response has already begun. const auto * notice = static_cast(first_result.get()); - preempt_prefix += preempt_notice_comment(*notice); + preempt_prefix = preempt_notice_comment(*notice); if (notice->parked) { parked_idx.insert(notice->index); } else { parked_idx.erase(notice->index); } - first_result = rd.next(req.should_stop); - } - if (first_result == nullptr) { - GGML_ASSERT(req.should_stop()); - return res; // connection is closed - } + first_result.reset(); + } else { + if (first_result == nullptr) { + GGML_ASSERT(req.should_stop()); + return res; // connection is closed + } - if (first_result->is_error()) { - res->error(first_result->to_json()); - return res; - } + if (first_result->is_error()) { + res->error(first_result->to_json()); + return res; + } - GGML_ASSERT( - dynamic_cast(first_result.get()) != nullptr || - dynamic_cast (first_result.get()) != nullptr - ); + GGML_ASSERT( + dynamic_cast(first_result.get()) != nullptr || + dynamic_cast (first_result.get()) != nullptr + ); + } - // next responses are streamed - // to be sent immediately - json first_result_json = first_result->to_json(); + json first_result_json = first_result ? first_result->to_json() : json(nullptr); if (first_result_json == nullptr) { res->data = preempt_prefix; // simply send HTTP headers and status code } else if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { @@ -5257,6 +5262,13 @@ std::unique_ptr server_routes::handle_completions_impl( || dynamic_cast(result.get()) != nullptr ); json res_json = result->to_json(); + if (res_json.is_null()) { + // [TAG_PREEMPT] the signal a prompt sends before its first token, so + // that the headers go out, carries no data. Normally it is the first + // result and only opens the stream; after a notice opened the stream + // it has nothing to add, and the sender skips an empty chunk. + return true; + } if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { output = format_anthropic_sse(res_json); } else if (res_type == TASK_RESPONSE_TYPE_OAI_RESP) { diff --git a/tools/server/tests/unit/test_preempt_notify.py b/tools/server/tests/unit/test_preempt_notify.py index 86a4e66a674..76f722f6865 100644 --- a/tools/server/tests/unit/test_preempt_notify.py +++ b/tools/server/tests/unit/test_preempt_notify.py @@ -1,5 +1,7 @@ import os import tempfile +import threading +import time import pytest import requests from utils import * @@ -166,3 +168,82 @@ def test_two_overflowing_streams_both_finish_and_the_parked_one_says_so(): announced += 1 assert ": resumed" in comments assert announced >= 1, [r[0] for r in results] + + +def test_a_stream_parked_before_its_first_token_starts_with_the_notice(): + # A request parked while still processing its prompt has no token to send yet. The + # response must not wait for one: it starts with the notice, so the client sees + # "paused" and gets the keepalive at once, instead of a silent connection that only + # opens when the slot resumes. + global server + # The resident keeps growing towards the whole pool; the newcomer's prompt is larger + # than what is free beside it, so the planner parks the newcomer before it has a token. + global server + server.n_ctx = 512 + server.n_batch = 512 # the whole prompt in one batch, so the planner sees its size at once + server.start() + url = f"http://{server.server_host}:{server.server_port}/completion" + first = _completion_payload(390) | {"prompt": " ".join(["Once upon a time there was a brave knight who"] * 6)} + second = _completion_payload(32) | {"prompt": " ".join(["The quick brown fox jumps over the lazy dog and"] * 14)} + + timeline = [] + lock = threading.Lock() + + def _run(name, payload, started=None): + res = requests.post(url, json=payload, stream=True) + assert res.status_code == 200 + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if not line: + continue + with lock: + timeline.append((time.monotonic(), name, line)) + if started is not None and line.startswith("data: "): + started.set() + + started = threading.Event() + t = threading.Thread(target=_run, args=("first", first, started)) + t.start() + assert started.wait(30) + _run("second", second) + t.join(60) + + second_lines = [(ts, line) for ts, name, line in timeline if name == "second"] + first_end = max(ts for ts, name, _ in timeline if name == "first") + # The notice is the very first thing on the wire, and it arrives while the other + # stream is still running, not when it has finished and the parked slot resumes. + assert second_lines[0][1] == ": preempted", second_lines[:3] + assert second_lines[0][0] < first_end + events = [line for _, line in second_lines if line in (": preempted", ": resumed") or line.startswith("data: ")] + assert events[0] == ": preempted" and events[1] == ": resumed" and events[2].startswith("data: "), events[:3] + datas = [line[6:] for _, line in second_lines if line.startswith("data: ")] + assert _content(datas) + final = json.loads([d for d in datas if d != "[DONE]"][-1]) + assert final["tokens_predicted"] == 32 + + +def test_a_resident_rotated_out_for_a_parked_head_is_told_so(): + # The rotation from test_preempt: a resident cycling through context shifts holds the + # pool, and after the head has waited its turn the resident is parked in its place. + # That park is a park like any other, so its stream must say so, and every notice + # must be paired: no stream ends with a park it was never told about. + global server + server.n_ctx = 256 + server.enable_ctx_shift = True + server.start() + n_predict = 12000 + p1 = _completion_payload(n_predict) | {"prompt": "Once upon a time there was a brave knight who"} + p2 = _completion_payload(n_predict) | {"prompt": "The quick brown fox jumps over the lazy dog and"} + results = parallel_function_calls([ + (_stream_raw, ("/completion", p1)), + (_stream_raw, ("/completion", p2)), + ]) + n_parked = 0 + for comments, datas in results: + final = json.loads([d for d in datas if d != "[DONE]"][-1]) + assert final["tokens_predicted"] == n_predict + seq = [c for c in comments if c in (": preempted", ": resumed")] + assert seq == [": preempted", ": resumed"] * (len(seq) // 2), seq + n_parked += len(seq) // 2 + # Both streams took turns: at least one park each, so at least two in all. + assert n_parked >= 2, [r[0] for r in results] From d4e3fc8a6995fee7fa81b414ad4beb4b20e6aad8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 18:58:00 +0000 Subject: [PATCH 50/81] exact concurrency: equal-count grouping stays on for recurrent and hybrid memories; a width the explicit bound cannot cover is refused The recurrent and hybrid memories passed the isolation figure to split_equal only when a prompt was in the batch, so a three-token verify step could be placed beside a two-token one as two tokens now and one later, and a memory that reduces over a chunk of tokens would sum in an order the solo run never had. The figure is passed whenever the mode is on. An explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS wins in the backend, so a width raised past it after a context exists would leave decodes batched above the bound. llama_set_exact_decode_width and llama_set_exact_decode_tokens now return false and change nothing when the bound cannot cover the width; the context constructor and common's setup treat that as the error it is. --- common/common.cpp | 16 +++++++---- include/llama.h | 13 ++++++--- src/llama-context.cpp | 4 ++- src/llama-impl.cpp | 50 +++++++++++++++++++++++++++++----- src/llama-impl.h | 2 +- src/llama-memory-hybrid.cpp | 4 ++- src/llama-memory-recurrent.cpp | 8 ++++-- 7 files changed, 75 insertions(+), 22 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 7b6fc8ad2e2..8847feaf9b2 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1472,9 +1472,6 @@ bool common_exact_concurrency_init(const common_params & params) { const int n_cols = common_exact_decode_width(params); - // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step is - llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel))); - const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); if (bound) { const int max_cols = atoi(bound); @@ -1488,9 +1485,16 @@ bool common_exact_concurrency_init(const common_params & params) { } } - // a context created later reports n_seq_max times the per-sequence width, which is this - // figure again; reporting it here as well covers a caller that decodes before that - llama_set_exact_decode_width((uint32_t) n_cols); + // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode + // step is; a context created later reports n_seq_max times that figure, which is n_cols + // again, and reporting n_cols here as well covers a caller that decodes before that. Both + // refuse a width the explicit bound above cannot cover, which the check above already + // caught for this process; contexts created earlier by the caller are covered here. + if (!llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel))) || + !llama_set_exact_decode_width((uint32_t) n_cols)) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: the decode width could not be reported, see the error above\n"); + return false; + } return true; } diff --git a/include/llama.h b/include/llama.h index 299249c5daf..1608be8a6a0 100644 --- a/include/llama.h +++ b/include/llama.h @@ -807,8 +807,11 @@ extern "C" { // 1 plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence // set with more tokens than this left to place is a prompt and is prefilled in a ubatch of its // own; a set at or below it is a decode step and stays grouped with the other decodes, so a - // speculative verify batch is not run once per sequence. Process-wide, default 1. - LLAMA_API void llama_set_exact_decode_tokens(uint32_t n_tokens); + // speculative verify batch is not run once per sequence. Process-wide, default 1. Raising it + // widens the decode step of every context that exists, and their width is re-reported with + // it; false, and no change, when an explicit column bound given to the backend cannot cover + // that width (see llama_set_exact_decode_width). + LLAMA_API bool llama_set_exact_decode_tokens(uint32_t n_tokens); LLAMA_API uint32_t llama_exact_decode_tokens(void); // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns: the @@ -816,8 +819,10 @@ extern "C" { // reports its own at creation and a backend keeps the widest it has heard, so a decode of any // context stays within the bound its kernels split at. A caller that builds wider steps than // the contexts imply (a draft of its own, say) reports the width itself, before creating the - // context or before the first decode. Never lowers what was reported. - LLAMA_API void llama_set_exact_decode_width(uint32_t n_cols); + // context or before the first decode. Never lowers what was reported. Returns false, and + // reports nothing, when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is set to a positive figure below + // the width: that bound wins in the backend, so decodes above it would be left batched. + LLAMA_API bool llama_set_exact_decode_width(uint32_t n_cols); LLAMA_API uint32_t llama_exact_decode_width(void); // diff --git a/src/llama-context.cpp b/src/llama-context.cpp index d14123779a5..e8ccf985477 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -123,7 +123,9 @@ llama_context::llama_context( } } - llama_exact_report_n_seq(cparams.n_seq_max); + if (!llama_exact_report_n_seq(cparams.n_seq_max)) { + throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); + } } cparams.n_rs_seq = params.n_rs_seq; diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index 51c50898bdb..20739bb7c96 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -192,23 +192,35 @@ static std::atomic g_exact_decode_tokens{1}; // context created under a narrower figure would batch above the bound it reported. static std::atomic g_exact_max_n_seq{0}; -void llama_exact_report_n_seq(uint32_t n_seq) { +bool llama_exact_report_n_seq(uint32_t n_seq) { + const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed)); + + if (!llama_set_exact_decode_width(n_seq_max * llama_exact_decode_tokens())) { + return false; + } + uint32_t cur = g_exact_max_n_seq.load(std::memory_order_relaxed); while (n_seq > cur && !g_exact_max_n_seq.compare_exchange_weak(cur, n_seq, std::memory_order_relaxed)) { } - llama_set_exact_decode_width(g_exact_max_n_seq.load(std::memory_order_relaxed) * llama_exact_decode_tokens()); + return true; } -void llama_set_exact_decode_tokens(uint32_t n_tokens) { - g_exact_decode_tokens.store(n_tokens > 0 ? n_tokens : 1, std::memory_order_relaxed); +bool llama_set_exact_decode_tokens(uint32_t n_tokens) { + n_tokens = n_tokens > 0 ? n_tokens : 1; + // every context that exists widens with the figure, so the width they will need is + // reported first; a figure the explicit bound cannot cover leaves the old one in place const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed); - if (n_seq > 0) { - llama_set_exact_decode_width(n_seq * llama_exact_decode_tokens()); + if (n_seq > 0 && !llama_set_exact_decode_width(n_seq * n_tokens)) { + return false; } + + g_exact_decode_tokens.store(n_tokens, std::memory_order_relaxed); + + return true; } uint32_t llama_exact_decode_tokens(void) { @@ -220,7 +232,29 @@ uint32_t llama_exact_decode_tokens(void) { // reached through the registry so that a backend that is absent or loaded late costs nothing. static std::atomic g_exact_decode_width{0}; -void llama_set_exact_decode_width(uint32_t n_cols) { +// an explicit column bound given to the CUDA backend wins over the reported width there, so a +// width above it would leave decodes batched past the bound with the mode still reporting itself +// on; a width the bound does not cover is refused instead of stored +static bool llama_exact_width_within_explicit_bound(uint32_t n_cols) { + static const int explicit_cols = []() { + const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); + return val ? atoi(val) : -1; + }(); + + if (explicit_cols > 0 && (uint32_t) explicit_cols < n_cols) { + LLAMA_LOG_ERROR("%s: GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at least %u columns for the decode step just requested; raise it, set it to 0 for no bound, or unset it\n", + __func__, explicit_cols, n_cols); + return false; + } + + return true; +} + +bool llama_set_exact_decode_width(uint32_t n_cols) { + if (!llama_exact_width_within_explicit_bound(n_cols)) { + return false; + } + uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed); while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { @@ -239,6 +273,8 @@ void llama_set_exact_decode_width(uint32_t n_cols) { fn((int) widest); } } + + return true; } uint32_t llama_exact_decode_width(void) { diff --git a/src/llama-impl.h b/src/llama-impl.h index 3c720cec150..de5c6a2d216 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -112,4 +112,4 @@ bool llama_exact_concurrency(); // [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so that the // decode width every context needs is known to the backend and follows llama_set_exact_decode_tokens -void llama_exact_report_n_seq(uint32_t n_seq); +bool llama_exact_report_n_seq(uint32_t n_seq); diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 3fea4bb0ba3..e8d80c770ba 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -98,7 +98,9 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // leaves a different gated delta net state than the same prompt processed alone. // Giving such a sequence a ubatch of its own removes that. A plain decode step, one // token per sequence, is already exact and stays batched. - const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0; + // The figure is passed whenever the mode is on, not only when a prompt is present: + // it also keeps sets of unequal token counts apart (see llama_batch_allocr::split_equal). + const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 8943faf316c..61463c72964 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -432,8 +432,12 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch // so that the rollback snapshots remain valid // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: a recurrent state that a - // prompt leaves behind depends on what shared its ubatch, so isolate the prompts - const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0; + // prompt leaves behind depends on what shared its ubatch, so isolate the prompts. + // The figure is passed whenever the mode is on, not only when a prompt is present: + // it also keeps sets of unequal token counts apart, and a three-token verify step + // placed beside a two-token one as two now and one later would be reduced in + // chunks the solo run never had. + const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } From 6af0d4e99ac85d51dc643f719b1d2b3764158821 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 18:58:53 +0000 Subject: [PATCH 51/81] server: a restore landing inside the context-shift wait sends the resume notice preempt_wait_for_shift() puts a restoring slot back into the state it was parked from, the same as the poll in update_preempt_copies(), but never sent the mirror of the park notice, so a stream that saw the park never saw the resume and kept the parked keepalive. It is sent there now. --- tools/server/server-context.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7943a28860a..3bc163b5f86 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3757,6 +3757,10 @@ struct server_context_impl { metrics.n_resume++; + // the restore landed here rather than in update_preempt_copies(), so the mirror + // of the park notice is sent here: no later poll sees this slot restoring + send_preempt_notice(slot, false); + SLT_WRN(slot, "restore completed after %.2f ms (waited for, a context shift is due): %d tokens back in the cache, kv %d/%d, preemptions %d\n", (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, slot.prompt.n_tokens(), From bf00ac38a0af5c26c8156b654967dad131d649dc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 19:30:46 +0000 Subject: [PATCH 52/81] exact concurrency: soft-capped attention refused at load; width reports serialised and monotonic in the backend; the token figure never lowered The paged attention kernel has no soft-capped variant and asserted on its first call; a model with attn_soft_cap is refused when the cache is created, with the other unsupported layouts. Two contexts created at once could hand the backend a narrower width after a wider one: the report is now made under a lock, and the CUDA setter keeps the widest figure it has heard whatever the order. A narrower context set up after a speculative one lowered the process-wide token figure and turned the existing context's verify steps into prompts, serialising them; the figure is never lowered now, as the width never was. --- ggml/src/ggml-cuda/ggml-cuda.cu | 6 +++++- include/llama.h | 3 ++- src/llama-impl.cpp | 12 ++++++++++++ src/llama-kv-cache.cpp | 8 ++++++++ 4 files changed, 27 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e5e66f95265..8c5b2a408c3 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1856,7 +1856,11 @@ int ggml_cuda_batch_invariant() { static std::atomic g_exact_decode_width{0}; void ggml_backend_cuda_set_exact_decode_width(int n_cols) { - g_exact_decode_width.store(n_cols > 0 ? n_cols : 0, std::memory_order_relaxed); + // monotonic: the widest figure ever reported stays, whatever order the reports arrive in + int cur = g_exact_decode_width.load(std::memory_order_relaxed); + + while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { + } } int ggml_cuda_batch_invariant_max_cols() { diff --git a/include/llama.h b/include/llama.h index 1608be8a6a0..09c331ff3e1 100644 --- a/include/llama.h +++ b/include/llama.h @@ -810,7 +810,8 @@ extern "C" { // speculative verify batch is not run once per sequence. Process-wide, default 1. Raising it // widens the decode step of every context that exists, and their width is re-reported with // it; false, and no change, when an explicit column bound given to the backend cannot cover - // that width (see llama_set_exact_decode_width). + // that width (see llama_set_exact_decode_width). Never lowers what was set: a narrower context + // set up later must not turn an existing context's verify steps into prompts. LLAMA_API bool llama_set_exact_decode_tokens(uint32_t n_tokens); LLAMA_API uint32_t llama_exact_decode_tokens(void); diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index 20739bb7c96..2a33c57e869 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -210,6 +211,12 @@ bool llama_exact_report_n_seq(uint32_t n_seq) { bool llama_set_exact_decode_tokens(uint32_t n_tokens) { n_tokens = n_tokens > 0 ? n_tokens : 1; + // never lowered: a narrower context set up later would otherwise turn the verify steps of + // an existing speculative context into prompts and serialise them + if (n_tokens <= g_exact_decode_tokens.load(std::memory_order_relaxed)) { + return true; + } + // every context that exists widens with the figure, so the width they will need is // reported first; a figure the explicit bound cannot cover leaves the old one in place const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed); @@ -255,6 +262,11 @@ bool llama_set_exact_decode_width(uint32_t n_cols) { return false; } + // one reporter at a time: the widest figure is read and handed to the backends below as + // one step, so a narrower report cannot overtake a wider one on its way to a backend + static std::mutex mutex; + std::lock_guard lock(mutex); + uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed); while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 1c31e57a761..bb947c9ed09 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -284,6 +284,14 @@ llama_kv_cache::llama_kv_cache( throw std::runtime_error("exact concurrency: unsupported attention head size"); } + // [TAG_EXACT_CONCURRENCY] the paged attention kernel has no soft-capped variant and would + // assert on its first call, so a soft-capped model is refused at load instead + if (exact_pages && hparams.attn_soft_cap) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but this model soft-caps its attention logits (%.1f), " + "which the paged attention kernel does not apply\n", __func__, hparams.f_attn_logit_softcapping); + throw std::runtime_error("exact concurrency: attention soft cap is not supported"); + } + if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, " "which has no paged attention: every layer must be offloaded to the CUDA backend " From 98fe86dfc0a42940906f57ea6a27cf9bdc3907ba Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 19:56:18 +0000 Subject: [PATCH 53/81] exact concurrency: the setup runs before anything is loaded; one lock for the token figure, the sequence count and the width common_init_from_params() checked the mode after the context existed, so a caller that skipped common_params_parse() could be handed a live context under a bound the setup had just refused. The check moved to the front of the init result's constructor, ahead of the fitting contexts and the model load; on failure nothing is loaded. The token figure, the widest sequence count and the width moved as three separate atomics, so a context reporting its count while the figure changed could leave the backend with a width that covered neither. One recursive lock now spans every transition. --- common/common.cpp | 14 +++++++++----- src/llama-impl.cpp | 15 +++++++++++---- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 8847feaf9b2..966620c7910 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1290,6 +1290,15 @@ struct common_init_result::impl { common_init_result::common_init_result(common_params & params, bool model_only) : pimpl(new impl{}) { + // [TAG_EXACT_CONCURRENCY] before any context exists, the fitting ones included: the + // per-sequence figure and the column bound are checked against the explicit bound first, + // so a context is never created under a figure the bound does not cover. A caller that + // skipped common_params_parse() gets the same check here; on failure nothing is loaded. + if (!model_only && !common_exact_concurrency_init(params)) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: refusing to load the model, see the error above\n"); + return; + } + auto mparams = common_model_params_to_llama(params); auto cparams = common_context_params_to_llama(params); @@ -1520,11 +1529,6 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode const llama_vocab * vocab = llama_model_get_vocab(model); - // [TAG_EXACT_CONCURRENCY] before the warmup, which is the first graph this process computes - if (!common_exact_concurrency_init(params)) { - return res; - } - if (params.ctx_shift && !llama_memory_can_shift(llama_get_memory(lctx))) { COM_WRN("%s", "KV cache shifting is not supported for this context, disabling KV cache shifting\n"); params.ctx_shift = false; diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index 2a33c57e869..266a14dd4d4 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -187,6 +187,12 @@ bool llama_exact_concurrency() { // [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h static std::atomic g_exact_decode_tokens{1}; +// one lock for the token figure, the sequence count and the width: the three move together +// (a context reports its count and the width that follows; a new token figure re-reports the +// width for every count seen), and a report interleaved with a change of figure could leave +// the backend with a width that covers neither. Recursive, since the setters call each other. +static std::recursive_mutex g_exact_mutex; + // the most sequences any context so far was created with. The tokens figure is process // wide, so raising it widens the decode step of every context that already exists; the // width those contexts reported at creation is re-reported here with the new figure, or a @@ -194,6 +200,8 @@ static std::atomic g_exact_decode_tokens{1}; static std::atomic g_exact_max_n_seq{0}; bool llama_exact_report_n_seq(uint32_t n_seq) { + std::lock_guard lock(g_exact_mutex); + const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed)); if (!llama_set_exact_decode_width(n_seq_max * llama_exact_decode_tokens())) { @@ -211,6 +219,8 @@ bool llama_exact_report_n_seq(uint32_t n_seq) { bool llama_set_exact_decode_tokens(uint32_t n_tokens) { n_tokens = n_tokens > 0 ? n_tokens : 1; + std::lock_guard lock(g_exact_mutex); + // never lowered: a narrower context set up later would otherwise turn the verify steps of // an existing speculative context into prompts and serialise them if (n_tokens <= g_exact_decode_tokens.load(std::memory_order_relaxed)) { @@ -262,10 +272,7 @@ bool llama_set_exact_decode_width(uint32_t n_cols) { return false; } - // one reporter at a time: the widest figure is read and handed to the backends below as - // one step, so a narrower report cannot overtake a wider one on its way to a backend - static std::mutex mutex; - std::lock_guard lock(mutex); + std::lock_guard lock(g_exact_mutex); uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed); From 02a3e11d43efaa6b94101c6c16ba2e9af731238c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 19:57:19 +0000 Subject: [PATCH 54/81] server: the abort sweep also leaves a slot whose copy is in flight alone --- tools/server/server-context.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d4961788d93..88b6d7f7e20 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3179,7 +3179,8 @@ struct server_context_impl { // [TAG_PREEMPT] a parked slot took no part in what failed: its sequence is in // host RAM, not in the cache, and it comes back when there is room, the same as // in the decode error sweep - if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + // [TAG_PREEMPT_ASYNC] a slot whose copy is in flight is out of the round as well + if (slot.is_processing() && !slot.preempt_is_out()) { send_error(slot, reason, ERROR_TYPE_SERVER); slot.release(); } From 5a13c675840827b6e3902ac43e01b77953b8e606 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 20:58:26 +0000 Subject: [PATCH 55/81] server: asynchronous copies and the rotation, the budget and the cache-reuse shift No rotation runs while a park is still copying: its cells are still held, the head would not fit yet, and the rotation would only park another resident on top. An asynchronous rotation park is not re-examined on the same pass either; the head is re-examined when the copy lands. The rotation's budget no longer counts an asynchronous head's bytes as leaving, since its pinned buffer is kept through the restore by design, and charges the resident only what it does not hold yet. A restored slot's buffer is returned when the pool is over its budget, so a buffer held by a running slot cannot keep every other slot from being parked. The cache-reuse shift is applied by the same in-place graph as a context shift, so it sets the flag that makes the round wait for copies in flight. --- tools/server/server-context.cpp | 53 +++++++++++++++++++++++++++++---- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 88b6d7f7e20..7a13443b4d5 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3322,16 +3322,38 @@ struct server_context_impl { // resident, so its bytes are on their way out and are not held against the resident. // A budget that holds one sequence but not two would otherwise refuse every rotation // and leave the head parked for as long as the resident cares to generate. + // [TAG_PREEMPT_ASYNC] an asynchronous head keeps its pinned buffer through the restore + // (see preempt_state_size), so nothing of it leaves; what the resident already holds is + // reused, as in preempt_fits_budget, and only the rest is charged. bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) const { if (params_base.preempt_ram_mib < 0) { return true; } + const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; + const size_t used = preempt_ram_used(); + const size_t leaving = head.preempt_is_async() ? 0 : std::min(used, head.preempt_state_size()); + const size_t held = slot.preempt_state_size(); + const size_t need = slot.preempt_state_required(); + const size_t extra = need > held ? need - held : 0; + + return used - leaving + extra <= budget; + } + + // [TAG_PREEMPT_ASYNC] a restored slot keeps its pinned buffer for its next park, which is + // worth it while the budget has room for it and not otherwise: over budget, a buffer held + // by a slot that is running would keep every other slot from being parked at all + void preempt_trim_ram(server_slot & slot) { + if (params_base.preempt_ram_mib < 0) { + return; + } + const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; - const size_t used = preempt_ram_used(); - const size_t leaving = std::min(used, head.preempt_state_size()); - return used - leaving + slot.preempt_state_required() <= budget; + if (preempt_ram_used() > budget && slot.preempt_state_size() > 0) { + SLT_INF(slot, "%.1f MiB of parked RAM returned: the pool is over its budget\n", slot.preempt_state_size() / (1024.0 * 1024.0)); + slot.preempt_state_free(); + } } // cells the slot will ask for on its next step once it is back in the pool @@ -3634,6 +3656,8 @@ struct server_context_impl { if (slot.preempt_restore_poll()) { metrics.n_resume++; + preempt_trim_ram(slot); + SLT_WRN(slot, "restore completed after %.2f ms: %d tokens back in the cache, kv %d/%d, preemptions %d\n", (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, slot.prompt.n_tokens(), @@ -3672,6 +3696,8 @@ struct server_context_impl { metrics.n_resume++; + preempt_trim_ram(slot); + SLT_WRN(slot, "restore completed after %.2f ms (waited for, a context shift is due): %d tokens back in the cache, kv %d/%d, preemptions %d\n", (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, slot.prompt.n_tokens(), @@ -3815,7 +3841,15 @@ struct server_context_impl { if (!best) { server_slot * head = parked.front(); - if (ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { + // [TAG_PREEMPT_ASYNC] a park still copying holds its cells, so the head would + // not fit yet and a rotation now would only park another resident on top + bool parking = false; + + for (const auto & slot : slots) { + parking = parking || slot.state == SLOT_STATE_PREEMPTING; + } + + if (!parking && ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { // the resident whose cells let the head in, the smallest of those; failing // one that does so alone, the largest, since it makes the most room. Taking // the first shifting resident in slot order could park one too small to @@ -3870,7 +3904,12 @@ struct server_context_impl { pick_enough ? "" : " (not enough room by itself)", slot.n_preempt); - best = head; // re-examined by the loop, which sees the room it just got + // [TAG_PREEMPT_ASYNC] a synchronous park has released its cells, so the + // head is re-examined now; an asynchronous one has not, and the head is + // re-examined on the pass that sees the copy land + if (slot.state == SLOT_STATE_PREEMPTED) { + best = head; + } } } @@ -4548,6 +4587,10 @@ struct server_context_impl { slot.mem.seq_rm (slot.id, head_p, head_c); slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift); + // [TAG_PREEMPT_ASYNC] applied inside the next llama_decode by the + // same in-place graph as a context shift, see preempt_wait_for_shift + preempt_shift_pending = true; + for (size_t i = 0; i < n_match; i++) { slot.prompt.tokens.set_token(head_p + i, slot.prompt.tokens[head_c + i]); n_past++; From 77d318522e4b24d3f8e2f74cce1d4fdc52455a1c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 20:58:28 +0000 Subject: [PATCH 56/81] exact concurrency: an isolated ubatch takes only sets that finish in it; one bound check at context creation Under isolation the equal-count guard chose which sets join a ubatch, but the expansion could still cut them all part way when the sets together exceeded n_ubatch, the chunking the guard exists to prevent. A set that would not finish in the ubatch waits for the next one. The context constructor checked the explicit column bound itself and then again through the report; the report's refusal is the error now. --- src/llama-batch.cpp | 4 ++++ src/llama-context.cpp | 14 +------------- 2 files changed, 5 insertions(+), 13 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 5db683db281..cc73d83963c 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -611,6 +611,10 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, n_left_first = n_left; } else if (n_left != n_left_first) { continue; + } else if ((cur_seq_set.size() + 1) * n_left_first > n_ubatch) { + // one more set would not finish in this ubatch: the expansion below would + // then cut every set part way, the chunking the guard exists to prevent + break; } } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index e8ccf985477..68caf978337 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -108,21 +108,9 @@ llama_context::llama_context( // The sequence count is what is reported: the tokens figure can be raised later for the // whole process, and the width then follows it for this context too. if (llama_exact_concurrency()) { - const uint32_t n_cols = cparams.n_seq_max * llama_exact_decode_tokens(); - // an explicit column bound wins over the reported width in the backend, so one below // this context's width would leave its decodes batched above the bound with the mode - // still reporting itself on; refuse it here, the way the server's setup does - if (const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS")) { - const int max_cols = atoi(bound); - - if (max_cols > 0 && (uint32_t) max_cols < n_cols) { - LLAMA_LOG_ERROR("%s: GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at least %u to cover a decode step of %u sequences; raise it, set it to 0 for no bound, or unset it\n", - __func__, max_cols, n_cols, cparams.n_seq_max); - throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); - } - } - + // still reporting itself on; the report refuses that, and the refusal is an error here if (!llama_exact_report_n_seq(cparams.n_seq_max)) { throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); } From 82f40df769c8153b5d2c602d973267c5b2807805 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 21:35:50 +0000 Subject: [PATCH 57/81] preempt: post no copies for a transfer that failed part way, park synchronously into pageable memory, offer no transfer for state that is not on a device The asynchronous state adapters posted their queued copies from the destructor whether or not serialisation had got to the end, so a buffer one byte short made llama_state_seq_copy_get() return 0 while 64 copies were still reading it, and the restore counterpart wrote into cells the failed restore had already given up. Both adapters commit only after the serialisation succeeds; a failed one posts nothing. A host buffer type may hand back pageable memory instead of failing, and a copy into or out of pageable memory blocks the thread that issued it. The server takes a one MiB buffer at load and looks at what it got; if it is pageable the slots park synchronously and say so. state_seq_copy_init() returned a transfer whenever a device could copy asynchronously, even when every state tensor lived in host memory (most layers on the CPU) and every copy took the synchronous branch. It returns NULL unless every non-empty memory buffer lives on one of its devices. --- src/llama-context.cpp | 69 ++++++++++++++++++++++++++++++++- tools/server/server-context.cpp | 28 ++++++++++--- 2 files changed, 90 insertions(+), 7 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 09e53123da6..c73ba5ddacf 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3332,7 +3332,19 @@ class llama_io_write_host_async : public llama_io_write_i { llama_io_write_host_async(uint8_t * p, size_t len, llama_state_seq_copy & cpy) : ptr(p), buf_size(len), cpy(cpy) {} + // The transfers are posted from the destructor, and only once serialisation has got to + // the end: a failure part way, a buffer one byte short say, is reported to the caller as + // a zero return, and a caller told that is free to reuse the buffer at once. Copies + // posted regardless would still be reading it. + void commit() { + committed = true; + } + ~llama_io_write_host_async() { + if (!committed) { + return; + } + llama_io_emit(winfos, 0, winfos.size(), [this](ggml_tensor * tensor, uint8_t * ptr, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { @@ -3385,6 +3397,8 @@ class llama_io_write_host_async : public llama_io_write_i { std::vector winfos; llama_state_seq_copy & cpy; + + bool committed = false; }; class llama_io_read_host_async : public llama_io_read_i { @@ -3392,7 +3406,18 @@ class llama_io_read_host_async : public llama_io_read_i { llama_io_read_host_async(const uint8_t * p, size_t len, llama_state_seq_copy & cpy) : ptr(p), buf_size(len), cpy(cpy) {} + // see llama_io_write_host_async::commit(): the restore that failed part way has already + // dropped the sequence, and copies posted for it would write into cells that are no + // longer its own + void commit() { + committed = true; + } + ~llama_io_read_host_async() { + if (!committed) { + return; + } + // No whole-tensor staging here, unlike the synchronous path above. Staging reads a // tensor, patches this sequence's bytes into the host copy and writes the whole // tensor back, which preserves the neighbours only while nothing else is touching @@ -3454,6 +3479,8 @@ class llama_io_read_host_async : public llama_io_read_i { std::vector rinfos; llama_state_seq_copy & cpy; + + bool committed = false; }; static constexpr uint32_t io_magic = 0xaf143cd8; @@ -3594,6 +3621,36 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { return nullptr; } + // The devices above are the ones the graphs run on, not necessarily the ones the state + // lives on: with most layers left on the CPU the KV cache is host memory, and a tensor + // there takes the synchronous branch of backend_for(). A transfer whose every copy would + // do that is not asynchronous, whatever it is called, and the caller is better served by + // the synchronous calls it already has and a log line that says so. + if (memory) { + bool on_device = false; + + for (const auto & [buft, size] : memory->memory_breakdown()) { + if (size == 0) { + continue; + } + + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); + + if (ggml_backend_buft_is_host(buft) || !dev || buft != ggml_backend_dev_buffer_type(dev) || + cpy->devs.find(dev) == cpy->devs.end()) { + LLAMA_LOG_INFO("%s: the sequence state is not all in device memory (%s), so it is copied synchronously\n", + __func__, ggml_backend_buft_name(buft)); + return nullptr; + } + + on_device = true; + } + + if (!on_device) { + return nullptr; + } + } + cpy->can_pin = cpy->host_buffer_type() != ggml_backend_cpu_buffer_type(); return cpy.release(); @@ -3636,7 +3693,11 @@ size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size io.write(&io_magic, sizeof(io_magic)); io.write(&seq_id, sizeof(seq_id)); - return state_seq_write_data(io, seq_id, flags); + const size_t n = state_seq_write_data(io, seq_id, flags); + + io.commit(); + + return n; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error saving state: %s\n", __func__, err.what()); return 0; @@ -3681,7 +3742,11 @@ size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size llama_seq_id seq_id_read; io.read(&seq_id_read, sizeof(seq_id_read)); - return state_seq_read_data(io, seq_id, flags); + const size_t n = state_seq_read_data(io, seq_id, flags); + + io.commit(); + + return n; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); return 0; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a0d1564405c..a4b40961842 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1725,11 +1725,29 @@ struct server_context_impl { if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { if (preempt_async_ok) { - // no buffer has been allocated yet, so this is what the backend offers, - // not what is held. What was actually got is reported by the first park, - // because a host buffer type may still hand back ordinary memory. - SRV_INF("preemption: parking and resuming asynchronously, backend offers %s host memory\n", - llama_state_seq_copy_buf_can_pin(slots[0].preempt_cpy_tgt.get()) ? "pinned" : "pageable"); + // Pinned host memory is what lets a copy run beside the decode: one into or + // out of pageable memory is staged by the driver and blocks the thread that + // issued it, which is the stall the asynchronous path exists to remove. A + // host buffer type is free to hand back ordinary memory instead of failing + // (GGML_CUDA_NO_PINNED, or a pinning limit), and that is only knowable from + // a buffer, so a small one is taken and looked at before the first park. + bool pinned = llama_state_seq_copy_buf_can_pin(slots[0].preempt_cpy_tgt.get()); + + if (pinned) { + auto * cpy = slots[0].preempt_cpy_tgt.get(); + + pinned = llama_state_seq_copy_buf_resize(cpy, 1u << 20) != nullptr && + llama_state_seq_copy_buf_is_pinned(cpy); + + llama_state_seq_copy_buf_free(cpy); + } + + if (pinned) { + SRV_INF("%s", "preemption: parking and resuming asynchronously through pinned host memory\n"); + } else { + SRV_WRN("%s", "preemption: the host memory on offer is pageable, so a copy would block the decode; parking and resuming synchronously\n"); + preempt_async_ok = false; + } } else { SRV_WRN("%s", "preemption: this backend cannot copy asynchronously, parking and resuming synchronously\n"); } From 2f2258dc0faa6b39be80fc20d0c1b45adecfa3fd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 21:35:50 +0000 Subject: [PATCH 58/81] exact concurrency: refuse a whole-context restore before it clears the cache, publish the width once construction succeeds, refuse the page table on every backend that ignores it A whole-context restore under LLAMA_EXACT_CONCURRENCY was refused inside state_read_meta(), after which the generic restore path cleared the live cache. It is refused at the top of state_read_data() and llama_kv_cache::state_read() now, before a byte is read, so llama_state_set_data() returns 0 and the sequences are untouched. The context reported its sequence count at the front of the constructor, so a construction that failed later on left a width behind that no context needed. The count is checked against the explicit column bound early and reported at the end, once nothing can fail any more. WebGPU, ExecuTorch, Hexagon, OpenVINO and RPC advertised FLASH_ATTN_EXT with the page table in src[5] that only the CUDA kernels read; they refuse it like CANN, Metal, OpenCL, SYCL and Vulkan already did. DSpark runs on the DFlash implementation and turns causal attention off on its draft context, so it is refused alongside DFlash. The decode width is computed in 64 bits and refused above INT32_MAX instead of wrapping. The probe's PROBE_A_PERM keeps prompt 0 on sequence 0, which phase B compares against. --- common/common.cpp | 23 ++++++++++++---- ggml/src/ggml-et/ggml-et.cpp | 6 ++++ ggml/src/ggml-hexagon/ggml-hexagon.cpp | 4 ++- ggml/src/ggml-openvino/ggml-openvino.cpp | 5 ++++ ggml/src/ggml-rpc/ggml-rpc.cpp | 6 +++- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 7 +++++ scripts/batchinv/probe.cpp | 4 +++ src/llama-context.cpp | 19 ++++++++++++- src/llama-impl.cpp | 35 ++++++++++++++++++++++-- src/llama-impl.h | 4 +++ src/llama-kv-cache.cpp | 17 +++++++----- 11 files changed, 112 insertions(+), 18 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 966620c7910..04ce8744359 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1455,13 +1455,17 @@ bool common_exact_concurrency() { // [TAG_EXACT_CONCURRENCY] int common_exact_decode_width(const common_params & params) { - const int n_slots = std::max(1, params.n_parallel); + const int64_t n_slots = std::max(1, params.n_parallel); // the draft tokens a slot carries into the verify ubatch alongside its accepted token, per // speculation type, from the same place the speculation code takes its own width - const int n_draft = std::max(0, (int) common_speculative_n_max(¶ms.speculative)); + const int64_t n_draft = std::max(0, (int) common_speculative_n_max(¶ms.speculative)); - return n_slots*(1 + n_draft); + // the product is what a backend is asked to split columns by, as an int; one that does not + // fit is reported as such rather than wrapped + const int64_t n_cols = n_slots*(1 + n_draft); + + return n_cols > INT32_MAX ? -1 : (int) n_cols; } // [TAG_EXACT_CONCURRENCY] @@ -1471,16 +1475,23 @@ bool common_exact_concurrency_init(const common_params & params) { } // DFlash drafting turns causal attention off on its draft context, and the paged - // attention the mode runs on needs it; say so instead of asserting in the graph + // attention the mode runs on needs it; say so instead of asserting in the graph. DSpark + // is the same implementation under another name, so it is refused with it. for (const auto type : params.speculative.types) { - if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) { - COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash: it disables causal attention, which the paged attention needs\n"); + if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash or draft-dspark: both disable causal attention on the draft, which the paged attention needs\n"); return false; } } const int n_cols = common_exact_decode_width(params); + if (n_cols < 0) { + COM_ERR("LLAMA_EXACT_CONCURRENCY: a decode step of %d slots with %d draft tokens each is too wide to report\n", + std::max(1, params.n_parallel), std::max(0, (int) common_speculative_n_max(¶ms.speculative))); + return false; + } + const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); if (bound) { const int max_cols = atoi(bound); diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp index b87b189a57a..a3792f3852b 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1266,6 +1266,12 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm (op->src[1]->ne[1] % op->src[4]->ne[1] == 0); break; case GGML_OP_FLASH_ATTN_EXT: + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only + // the CUDA backend reads + if (op->src[5]) { + supported = false; + break; + } if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] && (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && op->src[2] && (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr && diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index e8a5009b381..aa20083ec05 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -4157,7 +4157,9 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons break; case GGML_OP_FLASH_ATTN_EXT: - supp = ggml_hexagon_supported_flash_attn_ext(sess, op); + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only + // the CUDA backend reads + supp = op->src[5] == nullptr && ggml_hexagon_supported_flash_attn_ext(sess, op); break; case GGML_OP_SET_ROWS: diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index e299e16c778..dfc9f90926f 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -1128,6 +1128,11 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { break; } case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only + // the CUDA backend reads + if (op->src[5]) { + return true; + } float scale = 1.0f; float max_bias = 0.0f; float logit_softcap = 0.0f; diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 69a8a08ae17..6fb8851a904 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1915,7 +1915,11 @@ static ggml_backend_buffer_type_t ggml_backend_rpc_device_get_buffer_type(ggml_b static bool ggml_backend_rpc_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { GGML_UNUSED(dev); - GGML_UNUSED(op); + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the + // CUDA backend reads; the remote end is not asked, so it is not claimed here + if (op->op == GGML_OP_FLASH_ATTN_EXT && op->src[5]) { + return false; + } //TODO: call the remote backend and cache the results return true; } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 2434848a55a..70462a97f3c 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4408,6 +4408,13 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const break; case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only + // the CUDA backend reads + if (op->src[5]) { + supports_op = false; + break; + } + // conservative support checks for whether the more resource-intensive shader paths // can be used, to avoid cases where flash_attn is assigned to the CPU later on supports_op = src0->type == GGML_TYPE_F32 && diff --git a/scripts/batchinv/probe.cpp b/scripts/batchinv/probe.cpp index 151463cd768..c4e478f97c7 100644 --- a/scripts/batchinv/probe.cpp +++ b/scripts/batchinv/probe.cpp @@ -172,9 +172,13 @@ int main(int argc, char ** argv) { const int a_fill = getenv("PROBE_A_FILL") ? atoi(getenv("PROBE_A_FILL")) : 1; // PROBE_A_PERM reorders which prompt goes into which sequence in phase A. With the same // multiset of prompts the cache keeps its length but the masked cells hold different data. + // Phase B decodes prompt 0's first token on sequence 0, so the permutation may only move + // the neighbours: sequence 0 keeps prompt 0, or the two phases would compare different + // sequences. int a_perm[4] = {0, 1, 2, 3}; if (const char * perm = getenv("PROBE_A_PERM")) { for (int k = 0; k < 4 && perm[2*k]; ++k) a_perm[k] = perm[2*k] - '0'; + if (a_perm[0] != 0) { fprintf(stderr, "PROBE_A_PERM must keep prompt 0 on sequence 0\n"); return 1; } } { llama_context * ctx = make_ctx(); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 68caf978337..8d35922a47b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -107,11 +107,14 @@ llama_context::llama_context( // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width. // The sequence count is what is reported: the tokens figure can be raised later for the // whole process, and the width then follows it for this context too. + // Checked here and reported at the end of the constructor: the count is process-wide + // state that outlives a context, so a construction that fails later on, an unsupported + // cache layout say, must not leave a width behind that no context needs. if (llama_exact_concurrency()) { // an explicit column bound wins over the reported width in the backend, so one below // this context's width would leave its decodes batched above the bound with the mode // still reporting itself on; the report refuses that, and the refusal is an error here - if (!llama_exact_report_n_seq(cparams.n_seq_max)) { + if (!llama_exact_check_n_seq(cparams.n_seq_max)) { throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); } } @@ -498,6 +501,13 @@ llama_context::llama_context( sampling.token_ids_full_vocab[i] = i; } } + + // [TAG_EXACT_CONCURRENCY] nothing above can fail any more, so the width this context + // needs is published now; checked against the explicit bound at the top, so this + // cannot refuse unless the bound moved underneath it, which is an error all the same + if (llama_exact_concurrency() && !llama_exact_report_n_seq(cparams.n_seq_max)) { + throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); + } } llama_context::~llama_context() { @@ -3255,6 +3265,13 @@ size_t llama_context::state_write_data(llama_io_write_i & io) { } size_t llama_context::state_read_data(llama_io_read_i & io) { + // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical + // index, which the paged pool owns. Refused here, before anything is parsed, so that the + // cache the caller has is left as it was: the generic restore path clears it on failure. + if (memory && memory->alloc_granularity() > 1) { + throw std::runtime_error("whole-context restore is not supported with LLAMA_EXACT_CONCURRENCY, restore per sequence"); + } + LLAMA_LOG_DEBUG("%s: reading state\n", __func__); // read model info diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index 266a14dd4d4..0b218bc64e7 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -199,12 +199,41 @@ static std::recursive_mutex g_exact_mutex; // context created under a narrower figure would batch above the bound it reported. static std::atomic g_exact_max_n_seq{0}; +static bool llama_exact_width_within_explicit_bound(uint32_t n_cols); + +// the width is sequences times tokens, handed to a backend as an int; a product that does not +// fit is refused rather than wrapped +static bool llama_exact_width_of(uint32_t n_seq, uint32_t n_tokens, uint32_t & n_cols) { + const uint64_t w = (uint64_t) n_seq * (uint64_t) n_tokens; + + if (w > (uint64_t) INT32_MAX) { + LLAMA_LOG_ERROR("%s: a decode step of %u sequences with %u tokens each is too wide to report\n", __func__, n_seq, n_tokens); + return false; + } + + n_cols = (uint32_t) w; + + return true; +} + +bool llama_exact_check_n_seq(uint32_t n_seq) { + std::lock_guard lock(g_exact_mutex); + + const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed)); + + uint32_t n_cols = 0; + + return llama_exact_width_of(n_seq_max, llama_exact_decode_tokens(), n_cols) && llama_exact_width_within_explicit_bound(n_cols); +} + bool llama_exact_report_n_seq(uint32_t n_seq) { std::lock_guard lock(g_exact_mutex); const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed)); - if (!llama_set_exact_decode_width(n_seq_max * llama_exact_decode_tokens())) { + uint32_t n_cols = 0; + + if (!llama_exact_width_of(n_seq_max, llama_exact_decode_tokens(), n_cols) || !llama_set_exact_decode_width(n_cols)) { return false; } @@ -231,7 +260,9 @@ bool llama_set_exact_decode_tokens(uint32_t n_tokens) { // reported first; a figure the explicit bound cannot cover leaves the old one in place const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed); - if (n_seq > 0 && !llama_set_exact_decode_width(n_seq * n_tokens)) { + uint32_t n_cols = 0; + + if (n_seq > 0 && (!llama_exact_width_of(n_seq, n_tokens, n_cols) || !llama_set_exact_decode_width(n_cols))) { return false; } diff --git a/src/llama-impl.h b/src/llama-impl.h index de5c6a2d216..65d56a51d1d 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -113,3 +113,7 @@ bool llama_exact_concurrency(); // [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so that the // decode width every context needs is known to the backend and follows llama_set_exact_decode_tokens bool llama_exact_report_n_seq(uint32_t n_seq); + +// the same check without the report: whether a context of n_seq sequences could be reported +// under the explicit column bound, for a constructor that may still fail after asking +bool llama_exact_check_n_seq(uint32_t n_seq); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index bb947c9ed09..c384e46d520 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2362,6 +2362,14 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla } void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical + // index, which the paged pool owns. Refused before a byte is read, so that the failure + // path below, which clears the cache, is never entered for it. + if (exact_pages && seq_id == -1) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state restore only\n", __func__); + throw std::runtime_error("whole-cache restore is not supported with LLAMA_EXACT_CONCURRENCY"); + } + // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -2610,13 +2618,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 } else { // whole KV cache restore - // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical - // index, which the paged pool owns. Report it like every other failure in this function. - if (exact_pages) { - LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state " - "restore only\n", __func__); - return false; - } + // [TAG_EXACT_CONCURRENCY] refused at the top of state_read(), before anything is read + GGML_ASSERT(!exact_pages); if (cell_count > cells.size()) { LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__); From 2b4a6912299776ad5e335b2f22adaaf20f0ef823 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 21:36:26 +0000 Subject: [PATCH 59/81] tests: a state transfer that fails one byte short posts no copies, on both sides --- tests/test-state-seq-copy.cpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test-state-seq-copy.cpp b/tests/test-state-seq-copy.cpp index a4633954cee..dabc4db50c3 100644 --- a/tests/test-state-seq-copy.cpp +++ b/tests/test-state-seq-copy.cpp @@ -113,6 +113,14 @@ int main(int argc, char ** argv) { fprintf(stderr, "%s : oversized, empty and ON_DEVICE transfers are all refused\n", __func__); + // a transfer that fails part way, one byte short of the state, must post nothing: the + // caller is told it failed and is free to reuse the buffer at once + CHECK(llama_state_seq_copy_get(cpy, size - 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + CHECK(llama_state_seq_copy_n_copies(cpy) == 0); + CHECK(llama_state_seq_copy_done(cpy)); + + fprintf(stderr, "%s : a transfer one byte short is refused and posts no copies\n", __func__); + // the same call at the size the transfer does own still works, and round-trips std::vector before(llama_state_seq_get_size(ctx, seq_id)); CHECK(llama_state_seq_get_data(ctx, before.data(), before.size(), seq_id) == before.size()); @@ -122,6 +130,11 @@ int main(int argc, char ** argv) { llama_memory_seq_rm(llama_get_memory(ctx), seq_id, -1, -1); + // the restore side too: a buffer claimed one byte short is refused before a copy is posted + CHECK(llama_state_seq_copy_set(cpy, size - 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + CHECK(llama_state_seq_copy_n_copies(cpy) == 0); + CHECK(llama_state_seq_copy_done(cpy)); + CHECK(llama_state_seq_copy_set(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == size); llama_state_seq_copy_wait(cpy); From b4a5f18d813ffba42c94474d3258221e97ec697d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 21:36:59 +0000 Subject: [PATCH 60/81] preempt: say why the asynchronous runway is rounded once rather than per slot --- tools/server/server-context.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index caa287c3657..d1e21e99389 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3617,6 +3617,10 @@ struct server_context_impl { // page allocator any one of them can cost a whole page rather than a cell. Round the // runway up to a page so the park is issued with at least one page of real room // behind it; with a page size of 1 this is the figure it always was. + // The sum is rounded once, not once per running slot: every slot could cross a page + // boundary during the copy, but reserving a page for each of them would keep a page + // per slot out of the users' reach all the time, whereas the case it guards against + // costs one synchronous wait for a copy already in flight, in preempt_wait_in_flight(). return preempt_n_cells( PREEMPT_N_MARGIN + n_running * (1 + preempt_n_spec_max()) * PREEMPT_N_ASYNC_STEPS); } From 9e827b53d98aad78ac65994c4a4340039e623faf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 23:26:59 +0000 Subject: [PATCH 61/81] preempt: a slot whose park buffer comes back pageable parks synchronously from then on The load-time probe takes one MiB and looks at it, but a buffer many times larger can still come back pageable (a host-locking limit, say), and a copy into pageable memory blocks the thread that issued it. The park now looks at the buffer it actually got: a pageable one is given back with the slot's transfers, and the slot takes the synchronous path for this park and every later one. --- tools/server/server-context.cpp | 55 +++++++++++++++++++++------------ 1 file changed, 36 insertions(+), 19 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a4b40961842..2d79a259684 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -545,30 +545,47 @@ struct server_slot { return false; } - if (llama_state_seq_copy_get(preempt_cpy_tgt.get(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { - SLT_ERR(*this, "%s", "failed to issue the copy of the target sequence out of the KV cache\n"); - preempt_state_free(); - return false; - } + // [TAG_PREEMPT_ASYNC] the load-time probe saw pinned memory, but a buffer this much + // larger can still come back pageable (a host-locking limit, say): the host buffer + // type hands back ordinary memory rather than failing, and a copy into pageable + // memory blocks the thread that issued it, which is the stall this path exists to + // remove. Such a slot parks synchronously from now on: its transfers are given + // back and the plain path below takes over, for this park and every later one. + const bool pageable = !llama_state_seq_copy_buf_is_pinned(preempt_cpy_tgt.get()) || + (size_dft > 0 && !llama_state_seq_copy_buf_is_pinned(preempt_cpy_dft.get())); + + if (pageable) { + SLT_WRN(*this, "the host memory for a %.3f MiB park is pageable, so this slot parks synchronously from now on\n", + (size_tgt + size_dft) / (1024.0 * 1024.0)); - if (size_dft > 0 && - llama_state_seq_copy_get(preempt_cpy_dft.get(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { - SLT_ERR(*this, "%s", "failed to issue the copy of the draft sequence out of the KV cache\n"); - preempt_state_free(); - return false; - } + preempt_cpy_tgt.reset(); + preempt_cpy_dft.reset(); + } else { + if (llama_state_seq_copy_get(preempt_cpy_tgt.get(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { + SLT_ERR(*this, "%s", "failed to issue the copy of the target sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + if (size_dft > 0 && + llama_state_seq_copy_get(preempt_cpy_dft.get(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + SLT_ERR(*this, "%s", "failed to issue the copy of the draft sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } - preempt_detach(); + preempt_detach(); - // note: no mem.seq_rm() here. The copy is still reading these cells, so they are - // released in preempt_save_poll() once it has finished with them. - state_before_preempt = state; - state = SLOT_STATE_PREEMPTING; - t_preempt_us = ggml_time_us(); + // note: no mem.seq_rm() here. The copy is still reading these cells, so they are + // released in preempt_save_poll() once it has finished with them. + state_before_preempt = state; + state = SLOT_STATE_PREEMPTING; + t_preempt_us = ggml_time_us(); - n_preempt++; + n_preempt++; - return true; + return true; + } } try { From 918a8bf4f022023953c4e3f2da9cef279d33251c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 23:33:32 +0000 Subject: [PATCH 62/81] exact concurrency: ask the device whether it can run the paged attention, not only which backend it is The KV cache accepted a layer on any CUDA, ROCm or MUSA device by the registry name alone. A build without the flash attention kernels, or a device and head shape they do not cover, would then have the scheduler hand the paged op to the CPU, which accepts the page table as the reference for test-backend-ops and ignores it, and the mode would report itself on while attending in physical order. The constructor now builds the attention op the way the graph does, page table attached, at the widths of a decode step, a verify step and a prompt chunk, and asks the device; a refusal is a load error naming the layer and the types. --- src/llama-kv-cache.cpp | 67 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index c384e46d520..31154639ad9 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -85,6 +85,61 @@ static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) { return strcmp(name, "CUDA") == 0 || strcmp(name, "ROCm") == 0 || strcmp(name, "MUSA") == 0; } +// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a +// layer of this shape. The registry name says which backends carry the kernels; it does not +// say the build has them (FLASH_ATTN_AVAILABLE), nor that the device's architecture, the +// head width and the K/V types land on a kernel. Where they do not, the scheduler would hand +// the op to the CPU, which accepts the page table as the reference for test-backend-ops and +// ignores it, and the mode would report itself on while attending in physical order. So the +// op is built the way the graph builds it, at the widths a decode step, a verify step and a +// prompt chunk use, and the device is asked. +static bool llama_dev_supports_paged_attn( + ggml_backend_dev_t dev, + ggml_type type_k, ggml_type type_v, + uint32_t n_embd_head_k, uint32_t n_embd_head_v, + uint32_t n_head, uint32_t n_head_kv, + uint32_t n_cells, uint32_t page_size) { + if (!llama_dev_has_paged_attn(dev)) { + return false; + } + + ggml_init_params ip = { + /*.mem_size =*/ ggml_tensor_overhead()*16 + ggml_graph_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + ggml_context * ctx = ggml_init(ip); + if (!ctx) { + return false; + } + + bool res = true; + + const int64_t n_kv = page_size; + + for (const int64_t n_tokens : { (int64_t) 1, (int64_t) 4, (int64_t) 16, (int64_t) 512 }) { + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_embd_head_k, n_tokens, n_head, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, type_k, n_embd_head_k, n_kv, n_head_kv, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, type_v, n_embd_head_v, n_kv, n_head_kv, 1); + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_tokens, 1, 1); + + ggml_tensor * op = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf((float) n_embd_head_k), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(op, GGML_PREC_F32); + + op->src[5] = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1 + n_cells/page_size, n_tokens); + + if (!ggml_backend_dev_supports_op(dev, op)) { + res = false; + break; + } + } + + ggml_free(ctx); + + return res; +} + llama_kv_cache::llama_kv_cache( const llama_model & model, const llama_hparams & hparams, @@ -300,6 +355,18 @@ llama_kv_cache::llama_kv_cache( throw std::runtime_error("exact concurrency: KV cache layer is not on the CUDA backend"); } + // [TAG_EXACT_CONCURRENCY] the backend is the right one; ask it whether this layer's + // attention, with the page table attached, lands on one of its kernels at all + if (exact_pages && !llama_dev_supports_paged_attn(model.dev_layer(il), type_k, type_v, + hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), + hparams.n_head(il), hparams.n_head_kv(il), kv_size, exact_page_size)) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but %s cannot run the paged attention for layer %d " + "(K %s, V %s, %u-wide heads): the build or the device has no flash attention kernel for it, " + "and the op would fall to the CPU, which ignores the page table\n", + __func__, dev_name, il, ggml_type_name(type_k), ggml_type_name(type_v), hparams.n_embd_head_k(il)); + throw std::runtime_error("exact concurrency: the device cannot run the paged attention"); + } + ggml_context * ctx = ctx_for_buft(buft); if (!ctx) { throw std::runtime_error("failed to create ggml context for kv cache"); From a4b62f2be5dea11c887e24ea491e10c0d85d5c36 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 23:36:16 +0000 Subject: [PATCH 63/81] preempt: the graphs that follow a restore wait for its copies, on the device A restore writes cells of the KV cache on the copy stream while the other sequences keep decoding on the compute stream, and an attention that is not paged reads every cell up to n_kv, masked ones included, so the reads and the writes were unordered. After the copies are posted and their events recorded, every backend the graphs run on waits for the event of the transfer on its device: a stream wait, so the thread that issued the restore carries on and the next decode starts the moment the copy lands. A park needs none of this: it reads cells nobody writes until it has landed, behind the synchronize at the top of the issue. --- src/llama-context.cpp | 62 ++++++++++++++++++++++++++++++++----------- 1 file changed, 47 insertions(+), 15 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c73ba5ddacf..34d2f51399a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3221,6 +3221,30 @@ struct llama_state_seq_copy { } } + // Order the context's compute behind the copies just recorded, on the device: every + // backend the graphs run on waits for the event of the transfer on its device before + // the next graph it is given. This is a stream wait, not a host wait, so the caller's + // thread carries on and the decode it issues next starts the moment the copy lands. + // + // Needed for a restore and only a restore: its copies write cells of the KV cache + // while other sequences keep decoding, and an attention that is not paged reads every + // cell up to n_kv, masked ones included, so without this the reads and the writes are + // unordered. A park reads cells nobody writes until it has landed, and the decode that + // produced them has been drained by the synchronize() at the top of the issue. + void order_before(const std::vector & compute) { + for (auto & it : devs) { + if (!it.second.pending) { + continue; + } + + for (const auto & backend : compute) { + if (ggml_backend_get_device(backend.get()) == it.first) { + ggml_backend_event_wait(backend.get(), it.second.event); + } + } + } + } + bool done() { bool res = true; @@ -3730,27 +3754,35 @@ size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size cpy.n_copies = 0; - llama_io_read_host_async io(cpy.data, size, cpy); + size_t n = 0; - try { - uint32_t magic_read; - io.read(&magic_read, sizeof(magic_read)); - if (io_magic != magic_read) { - throw std::runtime_error("wrong sequence state magic"); - } + { + llama_io_read_host_async io(cpy.data, size, cpy); - llama_seq_id seq_id_read; - io.read(&seq_id_read, sizeof(seq_id_read)); + try { + uint32_t magic_read; + io.read(&magic_read, sizeof(magic_read)); + if (io_magic != magic_read) { + throw std::runtime_error("wrong sequence state magic"); + } - const size_t n = state_seq_read_data(io, seq_id, flags); + llama_seq_id seq_id_read; + io.read(&seq_id_read, sizeof(seq_id_read)); - io.commit(); + n = state_seq_read_data(io, seq_id, flags); - return n; - } catch (const std::exception & err) { - LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); - return 0; + io.commit(); + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); + return 0; + } } + + // the adapter has posted the copies and recorded the events on its way out; the + // graphs that follow on these devices wait for them, see order_before() + cpy.order_before(backends); + + return n; } bool llama_context::state_load_file(const char * filepath, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { From 294d2a9120a7e63c1188be2eeca2920b3cae4c38 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 00:23:45 +0000 Subject: [PATCH 64/81] preempt: idle parked RAM is given back when another slot needs to park; the copies wait for the compute stream on the device instead of draining it A restored slot keeps its pinned buffer for its next park, and that capacity was charged against --preempt-ram while it held nothing, so a budget that held one sequence was spent for good by the first restore: every later park was refused, and once the slot holding the buffer was the leader nothing could be parked at all. Both budget checks now give idle buffers back before deciding: largest first, never a buffer that still holds a parked sequence or has a copy in flight, never the candidate's own, and never the head of a rotation. state_seq_copy_get and state_seq_copy_set drained the host with synchronize() so that the copies were ordered behind the compute already queued. With the wait that order_before() puts on the compute stream for the previous restore, that drain blocked the calling thread until the previous copy had landed, and two restores issued in one pass ran one after the other with the whole transfer back on the decode loop. Each copy device now carries a second event, recorded on the compute stream and waited for on the copy stream, so the ordering is on the device and the host drains nothing. New server test: two sequences parked in turn under a budget that holds only one, both preempted, no context error, the idle-return line logged. --- src/llama-context.cpp | 48 ++++++++++++++++--- tools/server/server-context.cpp | 63 +++++++++++++++++++++++-- tools/server/tests/unit/test_preempt.py | 37 +++++++++++++++ 3 files changed, 137 insertions(+), 11 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 34d2f51399a..20e548220db 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3153,6 +3153,8 @@ struct llama_state_seq_copy { struct dev_copy { ggml_backend_ptr backend; ggml_backend_event_t event = nullptr; + // recorded on the compute stream and waited for on the copy stream, see order_after() + ggml_backend_event_t fence = nullptr; bool pending = false; }; @@ -3178,6 +3180,9 @@ struct llama_state_seq_copy { if (it.second.event) { ggml_backend_event_free(it.second.event); } + if (it.second.fence) { + ggml_backend_event_free(it.second.fence); + } } } @@ -3221,6 +3226,25 @@ struct llama_state_seq_copy { } } + // Order the copies about to be posted behind the compute already queued on each device: + // the decode that produced the cells a park reads, or that a restore's cells were + // carved out of, has to be finished before the copy touches them. Recorded on the + // compute backend's stream and waited for on the copy stream, so the host drains + // nothing. Draining it (synchronize()) is what this replaces: with the wait that + // order_before() queues on the compute stream for the previous restore, a host drain + // blocked this thread until that copy had landed, and two restores issued in one pass + // ran one after the other with the whole transfer back on the decode loop. + void order_after(const std::vector & compute) { + for (auto & it : devs) { + for (const auto & backend : compute) { + if (ggml_backend_get_device(backend.get()) == it.first) { + ggml_backend_event_record(it.second.fence, backend.get()); + ggml_backend_event_wait(it.second.backend.get(), it.second.fence); + } + } + } + } + // Order the context's compute behind the copies just recorded, on the device: every // backend the graphs run on waits for the event of the transfer on its device before // the next graph it is given. This is a stream wait, not a host wait, so the caller's @@ -3635,10 +3659,19 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } + ggml_backend_event_t fence = ggml_backend_event_new(dev); + + if (!fence) { + ggml_backend_event_free(event); + ggml_backend_free(backend_cpy); + continue; + } + auto & dc = cpy->devs[dev]; dc.backend.reset(backend_cpy); dc.event = event; + dc.fence = fence; } if (cpy->devs.empty()) { @@ -3700,13 +3733,11 @@ size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size return 0; } - // The copies run on their own stream and are ordered against nothing, so the decode that - // produced these cells has to be finished before they are read. This is the one part of - // the transfer that stays on the caller's thread, and it costs nothing where it is used: - // a caller preempting a sequence does it between two decodes, with the previous one - // already drained by the sampling that followed it. + // The copies run on their own stream, so the decode that produced these cells has to be + // finished before they are read: the copy stream waits for the compute stream, on the + // device, see order_after(). Nothing stays on the caller's thread. const int64_t t_sync = ggml_time_us(); - synchronize(); + cpy.order_after(backends); cpy.t_sync_us = ggml_time_us() - t_sync; cpy.n_copies = 0; @@ -3748,8 +3779,11 @@ size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size return 0; } + // the cells this restore was given may still be read by a graph in flight (masked, but + // read), so the copy stream waits for the compute stream before it writes them: on the + // device, see order_after(), rather than by draining the compute stream on this thread const int64_t t_sync = ggml_time_us(); - synchronize(); + cpy.order_after(backends); cpy.t_sync_us = ggml_time_us() - t_sync; cpy.n_copies = 0; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 2d79a259684..3745762a725 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3337,7 +3337,51 @@ struct server_context_impl { } // whether parking this slot stays under --preempt-ram - bool preempt_fits_budget(const server_slot & slot) const { + // [TAG_PREEMPT_ASYNC] a restored slot keeps its pinned buffer for its next park, and + // that capacity counts against the budget while it holds no state. When a park does + // not fit, that idle capacity is what to give back first: largest first, never a + // buffer that still holds a parked sequence or has a copy in flight, and never the + // candidate's own, which it reuses. Without this a budget that holds one sequence was + // spent for good by the first restore: every later park was refused, and once the + // slot holding the buffer was the leader nothing could be parked at all. + void preempt_reclaim_idle_ram(size_t budget, size_t extra, const server_slot & keep) { + for (;;) { + if (preempt_ram_used() + extra <= budget) { + return; + } + + server_slot * best = nullptr; + + for (auto & other : slots) { + if (&other == &keep) { + continue; + } + + if (other.state == SLOT_STATE_PREEMPTED || other.state == SLOT_STATE_PREEMPTING || other.state == SLOT_STATE_RESTORING) { + continue; + } + + if (other.preempt_state_size() == 0) { + continue; + } + + if (!best || other.preempt_state_size() > best->preempt_state_size()) { + best = &other; + } + } + + if (!best) { + return; + } + + SLT_INF(*best, "%.1f MiB of idle parked RAM returned so that another slot can park\n", + best->preempt_state_size() / (1024.0 * 1024.0)); + + best->preempt_state_free(); + } + } + + bool preempt_fits_budget(const server_slot & slot) { if (params_base.preempt_ram_mib < 0) { return true; } @@ -3350,6 +3394,8 @@ struct server_context_impl { const size_t need = slot.preempt_state_required(); const size_t extra = need > held ? need - held : 0; + preempt_reclaim_idle_ram(budget, extra, slot); + return preempt_ram_used() + extra <= budget; } @@ -3360,18 +3406,27 @@ struct server_context_impl { // [TAG_PREEMPT_ASYNC] an asynchronous head keeps its pinned buffer through the restore // (see preempt_state_size), so nothing of it leaves; what the resident already holds is // reused, as in preempt_fits_budget, and only the rest is charged. - bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) const { + bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) { if (params_base.preempt_ram_mib < 0) { return true; } const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; - const size_t used = preempt_ram_used(); - const size_t leaving = head.preempt_is_async() ? 0 : std::min(used, head.preempt_state_size()); const size_t held = slot.preempt_state_size(); const size_t need = slot.preempt_state_required(); const size_t extra = need > held ? need - held : 0; + // the head is parked, so it is never among the idle buffers given back here + { + const size_t used = preempt_ram_used(); + const size_t leaving = head.preempt_is_async() ? 0 : std::min(used, head.preempt_state_size()); + + preempt_reclaim_idle_ram(budget + leaving, extra, slot); + } + + const size_t used = preempt_ram_used(); + const size_t leaving = head.preempt_is_async() ? 0 : std::min(used, head.preempt_state_size()); + return used - leaving + extra <= budget; } diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 26c6e7579ac..de9da44c4ed 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -719,3 +719,40 @@ def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_ after = _complete(8) assert after.status_code == 200 assert after.body["timings"]["predicted_n"] == 8 + + +def test_a_restored_slot_gives_its_idle_buffer_back_when_another_slot_needs_to_park(): + # Under a finite --preempt-ram an asynchronous slot keeps its pinned buffer after a + # restore, for its next park, and that idle capacity counted against the budget. With + # a budget that holds one sequence, the first restore spent it for good: every later + # park of the other slot was refused. The idle buffer is given back when another slot + # needs the room, and both slots go on being parked. + global server + # a pool of 8192 cells, but the model's own window is 2048, so each generation stays + # under that; 1800 tokens of this model's state is about 1.1 MiB, so a budget of + # 2 MiB holds one sequence and not two + server.n_ctx = 8192 + server.n_gpu_layer = 99 + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "256" + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" + text = _start_async() + _require_async(text) + log = LogReader(server.log_path) + + n_predict = 1800 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "idle parked RAM returned" in text, "the idle buffer of a restored slot was never given back" + import re + parked = re.findall(r"id\s+(\d+) \| task \d+ \| preempted on request", text) + assert {"0", "1"} <= set(parked), f"only slots {sorted(set(parked))} were ever parked" + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False From 012ef75479ea0fd9d9c9e6032cd32794a53b5767 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 01:07:14 +0000 Subject: [PATCH 65/81] llama: the copies wait for a fence the context records after every decode, so restores in one pass do not wait for each other order_after() recorded its fence on the compute stream at the time of the copy, which put it behind the waits order_before() had queued for the restores issued earlier in the same pass: restore B's copies then waited for restore A's to land, and restores issued together ran one after the other on the device, though the host no longer blocked. The fence is now one event per device owned by the context, recorded on the compute stream at the end of every decode and encode once a transfer exists (and once when the first transfer is created), and every park or restore waits for that point instead of recording its own. --- src/llama-context.cpp | 74 ++++++++++++++++++++++++++++--------------- src/llama-context.h | 8 +++++ 2 files changed, 57 insertions(+), 25 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 20e548220db..21776af9afd 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -483,6 +483,10 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); + for (auto & it : state_copy_fences) { + ggml_backend_event_free(it.second); + } + if (!model.hparams.no_alloc) { for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; @@ -1578,6 +1582,10 @@ int llama_context::encode(const llama_batch & batch_inp) { } } + if (!state_copy_fences.empty()) { + state_seq_copy_fence(); + } + return 0; } @@ -2023,6 +2031,10 @@ int llama_context::decode(const llama_batch & batch_inp) { // wait for the computation to finish (automatically done when obtaining the model output) //synchronize(); + if (!state_copy_fences.empty()) { + state_seq_copy_fence(); + } + return 0; } @@ -3153,8 +3165,6 @@ struct llama_state_seq_copy { struct dev_copy { ggml_backend_ptr backend; ggml_backend_event_t event = nullptr; - // recorded on the compute stream and waited for on the copy stream, see order_after() - ggml_backend_event_t fence = nullptr; bool pending = false; }; @@ -3180,9 +3190,6 @@ struct llama_state_seq_copy { if (it.second.event) { ggml_backend_event_free(it.second.event); } - if (it.second.fence) { - ggml_backend_event_free(it.second.fence); - } } } @@ -3228,19 +3235,20 @@ struct llama_state_seq_copy { // Order the copies about to be posted behind the compute already queued on each device: // the decode that produced the cells a park reads, or that a restore's cells were - // carved out of, has to be finished before the copy touches them. Recorded on the - // compute backend's stream and waited for on the copy stream, so the host drains - // nothing. Draining it (synchronize()) is what this replaces: with the wait that - // order_before() queues on the compute stream for the previous restore, a host drain - // blocked this thread until that copy had landed, and two restores issued in one pass - // ran one after the other with the whole transfer back on the decode loop. - void order_after(const std::vector & compute) { + // carved out of, has to be finished before the copy touches them. The copy stream waits + // for the context's fence on its device, an event the context records on the compute + // stream at the end of every decode, so the host drains nothing. The fence is recorded + // there and not here: recorded here, it would land behind the waits that order_before() + // queued for the restores issued earlier in the same pass, and each restore would then + // wait for the previous one's copies. Draining the host (synchronize()) is what this + // replaced: with those same waits on the compute stream, a host drain blocked this thread + // until the previous restore had landed. + void order_after(const std::map & fences) { for (auto & it : devs) { - for (const auto & backend : compute) { - if (ggml_backend_get_device(backend.get()) == it.first) { - ggml_backend_event_record(it.second.fence, backend.get()); - ggml_backend_event_wait(it.second.backend.get(), it.second.fence); - } + const auto fence = fences.find(it.first); + + if (fence != fences.end()) { + ggml_backend_event_wait(it.second.backend.get(), fence->second); } } } @@ -3606,6 +3614,16 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr // [TAG_STATE_ASYNC] +void llama_context::state_seq_copy_fence() { + for (const auto & it : state_copy_fences) { + for (const auto & backend : backends) { + if (ggml_backend_get_device(backend.get()) == it.first) { + ggml_backend_event_record(it.second, backend.get()); + } + } + } +} + llama_state_seq_copy * llama_context::state_seq_copy_init() { std::unique_ptr cpy(new llama_state_seq_copy()); @@ -3659,25 +3677,31 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } - ggml_backend_event_t fence = ggml_backend_event_new(dev); + if (state_copy_fences.find(dev) == state_copy_fences.end()) { + ggml_backend_event_t fence = ggml_backend_event_new(dev); - if (!fence) { - ggml_backend_event_free(event); - ggml_backend_free(backend_cpy); - continue; + if (!fence) { + ggml_backend_event_free(event); + ggml_backend_free(backend_cpy); + continue; + } + + state_copy_fences[dev] = fence; } auto & dc = cpy->devs[dev]; dc.backend.reset(backend_cpy); dc.event = event; - dc.fence = fence; } if (cpy->devs.empty()) { return nullptr; } + // the fences say where the compute streams are now, before any transfer asks + state_seq_copy_fence(); + // The devices above are the ones the graphs run on, not necessarily the ones the state // lives on: with most layers left on the CPU the KV cache is host memory, and a tensor // there takes the synchronous branch of backend_for(). A transfer whose every copy would @@ -3737,7 +3761,7 @@ size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size // finished before they are read: the copy stream waits for the compute stream, on the // device, see order_after(). Nothing stays on the caller's thread. const int64_t t_sync = ggml_time_us(); - cpy.order_after(backends); + cpy.order_after(state_copy_fences); cpy.t_sync_us = ggml_time_us() - t_sync; cpy.n_copies = 0; @@ -3783,7 +3807,7 @@ size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size // read), so the copy stream waits for the compute stream before it writes them: on the // device, see order_after(), rather than by draining the compute stream on this thread const int64_t t_sync = ggml_time_us(); - cpy.order_after(backends); + cpy.order_after(state_copy_fences); cpy.t_sync_us = ggml_time_us() - t_sync; cpy.n_copies = 0; diff --git a/src/llama-context.h b/src/llama-context.h index 3d66f2a948a..efa69f33cdc 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -165,6 +165,10 @@ struct llama_context { size_t state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags); size_t state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id dest_seq_id, llama_state_seq_flags flags); + // [TAG_STATE_ASYNC] mark the point the compute streams have reached, for the copies to + // wait for; recorded after every decode and encode once a transfer exists + void state_seq_copy_fence(); + bool state_load_file( const char * filepath, llama_token * tokens_out, @@ -357,6 +361,10 @@ struct llama_context { ggml_backend_t backend_cpu = nullptr; std::vector backends; + // [TAG_STATE_ASYNC] one event per device that copies asynchronously, recorded on the + // compute stream at the end of every decode; see state_seq_copy_fence() + std::map state_copy_fences; + // training ggml_opt_context_t opt_ctx = nullptr; From 02a908bf2e107558986d827168f80fb0717eff8f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 01:44:37 +0000 Subject: [PATCH 66/81] server: a parked stream keeps a shorter ping interval the request asked for While parked the keepalive ran every two seconds whatever --sse-ping said, which lengthened the silence for a client that had asked for a ping every second, exactly while nothing else was coming. The parked interval is now the shorter of the two. --- tools/server/server-context.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d639ad40005..5c0381a2648 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -5363,9 +5363,12 @@ std::unique_ptr server_routes::handle_completions_impl( bool timeout = false; int64_t start_time = ggml_time_ms(); // [TAG_PREEMPT] a parked slot produces nothing for as long as the pool is - // full, so while parked the ping runs every 2 s regardless of --sse-ping and - // is named, so a client can tell "waiting for cells" from "slow". - const int64_t ping_ms = parked ? PREEMPT_KEEPALIVE_MS : (sse_ping_interval > 0 ? (int64_t) sse_ping_interval * 1000 : -1); + // full, so while parked the ping runs at least every 2 s whether or not + // --sse-ping asked for one, and is named, so a client can tell "waiting for + // cells" from "slow". A shorter interval the request asked for is kept: a + // client that wants a ping every second wants it most while nothing else comes. + const int64_t ping_cfg = sse_ping_interval > 0 ? (int64_t) sse_ping_interval * 1000 : -1; + const int64_t ping_ms = parked ? (ping_cfg > 0 ? std::min(ping_cfg, PREEMPT_KEEPALIVE_MS) : PREEMPT_KEEPALIVE_MS) : ping_cfg; auto result = rd.next([&timeout, &start_time, ping_ms, &effective_should_stop]() { if (effective_should_stop()) { return true; // should_stop condition met From 4b9174448466ec0b680a8611ce02078a5d892a53 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 01:45:28 +0000 Subject: [PATCH 67/81] server: a decode that fits goes ahead beside a park in flight; a restoring prompt counts in the page-boundary cap The planner waited for an outstanding park on the same predicate that had brought it into its loop, so a step short only of the asynchronous lookahead blocked on the copy although the step itself had room, which is the overlap the asynchronous path exists for. The wait is now for a step that does not fit; short only of the lookahead with a park in flight, the step goes ahead and the planner looks again once the copy has landed. Under page allocation the cap on prompt reservations kept one boundary per prompt slot, counting the resident ones; a slot restoring into the prompt phase joins the next prompt batch too and reserved its chunk, so it counts as well. --- tools/server/server-context.cpp | 38 ++++++++++++++++++++++++++------- 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 2fef99e71c0..32bb3d637e2 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3741,7 +3741,11 @@ struct server_context_impl { int32_t n_pmt = 0; for (const auto & slot : slots) { - if (slot.state == SLOT_STATE_STARTED || slot.state == SLOT_STATE_PROCESSING_PROMPT) { + // a slot restoring into the prompt phase joins the next prompt batch too, and + // reserves its chunk above, so it can cross a boundary of its own as well + const slot_state state = slot.state == SLOT_STATE_RESTORING ? slot.state_before_preempt : slot.state; + + if (state == SLOT_STATE_STARTED || state == SLOT_STATE_PROCESSING_PROMPT) { n_pmt++; } } @@ -3964,6 +3968,16 @@ struct server_context_impl { } } + bool preempt_copies_in_flight() const { + for (const auto & slot : slots) { + if (slot.state == SLOT_STATE_PREEMPTING) { + return true; + } + } + + return false; + } + // Wait for one outstanding park, the last thing tried before giving up on finding room. // It is what keeps a pool that fills faster than the copies drain no worse than the // synchronous path: the decode waits for the copy exactly as it used to. @@ -4295,13 +4309,21 @@ struct server_context_impl { continue; } - // [TAG_PREEMPT_ASYNC] Out of room for the step about to be built, rather than - // merely short of the lookahead the asynchronous path keeps. A park that has - // been issued but not landed is holding cells that are already spoken for, and - // waiting for it is both quicker and more useful than parking somebody else, - // whose cells would not come back this iteration either. - if (n_used + preempt_n_margin() > n_cells && preempt_wait_in_flight()) { - continue; + // [TAG_PREEMPT_ASYNC] A park that has been issued but not landed is holding + // cells that are already spoken for. Out of room for the step about to be + // built, waiting for it is both quicker and more useful than parking somebody + // else, whose cells would not come back this iteration either. Short only of + // the lookahead the asynchronous path keeps, the step itself fits: it goes + // ahead beside the copy, which is the overlap the path exists for, and the + // planner looks again once the copy has landed. + if (preempt_copies_in_flight()) { + if (n_used > n_cells) { + if (preempt_wait_in_flight()) { + continue; + } + } else { + break; + } } server_slot * victim = preempt_pick_victim(); From 7eff72650d2e90edb45812805f6d00879b14d50b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 01:48:16 +0000 Subject: [PATCH 68/81] llama: staging counts by what a buffer charges; fences installed after the layout check; transfers only where a park can happen The decision to stage a fragmented restore counted the calls the emitter makes, and on a buffer without 2-D copies one strided call expands into one synchronous transfer per row, so a regularly interleaved sequence that had been staged at 64 runs was no longer staged and paid for every row. ggml_backend_buffer_supports_2d() says which kind of buffer it is, and the count is by rows where a row is what a call costs. The per-device fences were installed before the check that refuses a transfer for a state not all in device memory, so a server that then fell back to synchronous copies recorded them after every decode for nobody. They are installed after the checks, and an install that fails is undone. The server made a transfer for every slot, fences included, even where no park can happen: one slot, no memory, a recurrent cache. The transfers and the banner are now gated on the same conditions as the planner. --- ggml/include/ggml-backend.h | 3 ++ ggml/src/ggml-backend.cpp | 4 +++ src/llama-context.cpp | 59 ++++++++++++++++++++++----------- tools/server/server-context.cpp | 16 +++++++-- 4 files changed, 59 insertions(+), 23 deletions(-) diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index d21bf40dd58..09ee64a6561 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -62,6 +62,9 @@ extern "C" { GGML_API size_t ggml_backend_buffer_get_alloc_size(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor); GGML_API void ggml_backend_buffer_clear (ggml_backend_buffer_t buffer, uint8_t value); GGML_API bool ggml_backend_buffer_is_host (ggml_backend_buffer_t buffer); + // whether the buffer copies a strided set of rows in one call (see ggml_backend_tensor_set_2d); + // without it the generic path issues one transfer per row + GGML_API bool ggml_backend_buffer_supports_2d (ggml_backend_buffer_t buffer); GGML_API void ggml_backend_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage); GGML_API enum ggml_backend_buffer_usage ggml_backend_buffer_get_usage (ggml_backend_buffer_t buffer); GGML_API ggml_backend_buffer_type_t ggml_backend_buffer_get_type (ggml_backend_buffer_t buffer); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index b13d9c811c4..1ac8ecad9f6 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -175,6 +175,10 @@ bool ggml_backend_buffer_is_host(ggml_backend_buffer_t buffer) { return ggml_backend_buft_is_host(ggml_backend_buffer_get_type(buffer)); } +bool ggml_backend_buffer_supports_2d(ggml_backend_buffer_t buffer) { + return buffer->iface.set_tensor_2d != NULL && buffer->iface.get_tensor_2d != NULL; +} + void ggml_backend_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) { GGML_ASSERT(buffer); buffer->usage = usage; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 21776af9afd..4c63f0bec2a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2769,14 +2769,19 @@ class llama_io_read_host : public llama_io_read_i { // matters is how many runs of adjacent cells they form, because that is how many // transfers they actually cost. Count the runs first, and only fall back to // staging the whole tensor when even the runs are too many. + const size_t tensor_bytes = ggml_nbytes(tensor); + auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + + // A strided set of rows is one transfer on a buffer that copies 2-D, and one + // per row on one that does not (the generic path expands it), so it is counted + // by what it costs on this buffer, not by the calls it makes. + const bool has_2d = ggml_backend_buffer_supports_2d(buffer); + size_t n_runs = 0; llama_io_emit(rinfos, i, end, - [&n_runs](ggml_tensor *, const uint8_t *, size_t, size_t, size_t, size_t, size_t) { - n_runs++; + [&n_runs, has_2d](ggml_tensor *, const uint8_t *, size_t, size_t, size_t n_copies, size_t, size_t) { + n_runs += has_2d ? 1 : n_copies; }); - - const size_t tensor_bytes = ggml_nbytes(tensor); - auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; // A fragmented sequence can require thousands of synchronous device // transfers per layer. For bounded tensors, stage the tensor once and // preserve every byte belonging to other sequences. Bound scratch RAM @@ -3677,18 +3682,6 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } - if (state_copy_fences.find(dev) == state_copy_fences.end()) { - ggml_backend_event_t fence = ggml_backend_event_new(dev); - - if (!fence) { - ggml_backend_event_free(event); - ggml_backend_free(backend_cpy); - continue; - } - - state_copy_fences[dev] = fence; - } - auto & dc = cpy->devs[dev]; dc.backend.reset(backend_cpy); @@ -3699,9 +3692,6 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { return nullptr; } - // the fences say where the compute streams are now, before any transfer asks - state_seq_copy_fence(); - // The devices above are the ones the graphs run on, not necessarily the ones the state // lives on: with most layers left on the CPU the KV cache is host memory, and a tensor // there takes the synchronous branch of backend_for(). A transfer whose every copy would @@ -3734,6 +3724,35 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { cpy->can_pin = cpy->host_buffer_type() != ggml_backend_cpu_buffer_type(); + // One fence per device, shared by every transfer on this context and recorded after + // every decode from now on. Installed only here, after the checks above: a transfer + // refused for its layout must leave nothing behind that every later decode would keep + // recording for nobody. + std::vector fences_new; + + for (const auto & it : cpy->devs) { + if (state_copy_fences.find(it.first) != state_copy_fences.end()) { + continue; + } + + ggml_backend_event_t fence = ggml_backend_event_new(it.first); + + if (!fence) { + for (auto dev : fences_new) { + ggml_backend_event_free(state_copy_fences[dev]); + state_copy_fences.erase(dev); + } + + return nullptr; + } + + state_copy_fences[it.first] = fence; + fences_new.push_back(it.first); + } + + // the fences say where the compute streams are now, before any transfer asks + state_seq_copy_fence(); + return cpy.release(); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index dcba6d0408e..5991a2ba3a3 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1698,8 +1698,11 @@ struct server_context_impl { }; // [TAG_PREEMPT_ASYNC] one transfer per context, made once and reused for every - // park and resume this slot ever does, because each owns a backend and a stream - if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { + // park and resume this slot ever does, because each owns a backend and a stream. + // Only where a park can happen at all (see update_preemption): a transfer also + // installs the fences the context records after every decode, which a server + // that will never park has no use for. + if (preempt_async_possible()) { slot.preempt_cpy_tgt = llama_state_seq_copy_make(ctx_tgt); if (slot.preempt_cpy_tgt && ctx_dft) { @@ -1742,7 +1745,7 @@ struct server_context_impl { preempt_async_ok = preempt_async_ok && slot.preempt_is_async(); } - if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { + if (preempt_async_possible()) { if (preempt_async_ok) { // Pinned host memory is what lets a copy run beside the decode: one into or // out of pageable memory is staged by the driver and blocks the thread that @@ -5069,6 +5072,13 @@ struct server_context_impl { // update_slots() rebuilds the batch from the survivors. The planner brings the parked // ones back as cells free up. A multimodal prompt has no boundary the cache can name, // so it keeps the old path. + // [TAG_PREEMPT_ASYNC] whether a park can happen and go asynchronously: the conditions + // update_preemption() gates on, and the asynchronous switch + bool preempt_async_possible() const { + return params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0 && + slots.size() >= 2 && llama_get_memory(ctx_tgt) && !llama_model_is_recurrent(model_tgt); + } + bool preempt_last_resort_possible() const { return params_base.kv_unified && params_base.preempt_ram_mib != 0 && !preempt_recurrent && slots.size() >= 2 && llama_get_memory(ctx_tgt); } From 3475eb0708382b60ecaac1a13268b6e5b98ef02a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 02:30:12 +0000 Subject: [PATCH 69/81] llama: the shift wait runs before the draft is asked for; the last transfer takes the fences with it The wait for copies in flight before a context shift is applied ran just before the target decode, but pre_decode() had already asked the draft context for its draft, a decode that applies that cache's pending shift in place while a park or restore may still be copying draft cells. The wait now follows update_preemption() and precedes pre_decode(), so both caches shift after the copies have landed. When the pinning probe found pageable memory the server gave its transfers up, but the fences a transfer installs stayed with the context and were recorded after every decode for nobody. The context counts its live transfers now and the last one to go frees the fences. --- src/llama-context.cpp | 22 ++++++++++++++++++++++ src/llama-context.h | 7 +++++++ tools/server/server-context.cpp | 8 ++++++-- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4c63f0bec2a..4e6132bb0a8 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -483,6 +483,7 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); + // transfers outlive nothing: the server frees its slots before the contexts for (auto & it : state_copy_fences) { ggml_backend_event_free(it.second); } @@ -3177,6 +3178,8 @@ struct llama_state_seq_copy { ggml_backend_buffer_ptr host_buf; + bool counted = false; // held in the context's count of live transfers + uint8_t * data = nullptr; size_t size = 0; // bytes the current transfer covers size_t capacity = 0; // bytes actually held, kept across transfers @@ -3189,6 +3192,10 @@ struct llama_state_seq_copy { int64_t t_sync_us = 0; ~llama_state_seq_copy() { + if (counted) { + ctx->state_seq_copy_release(); + } + wait(); for (auto & it : devs) { @@ -3619,6 +3626,18 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr // [TAG_STATE_ASYNC] +void llama_context::state_seq_copy_release() { + GGML_ASSERT(state_copy_live > 0); + + if (--state_copy_live == 0) { + for (auto & it : state_copy_fences) { + ggml_backend_event_free(it.second); + } + + state_copy_fences.clear(); + } +} + void llama_context::state_seq_copy_fence() { for (const auto & it : state_copy_fences) { for (const auto & backend : backends) { @@ -3753,6 +3772,9 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { // the fences say where the compute streams are now, before any transfer asks state_seq_copy_fence(); + state_copy_live++; + cpy->counted = true; + return cpy.release(); } diff --git a/src/llama-context.h b/src/llama-context.h index efa69f33cdc..db7ffc6b4c0 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -169,6 +169,9 @@ struct llama_context { // wait for; recorded after every decode and encode once a transfer exists void state_seq_copy_fence(); + // a transfer letting go of this context: the last one takes the fences with it + void state_seq_copy_release(); + bool state_load_file( const char * filepath, llama_token * tokens_out, @@ -365,6 +368,10 @@ struct llama_context { // compute stream at the end of every decode; see state_seq_copy_fence() std::map state_copy_fences; + // transfers alive on this context; the fences go when the last one does, so a server + // that made transfers and then gave them up records nothing after its decodes + int32_t state_copy_live = 0; + // training ggml_opt_context_t opt_ctx = nullptr; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 418d17580db..c3a8195bf0a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4244,6 +4244,12 @@ struct server_context_impl { pre_decode_shift(); update_preemption(); + // [TAG_PREEMPT_ASYNC] before pre_decode(), not only before the target decode: the + // draft it asks for is a decode on the draft context, which applies that cache's + // pending shift in place, and a park or restore still copying draft cells would + // read through it or be overwritten by it just the same + preempt_wait_for_shift(); + scoped_timer t(t_pre_decode, n_pre_decode); pre_decode(); batch.render(); @@ -4280,8 +4286,6 @@ struct server_context_impl { int32_t off_next = 0; int32_t n_batch = llama_n_batch(ctx_tgt); - preempt_wait_for_shift(); - for (int32_t off = 0; off < batch.size(); off = off_next) { const int32_t n_tokens = std::min(n_batch, batch.size() - off); try { From ebfa47b1678a2525840f82ad249f13cad19eda13 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 02:59:15 +0000 Subject: [PATCH 70/81] llama: a context freed with live transfers drains and disowns them; the shift wait runs before the decode as well synchronize() covers the graph backends, not the copy backend a transfer owns, so a context freed while a transfer was still copying could free the KV buffers under it, and freeing the transfer afterwards touched the dead context. The context keeps the set of its live transfers now: at teardown each is waited for and disowned, and its own free then touches nothing of the context. The wait for copies in flight before a shift is applied runs before pre_decode(), for the draft, and again before the decode: a shift that --cache-reuse asks for is found inside pre_decode(), after the first wait, and the decode applies it in place like any other. --- src/llama-context.cpp | 27 +++++++++++++++++++++------ src/llama-context.h | 11 ++++++++--- tools/server/server-context.cpp | 5 +++++ 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4e6132bb0a8..a5e9f269233 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -483,7 +483,12 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); - // transfers outlive nothing: the server frees its slots before the contexts + // A transfer still alive is drained first: synchronize() covers the graph backends, + // not the copy backend a transfer owns, and the KV buffers it may still be reading or + // writing are about to go. It is then let go of, so freeing it later touches nothing + // of this context. + state_seq_copies_drain(); + for (auto & it : state_copy_fences) { ggml_backend_event_free(it.second); } @@ -3193,7 +3198,7 @@ struct llama_state_seq_copy { ~llama_state_seq_copy() { if (counted) { - ctx->state_seq_copy_release(); + ctx->state_seq_copy_release(this); } wait(); @@ -3626,10 +3631,20 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr // [TAG_STATE_ASYNC] -void llama_context::state_seq_copy_release() { - GGML_ASSERT(state_copy_live > 0); +void llama_context::state_seq_copies_drain() { + for (auto * cpy : state_copies) { + cpy->wait(); + cpy->ctx = nullptr; + cpy->counted = false; + } + + state_copies.clear(); +} + +void llama_context::state_seq_copy_release(llama_state_seq_copy * cpy) { + GGML_ASSERT(state_copies.erase(cpy) == 1); - if (--state_copy_live == 0) { + if (state_copies.empty()) { for (auto & it : state_copy_fences) { ggml_backend_event_free(it.second); } @@ -3772,7 +3787,7 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { // the fences say where the compute streams are now, before any transfer asks state_seq_copy_fence(); - state_copy_live++; + state_copies.insert(cpy.get()); cpy->counted = true; return cpy.release(); diff --git a/src/llama-context.h b/src/llama-context.h index db7ffc6b4c0..106dc3922e6 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -12,6 +12,7 @@ #include "ggml-opt.h" #include +#include #include struct llama_model; @@ -170,7 +171,10 @@ struct llama_context { void state_seq_copy_fence(); // a transfer letting go of this context: the last one takes the fences with it - void state_seq_copy_release(); + void state_seq_copy_release(llama_state_seq_copy * cpy); + + // at teardown: wait for every live transfer and let it go + void state_seq_copies_drain(); bool state_load_file( const char * filepath, @@ -369,8 +373,9 @@ struct llama_context { std::map state_copy_fences; // transfers alive on this context; the fences go when the last one does, so a server - // that made transfers and then gave them up records nothing after its decodes - int32_t state_copy_live = 0; + // that made transfers and then gave them up records nothing after its decodes, and a + // context freed with transfers still alive drains them and lets them go first + std::set state_copies; // training ggml_opt_context_t opt_ctx = nullptr; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 6b24b528b2a..99751e77e46 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4290,6 +4290,11 @@ struct server_context_impl { int32_t off_next = 0; int32_t n_batch = llama_n_batch(ctx_tgt); + // [TAG_PREEMPT_ASYNC] and once more here: a shift --cache-reuse asks for is found + // inside pre_decode(), after the wait above, and the decode below applies it in + // place like any other + preempt_wait_for_shift(); + for (int32_t off = 0; off < batch.size(); off = off_next) { const int32_t n_tokens = std::min(n_batch, batch.size() - off); try { From b4b0f9bd71f1b411df4a3a00d373140272be9d91 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 04:06:23 +0000 Subject: [PATCH 71/81] exact concurrency: shorter comments Comment-only pass over the PR's diff: collapse the long explanations to one or two lines each and drop the ones the code already says. --- .github/workflows/unsloth-pin-preflight.yml | 57 ++-- .github/workflows/unsloth-pr-set-lint.yml | 6 +- .github/workflows/unsloth-prebuilt.yml | 23 +- common/common.cpp | 27 +- common/common.h | 12 +- ggml/include/ggml-cuda.h | 12 +- ggml/src/ggml-cann/ggml-cann.cpp | 3 +- ggml/src/ggml-cpu/ggml-cpu.cpp | 11 +- ggml/src/ggml-cuda/common.cuh | 4 +- ggml/src/ggml-cuda/fattn-common.cuh | 9 +- ggml/src/ggml-cuda/fattn-vec.cuh | 4 +- ggml/src/ggml-cuda/fattn.cu | 10 +- ggml/src/ggml-cuda/ggml-cuda.cu | 111 +++---- ggml/src/ggml-cuda/mmvq.cu | 16 +- ggml/src/ggml-cuda/mmvq.cuh | 5 +- ggml/src/ggml-et/ggml-et.cpp | 3 +- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 3 +- ggml/src/ggml-metal/ggml-metal-device.m | 4 +- ggml/src/ggml-opencl/ggml-opencl.cpp | 3 +- ggml/src/ggml-openvino/ggml-openvino.cpp | 3 +- ggml/src/ggml-rpc/ggml-rpc.cpp | 4 +- ggml/src/ggml-sycl/ggml-sycl.cpp | 3 +- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 3 +- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 3 +- include/llama.h | 36 +-- scripts/batchinv/divergence.py | 12 +- scripts/batchinv/probe.cpp | 30 +- scripts/batchinv/prompts.py | 2 +- scripts/unsloth/additive_merge.py | 43 +-- scripts/unsloth/feature_matrix.py | 14 +- scripts/unsloth/pin_contract.py | 36 +-- scripts/unsloth/test_additive_merge.py | 6 +- scripts/unsloth/test_pin_contract.py | 11 +- src/llama-batch.cpp | 23 +- src/llama-batch.h | 9 +- src/llama-context.cpp | 44 ++- src/llama-graph.cpp | 17 +- src/llama-impl.cpp | 40 +-- src/llama-impl.h | 14 +- src/llama-kv-cache.cpp | 110 +++---- src/llama-kv-cache.h | 9 +- src/llama-memory-hybrid.cpp | 19 +- src/llama-memory-recurrent.cpp | 10 +- src/llama-memory.h | 11 +- tests/test-backend-ops.cpp | 13 +- tests/test-state-restore-fragmented.cpp | 4 +- tools/server/server-context.cpp | 333 +++++++------------- tools/server/tests/unit/test_preempt.py | 111 +++---- 48 files changed, 497 insertions(+), 799 deletions(-) diff --git a/.github/workflows/unsloth-pin-preflight.yml b/.github/workflows/unsloth-pin-preflight.yml index f390f381e6f..494bca228c8 100644 --- a/.github/workflows/unsloth-pin-preflight.yml +++ b/.github/workflows/unsloth-pin-preflight.yml @@ -28,16 +28,10 @@ permissions: contents: write issues: write -# Two runs of the same ref probe the same pins against the same base, so the -# second adds nothing and just competes for runners. On 08-04 a dispatch and the -# schedule sat queued together for an hour. Newest wins: it sees the newest -# pr-set.json. -# -# Per ref, though, not globally. This file also runs on any push that touches -# pr-set.json, so with one shared group a push to a second branch cancelled the -# first branch's run: observed on 09-03, where the run that would have said -# whether a repin fixed the nightly was cancelled by an unrelated branch, and -# the PR was left showing the failure from before the fix. +# Two runs of the same ref probe the same pins against the same base, so the second only +# competes for runners; newest wins, since it sees the newest pr-set.json. Per ref, not +# globally: this also runs on any push touching pr-set.json, and with one shared group a push +# to a second branch cancelled the first branch's run (09-03). concurrency: group: unsloth-pin-preflight-${{ github.ref }} cancel-in-progress: true @@ -56,10 +50,8 @@ jobs: id: p run: | set -uo pipefail - # Everything below reports through `status`/`details`, so a death - # anywhere else leaves both empty and the alert blank: a red X on a - # scheduled run nobody opens. Report the abort through the same - # channel as a finding, so the repin bot sees a failure either way. + # everything below reports through `status`/`details`, so a death anywhere else + # leaves the alert blank; report the abort through the same channel as a finding trap 'rc=$?; if [ "$rc" != 0 ]; then { echo "status=failure" echo "details</dev/null PROBLEMS="${PROBLEMS}- \`${SRC}#${NUM}\` (\`${SHA:0:10}\`) does not merge onto \`${BASE}\` + the pins before it.\n\n Conflicting files:\n\n\`\`\`\n${FILES}\n\`\`\`\n\n
conflict hunks\n\n\`\`\`diff\n${HUNKS}\n\`\`\`\n\n
\n" @@ -193,11 +180,9 @@ jobs: PROBLEMS="${PROBLEMS}- the merged tree builds, but \`scripts/unsloth/merge_checks.py\` found a resolution that is silently wrong. See the run log for file and line.\n" fi - # The other half of that question. merge_checks.py asks whether the - # tree contains something wrong; this asks whether it still contains - # what each pin carries. A pin that has rotted into a no-op, or an - # arch registration a resolution quietly dropped, is invisible to - # every other check here and to the compiler. + # the other half: merge_checks.py asks whether the tree contains something + # wrong, this asks whether it still contains what each pin carries. A pin rotted + # into a no-op is invisible to every other check here and to the compiler. if ! python3 ../scripts/unsloth/pin_contract.py --root . --base "$BASE" \ --pr-set ../scripts/unsloth/pr-set.json --report "${RUNNER_TEMP}/pin_contract.json" ; then PROBLEMS="${PROBLEMS}- the merged tree is missing code a pin carries. See the run log for the pin and file.\n" @@ -207,12 +192,10 @@ jobs: PROBLEMS="${PROBLEMS}- pins upstream has taken over, safe to delete from \`pr-set.json\`:\n\n\`\`\`\n${NOTES}\n\`\`\`\n" fi - # A clean merge is not a compiling tree. On 09-03 ggml-org#27754 - # merged with no conflicts at all and did not compile: upstream had - # added a parameter to build_attn_mha and the pin's new - # build_attn_sparse still called the old signature. Nothing above - # can see that. CPU only and the `llama` target only, which is where - # that translation unit lives; 59s cold at -j4 with no ccache. + # a clean merge is not a compiling tree: on 09-03 ggml-org#27754 merged with no + # conflicts and did not compile, upstream having added a parameter to + # build_attn_mha that the pin's build_attn_sparse still called without. CPU only, + # 59s cold at -j4 with no ccache. GATE_OK=1 if ! cmake -B "${RUNNER_TEMP}/gate" -DCMAKE_BUILD_TYPE=Release \ -DGGML_CUDA=OFF -DLLAMA_BUILD_TESTS=ON -DLLAMA_BUILD_SERVER=OFF \ @@ -223,10 +206,8 @@ jobs: PROBLEMS="${PROBLEMS}- the pins merge cleanly and the merged tree does not compile. See the run log for the file and line; this is the failure that only shows up in the CUDA leg once the nightly has fanned out.\n" fi - # The last question, and the only one that needs a binary: does each - # feature we ship still work. Everything above is about the source. - # CPU only, because no runner in this pipeline has a GPU -- see the - # note in feature_matrix.py about what that does and does not prove. + # the only question that needs a binary: does each feature we ship still work. + # CPU only, since no runner here has a GPU; see the note in feature_matrix.py. if [ -n "$GATE_OK" ]; then if ! python3 ../scripts/unsloth/feature_matrix.py \ --build-dir "${RUNNER_TEMP}/gate" \ diff --git a/.github/workflows/unsloth-pr-set-lint.yml b/.github/workflows/unsloth-pr-set-lint.yml index 8a898539f02..99b06b74f34 100644 --- a/.github/workflows/unsloth-pr-set-lint.yml +++ b/.github/workflows/unsloth-pr-set-lint.yml @@ -128,9 +128,9 @@ jobs: done exit "$fail" - # A pin nobody decided about is the failure this whole file exists to stop. - # Being in `unchecked` with a reason is a fine answer; being in neither map - # is how DiffusionGemma went five weeks with no coverage and no record of it. + # a pin nobody decided about is the failure this file exists to stop: being in + # `unchecked` with a reason is fine, being in neither map is how DiffusionGemma went + # five weeks with no coverage and no record of it - name: Every pin is either checked or knowingly unchecked run: | set -euo pipefail diff --git a/.github/workflows/unsloth-prebuilt.yml b/.github/workflows/unsloth-prebuilt.yml index 834ff544909..353d88417ca 100644 --- a/.github/workflows/unsloth-prebuilt.yml +++ b/.github/workflows/unsloth-prebuilt.yml @@ -282,8 +282,8 @@ jobs: # .github/workflows, which upstream history routinely does). if [ "$EXISTS" != "true" ] || [ "${{ github.event_name }}" = "workflow_dispatch" ]; then git remote add upstream https://github.com/ggml-org/llama.cpp.git - # The upstream checkout below takes scripts/unsloth/ away. Copy the - # whole dir out, not file by file: see the note above the step. + # the upstream checkout below takes scripts/unsloth/ away; copy the whole dir + # out, not file by file, see the note above the step cp -r scripts/unsloth "${RUNNER_TEMP}/us" ADDITIVE_MERGE="${RUNNER_TEMP}/us/additive_merge.py" if [ "$(jq length <<<"$PRS")" != 0 ]; then @@ -446,8 +446,7 @@ jobs: # A bad pin resolution can still build fine, so it must be caught before the source artifact ships. See merge_checks.py. # Its own step, not more script in `resolve`: GitHub caps one workflow string at 21000 chars and that step is near it. See check_workflow_scalars.py. - # That is also why `resolve` copies all of scripts/unsloth/ to ${RUNNER_TEMP}/us in one line rather than one cp per script: every check added - # here would otherwise cost another line inside the capped block, and going over silently disables the whole workflow. + # That is also why `resolve` copies all of scripts/unsloth/ to ${RUNNER_TEMP}/us in one line rather than one cp per script: another line inside the capped block per check would eventually go over, which silently disables the whole workflow. - name: Check the merged tree for silently wrong resolutions if: ${{ env.MERGED_PINS == '1' }} run: | @@ -457,13 +456,11 @@ jobs: exit 1 fi - # merge_checks.py asserts the ABSENCE of two known-bad shapes. This asserts the PRESENCE of what each pin carries, which is a different question and - # the one that goes unanswered when a pin rots into a no-op or a resolution quietly drops an arch registration. Free, so it runs before the compile gate. + # merge_checks.py asserts the ABSENCE of two known-bad shapes; this asserts the PRESENCE of what each pin carries, the question that goes unanswered when a pin rots into a no-op. Free, so it runs before the compile gate. - name: Check every pin still contributes what it carries if: ${{ env.MERGED_PINS == '1' }} - # Through env, never interpolated into the script: `prs` carries PR - # titles, which are third-party text, and `${{ }}` pastes them into the - # shell source before bash ever sees it. + # through env, never interpolated: `prs` carries PR titles, which are third-party + # text, and `${{ }}` pastes them into the shell source before bash sees it env: PRS: ${{ steps.r.outputs.prs }} BASE: ${{ steps.r.outputs.base }} @@ -475,12 +472,8 @@ jobs: exit 1 fi - # The gap this closes, observed 09-03: ggml-org#27754 merged with zero conflicts and did not compile, because upstream had added a parameter to - # build_attn_mha and the pin's new build_attn_sparse still called the old signature. Nothing before this point can see that, and without it the release - # dies in the CUDA leg after the 38-job fan-out. CPU only: a cold `llama` build took 59s at -j4 with no ccache, against 20-60 minutes for a CUDA build. - # mtmd is in the gate because `llama` alone is not enough: observed 09-04, ggml-org#25731 built `llama` clean while tools/mtmd did not compile at all, - # upstream having made mtmd_image_preprocessor::preprocess const while the pin's Inkling subclass stayed non-const, so it overrode nothing and the - # vision and audio towers were abstract. Every vision pin lands in mtmd, so a gate that skips it cannot see the whole class. + # The gap this closes, observed 09-03: ggml-org#27754 merged with zero conflicts and did not compile, upstream having added a parameter to build_attn_mha that the pin's build_attn_sparse still called without. Without this the release dies in the CUDA leg after the 38-job fan-out. CPU only: a cold `llama` build took 59s at -j4, against 20-60 minutes for CUDA. + # mtmd is in the gate because `llama` alone is not enough: on 09-04 ggml-org#25731 built `llama` clean while tools/mtmd did not compile at all, upstream having made mtmd_image_preprocessor::preprocess const while the pin's Inkling subclass stayed non-const. Every vision pin lands in mtmd. - name: Compile gate (CPU, llama and mtmd targets) if: ${{ env.MERGED_PINS == '1' }} run: | diff --git a/common/common.cpp b/common/common.cpp index 04ce8744359..b74bee734ae 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1290,10 +1290,9 @@ struct common_init_result::impl { common_init_result::common_init_result(common_params & params, bool model_only) : pimpl(new impl{}) { - // [TAG_EXACT_CONCURRENCY] before any context exists, the fitting ones included: the - // per-sequence figure and the column bound are checked against the explicit bound first, - // so a context is never created under a figure the bound does not cover. A caller that - // skipped common_params_parse() gets the same check here; on failure nothing is loaded. + // [TAG_EXACT_CONCURRENCY] before any context exists, so one is never created under a figure + // the explicit bound does not cover; this also covers a caller that skipped + // common_params_parse(). On failure nothing is loaded. if (!model_only && !common_exact_concurrency_init(params)) { COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: refusing to load the model, see the error above\n"); return; @@ -1457,12 +1456,11 @@ bool common_exact_concurrency() { int common_exact_decode_width(const common_params & params) { const int64_t n_slots = std::max(1, params.n_parallel); - // the draft tokens a slot carries into the verify ubatch alongside its accepted token, per - // speculation type, from the same place the speculation code takes its own width + // draft tokens a slot carries into the verify ubatch, from the same place the speculation + // code takes its own width const int64_t n_draft = std::max(0, (int) common_speculative_n_max(¶ms.speculative)); - // the product is what a backend is asked to split columns by, as an int; one that does not - // fit is reported as such rather than wrapped + // the product is handed to a backend as an int; one that overflows is reported, not wrapped const int64_t n_cols = n_slots*(1 + n_draft); return n_cols > INT32_MAX ? -1 : (int) n_cols; @@ -1474,9 +1472,8 @@ bool common_exact_concurrency_init(const common_params & params) { return true; } - // DFlash drafting turns causal attention off on its draft context, and the paged - // attention the mode runs on needs it; say so instead of asserting in the graph. DSpark - // is the same implementation under another name, so it is refused with it. + // DFlash drafting turns causal attention off on its draft context, which the paged attention + // needs; say so instead of asserting in the graph. DSpark is the same implementation. for (const auto type : params.speculative.types) { if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash or draft-dspark: both disable causal attention on the draft, which the paged attention needs\n"); @@ -1505,11 +1502,9 @@ bool common_exact_concurrency_init(const common_params & params) { } } - // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode - // step is; a context created later reports n_seq_max times that figure, which is n_cols - // again, and reporting n_cols here as well covers a caller that decodes before that. Both - // refuse a width the explicit bound above cannot cover, which the check above already - // caught for this process; contexts created earlier by the caller are covered here. + // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step + // is. A context created later reports n_seq_max times that, which is n_cols again; reporting + // n_cols here too covers a caller that decodes first, or contexts it created earlier. if (!llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel))) || !llama_set_exact_decode_width((uint32_t) n_cols)) { COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: the decode width could not be reported, see the error above\n"); diff --git a/common/common.h b/common/common.h index 2be2fab6a8b..4616b73153e 100644 --- a/common/common.h +++ b/common/common.h @@ -931,17 +931,15 @@ using common_init_result_ptr = std::unique_ptr; common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false); -// [TAG_EXACT_CONCURRENCY] -// true when LLAMA_EXACT_CONCURRENCY is set for this process +// [TAG_EXACT_CONCURRENCY] true when LLAMA_EXACT_CONCURRENCY is set for this process bool common_exact_concurrency(); -// the widest ubatch a decode step can build with these parameters: one column per slot, times one -// plus the number of speculative draft tokens carried with it. Under exact mode this is what the -// CUDA column policy has to cover, and what its default bound is derived from. +// the widest ubatch a decode step can build here: one column per slot times one plus its draft +// tokens. Under exact mode the CUDA column policy has to cover this, and derives its bound from it. int common_exact_decode_width(const common_params & params); -// report that width to the CUDA backend, and refuse an explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS -// that is smaller than it. Returns false if the configuration must not run. +// report that width to the CUDA backend, refusing a smaller explicit +// GGML_CUDA_BATCH_INVARIANT_MAX_COLS; false if the configuration must not run bool common_exact_concurrency_init(const common_params & params); struct llama_model_params common_model_params_to_llama ( common_params & params); diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h index c3dd87c97b7..07131df327c 100644 --- a/ggml/include/ggml-cuda.h +++ b/ggml/include/ggml-cuda.h @@ -39,13 +39,11 @@ GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * f GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); -// [TAG_EXACT_CONCURRENCY] -// Report the widest ubatch a decode step of this process can build: one column per slot, times one -// plus the number of speculative draft tokens carried with it. Under LLAMA_EXACT_CONCURRENCY the -// column policy then defaults to that width instead of a fixed number, so --parallel or a wider -// draft cannot silently push a decode above the bound and leave it batched. An explicitly set -// GGML_CUDA_BATCH_INVARIANT_MAX_COLS still wins. Call before the first graph is computed. Also -// available through ggml_backend_reg_get_proc_address(). +// [TAG_EXACT_CONCURRENCY] report the widest ubatch a decode step of this process can build: one +// column per slot times one plus its draft tokens. Under LLAMA_EXACT_CONCURRENCY the column policy +// defaults to that instead of a fixed number, so --parallel or a wider draft cannot silently push a +// decode above the bound. An explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS still wins. Call before +// the first graph is computed; also available through ggml_backend_reg_get_proc_address(). GGML_BACKEND_API void ggml_backend_cuda_set_exact_decode_width(int n_cols); GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index 0e901ed0160..9b9213df856 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2656,8 +2656,7 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten return true; case GGML_OP_FLASH_ATTN_EXT: { - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only - // the CUDA backend reads + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads if (op->src[5]) { return false; } diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index a548b33bd71..8bed84e44bb 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -474,13 +474,10 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return ggml_is_contiguous(op->src[0]); case GGML_OP_SSM_SCAN: return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1; - // [TAG_EXACT_CONCURRENCY] note: GGML_OP_FLASH_ATTN_EXT with src[5] set, the - // exact-concurrency page table, is deliberately still accepted here. The CPU ignores the - // page table and attends in physical cell order, which is why every other backend refuses - // it, but the CPU is also the reference that test-backend-ops compares the paged CUDA - // kernel against, and that test builds a mask which selects exactly the listed cells. A KV - // cache layer cannot reach the CPU under the mode anyway: llama_kv_cache refuses to - // construct unless every KV layer is on the CUDA backend. + // [TAG_EXACT_CONCURRENCY] note: FLASH_ATTN_EXT with src[5], the page table, is deliberately + // still accepted. The CPU ignores it and attends in physical order, but it is also the + // reference test-backend-ops compares the paged CUDA kernel against, and that test's mask + // selects exactly the listed cells. A KV layer cannot reach the CPU under the mode anyway. default: return true; } diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 728aa08dcb5..d749b42ee48 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -51,8 +51,8 @@ #define GGML_CUDA_CC_DP4A 610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products // [TAG_BATCH_INVARIANT] 0 = off, 1 = split every batched matmul, 2 = split only where it changes bits int ggml_cuda_batch_invariant(); -// Widest batch the split is applied to, 0 = no bound. Prompt-sized batches cost far more to -// split than decode-sized ones, and only prompt-phase invariance is given up by bounding it. +// widest batch the split applies to, 0 = no bound; bounding it gives up prompt-phase invariance +// only, and prompt-sized batches cost far more to split int ggml_cuda_batch_invariant_max_cols(); #define GGML_CUDA_CC_VOLTA 700 diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index f6aa5b03ec1..10f28eb229b 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1091,8 +1091,8 @@ void launch_fattn( // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped. // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or // multiple sequences of possibly different lengths. - // [TAG_BATCH_INVARIANT] Without this scan the KV loop runs to K->ne[1], which grows with the - // other sequences sharing the cache. Scanning the mask bounds it by the sequence's own extent. + // [TAG_BATCH_INVARIANT] without this scan the KV loop runs to K->ne[1], which grows with the + // other sequences sharing the cache; scanning the mask bounds it by the sequence's own extent const bool batch_invariant_KV_max = ggml_cuda_batch_invariant() != 0; if (!dst->src[5] && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) { const int64_t s31 = mask->nb[1] / sizeof(half2); @@ -1152,9 +1152,8 @@ void launch_fattn( dst_tmp_meta.alloc((size_t(blocks_num.x) * ncols * (2 + DV/2))); } } else if (dst->src[5] || ggml_cuda_batch_invariant()) { - // [TAG_BATCH_INVARIANT] How the KV cache is split between blocks, and therefore the order - // in which the partial attention results are combined, follows K->ne[1]. That length grows - // with the other sequences sharing the cache, so pin the split to a single block per tile. + // [TAG_BATCH_INVARIANT] the KV split between blocks, and so the order the partials combine + // in, follows K->ne[1], which grows with the other sequences: pin it to one block per tile parallel_blocks = 1; blocks_num.x = ntiles_x; diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index f402795942c..4005d28879e 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -247,8 +247,8 @@ static __global__ void flash_attn_ext_vec( #endif // V_DOT2_F32_F16_AVAILABLE } - // In the paged specialization KV_max carries [count, physical page IDs...] per query. - // The loop and each warp's recurrence follow logical positions, never physical addresses. + // in the paged specialization KV_max carries [count, physical page IDs...] per query; the loop + // and each warp's recurrence follow logical positions, never physical addresses static_assert(!paged || ncols == 1, "paged attention has one query per block"); const int * pages = paged ? KV_max + (sequence*int(ne01.z) + ic0)*(1 + ne11/FATTN_KQ_STRIDE) : nullptr; const int k_VKQ_max = paged ? pages[0]*FATTN_KQ_STRIDE : (KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11); diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index eff18212272..915c2d04da8 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -457,9 +457,8 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path. const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0; - // [TAG_BATCH_INVARIANT] Every choice below switches on Q->ne[1] or on K->ne[1], and both - // grow with the other sequences in the batch and in the shared KV cache. Pin the kernel a - // batch of one would use so a request is never moved onto a different algorithm by its neighbours. + // [TAG_BATCH_INVARIANT] every choice below switches on Q->ne[1] or K->ne[1], both of which grow + // with the other sequences, so pin the kernel a batch of one would use if (ggml_cuda_batch_invariant() && can_use_vector_kernel && Q->ne[1] == 1) { return BEST_FATTN_KERNEL_VEC; } @@ -592,7 +591,7 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst return; } - // [TAG_BATCH_INVARIANT] Attend one query row at a time, as a batch of one would. + // [TAG_BATCH_INVARIANT] attend one query row at a time, as a batch of one would const int fattn_max_cols = ggml_cuda_batch_invariant_max_cols(); if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1 && (fattn_max_cols <= 0 || dst->src[0]->ne[1] <= fattn_max_cols)) { @@ -606,8 +605,7 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst ggml_tensor mask_row; ggml_tensor dst_row = *dst; - // ne[2] keeps running to the end of dst so that the scratch space for F16 copies of - // K and V, which is placed right behind dst, is still put in the same place. + // ne[2] runs to the end of dst so the F16 K/V scratch behind dst stays in place dst_row.ne[2] = dst->ne[2] - i; dst_row.data = (char *) dst->data + i*dst->nb[2]; dst_row.src[0] = &Q_row; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 8c5b2a408c3..76fa4dc66c6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1758,8 +1758,8 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, } static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { - // [TAG_BATCH_INVARIANT] mul_mat+GLU is only fused for a single destination column, so - // leaving it on would give a solo request a different code path from a batched one. + // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving + // it on would give a solo request a different code path from a batched one if (ggml_cuda_batch_invariant()) { return false; } @@ -1791,8 +1791,8 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { } static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { - // [TAG_BATCH_INVARIANT] mul_mat+GLU is only fused for a single destination column, so - // leaving it on would give a solo request a different code path from a batched one. + // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving + // it on would give a solo request a different code path from a batched one if (ggml_cuda_batch_invariant()) { return false; } @@ -1825,16 +1825,11 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } -// [TAG_BATCH_INVARIANT] -// The number of tokens in a batch picks both the matmul implementation below and, inside -// several of them, how the K loop is divided between threads. Both change the order in -// which the partial products of one destination element are summed, so the same request -// produces different bits depending on how many other requests decode alongside it. -// -// GGML_CUDA_BATCH_INVARIANT removes that dependency: -// 1 - compute every destination column on its own, exactly as a batch of one would. -// 2 - split off only the columns whose batch-of-one configuration differs from the -// batched one, leaving the already invariant matmuls batched. +// [TAG_BATCH_INVARIANT] the token count picks the matmul implementation and how its K loop is +// divided between threads, both of which change the summation order, so the same request produces +// different bits depending on how many others decode alongside it. GGML_CUDA_BATCH_INVARIANT: +// 1 - compute every destination column on its own, exactly as a batch of one would +// 2 - split off only the columns whose batch-of-one configuration differs from the batched one static bool ggml_cuda_exact_concurrency() { static const bool exact = []() { const char * value = getenv("LLAMA_EXACT_CONCURRENCY"); @@ -1864,9 +1859,8 @@ void ggml_backend_cuda_set_exact_decode_width(int n_cols) { } int ggml_cuda_batch_invariant_max_cols() { - // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill - // matmul shapes match its solo run, so exact mode no longer needs the column policy to be - // unbounded there. An explicit bound always wins, in either mode. + // [TAG_EXACT_CONCURRENCY] prompt ubatches hold one sequence, so a prefill already matches its + // solo run and needs no unbounded column policy. An explicit bound always wins. static const int explicit_cols = []() { const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); return val ? atoi(val) : -1; @@ -1880,24 +1874,19 @@ int ggml_cuda_batch_invariant_max_cols() { return 0; } - // Exact mode only has to cover the widest ubatch a decode step can build: one column per slot, - // times one plus the number of speculative draft tokens carried with it. Use that width when - // the caller reported it through ggml_backend_cuda_set_exact_decode_width(). Nothing reported - // it, so fall back to 16, which covers four slots at up to three tokens each, which is what - // --parallel 4 --spec-type draft-mtp --spec-draft-n-max 2 produces. Above the bound the column - // split does not fire, and ggml_cuda_warn_above_exact_bound() says so once. + // the widest ubatch a decode step can build: one column per slot times one plus its draft + // tokens, as reported by ggml_backend_cuda_set_exact_decode_width(). Failing a report, 16, + // which covers --parallel 4 --spec-type draft-mtp --spec-draft-n-max 2. Above the bound the + // column split does not fire and ggml_cuda_warn_above_exact_bound() says so once. const int width = g_exact_decode_width.load(std::memory_order_relaxed); return width > 0 ? width : 16; } -// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so its rows depend on the -// other rows in it and the mode does not hold for that op. Say so once, rather than never. -// -// Only when nothing reported a decode width. When one was reported the bound is derived from it, so -// the only batches above the bound are prompt ubatches, and those hold a single sequence under this -// mode: their exactness comes from that, not from the column policy, and leaving them batched is -// the whole point of having a bound at all. Warning on those would be crying wolf on every prefill. +// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so the mode does not hold +// for that op; say so once. Only when nothing reported a decode width: a reported one makes the +// batches above the bound prompt ubatches, which are exact by holding a single sequence, so +// warning on those would be crying wolf on every prefill. static void ggml_cuda_warn_above_exact_bound(const char * op, int64_t ncols, int max_cols) { if (!ggml_cuda_exact_concurrency()) { return; @@ -1929,7 +1918,7 @@ enum ggml_cuda_mm_path { GGML_CUDA_MM_CUBLAS, }; -// The implementation ggml_cuda_mul_mat would pick for a batch of ne11 columns. +// the implementation ggml_cuda_mul_mat would pick for a batch of ne11 columns static ggml_cuda_mm_path ggml_cuda_mul_mat_path( int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, int64_t ne11) { // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. @@ -1966,14 +1955,10 @@ static ggml_cuda_mm_path ggml_cuda_mul_mat_path( static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); -// [TAG_BATCH_INVARIANT] -// The widest slice of columns that can be recomputed in one launch while every column in it still -// sums the way a batch of one would. A column's result depends on the implementation and, for -// MMVQ, on the warp count of the launch, and neither depends on the values of the other columns -// in the launch, so a slice as wide as the batch-of-one configuration reaches gives each of its -// columns the batch-of-one value while reading the weights once for all of them instead of once -// per column. A twelve-column speculative decode over a table whose configuration holds up to four -// columns then costs three weight reads rather than twelve. Always below ncols_dst, so the +// [TAG_BATCH_INVARIANT] the widest slice of columns that can be recomputed in one launch while +// every column still sums as a batch of one would. A column's result depends on the implementation +// and, for MMVQ, the launch's warp count, never on the other columns, so a slice this wide reads +// the weights once for all of them instead of once per column. Always below ncols_dst, so the // recursive call cannot land back here with the same shape. static int64_t ggml_cuda_mul_mat_invariant_width( int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, @@ -1994,14 +1979,13 @@ static int64_t ggml_cuda_mul_mat_invariant_width( return 1; } -// Recompute dst in slices of columns so that each column sees the batch-of-one configuration. -// Returns false when the batched launch already gives every column that same value. +// recompute dst in slices of columns so each column sees the batch-of-one configuration; false +// when the batched launch already gives every column that same value static bool ggml_cuda_mul_mat_split_columns( ggml_backend_cuda_context & ctx, int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - // Recurrent-model output projections broadcast one weight matrix over sequence - // planes. These are token projections too, even though ne[2] or ne[3] is > 1. - // Normalize each plane before applying the existing selective column policy. + // recurrent output projections broadcast one weight matrix over sequence planes: these are + // token projections too, so normalize each plane before applying the column policy if (ggml_cuda_exact_concurrency() && src0->ne[2] == 1 && src0->ne[3] == 1 && (dst->ne[2] > 1 || dst->ne[3] > 1) && src1->ne[2] == dst->ne[2] && src1->ne[3] == dst->ne[3]) { @@ -2023,7 +2007,7 @@ static bool ggml_cuda_mul_mat_split_columns( if (ncols_dst <= 1 || src1->ne[1] != ncols_dst) { return false; } - // Only the token dimension is split, batched matmuls (attention) keep their shape. + // only the token dimension is split; batched matmuls (attention) keep their shape if (src1->ne[2] != 1 || src1->ne[3] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) { return false; } @@ -2033,14 +2017,14 @@ static bool ggml_cuda_mul_mat_split_columns( return false; } - // Mode 1 recomputes one column at a time. Mode 2 recomputes in the widest slices that keep the - // batch-of-one arithmetic, which is what the exact concurrency mode runs under. + // mode 1 recomputes one column at a time; mode 2, which exact concurrency runs under, uses the + // widest slices that keep the batch-of-one arithmetic int64_t width = 1; if (ggml_cuda_batch_invariant() >= 2) { const ggml_cuda_mm_path path_one = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1); const ggml_cuda_mm_path path_batched = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ncols_dst); if (path_one == path_batched) { - // Same implementation, but it still has to sum each destination element in the same order. + // same implementation, but it still has to sum in the same order if (path_batched == GGML_CUDA_MM_MMVF) { return false; // the block size follows K alone } @@ -2122,10 +2106,8 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor GGML_ABORT("fatal error"); } -// [TAG_BATCH_INVARIANT] -// True when the batch-invariant policy computes this MUL_MAT_ID one token at a time. -// Every expert product then reduces the way it would in a batch of one, whatever the -// rest of the ubatch routed to. +// [TAG_BATCH_INVARIANT] true when the policy computes this MUL_MAT_ID one token at a time, so +// every expert product reduces as it would in a batch of one static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) { if (!ggml_cuda_batch_invariant()) { return false; @@ -2152,8 +2134,8 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c return true; } - // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path - // that decides whether the stream is synchronized is the single-token one. + // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path that + // decides whether the stream is synchronized is the single-token one const int64_t ntokens = ggml_cuda_mul_mat_id_splits_tokens(dst) ? 1 : dst->ne[2]; if (ntokens <= MMVQ_MAX_BATCH_SIZE) { @@ -2179,12 +2161,9 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst); -// [TAG_BATCH_INVARIANT] -// Recompute dst one token at a time. Every implementation below groups the ubatch's tokens -// by the expert they routed to, so the column count of an expert's matmul, the tokens the -// per-expert copy gathers and the width the activations are quantized at all depend on what -// the other tokens in the ubatch picked. Handing each token its own call removes that: the -// callee sees the shapes a batch of one has, whatever the neighbours did. +// [TAG_BATCH_INVARIANT] recompute dst one token at a time. Every implementation below groups the +// ubatch's tokens by the expert they routed to, so shapes depend on what the other tokens picked; +// one call per token makes the callee see the shapes a batch of one has. static void ggml_cuda_mul_mat_id_split_tokens(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src1 = dst->src[1]; const ggml_tensor * ids = dst->src[2]; @@ -2232,9 +2211,8 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * // [TAG_BATCH_INVARIANT] if (ggml_cuda_mul_mat_id_splits_tokens(dst)) { GGML_ASSERT(ne3 == 1 && src1->ne[3] == 1 && ids->ne[2] == 1 && ids->ne[3] == 1); - // A quantized expert matrix takes the single-token MMVQ path for every token count, and - // that path can put the tokens on its sample axis in one launch rather than being - // re-entered once per token. Anything else is still recomputed one token at a time. + // a quantized expert matrix takes the single-token MMVQ path at every token count, and that + // path can put the tokens on its sample axis in one launch; anything else goes token by token if (ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); return; @@ -3629,10 +3607,9 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } // topk-moe - // [TAG_BATCH_INVARIANT] The routing fusion passes its memory-range check only when the ubatch - // holds one token, so a request decoding alone picks the fused warp-local top-k kernel and the - // same request decoding next to neighbours picks the softmax, argsort and normalize chain. - // Two algorithms for one set of routing weights is the batch dependence this mode removes. + // [TAG_BATCH_INVARIANT] the routing fusion passes its memory-range check only for a one-token + // ubatch, so a solo request takes the fused top-k kernel and a batched one takes the softmax, + // argsort and normalize chain: two algorithms for one set of routing weights if (!ggml_cuda_batch_invariant() && (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || cgraph->nodes[i]->op == GGML_OP_ARGSORT)) { diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index b4e5196f172..30f37c7722e 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -551,8 +551,8 @@ bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t n // There nwarps also depends on the K loop trip count, which the caller does not pass in. return ncols_dst == 1; } - // blocks_per_iter, which is what assigns K blocks to threads, is proportional to nwarps. - // rows_per_cuda_block only changes which rows a block owns, not the order within a row. + // blocks_per_iter, which assigns K blocks to threads, is proportional to nwarps; + // rows_per_cuda_block only changes which rows a block owns, not the order within a row return calc_nwarps(type, 1, table_id) == calc_nwarps(type, (int) ncols_dst, table_id); } @@ -594,9 +594,8 @@ static __global__ void mul_mat_vec_q( ggml_cuda_pdl_sync(); sample_dst = blockIdx.z; // [TAG_BATCH_INVARIANT] with ids, a sample is a token: the batch-invariant MUL_MAT_ID launch - // puts every token of the batch on the z axis of one single-column launch, so each (token, - // expert slot) block runs the exact single-token configuration. The stock single-token launch - // has one sample, where this indexing is ids[channel_dst] as before. + // puts every token on the z axis of one single-column launch, so each (token, expert slot) + // block runs the single-token configuration. The stock launch has one sample, as before. channel_x = ncols_dst == 1 && ids ? ids[sample_dst*ids_stride + channel_dst] : fastdiv(channel_dst, channel_ratio); channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; @@ -1285,10 +1284,9 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT( nb0 == ts_dst); GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); - // [TAG_BATCH_INVARIANT] under the knob a MUL_MAT_ID with several tokens is computed as one - // launch of the single-token configuration with the tokens on the sample axis, so every - // (token, expert slot) block reduces exactly as the token alone would. The token count is - // then not bounded by the column templates. + // [TAG_BATCH_INVARIANT] a multi-token MUL_MAT_ID becomes one launch of the single-token + // configuration with the tokens on the sample axis, so every (token, expert slot) block + // reduces as the token alone would and the count is not bounded by the column templates const bool tokens_as_samples = ids && ne2 > 1 && ggml_cuda_batch_invariant(); GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE || tokens_as_samples); diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 61a88b851ec..67d69b1415d 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -4,9 +4,8 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11); -// [TAG_BATCH_INVARIANT] -// True when an MMVQ launch of ncols_dst columns sums each destination element in the same -// order as a launch of a single column, i.e. when the column count leaves nwarps unchanged. +// [TAG_BATCH_INVARIANT] true when an MMVQ launch of ncols_dst columns sums each destination element +// in the same order as a single-column launch, i.e. when the column count leaves nwarps unchanged bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst); // Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID, diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp index a3792f3852b..b787108a7aa 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1266,8 +1266,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm (op->src[1]->ne[1] % op->src[4]->ne[1] == 0); break; case GGML_OP_FLASH_ATTN_EXT: - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only - // the CUDA backend reads + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads if (op->src[5]) { supported = false; break; diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index aa20083ec05..574276fed15 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -4157,8 +4157,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons break; case GGML_OP_FLASH_ATTN_EXT: - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only - // the CUDA backend reads + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads supp = op->src[5] == nullptr && ggml_hexagon_supported_flash_attn_ext(sess, op); break; diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 90873f2fab0..e5fc6a8063b 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1592,8 +1592,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_ROLL: return ggml_is_contiguous(op->src[0]); case GGML_OP_FLASH_ATTN_EXT: - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the - // CUDA backend reads; walking the pool in physical order here would be silently wrong + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; + // walking the pool in physical order here would be silently wrong if (op->src[5] != NULL) { return false; } diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index effd11714f9..9c2362c8c4b 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7842,8 +7842,7 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te case GGML_OP_MEAN: return op->src[0]->type == GGML_TYPE_F32; case GGML_OP_FLASH_ATTN_EXT: { - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the - // CUDA backend reads + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads if (op->src[5]) { return false; } diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index dfc9f90926f..1cec1583d51 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -1128,8 +1128,7 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { break; } case GGML_OP_FLASH_ATTN_EXT: { - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only - // the CUDA backend reads + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads if (op->src[5]) { return true; } diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 6fb8851a904..51b552acf75 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1915,8 +1915,8 @@ static ggml_backend_buffer_type_t ggml_backend_rpc_device_get_buffer_type(ggml_b static bool ggml_backend_rpc_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { GGML_UNUSED(dev); - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the - // CUDA backend reads; the remote end is not asked, so it is not claimed here + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; + // the remote end is not asked, so it is not claimed here if (op->op == GGML_OP_FLASH_ATTN_EXT && op->src[5]) { return false; } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 69a344ab790..a31a6a41ca8 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6342,8 +6342,7 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_SOLVE_TRI: return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K; case GGML_OP_FLASH_ATTN_EXT: - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the - // CUDA backend reads + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads return op->src[5] == nullptr && ggml_sycl_flash_attn_ext_supported(device, op); default: return false; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 4a2347219b3..eb855f24d17 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -18192,8 +18192,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } case GGML_OP_FLASH_ATTN_EXT: { - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only - // the CUDA backend reads + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads if (op->src[5]) { return false; } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 70462a97f3c..26717d8804d 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4408,8 +4408,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const break; case GGML_OP_FLASH_ATTN_EXT: { - // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only - // the CUDA backend reads + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads if (op->src[5]) { supports_op = false; break; diff --git a/include/llama.h b/include/llama.h index 09c331ff3e1..077f17ecabf 100644 --- a/include/llama.h +++ b/include/llama.h @@ -795,34 +795,26 @@ extern "C" { // Check if the memory supports shifting LLAMA_API bool llama_memory_can_shift(llama_memory_t mem); - // [TAG_EXACT_CONCURRENCY] Cells the memory allocates in one indivisible unit. - // - // 1 in every ordinary configuration. Larger where a mode places cells in blocks, and - // then a sequence of n tokens occupies round_up(n, granularity) cells. A caller that - // decides whether the pool has room by counting tokens has to round the same way, or it - // will believe there is space that cannot be handed out. + // [TAG_EXACT_CONCURRENCY] cells the memory allocates in one indivisible unit: 1 ordinarily, + // larger where a mode places cells in blocks, and then n tokens occupy round_up(n, granularity) + // cells. A caller deciding whether the pool has room must round the same way. LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem); - // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or - // 1 plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence - // set with more tokens than this left to place is a prompt and is prefilled in a ubatch of its - // own; a set at or below it is a decode step and stays grouped with the other decodes, so a - // speculative verify batch is not run once per sequence. Process-wide, default 1. Raising it - // widens the decode step of every context that exists, and their width is re-reported with - // it; false, and no change, when an explicit column bound given to the backend cannot cover - // that width (see llama_set_exact_decode_width). Never lowers what was set: a narrower context - // set up later must not turn an existing context's verify steps into prompts. + // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or 1 + // plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence + // set with more left to place is a prompt and is prefilled in a ubatch of its own; one at or + // below stays grouped with the other decodes. Process-wide, default 1, never lowered. Raising + // it widens every existing context's decode step and re-reports their width; returns false, + // and changes nothing, when an explicit column bound cannot cover that width. LLAMA_API bool llama_set_exact_decode_tokens(uint32_t n_tokens); LLAMA_API uint32_t llama_exact_decode_tokens(void); // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns: the - // sequences a context can hold times the tokens each contributes to a step. Every context - // reports its own at creation and a backend keeps the widest it has heard, so a decode of any - // context stays within the bound its kernels split at. A caller that builds wider steps than - // the contexts imply (a draft of its own, say) reports the width itself, before creating the - // context or before the first decode. Never lowers what was reported. Returns false, and - // reports nothing, when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is set to a positive figure below - // the width: that bound wins in the backend, so decodes above it would be left batched. + // sequences a context holds times the tokens each contributes. Every context reports its own + // at creation and a backend keeps the widest it has heard. A caller that builds wider steps + // reports the width itself, before the context or the first decode. Never lowered. Returns + // false, reporting nothing, when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is positive and below the + // width: that bound wins in the backend, so decodes above it would be left batched. LLAMA_API bool llama_set_exact_decode_width(uint32_t n_cols); LLAMA_API uint32_t llama_exact_decode_width(void); diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py index ec2c3a00467..198c2bdf881 100644 --- a/scripts/batchinv/divergence.py +++ b/scripts/batchinv/divergence.py @@ -5,8 +5,8 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from prompts import PROMPTS -# Environment recorded with every run. LLAMA_EXACT_CONCURRENCY inherited from the shell is what -# decides whether a run labelled as the mode-off reference actually was one, so it is not optional. +# recorded with every run; LLAMA_EXACT_CONCURRENCY inherited from the shell decides whether a run +# labelled as the mode-off reference actually was one, so it is not optional RECORDED_ENV = ("LLAMA_EXACT_CONCURRENCY", "GGML_CUDA_BATCH_INVARIANT", "GGML_CUDA_BATCH_INVARIANT_MAX_COLS", "LLAMA_SERVER_PREEMPT_EVERY", "LLAMA_KV_CACHE_DEBUG", "LLAMA_BATCH_DEBUG", "CUDA_VISIBLE_DEVICES") @@ -86,14 +86,13 @@ def __enter__(self): time.sleep(1.0) raise RuntimeError("server did not become healthy") except BaseException: - # __exit__ is not called when __enter__ raises, so a server that started but never + # __exit__ is not called when __enter__ raises, and a server that started but never # reported healthy would keep the GPU, the port and the log handle self.__exit__(None, None, None) raise def __exit__(self, *a): - # note: POSIX only. On Windows this needs CREATE_NEW_PROCESS_GROUP at Popen and - # terminate()/kill() here; the runs this harness backs are Linux only. + # note: POSIX only; Windows would need CREATE_NEW_PROCESS_GROUP at Popen if self.p is not None: print(f"[server] stopping pid={self.p.pid}", flush=True) try: @@ -141,8 +140,7 @@ def work(name): t.join() wall = time.time() - t0 - # a thread exception used to only print a traceback, so a run where P1..P3 failed and P0 - # succeeded was still reported as a clean four-way concurrency result + # without this a run where P1..P3 failed and P0 succeeded reads as a clean four-way result if errors: raise RuntimeError("concurrent requests failed: " + "; ".join(f"{n}: {type(e).__name__}: {e}" for n, e in errors)) diff --git a/scripts/batchinv/probe.cpp b/scripts/batchinv/probe.cpp index c4e478f97c7..204379c423b 100644 --- a/scripts/batchinv/probe.cpp +++ b/scripts/batchinv/probe.cpp @@ -1,6 +1,6 @@ -// Locate the first graph op whose sequence-0 output changes when the decode batch -// holds four sequences instead of one. Prompt KV for seq 0 is built identically in -// both phases, so the only difference is the width of the final decode ubatch. +// Locate the first graph op whose sequence-0 output changes when the decode batch holds four +// sequences instead of one. Seq 0's prompt KV is identical in both phases, so the only +// difference is the width of the final decode ubatch. #include "llama.h" #include "ggml.h" #include "ggml-backend.h" @@ -120,7 +120,7 @@ static llama_token greedy(llama_context * ctx, int32_t i, int n_vocab) { return best; } -// Feed a prompt as one decode call for one sequence, return the greedy next token. +// feed a prompt as one decode call for one sequence, return the greedy next token static llama_token feed(llama_context * ctx, const std::vector & p, llama_seq_id seq, int n_vocab) { batch_holder h; for (size_t i = 0; i < p.size(); ++i) { @@ -167,14 +167,12 @@ int main(int argc, char ** argv) { std::vector rec_a, rec_b; llama_token first_tok[4] = {0, 0, 0, 0}; - // Phase A: decode ubatch width 1. PROBE_A_FILL controls how many sequences are - // already in the shared KV cache, which is what sets K->ne[1] for attention. + // phase A: decode ubatch width 1. PROBE_A_FILL is how many sequences are already in the + // shared KV cache, which sets K->ne[1] for attention. const int a_fill = getenv("PROBE_A_FILL") ? atoi(getenv("PROBE_A_FILL")) : 1; - // PROBE_A_PERM reorders which prompt goes into which sequence in phase A. With the same - // multiset of prompts the cache keeps its length but the masked cells hold different data. - // Phase B decodes prompt 0's first token on sequence 0, so the permutation may only move - // the neighbours: sequence 0 keeps prompt 0, or the two phases would compare different - // sequences. + // PROBE_A_PERM reorders which prompt goes into which sequence in phase A, keeping the cache + // length but changing what the masked cells hold. It may only move the neighbours: sequence 0 + // keeps prompt 0, or the two phases would compare different sequences. int a_perm[4] = {0, 1, 2, 3}; if (const char * perm = getenv("PROBE_A_PERM")) { for (int k = 0; k < 4 && perm[2*k]; ++k) a_perm[k] = perm[2*k] - '0'; @@ -197,7 +195,7 @@ int main(int argc, char ** argv) { llama_free(ctx); } - // Phase B: same seq-0 prompt KV, then a decode ubatch holding n_seqs tokens. + // phase B: same seq-0 prompt KV, then a decode ubatch holding n_seqs tokens { llama_context * ctx = make_ctx(); if (prefill) { @@ -240,7 +238,7 @@ int main(int argc, char ** argv) { llama_free(ctx); } - // Optional: keep decoding and report the first step at which seq 0's token differs. + // optional: keep decoding and report the first step at which seq 0's token differs const int n_steps = getenv("PROBE_STEPS") ? atoi(getenv("PROBE_STEPS")) : 0; int first_bad_step = -1; if (n_steps > 0) { @@ -279,7 +277,7 @@ int main(int argc, char ** argv) { fprintf(stderr, "nodes: A=%zu B=%zu first tokens: %d %d %d %d\n", rec_a.size(), rec_b.size(), first_tok[0], first_tok[1], first_tok[2], first_tok[3]); - // Walk both node lists in order and compare seq 0's slice. + // walk both node lists in order and compare seq 0's slice FILE * out = out_path ? fopen(out_path, "w") : stdout; fprintf(out, "{\"n_seqs\":%d,\"first_bad_step\":%d,\"nodes_a\":%zu,\"nodes_b\":%zu,\"diffs\":[", n_seqs, first_bad_step, rec_a.size(), rec_b.size()); size_t n = rec_a.size() < rec_b.size() ? rec_a.size() : rec_b.size(); @@ -295,8 +293,8 @@ int main(int argc, char ** argv) { verdict = "misaligned"; } else if (A.op == "GATED_DELTA_NET" && A.gdn_tokens == B.gdn_tokens && !A.data.empty() && !B.data.empty()) { - // Packed GDN outputs put ALL token outputs before ALL sequence states. - // Sequence 0's state therefore moves when the number of sequences changes. + // packed GDN outputs put all token outputs before all sequence states, so seq 0's + // state moves when the number of sequences changes const size_t output = A.ne[0]*A.gdn_tokens; const size_t state = A.ne[0]*A.ne[1]/A.gdn_seqs - output; for (size_t k = 0; k < output + state; ++k) { diff --git a/scripts/batchinv/prompts.py b/scripts/batchinv/prompts.py index 860b65fe08d..efabc448a9d 100644 --- a/scripts/batchinv/prompts.py +++ b/scripts/batchinv/prompts.py @@ -1,4 +1,4 @@ -# Four distinct prompts, each about 300 tokens of raw text (no chat template). +# four distinct prompts, each about 300 tokens of raw text (no chat template) _BODIES = { "P0": """The history of numerical computing is a history of compromises between speed and exactness. Early machines used fixed point arithmetic because it was cheap, and programmers carried scaling diff --git a/scripts/unsloth/additive_merge.py b/scripts/unsloth/additive_merge.py index 5364dc1b7c1..d930158a542 100644 --- a/scripts/unsloth/additive_merge.py +++ b/scripts/unsloth/additive_merge.py @@ -94,18 +94,11 @@ def nonblank(lines: list[str]) -> list[str]: return [ln.strip() for ln in lines if ln.strip()] -# A line that closes or opens a block and nothing else. Two INDEPENDENT case -# arms in the same switch share these by construction -- `{`, `} break;`, `}` -# are what a case arm is made of, not what makes it that case arm -- so finding -# them on both sides says nothing about whether the two sides added the same -# construct. Matching them as "shared" is what refused the real add/add of -# PROJECTOR_TYPE_KIMIK3 next to PROJECTOR_TYPE_DEEPSEEK4V in tools/mtmd/clip.cpp -# with "one change made twice: {, } break;", when the two arms had no line of -# actual content in common. -# -# Deliberately narrow: braces, brackets, parens, semicolons and commas, around -# at most one bare block-terminating keyword. `break;` matches, `return true;` -# does not, and anything naming a type, a constant or a function does not. +# A line that only opens or closes a block. Two independent case arms share these by +# construction, so finding them on both sides says nothing about the two sides adding the +# same construct: treating them as shared is what refused the real PROJECTOR_TYPE_KIMIK3 / +# PROJECTOR_TYPE_DEEPSEEK4V add/add in tools/mtmd/clip.cpp. Deliberately narrow: brackets, +# semicolons and commas around at most one bare block-terminating keyword. STRUCTURAL = re.compile(r"^[\s{}()\[\];,]*(?:break|continue|return|pass)?[\s{}()\[\];,]*$") @@ -114,8 +107,7 @@ def identifying(lines: list[str]) -> set[str]: return {ln for ln in nonblank(lines) if not STRUCTURAL.match(ln)} -# `case FOO:`, `case FOO :`, `default:`. A fallthrough label may carry no body -# at all, which is the shape the nightly hits most often. +# `case FOO:`, `case FOO :`, `default:`; a fallthrough label may carry no body at all CASE_LABEL = re.compile(r"^(?:case\s+[^:]+|default\s*):") @@ -147,32 +139,23 @@ def resolve_region(ours: list[str], base: list[str], theirs: list[str]) -> list[ return list(ours) ours_arms, theirs_arms = case_arms(ours), case_arms(theirs) if ours_arms and theirs_arms and ours_arms.isdisjoint(theirs_arms): - # Both sides added case arms, and not one label is on both sides. Two - # arms of the same switch labelled differently are two constructs, so - # any line they happen to share is body text, not a duplicate: the real - # tools/mtmd/clip.cpp collision has a KIMIK3 arm and a DEEPSEEK4V arm - # that both set `hparams.rope_theta = 10000.0f;`, and refusing on that - # coincidence is what the shared-line check is for, backwards. - # - # The same change made twice would keep its label, so it lands in the - # check below instead. This is the one place where a shared line is - # allowed, and it is allowed because the labels prove the arms are - # distinct -- a duplicated label would not even compile. + # Both sides added case arms and no label is on both, so they are two constructs + # and any line they share is body text: the real clip.cpp collision has arms that + # both set `hparams.rope_theta = 10000.0f;`. The same change made twice would keep + # its label and land in the check below, so this is the one place a shared line is + # allowed - a duplicated label would not even compile. return list(theirs) + list(ours) shared = identifying(ours) & identifying(theirs) if shared: # Overlapping content is the signature of one construct added twice, # not two independent additions. Unioning it would duplicate code. - # Scaffolding lines are excluded above, so what is left is content both - # sides genuinely wrote, which is the thing that makes this a duplicate. + # scaffolding is excluded above, so what is left is content both sides wrote raise Unresolvable( "both sides add the same line(s), so this is one change made twice: " + ", ".join(sorted(shared)[:3]) ) if not identifying(ours) or not identifying(theirs): - # Everything one side added is scaffolding, so there is no content to - # tell the two additions apart and the exclusion above has nothing left - # to work with. Refuse rather than union braces onto braces. + # one side is all scaffolding, so there is no content to tell the additions apart raise Unresolvable( "one side adds only block scaffolding, so the two additions cannot " "be told apart" diff --git a/scripts/unsloth/feature_matrix.py b/scripts/unsloth/feature_matrix.py index 00b392ca1ec..d8fad55d469 100644 --- a/scripts/unsloth/feature_matrix.py +++ b/scripts/unsloth/feature_matrix.py @@ -37,7 +37,7 @@ import sys from pathlib import Path -# Output that means "this did not run" from a process that exited 0. +# output that means "this did not run" from a process that exited 0 SKIP_RE = re.compile(r"\bSKIP\b|not supported|unsupported|no tests|0 tests", re.I) @@ -78,7 +78,7 @@ def probe_arch(check: dict, b: Path, gpu: bool) -> str: rc, out = run([str(b / "test-llama-archs"), "-a", arch, "-s", "1234"], b, gpu) if rc != 0: raise Unproven(f"test-llama-archs -a {arch} exited {rc}") - # The arch's own rows, not the header and not another arch's. + # the arch's own rows, not the header and not another arch's rows = [ln for ln in out.splitlines() if ln.strip().startswith("|") and f"|{arch:>16}|" in ln or (ln.strip().startswith("|") and ln.split("|")[1].strip() == arch)] if not rows: @@ -119,8 +119,8 @@ def probe_mtmd(check: dict, b: Path, gpu: bool) -> str: m = re.search(r"assertions\s*:\s*(\d+)", out) if not m or int(m.group(1)) == 0: raise Unproven("test_projector_registry ran no assertions; the filter matched nothing") - # The registry test walks the whole enum, so it proves the table is sound. - # That the specific projector is IN the enum is pin_contract.py's job. + # the test walks the whole enum, so it proves the table is sound; that this projector + # is IN the enum is pin_contract.py's job return f"projector registry intact over {m.group(1)} assertions" @@ -173,8 +173,7 @@ def main() -> int: print(f"ok {name}: " + "; ".join(r["evidence"] for r in entry["results"]) + (f" [{len(entry['deferred'])} needs a GPU]" if entry["deferred"] else "")) else: - # Nothing was shown either way. Not a failure here, but it must not - # read as one of the ok lines. + # nothing shown either way: not a failure, but not an ok line either print(f"-- {name}: nothing provable without a GPU " f"({len(entry['deferred'])} check(s) deferred)") @@ -188,8 +187,7 @@ def main() -> int: if failed: print(f"\n{failed} feature(s) could not be shown to work", file=sys.stderr) return 1 - # Say what was NOT proven in the same breath as what was. A run that only - # ever prints a success line teaches the reader that green means covered. + # say what was NOT proven alongside what was, or green starts to read as covered tail = f", {deferred} check(s) need a GPU and were not run" if deferred else "" print(f"\nall {len(report['features'])} features demonstrated" + (" on GPU" if args.gpu else " on CPU") + tail) diff --git a/scripts/unsloth/pin_contract.py b/scripts/unsloth/pin_contract.py index bf07d58408e..d6a60746098 100644 --- a/scripts/unsloth/pin_contract.py +++ b/scripts/unsloth/pin_contract.py @@ -57,34 +57,28 @@ r"^https://github\.com/([^/]+)/llama\.cpp/pull/(\d+)/commits/([0-9a-f]{40})/?$" ) -# Identifier families that name a FEATURE. Deliberately not "every new symbol": -# a helper function renamed by a later upstream commit is not a lost feature, -# but a missing LLM_ARCH_ entry always is. These are the tables that decide -# whether an architecture, an op, a projector or a quant type exists at all. +# Identifier families that name a FEATURE, not every new symbol: a renamed helper is not a +# lost feature, but a missing LLM_ARCH_ entry always is. SYMBOL_FAMILIES = ( "LLM_ARCH_", "LLM_TENSOR_", "LLM_KV_", "LLM_TYPE_", "PROJECTOR_TYPE_", "GGML_OP_", "GGML_TYPE_", "LLAMA_FTYPE_", ) SYMBOL_RE = re.compile(r"\b(?:" + "|".join(SYMBOL_FAMILIES) + r")[A-Z0-9_]+\b") -# The subset that names a whole feature rather than one of its tensors. Used -# only to keep --emit readable; the check itself uses all of SYMBOL_FAMILIES. +# the subset naming a whole feature; only to keep --emit readable, the check uses them all HEADLINE = ("LLM_ARCH_", "GGML_OP_", "GGML_TYPE_", "PROJECTOR_TYPE_", "LLAMA_FTYPE_") -# A line worth tracking for survival. Comments and short punctuation drift with -# every reformat and would make the check noise; a substantial code line does -# not move on its own. +# a line worth tracking for survival: comments and short punctuation drift with every +# reformat, a substantial code line does not move on its own TRIVIAL_RE = re.compile(r"^\s*(?://|/\*|\*|\*/|#\s|$)") MIN_LINE = 12 -# Comments are stripped before anything is read off a line. A pin that merely -# NAMES an arch in a comment has not registered it, and holding the comment's -# wording as a contract fails the moment upstream rewords it. Observed on -# unslothai#70, whose comment mentions GGML_OP_SSM_SCAN to explain why it does -# NOT use it. +# Comments are stripped first: a pin that merely NAMES an arch in a comment has not +# registered it, and holding the wording as a contract fails when upstream rewords it. +# Observed on unslothai#70, whose comment mentions GGML_OP_SSM_SCAN to say it is not used. COMMENT_RE = re.compile(r"//.*$|/\*.*?\*/|(? dict: if code: symbols[cur].update(SYMBOL_RE.findall(code)) - # Only symbols the base does not ALREADY have in that file are evidence of - # this pin. Upstream naming an arch in a file the pin also touches is not - # something the pin is owed. + # only symbols the base does not already have in that file are evidence of this pin new_symbols: dict[str, list[str]] = {} for path, names in symbols.items(): fresh = sorted(n for n in names @@ -306,9 +298,8 @@ def main() -> int: if args.emit: report["pins"].append(entry) - # Only the families that NAME a feature are printed. Every symbol - # is still checked; a new file legitimately contributes a hundred - # LLM_TENSOR_ names and listing them buries the one that matters. + # only feature-naming families are printed; all are still checked, but a + # hundred LLM_TENSOR_ names would bury the one that matters sym = sorted({s for v in contract["symbols"].values() for s in v if s.startswith(HEADLINE)}) print(f"{name:>18} {entry['line_count']:>5} lines, " @@ -345,8 +336,7 @@ def main() -> int: if args.emit: return 0 - # Notices after the verdict lines, never mixed into them: "upstream took - # this, drop the entry" is housekeeping and must not read as a failure. + # notices after the verdict lines: housekeeping must not read as a failure for n in notices: print(f"note {n}") if failed: diff --git a/scripts/unsloth/test_additive_merge.py b/scripts/unsloth/test_additive_merge.py index 0f91e9afd12..dd2b22c6cac 100644 --- a/scripts/unsloth/test_additive_merge.py +++ b/scripts/unsloth/test_additive_merge.py @@ -97,8 +97,7 @@ def run(repo, *extra): reason.endswith('twice: log("same");'), reason) # --- 3b. two independent case arms: braces are shared, content is not ------- -# The real tools/mtmd/clip.cpp shape. Refusing this on `{` and `} break;` is -# what took the 09-02 nightly's last pin down. +# the real clip.cpp shape; refusing it on `{` and `} break;` took the 09-02 nightly down base = "switch (t) {\n}\n" ours = ("switch (t) {\n case PROJECTOR_TYPE_KIMIK3:\n {\n" " builder = std::make_unique(ctx, img);\n" @@ -115,8 +114,7 @@ def run(repo, *extra): txt.count("} break;") == 2 and txt.count("clip_graph_kimik3") == 1, txt) # --- 3b2. two case arms that share a body line, which is a coincidence ------ -# The clip.cpp shape after upstream landed DEEPSEEK4V: both arms set the same -# rope_theta, and refusing on that is the shared-line check backwards. +# clip.cpp after upstream landed DEEPSEEK4V: both arms set the same rope_theta base = "switch (t) {\n}\n" ours = ("switch (t) {\n case PROJECTOR_TYPE_KIMIK3:\n {\n" " hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;\n" diff --git a/scripts/unsloth/test_pin_contract.py b/scripts/unsloth/test_pin_contract.py index 3c7b58d62d5..e76586b1fc4 100644 --- a/scripts/unsloth/test_pin_contract.py +++ b/scripts/unsloth/test_pin_contract.py @@ -98,8 +98,7 @@ def run(repo, pr_set, *extra): check("intact merge reports no notices", rep["notices"] == [], rep) # --- 2. the arm is dropped from ONE file: a tree-wide grep would pass ------ -# The real shape: LLM_ARCH_INKLING survives in the enum and the dispatch arm -# that makes it do anything is gone. +# the real shape: LLM_ARCH_INKLING survives in the enum, the dispatch arm is gone repo, pr_set, sha = make_repo() p = repo / "src" / "llama-model.cpp" p.write_text(MODEL_CPP_BASE) @@ -128,8 +127,8 @@ def run(repo, pr_set, *extra): any("do_the_banded_thing" in x for x in rep["pins"][0]["problems"]), rep) # --- 5. redundancy: the base already has everything the pin adds ---------- -# Built the way it happens for real: upstream lands the same work, so the base -# tag has it and the pin is not an ancestor of anything. +# as it happens for real: upstream lands the same work, so the base tag has it and +# the pin is not an ancestor of anything d = Path(tempfile.mkdtemp(prefix="pc_")) git(d, "init", "-q", "-b", "main") (d / "src").mkdir() @@ -163,8 +162,8 @@ def run(repo, pr_set, *extra): rep["pins"][0]["added_files"] == ["src/inkling.cpp"], rep) # --- 7. a comment is not a contract --------------------------------------- -# unslothai#70 has a comment naming GGML_OP_SSM_SCAN to say it does NOT use it. -# Holding comment wording would fail the moment upstream rewords it. +# unslothai#70 names GGML_OP_SSM_SCAN in a comment to say it does NOT use it, and +# holding that wording would fail the moment upstream rewords it repo, pr_set, sha = make_repo() git(repo, "checkout", "-q", "pin") (repo / "src" / "note.cpp").write_text( diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index cc73d83963c..4fc1c808631 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -573,18 +573,13 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } if (add) { - // [TAG_EXACT_CONCURRENCY] a sequence set that still has more tokens to place than a - // decode step carries is a prompt, and a prompt shares its arithmetic with whatever - // else is in the ubatch, so give it a ubatch of its own. Sets at or below that width - // are decode steps, plain or speculative, whose columns the backend's column policy - // already keeps exact, so keep grouping those: isolating them too would make one - // prompt serialize every concurrent decode for the whole of the prefill, and would run - // a speculative verify step once per sequence. Grouped sets must have the same number - // of tokens left: the equal-length expansion below would otherwise place a three-token - // verify step beside a two-token one as two tokens now and one later, and a memory - // that reduces over a chunk of tokens (a chunked state space scan) would then sum in a - // different order than the solo run's single three-token ubatch. A set with a - // different count waits for a later ubatch. + // [TAG_EXACT_CONCURRENCY] a set with more tokens left than a decode step carries is a + // prompt, and a prompt shares its arithmetic with whatever else is in the ubatch, so + // give it one of its own. Sets at or below that width are decode steps, kept exact by + // the backend's column policy, so keep grouping them or one prompt would serialize + // every concurrent decode. Grouped sets must have the same number of tokens left, or + // the equal-length expansion below would cut a longer set in two and a memory that + // reduces over a chunk of tokens would sum in a different order than the solo run. if (isolate_seqs_above > 0) { uint32_t n_left = 0; @@ -612,8 +607,8 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } else if (n_left != n_left_first) { continue; } else if ((cur_seq_set.size() + 1) * n_left_first > n_ubatch) { - // one more set would not finish in this ubatch: the expansion below would - // then cut every set part way, the chunking the guard exists to prevent + // one more set would not finish here, and the expansion below would then cut + // every set part way: the chunking this guard exists to prevent break; } } diff --git a/src/llama-batch.h b/src/llama-batch.h index 52a375ad49e..f0ecc9d8407 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -105,12 +105,9 @@ class llama_batch_allocr { // make ubatches of equal-length sequences sets // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch - // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this many - // tokens left to place is a prompt and is given a ubatch of its own; sets at or - // below it are decode steps (one token, or one plus the speculative drafts) and - // stay grouped together, so a prompt next to three decodes costs one extra ubatch - // and does not serialize the three decodes, and a speculative verify batch is not - // run once per sequence + // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this + // many tokens left to place is a prompt and gets a ubatch of its own; sets at or + // below it are decode steps and stay grouped, so a prompt does not serialize them llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above = 0); // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than n_tokens left to place, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 8d35922a47b..16e090a0623 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -102,18 +102,14 @@ llama_context::llama_context( } // [TAG_EXACT_CONCURRENCY] the widest decode step this context can build: one column per - // sequence, times the tokens a sequence contributes to a step. Reported so that a backend - // splitting columns for exactness covers it without the caller having to know the bound; a - // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width. - // The sequence count is what is reported: the tokens figure can be raised later for the - // whole process, and the width then follows it for this context too. - // Checked here and reported at the end of the constructor: the count is process-wide - // state that outlives a context, so a construction that fails later on, an unsupported - // cache layout say, must not leave a width behind that no context needs. + // sequence times the tokens a sequence contributes, reported so a backend splitting columns + // covers it (a caller that builds wider steps uses llama_set_exact_decode_width). The + // sequence count is what is reported, so a later rise in the tokens figure follows it here + // too. Checked now but reported at the end of the constructor, so a construction that fails + // later does not leave a width behind that no context needs. if (llama_exact_concurrency()) { - // an explicit column bound wins over the reported width in the backend, so one below - // this context's width would leave its decodes batched above the bound with the mode - // still reporting itself on; the report refuses that, and the refusal is an error here + // an explicit column bound wins in the backend, so one below this context's width would + // leave decodes batched above it; the report refuses that, and that is an error here if (!llama_exact_check_n_seq(cparams.n_seq_max)) { throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); } @@ -412,8 +408,8 @@ llama_context::llama_context( memory.reset(model.create_memory(params_mem, cparams)); - // [TAG_EXACT_CONCURRENCY] the paged attention the mode runs on is causal; a context - // created non-causal with a cache would assert on its first graph, so it is refused here + // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a non-causal context with a + // cache would assert on its first graph if (llama_exact_concurrency() && memory && !cparams.causal_attn) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so it cannot be created with non-causal attention\n", __func__); throw std::runtime_error("exact concurrency: non-causal attention is not supported with a KV cache"); @@ -502,9 +498,8 @@ llama_context::llama_context( } } - // [TAG_EXACT_CONCURRENCY] nothing above can fail any more, so the width this context - // needs is published now; checked against the explicit bound at the top, so this - // cannot refuse unless the bound moved underneath it, which is an error all the same + // [TAG_EXACT_CONCURRENCY] nothing above can fail now, so publish the width; already checked + // against the explicit bound at the top, so a refusal here means the bound moved if (llama_exact_concurrency() && !llama_exact_report_n_seq(cparams.n_seq_max)) { throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); } @@ -1220,8 +1215,8 @@ void llama_context::set_causal_attn(bool value) { return; } - // [TAG_EXACT_CONCURRENCY] the paged attention the mode runs on is causal; a context with a - // cache under the mode keeps causal attention rather than asserting in the next graph + // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a context with a cache keeps + // causal attention rather than asserting in the next graph if (!value && memory && llama_exact_concurrency()) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so causal attention cannot be turned off; the change is refused\n", __func__); return; @@ -2664,17 +2659,16 @@ class llama_io_read_host : public llama_io_read_i { } const size_t tensor_bytes = ggml_nbytes(tensor); auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - // A fragmented sequence can require thousands of synchronous device - // transfers per layer. For bounded tensors, stage the tensor once and - // preserve every byte belonging to other sequences. Bound scratch RAM - // and leave ordinary contiguous transfers on their original fast path. + // a fragmented sequence can need thousands of synchronous device transfers per + // layer: stage a bounded tensor once instead, preserving other sequences' bytes and + // leaving ordinary contiguous transfers on their fast path if (end - i >= 64 && tensor_bytes <= 64 * 1024 * 1024 && !ggml_backend_buffer_is_host(buffer)) { std::vector staging; try { staging.resize(tensor_bytes); } catch (const std::bad_alloc &) { - // Fall back to the individual transfers below. + // fall back to the individual transfers below } if (!staging.empty()) { ggml_backend_tensor_get(tensor, staging.data(), 0, tensor_bytes); @@ -3266,8 +3260,8 @@ size_t llama_context::state_write_data(llama_io_write_i & io) { size_t llama_context::state_read_data(llama_io_read_i & io) { // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical - // index, which the paged pool owns. Refused here, before anything is parsed, so that the - // cache the caller has is left as it was: the generic restore path clears it on failure. + // index, which the paged pool owns. Refused before anything is parsed, so the caller's cache + // is left as it was: the generic restore path clears it on failure. if (memory && memory->alloc_granularity() > 1) { throw std::runtime_error("whole-context restore is not supported with LLAMA_EXACT_CONCURRENCY, restore per sequence"); } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index f8f5c83a2b6..dc02560fc66 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -27,18 +27,11 @@ // dedup helpers -// [TAG_EXACT_CONCURRENCY] -// The page table is wired into llm_graph_input_attn_kv only. The V-less layouts build their -// attention without one, so a model on one of those would get its cells placed in pages by the -// allocator and then attend in physical cell order anyway: the mode would report itself as on and -// lose the one invariant it exists for, which is the same silent failure the CUDA placement gate -// was added to stop. Refuse the context instead. -// -// Rejecting is the smaller correct change here. Wiring self_pages into llm_graph_input_attn_k alone -// is four lines, but it fixes only one of the four V-less input classes, and DeepSeek 3.2 uses two -// of them: its sparse layers rewrite the mask from a top-k selection and would still be unpaged, so -// the model would end up half paged, which is worse than refused. None of these architectures was -// measured, and the paged kernel additionally requires 256-dimensional K and V heads. +// [TAG_EXACT_CONCURRENCY] the page table is wired into llm_graph_input_attn_kv only, so a V-less +// layout would have its cells placed in pages and then attend in physical order anyway, with the +// mode reporting itself as on. Refuse the context instead: wiring self_pages into +// llm_graph_input_attn_k alone fixes one of the four V-less classes, and DeepSeek 3.2 uses two of +// them, so the model would end up half paged, which is worse than refused. static void llm_graph_reject_exact_concurrency(const char * layout) { if (!llama_exact_concurrency()) { return; diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index 0b218bc64e7..f577414a51c 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -187,22 +187,18 @@ bool llama_exact_concurrency() { // [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h static std::atomic g_exact_decode_tokens{1}; -// one lock for the token figure, the sequence count and the width: the three move together -// (a context reports its count and the width that follows; a new token figure re-reports the -// width for every count seen), and a report interleaved with a change of figure could leave -// the backend with a width that covers neither. Recursive, since the setters call each other. +// one lock for the token figure, the sequence count and the width, since the three move together +// and a report interleaved with a change of figure could leave the backend with a width that +// covers neither; recursive, since the setters call each other static std::recursive_mutex g_exact_mutex; -// the most sequences any context so far was created with. The tokens figure is process -// wide, so raising it widens the decode step of every context that already exists; the -// width those contexts reported at creation is re-reported here with the new figure, or a -// context created under a narrower figure would batch above the bound it reported. +// the most sequences any context was created with. The tokens figure is process wide, so raising +// it widens every existing context's decode step and their width is re-reported with it. static std::atomic g_exact_max_n_seq{0}; static bool llama_exact_width_within_explicit_bound(uint32_t n_cols); -// the width is sequences times tokens, handed to a backend as an int; a product that does not -// fit is refused rather than wrapped +// sequences times tokens, handed to a backend as an int; a product that overflows is refused static bool llama_exact_width_of(uint32_t n_seq, uint32_t n_tokens, uint32_t & n_cols) { const uint64_t w = (uint64_t) n_seq * (uint64_t) n_tokens; @@ -250,14 +246,14 @@ bool llama_set_exact_decode_tokens(uint32_t n_tokens) { std::lock_guard lock(g_exact_mutex); - // never lowered: a narrower context set up later would otherwise turn the verify steps of - // an existing speculative context into prompts and serialise them + // never lowered: a narrower context set up later would turn an existing speculative + // context's verify steps into prompts and serialise them if (n_tokens <= g_exact_decode_tokens.load(std::memory_order_relaxed)) { return true; } - // every context that exists widens with the figure, so the width they will need is - // reported first; a figure the explicit bound cannot cover leaves the old one in place + // every context widens with the figure, so report the width first; one the explicit bound + // cannot cover leaves the old figure in place const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed); uint32_t n_cols = 0; @@ -275,14 +271,13 @@ uint32_t llama_exact_decode_tokens(void) { return g_exact_decode_tokens.load(std::memory_order_relaxed); } -// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h. A backend that -// splits columns to make a decode exact reads it through ggml_backend_cuda_set_exact_decode_width, -// reached through the registry so that a backend that is absent or loaded late costs nothing. +// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h. Backends read it +// through ggml_backend_cuda_set_exact_decode_width, reached through the registry so an absent or +// late-loaded backend costs nothing. static std::atomic g_exact_decode_width{0}; -// an explicit column bound given to the CUDA backend wins over the reported width there, so a -// width above it would leave decodes batched past the bound with the mode still reporting itself -// on; a width the bound does not cover is refused instead of stored +// an explicit column bound wins in the CUDA backend, so a width above it would leave decodes +// batched past the bound: refuse such a width instead of storing it static bool llama_exact_width_within_explicit_bound(uint32_t n_cols) { static const int explicit_cols = []() { const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); @@ -310,9 +305,8 @@ bool llama_set_exact_decode_width(uint32_t n_cols) { while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { } - // The widest figure so far goes to every backend on every call, not only when it grew: a - // width reported before a backend was loaded would otherwise never reach it, and every - // context reports at creation, by which time the backends are there. + // the widest figure goes to every backend on every call, not only when it grew, or a width + // reported before a backend was loaded would never reach it const uint32_t widest = g_exact_decode_width.load(std::memory_order_relaxed); for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { diff --git a/src/llama-impl.h b/src/llama-impl.h index 65d56a51d1d..d0ab2e5cf1c 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -104,16 +104,14 @@ std::string llama_format_tensor_shape(const struct ggml_tensor * t); std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i); -// [TAG_EXACT_CONCURRENCY] -// opt-in mode under which a sequence's attention depends only on its own cells, in position order, -// so that its output does not change when other sequences share the KV cache. Off by default. -// Reads the same LLAMA_EXACT_CONCURRENCY variable as the paged KV cache and the CUDA backend. +// [TAG_EXACT_CONCURRENCY] opt-in mode under which a sequence's attention depends only on its own +// cells, in position order, so its output does not change when others share the KV cache. Off by +// default; reads the same variable as the paged KV cache and the CUDA backend. bool llama_exact_concurrency(); -// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so that the -// decode width every context needs is known to the backend and follows llama_set_exact_decode_tokens +// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so the backend +// knows the width every context needs and it follows llama_set_exact_decode_tokens bool llama_exact_report_n_seq(uint32_t n_seq); -// the same check without the report: whether a context of n_seq sequences could be reported -// under the explicit column bound, for a constructor that may still fail after asking +// the same check without the report, for a constructor that may still fail after asking bool llama_exact_check_n_seq(uint32_t n_seq); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 31154639ad9..428fbc6b567 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -62,11 +62,9 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { // llama_kv_cache // -// [TAG_EXACT_CONCURRENCY] -// The paged attention specialization that reads the logical page table lives in the CUDA backend -// sources, which are also built as the ROCm and MUSA backends. Every other backend ignores src[5] -// and walks the pool in physical cell order, so a KV layer placed there would silently lose the -// guarantee the mode exists to provide. +// [TAG_EXACT_CONCURRENCY] the paged attention specialization lives in the CUDA sources, which are +// also built as ROCm and MUSA. Every other backend ignores src[5] and walks the pool in physical +// cell order, so a KV layer placed there would silently lose the mode's guarantee. static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) { if (!dev) { return false; @@ -85,14 +83,11 @@ static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) { return strcmp(name, "CUDA") == 0 || strcmp(name, "ROCm") == 0 || strcmp(name, "MUSA") == 0; } -// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a -// layer of this shape. The registry name says which backends carry the kernels; it does not -// say the build has them (FLASH_ATTN_AVAILABLE), nor that the device's architecture, the -// head width and the K/V types land on a kernel. Where they do not, the scheduler would hand -// the op to the CPU, which accepts the page table as the reference for test-backend-ops and -// ignores it, and the mode would report itself on while attending in physical order. So the -// op is built the way the graph builds it, at the widths a decode step, a verify step and a -// prompt chunk use, and the device is asked. +// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a layer +// of this shape. The registry name only says which backends carry the kernels, not that the build +// has them or that this architecture, head width and K/V types land on one; otherwise the op +// falls to the CPU, which ignores the page table. So build the op as the graph does, at the +// widths a decode step, a verify step and a prompt chunk use, and ask the device. static bool llama_dev_supports_paged_attn( ggml_backend_dev_t dev, ggml_type type_k, ggml_type type_v, @@ -163,9 +158,8 @@ llama_kv_cache::llama_kv_cache( v_cells_impl(other ? other->v_cells_impl : std::make_shared()), v_cells(*v_cells_impl) { - // [TAG_EXACT_CONCURRENCY] read the knob through the one cached reader that the graph and the - // CUDA dispatcher also use, so a process that sets it between two context creations cannot end - // up with a paged cache on top of a dispatcher that is still in default mode + // [TAG_EXACT_CONCURRENCY] read the knob through the same cached reader the graph and the CUDA + // dispatcher use, so a mid-process change cannot leave the two disagreeing exact_pages = llama_exact_concurrency(); // shared cells view the source cache's K/V tensors, so the cell count @@ -181,9 +175,8 @@ llama_kv_cache::llama_kv_cache( GGML_ASSERT(kv_size % n_pad == 0); - // [TAG_EXACT_CONCURRENCY] - // Every one of these is reachable from the command line, so report which one failed by name - // instead of aborting on a bare assert that only prints a file and a line. + // [TAG_EXACT_CONCURRENCY] all of these are reachable from the command line, so name the one + // that failed instead of aborting on a bare assert if (exact_pages) { const char * unsupported = nullptr; @@ -328,10 +321,8 @@ llama_kv_cache::llama_kv_cache( LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); - // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order while - // the mode still reports itself as on, so refuse the load instead - // [TAG_EXACT_CONCURRENCY] the paged attention kernel handles 256-wide K and V heads only; - // any other width would run unpaged on the CPU while the mode reports itself as on + // [TAG_EXACT_CONCURRENCY] the paged kernel handles 256-wide K and V heads only; any other + // width would run unpaged while the mode reports itself as on if (exact_pages && (hparams.n_embd_head_k(il) != 256 || (!is_mla && hparams.n_embd_head_v(il) != 256) || is_mla)) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d has %u-wide K heads and %u-wide V heads%s, " "and the paged attention kernel supports 256-wide K and V heads only\n", @@ -339,14 +330,14 @@ llama_kv_cache::llama_kv_cache( throw std::runtime_error("exact concurrency: unsupported attention head size"); } - // [TAG_EXACT_CONCURRENCY] the paged attention kernel has no soft-capped variant and would - // assert on its first call, so a soft-capped model is refused at load instead + // [TAG_EXACT_CONCURRENCY] the paged kernel has no soft-capped variant and would assert if (exact_pages && hparams.attn_soft_cap) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but this model soft-caps its attention logits (%.1f), " "which the paged attention kernel does not apply\n", __func__, hparams.f_attn_logit_softcapping); throw std::runtime_error("exact concurrency: attention soft cap is not supported"); } + // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, " "which has no paged attention: every layer must be offloaded to the CUDA backend " @@ -355,8 +346,8 @@ llama_kv_cache::llama_kv_cache( throw std::runtime_error("exact concurrency: KV cache layer is not on the CUDA backend"); } - // [TAG_EXACT_CONCURRENCY] the backend is the right one; ask it whether this layer's - // attention, with the page table attached, lands on one of its kernels at all + // [TAG_EXACT_CONCURRENCY] right backend; ask whether this layer's attention, with the + // page table attached, lands on one of its kernels at all if (exact_pages && !llama_dev_supports_paged_attn(model.dev_layer(il), type_k, type_v, hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), hparams.n_head(il), hparams.n_head_kv(il), kv_size, exact_page_size)) { @@ -671,12 +662,9 @@ void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, ll return; } - // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so cells cannot be shared - // between two of them. Refuse the operation rather than abort the process: a server - // rejects the request that would reach here (n_cmpl > 1), and any caller this does not - // cover degrades to a failed copy it can report instead of killing every other request - // on the machine. Placed after the shared-cells return so a draft cache, which copies - // nothing of its own, is unaffected. + // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so cells cannot be shared between + // two. Refuse rather than abort the process, so an uncovered caller gets a failed copy it can + // report. After the shared-cells return, so a draft cache is unaffected. if (exact_pages && seq_id_src != seq_id_dst) { LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between " "sequences (%d -> %d); ignoring the copy\n", @@ -806,9 +794,8 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll return; } - // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo the - // page size, so shifting positions would put every cell of the sequence in the wrong - // place. Context shift is unsupported in exact mode; say so rather than abort. + // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo the page size, + // so shifting positions would misplace every cell; say so rather than abort if (exact_pages && shift != 0) { LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions " "(seq %d, shift %d); ignoring the shift\n", @@ -866,8 +853,7 @@ void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, in return; } - // [TAG_EXACT_CONCURRENCY] same reason as seq_add: dividing positions breaks the - // identity between a cell's position and its offset inside its page. + // [TAG_EXACT_CONCURRENCY] as in seq_add: dividing positions breaks the position/offset identity if (exact_pages && d != 1) { LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions " "(seq %d, d %d); ignoring the division\n", @@ -970,10 +956,9 @@ llama_memory_context_ptr llama_kv_cache::init_batch( std::vector ubatches; while (true) { - // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt tokens into one - // ubatch, so a sequence's prefill would run at a width its solo run never sees. Take - // the sequence-set split instead, which can give each prompt a ubatch of its own; a - // plain decode step has nothing to isolate and keeps taking split_simple. + // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt into one ubatch, + // so a prefill would run at a width its solo run never sees. The sequence-set split + // gives each prompt its own ubatch; a plain decode step keeps taking split_simple. const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0; auto ubatch = n_stream == 1 && !isolate @@ -1026,8 +1011,8 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector v_cells; // copy of the old cells, before placing the ubatch - // [TAG_EXACT_CONCURRENCY] page ownership before placing the ubatch, so that undoing the - // speculative placement does not force a rebuild from every cell on the next ubatch + // [TAG_EXACT_CONCURRENCY] page ownership before the ubatch, so undoing a speculative + // placement does not force a rebuild from every cell std::vector exact_page_owner_old; }; @@ -1078,9 +1063,8 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vectorv_heads_old[s]; } - // [TAG_EXACT_CONCURRENCY] the speculative placements are being undone behind the - // allocator's back. Put back what it knew before, unless something during the placement - // removed cells as well, in which case only the cells can say what is left. + // [TAG_EXACT_CONCURRENCY] put back what the allocator knew, unless the placement also + // removed cells, in which case only the cells can say what is left if (!exact_page_owner_dirty) { exact_page_owner = it->exact_page_owner_old; } @@ -1243,10 +1227,8 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, } if (exact_pages) { - // Page ownership is maintained as cells are placed and invalidated when they are removed, - // so the allocator reads one entry per physical page rather than scanning every cell. The - // claims this call makes are local: prepare() can still roll back its speculative - // placements, and empty pages stay immediately reusable. + // ownership is maintained as cells are placed, so this reads one entry per page rather + // than scanning every cell; the claims are local and prepare() can still roll them back const auto & cells = v_cells[0]; exact_pages_sync(); @@ -1269,7 +1251,7 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, const page_key key {ubatch.seq_id[i][0], ubatch.pos[i]/exact_page_size}; auto it = pages.find(key); if (it == pages.end()) { - // Round-robin free-page search deliberately permits nonmonotonic physical order. + // round-robin free-page search, deliberately nonmonotonic in physical order uint32_t page = v_heads[0]/exact_page_size; uint32_t tested = 0; while (tested < owner.size() && owner[page%owner.size()].seq >= 0) { ++page; ++tested; } @@ -1503,17 +1485,15 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & } uint32_t llama_kv_cache::alloc_granularity() const { - // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so a - // sequence holding n tokens holds round_up(n, exact_page_size) cells: its tail page is - // charged in full whether or not it is full. + // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so n + // tokens hold round_up(n, exact_page_size) cells: the tail page is charged in full return exact_pages ? exact_page_size : 1; } bool llama_kv_cache::get_can_shift() const { - // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so the - // paged pool cannot shift positions. Reporting it here is what makes the server disable - // --context-shift and --cache-reuse at load, with a warning, instead of accepting both and - // failing on the first request that needs them. + // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo 256, so the pool + // cannot shift positions. Reporting it here is what disables --context-shift and + // --cache-reuse at load rather than failing on the first request that needs them. if (exact_pages) { return false; } @@ -1605,7 +1585,7 @@ void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * uba auto * row = data.data() + i*dst->ne[0]; row[0] = 0; for (const auto & page : pages[ubatch->seq_id[i][0]]) { - // Exclude wholly future pages even when prefill includes later query rows. + // exclude wholly future pages even when prefill includes later query rows if (page.first*exact_page_size > uint32_t(ubatch->pos[i])) { break; } row[++row[0]] = page.second; } @@ -1622,8 +1602,8 @@ void llama_kv_cache_context::set_input_pages(ggml_tensor * dst, const llama_ubat } uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { - // The physical view spans the pool. The page map, independently padded per query, - // is the only loop bound for exact attention; neighbours cannot extend that loop. + // the physical view spans the pool; the per-query page map is the only loop bound for exact + // attention, so neighbours cannot extend it if (exact_pages) { return get_size(); } uint32_t result = 0; @@ -2429,9 +2409,9 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla } void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { - // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical - // index, which the paged pool owns. Refused before a byte is read, so that the failure - // path below, which clears the cache, is never entered for it. + // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical index, + // which the paged pool owns. Refused before a byte is read, so the clearing failure path below + // is never entered for it. if (exact_pages && seq_id == -1) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state restore only\n", __func__); throw std::runtime_error("whole-cache restore is not supported with LLAMA_EXACT_CONCURRENCY"); diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index af3a04be39d..c139ede57e0 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -243,11 +243,10 @@ class llama_kv_cache : public llama_memory_i { static constexpr uint32_t exact_page_size = 256; bool exact_pages = false; - // [TAG_EXACT_CONCURRENCY] - // Which (sequence, logical page) owns each physical page of the pool; seq < 0 means the page is - // free. Kept current as cells are placed, and marked dirty by the paths that remove cells, so - // that find_slot() and set_input_pages() read one entry per page instead of rebuilding the map - // from every live cell twice per ubatch. Mutable because set_input_pages() is const. + // [TAG_EXACT_CONCURRENCY] which (sequence, logical page) owns each physical page; seq < 0 means + // free. Kept current as cells are placed and marked dirty by removals, so find_slot() and + // set_input_pages() read one entry per page instead of rebuilding from every live cell twice per + // ubatch. Mutable because set_input_pages() is const. struct exact_page { llama_seq_id seq = -1; llama_pos lpg = -1; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index e8d80c770ba..d9e609fee20 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -93,13 +93,10 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // so that the rollback snapshots remain valid const uint32_t n_rs_seq = mem_recr->n_rs_seq; - // [TAG_EXACT_CONCURRENCY] the recurrent half of a hybrid model is not invariant to - // the shape of the ubatch: a prompt processed next to other sequences' prompt tokens - // leaves a different gated delta net state than the same prompt processed alone. - // Giving such a sequence a ubatch of its own removes that. A plain decode step, one - // token per sequence, is already exact and stays batched. - // The figure is passed whenever the mode is on, not only when a prompt is present: - // it also keeps sets of unequal token counts apart (see llama_batch_allocr::split_equal). + // [TAG_EXACT_CONCURRENCY] the recurrent half is not invariant to the ubatch shape: + // a prompt processed next to other prompts leaves a different gated delta net state, + // so it gets a ubatch of its own while plain decode steps stay batched. Passed + // whenever the mode is on, since it also keeps sets of unequal token counts apart. const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); @@ -152,8 +149,8 @@ bool llama_memory_hybrid::get_can_shift() const { } uint32_t llama_memory_hybrid::alloc_granularity() const { - // the recurrent half holds one state per sequence rather than per token, so the - // attention half is the one whose cells a caller is planning capacity for + // the recurrent half holds one state per sequence, so the attention half is the one whose + // cells a caller is planning capacity for return mem_attn->alloc_granularity(); } @@ -172,8 +169,8 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 } void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { - // [TAG_EXACT_CONCURRENCY] the attention half refuses this under exact mode; refuse it here - // before either half is touched, so the two halves cannot end up describing different states + // [TAG_EXACT_CONCURRENCY] the attention half refuses this, so refuse before either half is + // touched or the two could end up describing different states if (llama_exact_concurrency() && seq_id_src != seq_id_dst) { LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between sequences (%d -> %d); ignoring the copy\n", __func__, seq_id_src, seq_id_dst); diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 61463c72964..dadc0661344 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -431,12 +431,10 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & // [TAG_RECURRENT_ROLLBACK_SPLITS] // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch // so that the rollback snapshots remain valid - // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: a recurrent state that a - // prompt leaves behind depends on what shared its ubatch, so isolate the prompts. - // The figure is passed whenever the mode is on, not only when a prompt is present: - // it also keeps sets of unequal token counts apart, and a three-token verify step - // placed beside a two-token one as two now and one later would be reduced in - // chunks the solo run never had. + // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: the state a prompt leaves + // behind depends on what shared its ubatch, so isolate prompts. Passed whenever the + // mode is on, since it also keeps sets of unequal token counts apart, which would + // otherwise be reduced in chunks the solo run never had. const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); diff --git a/src/llama-memory.h b/src/llama-memory.h index 51539a03919..c0bc2ad7e30 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -100,13 +100,10 @@ struct llama_memory_i { // getters virtual bool get_can_shift() const = 0; - // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit. - // - // 1 for every module that allocates a cell per token, which is all of them unless a mode - // is on that allocates in larger blocks. Where it is larger, a sequence of n tokens - // occupies round_up(n, granularity) cells, and a caller that plans pool capacity by - // counting tokens will believe there is room that does not exist. Not pure, so a module - // that has never heard of this inherits the answer that has always been true of it. + // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit: 1 unless a mode + // that allocates in larger blocks is on, and then n tokens occupy round_up(n, granularity) + // cells, so a caller planning capacity in tokens would see room that does not exist. Not pure, + // so a module that has never heard of this inherits the answer that was always true of it. virtual uint32_t alloc_granularity() const { return 1; } // diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 685193be77d..3f08a060ef1 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7193,8 +7193,8 @@ struct test_flash_attn_ext : public test_case { } }; -// Same mathematical attention as the CPU mask reference, but visit nonadjacent pages -// in a different order. Covers a partial tail and different page counts per query. +// same attention as the CPU mask reference, but visiting nonadjacent pages in a different order; +// covers a partial tail and different page counts per query struct test_flash_attn_ext_pages : public test_flash_attn_ext { test_flash_attn_ext_pages(int64_t batch) : test_flash_attn_ext(256, 256, 2, {8, 1}, 1024, batch) {} @@ -9205,7 +9205,7 @@ static std::vector> make_test_cases_eval() { } } - // Shared weights over sequence planes, as in a recurrent-model output projection. + // shared weights over sequence planes, as in a recurrent-model output projection for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0}) { for (int n : {1, 17, 307}) { test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {4, 1})); @@ -9213,10 +9213,9 @@ static std::vector> make_test_cases_eval() { } } - // Mixture-of-experts projections at the token counts a decode ubatch forms. The gate and up - // projections broadcast one activation row over the expert list, the down projection carries - // one row per expert, and the mixed quantization of a real MoE gguf puts different types on - // the two. 17 tokens is past the width the exact-concurrency policy pins. + // MoE projections at the token counts a decode ubatch forms: gate and up broadcast one + // activation row over the expert list, down carries one row per expert, and a real MoE gguf + // puts different types on the two. 17 tokens is past the width exact concurrency pins. for (ggml_type type_a : {GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16}) { for (int n : {1, 2, 4, 8, 17}) { test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, true, 512, n, 2048)); diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp index 428a9252981..24b6359b05f 100644 --- a/tests/test-state-restore-fragmented.cpp +++ b/tests/test-state-restore-fragmented.cpp @@ -73,8 +73,8 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : saved seq 1 state, %zu bytes\n", __func__, ncopy); - // A fragmented restore may stage a whole device tensor. Check every - // sequence byte-for-byte, including the neighbours that must be preserved. + // a fragmented restore may stage a whole device tensor, so check every sequence + // byte-for-byte, including the neighbours that must be preserved std::vector> before(params.n_parallel); for (int s = 0; s < params.n_parallel; ++s) { before[s].resize(llama_state_seq_get_size(ctx, s)); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index ae47bec194c..6fe51f8cf1a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -38,10 +38,8 @@ constexpr int HTTP_POLLING_SECONDS = 1; -// [TAG_EXACT_CONCURRENCY] the knob is read from the environment by the KV cache, the batch -// splitter and the CUDA backend independently, because it has to be answered before a -// context exists. The server needs the same answer to refuse the one request shape the mode -// cannot serve, so it reads it the same way rather than growing a public API for it. +// [TAG_EXACT_CONCURRENCY] read from the env like the KV cache, batch splitter and CUDA +// backend do: the answer is needed before a context exists. static bool server_exact_concurrency() { static const bool enabled = []() { const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); @@ -78,28 +76,18 @@ enum slot_state { // [TAG_PREEMPT] server-side request preemption // -// With --kv-unified the cells are one pool shared by every slot, and each slot believes it -// has all of them. When the pool fills, llama_decode returns 1, the retry ladder halves -// n_batch down to 1, and the server ends EVERY conversation in flight with "Context size -// has been exceeded" -- including the ones nowhere near their own limit. Upstream marks the -// spot in decode(): "TODO: try to terminate only the largest active slot/sequence and -// continue with the rest". -// -// Nothing is terminated here. The cells of one slot are taken back and given to it again -// later: its sequence is copied to host RAM, its cells are released, and when the pool has -// room the copy goes back and the slot carries on with the same sampler, the same generated -// text and the same open stream. A streaming client sees a pause, not an error. +// With --kv-unified one full pool ends EVERY conversation in flight with "Context size has +// been exceeded", including the ones nowhere near their own limit. Instead of terminating, +// one slot's sequence is copied to host RAM and its cells released; when there is room the +// copy goes back and the slot carries on with the same sampler, text and open stream, so a +// streaming client sees a pause rather than an error. constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected -// [TAG_PREEMPT] The order parked slots come back in. Head of the line by park time, and nobody -// passes a head that does not fit yet: the head keeps the room the pool frees until it fits, so -// its wait is bounded by the slots ahead of it and not by how often a smaller slot can squeeze -// in, grow, and be parked again. Simulated over 60 seeds at eight chats this cuts the longest -// single wait by 2.5 to 3x for 0 to 3 percent of makespan at 8192 cells, and parks less often. -// LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order: most-preempted first, then longest -// parked, and a smaller slot may pass a head that does not fit. -// LLAMA_SERVER_PREEMPT_RESUME=head (the default) or pass; read once in load_model() and logged. +// [TAG_PREEMPT] the order parked slots come back in. Default (LLAMA_SERVER_PREEMPT_RESUME=head, +// read once in load_model()): head of the line by park time, and it keeps the room the pool +// frees until it fits, bounding its wait by the slots ahead of it. =pass keeps the previous +// order, most-preempted then longest parked, where a smaller slot may pass a head. static bool g_preempt_resume_head_of_line = true; static bool preempt_resume_head_of_line() { @@ -109,30 +97,22 @@ constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked constexpr int64_t PREEMPT_ROTATE_US = 2ll * 1000 * 1000; // a resident cycling through context shifts gives way to a parked head that has waited this long -// [TAG_EXACT_CONCURRENCY] The planner above counts cells, not tokens, because the two are not -// the same number under every mode. llama_memory_alloc_granularity() reports how many cells the -// pool hands out at a time: 1 in every ordinary configuration, and the exact concurrency page -// size when that mode is on, where one page belongs to one (sequence, position / page) pair and -// a sequence of n tokens therefore occupies round_up(n, page) cells. Four sequences can be -// holding up to 4 * (page - 1) cells that nobody else can be given, and a planner counting -// tokens sees room in the pool that find_slot cannot find in pages: it never reaches the -// threshold that would park anybody, the retry ladder halves n_batch to 1, and every request -// ends in the context error that preemption exists to remove. +// [TAG_EXACT_CONCURRENCY] the planner counts cells, not tokens: llama_memory_alloc_granularity() +// is 1 ordinarily but the page size under exact concurrency, where a sequence of n tokens +// occupies round_up(n, page) cells. A planner counting tokens would see room find_slot cannot +// find in pages, never park anybody, and end every request in the context error instead. // cells a run of n_tokens occupies when the pool allocates g at a time static constexpr int32_t preempt_n_cells_g(int32_t n_tokens, int32_t g) { return (g <= 1 || n_tokens <= 0) ? n_tokens : ((n_tokens + g - 1) / g) * g; } -// cells a run of n_tokens has to be given for a step of n_step more: nothing until the step -// crosses a page boundary, a whole page when it does +// cells a step of n_step more costs: nothing until it crosses a page boundary, a page when it does static constexpr int32_t preempt_n_cells_step_g(int32_t n_tokens, int32_t n_step, int32_t g) { return preempt_n_cells_g(n_tokens + n_step, g) - preempt_n_cells_g(n_tokens, g); } -// At a granularity of 1 both are the identity, so every figure the planner computes is exactly -// the arithmetic it did before it started asking the memory how it allocates, and nothing -// changes in any configuration that does not page. +// at a granularity of 1 both are the identity, so nothing changes in a configuration that does not page static_assert(preempt_n_cells_g(0, 1) == 0 && preempt_n_cells_g(1, 1) == 1 && preempt_n_cells_g(8191, 1) == 8191 && preempt_n_cells_g(-3, 1) == -3, "at a granularity of 1 a run of n tokens has to cost exactly n cells"); @@ -140,7 +120,7 @@ static_assert(preempt_n_cells_step_g(0, 1, 1) == 1 && preempt_n_cells_step_g(819 preempt_n_cells_step_g(1000, 512, 1) == 512, "at a granularity of 1 a step of n tokens has to cost exactly n cells"); -// and the page arithmetic itself, so the rounding cannot be changed by accident +// and the page arithmetic itself, so the rounding cannot change by accident static_assert(preempt_n_cells_g(1, 256) == 256 && preempt_n_cells_g(256, 256) == 256 && preempt_n_cells_g(257, 256) == 512, "a tail page is charged in full"); @@ -380,16 +360,13 @@ struct server_slot { prompt.clear(); } - // [TAG_PREEMPT] state of a slot whose cells were taken back - // - // Only the KV cells leave. The task, the sampler, the generated text and the position - // the stream has reached stay on the slot, so a resume is a memcpy and not a new - // request: no retokenisation, no replayed prompt, no seam in the output. + // [TAG_PREEMPT] state of a slot whose cells were taken back. Only the KV cells leave; + // the task, sampler, generated text and stream position stay, so a resume is a memcpy. slot_state state_before_preempt = SLOT_STATE_IDLE; std::vector preempt_state_tgt; std::vector preempt_state_dft; int32_t n_preempt = 0; // times the CURRENT task has been preempted - int32_t n_ctx_shift = 0; // context shifts the CURRENT task has made: it is at the pool's limit and cycling + int32_t n_ctx_shift = 0; // context shifts it has made: it is at the pool's limit and cycling int32_t n_preempt_fail = 0; // consecutive failed restores int64_t t_preempt_us = 0; // when it was parked @@ -438,10 +415,8 @@ struct server_slot { return false; } - // The draft is a prediction, not a result, so it goes with the cells. Preemption - // runs before the batch is built, so spec_i_batch is empty and prompt.tokens already - // holds exactly the tokens the state above covers -- including the rollback done by - // the checkpoint path when a draft was only partially accepted. + // the draft is a prediction, not a result, so it goes with the cells; prompt.tokens + // already covers exactly what the state above holds spec_draft.clear(); spec_i_batch.clear(); spec_ckpt.clear(); @@ -449,8 +424,7 @@ struct server_slot { i_batch = -1; - // note: prompt.tokens is deliberately kept. It is the mirror of the state just - // copied out, and the resume needs it to know how many cells to ask for. + // note: prompt.tokens is deliberately kept - the resume sizes its request from it mem.seq_rm(id, -1, -1); state_before_preempt = state; @@ -487,10 +461,8 @@ struct server_slot { state = state_before_preempt; - // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks - // that the draft context is where it should be, which the restore above ensures. - // A slot parked while still processing its prompt makes that transition itself - // once the prompt is done. + // same call the DONE_PROMPT -> GENERATING transition makes; a slot parked mid-prompt + // makes that transition itself once the prompt is done if (state == SLOT_STATE_GENERATING && can_speculate()) { common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); } @@ -498,12 +470,9 @@ struct server_slot { return true; } - // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds for this sequence. - // For a batch that is given up after it was built: the tokens added for this slot - // that were never decoded come off, the sampled token stays in `sampled` and goes into - // the next batch the way it went into this one, and a draft is a prediction that goes - // with them. A chunk that failed to decode left nothing in the cache, so the cache is - // the boundary. + // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds, for a batch given up + // after it was built: never-decoded tokens and the draft come off, `sampled` is kept for + // the next batch. A failed chunk left nothing in the cache, so the cache is the boundary. void rewind_to_cache() { const int32_t n_cached = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), id) + 1; @@ -511,8 +480,7 @@ struct server_slot { prompt.tokens.keep_first(n_cached); } - // a prompt whose last chunk was in the batch was marked done when the chunk was - // built; the chunk never ran, so the prompt is not done + // the last chunk was marked done when it was built but never ran, so it is not done if (state == SLOT_STATE_DONE_PROMPT && task && prompt.n_tokens() < task->n_tokens()) { state = SLOT_STATE_PROCESSING_PROMPT; } @@ -743,9 +711,8 @@ struct server_slot { t_last_used = ggml_time_us(); - // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be - // released while parked), so the mirror of them must not outlive them: the next - // task on this slot would otherwise take a prefix match against an empty cache + // [TAG_PREEMPT] the cells are already gone, so the mirror of them must not outlive + // them: the next task would take a prefix match against an empty cache if (state == SLOT_STATE_PREEMPTED) { preempt_state_free(); prompt_clear(); @@ -1499,15 +1466,13 @@ struct server_context_impl { } } - // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates, rather than assume a cell per - // token. 1 in every ordinary configuration, so this changes nothing unless a mode that - // places cells in blocks is on. + // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates rather than assume a cell per + // token; 1 unless a mode that places cells in blocks is on { preempt_alloc_granularity = (int32_t) std::max(1u, llama_memory_alloc_granularity(llama_get_memory(ctx_tgt))); - // a test knob: the paged attention kernel only supports a head size of 256, so a - // harness model cannot turn exact concurrency on, and this is the only way to reach - // the paged arithmetic of the planner from the server tests + // test knob: the paged kernel needs a head size of 256, so a harness model cannot + // turn exact concurrency on and this is the only way to reach the paged arithmetic const char * LLAMA_SERVER_PREEMPT_GRANULARITY = getenv("LLAMA_SERVER_PREEMPT_GRANULARITY"); if (LLAMA_SERVER_PREEMPT_GRANULARITY) { @@ -1537,9 +1502,8 @@ struct server_context_impl { preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; // LLAMA_SERVER_PREEMPT_POLICY: which non-leader the planner parks, for comparing - // policies against each other on the same workload. smallest (the default and the - // shipped one), largest, youngest (the most recent task, as vLLM's scheduler - // preempts), oldest. The leader is kept and the starvation guard applies under all. + // policies on the same workload: smallest (default), largest, youngest, oldest. + // The leader is kept and the starvation guard applies under all. const char * LLAMA_SERVER_PREEMPT_POLICY = getenv("LLAMA_SERVER_PREEMPT_POLICY"); preempt_test_policy = LLAMA_SERVER_PREEMPT_POLICY ? LLAMA_SERVER_PREEMPT_POLICY : "smallest"; @@ -2954,9 +2918,8 @@ struct server_context_impl { void abort_all_slots(const std::string & reason) { for (auto & slot : slots) { - // [TAG_PREEMPT] a parked slot took no part in what failed: its sequence is in - // host RAM, not in the cache, and it comes back when there is room, the same as - // in the decode error sweep + // [TAG_PREEMPT] a parked slot took no part in what failed and comes back when + // there is room, the same as in the decode error sweep if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { send_error(slot, reason, ERROR_TYPE_SERVER); slot.release(); @@ -3000,18 +2963,14 @@ struct server_context_impl { // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens, - // whether or not the pool is under pressure. It exists to answer the only question that - // matters about a resume: with one request on an idle server the batch has the same - // shape at every step, so a preempted continuation that is not byte-identical to an - // uninterrupted one is the preemption's fault and nothing else's. + // under pressure or not: on an idle server the batch shape is fixed, so a continuation + // that is not byte-identical to an uninterrupted one is the preemption's fault. int32_t preempt_test_every = 0; std::string preempt_test_policy = "smallest"; // LLAMA_SERVER_PREEMPT_POLICY, see load_model - // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load from the - // memory itself: 1 in every ordinary configuration, the page size under exact concurrency. - // Everything below plans in cells because of it. LLAMA_SERVER_PREEMPT_GRANULARITY overrides - // it, which is how the harness reaches the paged arithmetic on a model whose head size the - // paged attention kernel does not support. + // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load: 1 + // ordinarily, the page size under exact concurrency. Everything below plans in cells + // because of it. LLAMA_SERVER_PREEMPT_GRANULARITY overrides it for the harness. int32_t preempt_alloc_granularity = 1; // cells a slot holding n_tokens actually occupies @@ -3024,16 +2983,14 @@ struct server_context_impl { return preempt_n_cells_step_g(n_tokens, n_step, preempt_alloc_granularity); } - // Cells kept spare on top of the reservation. A step that crosses a page boundary costs a - // whole page rather than a cell, so a margin of a few cells is no margin at all under a page - // allocator: round it up to one page. With a granularity of 1 this is PREEMPT_N_MARGIN. + // cells kept spare on top of the reservation, rounded up to a page since a boundary + // crossing costs a whole one; PREEMPT_N_MARGIN at a granularity of 1 int32_t preempt_n_margin() const { return preempt_n_cells(PREEMPT_N_MARGIN); } - // env: LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, so - // the KV-full retry ladder and its last resort are the only thing between a full pool - // and the context error + // LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, leaving + // only the KV-full retry ladder and its last resort bool preempt_planner_off = false; // set by preempt_last_resort(): the batch being decoded was given up, stop the chunk loop @@ -3043,8 +3000,8 @@ struct server_context_impl { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } - // draft tokens this slot's next step can actually carry: the configured maximum, cut to - // what its context and its prediction budget leave, the way get_n_draft_max() cuts it + // draft tokens the next step can carry: the maximum cut to what context and prediction + // budget leave, the way get_n_draft_max() cuts it int32_t preempt_n_spec(const server_slot & slot) const { int32_t res = preempt_n_spec_max(); @@ -3083,10 +3040,8 @@ struct server_context_impl { return preempt_ram_used() + slot.preempt_state_required() <= budget; } - // the same for a rotation: the parked head is restored on the pass that parks the - // resident, so its bytes are on their way out and are not held against the resident. - // A budget that holds one sequence but not two would otherwise refuse every rotation - // and leave the head parked for as long as the resident cares to generate. + // the same for a rotation: the head is restored on the pass that parks the resident, so + // its bytes are on their way out and a one-sequence budget still allows the swap bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) const { if (params_base.preempt_ram_mib < 0) { return true; @@ -3106,9 +3061,8 @@ struct server_context_impl { if (slot.state_before_preempt == SLOT_STATE_GENERATING) { res += 1 + preempt_n_spec(slot); } else { - // a slot just given a task still mirrors the previous request's prompt; the batch - // builder keeps the prefix the two share and drops the rest, so what it holds and - // what it is about to ask for both count from that prefix, not from the old prompt + // a slot just given a task still mirrors the previous prompt; the batch builder + // keeps only the shared prefix, so count from that prefix if (slot.state == SLOT_STATE_STARTED && slot.task) { res = (int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens); } @@ -3119,21 +3073,17 @@ struct server_context_impl { } // [TAG_EXACT_CONCURRENCY] a restore takes fresh pages and its tail page is charged in - // full, so what the pool has to have free for this slot is the rounded figure. Under - // counting here is what admits a resume that find_slot then cannot satisfy. + // full; undercounting here admits a resume that find_slot cannot satisfy return preempt_n_cells(res); } - // Cells the pool is holding right now. A released slot keeps its prompt in the cache - // for the next request to reuse as a prefix, so idle slots count too: the first version - // of this counted only the running ones, decided a pool holding 8185 cached cells was - // empty, and every resume failed against a cache that was actually full. + // cells the pool is holding right now. A released slot keeps its prompt in the cache as a + // prefix for the next request, so idle slots count too or a full pool looks empty. int32_t preempt_kv_used() const { int32_t res = 0; - // n_cmpl > 1: the parent and its children share the prompt's cells through seq_cp, so - // the prompt is charged once per family, to whichever resident member comes first; - // the others are charged only what they generated on top of it + // n_cmpl > 1: a family shares the prompt's cells through seq_cp, so the prompt is + // charged once, to the first resident member; the others only for what they added std::vector charged; for (const auto & slot : slots) { @@ -3141,12 +3091,11 @@ struct server_context_impl { continue; // parked: its cells are in host RAM, not in the pool } - // [TAG_EXACT_CONCURRENCY] the slot's tail page is charged in full: it belongs to - // this sequence and cannot be given to anybody else, however little of it is used + // [TAG_EXACT_CONCURRENCY] the tail page is charged in full: it cannot be given to + // anybody else, however little of it is used - // a child waiting for its parent's prompt does not share anything yet: until - // copy_state_to() runs it still holds whatever the previous request left in its - // cells, so it is charged that on its own, outside the family + // a child waiting for its parent shares nothing until copy_state_to() runs, so it + // is charged its own stale cells, outside the family if (slot.state == SLOT_STATE_WAIT_OTHER) { res += preempt_n_cells(slot.prompt.n_tokens()); continue; @@ -3163,9 +3112,8 @@ struct server_context_impl { charged.push_back(family); } - // a slot just given a task still mirrors the previous request's prompt until the - // batch builder keeps the prefix the two share and drops the rest; what stays is - // the prefix, so that is what the pool holds for it + // a slot just given a task keeps only the prefix it shares with the new prompt, + // so that is what the pool holds for it if (slot.state == SLOT_STATE_STARTED && slot.task) { res += preempt_n_cells((int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens)); continue; @@ -3185,11 +3133,8 @@ struct server_context_impl { int32_t res_pmt = 0; // [TAG_EXACT_CONCURRENCY] each slot reserves the cells its next step ADDS, not the - // tokens it adds. preempt_kv_used() already charges every slot's tail page in full, so - // with a granularity of 1 these are the same number and nothing changes; with a larger - // one the step is free until it crosses a page boundary and costs a whole page when it - // does. Reserving tokens on top of a rounded used figure would miss exactly that - // crossing, which is the only moment the pool can actually run out. + // tokens: preempt_kv_used() already rounds up the tail page, so reserving tokens on + // top of it would miss the boundary crossing, the only moment the pool can run out for (const auto & slot : slots) { const int32_t n_cur = slot.prompt.n_tokens(); @@ -3211,11 +3156,9 @@ struct server_context_impl { } } - // one batch is all the prompt slots get between them, however many are waiting; in - // cells that batch can straddle one boundary more than it has tokens for // one batch is all the prompt slots get between them, however many are waiting; under - // page allocation each of them can still cross a page boundary of its own within that - // batch, so the cap keeps one boundary per prompt slot on top of the batch + // page allocation each can still cross a boundary of its own, so the cap allows one + // boundary per prompt slot on top of the batch int32_t n_pmt = 0; for (const auto & slot : slots) { @@ -3227,16 +3170,10 @@ struct server_context_impl { return res + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1)); } - // Keep the slot that is furthest along -- it is the closest to finishing and to giving - // its cells back -- and among the rest prefer one that has not been preempted - // PREEMPT_N_STARVED times already, then the smallest. - // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt - // until the batch builder keeps the prefix the two share and drops the rest (see the - // SLOT_STATE_STARTED block of update_slots). Parked as it is, it would be copied out, - // charged and sized by the old prompt, and a short unrelated request could exceed the - // budget or stay parked for room it will never use. Keeping only the shared prefix now - // is what the batch builder does anyway; the chunk reuse it can add on top is given up - // for a slot the planner has to touch, which is rare. + // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt until + // the batch builder drops it (see the SLOT_STATE_STARTED block of update_slots). Parked as + // it is, it would be copied out and sized by the old prompt. Keeping only the shared prefix + // now is what the batch builder does anyway, at the cost of the chunk reuse it can add. void preempt_normalize_started(server_slot & slot) { if (slot.state != SLOT_STATE_STARTED || !slot.task) { return; @@ -3248,10 +3185,8 @@ struct server_context_impl { return; } - // a memory that cannot remove part of a sequence (a recurrent state without rollback - // room for the stale suffix) aborts on a partial removal; for it the whole stale - // sequence goes, and the prompt is processed from the start on resume, as it would be - // without a usable checkpoint + // a memory that cannot remove part of a sequence aborts on a partial removal; for it + // the whole stale sequence goes and the prompt is reprocessed from the start const bool partial_ok = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART && (!ctx_dft || ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART); @@ -3269,10 +3204,8 @@ struct server_context_impl { server_slot * leader = nullptr; int32_t n_running = 0; - // a slot just given a task is measured by the prefix it keeps, not by the previous - // request's prompt it still mirrors: measured by the mirror, a short request over a - // large stale cache would be the never-parked leader while the longest live - // conversation was parked in its place + // measure a just-started slot by the prefix it keeps, not by the stale prompt it + // mirrors, or a short request over a large stale cache becomes the leader for (auto & slot : slots) { preempt_normalize_started(slot); } @@ -3288,19 +3221,16 @@ struct server_context_impl { } if (n_running < 2) { - // a single conversation that does not fit the pool on its own is a real context - // overflow and not a scheduling problem - leave it to the existing error path + // one conversation that does not fit alone is a real overflow, not a scheduling + // problem - leave it to the existing error path return nullptr; } server_slot * victim = nullptr; for (auto & slot : slots) { - // Before the batch is built every one of these is at a token boundary: a - // generating slot between two sampled tokens, a prompt-processing slot between - // two chunks of its prompt, a started slot with only a cached prefix (or - // nothing) in the pool. A slot holding no cells is still worth parking - it - // is about to ask for a whole batch of them. + // before the batch is built every one of these is at a token boundary. A slot + // holding no cells is still worth parking - it is about to ask for a batch. if (slot.state != SLOT_STATE_GENERATING && slot.state != SLOT_STATE_PROCESSING_PROMPT && slot.state != SLOT_STATE_STARTED) { @@ -3332,9 +3262,8 @@ struct server_context_impl { return victim; } - // is a the better victim of the two? the smallest slot under the shipped policy: it - // gives up the least work and its restore is the cheapest (see the PR's simulation); - // the other choices exist for the comparison runs behind LLAMA_SERVER_PREEMPT_POLICY + // is a the better victim? the smallest under the shipped policy, since it gives up the + // least work; the rest exist for comparison runs behind LLAMA_SERVER_PREEMPT_POLICY bool preempt_better_victim(const server_slot & a, const server_slot & b) const { if (preempt_test_policy == "largest") { return a.prompt.n_tokens() > b.prompt.n_tokens(); @@ -3351,10 +3280,8 @@ struct server_context_impl { return a.prompt.n_tokens() < b.prompt.n_tokens(); } - // called once per update_slots(), before the batch is built: at that point every slot is - // at a token boundary, prompt.tokens is exactly what the cache holds for it, and no - // draft is in flight, so a slot can be removed from the picture without unpicking a - // half-decoded batch + // called once per update_slots(), before the batch is built: every slot is then at a token + // boundary with no draft in flight, so one can be removed without unpicking a batch void update_preemption() { if (!params_base.kv_unified || slots.size() < 2) { return; // with a cache per slot, no slot can take another one's cells @@ -3370,10 +3297,7 @@ struct server_context_impl { const int32_t n_cells = n_ctx; - // Put back what fits, in the order preempt_resume_head_of_line() describes: by default - // the slot parked longest, and only that one until it fits; under - // LLAMA_SERVER_PREEMPT_RESUME=pass the most-preempted slot first, then the one parked - // longest, and a smaller slot may pass a head that does not fit. + // put back what fits, in the order preempt_resume_head_of_line() describes const bool head_of_line = preempt_resume_head_of_line(); for (;;) { @@ -3403,12 +3327,9 @@ struct server_context_impl { server_slot * best = nullptr; - // A parked slot whose sequence plus its next step would not fit an empty pool can - // never be restored, and would otherwise sit at the head of the line for ever - // without a restore ever being attempted: a prompt within n_ctx that was parked - // before it took any cells, but too close to n_ctx to leave room for its first - // batch. That is the single-conversation overflow the KV-full path reports, so - // report it the same way and rescan the line without it. + // a parked slot that would not fit an empty pool can never be restored and would + // sit at the head of the line for ever. That is the single-conversation overflow + // the KV-full path reports, so report it the same way and rescan without it. { server_slot * impossible = nullptr; @@ -3428,12 +3349,10 @@ struct server_context_impl { } } - // Room for the sequence AND for the next step of everything already running, - // so that a resume cannot immediately trigger the preemption of someone else. - // The margin is headroom for the others; with nothing resident there is nobody - // to keep it for, so a sequence that fits the pool exactly is let back in. - // A cached prompt on an idle slot is worth less than a conversation waiting to - // continue, so give those cells up first - same call the KV-full path makes. + // room for the sequence AND for the next step of everything running, so a resume + // cannot immediately preempt someone else. The margin is headroom for the others, + // so with nothing resident a sequence that fits exactly is let back in. Cached + // prompts on idle slots are worth less than a waiting conversation, so go first. for (;;) { const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); const int32_t margin = occupied == 0 ? 0 : preempt_n_margin(); @@ -3450,20 +3369,15 @@ struct server_context_impl { } } - // Nothing fits. A resident that has reached the pool's limit and is cycling - // through context shifts holds the room for as long as it likes to generate, - // and the head behind it would wait for ever. After the head has waited its - // turn, that resident is parked in its place: it is at a token boundary like - // any other park, and when it comes back it is the one waiting, so the two - // take turns instead of one taking everything. + // nothing fits. A resident cycling through context shifts holds the room for as + // long as it generates, so once the head has waited its turn that resident is + // parked in its place and the two take turns. if (!best) { server_slot * head = parked.front(); if (ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { - // the resident whose cells let the head in, the smallest of those; failing - // one that does so alone, the largest, since it makes the most room. Taking - // the first shifting resident in slot order could park one too small to - // matter, spend the park budget on it, and leave the head waiting anyway. + // the smallest resident whose cells let the head in; failing one that does + // so alone, the largest, since it makes the most room const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); const int32_t need = preempt_n_need(*head) + PREEMPT_N_MARGIN; @@ -3520,9 +3434,8 @@ struct server_context_impl { const int64_t t_start = ggml_time_us(); if (!best->preempt_restore()) { - // update_slots() runs in a tight loop while tasks are pending, so a counter - // alone burns its whole budget in a couple of milliseconds. Give up only on - // a slot that has been failing for a while, and keep the log quiet. + // update_slots() loops tightly, so a counter alone burns its budget in + // milliseconds: give up only on a slot failing for a while, and log quietly if (best->n_preempt_fail % 64 == 1) { SLT_WRN(*best, "resume failed (%d in a row, parked %.1f s), staying preempted\n", best->n_preempt_fail, (ggml_time_us() - best->t_preempt_us) / 1e6); @@ -3646,8 +3559,7 @@ struct server_context_impl { } } - // [TAG_PREEMPT] make the pool fit the step that is about to be built, measured after - // any context shift + // [TAG_PREEMPT] make the pool fit the step about to be built, measured after any shift pre_decode_shift(); update_preemption(); @@ -3700,8 +3612,7 @@ struct server_context_impl { #endif if (preempt_batch_abandoned) { - // [TAG_PREEMPT] the rest of this batch was never decoded and the slots no - // longer describe it; the next pass builds a new one + // [TAG_PREEMPT] the rest of this batch never ran; the next pass rebuilds it preempt_batch_abandoned = false; break; } @@ -3735,8 +3646,7 @@ struct server_context_impl { // apply context-shift if needed // TODO: simplify and improve - // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift, - // not with the cells the shift is about to give back + // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift void pre_decode_shift() { iterate(slots, [&](server_slot & slot) { if (slot.state == SLOT_STATE_GENERATING && slot.prompt.n_tokens() + 1 >= slot.n_ctx) { @@ -3949,8 +3859,7 @@ struct server_context_impl { return; // batch is full, skip remaining slots } - // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to - // batch; it takes no part in this pass until it is restored + // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to batch if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { return; } @@ -4477,14 +4386,11 @@ struct server_context_impl { } } - // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell. Upstream this is - // the context error for every slot in the batch. With a park budget the batch is given - // up instead: every resident slot is rewound to the token boundary the cache is at (a - // batch is applied one chunk at a time, and the chunk that failed left nothing behind), - // the smallest are parked until the planner's own bound holds again, and the next - // update_slots() rebuilds the batch from the survivors. The planner brings the parked - // ones back as cells free up. A multimodal prompt has no boundary the cache can name, - // so it keeps the old path. + // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell, which upstream is + // the context error for every slot in the batch. With a park budget the batch is given up + // instead: resident slots are rewound to the token boundary the cache is at, the smallest + // are parked until the planner's bound holds, and the next update_slots() rebuilds the + // batch. A multimodal prompt has no boundary the cache can name, so it keeps the old path. bool preempt_last_resort_possible() const { return params_base.kv_unified && params_base.preempt_ram_mib != 0 && slots.size() >= 2 && llama_get_memory(ctx_tgt); } @@ -4620,11 +4526,9 @@ struct server_context_impl { { std::string err; - // [TAG_PREEMPT] with speculation on, a slot's sampled token and its draft have - // to stay in one view: a narrower view splits the group and the verify step - // throws for the slot whose tokens straddle it. Halving is no help there, so - // after the idle slots the ladder goes to its last resort straight away. With - // no budget to park into the ladder is what it always was. + // [TAG_PREEMPT] a slot's sampled token and its draft have to stay in one view, + // so halving would split the group and make the verify step throw: after the + // idle slots the ladder goes straight to its last resort if (ret == 1 && n_batch > 1 && preempt_last_resort_possible() && batch_has_spec_groups()) { if (try_clear_idle_slots()) { SRV_WRN("%s", "failed to find free space in the KV cache, retrying after purging an idle slot\n"); @@ -4661,8 +4565,7 @@ struct server_context_impl { SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret); for (auto & slot : slots) { - // [TAG_PREEMPT] a parked slot has nothing in this batch and nothing in the - // cache; it is not part of this failure and comes back when there is room + // [TAG_PREEMPT] a parked slot is not part of this failure if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { send_error(slot, err); slot.release(); @@ -5270,11 +5173,9 @@ std::unique_ptr server_routes::handle_completions_impl( task.params.oaicompat_cmpl_id = completion_id; task.params.oaicompat_model = meta->model_name; - // [TAG_EXACT_CONCURRENCY] the children of an n_cmpl > 1 task are served by - // copying the parent's cells to another sequence id, and exact mode gives a KV - // page to one sequence, so there is nothing for that copy to land in. Refuse - // the request here, where it becomes a 400 the client can read, rather than - // letting it reach seq_cp with nothing to do. + // [TAG_EXACT_CONCURRENCY] children of an n_cmpl > 1 task are served by copying the + // parent's cells to another sequence id, and exact mode gives a page to a single + // sequence, so refuse here where it becomes a 400 rather than at seq_cp if (task.params.n_cmpl > 1 && server_exact_concurrency()) { throw std::runtime_error( "n > 1 is not supported while LLAMA_EXACT_CONCURRENCY is set: each " diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 7170efea080..238afa98f41 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -5,10 +5,9 @@ import pytest from utils import * -# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked -# (its sequence copied to host RAM, its cells released) instead of every slot being -# terminated. Both tests need more than one slot and --kv-unified, which is the only -# configuration where one slot can take another one's cells. +# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked (its +# sequence copied to host RAM, its cells released) instead of every slot being terminated. Needs +# more than one slot and --kv-unified, the only configuration where slots share cells. server = ServerPreset.tinyllama2() @@ -57,9 +56,8 @@ def _complete(n_predict: int, prompt: str = "Hi how are you"): def test_forced_preemption_does_not_change_the_output(): - # Park and restore the only running slot every 8 tokens. With one request the batch - # has the same shape at every step whether or not the slot was parked in between, so - # any difference in the output is the preemption's fault and nothing else's. + # park and restore the only running slot every 8 tokens: the batch shape is the same at every + # step, so any difference in the output is the preemption's fault global server server.n_ctx = 512 server.start() @@ -86,11 +84,9 @@ def test_forced_preemption_does_not_change_the_output(): def test_two_slots_that_overflow_the_pool_together_both_finish(): - # Each request alone fits in the pool: 8 prompt tokens plus 160 generated is well - # under 256. Together they do not, 336 against 256. Without preemption the retry - # ladder ends with "Context size has been exceeded" on every processing slot; with it - # the smaller slot is parked until the leader finishes and its cells are purged, and - # then it resumes from the token it was parked on. + # each request fits the pool alone (168 of 256 cells) but not together (336). Without + # preemption both end with "Context size has been exceeded"; with it the smaller is parked + # until the leader finishes, then resumes from the token it was parked on. global server server.n_ctx = 256 server.start() @@ -115,15 +111,11 @@ def test_two_slots_that_overflow_the_pool_together_both_finish(): def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages(): - # A pool that hands out cells in blocks gives a whole block to one sequence, so a sequence - # of n tokens occupies round_up(n, block) cells and holds the rest of its tail block against - # everybody else. The planner has to count those cells: counting tokens, it sees room the - # allocator cannot find, never parks anybody, and the retry ladder ends every request. - # - # llama_memory_alloc_granularity() reports the block size, and the only mode that returns - # more than 1 today is exact concurrency, whose paged attention kernel needs a head size this - # model does not have. LLAMA_SERVER_PREEMPT_GRANULARITY injects the figure instead: what is - # under test is the server's arithmetic, which is the same at 64 as at 256. + # a pool that allocates in blocks gives a whole block to one sequence, so n tokens occupy + # round_up(n, block) cells and the planner has to count cells: counting tokens it sees room + # the allocator cannot find, never parks anybody, and the retry ladder ends every request. + # LLAMA_SERVER_PREEMPT_GRANULARITY injects the block size, since the only mode that reports + # one needs a head size this model does not have; the arithmetic is the same at 64 as at 256. global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_GRANULARITY"] = "64" @@ -142,9 +134,8 @@ def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages(): assert "preempted:" in text assert "resumed after" in text - # every figure the planner logs is a whole number of blocks: "kv N/256" is what the pool is - # holding and "(wanted N)" is that plus what the next decode reserves. Counting tokens, both - # land wherever the sequences happen to be. + # every figure the planner logs is a whole number of blocks: "kv N/256" is what the pool holds + # and "(wanted N)" is that plus the next decode's reservation held = [int(n) for n in re.findall(r"kv (\d+)/256", text)] wanted = [int(n) for n in re.findall(r"\(wanted (\d+)\)", text)] assert held and wanted, f"the planner logged no figures:\n{text}" @@ -184,9 +175,8 @@ def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]: def test_two_prompts_that_overflow_the_pool_together_both_finish(): - # Neither slot ever generates before the pool is full: both are still processing their - # prompts. A prompt-processing slot is between two chunks of its prompt, which is as - # clean a boundary as the one between two sampled tokens, so it is parked the same way. + # neither slot generates before the pool is full: a slot between two chunks of its prompt is + # as clean a boundary as one between two sampled tokens, so it is parked the same way global server server.n_ctx = 256 server.start() @@ -215,20 +205,17 @@ def test_two_prompts_that_overflow_the_pool_together_both_finish(): def test_a_generating_slot_and_a_large_prompt_both_finish(): - # One slot is generating a long answer to a short prompt when a large prompt arrives - # beside it. Together they need far more than the pool has. The prompt is admitted - # chunk by chunk, whoever is smaller is parked when the pool fills, and both finish. - # This model produces a thousand tokens a second, so the second request is sent right - # behind the first rather than after a delay: its prompt takes several batches to - # process, which is enough for the two to overlap however fast the first one runs. + # a slot generating a long answer to a short prompt meets a large prompt arriving beside it, + # needing far more than the pool has: the prompt is admitted chunk by chunk, whoever is + # smaller is parked, and both finish. The second request follows immediately, since its + # prompt takes several batches and that is enough overlap however fast the first one runs. global server server.n_ctx = 256 server.start() log = LogReader(server.log_path) prompt_b, n_b = _prompt_of_about(150, "Charlie") - # b lives long enough for the two to collide: the first run of this used 16 tokens - # and b was finished and purged before a had grown into it + # b has to live long enough for the two to collide n_predict_a = 230 n_predict_b = 90 assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 @@ -254,8 +241,8 @@ def _late(n_predict, prompt): def test_preempt_ram_zero_disables_preemption(): - # --preempt-ram 0 is the switch back to the old behaviour: nothing is parked and the - # KV-full path ends the requests the way it always did. + # --preempt-ram 0 switches back to the old behaviour: nothing is parked and the KV-full path + # ends the requests the way it always did global server server.n_ctx = 256 os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" @@ -275,9 +262,8 @@ def test_preempt_ram_zero_disables_preemption(): def test_metrics_and_slots_report_the_parked_state(): - # A client that wants to tell a parked chat from a slow one reads /slots, and an - # operator reads /metrics. Both must show the preemption happening, and the counters - # must survive the requests finishing. + # /slots tells a parked chat from a slow one and /metrics reports it to an operator; both + # must show the preemption, and the counters must survive the requests finishing global server server.n_ctx = 256 server.server_metrics = True @@ -315,11 +301,9 @@ def test_metrics_and_slots_report_the_parked_state(): def test_two_prompts_near_the_context_size_both_complete(): - # Two prompts that each fit the context alone but not together. The second one is - # parked before it takes any cells, and it is close enough to n_ctx that its sequence - # plus its first batch would not leave the usual scheduling margin. It must still be - # restored once the first one finishes: with nothing resident there is nobody to keep - # the margin for. Before the fix it was parked for ever, with no restore ever tried. + # two prompts that each fit the context alone but not together. The second is parked before + # it takes any cells and is too close to n_ctx to leave the usual margin, but must still be + # restored once the first finishes: with nothing resident there is nobody to keep it for. global server server.n_ctx = 256 # the whole prompt in one batch, so the parked slot's first step is the whole prompt @@ -342,11 +326,9 @@ def test_two_prompts_near_the_context_size_both_complete(): def test_the_last_resort_parks_instead_of_ending_everyone(): - # With the planner off nothing is parked ahead of the decode, so two generations that - # fit alone but not together fill the pool until a single token finds no cell. That - # is where upstream ends every slot with the context error. Instead the batch is - # given up, the smaller slot is parked, the larger one finishes, and the parked one - # comes back and finishes too. + # with the planner off, two generations that fit alone but not together fill the pool until a + # single token finds no cell, where upstream ends every slot with the context error. Instead + # the batch is given up, the smaller slot is parked, and both finish. global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" @@ -400,10 +382,8 @@ def test_the_last_resort_works_with_an_unlimited_budget(): def test_the_last_resort_rewinds_a_prompt_in_flight(): - # Same, with a prompt being processed when the pool runs out: the chunk that failed - # is taken back off the slot's tokens and processed again after the resume, so the - # prompt is neither skipped nor fed twice. The prompt is far longer than a batch, so - # the failing chunk is a chunk of it, not its last token. + # same, with a prompt being processed when the pool runs out: the failed chunk comes back off + # the slot's tokens and is processed again after the resume, neither skipped nor fed twice global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" @@ -439,12 +419,10 @@ def _late(n_predict, prompt): def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_head(): - # Two generations that each outgrow the pool on their own, with context shift on. The - # resident reaches the limit, shifts, keeps about half the pool and would keep going - # for as long as it has tokens to make, while the parked one never fits beside it. - # After the head has waited its turn the resident is parked in its place, and the two - # take turns until both finish. Long enough that the resident is still going when the - # head's turn comes: this model makes a couple of thousand tokens a second. + # two generations that each outgrow the pool, with context shift on: the resident shifts and + # would hold half the pool for as long as it generates, while the parked one never fits + # beside it. After the head has waited its turn the resident is parked in its place and the + # two take turns. n_predict is large enough that the resident is still going by then. global server server.n_ctx = 256 server.enable_ctx_shift = True @@ -468,9 +446,8 @@ def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_hea def test_the_rotation_parks_a_resident_that_lets_the_head_in(): - # Three generations with no end in a 256-cell pool with context shift on: two residents - # cycle through shifts while the third waits parked. Every rotation must let the head - # in, so all three keep finishing their tokens and no stream ends short. + # three endless generations with context shift on: two residents cycle through shifts while + # the third waits parked, and every rotation must let the head in so no stream ends short global server server.n_slots = 3 server.n_ctx = 384 @@ -496,11 +473,9 @@ def test_the_rotation_parks_a_resident_that_lets_the_head_in(): def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_server_lives(): - # One request asking for two completions is one conversation in two slots: a parent - # and a child sharing the prompt. When the two together do not fit the pool there is - # nobody else to park, since the family is charged once and a member of it is not a - # victim for the other, so the request gets the context error it would get alone, and - # the server carries on serving. + # a two-completion request is one conversation in two slots, and a family member is not a + # victim for the other, so with nobody else to park it gets the context error it would get + # alone and the server carries on serving global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" From 90a5094d3f72fdef043f39e1b016d7e63c2718b2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 04:25:32 +0000 Subject: [PATCH 72/81] server : validate a started slot's prompt before it can be parked A slot just given a task has not yet passed the prompt checks the STARTED block runs, and the planner could park it first. The park notice opens the stream, so a prompt the checks reject came back as HTTP 200 with an in-stream error where the non-stream 4xx belongs. The checks are one helper now, slot_prompt_rejected(), run by the STARTED block as before and asked by the planner before a started slot can be chosen: a request about to be errored is never given a notice ahead of its error. --- tools/server/server-context.cpp | 108 ++++++++++++------ .../server/tests/unit/test_preempt_notify.py | 37 ++++++ 2 files changed, 110 insertions(+), 35 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7bebae701b8..f57baff24df 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3273,6 +3273,18 @@ struct server_context_impl { continue; // n_cmpl > 1 slots share one sequence, out of scope here } + // a started slot whose request the STARTED block is about to reject gets its + // error on its own pass, and nothing before it: a park notice would open the + // stream and turn that error into 200 plus an in-stream one + if (slot.state == SLOT_STATE_STARTED) { + std::string msg; + error_type type = ERROR_TYPE_SERVER; + + if (slot_prompt_rejected(slot, msg, type)) { + continue; + } + } + if (!preempt_fits_budget(slot)) { continue; } @@ -3596,6 +3608,60 @@ struct server_context_impl { } } + // the checks a slot's request has to pass before its prompt is processed, run from the + // SLOT_STATE_STARTED block below. true when the request is rejected, with the message and + // the type of the error it gets. The empty prompt is not here: it is a final response and + // not an error. + // [TAG_PREEMPT] the planner asks the same question before it parks a started slot, so a + // request that is about to be errored is never given a park notice ahead of its error: a + // notice opens the stream, and the client would get 200 plus an in-stream error where the + // non-stream 4xx belongs. + bool slot_prompt_rejected(const server_slot & slot, std::string & msg, error_type & type) const { + if (!slot.task) { + return false; + } + + // TODO: support memory-less logits computation + if (slot.task->need_logits() && !llama_get_memory(ctx_tgt)) { + msg = "the current context does not logits computation. skipping"; + type = ERROR_TYPE_SERVER; + return true; + } + + if (!slot.can_split()) { + const int32_t n_ubatch = llama_n_ubatch(ctx_tgt); + + if (slot.task->n_tokens() > n_ubatch) { + msg = string_format( + "input (%d tokens) is too large to process. increase the physical batch " + "size (current batch size: %d)", + slot.task->n_tokens(), n_ubatch); + type = ERROR_TYPE_SERVER; + return true; + } + + if (slot.task->n_tokens() > slot.n_ctx) { + msg = string_format( + "input (%d tokens) is larger than the max context size (%d tokens). skipping", + slot.task->n_tokens(), slot.n_ctx); + type = ERROR_TYPE_EXCEED_CONTEXT_SIZE; + return true; + } + + return false; + } + + if (slot.task->n_tokens() >= slot.n_ctx) { + msg = string_format( + "request (%d tokens) exceeds the available context size (%d tokens), try increasing it", + slot.task->n_tokens(), slot.n_ctx); + type = ERROR_TYPE_EXCEED_CONTEXT_SIZE; + return true; + } + + return false; + } + void update_slots() { #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; @@ -4002,46 +4068,18 @@ struct server_context_impl { return; } - // TODO: support memory-less logits computation - if (slot.task->need_logits() && !llama_get_memory(ctx_tgt)) { - send_error(slot, "the current context does not logits computation. skipping", ERROR_TYPE_SERVER); - slot.release(); - return; - } - - if (!slot.can_split()) { - if (slot.task->n_tokens() > n_ubatch) { - send_error(slot, - string_format( - "input (%d tokens) is too large to process. increase the physical batch " - "size (current batch size: %d)", - slot.task->n_tokens(), n_ubatch), - ERROR_TYPE_SERVER); - slot.release(); - return; - } + { + std::string msg; + error_type type = ERROR_TYPE_SERVER; - if (slot.task->n_tokens() > slot.n_ctx) { - send_error( - slot, - string_format( - "input (%d tokens) is larger than the max context size (%d tokens). skipping", - slot.task->n_tokens(), slot.n_ctx), - ERROR_TYPE_EXCEED_CONTEXT_SIZE); - slot.release(); - return; - } - } else { - if (slot.task->n_tokens() >= slot.n_ctx) { - send_error(slot, - string_format("request (%d tokens) exceeds the available context size (%d " - "tokens), try increasing it", - slot.task->n_tokens(), slot.n_ctx), - ERROR_TYPE_EXCEED_CONTEXT_SIZE); + if (slot_prompt_rejected(slot, msg, type)) { + send_error(slot, msg, type); slot.release(); return; } + } + if (slot.can_split()) { if (slot.task->params.cache_prompt) { // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); diff --git a/tools/server/tests/unit/test_preempt_notify.py b/tools/server/tests/unit/test_preempt_notify.py index 76f722f6865..6fa71c55473 100644 --- a/tools/server/tests/unit/test_preempt_notify.py +++ b/tools/server/tests/unit/test_preempt_notify.py @@ -247,3 +247,40 @@ def test_a_resident_rotated_out_for_a_parked_head_is_told_so(): n_parked += len(seq) // 2 # Both streams took turns: at least one park each, so at least two in all. assert n_parked >= 2, [r[0] for r in results] + + +def test_an_oversized_prompt_is_errored_instead_of_parked(): + # A slot that has just been given a task has not passed the prompt checks yet: they + # run on its first pass through update_slots. Parked before that, it would be told + # ": preempted" first, and the notice opens the stream, so a prompt larger than the + # context would come back as 200 plus an in-stream error instead of the plain error + # response it gets with nothing running. The planner leaves such a slot alone. + global server + server.n_ctx = 512 + server.n_batch = 512 # the whole prompt in one batch, so the planner sees its size at once + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + url = f"http://{server.server_host}:{server.server_port}/completion" + resident = _completion_payload(390) | {"prompt": " ".join(["Once upon a time there was a brave knight who"] * 6)} + oversized = _completion_payload(16) | {"prompt": " ".join(["The quick brown fox jumps over the lazy dog and"] * 80)} + + started = threading.Event() + + def _run_resident(): + res = requests.post(url, json=resident, stream=True) + assert res.status_code == 200 + for raw in res.iter_lines(): + if raw.decode("utf-8").startswith("data: "): + started.set() + + t = threading.Thread(target=_run_resident) + t.start() + try: + assert started.wait(60) + res = requests.post(url, json=oversized, stream=True) + body = res.text + assert res.status_code != 200, body + assert not body.lstrip().startswith(":"), body + assert "error" in json.loads(body), body + finally: + t.join(120) From f8fdf063a1e4337f0d8c92acc5c382655af959ff Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 05:37:28 +0000 Subject: [PATCH 73/81] server, llama, ggml : shorter comments across the preemption work Cuts the comments this branch adds by about 70 percent, keeping the invariants, the citations and the traps and dropping the restatements, banners and measured asides. Comments only; no code changes. --- common/arg.cpp | 3 +- common/common.cpp | 13 +- common/common.h | 5 +- ggml/include/ggml-backend.h | 10 +- ggml/include/ggml-cuda.h | 6 +- ggml/src/ggml-backend-impl.h | 4 +- ggml/src/ggml-cpu/ggml-cpu.cpp | 5 +- ggml/src/ggml-cuda/common.cuh | 3 +- ggml/src/ggml-cuda/fattn-common.cuh | 6 +- ggml/src/ggml-cuda/fattn-vec.cuh | 3 +- ggml/src/ggml-cuda/fattn.cu | 3 +- ggml/src/ggml-cuda/ggml-cuda.cu | 62 +- ggml/src/ggml-cuda/mmvq.cu | 13 +- ggml/src/ggml-cuda/mmvq.cuh | 3 +- ggml/src/ggml-metal/ggml-metal-device.m | 3 +- ggml/src/ggml-rpc/ggml-rpc.cpp | 3 +- include/llama.h | 59 +- scripts/batchinv/divergence.py | 7 +- scripts/batchinv/probe.cpp | 19 +- scripts/unsloth/additive_merge.py | 16 +- scripts/unsloth/feature_matrix.py | 7 +- scripts/unsloth/pin_contract.py | 16 +- scripts/unsloth/test_additive_merge.py | 1 - scripts/unsloth/test_pin_contract.py | 6 +- src/llama-batch.cpp | 14 +- src/llama-batch.h | 7 +- src/llama-context.cpp | 187 +----- src/llama-context.h | 12 +- src/llama-graph.cpp | 6 +- src/llama-impl.cpp | 23 +- src/llama-impl.h | 8 +- src/llama-kv-cache.cpp | 63 +- src/llama-kv-cache.h | 8 +- src/llama-memory-hybrid.cpp | 11 +- src/llama-memory-recurrent.cpp | 5 +- src/llama-memory.h | 6 +- tests/test-backend-ops.cpp | 8 +- tests/test-state-restore-fragmented.cpp | 3 +- tests/test-state-seq-copy.cpp | 23 +- tools/server/server-context.cpp | 587 +++--------------- tools/server/server-queue.cpp | 11 +- tools/server/server-task.h | 5 +- tools/server/tests/unit/test_preempt.py | 131 +--- .../server/tests/unit/test_preempt_notify.py | 32 +- 44 files changed, 255 insertions(+), 1171 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 0dea146fa49..633337da885 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1305,8 +1305,7 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e } params.lr.init(); - // [TAG_EXACT_CONCURRENCY] refuse a column bound that cannot cover a decode step before - // anything is loaded, rather than running with the guarantee quietly switched off + // [TAG_EXACT_CONCURRENCY] refuse a column bound that cannot cover a decode step before anything is loaded, rather than running with the guarantee quietly off if (!common_exact_concurrency_init(ctx_arg.params)) { ctx_arg.params = params_org; return false; diff --git a/common/common.cpp b/common/common.cpp index b74bee734ae..9d8918f4b80 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1290,9 +1290,7 @@ struct common_init_result::impl { common_init_result::common_init_result(common_params & params, bool model_only) : pimpl(new impl{}) { - // [TAG_EXACT_CONCURRENCY] before any context exists, so one is never created under a figure - // the explicit bound does not cover; this also covers a caller that skipped - // common_params_parse(). On failure nothing is loaded. + // [TAG_EXACT_CONCURRENCY] before any context exists, so one is never created under a figure the explicit bound does not cover if (!model_only && !common_exact_concurrency_init(params)) { COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: refusing to load the model, see the error above\n"); return; @@ -1456,8 +1454,6 @@ bool common_exact_concurrency() { int common_exact_decode_width(const common_params & params) { const int64_t n_slots = std::max(1, params.n_parallel); - // draft tokens a slot carries into the verify ubatch, from the same place the speculation - // code takes its own width const int64_t n_draft = std::max(0, (int) common_speculative_n_max(¶ms.speculative)); // the product is handed to a backend as an int; one that overflows is reported, not wrapped @@ -1472,8 +1468,7 @@ bool common_exact_concurrency_init(const common_params & params) { return true; } - // DFlash drafting turns causal attention off on its draft context, which the paged attention - // needs; say so instead of asserting in the graph. DSpark is the same implementation. + // DFlash drafting turns causal attention off on its draft context, which the paged attention needs; say so instead of asserting in the graph. DSpark is the same. for (const auto type : params.speculative.types) { if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash or draft-dspark: both disable causal attention on the draft, which the paged attention needs\n"); @@ -1502,9 +1497,7 @@ bool common_exact_concurrency_init(const common_params & params) { } } - // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step - // is. A context created later reports n_seq_max times that, which is n_cols again; reporting - // n_cols here too covers a caller that decodes first, or contexts it created earlier. + // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step is; this also covers a caller that decodes before creating a context if (!llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel))) || !llama_set_exact_decode_width((uint32_t) n_cols)) { COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: the decode width could not be reported, see the error above\n"); diff --git a/common/common.h b/common/common.h index ea84e0e7515..e22c07bad12 100644 --- a/common/common.h +++ b/common/common.h @@ -935,12 +935,9 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode // [TAG_EXACT_CONCURRENCY] true when LLAMA_EXACT_CONCURRENCY is set for this process bool common_exact_concurrency(); -// the widest ubatch a decode step can build here: one column per slot times one plus its draft -// tokens. Under exact mode the CUDA column policy has to cover this, and derives its bound from it. int common_exact_decode_width(const common_params & params); -// report that width to the CUDA backend, refusing a smaller explicit -// GGML_CUDA_BATCH_INVARIANT_MAX_COLS; false if the configuration must not run +// report that width to the CUDA backend, refusing a smaller explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS; false if the configuration must not run bool common_exact_concurrency_init(const common_params & params); struct llama_model_params common_model_params_to_llama ( common_params & params); diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 09ee64a6561..84a2f8458ea 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -62,8 +62,7 @@ extern "C" { GGML_API size_t ggml_backend_buffer_get_alloc_size(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor); GGML_API void ggml_backend_buffer_clear (ggml_backend_buffer_t buffer, uint8_t value); GGML_API bool ggml_backend_buffer_is_host (ggml_backend_buffer_t buffer); - // whether the buffer copies a strided set of rows in one call (see ggml_backend_tensor_set_2d); - // without it the generic path issues one transfer per row + // whether the buffer copies a strided set of rows in one call (see ggml_backend_tensor_set_2d); without it the generic path issues one transfer per row GGML_API bool ggml_backend_buffer_supports_2d (ggml_backend_buffer_t buffer); GGML_API void ggml_backend_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage); GGML_API enum ggml_backend_buffer_usage ggml_backend_buffer_get_usage (ggml_backend_buffer_t buffer); @@ -128,9 +127,7 @@ extern "C" { GGML_API void ggml_backend_event_free(ggml_backend_event_t event); GGML_API void ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backend); GGML_API void ggml_backend_event_synchronize(ggml_backend_event_t event); - // non-blocking: true once everything recorded before the event has completed. - // backends without a query implementation fall back to a blocking synchronize and return true, - // which ggml_backend_dev_supports_event_query() tells apart from a real non-blocking query. + // non-blocking: true once everything recorded before the event has completed. Backends without a query implementation fall back to a blocking synchronize. GGML_API bool ggml_backend_event_query(ggml_backend_event_t event); GGML_API void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event); @@ -197,8 +194,7 @@ extern "C" { GGML_API ggml_backend_buffer_t ggml_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void * ptr, size_t size, size_t max_tensor_size); GGML_API bool ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_tensor * op); - // whether ggml_backend_event_query() on this device really is non-blocking, i.e. whether - // the device implements it rather than falling back to a blocking synchronize + // whether ggml_backend_event_query() on this device really is non-blocking, rather than falling back to a blocking synchronize GGML_API bool ggml_backend_dev_supports_event_query(ggml_backend_dev_t device); GGML_API bool ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft); GGML_API bool ggml_backend_dev_offload_op(ggml_backend_dev_t device, const struct ggml_tensor * op); diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h index 07131df327c..897da6ca5f8 100644 --- a/ggml/include/ggml-cuda.h +++ b/ggml/include/ggml-cuda.h @@ -39,11 +39,7 @@ GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * f GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); -// [TAG_EXACT_CONCURRENCY] report the widest ubatch a decode step of this process can build: one -// column per slot times one plus its draft tokens. Under LLAMA_EXACT_CONCURRENCY the column policy -// defaults to that instead of a fixed number, so --parallel or a wider draft cannot silently push a -// decode above the bound. An explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS still wins. Call before -// the first graph is computed; also available through ggml_backend_reg_get_proc_address(). +// [TAG_EXACT_CONCURRENCY] report the widest ubatch a decode step of this process can build, so the column policy covers it; call before the first graph is computed GGML_BACKEND_API void ggml_backend_cuda_set_exact_decode_width(int n_cols); GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index bb3be31217b..e417b3cc448 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -201,9 +201,7 @@ extern "C" { void (*event_free) (ggml_backend_dev_t dev, ggml_backend_event_t event); void (*event_synchronize) (ggml_backend_dev_t dev, ggml_backend_event_t event); - // (optional) non-blocking completion test for an event. - // kept last so that backends that do not implement it need no change: a missing entry - // is NULL, and ggml_backend_event_query() then falls back to a blocking synchronize. + // (optional) non-blocking completion test for an event. Kept last: a missing entry is NULL and ggml_backend_event_query() then blocks instead. bool (*event_query) (ggml_backend_dev_t dev, ggml_backend_event_t event); }; diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 475f30820ba..7ea548bcbce 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -474,10 +474,7 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return ggml_is_contiguous(op->src[0]); case GGML_OP_SSM_SCAN: return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1; - // [TAG_EXACT_CONCURRENCY] note: FLASH_ATTN_EXT with src[5], the page table, is deliberately - // still accepted. The CPU ignores it and attends in physical order, but it is also the - // reference test-backend-ops compares the paged CUDA kernel against, and that test's mask - // selects exactly the listed cells. A KV layer cannot reach the CPU under the mode anyway. + // [TAG_EXACT_CONCURRENCY] note: FLASH_ATTN_EXT with src[5], the page table, is deliberately still accepted: the CPU ignores it, but it is the reference test-backend-ops uses default: return true; } diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index d749b42ee48..47ecee48c2f 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -51,8 +51,7 @@ #define GGML_CUDA_CC_DP4A 610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products // [TAG_BATCH_INVARIANT] 0 = off, 1 = split every batched matmul, 2 = split only where it changes bits int ggml_cuda_batch_invariant(); -// widest batch the split applies to, 0 = no bound; bounding it gives up prompt-phase invariance -// only, and prompt-sized batches cost far more to split +// widest batch the split applies to, 0 = no bound; bounding it gives up prompt-phase invariance only int ggml_cuda_batch_invariant_max_cols(); #define GGML_CUDA_CC_VOLTA 700 diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index 10f28eb229b..61aa0dc376a 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1091,8 +1091,7 @@ void launch_fattn( // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped. // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or // multiple sequences of possibly different lengths. - // [TAG_BATCH_INVARIANT] without this scan the KV loop runs to K->ne[1], which grows with the - // other sequences sharing the cache; scanning the mask bounds it by the sequence's own extent + // [TAG_BATCH_INVARIANT] without this scan the KV loop runs to K->ne[1], which grows with the other sequences; the mask bounds it by the sequence's own extent const bool batch_invariant_KV_max = ggml_cuda_batch_invariant() != 0; if (!dst->src[5] && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) { const int64_t s31 = mask->nb[1] / sizeof(half2); @@ -1152,8 +1151,7 @@ void launch_fattn( dst_tmp_meta.alloc((size_t(blocks_num.x) * ncols * (2 + DV/2))); } } else if (dst->src[5] || ggml_cuda_batch_invariant()) { - // [TAG_BATCH_INVARIANT] the KV split between blocks, and so the order the partials combine - // in, follows K->ne[1], which grows with the other sequences: pin it to one block per tile + // [TAG_BATCH_INVARIANT] the KV split between blocks, and so the order the partials combine in, follows K->ne[1]: pin it to one block per tile parallel_blocks = 1; blocks_num.x = ntiles_x; diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 4005d28879e..0dba4b80b9c 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -247,8 +247,7 @@ static __global__ void flash_attn_ext_vec( #endif // V_DOT2_F32_F16_AVAILABLE } - // in the paged specialization KV_max carries [count, physical page IDs...] per query; the loop - // and each warp's recurrence follow logical positions, never physical addresses + // in the paged specialization KV_max carries [count, physical page IDs...] per query; the loop and each warp's recurrence follow logical positions, never physical addresses static_assert(!paged || ncols == 1, "paged attention has one query per block"); const int * pages = paged ? KV_max + (sequence*int(ne01.z) + ic0)*(1 + ne11/FATTN_KQ_STRIDE) : nullptr; const int k_VKQ_max = paged ? pages[0]*FATTN_KQ_STRIDE : (KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11); diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 915c2d04da8..5a1303401d0 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -457,8 +457,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path. const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0; - // [TAG_BATCH_INVARIANT] every choice below switches on Q->ne[1] or K->ne[1], both of which grow - // with the other sequences, so pin the kernel a batch of one would use + // [TAG_BATCH_INVARIANT] every choice below switches on Q->ne[1] or K->ne[1], both of which grow with the other sequences, so pin the kernel a batch of one would use if (ggml_cuda_batch_invariant() && can_use_vector_kernel && Q->ne[1] == 1) { return BEST_FATTN_KERNEL_VEC; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 5ff2c5f759f..b41b704e784 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1758,8 +1758,7 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, } static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { - // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving - // it on would give a solo request a different code path from a batched one + // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving it on would give a solo request a different code path from a batched one if (ggml_cuda_batch_invariant()) { return false; } @@ -1791,8 +1790,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { } static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { - // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving - // it on would give a solo request a different code path from a batched one + // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving it on would give a solo request a different code path from a batched one if (ggml_cuda_batch_invariant()) { return false; } @@ -1825,9 +1823,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } -// [TAG_BATCH_INVARIANT] the token count picks the matmul implementation and how its K loop is -// divided between threads, both of which change the summation order, so the same request produces -// different bits depending on how many others decode alongside it. GGML_CUDA_BATCH_INVARIANT: +// [TAG_BATCH_INVARIANT] the token count picks the matmul and how its K loop is split, so the same request produces different bits. GGML_CUDA_BATCH_INVARIANT: // 1 - compute every destination column on its own, exactly as a batch of one would // 2 - split off only the columns whose batch-of-one configuration differs from the batched one static bool ggml_cuda_exact_concurrency() { @@ -1859,8 +1855,7 @@ void ggml_backend_cuda_set_exact_decode_width(int n_cols) { } int ggml_cuda_batch_invariant_max_cols() { - // [TAG_EXACT_CONCURRENCY] prompt ubatches hold one sequence, so a prefill already matches its - // solo run and needs no unbounded column policy. An explicit bound always wins. + // [TAG_EXACT_CONCURRENCY] prompt ubatches hold one sequence, so a prefill already matches its solo run; an explicit bound always wins static const int explicit_cols = []() { const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); return val ? atoi(val) : -1; @@ -1874,19 +1869,12 @@ int ggml_cuda_batch_invariant_max_cols() { return 0; } - // the widest ubatch a decode step can build: one column per slot times one plus its draft - // tokens, as reported by ggml_backend_cuda_set_exact_decode_width(). Failing a report, 16, - // which covers --parallel 4 --spec-type draft-mtp --spec-draft-n-max 2. Above the bound the - // column split does not fire and ggml_cuda_warn_above_exact_bound() says so once. const int width = g_exact_decode_width.load(std::memory_order_relaxed); return width > 0 ? width : 16; } -// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so the mode does not hold -// for that op; say so once. Only when nothing reported a decode width: a reported one makes the -// batches above the bound prompt ubatches, which are exact by holding a single sequence, so -// warning on those would be crying wolf on every prefill. +// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so say so once. Only when nothing reported a decode width: with one, wider batches are single-sequence prefills. static void ggml_cuda_warn_above_exact_bound(const char * op, int64_t ncols, int max_cols) { if (!ggml_cuda_exact_concurrency()) { return; @@ -1918,7 +1906,6 @@ enum ggml_cuda_mm_path { GGML_CUDA_MM_CUBLAS, }; -// the implementation ggml_cuda_mul_mat would pick for a batch of ne11 columns static ggml_cuda_mm_path ggml_cuda_mul_mat_path( int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, int64_t ne11) { // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. @@ -1955,11 +1942,7 @@ static ggml_cuda_mm_path ggml_cuda_mul_mat_path( static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); -// [TAG_BATCH_INVARIANT] the widest slice of columns that can be recomputed in one launch while -// every column still sums as a batch of one would. A column's result depends on the implementation -// and, for MMVQ, the launch's warp count, never on the other columns, so a slice this wide reads -// the weights once for all of them instead of once per column. Always below ncols_dst, so the -// recursive call cannot land back here with the same shape. +// [TAG_BATCH_INVARIANT] the widest slice of columns that can be recomputed in one launch while every column still sums as a batch of one; always below ncols_dst, so the recursion ends static int64_t ggml_cuda_mul_mat_invariant_width( int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, ggml_cuda_mm_path path_one, int64_t ncols_dst) { @@ -1979,13 +1962,11 @@ static int64_t ggml_cuda_mul_mat_invariant_width( return 1; } -// recompute dst in slices of columns so each column sees the batch-of-one configuration; false -// when the batched launch already gives every column that same value +// recompute dst in slices of columns so each column sees the batch-of-one configuration; false when the batched launch already gives every column that value static bool ggml_cuda_mul_mat_split_columns( ggml_backend_cuda_context & ctx, int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - // recurrent output projections broadcast one weight matrix over sequence planes: these are - // token projections too, so normalize each plane before applying the column policy + // recurrent output projections broadcast one weight matrix over sequence planes, so normalize each plane before applying the column policy if (ggml_cuda_exact_concurrency() && src0->ne[2] == 1 && src0->ne[3] == 1 && (dst->ne[2] > 1 || dst->ne[3] > 1) && src1->ne[2] == dst->ne[2] && src1->ne[3] == dst->ne[3]) { @@ -2007,7 +1988,6 @@ static bool ggml_cuda_mul_mat_split_columns( if (ncols_dst <= 1 || src1->ne[1] != ncols_dst) { return false; } - // only the token dimension is split; batched matmuls (attention) keep their shape if (src1->ne[2] != 1 || src1->ne[3] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) { return false; } @@ -2017,8 +1997,7 @@ static bool ggml_cuda_mul_mat_split_columns( return false; } - // mode 1 recomputes one column at a time; mode 2, which exact concurrency runs under, uses the - // widest slices that keep the batch-of-one arithmetic + // mode 1 recomputes one column at a time; mode 2, which exact concurrency runs under, uses the widest slices that keep the batch-of-one arithmetic int64_t width = 1; if (ggml_cuda_batch_invariant() >= 2) { const ggml_cuda_mm_path path_one = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1); @@ -2106,8 +2085,7 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor GGML_ABORT("fatal error"); } -// [TAG_BATCH_INVARIANT] true when the policy computes this MUL_MAT_ID one token at a time, so -// every expert product reduces as it would in a batch of one +// [TAG_BATCH_INVARIANT] true when the policy computes this MUL_MAT_ID one token at a time static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) { if (!ggml_cuda_batch_invariant()) { return false; @@ -2134,8 +2112,7 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c return true; } - // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path that - // decides whether the stream is synchronized is the single-token one + // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path that decides whether the stream is synchronized is the single-token one const int64_t ntokens = ggml_cuda_mul_mat_id_splits_tokens(dst) ? 1 : dst->ne[2]; if (ntokens <= MMVQ_MAX_BATCH_SIZE) { @@ -2161,9 +2138,7 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst); -// [TAG_BATCH_INVARIANT] recompute dst one token at a time. Every implementation below groups the -// ubatch's tokens by the expert they routed to, so shapes depend on what the other tokens picked; -// one call per token makes the callee see the shapes a batch of one has. +// [TAG_BATCH_INVARIANT] recompute dst one token at a time: every implementation below groups the ubatch's tokens by the expert they routed to, so shapes depend on the other tokens static void ggml_cuda_mul_mat_id_split_tokens(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src1 = dst->src[1]; const ggml_tensor * ids = dst->src[2]; @@ -2175,7 +2150,6 @@ static void ggml_cuda_mul_mat_id_split_tokens(ggml_backend_cuda_context & ctx, g ggml_tensor ids_token = *ids; ggml_tensor dst_token = *dst; - // src1 is [ne10, ne11, ntokens], one expert list per token in ids [n_expert_used, ntokens] src1_token.ne[2] = 1; src1_token.nb[3] = src1_token.nb[2]; src1_token.data = (char *) src1->data + i*src1->nb[2]; @@ -2211,8 +2185,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * // [TAG_BATCH_INVARIANT] if (ggml_cuda_mul_mat_id_splits_tokens(dst)) { GGML_ASSERT(ne3 == 1 && src1->ne[3] == 1 && ids->ne[2] == 1 && ids->ne[3] == 1); - // a quantized expert matrix takes the single-token MMVQ path at every token count, and that - // path can put the tokens on its sample axis in one launch; anything else goes token by token + // a quantized expert matrix takes the single-token MMVQ path at every token count and can put the tokens on its sample axis in one launch; anything else goes token by token if (ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); return; @@ -3606,10 +3579,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } - // topk-moe - // [TAG_BATCH_INVARIANT] the routing fusion passes its memory-range check only for a one-token - // ubatch, so a solo request takes the fused top-k kernel and a batched one takes the softmax, - // argsort and normalize chain: two algorithms for one set of routing weights + // [TAG_BATCH_INVARIANT] the routing fusion passes its memory-range check only for a one-token ubatch, so a solo request takes the fused top-k kernel and a batched one the long chain if (!ggml_cuda_batch_invariant() && (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || cgraph->nodes[i]->op == GGML_OP_ARGSORT)) { @@ -5693,9 +5663,7 @@ static bool ggml_backend_cuda_device_event_query(ggml_backend_dev_t dev, ggml_ba const cudaError_t err = cudaEventQuery((cudaEvent_t)event->context); - // not an error, and nothing to clear: cudaEventQuery() returns cudaErrorNotReady - // without recording it as the thread's last error, so collecting one here would only - // consume somebody else's, and a real launch failure would be swallowed + // not an error, and nothing to clear: cudaEventQuery() returns cudaErrorNotReady without recording it, so collecting one here would consume somebody else's if (err == cudaErrorNotReady) { return false; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 30f37c7722e..032964198a3 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -551,8 +551,7 @@ bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t n // There nwarps also depends on the K loop trip count, which the caller does not pass in. return ncols_dst == 1; } - // blocks_per_iter, which assigns K blocks to threads, is proportional to nwarps; - // rows_per_cuda_block only changes which rows a block owns, not the order within a row + // blocks_per_iter, which assigns K blocks to threads, is proportional to nwarps; rows_per_cuda_block only changes which rows a block owns, not the order within a row return calc_nwarps(type, 1, table_id) == calc_nwarps(type, (int) ncols_dst, table_id); } @@ -593,9 +592,8 @@ static __global__ void mul_mat_vec_q( ggml_cuda_pdl_sync(); sample_dst = blockIdx.z; - // [TAG_BATCH_INVARIANT] with ids, a sample is a token: the batch-invariant MUL_MAT_ID launch - // puts every token on the z axis of one single-column launch, so each (token, expert slot) - // block runs the single-token configuration. The stock launch has one sample, as before. + // [TAG_BATCH_INVARIANT] with ids, a sample is a token: every token goes on the z axis of one single-column + // launch, so each (token, expert slot) block runs the single-token configuration channel_x = ncols_dst == 1 && ids ? ids[sample_dst*ids_stride + channel_dst] : fastdiv(channel_dst, channel_ratio); channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; @@ -1284,9 +1282,8 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT( nb0 == ts_dst); GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); - // [TAG_BATCH_INVARIANT] a multi-token MUL_MAT_ID becomes one launch of the single-token - // configuration with the tokens on the sample axis, so every (token, expert slot) block - // reduces as the token alone would and the count is not bounded by the column templates + // [TAG_BATCH_INVARIANT] a multi-token MUL_MAT_ID becomes one launch of the single-token configuration + // with the tokens on the sample axis, so the count is not bounded by the column templates const bool tokens_as_samples = ids && ne2 > 1 && ggml_cuda_batch_invariant(); GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE || tokens_as_samples); diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 67d69b1415d..688c944c1f9 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -4,8 +4,7 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11); -// [TAG_BATCH_INVARIANT] true when an MMVQ launch of ncols_dst columns sums each destination element -// in the same order as a single-column launch, i.e. when the column count leaves nwarps unchanged +// [TAG_BATCH_INVARIANT] true when an MMVQ launch of ncols_dst columns sums each destination element in the same order as a single-column launch, i.e. when nwarps is unchanged bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst); // Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID, diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index e5fc6a8063b..bc831a5afac 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1592,8 +1592,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_ROLL: return ggml_is_contiguous(op->src[0]); case GGML_OP_FLASH_ATTN_EXT: - // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; - // walking the pool in physical order here would be silently wrong + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; walking the pool in physical order here would be silently wrong if (op->src[5] != NULL) { return false; } diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 6602027e688..7c51ddd8206 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1915,8 +1915,7 @@ static ggml_backend_buffer_type_t ggml_backend_rpc_device_get_buffer_type(ggml_b static bool ggml_backend_rpc_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { GGML_UNUSED(dev); - // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; - // the remote end is not asked, so it is not claimed here + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; the remote end is not asked, so it is not claimed here if (op->op == GGML_OP_FLASH_ATTN_EXT && op->src[5]) { return false; } diff --git a/include/llama.h b/include/llama.h index 772c5fe9e1c..13362a7f098 100644 --- a/include/llama.h +++ b/include/llama.h @@ -795,26 +795,14 @@ extern "C" { // Check if the memory supports shifting LLAMA_API bool llama_memory_can_shift(llama_memory_t mem); - // [TAG_EXACT_CONCURRENCY] cells the memory allocates in one indivisible unit: 1 ordinarily, - // larger where a mode places cells in blocks, and then n tokens occupy round_up(n, granularity) - // cells. A caller deciding whether the pool has room must round the same way. + // [TAG_EXACT_CONCURRENCY] cells the memory allocates in one indivisible unit: 1 ordinarily, larger where a mode places cells in blocks, when n tokens occupy round_up(n, granularity) LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem); - // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or 1 - // plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence - // set with more left to place is a prompt and is prefilled in a ubatch of its own; one at or - // below stays grouped with the other decodes. Process-wide, default 1, never lowered. Raising - // it widens every existing context's decode step and re-reports their width; returns false, - // and changes nothing, when an explicit column bound cannot cover that width. + // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or 1 plus the draft length. Never lowered; false when a column bound cannot cover it. LLAMA_API bool llama_set_exact_decode_tokens(uint32_t n_tokens); LLAMA_API uint32_t llama_exact_decode_tokens(void); - // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns: the - // sequences a context holds times the tokens each contributes. Every context reports its own - // at creation and a backend keeps the widest it has heard. A caller that builds wider steps - // reports the width itself, before the context or the first decode. Never lowered. Returns - // false, reporting nothing, when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is positive and below the - // width: that bound wins in the backend, so decodes above it would be left batched. + // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns; never lowered, and false when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is below it LLAMA_API bool llama_set_exact_decode_width(uint32_t n_cols); LLAMA_API uint32_t llama_exact_decode_width(void); @@ -950,53 +938,27 @@ extern "C" { llama_seq_id dest_seq_id, llama_state_seq_flags flags); - // [TAG_STATE_ASYNC] asynchronous per-sequence state transfer - // - // llama_state_seq_get_data_ext / set_data_ext do not return until every byte has moved, - // so a caller that copies a sequence out of the cache to make room stops doing anything - // else for as long as the copy takes. A transfer object issues the same copies on a - // stream of its own and hands back control immediately; the caller polls - // llama_state_seq_copy_done() and gets on with its other work in between. - // - // The transfer owns the host buffer it reads from or writes into. That buffer is pinned - // when the backend offers pinned memory, which is what makes the copy fast, and it - // cannot be freed while a copy is still using it. - // - // Between issuing and completion the caller must not touch the buffer, must not free or - // reuse the cells of a sequence being read, and must not decode a sequence being - // written. llama_state_seq_copy_free() waits for an outstanding copy first. + // [TAG_STATE_ASYNC] asynchronous per-sequence state transfer, polled with llama_state_seq_copy_done(). + // Until it completes the caller must not touch the buffer, free the cells read, or decode what is written. struct llama_state_seq_copy; - // NULL if the context's backends cannot copy asynchronously, or cannot say whether a - // copy has finished without waiting for it, which would put the stall straight back; the - // caller then uses the synchronous llama_state_seq_*_data_ext calls + // NULL when the backends cannot copy asynchronously, or cannot say whether a copy has finished without waiting for it; the caller then uses the synchronous calls LLAMA_API struct llama_state_seq_copy * llama_state_seq_copy_init(struct llama_context * ctx); LLAMA_API void llama_state_seq_copy_free(struct llama_state_seq_copy * cpy); - // Size the transfer's host buffer, keeping no contents; NULL on failure. Grow-only: - // page-locking host memory is far too slow to do once per transfer, so the memory is - // kept between them and only given back by llama_state_seq_copy_buf_free(). + // size the transfer's host buffer, keeping no contents; NULL on failure. Grow-only: page-locking is far too slow to redo per transfer, so only llama_state_seq_copy_buf_free() frees it. LLAMA_API uint8_t * llama_state_seq_copy_buf_resize (struct llama_state_seq_copy * cpy, size_t size); LLAMA_API uint8_t * llama_state_seq_copy_buf (struct llama_state_seq_copy * cpy); LLAMA_API size_t llama_state_seq_copy_buf_size (struct llama_state_seq_copy * cpy); - // host memory actually held, which is what a caller budgeting host RAM has to count LLAMA_API size_t llama_state_seq_copy_buf_capacity(struct llama_state_seq_copy * cpy); LLAMA_API void llama_state_seq_copy_buf_free (struct llama_state_seq_copy * cpy); - // true when the buffer that is held right now is page-locked, i.e. when the copies can - // really overlap. False while no buffer is held, since none is page-locked then: a - // caller asking before the first resize wants llama_state_seq_copy_buf_can_pin(). + // true when the buffer held right now is page-locked. False while no buffer is held: ask llama_state_seq_copy_buf_can_pin() instead. LLAMA_API bool llama_state_seq_copy_buf_is_pinned(struct llama_state_seq_copy * cpy); - // true when the backend offers pinned host memory at all. It is what the next resize - // will ask for, not what any buffer is: an allocation can still come back pageable. LLAMA_API bool llama_state_seq_copy_buf_can_pin(struct llama_state_seq_copy * cpy); - // Issue the copies; return the number of bytes covered, 0 on failure. size must be - // between 1 and llama_state_seq_copy_buf_size(): the buffer belongs to the transfer, and - // a size beyond it is refused rather than believed. LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is - // refused too, since these copies serialise through host memory; use - // llama_state_seq_get_data_ext / set_data_ext for that flag. + // issue the copies; the bytes covered, 0 on failure. size must be within llama_state_seq_copy_buf_size(), and LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is refused. LLAMA_API size_t llama_state_seq_copy_get( struct llama_state_seq_copy * cpy, size_t size, @@ -1009,13 +971,10 @@ extern "C" { llama_seq_id dest_seq_id, llama_state_seq_flags flags); - // transfers the last issue posted: one per run of adjacent cells, per tensor LLAMA_API size_t llama_state_seq_copy_n_copies(struct llama_state_seq_copy * cpy); - // microseconds the last issue spent waiting for the compute streams before it could start LLAMA_API int64_t llama_state_seq_copy_sync_us(struct llama_state_seq_copy * cpy); - // non-blocking completion test, and the blocking wait behind it LLAMA_API bool llama_state_seq_copy_done(struct llama_state_seq_copy * cpy); LLAMA_API void llama_state_seq_copy_wait(struct llama_state_seq_copy * cpy); diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py index 198c2bdf881..32e79be2fc1 100644 --- a/scripts/batchinv/divergence.py +++ b/scripts/batchinv/divergence.py @@ -5,8 +5,7 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from prompts import PROMPTS -# recorded with every run; LLAMA_EXACT_CONCURRENCY inherited from the shell decides whether a run -# labelled as the mode-off reference actually was one, so it is not optional +# recorded with every run: LLAMA_EXACT_CONCURRENCY inherited from the shell decides whether a reference run really was one RECORDED_ENV = ("LLAMA_EXACT_CONCURRENCY", "GGML_CUDA_BATCH_INVARIANT", "GGML_CUDA_BATCH_INVARIANT_MAX_COLS", "LLAMA_SERVER_PREEMPT_EVERY", "LLAMA_KV_CACHE_DEBUG", "LLAMA_BATCH_DEBUG", "CUDA_VISIBLE_DEVICES") @@ -60,7 +59,6 @@ def __init__(self, port, binary, extra, env_extra, log_path, spec, kv_unified=Tr self.env = dict(os.environ) self.env["CUDA_VISIBLE_DEVICES"] = "3" self.env.update(env_extra) - # what the server will actually see, not what this run meant to set self.env_resolved = {k: self.env[k] for k in RECORDED_ENV if k in self.env} self.p = None self.fh = None @@ -86,8 +84,6 @@ def __enter__(self): time.sleep(1.0) raise RuntimeError("server did not become healthy") except BaseException: - # __exit__ is not called when __enter__ raises, and a server that started but never - # reported healthy would keep the GPU, the port and the log handle self.__exit__(None, None, None) raise @@ -189,7 +185,6 @@ def main(): res["reference"] = a.reference res["solo"] = {"n_tokens": len(ref), "tok_per_s": solo["timings"]["predicted_per_second"], "text_sha": None} - # solo repeat, to prove solo itself is stable solo2 = completion(a.port, PROMPTS["P0"], a.n_predict) res["solo_repeat_first_diff"] = first_diff(ref, solo2["tokens"]) res["rounds"] = [] diff --git a/scripts/batchinv/probe.cpp b/scripts/batchinv/probe.cpp index 204379c423b..4e72d5b8a55 100644 --- a/scripts/batchinv/probe.cpp +++ b/scripts/batchinv/probe.cpp @@ -1,6 +1,4 @@ -// Locate the first graph op whose sequence-0 output changes when the decode batch holds four -// sequences instead of one. Seq 0's prompt KV is identical in both phases, so the only -// difference is the width of the final decode ubatch. +// Locate the first graph op whose sequence-0 output changes when the decode batch holds four sequences instead of one; seq 0's prompt KV is identical in both phases. #include "llama.h" #include "ggml.h" #include "ggml-backend.h" @@ -120,7 +118,6 @@ static llama_token greedy(llama_context * ctx, int32_t i, int n_vocab) { return best; } -// feed a prompt as one decode call for one sequence, return the greedy next token static llama_token feed(llama_context * ctx, const std::vector & p, llama_seq_id seq, int n_vocab) { batch_holder h; for (size_t i = 0; i < p.size(); ++i) { @@ -167,12 +164,9 @@ int main(int argc, char ** argv) { std::vector rec_a, rec_b; llama_token first_tok[4] = {0, 0, 0, 0}; - // phase A: decode ubatch width 1. PROBE_A_FILL is how many sequences are already in the - // shared KV cache, which sets K->ne[1] for attention. + // phase A: decode ubatch width 1. PROBE_A_FILL is how many sequences are already in the shared KV cache, which sets K->ne[1]. const int a_fill = getenv("PROBE_A_FILL") ? atoi(getenv("PROBE_A_FILL")) : 1; - // PROBE_A_PERM reorders which prompt goes into which sequence in phase A, keeping the cache - // length but changing what the masked cells hold. It may only move the neighbours: sequence 0 - // keeps prompt 0, or the two phases would compare different sequences. + // PROBE_A_PERM reorders which prompt goes into which sequence in phase A; it may only move the neighbours, since sequence 0 must keep prompt 0 int a_perm[4] = {0, 1, 2, 3}; if (const char * perm = getenv("PROBE_A_PERM")) { for (int k = 0; k < 4 && perm[2*k]; ++k) a_perm[k] = perm[2*k] - '0'; @@ -195,7 +189,6 @@ int main(int argc, char ** argv) { llama_free(ctx); } - // phase B: same seq-0 prompt KV, then a decode ubatch holding n_seqs tokens { llama_context * ctx = make_ctx(); if (prefill) { @@ -238,7 +231,6 @@ int main(int argc, char ** argv) { llama_free(ctx); } - // optional: keep decoding and report the first step at which seq 0's token differs const int n_steps = getenv("PROBE_STEPS") ? atoi(getenv("PROBE_STEPS")) : 0; int first_bad_step = -1; if (n_steps > 0) { @@ -277,7 +269,6 @@ int main(int argc, char ** argv) { fprintf(stderr, "nodes: A=%zu B=%zu first tokens: %d %d %d %d\n", rec_a.size(), rec_b.size(), first_tok[0], first_tok[1], first_tok[2], first_tok[3]); - // walk both node lists in order and compare seq 0's slice FILE * out = out_path ? fopen(out_path, "w") : stdout; fprintf(out, "{\"n_seqs\":%d,\"first_bad_step\":%d,\"nodes_a\":%zu,\"nodes_b\":%zu,\"diffs\":[", n_seqs, first_bad_step, rec_a.size(), rec_b.size()); size_t n = rec_a.size() < rec_b.size() ? rec_a.size() : rec_b.size(); @@ -293,8 +284,7 @@ int main(int argc, char ** argv) { verdict = "misaligned"; } else if (A.op == "GATED_DELTA_NET" && A.gdn_tokens == B.gdn_tokens && !A.data.empty() && !B.data.empty()) { - // packed GDN outputs put all token outputs before all sequence states, so seq 0's - // state moves when the number of sequences changes + // packed GDN outputs put all token outputs before all sequence states, so seq 0's state moves when the number of sequences changes const size_t output = A.ne[0]*A.gdn_tokens; const size_t state = A.ne[0]*A.ne[1]/A.gdn_seqs - output; for (size_t k = 0; k < output + state; ++k) { @@ -325,7 +315,6 @@ int main(int argc, char ** argv) { } else if (A.data.empty() || B.data.empty()) { verdict = "too-large"; } else { - // compare element (.., i_tdim = 0, ..) across all other indices int64_t st[4] = {1, A.ne[0], A.ne[0]*A.ne[1], A.ne[0]*A.ne[1]*A.ne[2]}; int64_t stb[4] = {1, B.ne[0], B.ne[0]*B.ne[1], B.ne[0]*B.ne[1]*B.ne[2]}; for (int64_t i3 = 0; i3 < A.ne[3]; ++i3) diff --git a/scripts/unsloth/additive_merge.py b/scripts/unsloth/additive_merge.py index d930158a542..13e4784b775 100644 --- a/scripts/unsloth/additive_merge.py +++ b/scripts/unsloth/additive_merge.py @@ -94,11 +94,8 @@ def nonblank(lines: list[str]) -> list[str]: return [ln.strip() for ln in lines if ln.strip()] -# A line that only opens or closes a block. Two independent case arms share these by -# construction, so finding them on both sides says nothing about the two sides adding the -# same construct: treating them as shared is what refused the real PROJECTOR_TYPE_KIMIK3 / -# PROJECTOR_TYPE_DEEPSEEK4V add/add in tools/mtmd/clip.cpp. Deliberately narrow: brackets, -# semicolons and commas around at most one bare block-terminating keyword. +# a line that only opens or closes a block. Two independent case arms share these by construction, so +# treating them as shared refused the real PROJECTOR_TYPE_KIMIK3 / _DEEPSEEK4V add/add in tools/mtmd/clip.cpp. STRUCTURAL = re.compile(r"^[\s{}()\[\];,]*(?:break|continue|return|pass)?[\s{}()\[\];,]*$") @@ -107,7 +104,6 @@ def identifying(lines: list[str]) -> set[str]: return {ln for ln in nonblank(lines) if not STRUCTURAL.match(ln)} -# `case FOO:`, `case FOO :`, `default:`; a fallthrough label may carry no body at all CASE_LABEL = re.compile(r"^(?:case\s+[^:]+|default\s*):") @@ -139,23 +135,17 @@ def resolve_region(ours: list[str], base: list[str], theirs: list[str]) -> list[ return list(ours) ours_arms, theirs_arms = case_arms(ours), case_arms(theirs) if ours_arms and theirs_arms and ours_arms.isdisjoint(theirs_arms): - # Both sides added case arms and no label is on both, so they are two constructs - # and any line they share is body text: the real clip.cpp collision has arms that - # both set `hparams.rope_theta = 10000.0f;`. The same change made twice would keep - # its label and land in the check below, so this is the one place a shared line is - # allowed - a duplicated label would not even compile. + # both sides added case arms and no label is on both, so any line they share is body text; the same change made twice would keep its label, so sharing is allowed only here return list(theirs) + list(ours) shared = identifying(ours) & identifying(theirs) if shared: # Overlapping content is the signature of one construct added twice, # not two independent additions. Unioning it would duplicate code. - # scaffolding is excluded above, so what is left is content both sides wrote raise Unresolvable( "both sides add the same line(s), so this is one change made twice: " + ", ".join(sorted(shared)[:3]) ) if not identifying(ours) or not identifying(theirs): - # one side is all scaffolding, so there is no content to tell the additions apart raise Unresolvable( "one side adds only block scaffolding, so the two additions cannot " "be told apart" diff --git a/scripts/unsloth/feature_matrix.py b/scripts/unsloth/feature_matrix.py index d8fad55d469..6c0d7f357cd 100644 --- a/scripts/unsloth/feature_matrix.py +++ b/scripts/unsloth/feature_matrix.py @@ -37,7 +37,6 @@ import sys from pathlib import Path -# output that means "this did not run" from a process that exited 0 SKIP_RE = re.compile(r"\bSKIP\b|not supported|unsupported|no tests|0 tests", re.I) @@ -78,7 +77,6 @@ def probe_arch(check: dict, b: Path, gpu: bool) -> str: rc, out = run([str(b / "test-llama-archs"), "-a", arch, "-s", "1234"], b, gpu) if rc != 0: raise Unproven(f"test-llama-archs -a {arch} exited {rc}") - # the arch's own rows, not the header and not another arch's rows = [ln for ln in out.splitlines() if ln.strip().startswith("|") and f"|{arch:>16}|" in ln or (ln.strip().startswith("|") and ln.split("|")[1].strip() == arch)] if not rows: @@ -119,8 +117,7 @@ def probe_mtmd(check: dict, b: Path, gpu: bool) -> str: m = re.search(r"assertions\s*:\s*(\d+)", out) if not m or int(m.group(1)) == 0: raise Unproven("test_projector_registry ran no assertions; the filter matched nothing") - # the test walks the whole enum, so it proves the table is sound; that this projector - # is IN the enum is pin_contract.py's job + # the test walks the whole enum, so it proves the table is sound; that this projector is IN the enum is pin_contract.py's job return f"projector registry intact over {m.group(1)} assertions" @@ -173,7 +170,6 @@ def main() -> int: print(f"ok {name}: " + "; ".join(r["evidence"] for r in entry["results"]) + (f" [{len(entry['deferred'])} needs a GPU]" if entry["deferred"] else "")) else: - # nothing shown either way: not a failure, but not an ok line either print(f"-- {name}: nothing provable without a GPU " f"({len(entry['deferred'])} check(s) deferred)") @@ -187,7 +183,6 @@ def main() -> int: if failed: print(f"\n{failed} feature(s) could not be shown to work", file=sys.stderr) return 1 - # say what was NOT proven alongside what was, or green starts to read as covered tail = f", {deferred} check(s) need a GPU and were not run" if deferred else "" print(f"\nall {len(report['features'])} features demonstrated" + (" on GPU" if args.gpu else " on CPU") + tail) diff --git a/scripts/unsloth/pin_contract.py b/scripts/unsloth/pin_contract.py index d6a60746098..9eebe824486 100644 --- a/scripts/unsloth/pin_contract.py +++ b/scripts/unsloth/pin_contract.py @@ -57,28 +57,22 @@ r"^https://github\.com/([^/]+)/llama\.cpp/pull/(\d+)/commits/([0-9a-f]{40})/?$" ) -# Identifier families that name a FEATURE, not every new symbol: a renamed helper is not a -# lost feature, but a missing LLM_ARCH_ entry always is. +# identifier families that name a FEATURE, not every new symbol: a renamed helper is not a lost feature, a missing LLM_ARCH_ entry always is SYMBOL_FAMILIES = ( "LLM_ARCH_", "LLM_TENSOR_", "LLM_KV_", "LLM_TYPE_", "PROJECTOR_TYPE_", "GGML_OP_", "GGML_TYPE_", "LLAMA_FTYPE_", ) SYMBOL_RE = re.compile(r"\b(?:" + "|".join(SYMBOL_FAMILIES) + r")[A-Z0-9_]+\b") -# the subset naming a whole feature; only to keep --emit readable, the check uses them all HEADLINE = ("LLM_ARCH_", "GGML_OP_", "GGML_TYPE_", "PROJECTOR_TYPE_", "LLAMA_FTYPE_") -# a line worth tracking for survival: comments and short punctuation drift with every -# reformat, a substantial code line does not move on its own TRIVIAL_RE = re.compile(r"^\s*(?://|/\*|\*|\*/|#\s|$)") MIN_LINE = 12 -# Comments are stripped first: a pin that merely NAMES an arch in a comment has not -# registered it, and holding the wording as a contract fails when upstream rewords it. -# Observed on unslothai#70, whose comment mentions GGML_OP_SSM_SCAN to say it is not used. +# comments are stripped first: a pin that merely NAMES an arch in a comment has not registered it, and +# holding the wording as a contract fails when upstream rewords it (unslothai#70) COMMENT_RE = re.compile(r"//.*$|/\*.*?\*/|(? dict: if code: symbols[cur].update(SYMBOL_RE.findall(code)) - # only symbols the base does not already have in that file are evidence of this pin new_symbols: dict[str, list[str]] = {} for path, names in symbols.items(): fresh = sorted(n for n in names @@ -298,8 +291,6 @@ def main() -> int: if args.emit: report["pins"].append(entry) - # only feature-naming families are printed; all are still checked, but a - # hundred LLM_TENSOR_ names would bury the one that matters sym = sorted({s for v in contract["symbols"].values() for s in v if s.startswith(HEADLINE)}) print(f"{name:>18} {entry['line_count']:>5} lines, " @@ -336,7 +327,6 @@ def main() -> int: if args.emit: return 0 - # notices after the verdict lines: housekeeping must not read as a failure for n in notices: print(f"note {n}") if failed: diff --git a/scripts/unsloth/test_additive_merge.py b/scripts/unsloth/test_additive_merge.py index dd2b22c6cac..70da225256d 100644 --- a/scripts/unsloth/test_additive_merge.py +++ b/scripts/unsloth/test_additive_merge.py @@ -114,7 +114,6 @@ def run(repo, *extra): txt.count("} break;") == 2 and txt.count("clip_graph_kimik3") == 1, txt) # --- 3b2. two case arms that share a body line, which is a coincidence ------ -# clip.cpp after upstream landed DEEPSEEK4V: both arms set the same rope_theta base = "switch (t) {\n}\n" ours = ("switch (t) {\n case PROJECTOR_TYPE_KIMIK3:\n {\n" " hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;\n" diff --git a/scripts/unsloth/test_pin_contract.py b/scripts/unsloth/test_pin_contract.py index e76586b1fc4..628ab85e49c 100644 --- a/scripts/unsloth/test_pin_contract.py +++ b/scripts/unsloth/test_pin_contract.py @@ -98,7 +98,6 @@ def run(repo, pr_set, *extra): check("intact merge reports no notices", rep["notices"] == [], rep) # --- 2. the arm is dropped from ONE file: a tree-wide grep would pass ------ -# the real shape: LLM_ARCH_INKLING survives in the enum, the dispatch arm is gone repo, pr_set, sha = make_repo() p = repo / "src" / "llama-model.cpp" p.write_text(MODEL_CPP_BASE) @@ -127,8 +126,6 @@ def run(repo, pr_set, *extra): any("do_the_banded_thing" in x for x in rep["pins"][0]["problems"]), rep) # --- 5. redundancy: the base already has everything the pin adds ---------- -# as it happens for real: upstream lands the same work, so the base tag has it and -# the pin is not an ancestor of anything d = Path(tempfile.mkdtemp(prefix="pc_")) git(d, "init", "-q", "-b", "main") (d / "src").mkdir() @@ -162,8 +159,7 @@ def run(repo, pr_set, *extra): rep["pins"][0]["added_files"] == ["src/inkling.cpp"], rep) # --- 7. a comment is not a contract --------------------------------------- -# unslothai#70 names GGML_OP_SSM_SCAN in a comment to say it does NOT use it, and -# holding that wording would fail the moment upstream rewords it +# unslothai#70 names GGML_OP_SSM_SCAN in a comment to say it does NOT use it, and holding that wording would fail on a reword repo, pr_set, sha = make_repo() git(repo, "checkout", "-q", "pin") (repo / "src" / "note.cpp").write_text( diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 4fc1c808631..62912596ce9 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -547,8 +547,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, llama_seq_id last_seq_id = -1; - // [TAG_EXACT_CONCURRENCY] tokens left in the first set taken, when isolating: only sets with - // the same count join it, so that every set in the ubatch finishes in this ubatch + // [TAG_EXACT_CONCURRENCY] tokens left in the first set taken, when isolating: only sets with the same count join it, so every set in the ubatch finishes in it uint32_t n_left_first = 0; // determine the non-overlapping sequence sets participating in this ubatch @@ -573,13 +572,8 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } if (add) { - // [TAG_EXACT_CONCURRENCY] a set with more tokens left than a decode step carries is a - // prompt, and a prompt shares its arithmetic with whatever else is in the ubatch, so - // give it one of its own. Sets at or below that width are decode steps, kept exact by - // the backend's column policy, so keep grouping them or one prompt would serialize - // every concurrent decode. Grouped sets must have the same number of tokens left, or - // the equal-length expansion below would cut a longer set in two and a memory that - // reduces over a chunk of tokens would sum in a different order than the solo run. + // [TAG_EXACT_CONCURRENCY] a set with more tokens left than a decode step carries is a prompt and gets an + // ubatch of its own; grouped sets need equal tokens left, or the expansion below changes their sum order if (isolate_seqs_above > 0) { uint32_t n_left = 0; @@ -607,8 +601,6 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } else if (n_left != n_left_first) { continue; } else if ((cur_seq_set.size() + 1) * n_left_first > n_ubatch) { - // one more set would not finish here, and the expansion below would then cut - // every set part way: the chunking this guard exists to prevent break; } } diff --git a/src/llama-batch.h b/src/llama-batch.h index f0ecc9d8407..ddf05843d10 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -105,13 +105,10 @@ class llama_batch_allocr { // make ubatches of equal-length sequences sets // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch - // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this - // many tokens left to place is a prompt and gets a ubatch of its own; sets at or - // below it are decode steps and stay grouped, so a prompt does not serialize them + // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this many tokens left is a prompt and gets a ubatch of its own llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above = 0); - // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than n_tokens left to place, - // i.e. what remains of the batch holds a prompt rather than decode steps only + // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than n_tokens left to place, i.e. what remains of the batch holds a prompt bool has_seq_wider_than(uint32_t n_tokens) const; // [TAG_EXACT_CONCURRENCY] true if some token carries more than one sequence id diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 30353028f26..7f81f570ad1 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -102,15 +102,9 @@ llama_context::llama_context( throw std::runtime_error("n_seq_max must be <= " + std::to_string(LLAMA_MAX_SEQ)); } - // [TAG_EXACT_CONCURRENCY] the widest decode step this context can build: one column per - // sequence times the tokens a sequence contributes, reported so a backend splitting columns - // covers it (a caller that builds wider steps uses llama_set_exact_decode_width). The - // sequence count is what is reported, so a later rise in the tokens figure follows it here - // too. Checked now but reported at the end of the constructor, so a construction that fails - // later does not leave a width behind that no context needs. + // [TAG_EXACT_CONCURRENCY] the widest decode step this context can build, reported so a backend that splits columns covers it; reported at the end of the constructor if (llama_exact_concurrency()) { - // an explicit column bound wins in the backend, so one below this context's width would - // leave decodes batched above it; the report refuses that, and that is an error here + // an explicit column bound below this context's width would leave decodes batched above it, so the report refuses it if (!llama_exact_check_n_seq(cparams.n_seq_max)) { throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); } @@ -409,8 +403,7 @@ llama_context::llama_context( memory.reset(model.create_memory(params_mem, cparams)); - // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a non-causal context with a - // cache would assert on its first graph + // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a non-causal context with a cache would assert on its first graph if (llama_exact_concurrency() && memory && !cparams.causal_attn) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so it cannot be created with non-causal attention\n", __func__); throw std::runtime_error("exact concurrency: non-causal attention is not supported with a KV cache"); @@ -499,8 +492,7 @@ llama_context::llama_context( } } - // [TAG_EXACT_CONCURRENCY] nothing above can fail now, so publish the width; already checked - // against the explicit bound at the top, so a refusal here means the bound moved + // [TAG_EXACT_CONCURRENCY] nothing above can fail now, so publish the width; a refusal here means the bound moved if (llama_exact_concurrency() && !llama_exact_report_n_seq(cparams.n_seq_max)) { throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); } @@ -510,10 +502,7 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); - // A transfer still alive is drained first: synchronize() covers the graph backends, - // not the copy backend a transfer owns, and the KV buffers it may still be reading or - // writing are about to go. It is then let go of, so freeing it later touches nothing - // of this context. + // a transfer still alive is drained first: synchronize() covers the graph backends, not the copy backend a transfer owns, and its KV buffers are about to go state_seq_copies_drain(); for (auto & it : state_copy_fences) { @@ -1226,8 +1215,7 @@ void llama_context::set_causal_attn(bool value) { return; } - // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a context with a cache keeps - // causal attention rather than asserting in the next graph + // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a context with a cache keeps causal attention rather than asserting in the next graph if (!value && memory && llama_exact_concurrency()) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so causal attention cannot be turned off; the change is refused\n", __func__); return; @@ -2611,14 +2599,7 @@ class llama_io_write_dummy : public llama_io_write_i { size_t size_written = 0; }; -// [TAG_STATE_COALESCE] one transfer per run of cells, not one per cell -// -// A sequence's state is emitted in cell order, so a run of cells that is contiguous in the -// cache is contiguous both in the tensor and in the host buffer, and the fragments covering -// it are one transfer. The save side already coalesces its cells into ranges before it emits -// them; the restore side does not, and asks for one transfer per cell even when the cells it -// was given are a handful of long runs. Merging here fixes both sides at once, and covers -// the transposed V layout, where the same runs are emitted once per embedding row. +// [TAG_STATE_COALESCE] one transfer per run of cells, not one per cell; the restore side asks for one per cell, and the transposed V layout repeats every run once per row template static size_t llama_io_run_end(const std::vector & infos, size_t i) { size_t end = i + 1; @@ -2644,19 +2625,9 @@ static size_t llama_io_run_size(const std::vector & infos, size_t i, siz return size; } -// [TAG_STATE_COALESCE] runs of one length at a constant stride are a single strided copy -// -// Sequences sharing a unified cache take their cells in turn, so a sequence's cells are not -// one block but a regular comb: a few cells, a gap, a few cells, for as long as the sequence -// is. Merging adjacent cells still leaves hundreds of runs per tensor, and at a few -// microseconds to post each one that is tens of milliseconds spent issuing copies. A comb is -// exactly what a strided copy describes, so one call replaces a whole group of runs. -// -// emit(tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); n_copies == 1 means -// an ordinary contiguous transfer and the strides are not meaningful. +// [TAG_STATE_COALESCE] a comb of equal runs at a constant stride is one strided copy: sequences sharing a unified cache take their cells in turn template static void llama_io_emit(const std::vector & infos, size_t first, size_t last, emit_t emit) { - // the runs of adjacent cells, as index ranges into infos std::vector> runs; for (size_t i = first; i < last; ) { @@ -2805,16 +2776,10 @@ class llama_io_read_host : public llama_io_read_i { while (end < rinfos.size() && rinfos[end].tensor == tensor) { end++; } - // [TAG_STATE_COALESCE] the fragments the restore emits are one per cell; what - // matters is how many runs of adjacent cells they form, because that is how many - // transfers they actually cost. Count the runs first, and only fall back to - // staging the whole tensor when even the runs are too many. + // [TAG_STATE_COALESCE] the restore emits one fragment per cell, but the cost is the number of runs of adjacent cells, so count runs before falling back to staging const size_t tensor_bytes = ggml_nbytes(tensor); auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - // A strided set of rows is one transfer on a buffer that copies 2-D, and one - // per row on one that does not (the generic path expands it), so it is counted - // by what it costs on this buffer, not by the calls it makes. const bool has_2d = ggml_backend_buffer_supports_2d(buffer); size_t n_runs = 0; @@ -2822,17 +2787,12 @@ class llama_io_read_host : public llama_io_read_i { [&n_runs, has_2d](ggml_tensor *, const uint8_t *, size_t, size_t, size_t n_copies, size_t, size_t) { n_runs += has_2d ? 1 : n_copies; }); - // A fragmented sequence can require thousands of synchronous device - // transfers per layer. For bounded tensors, stage the tensor once and - // preserve every byte belonging to other sequences. Bound scratch RAM - // and leave ordinary contiguous transfers on their original fast path. if (n_runs >= 64 && tensor_bytes <= 64 * 1024 * 1024 && !ggml_backend_buffer_is_host(buffer)) { std::vector staging; try { staging.resize(tensor_bytes); } catch (const std::bad_alloc &) { - // fall back to the individual transfers below } if (!staging.empty()) { ggml_backend_tensor_get(tensor, staging.data(), 0, tensor_bytes); @@ -3198,12 +3158,7 @@ size_t llama_context::state_set_data(const uint8_t * src, size_t size) { } } -// [TAG_STATE_ASYNC] a sequence state transfer that runs beside the decode instead of in it -// -// Everything the transfer needs to outlive the call that issued it lives here: the host -// buffer the bytes land in or come from, one backend per device holding part of the cache -// (each with a stream of its own, so the copies never queue behind the graphs), and one -// event per device to tell the caller when its half is finished. +// [TAG_STATE_ASYNC] a sequence state transfer that runs beside the decode instead of in it: the host buffer, one backend per device, each with its own stream, and one event per device struct llama_state_seq_copy { llama_context * ctx = nullptr; @@ -3225,9 +3180,7 @@ struct llama_state_seq_copy { bool pinned = false; bool can_pin = false; - // transfers the last issue actually posted, i.e. runs of adjacent cells over all tensors size_t n_copies = 0; - // microseconds the last issue spent draining the compute streams before it could start int64_t t_sync_us = 0; ~llama_state_seq_copy() { @@ -3244,10 +3197,7 @@ struct llama_state_seq_copy { } } - // The stream this tensor is copied on, or null when it needs no stream: tensors already - // in host memory are a memcpy, and a tensor in a split or otherwise non-default buffer - // fails the buffer check every backend's async copy asserts, so both take the plain - // synchronous path. Handing a backend out marks it, so record() knows which ones ran. + // the stream this tensor is copied on, or null when it needs none: a host tensor is a memcpy, and a split buffer fails every backend's async copy assert ggml_backend_t backend_for(const ggml_tensor * t) { ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; @@ -3274,7 +3224,6 @@ struct llama_state_seq_copy { return it->second.backend.get(); } - // close every stream the transfer just used void record() { for (auto & it : devs) { @@ -3284,16 +3233,7 @@ struct llama_state_seq_copy { } } - // Order the copies about to be posted behind the compute already queued on each device: - // the decode that produced the cells a park reads, or that a restore's cells were - // carved out of, has to be finished before the copy touches them. The copy stream waits - // for the context's fence on its device, an event the context records on the compute - // stream at the end of every decode, so the host drains nothing. The fence is recorded - // there and not here: recorded here, it would land behind the waits that order_before() - // queued for the restores issued earlier in the same pass, and each restore would then - // wait for the previous one's copies. Draining the host (synchronize()) is what this - // replaced: with those same waits on the compute stream, a host drain blocked this thread - // until the previous restore had landed. + // order the copies behind the compute already queued on each device: the copy stream waits for the context's fence, recorded at the end of every decode void order_after(const std::map & fences) { for (auto & it : devs) { const auto fence = fences.find(it.first); @@ -3304,16 +3244,7 @@ struct llama_state_seq_copy { } } - // Order the context's compute behind the copies just recorded, on the device: every - // backend the graphs run on waits for the event of the transfer on its device before - // the next graph it is given. This is a stream wait, not a host wait, so the caller's - // thread carries on and the decode it issues next starts the moment the copy lands. - // - // Needed for a restore and only a restore: its copies write cells of the KV cache - // while other sequences keep decoding, and an attention that is not paged reads every - // cell up to n_kv, masked ones included, so without this the reads and the writes are - // unordered. A park reads cells nobody writes until it has landed, and the decode that - // produced them has been drained by the synchronize() at the top of the issue. + // order the context's compute behind the copies just recorded, for a restore only: its copies write KV cells while other sequences read every cell up to n_kv void order_before(const std::vector & compute) { for (auto & it : devs) { if (!it.second.pending) { @@ -3358,11 +3289,7 @@ struct llama_state_seq_copy { } } - // Grow-only. Pinning host memory is expensive -- a hundred MiB of it costs about as long - // as the copy it is for -- and a caller that parks the same sequence over and over asks - // for a slightly different size every time, so freeing between transfers would put that - // cost back on the very loop this is keeping clear. The memory is given back by - // buf_free() when the caller is finished with the slot, not between two of its parks. + // grow-only: pinning host memory costs about as long as the copy it is for, and a caller parking the same sequence asks for a slightly different size each time uint8_t * buf_resize(size_t size_new) { if (size_new <= capacity) { size = size_new; @@ -3400,8 +3327,7 @@ struct llama_state_seq_copy { data = base; size = size_new; capacity = size_new; - // a host buffer type may quietly hand back ordinary memory when pinning is turned - // off, so believe the buffer that came back rather than the type that was asked + // a host buffer type may quietly hand back ordinary memory when pinning is off, so believe the buffer that came back rather than the type pinned = can_pin && ggml_backend_buffer_get_type(buf) == host_buft; return data; @@ -3418,9 +3344,6 @@ struct llama_state_seq_copy { pinned = false; } - // Pinned host memory is the point of allocating through the backend at all: a copy in or - // out of pageable memory is staged through a pinned bounce buffer by the driver and - // blocks, which is exactly the stall being removed here. ggml_backend_buffer_type_t host_buffer_type() { for (auto & it : devs) { ggml_backend_buffer_type_t buft = ggml_backend_dev_host_buffer_type(it.first); @@ -3439,10 +3362,7 @@ class llama_io_write_host_async : public llama_io_write_i { llama_io_write_host_async(uint8_t * p, size_t len, llama_state_seq_copy & cpy) : ptr(p), buf_size(len), cpy(cpy) {} - // The transfers are posted from the destructor, and only once serialisation has got to - // the end: a failure part way, a buffer one byte short say, is reported to the caller as - // a zero return, and a caller told that is free to reuse the buffer at once. Copies - // posted regardless would still be reading it. + // posted from the destructor, and only once serialisation reached the end: a caller told of a partial failure by a zero return is free to reuse the buffer at once void commit() { committed = true; } @@ -3513,9 +3433,7 @@ class llama_io_read_host_async : public llama_io_read_i { llama_io_read_host_async(const uint8_t * p, size_t len, llama_state_seq_copy & cpy) : ptr(p), buf_size(len), cpy(cpy) {} - // see llama_io_write_host_async::commit(): the restore that failed part way has already - // dropped the sequence, and copies posted for it would write into cells that are no - // longer its own + // see llama_io_write_host_async::commit(): a restore that failed part way has dropped the sequence, and copies posted for it would write cells that are no longer its own void commit() { committed = true; } @@ -3525,15 +3443,7 @@ class llama_io_read_host_async : public llama_io_read_i { return; } - // No whole-tensor staging here, unlike the synchronous path above. Staging reads a - // tensor, patches this sequence's bytes into the host copy and writes the whole - // tensor back, which preserves the neighbours only while nothing else is touching - // the cache. These copies are issued precisely so that decoding can carry on beside - // them, so a write-back would undo whatever the sequences sharing the tensor wrote - // to their own cells in the meantime. Writing only this sequence's runs cannot: - // every byte in them belongs to the sequence being restored. That is affordable - // because the runs have been coalesced -- one transfer per run of adjacent cells, - // which is what staging was working around in the first place. + // no whole-tensor staging here, unlike the synchronous path above: a write-back would undo whatever the sequences sharing the tensor wrote while these copies ran llama_io_emit(rinfos, 0, rinfos.size(), [this](ggml_tensor * tensor, const uint8_t * ptr, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { @@ -3716,13 +3626,7 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } - // A device that advertises events but does not implement event_query is no use - // here. ggml_backend_event_query() then answers the only way it can, by waiting for - // the event, so the first poll of a transfer blocks the caller for the whole copy -- - // the very stall this exists to remove, except that the caller has been told the - // copy is asynchronous and has stopped looking for it. Such a device is left out, so - // that state_seq_copy_init() returns NULL and the caller keeps the synchronous calls - // it already had. + // a device that advertises events but does not implement event_query makes the first poll wait for the whole copy, so leave it out and let state_seq_copy_init() return NULL if (!ggml_backend_dev_supports_event_query(dev)) { static std::atomic warned(false); @@ -3734,9 +3638,7 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } - // a backend of its own, not the one the graphs are computed on: that one moves its - // copies to whichever stream it is currently using, so a transfer posted to it could - // end up ordered behind a graph -- which is the stall this exists to avoid + // a backend of its own, not the one the graphs are computed on: that one moves its copies to whichever stream it is using, so a transfer could end up ordered behind a graph ggml_backend_t backend_cpy = ggml_backend_dev_init(dev, nullptr); if (!backend_cpy) { @@ -3760,11 +3662,7 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { return nullptr; } - // The devices above are the ones the graphs run on, not necessarily the ones the state - // lives on: with most layers left on the CPU the KV cache is host memory, and a tensor - // there takes the synchronous branch of backend_for(). A transfer whose every copy would - // do that is not asynchronous, whatever it is called, and the caller is better served by - // the synchronous calls it already has and a log line that says so. + // the devices above are the ones the graphs run on, not the ones the state lives on: with most layers on the CPU every copy takes the synchronous branch of backend_for() if (memory) { bool on_device = false; @@ -3792,10 +3690,7 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { cpy->can_pin = cpy->host_buffer_type() != ggml_backend_cpu_buffer_type(); - // One fence per device, shared by every transfer on this context and recorded after - // every decode from now on. Installed only here, after the checks above: a transfer - // refused for its layout must leave nothing behind that every later decode would keep - // recording for nobody. + // one fence per device, shared by every transfer and recorded after every decode; installed only after the checks above, so a refused transfer leaves nothing behind std::vector fences_new; for (const auto & it : cpy->devs) { @@ -3818,7 +3713,6 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { fences_new.push_back(it.first); } - // the fences say where the compute streams are now, before any transfer asks state_seq_copy_fence(); state_copies.insert(cpy.get()); @@ -3828,28 +3722,18 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { } size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { - // Unlike the legacy API the library owns this buffer, so the extent the io object is - // built with can be checked instead of believed. Every bounds check inside that object - // validates against the extent it was given, so a size larger than the allocation makes - // all of them agree with the caller and the copy runs past the buffer. + // the library owns this buffer, so the extent can be checked instead of believed: every bounds check validates against it, so an oversized one agrees and the copy overruns if (!cpy.data || size == 0 || size > cpy.size) { LLAMA_LOG_ERROR("%s: cannot cover %zu bytes, the transfer's buffer holds %zu\n", __func__, size, cpy.size); return 0; } - // LLAMA_STATE_SEQ_FLAGS_ON_DEVICE asks for the tensor data to be left in device buffers, - // and this path has nowhere to leave it: it serialises through the host buffer it owns, - // which is the whole point of it. llama_state_seq_get_size_ext() with that flag reports - // a metadata-sized state, so a caller pairing the two would size a buffer for one thing - // and fill it with another; the synchronous calls serve that flag. + // LLAMA_STATE_SEQ_FLAGS_ON_DEVICE has nowhere to leave the data here, and get_size_ext() with that flag reports a metadata-sized state, so the two cannot be paired if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { LLAMA_LOG_ERROR("%s: LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is not supported here, the copies go through host memory\n", __func__); return 0; } - // The copies run on their own stream, so the decode that produced these cells has to be - // finished before they are read: the copy stream waits for the compute stream, on the - // device, see order_after(). Nothing stays on the caller's thread. const int64_t t_sync = ggml_time_us(); cpy.order_after(state_copy_fences); cpy.t_sync_us = ggml_time_us() - t_sync; @@ -3874,28 +3758,17 @@ size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size } size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { - // Unlike the legacy API the library owns this buffer, so the extent the io object is - // built with can be checked instead of believed. Every bounds check inside that object - // validates against the extent it was given, so a size larger than the allocation makes - // all of them agree with the caller and the copy runs past the buffer. if (!cpy.data || size == 0 || size > cpy.size) { LLAMA_LOG_ERROR("%s: cannot cover %zu bytes, the transfer's buffer holds %zu\n", __func__, size, cpy.size); return 0; } - // LLAMA_STATE_SEQ_FLAGS_ON_DEVICE asks for the tensor data to be left in device buffers, - // and this path has nowhere to leave it: it serialises through the host buffer it owns, - // which is the whole point of it. llama_state_seq_get_size_ext() with that flag reports - // a metadata-sized state, so a caller pairing the two would size a buffer for one thing - // and fill it with another; the synchronous calls serve that flag. if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { LLAMA_LOG_ERROR("%s: LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is not supported here, the copies go through host memory\n", __func__); return 0; } - // the cells this restore was given may still be read by a graph in flight (masked, but - // read), so the copy stream waits for the compute stream before it writes them: on the - // device, see order_after(), rather than by draining the compute stream on this thread + // the cells this restore was given may still be read, masked, by a graph in flight, so the copy stream waits for the compute stream on the device, see order_after() const int64_t t_sync = ggml_time_us(); cpy.order_after(state_copy_fences); cpy.t_sync_us = ggml_time_us() - t_sync; @@ -3926,8 +3799,6 @@ size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size } } - // the adapter has posted the copies and recorded the events on its way out; the - // graphs that follow on these devices wait for them, see order_before() cpy.order_before(backends); return n; @@ -4088,9 +3959,7 @@ size_t llama_context::state_write_data(llama_io_write_i & io) { } size_t llama_context::state_read_data(llama_io_read_i & io) { - // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical - // index, which the paged pool owns. Refused before anything is parsed, so the caller's cache - // is left as it was: the generic restore path clears it on failure. + // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical index, which the paged pool owns; refused before anything is parsed if (memory && memory->alloc_granularity() > 1) { throw std::runtime_error("whole-context restore is not supported with LLAMA_EXACT_CONCURRENCY, restore per sequence"); } @@ -5041,9 +4910,7 @@ void llama_state_seq_copy_buf_free(llama_state_seq_copy * cpy) { } bool llama_state_seq_copy_buf_is_pinned(llama_state_seq_copy * cpy) { - // what was allocated, not what could be: a host buffer type is free to hand back - // ordinary memory, which is what CUDA does under GGML_CUDA_NO_PINNED, and there is - // nothing page-locked before the first resize or after buf_free() + // what was allocated, not what could be: a host buffer type is free to hand back ordinary memory, as CUDA does under GGML_CUDA_NO_PINNED return cpy->pinned; } diff --git a/src/llama-context.h b/src/llama-context.h index 106dc3922e6..496ddcb05a2 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -166,14 +166,11 @@ struct llama_context { size_t state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags); size_t state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id dest_seq_id, llama_state_seq_flags flags); - // [TAG_STATE_ASYNC] mark the point the compute streams have reached, for the copies to - // wait for; recorded after every decode and encode once a transfer exists + // [TAG_STATE_ASYNC] mark the point the compute streams have reached, for the copies to wait for; recorded after every decode and encode once a transfer exists void state_seq_copy_fence(); - // a transfer letting go of this context: the last one takes the fences with it void state_seq_copy_release(llama_state_seq_copy * cpy); - // at teardown: wait for every live transfer and let it go void state_seq_copies_drain(); bool state_load_file( @@ -368,13 +365,10 @@ struct llama_context { ggml_backend_t backend_cpu = nullptr; std::vector backends; - // [TAG_STATE_ASYNC] one event per device that copies asynchronously, recorded on the - // compute stream at the end of every decode; see state_seq_copy_fence() + // [TAG_STATE_ASYNC] one event per device that copies asynchronously, recorded on the compute stream at the end of every decode; see state_seq_copy_fence() std::map state_copy_fences; - // transfers alive on this context; the fences go when the last one does, so a server - // that made transfers and then gave them up records nothing after its decodes, and a - // context freed with transfers still alive drains them and lets them go first + // transfers alive on this context; the fences go when the last one does, and a context freed with transfers still alive drains them and lets them go first std::set state_copies; // training diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index dc02560fc66..e9fd4ecd6e8 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -27,11 +27,7 @@ // dedup helpers -// [TAG_EXACT_CONCURRENCY] the page table is wired into llm_graph_input_attn_kv only, so a V-less -// layout would have its cells placed in pages and then attend in physical order anyway, with the -// mode reporting itself as on. Refuse the context instead: wiring self_pages into -// llm_graph_input_attn_k alone fixes one of the four V-less classes, and DeepSeek 3.2 uses two of -// them, so the model would end up half paged, which is worse than refused. +// [TAG_EXACT_CONCURRENCY] the page table is wired into llm_graph_input_attn_kv only, so a V-less layout would attend in physical order with the mode reporting itself on static void llm_graph_reject_exact_concurrency(const char * layout) { if (!llama_exact_concurrency()) { return; diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index f577414a51c..ee35b914d72 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -187,18 +187,14 @@ bool llama_exact_concurrency() { // [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h static std::atomic g_exact_decode_tokens{1}; -// one lock for the token figure, the sequence count and the width, since the three move together -// and a report interleaved with a change of figure could leave the backend with a width that -// covers neither; recursive, since the setters call each other +// one lock for the token figure, the sequence count and the width: a report interleaved with a change of figure could leave the backend with a width that covers neither static std::recursive_mutex g_exact_mutex; -// the most sequences any context was created with. The tokens figure is process wide, so raising -// it widens every existing context's decode step and their width is re-reported with it. +// the most sequences any context was created with; the tokens figure is process wide, so raising it re-reports every context's width static std::atomic g_exact_max_n_seq{0}; static bool llama_exact_width_within_explicit_bound(uint32_t n_cols); -// sequences times tokens, handed to a backend as an int; a product that overflows is refused static bool llama_exact_width_of(uint32_t n_seq, uint32_t n_tokens, uint32_t & n_cols) { const uint64_t w = (uint64_t) n_seq * (uint64_t) n_tokens; @@ -246,14 +242,12 @@ bool llama_set_exact_decode_tokens(uint32_t n_tokens) { std::lock_guard lock(g_exact_mutex); - // never lowered: a narrower context set up later would turn an existing speculative - // context's verify steps into prompts and serialise them + // never lowered: a narrower context set up later would turn an existing speculative context's verify steps into prompts if (n_tokens <= g_exact_decode_tokens.load(std::memory_order_relaxed)) { return true; } - // every context widens with the figure, so report the width first; one the explicit bound - // cannot cover leaves the old figure in place + // every context widens with the figure, so report the width first; one the explicit bound cannot cover leaves the old figure in place const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed); uint32_t n_cols = 0; @@ -271,13 +265,10 @@ uint32_t llama_exact_decode_tokens(void) { return g_exact_decode_tokens.load(std::memory_order_relaxed); } -// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h. Backends read it -// through ggml_backend_cuda_set_exact_decode_width, reached through the registry so an absent or -// late-loaded backend costs nothing. +// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h; reached through the registry so an absent or late-loaded backend costs nothing static std::atomic g_exact_decode_width{0}; -// an explicit column bound wins in the CUDA backend, so a width above it would leave decodes -// batched past the bound: refuse such a width instead of storing it +// an explicit column bound wins in the CUDA backend, so a width above it would leave decodes batched past the bound static bool llama_exact_width_within_explicit_bound(uint32_t n_cols) { static const int explicit_cols = []() { const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); @@ -305,8 +296,6 @@ bool llama_set_exact_decode_width(uint32_t n_cols) { while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { } - // the widest figure goes to every backend on every call, not only when it grew, or a width - // reported before a backend was loaded would never reach it const uint32_t widest = g_exact_decode_width.load(std::memory_order_relaxed); for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { diff --git a/src/llama-impl.h b/src/llama-impl.h index d0ab2e5cf1c..01970f2a88c 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -104,14 +104,10 @@ std::string llama_format_tensor_shape(const struct ggml_tensor * t); std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i); -// [TAG_EXACT_CONCURRENCY] opt-in mode under which a sequence's attention depends only on its own -// cells, in position order, so its output does not change when others share the KV cache. Off by -// default; reads the same variable as the paged KV cache and the CUDA backend. +// [TAG_EXACT_CONCURRENCY] opt-in mode under which a sequence's attention depends only on its own cells, so its output does not change when others share the KV cache bool llama_exact_concurrency(); -// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so the backend -// knows the width every context needs and it follows llama_set_exact_decode_tokens +// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so the backend knows the width every context needs bool llama_exact_report_n_seq(uint32_t n_seq); -// the same check without the report, for a constructor that may still fail after asking bool llama_exact_check_n_seq(uint32_t n_seq); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 428fbc6b567..e7511dfd39e 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -62,9 +62,7 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { // llama_kv_cache // -// [TAG_EXACT_CONCURRENCY] the paged attention specialization lives in the CUDA sources, which are -// also built as ROCm and MUSA. Every other backend ignores src[5] and walks the pool in physical -// cell order, so a KV layer placed there would silently lose the mode's guarantee. +// [TAG_EXACT_CONCURRENCY] the paged specialization lives in the CUDA sources; every other backend ignores src[5] and walks the pool in physical cell order static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) { if (!dev) { return false; @@ -83,11 +81,7 @@ static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) { return strcmp(name, "CUDA") == 0 || strcmp(name, "ROCm") == 0 || strcmp(name, "MUSA") == 0; } -// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a layer -// of this shape. The registry name only says which backends carry the kernels, not that the build -// has them or that this architecture, head width and K/V types land on one; otherwise the op -// falls to the CPU, which ignores the page table. So build the op as the graph does, at the -// widths a decode step, a verify step and a prompt chunk use, and ask the device. +// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a layer of this shape: the registry name only says which backends carry the kernels static bool llama_dev_supports_paged_attn( ggml_backend_dev_t dev, ggml_type type_k, ggml_type type_v, @@ -158,8 +152,7 @@ llama_kv_cache::llama_kv_cache( v_cells_impl(other ? other->v_cells_impl : std::make_shared()), v_cells(*v_cells_impl) { - // [TAG_EXACT_CONCURRENCY] read the knob through the same cached reader the graph and the CUDA - // dispatcher use, so a mid-process change cannot leave the two disagreeing + // [TAG_EXACT_CONCURRENCY] read the knob through the same cached reader the graph and the CUDA dispatcher use, so a mid-process change cannot leave them disagreeing exact_pages = llama_exact_concurrency(); // shared cells view the source cache's K/V tensors, so the cell count @@ -175,8 +168,7 @@ llama_kv_cache::llama_kv_cache( GGML_ASSERT(kv_size % n_pad == 0); - // [TAG_EXACT_CONCURRENCY] all of these are reachable from the command line, so name the one - // that failed instead of aborting on a bare assert + // [TAG_EXACT_CONCURRENCY] all of these are reachable from the command line, so name the one that failed instead of aborting on a bare assert if (exact_pages) { const char * unsupported = nullptr; @@ -321,8 +313,7 @@ llama_kv_cache::llama_kv_cache( LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); - // [TAG_EXACT_CONCURRENCY] the paged kernel handles 256-wide K and V heads only; any other - // width would run unpaged while the mode reports itself as on + // [TAG_EXACT_CONCURRENCY] the paged kernel handles 256-wide K and V heads only; any other width would run unpaged while the mode reports itself as on if (exact_pages && (hparams.n_embd_head_k(il) != 256 || (!is_mla && hparams.n_embd_head_v(il) != 256) || is_mla)) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d has %u-wide K heads and %u-wide V heads%s, " "and the paged attention kernel supports 256-wide K and V heads only\n", @@ -346,8 +337,7 @@ llama_kv_cache::llama_kv_cache( throw std::runtime_error("exact concurrency: KV cache layer is not on the CUDA backend"); } - // [TAG_EXACT_CONCURRENCY] right backend; ask whether this layer's attention, with the - // page table attached, lands on one of its kernels at all + // [TAG_EXACT_CONCURRENCY] right backend; ask whether this layer's attention, with the page table attached, lands on one of its kernels at all if (exact_pages && !llama_dev_supports_paged_attn(model.dev_layer(il), type_k, type_v, hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), hparams.n_head(il), hparams.n_head_kv(il), kv_size, exact_page_size)) { @@ -540,7 +530,6 @@ void llama_kv_cache::exact_pages_sync() const { } if (debug > 0) { - // the incrementally maintained ownership has to say what the cells say const auto kept = exact_page_owner; exact_pages_rebuild(); @@ -556,7 +545,6 @@ void llama_kv_cache::exact_pages_sync() const { // [TAG_EXACT_CONCURRENCY] void llama_kv_cache::exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos) { if (exact_page_owner_dirty || exact_page_owner.empty()) { - // the next sync rebuilds from the cells anyway return; } @@ -662,9 +650,7 @@ void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, ll return; } - // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so cells cannot be shared between - // two. Refuse rather than abort the process, so an uncovered caller gets a failed copy it can - // report. After the shared-cells return, so a draft cache is unaffected. + // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so refuse a copy that would share cells rather than abort. After the shared-cells return, so a draft cache is unaffected. if (exact_pages && seq_id_src != seq_id_dst) { LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between " "sequences (%d -> %d); ignoring the copy\n", @@ -794,8 +780,7 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll return; } - // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo the page size, - // so shifting positions would misplace every cell; say so rather than abort + // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo the page size, so shifting positions would misplace every cell if (exact_pages && shift != 0) { LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions " "(seq %d, shift %d); ignoring the shift\n", @@ -944,8 +929,7 @@ llama_memory_context_ptr llama_kv_cache::init_batch( GGML_UNUSED(embd_all); do { - // [TAG_EXACT_CONCURRENCY] a token shared by several sequences would be one cell in a page - // that belongs to one sequence; the placement asserts on it later, so refuse it here + // [TAG_EXACT_CONCURRENCY] a token shared by several sequences would be one cell in a page that belongs to one sequence, so refuse it here rather than assert at placement if (exact_pages && balloc.has_shared_tokens()) { LLAMA_LOG_ERROR("%s: exact concurrency does not support tokens shared by several sequence ids; " "give every token exactly one sequence id\n", __func__); @@ -956,9 +940,7 @@ llama_memory_context_ptr llama_kv_cache::init_batch( std::vector ubatches; while (true) { - // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt into one ubatch, - // so a prefill would run at a width its solo run never sees. The sequence-set split - // gives each prompt its own ubatch; a plain decode step keeps taking split_simple. + // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt into one ubatch, so a prefill would run at a width its solo run never sees; the set split gives each its own const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0; auto ubatch = n_stream == 1 && !isolate @@ -1011,8 +993,7 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector v_cells; // copy of the old cells, before placing the ubatch - // [TAG_EXACT_CONCURRENCY] page ownership before the ubatch, so undoing a speculative - // placement does not force a rebuild from every cell + // [TAG_EXACT_CONCURRENCY] page ownership before the ubatch, so undoing a speculative placement does not force a rebuild from every cell std::vector exact_page_owner_old; }; @@ -1063,8 +1044,7 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vectorv_heads_old[s]; } - // [TAG_EXACT_CONCURRENCY] put back what the allocator knew, unless the placement also - // removed cells, in which case only the cells can say what is left + // [TAG_EXACT_CONCURRENCY] put back what the allocator knew, unless the placement also removed cells, when only the cells can say what is left if (!exact_page_owner_dirty) { exact_page_owner = it->exact_page_owner_old; } @@ -1227,8 +1207,7 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, } if (exact_pages) { - // ownership is maintained as cells are placed, so this reads one entry per page rather - // than scanning every cell; the claims are local and prepare() can still roll them back + // ownership is maintained as cells are placed, so this reads one entry per page rather than scanning every cell const auto & cells = v_cells[0]; exact_pages_sync(); @@ -1251,7 +1230,6 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, const page_key key {ubatch.seq_id[i][0], ubatch.pos[i]/exact_page_size}; auto it = pages.find(key); if (it == pages.end()) { - // round-robin free-page search, deliberately nonmonotonic in physical order uint32_t page = v_heads[0]/exact_page_size; uint32_t tested = 0; while (tested < owner.size() && owner[page%owner.size()].seq >= 0) { ++page; ++tested; } @@ -1485,15 +1463,12 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & } uint32_t llama_kv_cache::alloc_granularity() const { - // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so n - // tokens hold round_up(n, exact_page_size) cells: the tail page is charged in full + // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so n tokens hold round_up(n, exact_page_size) cells: the tail page is charged in full return exact_pages ? exact_page_size : 1; } bool llama_kv_cache::get_can_shift() const { - // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo 256, so the pool - // cannot shift positions. Reporting it here is what disables --context-shift and - // --cache-reuse at load rather than failing on the first request that needs them. + // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo 256, so the pool cannot shift positions; reporting it disables --context-shift and --cache-reuse at load if (exact_pages) { return false; } @@ -1585,7 +1560,6 @@ void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * uba auto * row = data.data() + i*dst->ne[0]; row[0] = 0; for (const auto & page : pages[ubatch->seq_id[i][0]]) { - // exclude wholly future pages even when prefill includes later query rows if (page.first*exact_page_size > uint32_t(ubatch->pos[i])) { break; } row[++row[0]] = page.second; } @@ -1602,8 +1576,7 @@ void llama_kv_cache_context::set_input_pages(ggml_tensor * dst, const llama_ubat } uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { - // the physical view spans the pool; the per-query page map is the only loop bound for exact - // attention, so neighbours cannot extend it + // the per-query page map is the only loop bound for exact attention, so neighbours cannot extend it if (exact_pages) { return get_size(); } uint32_t result = 0; @@ -2409,9 +2382,7 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla } void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { - // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical index, - // which the paged pool owns. Refused before a byte is read, so the clearing failure path below - // is never entered for it. + // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical index, which the paged pool owns; refused before a byte is read if (exact_pages && seq_id == -1) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state restore only\n", __func__); throw std::runtime_error("whole-cache restore is not supported with LLAMA_EXACT_CONCURRENCY"); diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index c139ede57e0..c36be8c0154 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -243,10 +243,7 @@ class llama_kv_cache : public llama_memory_i { static constexpr uint32_t exact_page_size = 256; bool exact_pages = false; - // [TAG_EXACT_CONCURRENCY] which (sequence, logical page) owns each physical page; seq < 0 means - // free. Kept current as cells are placed and marked dirty by removals, so find_slot() and - // set_input_pages() read one entry per page instead of rebuilding from every live cell twice per - // ubatch. Mutable because set_input_pages() is const. + // [TAG_EXACT_CONCURRENCY] which (sequence, logical page) owns each physical page; seq < 0 means free, and it is kept current as cells are placed and dirtied by removals struct exact_page { llama_seq_id seq = -1; llama_pos lpg = -1; @@ -255,13 +252,10 @@ class llama_kv_cache : public llama_memory_i { mutable std::vector exact_page_owner; mutable bool exact_page_owner_dirty = true; - // bring exact_page_owner up to date; rebuilds only when a removal marked it dirty void exact_pages_sync() const; - // recompute it from the live cells void exact_pages_rebuild() const; - // record that a cell of (seq, pos) now lives at physical cell idx void exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos); bool v_trans = true; // the value tensor is transposed diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index d9e609fee20..a596f35bd8d 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -93,10 +93,7 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // so that the rollback snapshots remain valid const uint32_t n_rs_seq = mem_recr->n_rs_seq; - // [TAG_EXACT_CONCURRENCY] the recurrent half is not invariant to the ubatch shape: - // a prompt processed next to other prompts leaves a different gated delta net state, - // so it gets a ubatch of its own while plain decode steps stay batched. Passed - // whenever the mode is on, since it also keeps sets of unequal token counts apart. + // [TAG_EXACT_CONCURRENCY] the recurrent half is not invariant to the ubatch shape, so a prompt gets a ubatch of its own const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); @@ -149,8 +146,7 @@ bool llama_memory_hybrid::get_can_shift() const { } uint32_t llama_memory_hybrid::alloc_granularity() const { - // the recurrent half holds one state per sequence, so the attention half is the one whose - // cells a caller is planning capacity for + // the recurrent half holds one state per sequence, so the attention half is the one whose cells a caller is planning capacity for return mem_attn->alloc_granularity(); } @@ -169,8 +165,7 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 } void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { - // [TAG_EXACT_CONCURRENCY] the attention half refuses this, so refuse before either half is - // touched or the two could end up describing different states + // [TAG_EXACT_CONCURRENCY] the attention half refuses this, so refuse before either half is touched or the two could end up describing different states if (llama_exact_concurrency() && seq_id_src != seq_id_dst) { LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between sequences (%d -> %d); ignoring the copy\n", __func__, seq_id_src, seq_id_dst); diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index dadc0661344..d6ac2e55358 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -431,10 +431,7 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & // [TAG_RECURRENT_ROLLBACK_SPLITS] // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch // so that the rollback snapshots remain valid - // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: the state a prompt leaves - // behind depends on what shared its ubatch, so isolate prompts. Passed whenever the - // mode is on, since it also keeps sets of unequal token counts apart, which would - // otherwise be reduced in chunks the solo run never had. + // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: the state a prompt leaves behind depends on what shared its ubatch, so isolate prompts const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); diff --git a/src/llama-memory.h b/src/llama-memory.h index c0bc2ad7e30..8f85aac4130 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -100,10 +100,8 @@ struct llama_memory_i { // getters virtual bool get_can_shift() const = 0; - // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit: 1 unless a mode - // that allocates in larger blocks is on, and then n tokens occupy round_up(n, granularity) - // cells, so a caller planning capacity in tokens would see room that does not exist. Not pure, - // so a module that has never heard of this inherits the answer that was always true of it. + // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit: 1 unless a mode allocates in + // larger blocks, when n tokens occupy round_up(n, granularity) cells. Not pure, so old modules inherit 1. virtual uint32_t alloc_granularity() const { return 1; } // diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 3f08a060ef1..b86f2703aea 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7193,8 +7193,7 @@ struct test_flash_attn_ext : public test_case { } }; -// same attention as the CPU mask reference, but visiting nonadjacent pages in a different order; -// covers a partial tail and different page counts per query +// same attention as the CPU mask reference, but visiting nonadjacent pages in a different order struct test_flash_attn_ext_pages : public test_flash_attn_ext { test_flash_attn_ext_pages(int64_t batch) : test_flash_attn_ext(256, 256, 2, {8, 1}, 1024, batch) {} @@ -9205,7 +9204,6 @@ static std::vector> make_test_cases_eval() { } } - // shared weights over sequence planes, as in a recurrent-model output projection for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0}) { for (int n : {1, 17, 307}) { test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {4, 1})); @@ -9213,9 +9211,7 @@ static std::vector> make_test_cases_eval() { } } - // MoE projections at the token counts a decode ubatch forms: gate and up broadcast one - // activation row over the expert list, down carries one row per expert, and a real MoE gguf - // puts different types on the two. 17 tokens is past the width exact concurrency pins. + // MoE projections at the token counts a decode ubatch forms; 17 tokens is past the width exact concurrency pins for (ggml_type type_a : {GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16}) { for (int n : {1, 2, 4, 8, 17}) { test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, true, 512, n, 2048)); diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp index 24b6359b05f..5a1502f747f 100644 --- a/tests/test-state-restore-fragmented.cpp +++ b/tests/test-state-restore-fragmented.cpp @@ -73,8 +73,7 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : saved seq 1 state, %zu bytes\n", __func__, ncopy); - // a fragmented restore may stage a whole device tensor, so check every sequence - // byte-for-byte, including the neighbours that must be preserved + // a fragmented restore may stage a whole device tensor, so check every sequence byte-for-byte, neighbours included std::vector> before(params.n_parallel); for (int s = 0; s < params.n_parallel; ++s) { before[s].resize(llama_state_seq_get_size(ctx, s)); diff --git a/tests/test-state-seq-copy.cpp b/tests/test-state-seq-copy.cpp index dabc4db50c3..488a268ee07 100644 --- a/tests/test-state-seq-copy.cpp +++ b/tests/test-state-seq-copy.cpp @@ -1,14 +1,4 @@ -// [TAG_STATE_ASYNC] guards on the asynchronous per-sequence state transfer -// -// llama_state_seq_copy_get / _set take a size and a flags word from the caller and hand both -// to an io object that validates everything else against them. The buffer belongs to the -// transfer, so a size larger than it is refused rather than believed, and -// LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is refused because these copies serialise through host -// memory. This also checks that the buffer reports itself as page-locked only while it holds -// memory that is. -// -// Skipped, not failed, on a backend that cannot copy asynchronously: there is no transfer to -// make and the synchronous calls are what a caller uses there. +// [TAG_STATE_ASYNC] guards on the asynchronous state transfer: the buffer belongs to the transfer, so an oversized size is refused, and ON_DEVICE is refused as these go via the host #include "arg.h" #include "common.h" @@ -52,8 +42,7 @@ int main(int argc, char ** argv) { return 1; } - // put something in the cache to copy: two sequences interleaved, so the cells of each - // are a comb rather than one block, which is what the transfer is built for + // two sequences interleaved, so the cells of each are a comb rather than one block, which is what the transfer is built for std::vector tokens(60, 1); llama_batch batch = llama_batch_init(params.n_parallel*tokens.size(), 0, 1); @@ -100,28 +89,23 @@ int main(int argc, char ** argv) { CHECK(llama_state_seq_copy_get(cpy, size + 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); CHECK(llama_state_seq_copy_set(cpy, size + 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); - // and so is an empty one, which cannot even hold the header CHECK(llama_state_seq_copy_get(cpy, 0, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); CHECK(llama_state_seq_copy_set(cpy, 0, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); - // ON_DEVICE keeps the tensor data off the host, which is where these copies go CHECK(llama_state_seq_copy_get(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) == 0); CHECK(llama_state_seq_copy_set(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) == 0); - // none of that may have posted anything CHECK(llama_state_seq_copy_done(cpy)); fprintf(stderr, "%s : oversized, empty and ON_DEVICE transfers are all refused\n", __func__); - // a transfer that fails part way, one byte short of the state, must post nothing: the - // caller is told it failed and is free to reuse the buffer at once + // a transfer that fails part way must post nothing: the caller is told it failed and is free to reuse the buffer at once CHECK(llama_state_seq_copy_get(cpy, size - 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); CHECK(llama_state_seq_copy_n_copies(cpy) == 0); CHECK(llama_state_seq_copy_done(cpy)); fprintf(stderr, "%s : a transfer one byte short is refused and posts no copies\n", __func__); - // the same call at the size the transfer does own still works, and round-trips std::vector before(llama_state_seq_get_size(ctx, seq_id)); CHECK(llama_state_seq_get_data(ctx, before.data(), before.size(), seq_id) == before.size()); @@ -146,7 +130,6 @@ int main(int argc, char ** argv) { fprintf(stderr, "%s : a transfer at the buffer's own size round-trips seq %d byte-for-byte\n", __func__, seq_id); - // giving the memory back leaves nothing page-locked to report llama_state_seq_copy_buf_free(cpy); CHECK(llama_state_seq_copy_buf_is_pinned(cpy) == false); CHECK(llama_state_seq_copy_buf_capacity(cpy) == 0); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5c07c979280..8cff94e34a1 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -39,8 +39,7 @@ constexpr int HTTP_POLLING_SECONDS = 1; -// [TAG_EXACT_CONCURRENCY] read from the env like the KV cache, batch splitter and CUDA -// backend do: the answer is needed before a context exists. +// [TAG_EXACT_CONCURRENCY] read from the env: the answer is needed before a context exists static bool server_exact_concurrency() { static const bool enabled = []() { const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); @@ -77,28 +76,13 @@ enum slot_state { SLOT_STATE_RESTORING, // [TAG_PREEMPT_ASYNC] the copy back in is running; the cells are allocated but not yet filled }; -// [TAG_PREEMPT] server-side request preemption -// -// With --kv-unified one full pool ends EVERY conversation in flight with "Context size has -// been exceeded", including the ones nowhere near their own limit. Instead of terminating, -// one slot's sequence is copied to host RAM and its cells released; when there is room the -// copy goes back and the slot carries on with the same sampler, text and open stream, so a -// streaming client sees a pause rather than an error. +// [TAG_PREEMPT] server-side request preemption: instead of ending every conversation in flight with a context error, one slot's sequence is copied to host RAM and back when there is room constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation constexpr int64_t PREEMPT_KEEPALIVE_MS = 2000; // SSE keepalive period while a streaming slot is parked constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected -// [TAG_PREEMPT] The order parked slots come back in. Head of the line by park time, and nobody -// passes a head that does not fit yet: the head keeps the room the pool frees until it fits, so -// its wait is bounded by the slots ahead of it and not by how often a smaller slot can squeeze -// in, grow, and be parked again. Simulated over 60 seeds at eight chats this cuts the longest -// single wait by 2.5 to 3x for 0 to 3 percent of makespan at 8192 cells, and parks less often. -// LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order: most-preempted first, then longest -// parked, and a smaller slot may pass a head that does not fit. -// LLAMA_SERVER_PREEMPT_RESUME=head (the default) or pass; read once in load_model() and logged. -// [TAG_PREEMPT] the SSE comment for a park or a resume. A request with several prompts -// streams them through one reader, so the comment names the prompt it is about, except for -// prompt 0, whose comment stays the bare form a single-prompt client matches on. +// [TAG_PREEMPT] LLAMA_SERVER_PREEMPT_RESUME: head (the default) resumes by park time and lets nobody pass a head that does not fit; pass restores most-preempted-first +// [TAG_PREEMPT] the SSE comment for a park or resume; prompt 0 keeps the bare form single-prompt clients match on static std::string preempt_notice_comment(const server_task_result_preempt_notice & notice) { std::string res = notice.parked ? ": preempted" : ": resumed"; @@ -112,20 +96,7 @@ constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked constexpr int64_t PREEMPT_ROTATE_US = 2ll * 1000 * 1000; // a resident cycling through context shifts gives way to a parked head that has waited this long -// [TAG_PREEMPT_ASYNC] how far ahead of the pool filling an asynchronous park is triggered -// -// A synchronous park hands the cells back before update_slots() goes on, so it only has to -// fire once the next step would not fit. An asynchronous one does not: the copy is still -// reading the cells, and they are only released when it lands. The slots that keep decoding -// in the meantime need somewhere to put their tokens, so the park has to be triggered this -// many decode steps before the pool would actually have run out. Too small and the pool -// fills while the copy is still running, and the decode ends up waiting for it after all, -// which is no worse than not doing this at all but no better either. Too large and slots -// are parked, and so parked again, earlier and more often than they need to be, which costs -// more in total than the one late park it avoided. -// -// Eight steps of every running slot is about a tenth of a second of runway at the speeds a -// handful of parallel chats decode at, which is the order a copy of one sequence takes. +// [TAG_PREEMPT_ASYNC] an asynchronous park only releases its cells when its copy lands, so it must fire this many decode steps before the pool would run out constexpr int32_t PREEMPT_N_ASYNC_STEPS = 8; struct llama_state_seq_copy_deleter { @@ -140,27 +111,16 @@ static llama_state_seq_copy_ptr llama_state_seq_copy_make(llama_context * ctx) { return cpy ? llama_state_seq_copy_ptr(cpy, llama_state_seq_copy_deleter{}) : llama_state_seq_copy_ptr(); } -// [TAG_EXACT_CONCURRENCY] The planner above counts cells, not tokens, because the two are not -// the same number under every mode. llama_memory_alloc_granularity() reports how many cells the -// pool hands out at a time: 1 in every ordinary configuration, and the exact concurrency page -// size when that mode is on, where one page belongs to one (sequence, position / page) pair and -// a sequence of n tokens therefore occupies round_up(n, page) cells. Four sequences can be -// holding up to 4 * (page - 1) cells that nobody else can be given, and a planner counting -// tokens sees room in the pool that find_slot cannot find in pages: it never reaches the -// threshold that would park anybody, the retry ladder halves n_batch to 1, and every request -// ends in the context error that preemption exists to remove. - -// cells a run of n_tokens occupies when the pool allocates g at a time +// [TAG_EXACT_CONCURRENCY] the planner counts cells, not tokens: a page belongs to one sequence, so a token count sees room find_slot cannot find and nobody is ever parked + static constexpr int32_t preempt_n_cells_g(int32_t n_tokens, int32_t g) { return (g <= 1 || n_tokens <= 0) ? n_tokens : ((n_tokens + g - 1) / g) * g; } -// cells a step of n_step more costs: nothing until it crosses a page boundary, a page when it does static constexpr int32_t preempt_n_cells_step_g(int32_t n_tokens, int32_t n_step, int32_t g) { return preempt_n_cells_g(n_tokens + n_step, g) - preempt_n_cells_g(n_tokens, g); } -// at a granularity of 1 both are the identity, so nothing changes in a configuration that does not page static_assert(preempt_n_cells_g(0, 1) == 0 && preempt_n_cells_g(1, 1) == 1 && preempt_n_cells_g(8191, 1) == 8191 && preempt_n_cells_g(-3, 1) == -3, "at a granularity of 1 a run of n tokens has to cost exactly n cells"); @@ -168,7 +128,6 @@ static_assert(preempt_n_cells_step_g(0, 1, 1) == 1 && preempt_n_cells_step_g(819 preempt_n_cells_step_g(1000, 512, 1) == 512, "at a granularity of 1 a step of n tokens has to cost exactly n cells"); -// and the page arithmetic itself, so the rounding cannot change by accident static_assert(preempt_n_cells_g(1, 256) == 256 && preempt_n_cells_g(256, 256) == 256 && preempt_n_cells_g(257, 256) == 512, "a tail page is charged in full"); @@ -408,18 +367,12 @@ struct server_slot { prompt.clear(); } - // [TAG_PREEMPT] state of a slot whose cells were taken back. Only the KV cells leave; - // the task, sampler, generated text and stream position stay, so a resume is a memcpy. + // [TAG_PREEMPT] state of a slot whose cells were taken back; the task, sampler, text and stream stay, so a resume is a memcpy slot_state state_before_preempt = SLOT_STATE_IDLE; std::vector preempt_state_tgt; std::vector preempt_state_dft; - // [TAG_PREEMPT_ASYNC] the two transfers this slot parks and resumes through - // - // They own the pinned host buffers the sequence lives in while it is parked, so when - // they exist the std::vectors above stay empty and the state is in the transfers. Held - // by shared_ptr only because the slots are built with emplace_back into a vector that - // reallocates as it grows, and a slot must survive being moved. + // [TAG_PREEMPT_ASYNC] the two transfers this slot parks and resumes through; they own the pinned host buffers, and are shared_ptr only so a slot survives the vector's reallocation llama_state_seq_copy_ptr preempt_cpy_tgt; llama_state_seq_copy_ptr preempt_cpy_dft; @@ -427,7 +380,6 @@ struct server_slot { return (bool) preempt_cpy_tgt; } - // microseconds the last park or resume spent draining the compute streams int64_t preempt_sync_us() const { if (!preempt_is_async()) { return 0; @@ -437,7 +389,6 @@ struct server_slot { (preempt_cpy_dft ? llama_state_seq_copy_sync_us(preempt_cpy_dft.get()) : 0); } - // transfers the last park or resume posted, which is what its issue cost is made of size_t preempt_n_copies() const { if (!preempt_is_async()) { return 0; @@ -447,13 +398,11 @@ struct server_slot { (preempt_cpy_dft ? llama_state_seq_copy_n_copies(preempt_cpy_dft.get()) : 0); } - // [TAG_PREEMPT_ASYNC] a copy is running for this slot: it is not decoding and must not be - // scheduled, but it still owns cells, so it is neither running nor parked + // [TAG_PREEMPT_ASYNC] a copy is running: the slot must not be scheduled but still owns cells, so it is neither running nor parked bool preempt_in_flight() const { return state == SLOT_STATE_PREEMPTING || state == SLOT_STATE_RESTORING; } - // parked, or on its way out or back in: in none of these does the slot take part in a decode bool preempt_is_out() const { return state == SLOT_STATE_PREEMPTED || preempt_in_flight(); } @@ -466,9 +415,7 @@ struct server_slot { size_t preempt_state_size() const { if (preempt_is_async()) { - // the capacity, not the live size: the pinned buffers are kept between two parks - // of the same task because page-locking them again would cost as much as the - // copy, so what --preempt-ram has to bound is what is held, not what is in use + // the capacity, not the live size: the pinned buffers are kept between parks, so --preempt-ram has to bound what is held return llama_state_seq_copy_buf_capacity(preempt_cpy_tgt.get()) + (preempt_cpy_dft ? llama_state_seq_copy_buf_capacity(preempt_cpy_dft.get()) : 0); } @@ -477,9 +424,7 @@ struct server_slot { } void preempt_state_free() { - // resizing waits for anything still in flight first: this is called from release(), - // which a cancelled request reaches while its copy may still be reading or writing - // the buffer, and freeing it underneath a running transfer would be a use-after-free + // wait for anything in flight first: release() is reached with a copy possibly still using the buffer if (preempt_cpy_tgt) { llama_state_seq_copy_buf_free(preempt_cpy_tgt.get()); } @@ -494,7 +439,6 @@ struct server_slot { preempt_state_dft.shrink_to_fit(); } - // [TAG_PREEMPT_ASYNC] give up on an outstanding copy without using its result void preempt_copy_wait() { if (preempt_cpy_tgt) { llama_state_seq_copy_wait(preempt_cpy_tgt.get()); @@ -505,18 +449,12 @@ struct server_slot { } } - // bytes preempt_save() would need for this slot right now size_t preempt_state_required() const { return llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) + (ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0); } - // take the slot out of the step that is about to be built - // - // The draft is a prediction, not a result, so it goes with the cells. Preemption runs - // before the batch is built, so spec_i_batch is empty and prompt.tokens already holds - // exactly the tokens the state covers -- including the rollback done by the checkpoint - // path when a draft was only partially accepted. + // take the slot out of the step that is about to be built; the draft is a prediction, not a result, so it goes with the cells void preempt_detach() { spec_draft.clear(); spec_i_batch.clear(); @@ -526,7 +464,6 @@ struct server_slot { i_batch = -1; } - // [TAG_PREEMPT_ASYNC] has the copy out finished? if so, the cells can finally go bool preempt_save_poll() { if (!llama_state_seq_copy_done(preempt_cpy_tgt.get())) { return false; @@ -543,7 +480,6 @@ struct server_slot { return true; } - // [TAG_PREEMPT_ASYNC] has the copy back in finished? if so, the slot can decode again bool preempt_restore_poll() { if (!llama_state_seq_copy_done(preempt_cpy_tgt.get())) { return false; @@ -553,8 +489,6 @@ struct server_slot { return false; } - // the state is back in the cache, so the buffers hold nothing that matters; the - // memory itself is kept for the next park of this task and handed back by release() llama_state_seq_copy_buf_resize(preempt_cpy_tgt.get(), 0); if (preempt_cpy_dft) { @@ -565,8 +499,6 @@ struct server_slot { state = state_before_preempt; - // same call the DONE_PROMPT -> GENERATING transition makes; it reads the restored - // sequence, so it has to wait for the copy like everything else if (state == SLOT_STATE_GENERATING && can_speculate()) { common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); } @@ -574,11 +506,7 @@ struct server_slot { return true; } - // copy the sequence out of the cache and release its cells - // - // [TAG_PREEMPT_ASYNC] With a transfer this returns as soon as the copy has been issued, - // leaving the slot PREEMPTING: the cells are still its own, because the copy is still - // reading them, and nobody may take them until preempt_save_poll() says the copy landed. + // [TAG_PREEMPT_ASYNC] copy the sequence out and release its cells; with a transfer this returns once the copy is issued and the cells stay the slot's until preempt_save_poll() sees it land bool preempt_save() { const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); const size_t size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; @@ -593,12 +521,7 @@ struct server_slot { return false; } - // [TAG_PREEMPT_ASYNC] the load-time probe saw pinned memory, but a buffer this much - // larger can still come back pageable (a host-locking limit, say): the host buffer - // type hands back ordinary memory rather than failing, and a copy into pageable - // memory blocks the thread that issued it, which is the stall this path exists to - // remove. Such a slot parks synchronously from now on: its transfers are given - // back and the plain path below takes over, for this park and every later one. + // [TAG_PREEMPT_ASYNC] the load-time probe saw pinned memory, but a larger buffer can still come back pageable, and a copy into pageable memory blocks; such a slot parks synchronously from now on const bool pageable = !llama_state_seq_copy_buf_is_pinned(preempt_cpy_tgt.get()) || (size_dft > 0 && !llama_state_seq_copy_buf_is_pinned(preempt_cpy_dft.get())); @@ -624,8 +547,7 @@ struct server_slot { preempt_detach(); - // note: no mem.seq_rm() here. The copy is still reading these cells, so they are - // released in preempt_save_poll() once it has finished with them. + // note: no mem.seq_rm() here. The copy is still reading these cells; preempt_save_poll() releases them. state_before_preempt = state; state = SLOT_STATE_PREEMPTING; t_preempt_us = ggml_time_us(); @@ -674,12 +596,7 @@ struct server_slot { return true; } - // put the sequence back; the slot then continues from the token it was about to decode - // - // [TAG_PREEMPT_ASYNC] With a transfer this returns as soon as the copy has been issued, - // leaving the slot RESTORING: the cells are allocated and owned by this sequence, so - // nobody else can take them, but they do not hold its state until the copy lands, which - // is why the slot is not scheduled until preempt_restore_poll() says so. + // [TAG_PREEMPT_ASYNC] put the sequence back; with a transfer this returns once the copy is issued, leaving the slot RESTORING: it owns the cells, but they hold no state until the copy lands bool preempt_restore() { if (preempt_is_async()) { const size_t size_tgt = llama_state_seq_copy_buf_size(preempt_cpy_tgt.get()); @@ -688,9 +605,7 @@ struct server_slot { if (llama_state_seq_copy_set(preempt_cpy_tgt.get(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt || (size_dft > 0 && llama_state_seq_copy_set(preempt_cpy_dft.get(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft)) { - // no room after all: let whatever was already issued finish before the - // half-written sequence is dropped, or the cells would go while a copy is - // still writing into them + // no room after all: let what was issued finish before the half-written sequence is dropped, or cells go while a copy still writes them preempt_copy_wait(); mem.seq_rm(id, -1, -1); n_preempt_fail++; @@ -706,7 +621,6 @@ struct server_slot { const size_t size_dft = preempt_state_dft.size(); if (llama_state_seq_set_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { - // no room after all: drop the half-written sequence and stay parked mem.seq_rm(id, -1, -1); n_preempt_fail++; return false; @@ -725,8 +639,6 @@ struct server_slot { state = state_before_preempt; - // same call the DONE_PROMPT -> GENERATING transition makes; a slot parked mid-prompt - // makes that transition itself once the prompt is done if (state == SLOT_STATE_GENERATING && can_speculate()) { common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); } @@ -734,9 +646,7 @@ struct server_slot { return true; } - // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds, for a batch given up - // after it was built: never-decoded tokens and the draft come off, `sampled` is kept for - // the next batch. A failed chunk left nothing in the cache, so the cache is the boundary. + // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds, for a batch given up after it was built: never-decoded tokens and the draft come off, `sampled` is kept void rewind_to_cache() { const int32_t n_cached = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), id) + 1; @@ -975,14 +885,8 @@ struct server_slot { t_last_used = ggml_time_us(); - // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be - // released while parked), so the mirror of them must not outlive them: the next - // task on this slot would otherwise take a prefix match against an empty cache - // - // [TAG_PREEMPT_ASYNC] a slot can also be released with a copy still running, by - // a cancelled request or by the error paths. Wait for it before anything else: - // the host buffer is about to be freed and the cells about to be handed to the - // next task, and a transfer still reading or writing either would outlive both. + // [TAG_PREEMPT] a parked slot's cells are already gone, so the mirror must not outlive them or the next task prefix-matches an empty cache + // [TAG_PREEMPT_ASYNC] wait for any copy first: its host buffer and its cells are about to be handed on if (preempt_is_out()) { preempt_copy_wait(); preempt_state_free(); @@ -1376,25 +1280,14 @@ struct server_context_impl { int64_t t_last_load_progress_ms = 0; void destroy() { - // [TAG_PREEMPT_ASYNC] the slots outlive this call -- they are declared after - // llama_init, so they are still there when it is reset here, and load_model() clears - // them only after the next context exists -- and any one of them may be holding a - // park or a resume that is still reading or writing KV tensors of the contexts about - // to be freed. release() makes the same wait for a single slot; this is the one that - // covers all of them, and on the sleeping-state path it is the only one there is, - // because the server carries on running afterwards. + // [TAG_PREEMPT_ASYNC] the slots outlive this call and may hold a copy reading or writing KV tensors of the contexts about to be freed; release() makes the same wait for one slot for (auto & slot : slots) { slot.preempt_copy_wait(); - // the transfer holds a backend and an event of its own, and its host buffer is - // no use to the context that comes back: let go of both before that context's - // successor makes new ones slot.preempt_cpy_tgt.reset(); slot.preempt_cpy_dft.reset(); } - // the next context allocates its own host buffers, so say again what they turn out - // to be preempt_ram_kind_logged = false; spec.reset(); @@ -1737,11 +1630,7 @@ struct server_context_impl { } }; - // [TAG_PREEMPT_ASYNC] one transfer per context, made once and reused for every - // park and resume this slot ever does, because each owns a backend and a stream. - // Only where a park can happen at all (see update_preemption): a transfer also - // installs the fences the context records after every decode, which a server - // that will never park has no use for. + // [TAG_PREEMPT_ASYNC] one transfer per context, reused for every park and resume, because each owns a backend and installs the fences the context records after every decode if (preempt_async_possible()) { slot.preempt_cpy_tgt = llama_state_seq_copy_make(ctx_tgt); @@ -1749,8 +1638,7 @@ struct server_context_impl { slot.preempt_cpy_dft = llama_state_seq_copy_make(ctx_dft); if (!slot.preempt_cpy_dft) { - // a draft that cannot go asynchronously would have to be waited for - // in the middle of the park, so the whole slot stays synchronous + // a draft that cannot go asynchronously would have to be waited for mid-park, so the whole slot stays synchronous slot.preempt_cpy_tgt.reset(); } } @@ -1787,12 +1675,7 @@ struct server_context_impl { if (preempt_async_possible()) { if (preempt_async_ok) { - // Pinned host memory is what lets a copy run beside the decode: one into or - // out of pageable memory is staged by the driver and blocks the thread that - // issued it, which is the stall the asynchronous path exists to remove. A - // host buffer type is free to hand back ordinary memory instead of failing - // (GGML_CUDA_NO_PINNED, or a pinning limit), and that is only knowable from - // a buffer, so a small one is taken and looked at before the first park. + // a copy into pageable memory is staged by the driver and blocks the thread that issued it, and a host buffer type is free to hand back pageable memory rather than fail bool pinned = llama_state_seq_copy_buf_can_pin(slots[0].preempt_cpy_tgt.get()); if (pinned) { @@ -1823,14 +1706,11 @@ struct server_context_impl { } } - // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates, rather than assume a cell per - // token. 1 in every ordinary configuration, so this changes nothing unless a mode that - // places cells in blocks is on. + // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates rather than assume a cell per token { preempt_alloc_granularity = (int32_t) std::max(1u, llama_memory_alloc_granularity(llama_get_memory(ctx_tgt))); - // test knob: the paged kernel needs a head size of 256, so a harness model cannot - // turn exact concurrency on and this is the only way to reach the paged arithmetic + // test knob: the paged kernel needs a head size of 256, so a harness model cannot reach the paged arithmetic otherwise const char * LLAMA_SERVER_PREEMPT_GRANULARITY = getenv("LLAMA_SERVER_PREEMPT_GRANULARITY"); if (LLAMA_SERVER_PREEMPT_GRANULARITY) { @@ -1845,8 +1725,6 @@ struct server_context_impl { } { - // read on every load and kept on this context, so a reload after the variable - // changed, or another context loaded in the same process, has an order of its own preempt_resume_head = true; const char * LLAMA_SERVER_PREEMPT_RESUME = getenv("LLAMA_SERVER_PREEMPT_RESUME"); @@ -1863,9 +1741,7 @@ struct server_context_impl { const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; - // LLAMA_SERVER_PREEMPT_POLICY: which non-leader the planner parks, for comparing - // policies on the same workload: smallest (default), largest, youngest, oldest. - // The leader is kept and the starvation guard applies under all. + // LLAMA_SERVER_PREEMPT_POLICY: which non-leader the planner parks; smallest (default), largest, youngest, oldest const char * LLAMA_SERVER_PREEMPT_POLICY = getenv("LLAMA_SERVER_PREEMPT_POLICY"); preempt_test_policy = LLAMA_SERVER_PREEMPT_POLICY ? LLAMA_SERVER_PREEMPT_POLICY : "smallest"; @@ -1885,8 +1761,7 @@ struct server_context_impl { SRV_WRN("%s", "LLAMA_SERVER_PREEMPT_PLANNER = off (test knob: nothing is parked ahead of the decode, only as a last resort)\n"); } - // assigned, not only set: the same context reloaded with an attention model after - // a recurrent one gets its preemption back + // assigned, not only set: a context reloaded with an attention model after a recurrent one gets preemption back preempt_recurrent = llama_model_is_recurrent(model_tgt); if (preempt_recurrent) { @@ -2588,9 +2463,7 @@ struct server_context_impl { queue_results.send(std::move(res)); } - // [TAG_PREEMPT] tell a streaming client that its slot was parked or restored. The - // HTTP layer turns this into an SSE comment, so a client that does not know about - // preemption sees nothing, and one that does can show a pause instead of a stall. + // [TAG_PREEMPT] tell a streaming client its slot was parked or restored; the HTTP layer sends an SSE comment, which a client that does not know about preemption never sees void send_preempt_notice(server_slot & slot, bool parked) { if (!slot.task || !slot.task->params.stream) { return; @@ -3307,10 +3180,7 @@ struct server_context_impl { void abort_all_slots(const std::string & reason) { for (auto & slot : slots) { - // [TAG_PREEMPT] a parked slot took no part in what failed: its sequence is in - // host RAM, not in the cache, and it comes back when there is room, the same as - // in the decode error sweep - // [TAG_PREEMPT_ASYNC] a slot whose copy is in flight is out of the round as well + // [TAG_PREEMPT] a parked slot, or one with a copy in flight, took no part in what failed and comes back when there is room if (slot.is_processing() && !slot.preempt_is_out()) { send_error(slot, reason, ERROR_TYPE_SERVER); slot.release(); @@ -3348,69 +3218,46 @@ struct server_context_impl { }; #endif - // // [TAG_PREEMPT] server-side request preemption - // - // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens, - // under pressure or not: on an idle server the batch shape is fixed, so a continuation - // that is not byte-identical to an uninterrupted one is the preemption's fault. + // LLAMA_SERVER_PREEMPT_EVERY=N: preempt every generating slot every N tokens, pressure or not, so the determinism test can blame any difference on the preemption int32_t preempt_test_every = 0; std::string preempt_test_policy = "smallest"; // LLAMA_SERVER_PREEMPT_POLICY, see load_model - // [TAG_PREEMPT_ASYNC] whether the slots park and resume through a transfer. False when - // --no-preempt-async was given, or when the backend cannot copy asynchronously, in which - // case every park and resume is the synchronous one it always was. + // [TAG_PREEMPT_ASYNC] whether parks go through a transfer; false with --no-preempt-async or a backend that cannot copy asynchronously bool preempt_async_ok = false; - // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load from the - // memory itself: 1 in every ordinary configuration, the page size under exact concurrency. - // Everything below plans in cells because of it. LLAMA_SERVER_PREEMPT_GRANULARITY overrides - // it, which is how the harness reaches the paged arithmetic on a model whose head size the - // paged attention kernel does not support. + // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load: 1 ordinarily, the page size under exact concurrency. LLAMA_SERVER_PREEMPT_GRANULARITY overrides it. int32_t preempt_alloc_granularity = 1; - // cells a slot holding n_tokens actually occupies int32_t preempt_n_cells(int32_t n_tokens) const { return preempt_n_cells_g(n_tokens, preempt_alloc_granularity); } - // cells a slot holding n_tokens has to be given for a step of n_step more int32_t preempt_n_cells_step(int32_t n_tokens, int32_t n_step) const { return preempt_n_cells_step_g(n_tokens, n_step, preempt_alloc_granularity); } - // env: LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, so - // the KV-full retry ladder and its last resort are the only thing between a full pool - // and the context error + // LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, leaving only the KV-full retry ladder bool preempt_planner_off = false; - // LLAMA_SERVER_PREEMPT_RESUME: head (the default) puts parked slots back in the order they - // were parked and only the first until it fits; pass lets a smaller slot pass a head - // that does not fit. Read at load, per context. + // LLAMA_SERVER_PREEMPT_RESUME=head or pass, read at load, per context bool preempt_resume_head = true; - // a recurrent cache holds one state per sequence whatever its length: no cell pool, - // nothing to run out of, and the token count the planner measures says nothing about - // it. Preemption is off for those models; a hybrid keeps its attention cache and stays on. + // a recurrent cache has no cell pool to run out of, so preemption is off for those models; a hybrid keeps its attention cache and stays on bool preempt_recurrent = false; - // set by preempt_last_resort(): the batch being decoded was given up, stop the chunk loop bool preempt_batch_abandoned = false; - // [TAG_PREEMPT_ASYNC] a context shift was recorded this round; it is applied inside the - // next llama_decode as one graph over the whole K cache, in place + // [TAG_PREEMPT_ASYNC] a context shift was recorded this round; it is applied in place inside the next llama_decode bool preempt_shift_pending = false; int32_t preempt_n_spec_max() const { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } - // [TAG_PREEMPT_ASYNC] whether the kind of host memory the parks actually got has been - // reported. It is only knowable once a buffer exists, and it is worth knowing: pinned - // memory is what lets the copies overlap, and a host buffer type is free to hand back - // ordinary memory instead of failing. + // [TAG_PREEMPT_ASYNC] whether the kind of host memory the parks got has been reported; only knowable once a buffer exists bool preempt_ram_kind_logged = false; void preempt_log_ram_kind(const server_slot & slot) { @@ -3428,8 +3275,6 @@ struct server_context_impl { llama_state_seq_copy_buf_is_pinned(slot.preempt_cpy_tgt.get()) ? "pinned" : "pageable"); } - // draft tokens this slot's next step can actually carry: the configured maximum, cut to - // what its context and its prediction budget leave, the way get_n_draft_max() cuts it int32_t preempt_n_spec(const server_slot & slot) const { int32_t res = preempt_n_spec_max(); @@ -3446,7 +3291,6 @@ struct server_context_impl { return std::max(0, res); } - // host RAM the parked sequences hold right now size_t preempt_ram_used() const { size_t res = 0; @@ -3457,14 +3301,7 @@ struct server_context_impl { return res; } - // whether parking this slot stays under --preempt-ram - // [TAG_PREEMPT_ASYNC] a restored slot keeps its pinned buffer for its next park, and - // that capacity counts against the budget while it holds no state. When a park does - // not fit, that idle capacity is what to give back first: largest first, never a - // buffer that still holds a parked sequence or has a copy in flight, and never the - // candidate's own, which it reuses. Without this a budget that holds one sequence was - // spent for good by the first restore: every later park was refused, and once the - // slot holding the buffer was the leader nothing could be parked at all. + // [TAG_PREEMPT_ASYNC] whether parking this slot stays under --preempt-ram; a restored slot keeps its pinned buffer, so idle capacity is given back largest first when a park does not fit void preempt_reclaim_idle_ram(size_t budget, size_t extra, const server_slot & keep) { for (;;) { if (preempt_ram_used() + extra <= budget) { @@ -3509,8 +3346,7 @@ struct server_context_impl { const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; - // whatever this slot already holds is counted by preempt_ram_used() and will be - // reused, so parking it again only costs what it does not have yet + // what this slot already holds is counted by preempt_ram_used() and reused, so a park costs only the rest const size_t held = slot.preempt_state_size(); const size_t need = slot.preempt_state_required(); const size_t extra = need > held ? need - held : 0; @@ -3520,9 +3356,7 @@ struct server_context_impl { return preempt_ram_used() + extra <= budget; } - // [TAG_PREEMPT_ASYNC] a restored slot keeps its pinned buffer for its next park, which is - // worth it while the budget has room for it and not otherwise: over budget, a buffer held - // by a slot that is running would keep every other slot from being parked at all + // [TAG_PREEMPT_ASYNC] over budget, a buffer held by a running slot would keep every other slot from being parked at all void preempt_trim_ram(server_slot & slot) { if (params_base.preempt_ram_mib < 0) { return; @@ -3536,9 +3370,6 @@ struct server_context_impl { } } - // cells of the mirrored prompt that a started slot's request keeps, by the rule the batch - // builder applies when it takes the slot: nothing when the request does not cache its - // prompt, otherwise the prefix the two share, cut short of an aLoRA invocation size_t preempt_n_keep(const server_slot & slot) const { if (!slot.task->params.cache_prompt) { return 0; @@ -3553,10 +3384,7 @@ struct server_context_impl { return n_keep; } - // cells of the slot's that its next step keeps: a slot just given a task still mirrors - // the previous request's prompt until the batch builder keeps what preempt_n_keep() - // says and drops the rest, so what it holds, and what it is about to ask for, both - // count from that + // a slot just given a task still mirrors the previous request's prompt, so what it holds and what it asks for both count from preempt_n_keep() int32_t preempt_n_retained(const server_slot & slot) const { if (slot.state == SLOT_STATE_STARTED && slot.task) { return (int32_t) preempt_n_keep(slot); @@ -3565,7 +3393,6 @@ struct server_context_impl { return slot.prompt.n_tokens(); } - // cells the slot will ask for on its next step once it is back in the pool int32_t preempt_n_need(const server_slot & slot) const { int32_t res = preempt_n_retained(slot); @@ -3577,18 +3404,14 @@ struct server_context_impl { res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); } - // [TAG_EXACT_CONCURRENCY] a restore takes fresh pages and its tail page is charged in - // full; undercounting here admits a resume that find_slot cannot satisfy + // [TAG_EXACT_CONCURRENCY] a restore takes fresh pages and its tail page is charged in full; undercounting admits a resume find_slot cannot satisfy return preempt_n_cells(res); } - // cells the pool is holding right now. A released slot keeps its prompt in the cache as a - // prefix for the next request, so idle slots count too or a full pool looks empty. int32_t preempt_kv_used() const { int32_t res = 0; - // n_cmpl > 1: a family shares the prompt's cells through seq_cp, so the prompt is - // charged once, to the first resident member; the others only for what they added + // n_cmpl > 1: a family shares the prompt's cells through seq_cp, so it is charged once, to the first resident member std::vector charged; for (const auto & slot : slots) { @@ -3596,16 +3419,10 @@ struct server_context_impl { continue; // parked: its cells are in host RAM, not in the pool } - // [TAG_PREEMPT_ASYNC] deliberately not skipped: a slot whose copy is still - // running holds cells either way. One on its way out has not released them yet - // because the copy is still reading them, and one on its way back in has already - // been given them. Skipping either would hand the same cells out twice. + // [TAG_PREEMPT_ASYNC] deliberately not skipped: a slot with a copy in flight holds cells either way, and skipping it would hand the same cells out twice - // [TAG_EXACT_CONCURRENCY] the slot's tail page is charged in full: it belongs to - // this sequence and cannot be given to anybody else, however little of it is used + // [TAG_EXACT_CONCURRENCY] the tail page is charged in full: it cannot be given to anybody else - // a child waiting for its parent shares nothing until copy_state_to() runs, so it - // is charged its own stale cells, outside the family if (slot.state == SLOT_STATE_WAIT_OTHER) { res += preempt_n_cells(slot.prompt.n_tokens()); continue; @@ -3622,74 +3439,39 @@ struct server_context_impl { charged.push_back(family); } - // what the pool holds now, the previous request's prompt included for a slot just - // given a task: the batch builder trims that to the prefix the two share, but - // not until the slot is built into a batch, and with continuous batching off that - // can be a long time behind a running generation. Measured by the prefix, a - // restore was found to fit and attempted against cells still occupied. Under - // pressure the planner trims such slots itself, see preempt_normalize_started_all() res += preempt_n_cells(slot.prompt.n_tokens()); } return res; } - // [TAG_PREEMPT_ASYNC] the room the pool is kept clear of - // - // A synchronous park releases the cells before update_slots() carries on, so it only has - // to fire once the next step would not fit. An asynchronous one leaves them held until - // its copy lands, so it has to fire early enough that everything still decoding has - // somewhere to put its tokens until then. The same figure gates a resume, so that a slot - // is not put back into a pool it would immediately have to be taken out of again. - // - // n_additional_running is for slots that are not running yet but are about to be: a - // resume candidate is still PREEMPTED while it is being considered, so it is not counted - // by the loop below, yet the moment it is admitted it starts decoding and needs the same - // runway as everybody else. Admitting it without charging it that runway is what the - // margin exists to prevent, and it showed up as a slot restored and parked again a few - // iterations later, over and over. + // [TAG_PREEMPT_ASYNC] the room the pool is kept clear of, so everything still decoding has somewhere to put its tokens until a park lands; a resume candidate is charged the same runway int32_t preempt_n_margin(int32_t n_additional_running = 0) const { if (!preempt_async_active()) { - // [TAG_EXACT_CONCURRENCY] a margin of eight cells is no margin at all where a - // step can cost a whole page, so the synchronous figure is rounded as well + // [TAG_EXACT_CONCURRENCY] a margin of eight cells is no margin where a step can cost a whole page return preempt_n_cells(PREEMPT_N_MARGIN); } int32_t n_running = n_additional_running; for (const auto & slot : slots) { - // A slot on its way back in already holds its cells and starts decoding the - // moment its copy lands, so it needs the runway now; one on its way out does not. if (slot.is_processing() && (!slot.preempt_is_out() || slot.state == SLOT_STATE_RESTORING)) { n_running++; } } - // [TAG_EXACT_CONCURRENCY] the lookahead is a number of decode steps, and under a - // page allocator any one of them can cost a whole page rather than a cell. Round the - // runway up to a page so the park is issued with at least one page of real room - // behind it; with a page size of 1 this is the figure it always was. - // The sum is rounded once, not once per running slot: every slot could cross a page - // boundary during the copy, but reserving a page for each of them would keep a page - // per slot out of the users' reach all the time, whereas the case it guards against - // costs one synchronous wait for a copy already in flight, in preempt_wait_in_flight(). + // [TAG_EXACT_CONCURRENCY] round the runway up to a page, once and not per slot, which would keep a page per slot out of the users' reach return preempt_n_cells( PREEMPT_N_MARGIN + n_running * (1 + preempt_n_spec_max()) * PREEMPT_N_ASYNC_STEPS); } - // cells those slots are about to ask for on the next decode int32_t preempt_kv_reserve() const { const int32_t n_batch = llama_n_batch(ctx_tgt); int32_t res = 0; int32_t res_pmt = 0; - // [TAG_EXACT_CONCURRENCY] each slot reserves the cells its next step ADDS, not the - // tokens it adds. preempt_kv_used() already charges every slot's tail page in full, - // so with a page size of 1 these are the same number and nothing changes; with a - // larger one the step is free until it crosses a page boundary and costs a whole - // page when it does. Reserving tokens on top of a rounded used figure would miss - // exactly that crossing, which is the only moment the pool can actually run out. + // [TAG_EXACT_CONCURRENCY] reserve the cells the next step ADDS, not its tokens: the used figure already rounds every tail page up, and only a page crossing can empty the pool for (const auto & slot : slots) { const int32_t n_cur = slot.prompt.n_tokens(); @@ -3702,10 +3484,6 @@ struct server_context_impl { case SLOT_STATE_STARTED: case SLOT_STATE_PROCESSING_PROMPT: { - // from the prefix a started slot keeps, not from the prompt it still - // mirrors: measured by the mirror, a request shorter than the last one - // reserved one cell for a chunk of hundreds; the step starts from that - // prefix too, since that is what the used figure charges for it const int32_t n_have = preempt_n_retained(slot); const int32_t n_left = slot.task ? slot.task->n_tokens() - n_have : 0; @@ -3713,10 +3491,7 @@ struct server_context_impl { } break; case SLOT_STATE_RESTORING: { - // [TAG_PREEMPT_ASYNC] its cells are already counted by preempt_kv_used(), - // but it starts decoding as soon as its copy lands, so the step it will - // take has to be reserved now -- otherwise the pool is handed out from - // under it and its first step preempts somebody else straight away + // [TAG_PREEMPT_ASYNC] it starts decoding as soon as its copy lands, so its step has to be reserved now, or its first step preempts somebody else if (slot.state_before_preempt == SLOT_STATE_GENERATING) { res += preempt_n_cells_step(n_cur, 1 + preempt_n_spec(slot)); } else { @@ -3725,20 +3500,14 @@ struct server_context_impl { res_pmt += preempt_n_cells_step(n_cur, std::max(1, std::min(n_batch, n_left))); } } break; - // a preempting slot is on its way out and will not decode: nothing to reserve default: break; } } - // one batch is all the prompt slots get between them, however many are waiting; under - // page allocation each can still cross a boundary of its own, so the cap allows one - // boundary per prompt slot on top of the batch int32_t n_pmt = 0; for (const auto & slot : slots) { - // a slot restoring into the prompt phase joins the next prompt batch too, and - // reserves its chunk above, so it can cross a boundary of its own as well const slot_state state = slot.state == SLOT_STATE_RESTORING ? slot.state_before_preempt : slot.state; if (state == SLOT_STATE_STARTED || state == SLOT_STATE_PROCESSING_PROMPT) { @@ -3749,17 +3518,8 @@ struct server_context_impl { return res + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1)); } - // Keep the slot that is furthest along -- it is the closest to finishing and to giving - // its cells back -- and among the rest prefer one that has not been preempted - // PREEMPT_N_STARVED times already, then the smallest. - // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt - // until the batch builder keeps the prefix the two share and drops the rest (see the - // SLOT_STATE_STARTED block of update_slots). Parked as it is, it would be copied out, - // charged and sized by the old prompt, and a short unrelated request could exceed the - // budget or stay parked for room it will never use. Keeping only the shared prefix now - // is what the batch builder does anyway; the chunk reuse it can add on top is given up - // for a slot the planner has to touch, which is rare. - // every started slot, when the pool is short: true when any of them gave cells up + // keep the slot furthest along, it is the closest to giving its cells back; among the rest prefer one not preempted PREEMPT_N_STARVED times, then the smallest + // [TAG_PREEMPT] trim a just-started slot to the prefix it keeps first, or it is copied out, charged and sized by the previous request's prompt bool preempt_normalize_started_all() { bool res = false; @@ -3793,8 +3553,7 @@ struct server_context_impl { return; } - // a memory that cannot remove part of a sequence aborts on a partial removal; for it - // the whole stale sequence goes and the prompt is reprocessed from the start + // a memory that cannot remove part of a sequence aborts on a partial removal, so drop the whole stale sequence const bool partial_ok = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART && (!ctx_dft || ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART); @@ -3812,8 +3571,6 @@ struct server_context_impl { server_slot * leader = nullptr; int32_t n_running = 0; - // measure a just-started slot by the prefix it keeps, not by the stale prompt it - // mirrors, or a short request over a large stale cache becomes the leader for (auto & slot : slots) { preempt_normalize_started(slot); } @@ -3829,16 +3586,14 @@ struct server_context_impl { } if (n_running < 2) { - // one conversation that does not fit alone is a real overflow, not a scheduling - // problem - leave it to the existing error path + // one conversation that does not fit alone is a real overflow, not a scheduling problem return nullptr; } server_slot * victim = nullptr; for (auto & slot : slots) { - // before the batch is built every one of these is at a token boundary. A slot - // holding no cells is still worth parking - it is about to ask for a batch. + // before the batch is built every slot is at a token boundary; one holding no cells is still worth parking if (slot.state != SLOT_STATE_GENERATING && slot.state != SLOT_STATE_PROCESSING_PROMPT && slot.state != SLOT_STATE_STARTED) { @@ -3853,9 +3608,7 @@ struct server_context_impl { continue; // n_cmpl > 1 slots share one sequence, out of scope here } - // a started slot whose request the STARTED block is about to reject gets its - // error on its own pass, and nothing before it: a park notice would open the - // stream and turn that error into 200 plus an in-stream one + // a started slot the STARTED block is about to reject gets its error on its own pass: a park notice would open the stream and turn that 4xx into 200 plus an in-stream error if (slot.state == SLOT_STATE_STARTED) { std::string msg; error_type type = ERROR_TYPE_SERVER; @@ -3882,8 +3635,6 @@ struct server_context_impl { return victim; } - // is a the better victim? the smallest under the shipped policy, since it gives up the - // least work; the rest exist for comparison runs behind LLAMA_SERVER_PREEMPT_POLICY bool preempt_better_victim(const server_slot & a, const server_slot & b) const { if (preempt_test_policy == "largest") { return a.prompt.n_tokens() > b.prompt.n_tokens(); @@ -3900,18 +3651,12 @@ struct server_context_impl { return a.prompt.n_tokens() < b.prompt.n_tokens(); } - // called once per update_slots(), before the batch is built: at that point every slot is - // at a token boundary, prompt.tokens is exactly what the cache holds for it, and no - // draft is in flight, so a slot can be removed from the picture without unpicking a - // half-decoded batch + // called once per update_slots(), before the batch is built: every slot is then at a token boundary with no draft in flight, so it can be removed whole // [TAG_PREEMPT_ASYNC] is any slot parking or resuming through a transfer right now bool preempt_async_active() const { return preempt_async_ok; } - // Pick up the copies that have landed since the last iteration. This runs before - // anything reads preempt_kv_used(), so a park whose cells came back is seen as free - // room straight away and a resume that landed can be scheduled in the same iteration. void update_preempt_copies() { for (auto & slot : slots) { if (slot.state == SLOT_STATE_PREEMPTING) { @@ -3930,13 +3675,7 @@ struct server_context_impl { preempt_trim_ram(slot); - // [TAG_PREEMPT] the mirror of the park notice, and the reason it is - // here rather than where the restore was issued: preempt_restore_poll() - // is what puts the slot back into the state it was parked from, so this - // is the first moment it can be scheduled again. Announcing it at the - // issue would tell the client it had resumed while it was still - // RESTORING, unscheduled and producing nothing, and the keepalive would - // have stopped for that window. + // [TAG_PREEMPT] announced here rather than where the restore was issued: preempt_restore_poll() is the first moment the slot can be scheduled again send_preempt_notice(slot, false); SLT_WRN(slot, "restore completed after %.2f ms: %d tokens back in the cache, kv %d/%d, preemptions %d\n", @@ -3949,11 +3688,7 @@ struct server_context_impl { } } - // [TAG_PREEMPT_ASYNC] wait for every copy in flight, parks and restores alike. The - // context shift a slot recorded this round is applied inside the next llama_decode as one - // graph over the whole K cache, in place: a restore still writing its cells on its own - // stream could be read half done and written back stale, and a park still reading its - // cells would read through the rewrite. Shifts are rare, so this round waits. + // [TAG_PREEMPT_ASYNC] wait for every copy in flight before a shift: the shift is one in-place graph over the whole K cache, so a copy beside it reads or writes half-shifted cells void preempt_wait_for_shift() { if (!preempt_shift_pending) { return; @@ -3979,8 +3714,6 @@ struct server_context_impl { preempt_trim_ram(slot); - // the restore landed here rather than in update_preempt_copies(), so the mirror - // of the park notice is sent here: no later poll sees this slot restoring send_preempt_notice(slot, false); SLT_WRN(slot, "restore completed after %.2f ms (waited for, a context shift is due): %d tokens back in the cache, kv %d/%d, preemptions %d\n", @@ -4001,9 +3734,7 @@ struct server_context_impl { return false; } - // Wait for one outstanding park, the last thing tried before giving up on finding room. - // It is what keeps a pool that fills faster than the copies drain no worse than the - // synchronous path: the decode waits for the copy exactly as it used to. + // wait for one outstanding park, the last thing tried before giving up on room: the decode then waits for the copy exactly as the synchronous path did bool preempt_wait_in_flight() { for (auto & slot : slots) { if (slot.state != SLOT_STATE_PREEMPTING) { @@ -4046,10 +3777,6 @@ struct server_context_impl { const int32_t n_cells = n_ctx; - // Put back what fits, in the order preempt_resume_head describes: by default - // the slot parked longest, and only that one until it fits; under - // LLAMA_SERVER_PREEMPT_RESUME=pass the most-preempted slot first, then the one parked - // longest, and a smaller slot may pass a head that does not fit. const bool head_of_line = preempt_resume_head; for (;;) { @@ -4079,9 +3806,7 @@ struct server_context_impl { server_slot * best = nullptr; - // a parked slot that would not fit an empty pool can never be restored and would - // sit at the head of the line for ever. That is the single-conversation overflow - // the KV-full path reports, so report it the same way and rescan without it. + // a parked slot that would not fit an empty pool can never be restored, so report it as the single-conversation overflow and rescan without it { server_slot * impossible = nullptr; @@ -4101,13 +3826,7 @@ struct server_context_impl { } } - // Room for the sequence AND for the next step of everything already running, - // and for the lookahead of the candidate itself, which is about to become one of - // them: a resume must not immediately trigger the preemption of someone else, or - // of itself. The margin is headroom for the others; with nothing resident there - // is nobody to keep it for, so a sequence that fits the pool exactly is let back in. - // A cached prompt on an idle slot is worth less than a conversation waiting to - // continue, so give those cells up first - same call the KV-full path makes. + // room for the sequence and for the next step of everything running, the candidate included, or a resume immediately preempts somebody; with nobody resident an exact fit is let in for (;;) { const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); const int32_t margin = occupied == 0 ? 0 : preempt_n_margin(1); @@ -4123,9 +3842,6 @@ struct server_context_impl { break; } - // a slot just given a task still holds the previous request's prompt until - // the batch builder trims it; trimmed here instead, the cells it will not - // keep are counted out and a parked slot that fits without them comes back if (preempt_normalize_started_all()) { continue; } @@ -4135,14 +3851,11 @@ struct server_context_impl { } } - // nothing fits. A resident cycling through context shifts holds the room for as - // long as it generates, so once the head has waited its turn that resident is - // parked in its place and the two take turns. + // nothing fits: a resident cycling through context shifts holds the room for as long as it generates, so it is parked once the head has waited its turn if (!best) { server_slot * head = parked.front(); - // [TAG_PREEMPT_ASYNC] a park still copying holds its cells, so the head would - // not fit yet and a rotation now would only park another resident on top + // [TAG_PREEMPT_ASYNC] a park still copying holds its cells, so a rotation now would only park another resident on top bool parking = false; for (const auto & slot : slots) { @@ -4150,10 +3863,6 @@ struct server_context_impl { } if (!parking && ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { - // the resident whose cells let the head in, the smallest of those; failing - // one that does so alone, the largest, since it makes the most room. Taking - // the first shifting resident in slot order could park one too small to - // matter, spend the park budget on it, and leave the head waiting anyway. const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); const int32_t need = preempt_n_need(*head) + preempt_n_margin(1); @@ -4170,11 +3879,7 @@ struct server_context_impl { continue; } - // The head's own bytes are not credited as leaving: the resident is - // parked before the head is restored and freed, so both states are - // held at once, and the cap is a cap on what is held. A budget that - // holds one sequence but not two does not rotate, and the head waits - // for a resident to finish, which is said once per park below. + // the head's own bytes are not credited as leaving: the resident is parked before the head is restored and freed, so both states are held at once if (!preempt_fits_budget(slot)) { budget_refused = true; continue; @@ -4205,8 +3910,6 @@ struct server_context_impl { slot.t_preempt_copy_us = t_start; - // [TAG_PREEMPT_ASYNC] an asynchronous park is counted when its copy - // lands, and the head is re-examined on the pass that sees the room if (slot.state != SLOT_STATE_PREEMPTING) { metrics.n_preempt++; } @@ -4220,9 +3923,7 @@ struct server_context_impl { pick_enough ? "" : " (not enough room by itself)", slot.n_preempt); - // [TAG_PREEMPT_ASYNC] a synchronous park has released its cells, so the - // head is re-examined now; an asynchronous one has not, and the head is - // re-examined on the pass that sees the copy land + // [TAG_PREEMPT_ASYNC] a synchronous park has released its cells, so the head is re-examined now; an asynchronous one on the pass that sees the copy land if (slot.state == SLOT_STATE_PREEMPTED) { best = head; } @@ -4241,8 +3942,6 @@ struct server_context_impl { best->t_preempt_copy_us = t_start; if (!best->preempt_restore()) { - // update_slots() loops tightly, so a counter alone burns its budget in - // milliseconds: give up only on a slot failing for a while, and log quietly if (best->n_preempt_fail % 64 == 1) { SLT_WRN(*best, "resume failed (%d in a row, parked %.1f s), staying preempted\n", best->n_preempt_fail, (ggml_time_us() - best->t_preempt_us) / 1e6); @@ -4257,8 +3956,7 @@ struct server_context_impl { break; } - // [TAG_PREEMPT_ASYNC] with a transfer the copy has only been issued; the slot is - // RESTORING and update_preempt_copies() counts it and logs it when it lands + // [TAG_PREEMPT_ASYNC] with a transfer the copy has only been issued; update_preempt_copies() counts and logs it when it lands if (best->state == SLOT_STATE_RESTORING) { SLT_WRN(*best, "resumed after %.2f s: %d tokens, restore issued in %.2f ms (%zu transfers, %.2f ms sync), kv %d/%d, preemptions %d\n", (ggml_time_us() - best->t_preempt_us) / 1e6, @@ -4268,16 +3966,12 @@ struct server_context_impl { preempt_kv_used(), n_cells, best->n_preempt); - // it holds cells now but is not decoding yet, so there is nothing more to - // decide about it this iteration continue; } metrics.n_resume++; - // [TAG_PREEMPT] the synchronous restore returns with the slot already back in - // the state it was parked from, so here the issue and the landing are the same - // moment; the asynchronous one announces from update_preempt_copies() instead. + // [TAG_PREEMPT] the synchronous restore returns with the slot already back in its old state, so issue and landing are the same moment here send_preempt_notice(*best, false); SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", @@ -4288,7 +3982,6 @@ struct server_context_impl { best->n_preempt); } - // forced preemption, for the determinism test only if (preempt_test_every > 0) { for (auto & slot : slots) { if (slot.state == SLOT_STATE_GENERATING && @@ -4299,17 +3992,11 @@ struct server_context_impl { if (slot.preempt_save()) { preempt_log_ram_kind(slot); - // [TAG_PREEMPT_ASYNC] a slot left PREEMPTING is counted by - // update_preempt_copies() when its copy lands, not here + // [TAG_PREEMPT_ASYNC] a slot left PREEMPTING is counted by update_preempt_copies() when its copy lands if (slot.state == SLOT_STATE_PREEMPTED) { metrics.n_preempt++; } - // [TAG_PREEMPT] the notice goes with the save, not with the cell - // release: preempt_save() has already detached the slot, so from - // here it takes no part in a decode and the stream is silent - // whether the cells have gone (PREEMPTED) or the copy still holds - // them (PREEMPTING). send_preempt_notice(slot, true); SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", @@ -4323,7 +4010,6 @@ struct server_context_impl { return; // test knob: leave the pool to the retry ladder and its last resort } - // and take cells back until the next decode fits for (;;) { const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); @@ -4331,18 +4017,11 @@ struct server_context_impl { break; } - // a prompt cached on an idle slot is the cheapest thing in the pool to give up if (try_clear_idle_slots()) { continue; } - // [TAG_PREEMPT_ASYNC] A park that has been issued but not landed is holding - // cells that are already spoken for. Out of room for the step about to be - // built, waiting for it is both quicker and more useful than parking somebody - // else, whose cells would not come back this iteration either. Short only of - // the lookahead the asynchronous path keeps, the step itself fits: it goes - // ahead beside the copy, which is the overlap the path exists for, and the - // planner looks again once the copy has landed. + // [TAG_PREEMPT_ASYNC] a park issued and not landed holds cells that are already spoken for, so waiting for it is quicker than parking somebody else if (preempt_copies_in_flight()) { if (n_used > n_cells) { if (preempt_wait_in_flight()) { @@ -4372,19 +4051,10 @@ struct server_context_impl { preempt_log_ram_kind(*victim); - // [TAG_PREEMPT] the notice goes with the save, not with the cell release. - // preempt_save() has already detached the victim, so from here it takes no part - // in a decode and its stream is silent -- whether the cells went with the save - // (PREEMPTED) or the copy still holds them (PREEMPTING). Announcing it at the - // release instead would leave the client one copy (~40 ms, and unbounded if the - // pool never frees) of silence with no explanation, which is the thing this is - // for. + // [TAG_PREEMPT] the notice goes with the save, not the cell release: preempt_save() has already detached the victim, so a release-time notice would leave the copy's silence unexplained send_preempt_notice(*victim, true); - // [TAG_PREEMPT_ASYNC] the copy has only been issued; the cells are still the - // victim's until it lands, so nothing further can be decided about the pool this - // iteration. update_preempt_copies() picks it up on the next one, and the step - // that wanted the room is built from whatever is free right now. + // [TAG_PREEMPT_ASYNC] the copy has only been issued and the cells are still the victim's, so nothing further can be decided about the pool this iteration if (victim->state == SLOT_STATE_PREEMPTING) { SLT_WRN(*victim, "preempted: %d cells, park issued in %.2f ms (%zu transfers, %.2f ms sync), %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, @@ -4394,22 +4064,8 @@ struct server_context_impl { preempt_kv_used(), n_cells, n_used, victim->n_preempt); - // [TAG_PREEMPT_ASYNC] Whether we may leave now depends on which of the two - // thresholds we are under. - // - // Short of the lookahead only: there is still room for the step about to be - // built, the park is early by design, and leaving is the whole point -- the - // copy runs beside the decode and update_preempt_copies() collects it next - // iteration. - // - // Out of room for the step itself: the cells are held until the copy lands, - // so leaving now builds a batch into a pool that has not got smaller. The - // decode fails, and the retry ladder halves n_batch to 1 without ever - // polling the copy, ending in "Context size has been exceeded" for every - // slot. The synchronous path did not have this problem because it returned - // the cells before it returned. Go round instead: the next pass reaches - // preempt_wait_in_flight() and waits for the park just issued, which is no - // worse than the synchronous path and is what it was written for. + // [TAG_PREEMPT_ASYNC] short of the lookahead only, the step still fits and leaving is the point; out of room + // for it the cells are held until the copy lands, so the retry ladder ends every request instead of waiting if (n_used + preempt_n_margin() > n_cells) { continue; } @@ -4428,14 +4084,8 @@ struct server_context_impl { } } - // the checks a slot's request has to pass before its prompt is processed, run from the - // SLOT_STATE_STARTED block below. true when the request is rejected, with the message and - // the type of the error it gets. The empty prompt is not here: it is a final response and - // not an error. - // [TAG_PREEMPT] the planner asks the same question before it parks a started slot, so a - // request that is about to be errored is never given a park notice ahead of its error: a - // notice opens the stream, and the client would get 200 plus an in-stream error where the - // non-stream 4xx belongs. + // the checks a request has to pass before its prompt is processed; true when it is rejected. An empty prompt is not here: it is a final response, not an error. + // [TAG_PREEMPT] the planner asks the same question before parking a started slot: a notice opens the stream, and a rejected request would get 200 plus an in-stream error bool slot_prompt_rejected(const server_slot & slot, std::string & msg, error_type & type) const { if (!slot.task) { return false; @@ -4496,7 +4146,6 @@ struct server_context_impl { } #endif - // check if all slots are idle { bool all_idle = true; @@ -4524,17 +4173,11 @@ struct server_context_impl { } try { - // [TAG_PREEMPT] make the pool fit the step that is about to be built, measured - // after any context shift. Inside the guard with the rest of the step: a shift - // rebuilds a slot's tokens and a park allocates, and either can throw, which the - // slots are told about rather than the loop ending on an uncaught exception + // [TAG_PREEMPT] make the pool fit the step about to be built, measured after any context shift; inside the guard because a shift or a park can throw pre_decode_shift(); update_preemption(); - // [TAG_PREEMPT_ASYNC] before pre_decode(), not only before the target decode: the - // draft it asks for is a decode on the draft context, which applies that cache's - // pending shift in place, and a park or restore still copying draft cells would - // read through it or be overwritten by it just the same + // [TAG_PREEMPT_ASYNC] before pre_decode(), not only the target decode: the draft it asks for applies the draft cache's pending shift in place preempt_wait_for_shift(); scoped_timer t(t_pre_decode, n_pre_decode); @@ -4573,9 +4216,7 @@ struct server_context_impl { int32_t off_next = 0; int32_t n_batch = llama_n_batch(ctx_tgt); - // [TAG_PREEMPT_ASYNC] and once more here: a shift --cache-reuse asks for is found - // inside pre_decode(), after the wait above, and the decode below applies it in - // place like any other + // [TAG_PREEMPT_ASYNC] and once more here: a shift --cache-reuse asks for is found inside pre_decode(), after the wait above preempt_wait_for_shift(); for (int32_t off = 0; off < batch.size(); off = off_next) { @@ -4839,8 +4480,7 @@ struct server_context_impl { return; // batch is full, skip remaining slots } - // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to - // batch; it takes no part in this pass until it is restored + // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to batch until it is restored if (!slot.is_processing() || slot.preempt_is_out()) { return; } @@ -4966,8 +4606,7 @@ struct server_context_impl { slot.mem.seq_rm (slot.id, head_p, head_c); slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift); - // [TAG_PREEMPT_ASYNC] applied inside the next llama_decode by the - // same in-place graph as a context shift, see preempt_wait_for_shift + // [TAG_PREEMPT_ASYNC] applied in place inside the next llama_decode, like a context shift preempt_shift_pending = true; for (size_t i = 0; i < n_match; i++) { @@ -5343,16 +4982,8 @@ struct server_context_impl { } } - // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell. Upstream this is - // the context error for every slot in the batch. With a park budget the batch is given - // up instead: every resident slot is rewound to the token boundary the cache is at (a - // batch is applied one chunk at a time, and the chunk that failed left nothing behind), - // the smallest are parked until the planner's own bound holds again, and the next - // update_slots() rebuilds the batch from the survivors. The planner brings the parked - // ones back as cells free up. A multimodal prompt has no boundary the cache can name, - // so it keeps the old path. - // [TAG_PREEMPT_ASYNC] whether a park can happen and go asynchronously: the conditions - // update_preemption() gates on, and the asynchronous switch + // [TAG_PREEMPT] the retry ladder ran out: give the batch up, rewind every resident to the token boundary the cache is at and park the smallest. Multimodal keeps the old path. + // [TAG_PREEMPT_ASYNC] whether a park can happen at all and go asynchronously bool preempt_async_possible() const { return params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0 && slots.size() >= 2 && llama_get_memory(ctx_tgt) && !llama_model_is_recurrent(model_tgt); @@ -5424,8 +5055,7 @@ struct server_context_impl { send_preempt_notice(*victim, true); - // [TAG_PREEMPT_ASYNC] the cells are wanted now, not next iteration: wait for the - // copy to land, which releases them and logs the park the way the planner does + // [TAG_PREEMPT_ASYNC] the cells are wanted now, not next iteration: wait for the copy, which releases them if (victim->state == SLOT_STATE_PREEMPTING) { while (preempt_wait_in_flight()) { } @@ -5510,13 +5140,7 @@ struct server_context_impl { }); if (ret != 0) { - // [TAG_PREEMPT_ASYNC] Before giving up any batch width, and before the last resort: - // a park that has been issued and not yet landed is holding cells that are already - // spoken for, and waiting for it returns them. Halving the batch returns nothing, - // so without this the ladder can run all the way down to n_batch == 1 and end - // every request while the room it needed was moments from arriving. Safe from - // here because the slot was detached before this batch was built, so completing - // its park cannot change what the batch about to be retried contains. + // [TAG_PREEMPT_ASYNC] halving the batch returns no cells, so wait for an issued park first, or the ladder runs down to n_batch == 1 and ends every request if (ret == 1 && preempt_wait_in_flight()) { SRV_WRN("%s", "waited for an in-flight park before retrying the decode\n"); return false; // retry at the same batch size, with the cells it freed @@ -5525,9 +5149,7 @@ struct server_context_impl { { std::string err; - // [TAG_PREEMPT] a slot's sampled token and its draft have to stay in one view, - // so halving would split the group and make the verify step throw: after the - // idle slots the ladder goes straight to its last resort + // [TAG_PREEMPT] a slot's sampled token and its draft have to stay in one view, so halving would split the group and make the verify step throw if (ret == 1 && n_batch > 1 && preempt_last_resort_possible() && batch_has_spec_groups()) { if (try_clear_idle_slots()) { SRV_WRN("%s", "failed to find free space in the KV cache, retrying after purging an idle slot\n"); @@ -5564,8 +5186,7 @@ struct server_context_impl { SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret); for (auto & slot : slots) { - // [TAG_PREEMPT] a parked slot has nothing in this batch and nothing in the - // cache; it is not part of this failure and comes back when there is room + // [TAG_PREEMPT] a parked slot is not part of this failure and comes back when there is room if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED && !slot.preempt_in_flight()) { send_error(slot, err); slot.release(); @@ -6173,9 +5794,7 @@ std::unique_ptr server_routes::handle_completions_impl( task.params.oaicompat_cmpl_id = completion_id; task.params.oaicompat_model = meta->model_name; - // [TAG_EXACT_CONCURRENCY] children of an n_cmpl > 1 task are served by copying the - // parent's cells to another sequence id, and exact mode gives a page to a single - // sequence, so refuse here where it becomes a 400 rather than at seq_cp + // [TAG_EXACT_CONCURRENCY] exact mode gives a page to a single sequence, so refuse an n_cmpl > 1 child here, where it becomes a 400 rather than at seq_cp if (task.params.n_cmpl > 1 && server_exact_concurrency()) { throw std::runtime_error( "n > 1 is not supported while LLAMA_EXACT_CONCURRENCY is set: each " @@ -6236,17 +5855,12 @@ std::unique_ptr server_routes::handle_completions_impl( // in streaming mode, the first error must be treated as non-stream response // this is to match the OAI API behavior // ref: https://github.com/ggml-org/llama.cpp/pull/16486#discussion_r2419657309 - // [TAG_PREEMPT] a slot can be parked while still processing its prompt, before any - // token exists. Those notices arrive ahead of the first real result; keep them and - // send them in front of it, so the client learns about the wait it just had. + // [TAG_PREEMPT] a slot can be parked before any token exists, so those notices are kept and sent in front of the first real result std::string preempt_prefix; std::set parked_idx; // prompts of this request that are parked right now auto first_result = rd.next(req.should_stop); if (first_result != nullptr && dynamic_cast(first_result.get()) != nullptr) { - // [TAG_PREEMPT] parked before any token exists. The stream starts now, with the - // notice, so the parked keepalive runs through the wait instead of the client - // seeing nothing until the slot resumes; the first ordinary result follows in - // the stream, an error included, since the response has already begun. + // [TAG_PREEMPT] the stream starts now, with the notice, so the parked keepalive runs through the wait instead of the client seeing nothing const auto * notice = static_cast(first_result.get()); preempt_prefix = preempt_notice_comment(*notice); if (notice->parked) { @@ -6285,8 +5899,7 @@ std::unique_ptr server_routes::handle_completions_impl( res->status = 200; res->content_type = "text/event-stream"; res->set_next([res_this = res.get(), res_type, sse_ping_interval, parked_idx](std::string & output) mutable -> bool { - // [TAG_PREEMPT] the keepalive runs while ANY prompt of the request is parked: with - // several prompts in one stream, one resuming does not mean the others did + // [TAG_PREEMPT] the keepalive runs while ANY prompt of the request is parked const bool parked = !parked_idx.empty(); static auto format_error = [](task_response_type res_type, const json & res_json) { @@ -6339,11 +5952,7 @@ std::unique_ptr server_routes::handle_completions_impl( // receive subsequent results bool timeout = false; int64_t start_time = ggml_time_ms(); - // [TAG_PREEMPT] a parked slot produces nothing for as long as the pool is - // full, so while parked the ping runs at least every 2 s whether or not - // --sse-ping asked for one, and is named, so a client can tell "waiting for - // cells" from "slow". A shorter interval the request asked for is kept: a - // client that wants a ping every second wants it most while nothing else comes. + // [TAG_PREEMPT] a parked slot produces nothing, so ping at least every 2 s whether or not --sse-ping asked for one, and name it; a shorter interval asked for is kept const int64_t ping_cfg = sse_ping_interval > 0 ? (int64_t) sse_ping_interval * 1000 : -1; const int64_t ping_ms = parked ? (ping_cfg > 0 ? std::min(ping_cfg, PREEMPT_KEEPALIVE_MS) : PREEMPT_KEEPALIVE_MS) : ping_cfg; auto result = rd.next([&timeout, &start_time, ping_ms, &effective_should_stop]() { @@ -6376,8 +5985,7 @@ std::unique_ptr server_routes::handle_completions_impl( SRV_DBG("%s", "error received during streaming, terminating stream\n"); return false; // terminate on error } else if (const auto * notice = dynamic_cast(result.get())) { - // [TAG_PREEMPT] an SSE comment: invisible to clients that do not know - // about preemption, a pause indicator for the ones that do + // [TAG_PREEMPT] an SSE comment: invisible to clients that do not know about preemption if (notice->parked) { parked_idx.insert(notice->index); } else { @@ -6391,10 +5999,7 @@ std::unique_ptr server_routes::handle_completions_impl( ); json res_json = result->to_json(); if (res_json.is_null()) { - // [TAG_PREEMPT] the signal a prompt sends before its first token, so - // that the headers go out, carries no data. Normally it is the first - // result and only opens the stream; after a notice opened the stream - // it has nothing to add, and the sender skips an empty chunk. + // [TAG_PREEMPT] the empty signal a prompt sends before its first token has nothing to add once a notice has opened the stream return true; } if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index bd74db89408..b5c8ab4a8ac 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -448,16 +448,7 @@ server_task_result_ptr server_response::recv(const std::unordered_set & id_ } server_task_result_ptr server_response::recv_with_timeout(const std::unordered_set & id_tasks, int timeout) { - // [TAG_PREEMPT] The timeout is a deadline, not a per-wait duration. - // - // send() notify_all()s on one condition variable for every result of every task, so a - // reader waiting on a task that is producing nothing is woken by every token every other - // task produces. With wait_for() each of those wakeups restarted the wait, and on a busy - // server the timeout was never reached at all: whoever was waiting for a quiet task - // waited forever, however small the timeout they asked for. That is exactly the - // situation of a parked slot, which by definition exists because the others are busy, so - // neither its 2 s keepalive nor the ordinary --sse-ping could ever fire for it. Waiting - // until a fixed point instead makes the timeout mean what every caller reads it as. + // [TAG_PREEMPT] the timeout is a deadline, not a per-wait duration: send() notify_all()s for every result of every task, and with wait_for() each wakeup restarted the wait const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout); while (true) { diff --git a/tools/server/server-task.h b/tools/server/server-task.h index f20f891655f..f731378b1f6 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -392,10 +392,7 @@ struct server_task_result_cmpl_final : server_task_result { json to_json_anthropic_stream(); }; -// [TAG_PREEMPT] out-of-band notice for a streaming task whose slot was parked or restored. -// Serialised as an SSE comment (": preempted", ": resumed"), which every existing client -// ignores, so the body of the response is unchanged by preemption. Never sent to a -// non-streaming task. +// [TAG_PREEMPT] out-of-band notice for a streaming task whose slot was parked or restored, sent as an SSE comment (": preempted", ": resumed") every existing client ignores struct server_task_result_preempt_notice : server_task_result { bool parked = false; // true when the slot was just parked, false when restored int32_t n_preempt = 0; // how many times this task has been parked so far diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 6a589f3cb06..56493ed3048 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -5,9 +5,7 @@ import pytest from utils import * -# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked (its -# sequence copied to host RAM, its cells released) instead of every slot being terminated. Needs -# more than one slot and --kv-unified, the only configuration where slots share cells. +# Preemption on a unified KV pool: one slot is parked, its sequence copied to host RAM and its cells released, instead of every slot being terminated. Needs --kv-unified. server = ServerPreset.tinyllama2() @@ -57,8 +55,7 @@ def _complete(n_predict: int, prompt: str = "Hi how are you"): def test_forced_preemption_does_not_change_the_output(): - # park and restore the only running slot every 8 tokens: the batch shape is the same at every - # step, so any difference in the output is the preemption's fault + # park and restore the only running slot every 8 tokens: the batch shape is the same at every step, so any difference in the output is the preemption's fault global server server.n_ctx = 512 server.start() @@ -85,9 +82,7 @@ def test_forced_preemption_does_not_change_the_output(): def test_two_slots_that_overflow_the_pool_together_both_finish(): - # each request fits the pool alone (168 of 256 cells) but not together (336). Without - # preemption both end with "Context size has been exceeded"; with it the smaller is parked - # until the leader finishes, then resumes from the token it was parked on. + # each request fits the pool alone (168 of 256 cells) but not together; without preemption both end with "Context size has been exceeded" global server server.n_ctx = 256 server.start() @@ -112,11 +107,7 @@ def test_two_slots_that_overflow_the_pool_together_both_finish(): def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages(): - # a pool that allocates in blocks gives a whole block to one sequence, so n tokens occupy - # round_up(n, block) cells and the planner has to count cells: counting tokens it sees room - # the allocator cannot find, never parks anybody, and the retry ladder ends every request. - # LLAMA_SERVER_PREEMPT_GRANULARITY injects the block size, since the only mode that reports - # one needs a head size this model does not have; the arithmetic is the same at 64 as at 256. + # a block allocator gives a whole block to one sequence, so the planner has to count cells: counting tokens it sees room the allocator cannot find. GRANULARITY injects the size. global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_GRANULARITY"] = "64" @@ -135,8 +126,6 @@ def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages(): assert "preempted:" in text assert "resumed after" in text - # every figure the planner logs is a whole number of blocks: "kv N/256" is what the pool holds - # and "(wanted N)" is that plus the next decode's reservation held = [int(n) for n in re.findall(r"kv (\d+)/256", text)] wanted = [int(n) for n in re.findall(r"\(wanted (\d+)\)", text)] assert held and wanted, f"the planner logged no figures:\n{text}" @@ -170,14 +159,11 @@ def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]: if n <= n_tokens: assert n >= n_tokens - 12, f"could not land near {n_tokens} tokens, got {n}" return text, n - # about four tokens per word on this model's vocabulary words = words[: len(words) - max(1, (n - n_tokens) // 8)] raise AssertionError("empty prompt") def test_two_prompts_that_overflow_the_pool_together_both_finish(): - # neither slot generates before the pool is full: a slot between two chunks of its prompt is - # as clean a boundary as one between two sampled tokens, so it is parked the same way global server server.n_ctx = 256 server.start() @@ -206,17 +192,13 @@ def test_two_prompts_that_overflow_the_pool_together_both_finish(): def test_a_generating_slot_and_a_large_prompt_both_finish(): - # a slot generating a long answer to a short prompt meets a large prompt arriving beside it, - # needing far more than the pool has: the prompt is admitted chunk by chunk, whoever is - # smaller is parked, and both finish. The second request follows immediately, since its - # prompt takes several batches and that is enough overlap however fast the first one runs. + # a long generation meets a large prompt arriving beside it: the prompt is admitted chunk by chunk, whoever is smaller is parked, and both finish global server server.n_ctx = 256 server.start() log = LogReader(server.log_path) prompt_b, n_b = _prompt_of_about(150, "Charlie") - # b has to live long enough for the two to collide n_predict_a = 230 n_predict_b = 90 assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 @@ -242,8 +224,7 @@ def _late(n_predict, prompt): def test_preempt_ram_zero_disables_preemption(): - # --preempt-ram 0 switches back to the old behaviour: nothing is parked and the KV-full path - # ends the requests the way it always did + # --preempt-ram 0 switches back to the old behaviour: nothing is parked and the KV-full path ends the requests global server server.n_ctx = 256 os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" @@ -263,8 +244,6 @@ def test_preempt_ram_zero_disables_preemption(): def test_metrics_and_slots_report_the_parked_state(): - # /slots tells a parked chat from a slow one and /metrics reports it to an operator; both - # must show the preemption, and the counters must survive the requests finishing global server server.n_ctx = 256 server.server_metrics = True @@ -301,11 +280,7 @@ def test_metrics_and_slots_report_the_parked_state(): assert sum(slot["n_preempt"] for slot in res.body) == 0, "n_preempt is per task and resets with the slot" -# [TAG_PREEMPT_ASYNC] parking and resuming on a stream of their own -# -# The copies are only asynchronous on a backend that can copy asynchronously and signal an -# event, which today means a GPU one. On a CPU-only build the server says so and falls back -# to the synchronous path, and the tests below that need the asynchronous one skip. +# [TAG_PREEMPT_ASYNC] parking and resuming on a stream of their own, only on a backend that can copy asynchronously and signal an event; a CPU-only build falls back and these skip _ASYNC_BANNER = "parking and resuming asynchronously" @@ -326,10 +301,7 @@ def _require_async(text: str): def test_async_preemption_does_not_change_the_output(): - # The same question the synchronous determinism test asks, of the asynchronous path: - # with one request the batch has the same shape at every step, so a continuation that - # was parked and resumed through a transfer and is not byte-identical to an - # uninterrupted one is the transfer's fault and nothing else's. + # the synchronous determinism question asked of the asynchronous path: with one request the batch shape is fixed, so a continuation that is not byte-identical is the transfer's fault global server server.n_ctx = 512 server.n_gpu_layer = 99 @@ -351,8 +323,6 @@ def test_async_preemption_does_not_change_the_output(): _require_async(text) assert text.count("preempted on request") >= 6 assert text.count("resumed after") >= 6 - # the asynchronous path is the one that ran, not the synchronous fallback: only it - # splits a park and a resume into an issue and a completion assert "park completed after" in text assert "restore issued in" in text assert "restore completed after" in text @@ -362,9 +332,6 @@ def test_async_preemption_does_not_change_the_output(): def test_async_preemption_under_load_keeps_every_slot_and_its_output(): - # Two requests that do not fit the pool together, parked and resumed asynchronously - # while the other one keeps decoding. Every slot must finish, and finish with exactly - # the tokens it produces when it has the pool to itself. global server server.n_ctx = 256 server.n_gpu_layer = 99 @@ -377,7 +344,6 @@ def test_async_preemption_under_load_keeps_every_slot_and_its_output(): "The quick brown fox jumps over the lazy dog and", ] - # each one alone, for the reference tokens alone = [_complete(n_predict, prompt) for prompt in prompts] for res in alone: assert res.status_code == 200 @@ -415,16 +381,10 @@ def _cancel_soon(n_predict: int, prompt: str, timeout: float): def test_cancel_while_a_copy_is_in_flight_frees_the_slot(): - # A cancelled request can reach release() with a park or a resume still running, which - # is where the host buffer is freed and the cells are handed on. Both have to wait for - # the copy first. LLAMA_SERVER_PREEMPT_EVERY keeps every slot cycling between the two - # states, so cancelling at a spread of moments lands in both; what is asserted is that - # the server survives it, the slots come back, and it still answers correctly. + # a cancelled request can reach release() with a park or a resume still running, where the host buffer is freed and the cells handed on, so both have to wait for the copy global server server.n_ctx = 512 server.n_gpu_layer = 99 - # every 8 tokens, so a slot spends most of its life in one of the two copy states, but - # not so often that the abandoned requests take minutes to drain os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" text = _start_async() _require_async(text) @@ -432,7 +392,6 @@ def test_cancel_while_a_copy_is_in_flight_frees_the_slot(): for i in range(4): _cancel_soon(96, "Once upon a time there was a brave knight who", 0.05 + 0.1 * i) - # every slot back, and none of them still holding a parked sequence deadline = time.time() + 120 while time.time() < deadline: res = server.make_request("GET", "/slots") @@ -452,7 +411,6 @@ def test_cancel_while_a_copy_is_in_flight_frees_the_slot(): if line.startswith("llamacpp:preempt_ram_bytes"): assert float(line.split(" ", 1)[1]) == 0, "a cancelled slot kept its parked memory" - # and the server still works res = _complete(16) assert res.status_code == 200 assert res.body["timings"]["predicted_n"] == 16 @@ -475,24 +433,12 @@ def test_no_preempt_async_falls_back_to_the_synchronous_path(): assert _ASYNC_BANNER not in text assert "park issued in" not in text assert "restore issued in" not in text - # the synchronous path still parks and resumes assert text.count("preempted on request") >= 6 assert text.count("resumed after") >= 6 def test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_everything(): - # [TAG_PREEMPT_ASYNC] The case the async path made worse than the synchronous one, and - # that the existing tests miss because their victim holds almost no cells. - # - # Three slots are well into generating when a fourth request arrives whose prompt does - # not fit in what is left. update_preemption() picks a victim and issues its park, but - # an asynchronous park does not return the cells before update_slots() carries on. If - # the loop leaves at that point, the batch is built into a pool that has not got any - # smaller, llama_decode returns 1, and the retry ladder halves n_batch to 1 in - # microseconds without ever polling the copy -- ending every request with "Context size - # has been exceeded" while the room it wanted was one event query away. - # - # Pass is what the synchronous path gave: a park, and all four requests finish. + # [TAG_PREEMPT_ASYNC] the case the async path made worse than the synchronous one: an asynchronous park does not return the cells before update_slots() carries on global server server.n_ctx = 512 server.n_gpu_layer = 99 @@ -505,17 +451,12 @@ def test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_everyt prompt_c, n_c = _prompt_of_about(100, "Charlie") prompt_d, n_d = _prompt_of_about(150, "Delta") - # A, B and C oversubscribe the pool between them, so the pressure does not depend on - # when D arrives, and every occupant is holding real cells rather than the handful the - # other tests park. Each of the four still fits on its own. n_predict_abc = 130 n_predict_d = 40 assert max(n_a, n_b, n_c) + n_predict_abc < 512 and n_d + n_predict_d < 512 assert n_a + n_b + n_c + 3 * n_predict_abc > 512 def _late(n_predict, prompt): - # D's prompt arrives into a pool the other three have already grown into; this - # model decodes about 120 tokens a second, so they are all still running time.sleep(0.25) return _complete(n_predict, prompt) @@ -538,17 +479,13 @@ def _late(n_predict, prompt): def test_two_prompts_near_the_context_size_both_complete(): - # two prompts that each fit the context alone but not together. The second is parked before - # it takes any cells and is too close to n_ctx to leave the usual margin, but must still be - # restored once the first finishes: with nothing resident there is nobody to keep it for. + # the second prompt is parked before it takes any cells and is too close to n_ctx to leave the usual margin, but must still be restored once the first finishes global server server.n_ctx = 256 - # the whole prompt in one batch, so the parked slot's first step is the whole prompt server.n_batch = 256 server.start() log = LogReader(server.log_path) - # sized in tokens, not words: the prompt is the token ids of a short sentence repeated base = server.make_request("POST", "/tokenize", data={"content": "Once upon a time there was a little girl"}).body["tokens"] long_prompt = (base * 64)[:240] n_predict = 4 @@ -563,9 +500,6 @@ def test_two_prompts_near_the_context_size_both_complete(): def test_the_last_resort_parks_instead_of_ending_everyone(): - # with the planner off, two generations that fit alone but not together fill the pool until a - # single token finds no cell, where upstream ends every slot with the context error. Instead - # the batch is given up, the smaller slot is parked, and both finish. global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" @@ -594,8 +528,7 @@ def test_the_last_resort_parks_instead_of_ending_everyone(): def test_the_last_resort_works_with_an_unlimited_budget(): - # --preempt-ram -1 is the documented unlimited setting; it must enable the last resort - # the same as any positive budget does + # --preempt-ram -1 is the documented unlimited setting and must enable the last resort too global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" @@ -619,8 +552,7 @@ def test_the_last_resort_works_with_an_unlimited_budget(): def test_the_last_resort_rewinds_a_prompt_in_flight(): - # same, with a prompt being processed when the pool runs out: the failed chunk comes back off - # the slot's tokens and is processed again after the resume, neither skipped nor fed twice + # the failed chunk comes back off the slot's tokens and is processed again after the resume, neither skipped nor fed twice global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" @@ -650,16 +582,12 @@ def _late(n_predict, prompt): assert results[0].body["timings"]["predicted_n"] == n_predict_a assert results[1].status_code == 200 assert results[1].body["timings"]["predicted_n"] == n_predict_b - # the chunk that was in the batch given up is processed once, after the rewind, and - # the count is the prompt plus the BOS the server adds + # the chunk in the batch given up is processed once after the rewind; the count is the prompt plus the BOS the server adds assert results[1].body["timings"]["prompt_n"] == n_b + 1 def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_head(): - # two generations that each outgrow the pool, with context shift on: the resident shifts and - # would hold half the pool for as long as it generates, while the parked one never fits - # beside it. After the head has waited its turn the resident is parked in its place and the - # two take turns. n_predict is large enough that the resident is still going by then. + # with context shift on the resident would hold half the pool for as long as it generates, so once the head has waited its turn the resident is parked and the two take turns global server server.n_ctx = 256 server.enable_ctx_shift = True @@ -683,8 +611,6 @@ def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_hea def test_the_rotation_parks_a_resident_that_lets_the_head_in(): - # three endless generations with context shift on: two residents cycle through shifts while - # the third waits parked, and every rotation must let the head in so no stream ends short global server server.n_slots = 3 server.n_ctx = 384 @@ -710,9 +636,7 @@ def test_the_rotation_parks_a_resident_that_lets_the_head_in(): def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_server_lives(): - # a two-completion request is one conversation in two slots, and a family member is not a - # victim for the other, so with nobody else to park it gets the context error it would get - # alone and the server carries on serving + # a family member is not a victim for the other, so a two-completion request gets the context error it would get alone and the server carries on global server server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" @@ -741,15 +665,8 @@ def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_ def test_a_restored_slot_gives_its_idle_buffer_back_when_another_slot_needs_to_park(): - # Under a finite --preempt-ram an asynchronous slot keeps its pinned buffer after a - # restore, for its next park, and that idle capacity counted against the budget. With - # a budget that holds one sequence, the first restore spent it for good: every later - # park of the other slot was refused. The idle buffer is given back when another slot - # needs the room, and both slots go on being parked. + # an asynchronous slot keeps its pinned buffer after a restore, and that idle capacity counts against --preempt-ram: unless it is given back, the first restore spends the budget global server - # a pool of 8192 cells, but the model's own window is 2048, so each generation stays - # under that; 1800 tokens of this model's state is about 1.1 MiB, so a budget of - # 2 MiB holds one sequence and not two server.n_ctx = 8192 server.n_gpu_layer = 99 os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "256" @@ -778,21 +695,13 @@ def test_a_restored_slot_gives_its_idle_buffer_back_when_another_slot_needs_to_p def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_when_a_resident_finishes(): - # Three generations with no end in a pool one of them fills, with context shift on, - # under a --preempt-ram that holds the two parked heads but not a head and the resident - # at once. The resident is parked before the head is restored and freed, so a rotation - # holds both states together: under this budget the first one asked for is refused and - # said so, and the heads come back when the resident finishes instead. Every stream - # still finishes its tokens and nothing gets the context error. + # a rotation holds both states at once, since the resident is parked before the head is restored and freed, so a budget for two heads but not a head plus the resident must refuse global server server.n_slots = 3 server.n_ctx = 2048 server.enable_ctx_shift = True os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" server.start() - # long enough that the resident is still cycling through shifts two seconds after the - # heads were parked, which is when a rotation is first asked for: at 6000 this model - # finished in under three seconds on a fast host and nothing was ever refused n_predict = 12000 prompts = [ "Once upon a time there was a brave knight who", @@ -814,9 +723,7 @@ def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_w def test_a_recurrent_model_is_served_without_preemption(): - # A recurrent cache holds one state per sequence whatever its length, so the token - # count the planner measures says nothing about it: preemption is off for such a - # model, said so at load, and the forced-park knob parks nothing. + # a recurrent cache holds one state per sequence whatever its length, so preemption is off for such a model and the forced-park knob parks nothing global server path = os.environ.get("LLAMA_SERVER_TEST_RECURRENT_MODEL") if path: diff --git a/tools/server/tests/unit/test_preempt_notify.py b/tools/server/tests/unit/test_preempt_notify.py index 6fa71c55473..796c22b1930 100644 --- a/tools/server/tests/unit/test_preempt_notify.py +++ b/tools/server/tests/unit/test_preempt_notify.py @@ -6,11 +6,7 @@ import requests from utils import * -# [TAG_PREEMPT] A streaming client is told when its slot is parked and when it is -# restored, as SSE comments, and the body is byte for byte what it is without any park. -# Comments are legal SSE that every existing client ignores; a client that knows about -# preemption can show "paused" instead of a dead stream, and a keepalive every 2 s while -# parked keeps proxies and read timeouts from giving up on a wait that is by design long. +# [TAG_PREEMPT] a streaming client is told when its slot is parked and restored, as SSE comments every existing client ignores; a keepalive every 2 s keeps proxies from giving up server = ServerPreset.tinyllama2() @@ -102,11 +98,8 @@ def test_a_stream_announces_its_parks_and_the_body_is_unchanged(): resumed = [c for c in comments if c == ": resumed"] assert len(parked) >= 6, comments assert len(resumed) == len(parked), comments - # Every park is followed by its resume before the next park. seq = [c for c in comments if c in (": preempted", ": resumed")] assert seq == [": preempted", ": resumed"] * len(parked), seq - # The generated text is byte for byte the unparked text, token by token. Only the - # final chunk's wall-clock timings differ between the two runs. def _pieces(ds): return [json.loads(d).get("content") for d in ds if d != "[DONE]"] @@ -145,9 +138,6 @@ def test_non_streaming_requests_see_nothing(): def test_two_overflowing_streams_both_finish_and_the_parked_one_says_so(): - # The pair from test_preempt: each alone fits, together they do not, so one is - # parked until the other finishes. The parked stream must carry the comments and - # finish with its full output. global server server.n_ctx = 256 server.start() @@ -171,13 +161,8 @@ def test_two_overflowing_streams_both_finish_and_the_parked_one_says_so(): def test_a_stream_parked_before_its_first_token_starts_with_the_notice(): - # A request parked while still processing its prompt has no token to send yet. The - # response must not wait for one: it starts with the notice, so the client sees - # "paused" and gets the keepalive at once, instead of a silent connection that only - # opens when the slot resumes. + # A request parked while still processing its prompt has no token to send yet, so the response starts with the notice instead of a silent connection. global server - # The resident keeps growing towards the whole pool; the newcomer's prompt is larger - # than what is free beside it, so the planner parks the newcomer before it has a token. global server server.n_ctx = 512 server.n_batch = 512 # the whole prompt in one batch, so the planner sees its size at once @@ -210,8 +195,6 @@ def _run(name, payload, started=None): second_lines = [(ts, line) for ts, name, line in timeline if name == "second"] first_end = max(ts for ts, name, _ in timeline if name == "first") - # The notice is the very first thing on the wire, and it arrives while the other - # stream is still running, not when it has finished and the parked slot resumes. assert second_lines[0][1] == ": preempted", second_lines[:3] assert second_lines[0][0] < first_end events = [line for _, line in second_lines if line in (": preempted", ": resumed") or line.startswith("data: ")] @@ -223,10 +206,6 @@ def _run(name, payload, started=None): def test_a_resident_rotated_out_for_a_parked_head_is_told_so(): - # The rotation from test_preempt: a resident cycling through context shifts holds the - # pool, and after the head has waited its turn the resident is parked in its place. - # That park is a park like any other, so its stream must say so, and every notice - # must be paired: no stream ends with a park it was never told about. global server server.n_ctx = 256 server.enable_ctx_shift = True @@ -245,16 +224,11 @@ def test_a_resident_rotated_out_for_a_parked_head_is_told_so(): seq = [c for c in comments if c in (": preempted", ": resumed")] assert seq == [": preempted", ": resumed"] * (len(seq) // 2), seq n_parked += len(seq) // 2 - # Both streams took turns: at least one park each, so at least two in all. assert n_parked >= 2, [r[0] for r in results] def test_an_oversized_prompt_is_errored_instead_of_parked(): - # A slot that has just been given a task has not passed the prompt checks yet: they - # run on its first pass through update_slots. Parked before that, it would be told - # ": preempted" first, and the notice opens the stream, so a prompt larger than the - # context would come back as 200 plus an in-stream error instead of the plain error - # response it gets with nothing running. The planner leaves such a slot alone. + # A slot just given a task has not passed the prompt checks yet, and a notice opens the stream, so parking it would turn a plain error response into 200 plus an in-stream one. global server server.n_ctx = 512 server.n_batch = 512 # the whole prompt in one batch, so the planner sees its size at once From 951a21931f3ace23595d5232279cc4a4b332dc08 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 08:23:28 +0000 Subject: [PATCH 74/81] server : drop the unrelated tooling edits from the preemption change The scripts/unsloth and .github diffs are comment-only edits that a comment pass on a sibling branch made; they are not part of preemption. scripts/batchinv is a standalone divergence harness with no reference from the tree. --- .github/workflows/unsloth-pin-preflight.yml | 57 ++-- .github/workflows/unsloth-pr-set-lint.yml | 6 +- .github/workflows/unsloth-prebuilt.yml | 23 +- scripts/batchinv/README.md | 71 ---- scripts/batchinv/bench.py | 53 --- scripts/batchinv/divergence.py | 215 ------------ scripts/batchinv/probe.cpp | 359 -------------------- scripts/batchinv/prompts.py | 53 --- scripts/unsloth/additive_merge.py | 33 +- scripts/unsloth/feature_matrix.py | 9 +- scripts/unsloth/pin_contract.py | 26 +- scripts/unsloth/test_additive_merge.py | 5 +- scripts/unsloth/test_pin_contract.py | 7 +- 13 files changed, 127 insertions(+), 790 deletions(-) delete mode 100644 scripts/batchinv/README.md delete mode 100644 scripts/batchinv/bench.py delete mode 100644 scripts/batchinv/divergence.py delete mode 100644 scripts/batchinv/probe.cpp delete mode 100644 scripts/batchinv/prompts.py diff --git a/.github/workflows/unsloth-pin-preflight.yml b/.github/workflows/unsloth-pin-preflight.yml index 494bca228c8..f390f381e6f 100644 --- a/.github/workflows/unsloth-pin-preflight.yml +++ b/.github/workflows/unsloth-pin-preflight.yml @@ -28,10 +28,16 @@ permissions: contents: write issues: write -# Two runs of the same ref probe the same pins against the same base, so the second only -# competes for runners; newest wins, since it sees the newest pr-set.json. Per ref, not -# globally: this also runs on any push touching pr-set.json, and with one shared group a push -# to a second branch cancelled the first branch's run (09-03). +# Two runs of the same ref probe the same pins against the same base, so the +# second adds nothing and just competes for runners. On 08-04 a dispatch and the +# schedule sat queued together for an hour. Newest wins: it sees the newest +# pr-set.json. +# +# Per ref, though, not globally. This file also runs on any push that touches +# pr-set.json, so with one shared group a push to a second branch cancelled the +# first branch's run: observed on 09-03, where the run that would have said +# whether a repin fixed the nightly was cancelled by an unrelated branch, and +# the PR was left showing the failure from before the fix. concurrency: group: unsloth-pin-preflight-${{ github.ref }} cancel-in-progress: true @@ -50,8 +56,10 @@ jobs: id: p run: | set -uo pipefail - # everything below reports through `status`/`details`, so a death anywhere else - # leaves the alert blank; report the abort through the same channel as a finding + # Everything below reports through `status`/`details`, so a death + # anywhere else leaves both empty and the alert blank: a red X on a + # scheduled run nobody opens. Report the abort through the same + # channel as a finding, so the repin bot sees a failure either way. trap 'rc=$?; if [ "$rc" != 0 ]; then { echo "status=failure" echo "details</dev/null PROBLEMS="${PROBLEMS}- \`${SRC}#${NUM}\` (\`${SHA:0:10}\`) does not merge onto \`${BASE}\` + the pins before it.\n\n Conflicting files:\n\n\`\`\`\n${FILES}\n\`\`\`\n\n
conflict hunks\n\n\`\`\`diff\n${HUNKS}\n\`\`\`\n\n
\n" @@ -180,9 +193,11 @@ jobs: PROBLEMS="${PROBLEMS}- the merged tree builds, but \`scripts/unsloth/merge_checks.py\` found a resolution that is silently wrong. See the run log for file and line.\n" fi - # the other half: merge_checks.py asks whether the tree contains something - # wrong, this asks whether it still contains what each pin carries. A pin rotted - # into a no-op is invisible to every other check here and to the compiler. + # The other half of that question. merge_checks.py asks whether the + # tree contains something wrong; this asks whether it still contains + # what each pin carries. A pin that has rotted into a no-op, or an + # arch registration a resolution quietly dropped, is invisible to + # every other check here and to the compiler. if ! python3 ../scripts/unsloth/pin_contract.py --root . --base "$BASE" \ --pr-set ../scripts/unsloth/pr-set.json --report "${RUNNER_TEMP}/pin_contract.json" ; then PROBLEMS="${PROBLEMS}- the merged tree is missing code a pin carries. See the run log for the pin and file.\n" @@ -192,10 +207,12 @@ jobs: PROBLEMS="${PROBLEMS}- pins upstream has taken over, safe to delete from \`pr-set.json\`:\n\n\`\`\`\n${NOTES}\n\`\`\`\n" fi - # a clean merge is not a compiling tree: on 09-03 ggml-org#27754 merged with no - # conflicts and did not compile, upstream having added a parameter to - # build_attn_mha that the pin's build_attn_sparse still called without. CPU only, - # 59s cold at -j4 with no ccache. + # A clean merge is not a compiling tree. On 09-03 ggml-org#27754 + # merged with no conflicts at all and did not compile: upstream had + # added a parameter to build_attn_mha and the pin's new + # build_attn_sparse still called the old signature. Nothing above + # can see that. CPU only and the `llama` target only, which is where + # that translation unit lives; 59s cold at -j4 with no ccache. GATE_OK=1 if ! cmake -B "${RUNNER_TEMP}/gate" -DCMAKE_BUILD_TYPE=Release \ -DGGML_CUDA=OFF -DLLAMA_BUILD_TESTS=ON -DLLAMA_BUILD_SERVER=OFF \ @@ -206,8 +223,10 @@ jobs: PROBLEMS="${PROBLEMS}- the pins merge cleanly and the merged tree does not compile. See the run log for the file and line; this is the failure that only shows up in the CUDA leg once the nightly has fanned out.\n" fi - # the only question that needs a binary: does each feature we ship still work. - # CPU only, since no runner here has a GPU; see the note in feature_matrix.py. + # The last question, and the only one that needs a binary: does each + # feature we ship still work. Everything above is about the source. + # CPU only, because no runner in this pipeline has a GPU -- see the + # note in feature_matrix.py about what that does and does not prove. if [ -n "$GATE_OK" ]; then if ! python3 ../scripts/unsloth/feature_matrix.py \ --build-dir "${RUNNER_TEMP}/gate" \ diff --git a/.github/workflows/unsloth-pr-set-lint.yml b/.github/workflows/unsloth-pr-set-lint.yml index 99b06b74f34..8a898539f02 100644 --- a/.github/workflows/unsloth-pr-set-lint.yml +++ b/.github/workflows/unsloth-pr-set-lint.yml @@ -128,9 +128,9 @@ jobs: done exit "$fail" - # a pin nobody decided about is the failure this file exists to stop: being in - # `unchecked` with a reason is fine, being in neither map is how DiffusionGemma went - # five weeks with no coverage and no record of it + # A pin nobody decided about is the failure this whole file exists to stop. + # Being in `unchecked` with a reason is a fine answer; being in neither map + # is how DiffusionGemma went five weeks with no coverage and no record of it. - name: Every pin is either checked or knowingly unchecked run: | set -euo pipefail diff --git a/.github/workflows/unsloth-prebuilt.yml b/.github/workflows/unsloth-prebuilt.yml index 353d88417ca..834ff544909 100644 --- a/.github/workflows/unsloth-prebuilt.yml +++ b/.github/workflows/unsloth-prebuilt.yml @@ -282,8 +282,8 @@ jobs: # .github/workflows, which upstream history routinely does). if [ "$EXISTS" != "true" ] || [ "${{ github.event_name }}" = "workflow_dispatch" ]; then git remote add upstream https://github.com/ggml-org/llama.cpp.git - # the upstream checkout below takes scripts/unsloth/ away; copy the whole dir - # out, not file by file, see the note above the step + # The upstream checkout below takes scripts/unsloth/ away. Copy the + # whole dir out, not file by file: see the note above the step. cp -r scripts/unsloth "${RUNNER_TEMP}/us" ADDITIVE_MERGE="${RUNNER_TEMP}/us/additive_merge.py" if [ "$(jq length <<<"$PRS")" != 0 ]; then @@ -446,7 +446,8 @@ jobs: # A bad pin resolution can still build fine, so it must be caught before the source artifact ships. See merge_checks.py. # Its own step, not more script in `resolve`: GitHub caps one workflow string at 21000 chars and that step is near it. See check_workflow_scalars.py. - # That is also why `resolve` copies all of scripts/unsloth/ to ${RUNNER_TEMP}/us in one line rather than one cp per script: another line inside the capped block per check would eventually go over, which silently disables the whole workflow. + # That is also why `resolve` copies all of scripts/unsloth/ to ${RUNNER_TEMP}/us in one line rather than one cp per script: every check added + # here would otherwise cost another line inside the capped block, and going over silently disables the whole workflow. - name: Check the merged tree for silently wrong resolutions if: ${{ env.MERGED_PINS == '1' }} run: | @@ -456,11 +457,13 @@ jobs: exit 1 fi - # merge_checks.py asserts the ABSENCE of two known-bad shapes; this asserts the PRESENCE of what each pin carries, the question that goes unanswered when a pin rots into a no-op. Free, so it runs before the compile gate. + # merge_checks.py asserts the ABSENCE of two known-bad shapes. This asserts the PRESENCE of what each pin carries, which is a different question and + # the one that goes unanswered when a pin rots into a no-op or a resolution quietly drops an arch registration. Free, so it runs before the compile gate. - name: Check every pin still contributes what it carries if: ${{ env.MERGED_PINS == '1' }} - # through env, never interpolated: `prs` carries PR titles, which are third-party - # text, and `${{ }}` pastes them into the shell source before bash sees it + # Through env, never interpolated into the script: `prs` carries PR + # titles, which are third-party text, and `${{ }}` pastes them into the + # shell source before bash ever sees it. env: PRS: ${{ steps.r.outputs.prs }} BASE: ${{ steps.r.outputs.base }} @@ -472,8 +475,12 @@ jobs: exit 1 fi - # The gap this closes, observed 09-03: ggml-org#27754 merged with zero conflicts and did not compile, upstream having added a parameter to build_attn_mha that the pin's build_attn_sparse still called without. Without this the release dies in the CUDA leg after the 38-job fan-out. CPU only: a cold `llama` build took 59s at -j4, against 20-60 minutes for CUDA. - # mtmd is in the gate because `llama` alone is not enough: on 09-04 ggml-org#25731 built `llama` clean while tools/mtmd did not compile at all, upstream having made mtmd_image_preprocessor::preprocess const while the pin's Inkling subclass stayed non-const. Every vision pin lands in mtmd. + # The gap this closes, observed 09-03: ggml-org#27754 merged with zero conflicts and did not compile, because upstream had added a parameter to + # build_attn_mha and the pin's new build_attn_sparse still called the old signature. Nothing before this point can see that, and without it the release + # dies in the CUDA leg after the 38-job fan-out. CPU only: a cold `llama` build took 59s at -j4 with no ccache, against 20-60 minutes for a CUDA build. + # mtmd is in the gate because `llama` alone is not enough: observed 09-04, ggml-org#25731 built `llama` clean while tools/mtmd did not compile at all, + # upstream having made mtmd_image_preprocessor::preprocess const while the pin's Inkling subclass stayed non-const, so it overrode nothing and the + # vision and audio towers were abstract. Every vision pin lands in mtmd, so a gate that skips it cannot see the whole class. - name: Compile gate (CPU, llama and mtmd targets) if: ${{ env.MERGED_PINS == '1' }} run: | diff --git a/scripts/batchinv/README.md b/scripts/batchinv/README.md deleted file mode 100644 index 82b769b6d04..00000000000 --- a/scripts/batchinv/README.md +++ /dev/null @@ -1,71 +0,0 @@ -# Exact concurrency experiment p - -Opt in before loading the model with `LLAMA_EXACT_CONCURRENCY=1`. This also forces -`GGML_CUDA_BATCH_INVARIANT=2` and gives `GGML_CUDA_BATCH_INVARIANT_MAX_COLS` a -default. The column policy only has to cover the widest ubatch a decode step can -build, one column per slot times one plus the draft length, because a prompt -ubatch is kept to one sequence and gets its exactness from that instead. Tools -built on `common` report that width, so the default is `--parallel` times one plus -`--spec-draft-n-max`, and an explicitly set `GGML_CUDA_BATCH_INVARIANT_MAX_COLS` -smaller than it is refused at startup. Nothing reported a width, the default is 16 -and the dispatcher warns once the first time a `MUL_MAT` or `MUL_MAT_ID` above the -bound is left unsplit. Set the variable to `0` for no bound; above the bound the -column policy does not fire, including during prefill. - -The experimental policy supports unified, offloaded F16 K/V, causal flash attention, -256-dimensional K and V heads, no attention soft cap, and no sliding window. -Shared-weight matmuls over multiple sequence planes are normalized to one plane -before the inherited selective column dispatcher. Without this, the recurrent -output projection bypasses batch invariance during concurrent prefill. -It is measured on text prompts with Qwen3.5-4B on one B200. Every KV layer has to -be on the CUDA backend, since no other backend reads the page table; a partial or -absent offload fails the load naming the layer. The V-less attention layouts have -no page table either, and a model on one of those is refused at context creation. -Context shifting, position division, cross-sequence prefix copies, shared-prefix -input tokens, and whole-context state loading are unsupported. Per-sequence state -save and restore is supported. Unsupported cache transformations are refused with -a logged error and leave the cells untouched; `--context-shift` and -`--cache-reuse` are reported as unsupported at load and disabled there. - -The allocator owns pages of 256 cells on behalf of one (sequence, position/256). -Position modulo 256 fixes the cell offset. Empty pages remain in the unified pool -and can be allocated by any sequence. The metadata is derived from live cells so -allocation rollback, tail removal, and sequence removal do not need another -transaction log. Restoring a sequence allocates free pages from this same pool. -The cost is up to 255 reserved cells per active sequence tail, plus holes introduced -by partial range removal. - -Attention receives an I32 page table in source 5: `[count, physical page IDs...]` -for each query, sorted by logical position. The physical K/V view and mask span -the pool, but the attention loop only visits the query's logical pages. The -final page is padded to 256 cells using the existing causal mask. Wholly future -pages in a prefill ubatch are excluded from the query's table. - -The vector attention specialization runs one query per block, four warps, and -`parallel_blocks=1`. It reads K/V directly from the physical pages, with two -128-cell softmax iterations per page, in logical page order. There is no K/V -gather and no split-K combine. The default vector specialization has no page -lookup. The ordinary path allocates no page metadata and launches no extra kernels. - -`FATTN_KQ_STRIDE=256` is a mask-scan stride, not the actual MMA rescaling tile. -For K/V head size 256, the Ampere-or-newer MMA configurations use 64 KV rows at -8 query/head columns and 32 KV rows at 16/32/64 columns. The retained vector path -has a 128-cell iteration. Both divide the 256-cell placement page. - -The probe is adapted from the existing batch-invariant harness and rejects -nonfinite logits and attention. `PROBE_B_REVERSE=1` fills neighbours before P0; -`PROBE_RESTORE=1` parks P0, releases a neighbour, restores P0, then rebuilds the -neighbour. Compute rows can be compared across this relocation; physical cache -views and index tensors must not be mistaken for sequence-0 compute outputs. - -`divergence.py --reference FILE` compares against an existing unparked solo token -reference. `bench.py --modes 0,1 --pairs 3` measures default off against exact mode -on, with 256 predicted tokens. Set `UNSLOTH_WORKSPACE` to the model parent workspace -and `LD_LIBRARY_PATH` to this build's bin directory. The harness uses GPU 3; -select a port in 9601-9610 explicitly. - -A concurrent request that fails now fails the run instead of being dropped from -the result, and every run records the environment the server actually inherited, -including `LLAMA_EXACT_CONCURRENCY`, under `env` in the JSON and in the server log -header, so a run labelled as the mode-off reference can be checked rather than -trusted. Teardown is POSIX only. diff --git a/scripts/batchinv/bench.py b/scripts/batchinv/bench.py deleted file mode 100644 index 64d8b32244b..00000000000 --- a/scripts/batchinv/bench.py +++ /dev/null @@ -1,53 +0,0 @@ -#!/usr/bin/env python3 -"""Cost of the knob: solo tok/s and four-chat aggregate tok/s, knob off and on, back to back.""" -import argparse, json, os, sys, time - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from divergence import Server, completion, run_concurrent -from prompts import PROMPTS - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--binary", required=True) - ap.add_argument("--spec", default="none") - ap.add_argument("--n-predict", type=int, default=256) - ap.add_argument("--pairs", type=int, default=3) - ap.add_argument("--port", type=int, default=9602) - ap.add_argument("--modes", default="0,1") - ap.add_argument("--out", required=True) - a = ap.parse_args() - - modes = a.modes.split(",") - rows = [] - for pair in range(a.pairs): - for mode in modes: - env = {"LLAMA_EXACT_CONCURRENCY": mode, "GGML_CUDA_BATCH_INVARIANT": "0" if mode == "0" else "2"} - with Server(a.port, a.binary, [], env, a.out + ".server.log", a.spec) as s: - completion(a.port, PROMPTS["P0"], 32) # warm - solo = completion(a.port, PROMPTS["P0"], a.n_predict) - outs, wall = run_concurrent(a.port, ["P0", "P1", "P2", "P3"], a.n_predict) - row = { - "pair": pair, "mode": mode, "spec": a.spec, - "solo_tok_per_s": solo["timings"]["predicted_per_second"], - "solo_prompt_tok_per_s": solo["timings"]["prompt_per_second"], - "four_aggregate_tok_per_s": sum(o["timings"]["predicted_per_second"] for o in outs.values()), - "four_wall_s": wall, - "four_total_tokens": sum(len(o["tokens"]) for o in outs.values()), - } - row["four_wall_tok_per_s"] = row["four_total_tokens"] / wall - rows.append(row) - print(json.dumps(row), flush=True) - with open(a.out, "w") as f: - json.dump(rows, f, indent=2) - - print("\n=== summary ===", flush=True) - for mode in modes: - rs = [r for r in rows if r["mode"] == mode] - for k in ("solo_tok_per_s", "four_aggregate_tok_per_s", "four_wall_tok_per_s", "solo_prompt_tok_per_s"): - vals = sorted(r[k] for r in rs) - print(f"mode={mode} {k}: median {vals[len(vals)//2]:.1f} values {[round(v,1) for v in vals]}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py deleted file mode 100644 index 32e79be2fc1..00000000000 --- a/scripts/batchinv/divergence.py +++ /dev/null @@ -1,215 +0,0 @@ -#!/usr/bin/env python3 -"""Baseline / patched divergence harness: solo P0 vs P0 sharing batches with P1..P3.""" -import argparse, json, os, signal, subprocess, sys, threading, time, urllib.request, urllib.error - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from prompts import PROMPTS - -# recorded with every run: LLAMA_EXACT_CONCURRENCY inherited from the shell decides whether a reference run really was one -RECORDED_ENV = ("LLAMA_EXACT_CONCURRENCY", "GGML_CUDA_BATCH_INVARIANT", - "GGML_CUDA_BATCH_INVARIANT_MAX_COLS", "LLAMA_SERVER_PREEMPT_EVERY", - "LLAMA_KV_CACHE_DEBUG", "LLAMA_BATCH_DEBUG", "CUDA_VISIBLE_DEVICES") - -MODEL_REL = "models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf" - - -def model_path(): - """Resolved when the server args are built, so --help works without the variable set.""" - ws = os.environ.get("UNSLOTH_WORKSPACE") - if not ws: - raise RuntimeError("UNSLOTH_WORKSPACE is not set; it must point at the workspace holding " - + MODEL_REL) - return os.path.join(ws, MODEL_REL) - - -def post(port, path, payload, timeout=1800): - req = urllib.request.Request(f"http://127.0.0.1:{port}{path}", - data=json.dumps(payload).encode(), - headers={"Content-Type": "application/json"}) - with urllib.request.urlopen(req, timeout=timeout) as r: - return json.loads(r.read().decode()) - - -def get(port, path, timeout=10): - with urllib.request.urlopen(f"http://127.0.0.1:{port}{path}", timeout=timeout) as r: - return json.loads(r.read().decode()) - - -def completion(port, prompt, n_predict): - return post(port, "/completion", { - "prompt": prompt, "n_predict": n_predict, "temperature": 0.0, "top_k": 1, - "top_p": 1.0, "min_p": 0.0, "typical_p": 1.0, "seed": 0, - "repeat_penalty": 1.0, "presence_penalty": 0.0, "frequency_penalty": 0.0, - "cache_prompt": False, "return_tokens": True, "samplers": ["top_k", "temperature"], - }) - - -class Server: - def __init__(self, port, binary, extra, env_extra, log_path, spec, kv_unified=True): - self.port, self.log_path = port, log_path - self.args = [binary, "-m", model_path(), "--port", str(port), "--host", "127.0.0.1", - "--parallel", "4", "-c", "8192", - "--flash-attn", "on", "--metrics", "-ngl", "99", "--no-warmup", - "--seed", "0", "--spec-type", spec] - if kv_unified: - self.args += ["--kv-unified"] - if spec == "draft-mtp": - self.args += ["--spec-draft-n-max", "2"] - self.args += extra - self.env = dict(os.environ) - self.env["CUDA_VISIBLE_DEVICES"] = "3" - self.env.update(env_extra) - self.env_resolved = {k: self.env[k] for k in RECORDED_ENV if k in self.env} - self.p = None - self.fh = None - - def __enter__(self): - self.fh = open(self.log_path, "ab") - self.fh.write(("\n=== " + " ".join(self.args) + "\n=== env " + - json.dumps(self.env_resolved) + "\n").encode()) - self.fh.flush() - self.p = subprocess.Popen(self.args, stdout=self.fh, stderr=subprocess.STDOUT, - env=self.env, start_new_session=True) - print(f"[server] pid={self.p.pid} port={self.port} log={self.log_path}", flush=True) - try: - deadline = time.time() + 600 - while time.time() < deadline: - if self.p.poll() is not None: - raise RuntimeError(f"server died rc={self.p.returncode}, see {self.log_path}") - try: - if get(self.port, "/health").get("status") == "ok": - print("[server] ready", flush=True) - return self - except Exception: - time.sleep(1.0) - raise RuntimeError("server did not become healthy") - except BaseException: - self.__exit__(None, None, None) - raise - - def __exit__(self, *a): - # note: POSIX only; Windows would need CREATE_NEW_PROCESS_GROUP at Popen - if self.p is not None: - print(f"[server] stopping pid={self.p.pid}", flush=True) - try: - os.killpg(os.getpgid(self.p.pid), signal.SIGTERM) - self.p.wait(timeout=60) - except Exception: - try: - os.killpg(os.getpgid(self.p.pid), signal.SIGKILL) - except Exception: - pass - try: - self.p.wait(timeout=60) - except Exception: - pass - self.p = None - if self.fh is not None: - self.fh.close() - self.fh = None - - -def run_concurrent(port, names, n_predict): - barrier = threading.Barrier(len(names)) - lock = threading.Lock() - out = {} - errors = [] - - def work(name): - try: - barrier.wait() - res = completion(port, PROMPTS[name], n_predict) - except BaseException as e: - with lock: - errors.append((name, e)) - # release the others rather than let them block on a barrier that will never fill - barrier.abort() - return - with lock: - out[name] = res - - ts = [threading.Thread(target=work, args=(n,)) for n in names] - t0 = time.time() - for t in ts: - t.start() - for t in ts: - t.join() - wall = time.time() - t0 - - # without this a run where P1..P3 failed and P0 succeeded reads as a clean four-way result - if errors: - raise RuntimeError("concurrent requests failed: " + - "; ".join(f"{n}: {type(e).__name__}: {e}" for n, e in errors)) - missing = set(names) - set(out) - if missing: - raise RuntimeError(f"concurrent requests produced no result for {sorted(missing)}") - - return out, wall - - -def first_diff(a, b): - for i, (x, y) in enumerate(zip(a, b)): - if x != y: - return i - return None if len(a) == len(b) else min(len(a), len(b)) - - -def main(): - ap = argparse.ArgumentParser() - ap.add_argument("--reference") - ap.add_argument("--label", required=True) - ap.add_argument("--port", type=int, default=9601) - ap.add_argument("--binary", required=True) - ap.add_argument("--spec", default="none") - ap.add_argument("--n-predict", type=int, default=512) - ap.add_argument("--repeats", type=int, default=3) - ap.add_argument("--env", action="append", default=[]) - ap.add_argument("--extra", action="append", default=[]) - ap.add_argument("--out", required=True) - ap.add_argument("--no-kv-unified", action="store_true") - a = ap.parse_args() - - env_extra = dict(kv.split("=", 1) for kv in a.env) - server = Server(a.port, a.binary, a.extra, env_extra, a.out + ".server.log", a.spec, - kv_unified=not a.no_kv_unified) - res = {"label": a.label, "spec": a.spec, "n_predict": a.n_predict, - "env_requested": env_extra, "env": server.env_resolved, - "model": server.args[2], "args": server.args, - "extra": a.extra, "binary": a.binary, - "kv_unified": not a.no_kv_unified} - - with server as s: - solo = completion(a.port, PROMPTS["P0"], a.n_predict) - ref = json.load(open(a.reference))["tokens"] if a.reference else solo["tokens"] - res["solo_first_diff"] = first_diff(ref, solo["tokens"]) - res["reference"] = a.reference - res["solo"] = {"n_tokens": len(ref), "tok_per_s": solo["timings"]["predicted_per_second"], - "text_sha": None} - solo2 = completion(a.port, PROMPTS["P0"], a.n_predict) - res["solo_repeat_first_diff"] = first_diff(ref, solo2["tokens"]) - res["rounds"] = [] - for r in range(a.repeats): - outs, wall = run_concurrent(a.port, ["P0", "P1", "P2", "P3"], a.n_predict) - p0 = outs["P0"]["tokens"] - fd = first_diff(ref, p0) - agg = sum(outs[n]["timings"]["predicted_per_second"] for n in outs) - row = {"round": r, "first_diff": fd, "n_tokens": len(p0), - "identical": fd is None, "wall_s": wall, - "p0_tok_per_s": outs["P0"]["timings"]["predicted_per_second"], - "aggregate_tok_per_s": agg, - "per_req_n": {n: len(outs[n]["tokens"]) for n in outs}, "p0_tokens": p0} - res["rounds"].append(row) - print(f"[round {r}] first_diff={fd} identical={fd is None} wall={wall:.1f}s agg={agg:.1f} tok/s", flush=True) - with urllib.request.urlopen(f"http://127.0.0.1:{a.port}/metrics") as response: - res["metrics"] = response.read().decode() - with open(a.out + ".p0_solo.json", "w") as f: - json.dump({"tokens": ref, "content": solo["content"]}, f) - - with open(a.out, "w") as f: - json.dump(res, f, indent=2) - print(json.dumps({k: v for k, v in res.items() if k != "rounds"}, indent=2), flush=True) - print(json.dumps(res["rounds"], indent=2), flush=True) - - -if __name__ == "__main__": - main() diff --git a/scripts/batchinv/probe.cpp b/scripts/batchinv/probe.cpp deleted file mode 100644 index 4e72d5b8a55..00000000000 --- a/scripts/batchinv/probe.cpp +++ /dev/null @@ -1,359 +0,0 @@ -// Locate the first graph op whose sequence-0 output changes when the decode batch holds four sequences instead of one; seq 0's prompt KV is identical in both phases. -#include "llama.h" -#include "ggml.h" -#include "ggml-backend.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -struct node_rec { - std::string name; - std::string op; - std::string tname; - int64_t ne[4]; - int64_t gdn_tokens = 0, gdn_seqs = 0; - size_t esize = 0; // bytes per element, 0 = not byte comparable - std::vector data; // empty when skipped - bool contiguous = false; - uint64_t hash = 0; -}; - -static bool g_record = false; -static std::vector * g_sink = nullptr; - -static uint64_t fnv1a(const uint8_t * p, size_t n) { - uint64_t h = 1469598103934665603ULL; - for (size_t i = 0; i < n; ++i) { h ^= p[i]; h *= 1099511628211ULL; } - return h; -} - -static bool eval_cb(struct ggml_tensor * t, bool ask, void * /*ud*/) { - if (!g_record) return false; - if (ask) return true; - - node_rec r; - r.name = ggml_get_name(t); - r.tname = ggml_type_name(t->type); - r.op = t->op == GGML_OP_NONE ? "LEAF" : ggml_op_name(t->op); - if (t->op == GGML_OP_UNARY) r.op = std::string("UNARY_") + ggml_unary_op_name(ggml_get_unary_op(t)); - if (t->op == GGML_OP_GLU) r.op = std::string("GLU_") + ggml_glu_op_name(ggml_get_glu_op(t)); - for (int i = 0; i < 4; ++i) r.ne[i] = t->ne[i]; - r.contiguous = ggml_is_contiguous(t); - if (r.name == "linear_attn_out-0") { - fprintf(stderr, "linear_attn_out-0: weight=%s input=[%lld,%lld,%lld,%lld]\n", - ggml_type_name(t->src[0]->type), (long long)t->src[1]->ne[0], - (long long)t->src[1]->ne[1], (long long)t->src[1]->ne[2], (long long)t->src[1]->ne[3]); - } - if (t->op == GGML_OP_GATED_DELTA_NET) { - r.gdn_tokens = t->src[2]->ne[2]; - r.gdn_seqs = t->src[2]->ne[3]; - } - - const size_t nbytes = ggml_nbytes(t); - if (r.contiguous && ggml_blck_size(t->type) == 1 && nbytes <= (256u << 20)) { - r.esize = ggml_type_size(t->type); - r.data.resize(nbytes); - ggml_backend_tensor_get(t, r.data.data(), 0, nbytes); - if (t->op == GGML_OP_FLASH_ATTN_EXT) { - for (size_t i = 0; i < nbytes/sizeof(float); ++i) { - float v; memcpy(&v, r.data.data() + i*sizeof(float), sizeof(float)); - if (!std::isfinite(v)) { fprintf(stderr, "nonfinite attention: %s\n", t->name); exit(5); } - } - } - r.hash = fnv1a(r.data.data(), nbytes); - if (nbytes > (64u << 20)) { r.data.clear(); } // keep the hash only for the big ones - } - g_sink->push_back(std::move(r)); - return true; -} - -static std::string slurp(const char * path) { - std::ifstream f(path); - std::stringstream ss; ss << f.rdbuf(); - std::string s = ss.str(); - while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back(); - return s; -} - -static std::vector tokenize(const llama_vocab * v, const std::string & s) { - std::vector out(s.size() + 16); - int n = llama_tokenize(v, s.c_str(), (int) s.size(), out.data(), (int) out.size(), true, false); - if (n < 0) { out.resize(-n); n = llama_tokenize(v, s.c_str(), (int) s.size(), out.data(), (int) out.size(), true, false); } - out.resize(n); - return out; -} - -struct batch_holder { - std::vector tok; - std::vector pos; - std::vector nsid; - std::vector sid; - std::vector sidp; - std::vector out; - llama_batch get() { - sidp.resize(tok.size()); - for (size_t i = 0; i < tok.size(); ++i) sidp[i] = &sid[i]; - llama_batch b{}; - b.n_tokens = (int32_t) tok.size(); - b.token = tok.data(); b.pos = pos.data(); b.n_seq_id = nsid.data(); - b.seq_id = sidp.data(); b.logits = out.data(); - return b; - } -}; - -static llama_token greedy(llama_context * ctx, int32_t i, int n_vocab) { - const float * l = llama_get_logits_ith(ctx, i); - for (int k = 0; k < n_vocab; ++k) { - if (!std::isfinite(l[k])) { fprintf(stderr, "nonfinite logits at %d\n", k); exit(4); } - } - int best = 0; - for (int k = 1; k < n_vocab; ++k) if (l[k] > l[best]) best = k; - return best; -} - -static llama_token feed(llama_context * ctx, const std::vector & p, llama_seq_id seq, int n_vocab) { - batch_holder h; - for (size_t i = 0; i < p.size(); ++i) { - h.tok.push_back(p[i]); h.pos.push_back((llama_pos) i); - h.nsid.push_back(1); h.sid.push_back(seq); - h.out.push_back(i + 1 == p.size()); - } - llama_batch b = h.get(); - if (llama_decode(ctx, b) != 0) { fprintf(stderr, "decode failed\n"); exit(1); } - return greedy(ctx, (int32_t) p.size() - 1, n_vocab); -} - -int main(int argc, char ** argv) { - const bool prefill = getenv("PROBE_PREFILL") != nullptr; - const char * model_path = argv[1]; - const int n_seqs = argc > 2 ? atoi(argv[2]) : 4; // width of the probed decode batch - const char * out_path = argc > 3 ? argv[3] : nullptr; - std::vector prompts; - for (int i = 4; i < argc; ++i) prompts.push_back(slurp(argv[i])); - - llama_backend_init(); - llama_model_params mp = llama_model_default_params(); - mp.n_gpu_layers = 99; - llama_model * model = llama_model_load_from_file(model_path, mp); - if (!model) { fprintf(stderr, "model load failed\n"); return 1; } - const llama_vocab * vocab = llama_model_get_vocab(model); - const int n_vocab = llama_vocab_n_tokens(vocab); - - std::vector> ptok; - for (auto & s : prompts) ptok.push_back(tokenize(vocab, s)); - for (size_t i = 0; i < ptok.size(); ++i) fprintf(stderr, "prompt %zu: %zu tokens\n", i, ptok[i].size()); - - auto make_ctx = [&]() { - llama_context_params cp = llama_context_default_params(); - cp.n_ctx = 8192; cp.n_batch = 2048; cp.n_ubatch = 512; - if (prefill) { cp.n_ubatch = 2048; } - cp.n_seq_max = 4; cp.kv_unified = true; - cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; - cp.cb_eval = eval_cb; cp.cb_eval_user_data = nullptr; - cp.no_perf = true; - return llama_init_from_model(model, cp); - }; - - std::vector rec_a, rec_b; - llama_token first_tok[4] = {0, 0, 0, 0}; - - // phase A: decode ubatch width 1. PROBE_A_FILL is how many sequences are already in the shared KV cache, which sets K->ne[1]. - const int a_fill = getenv("PROBE_A_FILL") ? atoi(getenv("PROBE_A_FILL")) : 1; - // PROBE_A_PERM reorders which prompt goes into which sequence in phase A; it may only move the neighbours, since sequence 0 must keep prompt 0 - int a_perm[4] = {0, 1, 2, 3}; - if (const char * perm = getenv("PROBE_A_PERM")) { - for (int k = 0; k < 4 && perm[2*k]; ++k) a_perm[k] = perm[2*k] - '0'; - if (a_perm[0] != 0) { fprintf(stderr, "PROBE_A_PERM must keep prompt 0 on sequence 0\n"); return 1; } - } - { - llama_context * ctx = make_ctx(); - g_sink = &rec_a; g_record = prefill; - for (int s = 0; s < a_fill; ++s) { - const llama_token t = feed(ctx, ptok[a_perm[s]], s, n_vocab); - if (a_perm[s] == 0) first_tok[0] = t; - } - batch_holder h; - h.tok = {first_tok[0]}; h.pos = {(llama_pos) ptok[0].size()}; - h.nsid = {1}; h.sid = {0}; h.out = {1}; - llama_batch b = h.get(); - g_sink = &rec_a; g_record = !prefill; - if (llama_decode(ctx, b) != 0) { fprintf(stderr, "A decode failed\n"); return 1; } - g_record = false; - llama_free(ctx); - } - - { - llama_context * ctx = make_ctx(); - if (prefill) { - batch_holder h; - for (int seq = 0; seq < n_seqs; ++seq) { - for (size_t i = 0; i < ptok[0].size(); ++i) { - h.tok.push_back(ptok[seq][i%ptok[seq].size()]); h.pos.push_back(i); - h.nsid.push_back(1); h.sid.push_back(seq); h.out.push_back(i+1 == ptok[0].size()); - } - } - auto b = h.get(); - g_sink = &rec_b; g_record = true; - if (llama_decode(ctx, b) != 0) { return 6; } - g_record = false; - for (int seq = 0; seq < n_seqs; ++seq) { - first_tok[seq] = greedy(ctx, (seq+1)*ptok[0].size()-1, n_vocab); - } - } else for (int k = 0; k < n_seqs; ++k) { - const int s = getenv("PROBE_B_REVERSE") ? n_seqs - 1 - k : k; - first_tok[s] = feed(ctx, ptok[s], s, n_vocab); - } - if (getenv("PROBE_RESTORE")) { - std::vector state(llama_state_seq_get_size(ctx, 0)); - if (llama_state_seq_get_data(ctx, state.data(), state.size(), 0) != state.size()) { return 2; } - llama_memory_seq_rm(llama_get_memory(ctx), 0, -1, -1); - llama_memory_seq_rm(llama_get_memory(ctx), 1, -1, -1); - if (llama_state_seq_set_data(ctx, state.data(), state.size(), 0) != state.size()) { return 3; } - first_tok[1] = feed(ctx, ptok[1], 1, n_vocab); - } - if (first_tok[0] != 0 && rec_a.size()) {} - batch_holder h; - for (int s = 0; s < n_seqs; ++s) { - h.tok.push_back(first_tok[s]); h.pos.push_back((llama_pos) ptok[s].size()); - h.nsid.push_back(1); h.sid.push_back(s); h.out.push_back(1); - } - llama_batch b = h.get(); - g_sink = &rec_b; g_record = !prefill; - if (!prefill && llama_decode(ctx, b) != 0) { fprintf(stderr, "B decode failed\n"); return 1; } - g_record = false; - llama_free(ctx); - } - - const int n_steps = getenv("PROBE_STEPS") ? atoi(getenv("PROBE_STEPS")) : 0; - int first_bad_step = -1; - if (n_steps > 0) { - std::vector tok_a, tok_b; - for (int phase = 0; phase < 2; ++phase) { - const int fill = phase == 0 ? a_fill : n_seqs; - const int width = phase == 0 ? 1 : n_seqs; - std::vector & out = phase == 0 ? tok_a : tok_b; - llama_context * ctx = make_ctx(); - std::vector next(4, 0); - std::vector pos(4, 0); - for (int s = 0; s < fill; ++s) { - const int p = phase == 0 ? a_perm[s] : s; - next[s] = feed(ctx, ptok[p], s, n_vocab); - pos[s] = (llama_pos) ptok[p].size(); - } - for (int step = 0; step < n_steps; ++step) { - batch_holder h; - for (int s = 0; s < width; ++s) { - h.tok.push_back(next[s]); h.pos.push_back(pos[s]); - h.nsid.push_back(1); h.sid.push_back(s); h.out.push_back(1); - } - llama_batch b = h.get(); - if (llama_decode(ctx, b) != 0) { fprintf(stderr, "step decode failed\n"); exit(1); } - out.push_back(next[0]); - for (int s = 0; s < width; ++s) { next[s] = greedy(ctx, s, n_vocab); pos[s] += 1; } - } - llama_free(ctx); - } - for (int i = 0; i < n_steps; ++i) { - if (tok_a[i] != tok_b[i]) { first_bad_step = i; break; } - } - fprintf(stderr, "steps: %d first differing step: %d\n", n_steps, first_bad_step); - } - - fprintf(stderr, "nodes: A=%zu B=%zu first tokens: %d %d %d %d\n", - rec_a.size(), rec_b.size(), first_tok[0], first_tok[1], first_tok[2], first_tok[3]); - - FILE * out = out_path ? fopen(out_path, "w") : stdout; - fprintf(out, "{\"n_seqs\":%d,\"first_bad_step\":%d,\"nodes_a\":%zu,\"nodes_b\":%zu,\"diffs\":[", n_seqs, first_bad_step, rec_a.size(), rec_b.size()); - size_t n = rec_a.size() < rec_b.size() ? rec_a.size() : rec_b.size(); - int emitted = 0; - for (size_t i = 0; i < n; ++i) { - const node_rec & A = rec_a[i]; - const node_rec & B = rec_b[i]; - const char * verdict = nullptr; - double max_abs = 0.0; - size_t ndiff = 0, ncmp = 0; - - if (A.name != B.name || A.op != B.op) { - verdict = "misaligned"; - } else if (A.op == "GATED_DELTA_NET" && A.gdn_tokens == B.gdn_tokens && - !A.data.empty() && !B.data.empty()) { - // packed GDN outputs put all token outputs before all sequence states, so seq 0's state moves when the number of sequences changes - const size_t output = A.ne[0]*A.gdn_tokens; - const size_t state = A.ne[0]*A.ne[1]/A.gdn_seqs - output; - for (size_t k = 0; k < output + state; ++k) { - const size_t ia = k < output ? k : A.gdn_seqs*output + k-output; - const size_t ib = k < output ? k : B.gdn_seqs*output + k-output; - float va, vb; - memcpy(&va, A.data.data()+ia*4, 4); memcpy(&vb, B.data.data()+ib*4, 4); - ++ncmp; - if (memcmp(&va, &vb, 4)) { - ++ndiff; - if (std::abs(double(va)-vb) > max_abs) { max_abs = std::abs(double(va)-vb); } - } - } - verdict = ndiff ? "row-differs" : nullptr; - } else if (A.esize == 0 || B.esize == 0 || A.esize != B.esize) { - verdict = "skipped"; - } else { - int tdim = -1; bool same = true; - for (int d = 0; d < 4; ++d) { - if (A.ne[d] == B.ne[d]) continue; - same = false; - if (B.ne[d] == n_seqs*A.ne[d] && tdim < 0) tdim = d; else { tdim = -2; break; } - } - if (tdim == -2) { - verdict = "shape-incomparable"; - } else if (same) { - verdict = (A.hash == B.hash) ? nullptr : "whole-tensor-differs"; - } else if (A.data.empty() || B.data.empty()) { - verdict = "too-large"; - } else { - int64_t st[4] = {1, A.ne[0], A.ne[0]*A.ne[1], A.ne[0]*A.ne[1]*A.ne[2]}; - int64_t stb[4] = {1, B.ne[0], B.ne[0]*B.ne[1], B.ne[0]*B.ne[1]*B.ne[2]}; - for (int64_t i3 = 0; i3 < A.ne[3]; ++i3) - for (int64_t i2 = 0; i2 < A.ne[2]; ++i2) - for (int64_t i1 = 0; i1 < A.ne[1]; ++i1) - for (int64_t i0 = 0; i0 < A.ne[0]; ++i0) { - int64_t idx[4] = {i0, i1, i2, i3}; - - size_t oa = 0, ob = 0; - for (int d = 0; d < 4; ++d) { oa += idx[d]*st[d]; ob += idx[d]*stb[d]; } - ncmp++; - const uint8_t * pa = A.data.data() + oa*A.esize; - const uint8_t * pb = B.data.data() + ob*B.esize; - if (memcmp(pa, pb, A.esize) != 0) { - ndiff++; - if (A.esize == 4) { - float fa, fb; memcpy(&fa, pa, 4); memcpy(&fb, pb, 4); - double d2 = fa - fb; if (d2 < 0) d2 = -d2; - if (d2 > max_abs) max_abs = d2; - } - } - } - verdict = ndiff ? "row-differs" : nullptr; - } - } - { - if (emitted++) fprintf(out, ","); - fprintf(out, "\n{\"i\":%zu,\"name\":\"%s\",\"op\":\"%s\",\"ne_a\":[%lld,%lld,%lld,%lld]," - "\"ne_b\":[%lld,%lld,%lld,%lld],\"type\":\"%s\",\"verdict\":\"%s\",\"ndiff\":%zu,\"ncmp\":%zu,\"max_abs\":%.6g}", - i, A.name.c_str(), A.op.c_str(), - (long long)A.ne[0],(long long)A.ne[1],(long long)A.ne[2],(long long)A.ne[3], - (long long)B.ne[0],(long long)B.ne[1],(long long)B.ne[2],(long long)B.ne[3], - A.tname.c_str(), verdict ? verdict : "same", ndiff, ncmp, max_abs); - } - } - fprintf(out, "\n]}\n"); - if (out_path) fclose(out); - - llama_model_free(model); - llama_backend_free(); - return 0; -} diff --git a/scripts/batchinv/prompts.py b/scripts/batchinv/prompts.py deleted file mode 100644 index efabc448a9d..00000000000 --- a/scripts/batchinv/prompts.py +++ /dev/null @@ -1,53 +0,0 @@ -# four distinct prompts, each about 300 tokens of raw text (no chat template) -_BODIES = { -"P0": """The history of numerical computing is a history of compromises between speed and exactness. -Early machines used fixed point arithmetic because it was cheap, and programmers carried scaling -factors in their heads. Floating point hardware moved the bookkeeping into silicon, but it did not -remove the compromise, it only hid it. Addition of floating point numbers is commutative but it is -not associative, so the order in which a long sum is accumulated changes the last few bits of the -result. On a single processor that order is fixed by the program text and nobody notices. On a -parallel processor the order is fixed by how the work was divided, and the division is chosen for -speed, not for reproducibility. A reduction split across two warps sums a different set of partial -products than the same reduction split across four warps, and the two answers differ in the low -bits. Nothing is wrong with either answer. Both are within a fraction of an ulp of the exact value. -The trouble begins when a downstream decision is discrete. A comparison, a rounding to an integer, -or the selection of the largest element of a vector turns a difference of one bit into a difference -of one branch, and from there the two computations walk away from each other and never come back. -Explain, carefully and at length, why this matters for a system that serves many users at once, -and what an engineer would have to give up to make the answer depend only on the request and not -on what else the machine happened to be doing at the time. Discuss the cost.""", -"P1": """Consider a public library that lends physical books and must decide how many copies of a -popular title to buy. The librarian has a fixed budget, a waiting list that grows and shrinks, and -a shelf that is already full. Every copy purchased shortens the queue for that title and lengthens -the queue for every other title, because the money and the shelf space are shared. The obvious -policy, buy copies of whatever has the longest queue, is unstable, because a title that briefly -becomes fashionable will absorb the whole budget and then sit unread for a decade. A better policy -has to weigh how long the demand is likely to last against how long the book will remain useful, -and it has to do this with almost no information. Describe in detail how you would design such a -policy, what data you would collect, how you would test it without harming readers, and how you -would know whether it was working. Consider what happens when the budget is cut in half without -warning, when a title is suddenly assigned as required reading by a local school, and when the -shelf itself must shrink because the building is being renovated. Explain the tradeoffs plainly.""", -"P2": """A small coastal town has one bridge to the mainland and it is failing. The engineers say it -has perhaps eight years left. Replacing it costs more than the town has ever spent on anything. -Repairing it buys maybe four years and costs a third as much, and the repair work closes the bridge -for two months in the summer, which is when the town earns most of its money. Doing nothing is -free until the day it is not. The town council is split, the ferry operator has opinions, and the -regional government will match funds only for a replacement, only if construction begins within -three years, and only if the town covers the first quarter of the cost itself. Write a long and -careful analysis of the options available to the council. Identify the assumptions that matter -most, the ones where being wrong changes the recommendation, and say how the council could cheaply -find out whether those assumptions hold. Then give a recommendation and state honestly what would -have to be true for the recommendation to be wrong. Do not hedge. Commit to an answer at the end.""", -"P3": """Describe the process by which a large body of water freezes over in winter, beginning with -the surface layer and working downward, and explain why the ice floats rather than sinking, why a -deep lake takes much longer to freeze than a shallow one of the same surface area, and why the -temperature at the bottom of a frozen lake settles near four degrees Celsius rather than at zero. -Then explain what this means for the animals that live there, how fish survive a winter under a -solid lid, why a heavy snowfall on top of the ice can be more dangerous to them than the cold -itself, and what happens in the spring when the whole column overturns. Use plain language and -avoid equations. Where a common explanation is wrong or incomplete, say so and give the better one. -Be thorough. Assume the reader is curious and patient but has no training in physics or biology, -and would rather understand one thing properly than be told five things quickly.""", -} -PROMPTS = {k: " ".join(v.split()) for k, v in _BODIES.items()} diff --git a/scripts/unsloth/additive_merge.py b/scripts/unsloth/additive_merge.py index 13e4784b775..5364dc1b7c1 100644 --- a/scripts/unsloth/additive_merge.py +++ b/scripts/unsloth/additive_merge.py @@ -94,8 +94,18 @@ def nonblank(lines: list[str]) -> list[str]: return [ln.strip() for ln in lines if ln.strip()] -# a line that only opens or closes a block. Two independent case arms share these by construction, so -# treating them as shared refused the real PROJECTOR_TYPE_KIMIK3 / _DEEPSEEK4V add/add in tools/mtmd/clip.cpp. +# A line that closes or opens a block and nothing else. Two INDEPENDENT case +# arms in the same switch share these by construction -- `{`, `} break;`, `}` +# are what a case arm is made of, not what makes it that case arm -- so finding +# them on both sides says nothing about whether the two sides added the same +# construct. Matching them as "shared" is what refused the real add/add of +# PROJECTOR_TYPE_KIMIK3 next to PROJECTOR_TYPE_DEEPSEEK4V in tools/mtmd/clip.cpp +# with "one change made twice: {, } break;", when the two arms had no line of +# actual content in common. +# +# Deliberately narrow: braces, brackets, parens, semicolons and commas, around +# at most one bare block-terminating keyword. `break;` matches, `return true;` +# does not, and anything naming a type, a constant or a function does not. STRUCTURAL = re.compile(r"^[\s{}()\[\];,]*(?:break|continue|return|pass)?[\s{}()\[\];,]*$") @@ -104,6 +114,8 @@ def identifying(lines: list[str]) -> set[str]: return {ln for ln in nonblank(lines) if not STRUCTURAL.match(ln)} +# `case FOO:`, `case FOO :`, `default:`. A fallthrough label may carry no body +# at all, which is the shape the nightly hits most often. CASE_LABEL = re.compile(r"^(?:case\s+[^:]+|default\s*):") @@ -135,17 +147,32 @@ def resolve_region(ours: list[str], base: list[str], theirs: list[str]) -> list[ return list(ours) ours_arms, theirs_arms = case_arms(ours), case_arms(theirs) if ours_arms and theirs_arms and ours_arms.isdisjoint(theirs_arms): - # both sides added case arms and no label is on both, so any line they share is body text; the same change made twice would keep its label, so sharing is allowed only here + # Both sides added case arms, and not one label is on both sides. Two + # arms of the same switch labelled differently are two constructs, so + # any line they happen to share is body text, not a duplicate: the real + # tools/mtmd/clip.cpp collision has a KIMIK3 arm and a DEEPSEEK4V arm + # that both set `hparams.rope_theta = 10000.0f;`, and refusing on that + # coincidence is what the shared-line check is for, backwards. + # + # The same change made twice would keep its label, so it lands in the + # check below instead. This is the one place where a shared line is + # allowed, and it is allowed because the labels prove the arms are + # distinct -- a duplicated label would not even compile. return list(theirs) + list(ours) shared = identifying(ours) & identifying(theirs) if shared: # Overlapping content is the signature of one construct added twice, # not two independent additions. Unioning it would duplicate code. + # Scaffolding lines are excluded above, so what is left is content both + # sides genuinely wrote, which is the thing that makes this a duplicate. raise Unresolvable( "both sides add the same line(s), so this is one change made twice: " + ", ".join(sorted(shared)[:3]) ) if not identifying(ours) or not identifying(theirs): + # Everything one side added is scaffolding, so there is no content to + # tell the two additions apart and the exclusion above has nothing left + # to work with. Refuse rather than union braces onto braces. raise Unresolvable( "one side adds only block scaffolding, so the two additions cannot " "be told apart" diff --git a/scripts/unsloth/feature_matrix.py b/scripts/unsloth/feature_matrix.py index 6c0d7f357cd..00b392ca1ec 100644 --- a/scripts/unsloth/feature_matrix.py +++ b/scripts/unsloth/feature_matrix.py @@ -37,6 +37,7 @@ import sys from pathlib import Path +# Output that means "this did not run" from a process that exited 0. SKIP_RE = re.compile(r"\bSKIP\b|not supported|unsupported|no tests|0 tests", re.I) @@ -77,6 +78,7 @@ def probe_arch(check: dict, b: Path, gpu: bool) -> str: rc, out = run([str(b / "test-llama-archs"), "-a", arch, "-s", "1234"], b, gpu) if rc != 0: raise Unproven(f"test-llama-archs -a {arch} exited {rc}") + # The arch's own rows, not the header and not another arch's. rows = [ln for ln in out.splitlines() if ln.strip().startswith("|") and f"|{arch:>16}|" in ln or (ln.strip().startswith("|") and ln.split("|")[1].strip() == arch)] if not rows: @@ -117,7 +119,8 @@ def probe_mtmd(check: dict, b: Path, gpu: bool) -> str: m = re.search(r"assertions\s*:\s*(\d+)", out) if not m or int(m.group(1)) == 0: raise Unproven("test_projector_registry ran no assertions; the filter matched nothing") - # the test walks the whole enum, so it proves the table is sound; that this projector is IN the enum is pin_contract.py's job + # The registry test walks the whole enum, so it proves the table is sound. + # That the specific projector is IN the enum is pin_contract.py's job. return f"projector registry intact over {m.group(1)} assertions" @@ -170,6 +173,8 @@ def main() -> int: print(f"ok {name}: " + "; ".join(r["evidence"] for r in entry["results"]) + (f" [{len(entry['deferred'])} needs a GPU]" if entry["deferred"] else "")) else: + # Nothing was shown either way. Not a failure here, but it must not + # read as one of the ok lines. print(f"-- {name}: nothing provable without a GPU " f"({len(entry['deferred'])} check(s) deferred)") @@ -183,6 +188,8 @@ def main() -> int: if failed: print(f"\n{failed} feature(s) could not be shown to work", file=sys.stderr) return 1 + # Say what was NOT proven in the same breath as what was. A run that only + # ever prints a success line teaches the reader that green means covered. tail = f", {deferred} check(s) need a GPU and were not run" if deferred else "" print(f"\nall {len(report['features'])} features demonstrated" + (" on GPU" if args.gpu else " on CPU") + tail) diff --git a/scripts/unsloth/pin_contract.py b/scripts/unsloth/pin_contract.py index 9eebe824486..bf07d58408e 100644 --- a/scripts/unsloth/pin_contract.py +++ b/scripts/unsloth/pin_contract.py @@ -57,22 +57,34 @@ r"^https://github\.com/([^/]+)/llama\.cpp/pull/(\d+)/commits/([0-9a-f]{40})/?$" ) -# identifier families that name a FEATURE, not every new symbol: a renamed helper is not a lost feature, a missing LLM_ARCH_ entry always is +# Identifier families that name a FEATURE. Deliberately not "every new symbol": +# a helper function renamed by a later upstream commit is not a lost feature, +# but a missing LLM_ARCH_ entry always is. These are the tables that decide +# whether an architecture, an op, a projector or a quant type exists at all. SYMBOL_FAMILIES = ( "LLM_ARCH_", "LLM_TENSOR_", "LLM_KV_", "LLM_TYPE_", "PROJECTOR_TYPE_", "GGML_OP_", "GGML_TYPE_", "LLAMA_FTYPE_", ) SYMBOL_RE = re.compile(r"\b(?:" + "|".join(SYMBOL_FAMILIES) + r")[A-Z0-9_]+\b") +# The subset that names a whole feature rather than one of its tensors. Used +# only to keep --emit readable; the check itself uses all of SYMBOL_FAMILIES. HEADLINE = ("LLM_ARCH_", "GGML_OP_", "GGML_TYPE_", "PROJECTOR_TYPE_", "LLAMA_FTYPE_") +# A line worth tracking for survival. Comments and short punctuation drift with +# every reformat and would make the check noise; a substantial code line does +# not move on its own. TRIVIAL_RE = re.compile(r"^\s*(?://|/\*|\*|\*/|#\s|$)") MIN_LINE = 12 -# comments are stripped first: a pin that merely NAMES an arch in a comment has not registered it, and -# holding the wording as a contract fails when upstream rewords it (unslothai#70) +# Comments are stripped before anything is read off a line. A pin that merely +# NAMES an arch in a comment has not registered it, and holding the comment's +# wording as a contract fails the moment upstream rewords it. Observed on +# unslothai#70, whose comment mentions GGML_OP_SSM_SCAN to explain why it does +# NOT use it. COMMENT_RE = re.compile(r"//.*$|/\*.*?\*/|(? dict: if code: symbols[cur].update(SYMBOL_RE.findall(code)) + # Only symbols the base does not ALREADY have in that file are evidence of + # this pin. Upstream naming an arch in a file the pin also touches is not + # something the pin is owed. new_symbols: dict[str, list[str]] = {} for path, names in symbols.items(): fresh = sorted(n for n in names @@ -291,6 +306,9 @@ def main() -> int: if args.emit: report["pins"].append(entry) + # Only the families that NAME a feature are printed. Every symbol + # is still checked; a new file legitimately contributes a hundred + # LLM_TENSOR_ names and listing them buries the one that matters. sym = sorted({s for v in contract["symbols"].values() for s in v if s.startswith(HEADLINE)}) print(f"{name:>18} {entry['line_count']:>5} lines, " @@ -327,6 +345,8 @@ def main() -> int: if args.emit: return 0 + # Notices after the verdict lines, never mixed into them: "upstream took + # this, drop the entry" is housekeeping and must not read as a failure. for n in notices: print(f"note {n}") if failed: diff --git a/scripts/unsloth/test_additive_merge.py b/scripts/unsloth/test_additive_merge.py index 70da225256d..0f91e9afd12 100644 --- a/scripts/unsloth/test_additive_merge.py +++ b/scripts/unsloth/test_additive_merge.py @@ -97,7 +97,8 @@ def run(repo, *extra): reason.endswith('twice: log("same");'), reason) # --- 3b. two independent case arms: braces are shared, content is not ------- -# the real clip.cpp shape; refusing it on `{` and `} break;` took the 09-02 nightly down +# The real tools/mtmd/clip.cpp shape. Refusing this on `{` and `} break;` is +# what took the 09-02 nightly's last pin down. base = "switch (t) {\n}\n" ours = ("switch (t) {\n case PROJECTOR_TYPE_KIMIK3:\n {\n" " builder = std::make_unique(ctx, img);\n" @@ -114,6 +115,8 @@ def run(repo, *extra): txt.count("} break;") == 2 and txt.count("clip_graph_kimik3") == 1, txt) # --- 3b2. two case arms that share a body line, which is a coincidence ------ +# The clip.cpp shape after upstream landed DEEPSEEK4V: both arms set the same +# rope_theta, and refusing on that is the shared-line check backwards. base = "switch (t) {\n}\n" ours = ("switch (t) {\n case PROJECTOR_TYPE_KIMIK3:\n {\n" " hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;\n" diff --git a/scripts/unsloth/test_pin_contract.py b/scripts/unsloth/test_pin_contract.py index 628ab85e49c..3c7b58d62d5 100644 --- a/scripts/unsloth/test_pin_contract.py +++ b/scripts/unsloth/test_pin_contract.py @@ -98,6 +98,8 @@ def run(repo, pr_set, *extra): check("intact merge reports no notices", rep["notices"] == [], rep) # --- 2. the arm is dropped from ONE file: a tree-wide grep would pass ------ +# The real shape: LLM_ARCH_INKLING survives in the enum and the dispatch arm +# that makes it do anything is gone. repo, pr_set, sha = make_repo() p = repo / "src" / "llama-model.cpp" p.write_text(MODEL_CPP_BASE) @@ -126,6 +128,8 @@ def run(repo, pr_set, *extra): any("do_the_banded_thing" in x for x in rep["pins"][0]["problems"]), rep) # --- 5. redundancy: the base already has everything the pin adds ---------- +# Built the way it happens for real: upstream lands the same work, so the base +# tag has it and the pin is not an ancestor of anything. d = Path(tempfile.mkdtemp(prefix="pc_")) git(d, "init", "-q", "-b", "main") (d / "src").mkdir() @@ -159,7 +163,8 @@ def run(repo, pr_set, *extra): rep["pins"][0]["added_files"] == ["src/inkling.cpp"], rep) # --- 7. a comment is not a contract --------------------------------------- -# unslothai#70 names GGML_OP_SSM_SCAN in a comment to say it does NOT use it, and holding that wording would fail on a reword +# unslothai#70 has a comment naming GGML_OP_SSM_SCAN to say it does NOT use it. +# Holding comment wording would fail the moment upstream rewords it. repo, pr_set, sha = make_repo() git(repo, "checkout", "-q", "pin") (repo / "src" / "note.cpp").write_text( From 57efd21c65a1e2d362b80c9c4ed44315fc73c855 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 08:36:44 +0000 Subject: [PATCH 75/81] server : share the park and resume bookkeeping between the copy paths One preempt_park() for the four sites that park a slot, one preempt_parked() and preempt_restored() for the three that finish a copy, one preempt_copy_done() and preempt_resumed() inside the slot. rewind_to_cache() drops the batch through preempt_detach(), which already did exactly that. The only text that moves is the waited-for park line, which now carries the parked MiB like the other one. --- tools/server/server-context.cpp | 214 +++++++++++++------------------- 1 file changed, 83 insertions(+), 131 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 8cff94e34a1..7b7e03a3818 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -464,12 +464,26 @@ struct server_slot { i_batch = -1; } - bool preempt_save_poll() { - if (!llama_state_seq_copy_done(preempt_cpy_tgt.get())) { - return false; + bool preempt_copy_done() { + return llama_state_seq_copy_done(preempt_cpy_tgt.get()) && + (!preempt_cpy_dft || llama_state_seq_copy_done(preempt_cpy_dft.get())); + } + + // back in the state it was parked from, with a speculative context to match: the draft went out with the cells + bool preempt_resumed() { + n_preempt_fail = 0; + + state = state_before_preempt; + + if (state == SLOT_STATE_GENERATING && can_speculate()) { + common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); } - if (preempt_cpy_dft && !llama_state_seq_copy_done(preempt_cpy_dft.get())) { + return true; + } + + bool preempt_save_poll() { + if (!preempt_copy_done()) { return false; } @@ -481,11 +495,7 @@ struct server_slot { } bool preempt_restore_poll() { - if (!llama_state_seq_copy_done(preempt_cpy_tgt.get())) { - return false; - } - - if (preempt_cpy_dft && !llama_state_seq_copy_done(preempt_cpy_dft.get())) { + if (!preempt_copy_done()) { return false; } @@ -495,15 +505,7 @@ struct server_slot { llama_state_seq_copy_buf_resize(preempt_cpy_dft.get(), 0); } - n_preempt_fail = 0; - - state = state_before_preempt; - - if (state == SLOT_STATE_GENERATING && can_speculate()) { - common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); - } - - return true; + return preempt_resumed(); } // [TAG_PREEMPT_ASYNC] copy the sequence out and release its cells; with a transfer this returns once the copy is issued and the cells stay the slot's until preempt_save_poll() sees it land @@ -620,14 +622,9 @@ struct server_slot { const size_t size_tgt = preempt_state_tgt.size(); const size_t size_dft = preempt_state_dft.size(); - if (llama_state_seq_set_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { - mem.seq_rm(id, -1, -1); - n_preempt_fail++; - return false; - } - - if (size_dft > 0 && - llama_state_seq_set_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + if (llama_state_seq_set_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt || + (size_dft > 0 && + llama_state_seq_set_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft)) { mem.seq_rm(id, -1, -1); n_preempt_fail++; return false; @@ -635,15 +632,7 @@ struct server_slot { preempt_state_free(); - n_preempt_fail = 0; - - state = state_before_preempt; - - if (state == SLOT_STATE_GENERATING && can_speculate()) { - common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); - } - - return true; + return preempt_resumed(); } // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds, for a batch given up after it was built: never-decoded tokens and the draft come off, `sampled` is kept @@ -659,12 +648,7 @@ struct server_slot { state = SLOT_STATE_PROCESSING_PROMPT; } - spec_draft.clear(); - spec_i_batch.clear(); - spec_ckpt.clear(); - spec_is_replay = false; - - i_batch = -1; + preempt_detach(); } std::vector lora; @@ -3657,32 +3641,61 @@ struct server_context_impl { return preempt_async_ok; } + // [TAG_PREEMPT] park a slot: a synchronous park is finished here, an asynchronous one only issued, and update_preempt_copies() counts it when its copy lands. + // The notice goes with the save, not the cell release: preempt_save() has already detached the slot, so a release-time notice would leave the copy's silence unexplained. + bool preempt_park(server_slot & slot, int64_t t_start) { + slot.t_preempt_copy_us = t_start; + + if (!slot.preempt_save()) { + return false; + } + + preempt_log_ram_kind(slot); + + if (slot.state == SLOT_STATE_PREEMPTED) { + metrics.n_preempt++; + } + + send_preempt_notice(slot, true); + + return true; + } + + // [TAG_PREEMPT_ASYNC] a park whose copy has landed; `note` says how it was waited for, if it was + void preempt_parked(server_slot & slot, const char * note) { + metrics.n_preempt++; + + SLT_WRN(slot, "park completed after %.2f ms%s: %d cells released, %.1f MiB parked, kv %d/%d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, note, + slot.prompt.n_tokens(), + slot.preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_ctx); + } + + // [TAG_PREEMPT] a resume whose copy has landed; announced here rather than where the restore was issued, this being the first moment the slot can be scheduled again + void preempt_restored(server_slot & slot, const char * note) { + metrics.n_resume++; + + preempt_trim_ram(slot); + + send_preempt_notice(slot, false); + + SLT_WRN(slot, "restore completed after %.2f ms%s: %d tokens back in the cache, kv %d/%d, preemptions %d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, note, + slot.prompt.n_tokens(), + preempt_kv_used(), n_ctx, + slot.n_preempt); + } + void update_preempt_copies() { for (auto & slot : slots) { if (slot.state == SLOT_STATE_PREEMPTING) { if (slot.preempt_save_poll()) { - metrics.n_preempt++; - - SLT_WRN(slot, "park completed after %.2f ms: %d cells released, %.1f MiB parked, kv %d/%d\n", - (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, - slot.prompt.n_tokens(), - slot.preempt_state_size() / (1024.0 * 1024.0), - preempt_kv_used(), n_ctx); + preempt_parked(slot, ""); } } else if (slot.state == SLOT_STATE_RESTORING) { if (slot.preempt_restore_poll()) { - metrics.n_resume++; - - preempt_trim_ram(slot); - - // [TAG_PREEMPT] announced here rather than where the restore was issued: preempt_restore_poll() is the first moment the slot can be scheduled again - send_preempt_notice(slot, false); - - SLT_WRN(slot, "restore completed after %.2f ms: %d tokens back in the cache, kv %d/%d, preemptions %d\n", - (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, - slot.prompt.n_tokens(), - preempt_kv_used(), n_ctx, - slot.n_preempt); + preempt_restored(slot, ""); } } } @@ -3706,21 +3719,9 @@ struct server_context_impl { slot.preempt_copy_wait(); - if (!slot.preempt_restore_poll()) { - continue; + if (slot.preempt_restore_poll()) { + preempt_restored(slot, " (waited for, a context shift is due)"); } - - metrics.n_resume++; - - preempt_trim_ram(slot); - - send_preempt_notice(slot, false); - - SLT_WRN(slot, "restore completed after %.2f ms (waited for, a context shift is due): %d tokens back in the cache, kv %d/%d, preemptions %d\n", - (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, - slot.prompt.n_tokens(), - preempt_kv_used(), n_ctx, - slot.n_preempt); } } @@ -3747,12 +3748,7 @@ struct server_context_impl { continue; } - metrics.n_preempt++; - - SLT_WRN(slot, "park completed after %.2f ms (waited for): %d cells released, kv %d/%d\n", - (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, - slot.prompt.n_tokens(), - preempt_kv_used(), n_ctx); + preempt_parked(slot, " (waited for)"); return true; } @@ -3856,13 +3852,7 @@ struct server_context_impl { server_slot * head = parked.front(); // [TAG_PREEMPT_ASYNC] a park still copying holds its cells, so a rotation now would only park another resident on top - bool parking = false; - - for (const auto & slot : slots) { - parking = parking || slot.state == SLOT_STATE_PREEMPTING; - } - - if (!parking && ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { + if (!preempt_copies_in_flight() && ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); const int32_t need = preempt_n_need(*head) + preempt_n_margin(1); @@ -3905,17 +3895,9 @@ struct server_context_impl { params_base.preempt_ram_mib); } - if (pick && pick->preempt_save()) { + if (pick && preempt_park(*pick, t_start)) { server_slot & slot = *pick; - slot.t_preempt_copy_us = t_start; - - if (slot.state != SLOT_STATE_PREEMPTING) { - metrics.n_preempt++; - } - - send_preempt_notice(slot, true); - SLT_WRN(slot, "rotated out after %d context shifts: %d cells released, %.1f MiB parked, a head parked %.1f s takes its turn%s, preemptions %d\n", slot.n_ctx_shift, slot.prompt.n_tokens(), slot.preempt_state_size() / (1024.0 * 1024.0), @@ -3986,22 +3968,9 @@ struct server_context_impl { for (auto & slot : slots) { if (slot.state == SLOT_STATE_GENERATING && (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && - preempt_fits_budget(slot)) { - slot.t_preempt_copy_us = ggml_time_us(); - - if (slot.preempt_save()) { - preempt_log_ram_kind(slot); - - // [TAG_PREEMPT_ASYNC] a slot left PREEMPTING is counted by update_preempt_copies() when its copy lands - if (slot.state == SLOT_STATE_PREEMPTED) { - metrics.n_preempt++; - } - - send_preempt_notice(slot, true); - - SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", - (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); - } + preempt_fits_budget(slot) && preempt_park(slot, ggml_time_us())) { + SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", + (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); } } } @@ -4043,17 +4012,10 @@ struct server_context_impl { const int32_t n_tokens = victim->prompt.n_tokens(); const int64_t t_start = ggml_time_us(); - victim->t_preempt_copy_us = t_start; - - if (!victim->preempt_save()) { + if (!preempt_park(*victim, t_start)) { break; // could not park it; the existing retry ladder is still behind us } - preempt_log_ram_kind(*victim); - - // [TAG_PREEMPT] the notice goes with the save, not the cell release: preempt_save() has already detached the victim, so a release-time notice would leave the copy's silence unexplained - send_preempt_notice(*victim, true); - // [TAG_PREEMPT_ASYNC] the copy has only been issued and the cells are still the victim's, so nothing further can be decided about the pool this iteration if (victim->state == SLOT_STATE_PREEMPTING) { SLT_WRN(*victim, "preempted: %d cells, park issued in %.2f ms (%zu transfers, %.2f ms sync), %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", @@ -4073,8 +4035,6 @@ struct server_context_impl { break; } - metrics.n_preempt++; - SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, (ggml_time_us() - t_start) / 1e3, @@ -5043,18 +5003,12 @@ struct server_context_impl { const int32_t n_tokens = victim->prompt.n_tokens(); const int64_t t_start = ggml_time_us(); - victim->t_preempt_copy_us = t_start; - - if (!victim->preempt_save()) { + if (!preempt_park(*victim, t_start)) { break; } - preempt_log_ram_kind(*victim); - n_parked++; - send_preempt_notice(*victim, true); - // [TAG_PREEMPT_ASYNC] the cells are wanted now, not next iteration: wait for the copy, which releases them if (victim->state == SLOT_STATE_PREEMPTING) { while (preempt_wait_in_flight()) { @@ -5065,8 +5019,6 @@ struct server_context_impl { continue; } - metrics.n_preempt++; - SLT_WRN(*victim, "preempted as a last resort: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, (ggml_time_us() - t_start) / 1e3, From 7e424ad8da8d2bd49f49f26ddfc3a7f879183713 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 08:40:49 +0000 Subject: [PATCH 76/81] server : count the kv reserve in one pass over the slots A restoring slot is charged the step of the state it goes back to, which is what the second pass already did to count the prompts, so mapping the state once at the top of the loop removes both the RESTORING arm and the pass after it. state_before_preempt is only ever GENERATING, PROCESSING_PROMPT or STARTED, since those are the states preempt_pick_victim() parks from. --- tools/server/server-context.cpp | 29 +++++++---------------------- 1 file changed, 7 insertions(+), 22 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7b7e03a3818..f6a081fc7c4 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3454,12 +3454,16 @@ struct server_context_impl { int32_t res = 0; int32_t res_pmt = 0; + int32_t n_pmt = 0; // [TAG_EXACT_CONCURRENCY] reserve the cells the next step ADDS, not its tokens: the used figure already rounds every tail page up, and only a page crossing can empty the pool for (const auto & slot : slots) { const int32_t n_cur = slot.prompt.n_tokens(); - switch (slot.state) { + // [TAG_PREEMPT_ASYNC] a restoring slot decodes as soon as its copy lands, so it is charged the step of the state it goes back to, or that first step preempts somebody else + const slot_state state = slot.state == SLOT_STATE_RESTORING ? slot.state_before_preempt : slot.state; + + switch (state) { case SLOT_STATE_GENERATING: case SLOT_STATE_DONE_PROMPT: { @@ -3468,37 +3472,18 @@ struct server_context_impl { case SLOT_STATE_STARTED: case SLOT_STATE_PROCESSING_PROMPT: { + // preempt_n_retained() reads the live state, so a restoring slot is charged from what it holds const int32_t n_have = preempt_n_retained(slot); const int32_t n_left = slot.task ? slot.task->n_tokens() - n_have : 0; res_pmt += preempt_n_cells_step(n_have, std::max(1, std::min(n_batch, n_left))); - } break; - case SLOT_STATE_RESTORING: - { - // [TAG_PREEMPT_ASYNC] it starts decoding as soon as its copy lands, so its step has to be reserved now, or its first step preempts somebody else - if (slot.state_before_preempt == SLOT_STATE_GENERATING) { - res += preempt_n_cells_step(n_cur, 1 + preempt_n_spec(slot)); - } else { - const int32_t n_left = slot.task ? slot.task->n_tokens() - n_cur : 0; - - res_pmt += preempt_n_cells_step(n_cur, std::max(1, std::min(n_batch, n_left))); - } + n_pmt++; } break; default: break; } } - int32_t n_pmt = 0; - - for (const auto & slot : slots) { - const slot_state state = slot.state == SLOT_STATE_RESTORING ? slot.state_before_preempt : slot.state; - - if (state == SLOT_STATE_STARTED || state == SLOT_STATE_PROCESSING_PROMPT) { - n_pmt++; - } - } - return res + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1)); } From 7549c6806ccc52889545326db523afe36aa368d4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 08:48:11 +0000 Subject: [PATCH 77/81] tests : share the preemption server setup and the prompts _start() sets the attributes a test needs and hands back a reader of the log, _complete_all() and _stream_both() send the pair of prompts every second test sends, _assert_completed() holds the completion assertions. Every test and every assertion is unchanged; only the plumbing moved. --- tools/server/tests/unit/test_preempt.py | 308 ++++++------------ .../server/tests/unit/test_preempt_notify.py | 87 +++-- 2 files changed, 135 insertions(+), 260 deletions(-) diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 56493ed3048..72aeb44f693 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -54,19 +54,57 @@ def _complete(n_predict: int, prompt: str = "Hi how are you"): return res +_PROMPT_A = "Once upon a time there was a brave knight who" +_PROMPT_B = "The quick brown fox jumps over the lazy dog and" +_PROMPT_C = "In a small village by the sea there lived a fisherman who" + + +def _start(**kwargs) -> LogReader: + """Start the server with these settings, and read its log from the first line.""" + for key, value in kwargs.items(): + setattr(server, key, value) + server.start() + return LogReader(server.log_path) + + +def _late(n_predict: int, prompt: str, delay: float = 0.02): + time.sleep(delay) + return _complete(n_predict, prompt) + + +def _complete_all(n_predict: int, prompts=(_PROMPT_A, _PROMPT_B)): + return parallel_function_calls([(_complete, (n_predict, prompt)) for prompt in prompts]) + + +def _complete_all_raw(n_predict: int, prompts): + """As _complete_all, without return_tokens: these ask for thousands of tokens.""" + return parallel_function_calls([ + (server.make_request, ("POST", "/completion", { + "prompt": prompt, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, + })) for prompt in prompts + ]) + + +def _assert_completed(results, n_predict: int, whole: bool = False): + """Every request generated what it asked for; `whole` also pins the untruncated body.""" + for res in results: + assert res.status_code == 200, res.body + assert res.body["timings"]["predicted_n"] == n_predict + if whole: + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_predict + + def test_forced_preemption_does_not_change_the_output(): # park and restore the only running slot every 8 tokens: the batch shape is the same at every step, so any difference in the output is the preemption's fault - global server - server.n_ctx = 512 - server.start() + _start(n_ctx=512) reference = _complete(64) assert reference.status_code == 200 assert reference.body["timings"]["predicted_n"] == 64 server.stop() os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" - server.start() - log = LogReader(server.log_path) + log = _start() assert "LLAMA_SERVER_PREEMPT_EVERY = 8" in log.drain() preempted = _complete(64) @@ -83,43 +121,27 @@ def test_forced_preemption_does_not_change_the_output(): def test_two_slots_that_overflow_the_pool_together_both_finish(): # each request fits the pool alone (168 of 256 cells) but not together; without preemption both end with "Context size has been exceeded" - global server - server.n_ctx = 256 - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) n_predict = 160 - results = parallel_function_calls([ - (_complete, (n_predict, "Once upon a time there was a brave knight who")), - (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text assert "preempted:" in text assert "resumed after" in text - for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict - assert res.body["truncated"] is False - assert len(res.body["tokens"]) == n_predict + _assert_completed(results, n_predict, whole=True) def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages(): # a block allocator gives a whole block to one sequence, so the planner has to count cells: counting tokens it sees room the allocator cannot find. GRANULARITY injects the size. - global server - server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_GRANULARITY"] = "64" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) assert "LLAMA_SERVER_PREEMPT_GRANULARITY = 64" in log.drain() n_predict = 160 - results = parallel_function_calls([ - (_complete, (n_predict, "Once upon a time there was a brave knight who")), - (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text @@ -131,11 +153,7 @@ def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages(): assert held and wanted, f"the planner logged no figures:\n{text}" assert all(n % 64 == 0 for n in held + wanted), f"not whole blocks: {held} {wanted}" - for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict - assert res.body["truncated"] is False - assert len(res.body["tokens"]) == n_predict + _assert_completed(results, n_predict, whole=True) _WORDS = ( @@ -164,10 +182,7 @@ def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]: def test_two_prompts_that_overflow_the_pool_together_both_finish(): - global server - server.n_ctx = 256 - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) prompt_a, n_a = _prompt_of_about(150, "Alpha") prompt_b, n_b = _prompt_of_about(150, "Bravo") @@ -175,28 +190,21 @@ def test_two_prompts_that_overflow_the_pool_together_both_finish(): assert n_a + n_predict <= 256 and n_b + n_predict <= 256 assert n_a + n_b + 2 * n_predict > 256 - results = parallel_function_calls([ - (_complete, (n_predict, prompt_a)), - (_complete, (n_predict, prompt_b)), - ]) + results = _complete_all(n_predict, [prompt_a, prompt_b]) text = log.drain() assert "Context size has been exceeded" not in text assert "preempted:" in text assert "resumed after" in text + _assert_completed(results, n_predict) for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict assert len(res.body["tokens"]) == n_predict def test_a_generating_slot_and_a_large_prompt_both_finish(): # a long generation meets a large prompt arriving beside it: the prompt is admitted chunk by chunk, whoever is smaller is parked, and both finish - global server - server.n_ctx = 256 - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) prompt_b, n_b = _prompt_of_about(150, "Charlie") n_predict_a = 230 @@ -204,10 +212,6 @@ def test_a_generating_slot_and_a_large_prompt_both_finish(): assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 assert 8 + n_predict_a + n_b + n_predict_b > 256 - def _late(n_predict, prompt): - time.sleep(0.02) - return _complete(n_predict, prompt) - results = parallel_function_calls([ (_complete, (n_predict_a, "Hi how are you")), (_late, (n_predict_b, prompt_b)), @@ -225,17 +229,11 @@ def _late(n_predict, prompt): def test_preempt_ram_zero_disables_preemption(): # --preempt-ram 0 switches back to the old behaviour: nothing is parked and the KV-full path ends the requests - global server - server.n_ctx = 256 os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) n_predict = 160 - results = parallel_function_calls([ - (_complete, (n_predict, "Once upon a time there was a brave knight who")), - (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "preempted:" not in text @@ -244,10 +242,7 @@ def test_preempt_ram_zero_disables_preemption(): def test_metrics_and_slots_report_the_parked_state(): - global server - server.n_ctx = 256 - server.server_metrics = True - server.start() + _start(n_ctx=256, server_metrics=True) res = server.make_request("GET", "/slots") assert res.status_code == 200 @@ -256,10 +251,7 @@ def test_metrics_and_slots_report_the_parked_state(): assert slot["n_preempt"] == 0 n_predict = 160 - results = parallel_function_calls([ - (_complete, (n_predict, "Once upon a time there was a brave knight who")), - (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), - ]) + results = _complete_all(n_predict) for res in results: assert res.status_code == 200 @@ -288,11 +280,7 @@ def test_metrics_and_slots_report_the_parked_state(): def _start_async(**kwargs) -> str: """Start the server with the asynchronous path asked for, and return its log so far.""" os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "1" - for key, value in kwargs.items(): - setattr(server, key, value) - server.start() - with open(server.log_path) as f: - return f.read() + return _start(**kwargs).drain() def _require_async(text: str): @@ -302,10 +290,7 @@ def _require_async(text: str): def test_async_preemption_does_not_change_the_output(): # the synchronous determinism question asked of the asynchronous path: with one request the batch shape is fixed, so a continuation that is not byte-identical is the transfer's fault - global server - server.n_ctx = 512 - server.n_gpu_layer = 99 - text = _start_async() + text = _start_async(n_ctx=512, n_gpu_layer=99) _require_async(text) res_plain = _complete(64) @@ -313,8 +298,7 @@ def test_async_preemption_does_not_change_the_output(): server.stop() os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" - server.start() - log = LogReader(server.log_path) + log = _start() res_preempted = _complete(64) assert res_preempted.status_code == 200 @@ -332,27 +316,19 @@ def test_async_preemption_does_not_change_the_output(): def test_async_preemption_under_load_keeps_every_slot_and_its_output(): - global server - server.n_ctx = 256 - server.n_gpu_layer = 99 - text = _start_async() + text = _start_async(n_ctx=256, n_gpu_layer=99) _require_async(text) n_predict = 160 - prompts = [ - "Once upon a time there was a brave knight who", - "The quick brown fox jumps over the lazy dog and", - ] - alone = [_complete(n_predict, prompt) for prompt in prompts] + alone = [_complete(n_predict, prompt) for prompt in (_PROMPT_A, _PROMPT_B)] for res in alone: assert res.status_code == 200 server.stop() - server.start() - log = LogReader(server.log_path) + log = _start() - together = parallel_function_calls([(_complete, (n_predict, prompt)) for prompt in prompts]) + together = _complete_all(n_predict) text = log.drain() _require_async(text) @@ -360,9 +336,8 @@ def test_async_preemption_under_load_keeps_every_slot_and_its_output(): assert "preempted:" in text assert "resumed after" in text + _assert_completed(together, n_predict) for res, ref in zip(together, alone): - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict assert res.body["truncated"] is False assert res.body["tokens"] == ref.body["tokens"] @@ -382,15 +357,12 @@ def _cancel_soon(n_predict: int, prompt: str, timeout: float): def test_cancel_while_a_copy_is_in_flight_frees_the_slot(): # a cancelled request can reach release() with a park or a resume still running, where the host buffer is freed and the cells handed on, so both have to wait for the copy - global server - server.n_ctx = 512 - server.n_gpu_layer = 99 os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" - text = _start_async() + text = _start_async(n_ctx=512, n_gpu_layer=99) _require_async(text) for i in range(4): - _cancel_soon(96, "Once upon a time there was a brave knight who", 0.05 + 0.1 * i) + _cancel_soon(96, _PROMPT_A, 0.05 + 0.1 * i) deadline = time.time() + 120 while time.time() < deadline: @@ -418,13 +390,9 @@ def test_cancel_while_a_copy_is_in_flight_frees_the_slot(): def test_no_preempt_async_falls_back_to_the_synchronous_path(): # The flag has to really switch it off, so that the two can be compared on one binary. - global server - server.n_ctx = 512 - server.n_gpu_layer = 99 os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "0" os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=512, n_gpu_layer=99) res = _complete(64) assert res.status_code == 200 @@ -439,12 +407,7 @@ def test_no_preempt_async_falls_back_to_the_synchronous_path(): def test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_everything(): # [TAG_PREEMPT_ASYNC] the case the async path made worse than the synchronous one: an asynchronous park does not return the cells before update_slots() carries on - global server - server.n_ctx = 512 - server.n_gpu_layer = 99 - server.n_slots = 4 - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=512, n_gpu_layer=99, n_slots=4) prompt_a, n_a = _prompt_of_about(100, "Alpha") prompt_b, n_b = _prompt_of_about(100, "Bravo") @@ -456,15 +419,11 @@ def test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_everyt assert max(n_a, n_b, n_c) + n_predict_abc < 512 and n_d + n_predict_d < 512 assert n_a + n_b + n_c + 3 * n_predict_abc > 512 - def _late(n_predict, prompt): - time.sleep(0.25) - return _complete(n_predict, prompt) - results = parallel_function_calls([ (_complete, (n_predict_abc, prompt_a)), (_complete, (n_predict_abc, prompt_b)), (_complete, (n_predict_abc, prompt_c)), - (_late, (n_predict_d, prompt_d)), + (_late, (n_predict_d, prompt_d, 0.25)), ]) text = log.drain() @@ -480,38 +439,26 @@ def _late(n_predict, prompt): def test_two_prompts_near_the_context_size_both_complete(): # the second prompt is parked before it takes any cells and is too close to n_ctx to leave the usual margin, but must still be restored once the first finishes - global server - server.n_ctx = 256 - server.n_batch = 256 - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256, n_batch=256) base = server.make_request("POST", "/tokenize", data={"content": "Once upon a time there was a little girl"}).body["tokens"] long_prompt = (base * 64)[:240] n_predict = 4 - together = parallel_function_calls([(_complete, (n_predict, long_prompt)) for _ in range(2)]) + together = _complete_all(n_predict, [long_prompt, long_prompt]) text = log.drain() assert "cannot fit the pool" not in text - for res in together: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict + _assert_completed(together, n_predict) def test_the_last_resort_parks_instead_of_ending_everyone(): - global server - server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) assert "LLAMA_SERVER_PREEMPT_PLANNER = off" in log.drain() n_predict = 160 - results = parallel_function_calls([ - (_complete, (n_predict, "Once upon a time there was a brave knight who")), - (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text @@ -520,44 +467,29 @@ def test_the_last_resort_parks_instead_of_ending_everyone(): assert "last resort: batch given up" in text assert "resumed after" in text - for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict - assert res.body["truncated"] is False - assert len(res.body["tokens"]) == n_predict + _assert_completed(results, n_predict, whole=True) def test_the_last_resort_works_with_an_unlimited_budget(): # --preempt-ram -1 is the documented unlimited setting and must enable the last resort too - global server - server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" os.environ["LLAMA_ARG_PREEMPT_RAM"] = "-1" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) n_predict = 160 - results = parallel_function_calls([ - (_complete, (n_predict, "Once upon a time there was a brave knight who")), - (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text assert "preempted as a last resort" in text - for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict + _assert_completed(results, n_predict) def test_the_last_resort_rewinds_a_prompt_in_flight(): # the failed chunk comes back off the slot's tokens and is processed again after the resume, neither skipped nor fed twice - global server - server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) prompt_b, n_b = _prompt_of_about(150, "Charlie") n_predict_a = 230 @@ -565,10 +497,6 @@ def test_the_last_resort_rewinds_a_prompt_in_flight(): assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 assert 8 + n_predict_a + n_b + n_predict_b > 256 - def _late(n_predict, prompt): - time.sleep(0.02) - return _complete(n_predict, prompt) - results = parallel_function_calls([ (_complete, (n_predict_a, "Hi how are you")), (_late, (n_predict_b, prompt_b)), @@ -588,45 +516,23 @@ def _late(n_predict, prompt): def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_head(): # with context shift on the resident would hold half the pool for as long as it generates, so once the head has waited its turn the resident is parked and the two take turns - global server - server.n_ctx = 256 - server.enable_ctx_shift = True - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256, enable_ctx_shift=True) n_predict = 12000 - results = parallel_function_calls([ - (_complete, (n_predict, "Once upon a time there was a brave knight who")), - (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text assert "slot context shift" in text assert "rotated out after" in text - for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict + _assert_completed(results, n_predict) def test_the_rotation_parks_a_resident_that_lets_the_head_in(): - global server - server.n_slots = 3 - server.n_ctx = 384 - server.enable_ctx_shift = True - server.start() + _start(n_slots=3, n_ctx=384, enable_ctx_shift=True) n_predict = 9000 - prompts = [ - "Once upon a time there was a brave knight who", - "The quick brown fox jumps over the lazy dog and", - "In a small village by the sea there lived a fisherman who", - ] - results = parallel_function_calls([ - (server.make_request, ("POST", "/completion", { - "prompt": p, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, - })) for p in prompts - ]) + results = _complete_all_raw(n_predict, (_PROMPT_A, _PROMPT_B, _PROMPT_C)) for res in results: assert res.status_code == 200, res.body assert res.body["tokens_predicted"] == n_predict @@ -637,16 +543,13 @@ def test_the_rotation_parks_a_resident_that_lets_the_head_in(): def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_server_lives(): # a family member is not a victim for the other, so a two-completion request gets the context error it would get alone and the server carries on - global server - server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) res = server.make_request("POST", "/completion", data={ "n_predict": 160, "n_cmpl": 2, - "prompt": "Once upon a time there was a brave knight who", + "prompt": _PROMPT_A, "ignore_eos": True, "return_tokens": True, "temperature": 0.0, @@ -666,53 +569,32 @@ def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_ def test_a_restored_slot_gives_its_idle_buffer_back_when_another_slot_needs_to_park(): # an asynchronous slot keeps its pinned buffer after a restore, and that idle capacity counts against --preempt-ram: unless it is given back, the first restore spends the budget - global server - server.n_ctx = 8192 - server.n_gpu_layer = 99 os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "256" os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" - text = _start_async() + text = _start_async(n_ctx=8192, n_gpu_layer=99) _require_async(text) log = LogReader(server.log_path) n_predict = 1800 - results = parallel_function_calls([ - (_complete, (n_predict, "Once upon a time there was a brave knight who")), - (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text assert "idle parked RAM returned" in text, "the idle buffer of a restored slot was never given back" - import re parked = re.findall(r"id\s+(\d+) \| task \d+ \| preempted on request", text) assert {"0", "1"} <= set(parked), f"only slots {sorted(set(parked))} were ever parked" + _assert_completed(results, n_predict) for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict assert res.body["truncated"] is False def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_when_a_resident_finishes(): # a rotation holds both states at once, since the resident is parked before the head is restored and freed, so a budget for two heads but not a head plus the resident must refuse - global server - server.n_slots = 3 - server.n_ctx = 2048 - server.enable_ctx_shift = True os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" - server.start() + _start(n_slots=3, n_ctx=2048, enable_ctx_shift=True) n_predict = 12000 - prompts = [ - "Once upon a time there was a brave knight who", - "The quick brown fox jumps over the lazy dog and", - "In a small village by the sea there lived a fisherman who", - ] - results = parallel_function_calls([ - (server.make_request, ("POST", "/completion", { - "prompt": p, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, - })) for p in prompts - ]) + results = _complete_all_raw(n_predict, (_PROMPT_A, _PROMPT_B, _PROMPT_C)) for res in results: assert res.status_code == 200, res.body assert res.body["tokens_predicted"] == n_predict @@ -724,7 +606,6 @@ def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_w def test_a_recurrent_model_is_served_without_preemption(): # a recurrent cache holds one state per sequence whatever its length, so preemption is off for such a model and the forced-park knob parks nothing - global server path = os.environ.get("LLAMA_SERVER_TEST_RECURRENT_MODEL") if path: server.model_file = path @@ -736,10 +617,7 @@ def test_a_recurrent_model_is_served_without_preemption(): server.n_ctx = 1024 os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" server.start(timeout_seconds=300) - results = parallel_function_calls([ - (_complete, (64, "Once upon a time")), - (_complete, (64, "The quick brown fox")), - ]) + results = _complete_all(64, ["Once upon a time", "The quick brown fox"]) for res in results: assert res.status_code == 200, res.body assert res.body["tokens_predicted"] == 64 diff --git a/tools/server/tests/unit/test_preempt_notify.py b/tools/server/tests/unit/test_preempt_notify.py index 796c22b1930..485a6bbc336 100644 --- a/tools/server/tests/unit/test_preempt_notify.py +++ b/tools/server/tests/unit/test_preempt_notify.py @@ -81,10 +81,32 @@ def _chat_payload(n_predict: int) -> dict: } -def test_a_stream_announces_its_parks_and_the_body_is_unchanged(): - global server - server.n_ctx = 512 +_PROMPT_A = "Once upon a time there was a brave knight who" +_PROMPT_B = "The quick brown fox jumps over the lazy dog and" + + +def _start(**kwargs): + """Start the server with these settings.""" + for key, value in kwargs.items(): + setattr(server, key, value) server.start() + + +def _final(datas: list[str]) -> dict: + """The last response object of a finished stream, past the [DONE] marker.""" + return json.loads([d for d in datas if d != "[DONE]"][-1]) + + +def _stream_both(n_predict: int): + """One streaming completion per prompt, both at once.""" + return parallel_function_calls([ + (_stream_raw, ("/completion", _completion_payload(n_predict) | {"prompt": prompt})) + for prompt in (_PROMPT_A, _PROMPT_B) + ]) + + +def test_a_stream_announces_its_parks_and_the_body_is_unchanged(): + _start(n_ctx=512) ref_comments, ref_datas = _stream_raw("/completion", _completion_payload(64)) assert not any(c.startswith(": preempted") or c.startswith(": resumed") for c in ref_comments) assert _content(ref_datas) @@ -110,10 +132,8 @@ def _pieces(ds): def test_the_oai_chat_stream_carries_the_same_comments(): - global server - server.n_ctx = 512 os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" - server.start() + _start(n_ctx=512) comments, datas = _stream_raw("/v1/chat/completions", _chat_payload(48)) assert ": preempted" in comments and ": resumed" in comments assert datas[-1] == "[DONE]" @@ -121,10 +141,8 @@ def test_the_oai_chat_stream_carries_the_same_comments(): def test_non_streaming_requests_see_nothing(): - global server - server.n_ctx = 512 os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" - server.start() + _start(n_ctx=512) res = server.make_request("POST", "/completion", data={ "n_predict": 32, "prompt": "Hi how are you", @@ -138,20 +156,13 @@ def test_non_streaming_requests_see_nothing(): def test_two_overflowing_streams_both_finish_and_the_parked_one_says_so(): - global server - server.n_ctx = 256 - server.start() + _start(n_ctx=256) n_predict = 160 - p1 = _completion_payload(n_predict) | {"prompt": "Once upon a time there was a brave knight who"} - p2 = _completion_payload(n_predict) | {"prompt": "The quick brown fox jumps over the lazy dog and"} - results = parallel_function_calls([ - (_stream_raw, ("/completion", p1)), - (_stream_raw, ("/completion", p2)), - ]) + results = _stream_both(n_predict) announced = 0 for comments, datas in results: - final = json.loads([d for d in datas if d != "[DONE]"][-1]) + final = _final(datas) assert final["timings"]["predicted_n"] == n_predict assert final["truncated"] is False if ": preempted" in comments: @@ -162,14 +173,11 @@ def test_two_overflowing_streams_both_finish_and_the_parked_one_says_so(): def test_a_stream_parked_before_its_first_token_starts_with_the_notice(): # A request parked while still processing its prompt has no token to send yet, so the response starts with the notice instead of a silent connection. - global server - global server - server.n_ctx = 512 - server.n_batch = 512 # the whole prompt in one batch, so the planner sees its size at once - server.start() + # n_batch: the whole prompt in one batch, so the planner sees its size at once + _start(n_ctx=512, n_batch=512) url = f"http://{server.server_host}:{server.server_port}/completion" - first = _completion_payload(390) | {"prompt": " ".join(["Once upon a time there was a brave knight who"] * 6)} - second = _completion_payload(32) | {"prompt": " ".join(["The quick brown fox jumps over the lazy dog and"] * 14)} + first = _completion_payload(390) | {"prompt": " ".join([_PROMPT_A] * 6)} + second = _completion_payload(32) | {"prompt": " ".join([_PROMPT_B] * 14)} timeline = [] lock = threading.Lock() @@ -201,25 +209,16 @@ def _run(name, payload, started=None): assert events[0] == ": preempted" and events[1] == ": resumed" and events[2].startswith("data: "), events[:3] datas = [line[6:] for _, line in second_lines if line.startswith("data: ")] assert _content(datas) - final = json.loads([d for d in datas if d != "[DONE]"][-1]) - assert final["tokens_predicted"] == 32 + assert _final(datas)["tokens_predicted"] == 32 def test_a_resident_rotated_out_for_a_parked_head_is_told_so(): - global server - server.n_ctx = 256 - server.enable_ctx_shift = True - server.start() + _start(n_ctx=256, enable_ctx_shift=True) n_predict = 12000 - p1 = _completion_payload(n_predict) | {"prompt": "Once upon a time there was a brave knight who"} - p2 = _completion_payload(n_predict) | {"prompt": "The quick brown fox jumps over the lazy dog and"} - results = parallel_function_calls([ - (_stream_raw, ("/completion", p1)), - (_stream_raw, ("/completion", p2)), - ]) + results = _stream_both(n_predict) n_parked = 0 for comments, datas in results: - final = json.loads([d for d in datas if d != "[DONE]"][-1]) + final = _final(datas) assert final["tokens_predicted"] == n_predict seq = [c for c in comments if c in (": preempted", ": resumed")] assert seq == [": preempted", ": resumed"] * (len(seq) // 2), seq @@ -229,14 +228,12 @@ def test_a_resident_rotated_out_for_a_parked_head_is_told_so(): def test_an_oversized_prompt_is_errored_instead_of_parked(): # A slot just given a task has not passed the prompt checks yet, and a notice opens the stream, so parking it would turn a plain error response into 200 plus an in-stream one. - global server - server.n_ctx = 512 - server.n_batch = 512 # the whole prompt in one batch, so the planner sees its size at once os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" - server.start() + # n_batch: the whole prompt in one batch, so the planner sees its size at once + _start(n_ctx=512, n_batch=512) url = f"http://{server.server_host}:{server.server_port}/completion" - resident = _completion_payload(390) | {"prompt": " ".join(["Once upon a time there was a brave knight who"] * 6)} - oversized = _completion_payload(16) | {"prompt": " ".join(["The quick brown fox jumps over the lazy dog and"] * 80)} + resident = _completion_payload(390) | {"prompt": " ".join([_PROMPT_A] * 6)} + oversized = _completion_payload(16) | {"prompt": " ".join([_PROMPT_B] * 80)} started = threading.Event() From 4cf9baf942c2d181c5c836fd212ece56fedf3333 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 08:52:51 +0000 Subject: [PATCH 78/81] llama : build the asynchronous state io on the host io llama_io_write_host_async and llama_io_read_host_async repeated the whole buffer walk of llama_io_write_host and llama_io_read_host to change only what the destructor does with the tensors it collected. They derive from them now, and the base skips its own flush when a derived class says it posts the copies itself. --- src/llama-context.cpp | 119 +++++++++++------------------------------- 1 file changed, 30 insertions(+), 89 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 7f81f570ad1..322d968708f 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2716,6 +2716,10 @@ class llama_io_write_host : public llama_io_write_i { uint8_t * p, size_t len) : ptr(p), buf_size(len) {} ~llama_io_write_host() { + if (deferred) { + return; // [TAG_STATE_ASYNC] the derived class posts the copies itself + } + llama_io_emit(winfos, 0, winfos.size(), [](ggml_tensor * tensor, uint8_t * ptr, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { @@ -2750,10 +2754,8 @@ class llama_io_write_host : public llama_io_write_i { return size_written; } -private: - uint8_t * ptr; - size_t buf_size = 0; - size_t size_written = 0; +protected: + llama_io_write_host(uint8_t * p, size_t len, bool deferred) : ptr(p), buf_size(len), deferred(deferred) {} struct write_info { ggml_tensor * tensor; @@ -2762,6 +2764,12 @@ class llama_io_write_host : public llama_io_write_i { size_t offset; }; std::vector winfos; + +private: + uint8_t * ptr; + size_t buf_size = 0; + size_t size_written = 0; + const bool deferred = false; }; class llama_io_read_host : public llama_io_read_i { @@ -2769,6 +2777,10 @@ class llama_io_read_host : public llama_io_read_i { llama_io_read_host(const uint8_t * p, size_t len) : ptr(p), buf_size(len) {} ~llama_io_read_host() { + if (deferred) { + return; // [TAG_STATE_ASYNC] the derived class posts the copies itself + } + // flush the reads for (size_t i = 0; i < rinfos.size();) { auto * tensor = rinfos[i].tensor; @@ -2843,10 +2855,8 @@ class llama_io_read_host : public llama_io_read_i { return size_read; } -private: - const uint8_t * ptr; - size_t buf_size = 0; - size_t size_read = 0; +protected: + llama_io_read_host(const uint8_t * p, size_t len, bool deferred) : ptr(p), buf_size(len), deferred(deferred) {} struct read_info { ggml_tensor * tensor; @@ -2855,6 +2865,12 @@ class llama_io_read_host : public llama_io_read_i { size_t offset; }; std::vector rinfos; + +private: + const uint8_t * ptr; + size_t buf_size = 0; + size_t size_read = 0; + const bool deferred = false; }; class llama_io_write_file : public llama_io_write_i { @@ -3357,10 +3373,11 @@ struct llama_state_seq_copy { } }; -class llama_io_write_host_async : public llama_io_write_i { +// [TAG_STATE_ASYNC] the buffer walk of llama_io_write_host, with the copies posted on the transfer's stream instead of made here +class llama_io_write_host_async : public llama_io_write_host { public: llama_io_write_host_async(uint8_t * p, size_t len, llama_state_seq_copy & cpy) : - ptr(p), buf_size(len), cpy(cpy) {} + llama_io_write_host(p, len, true), cpy(cpy) {} // posted from the destructor, and only once serialisation reached the end: a caller told of a partial failure by a zero return is free to reuse the buffer at once void commit() { @@ -3384,54 +3401,17 @@ class llama_io_write_host_async : public llama_io_write_i { cpy.record(); } - void write(const void * src, size_t size) override { - if (size > buf_size) { - throw std::runtime_error("unexpectedly reached end of buffer"); - } - memcpy(ptr, src, size); - ptr += size; - size_written += size; - buf_size -= size; - } - - void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { - if (size > buf_size) { - throw std::runtime_error("unexpectedly reached end of buffer"); - } - - winfos.push_back({tensor, ptr, size, offset}); - - ptr += size; - size_written += size; - buf_size -= size; - } - - size_t n_bytes() override { - return size_written; - } - private: - uint8_t * ptr; - size_t buf_size = 0; - size_t size_written = 0; - - struct write_info { - ggml_tensor * tensor; - uint8_t * ptr; - size_t size; - size_t offset; - }; - std::vector winfos; - llama_state_seq_copy & cpy; bool committed = false; }; -class llama_io_read_host_async : public llama_io_read_i { +// [TAG_STATE_ASYNC] the read half of the same, without llama_io_read_host's whole-tensor staging: a write-back would undo whatever the sequences sharing the tensor wrote while these copies ran +class llama_io_read_host_async : public llama_io_read_host { public: llama_io_read_host_async(const uint8_t * p, size_t len, llama_state_seq_copy & cpy) : - ptr(p), buf_size(len), cpy(cpy) {} + llama_io_read_host(p, len, true), cpy(cpy) {} // see llama_io_write_host_async::commit(): a restore that failed part way has dropped the sequence, and copies posted for it would write cells that are no longer its own void commit() { @@ -3443,7 +3423,6 @@ class llama_io_read_host_async : public llama_io_read_i { return; } - // no whole-tensor staging here, unlike the synchronous path above: a write-back would undo whatever the sequences sharing the tensor wrote while these copies ran llama_io_emit(rinfos, 0, rinfos.size(), [this](ggml_tensor * tensor, const uint8_t * ptr, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { @@ -3456,45 +3435,7 @@ class llama_io_read_host_async : public llama_io_read_i { cpy.record(); } - void read(void * dst, size_t size) override { - if (size > buf_size) { - throw std::runtime_error("unexpectedly reached end of buffer"); - } - memcpy(dst, ptr, size); - ptr += size; - size_read += size; - buf_size -= size; - } - - void read_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { - if (size > buf_size) { - throw std::runtime_error("unexpectedly reached end of buffer"); - } - - rinfos.push_back({tensor, ptr, size, offset}); - - ptr += size; - size_read += size; - buf_size -= size; - } - - size_t n_bytes() override { - return size_read; - } - private: - const uint8_t * ptr; - size_t buf_size = 0; - size_t size_read = 0; - - struct read_info { - ggml_tensor * tensor; - const uint8_t * ptr; - size_t size; - size_t offset; - }; - std::vector rinfos; - llama_state_seq_copy & cpy; bool committed = false; From 00855c230188a2122d12302f0f97166d38c6012b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 08:58:42 +0000 Subject: [PATCH 79/81] server : drop the wrappers that only ever had one reader server_exact_concurrency() was a third copy of the LLAMA_EXACT_CONCURRENCY reader that common already has, preempt_async_active() returned a flag, llama_state_seq_copy_free is a deleter on its own, preempt_normalize_started() already refuses what its caller's guard refused, and --preempt-ram becomes a ceiling in one place instead of two. --- tools/server/server-context.cpp | 48 +++++++-------------------------- 1 file changed, 10 insertions(+), 38 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f6a081fc7c4..01a4bda7914 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -39,16 +39,6 @@ constexpr int HTTP_POLLING_SECONDS = 1; -// [TAG_EXACT_CONCURRENCY] read from the env: the answer is needed before a context exists -static bool server_exact_concurrency() { - static const bool enabled = []() { - const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); - return val && atoi(val) != 0; - }(); - - return enabled; -} - static common_speculative_output_limits server_output_limits(const common_params & params) { if (params.embedding || (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) { @@ -99,16 +89,12 @@ constexpr int64_t PREEMPT_ROTATE_US = 2ll * 1000 * 1000; // a resident cyclin // [TAG_PREEMPT_ASYNC] an asynchronous park only releases its cells when its copy lands, so it must fire this many decode steps before the pool would run out constexpr int32_t PREEMPT_N_ASYNC_STEPS = 8; -struct llama_state_seq_copy_deleter { - void operator()(llama_state_seq_copy * cpy) const { llama_state_seq_copy_free(cpy); } -}; - using llama_state_seq_copy_ptr = std::shared_ptr; static llama_state_seq_copy_ptr llama_state_seq_copy_make(llama_context * ctx) { llama_state_seq_copy * cpy = ctx ? llama_state_seq_copy_init(ctx) : nullptr; - return cpy ? llama_state_seq_copy_ptr(cpy, llama_state_seq_copy_deleter{}) : llama_state_seq_copy_ptr(); + return cpy ? llama_state_seq_copy_ptr(cpy, llama_state_seq_copy_free) : llama_state_seq_copy_ptr(); } // [TAG_EXACT_CONCURRENCY] the planner counts cells, not tokens: a page belongs to one sequence, so a token count sees room find_slot cannot find and nobody is ever parked @@ -3323,12 +3309,13 @@ struct server_context_impl { } } - bool preempt_fits_budget(const server_slot & slot) { - if (params_base.preempt_ram_mib < 0) { - return true; - } + // the --preempt-ram ceiling in bytes; the unlimited setting is a ceiling nothing reaches + size_t preempt_ram_budget() const { + return params_base.preempt_ram_mib < 0 ? SIZE_MAX : (size_t) params_base.preempt_ram_mib * 1024 * 1024; + } - const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; + bool preempt_fits_budget(const server_slot & slot) { + const size_t budget = preempt_ram_budget(); // what this slot already holds is counted by preempt_ram_used() and reused, so a park costs only the rest const size_t held = slot.preempt_state_size(); @@ -3342,13 +3329,7 @@ struct server_context_impl { // [TAG_PREEMPT_ASYNC] over budget, a buffer held by a running slot would keep every other slot from being parked at all void preempt_trim_ram(server_slot & slot) { - if (params_base.preempt_ram_mib < 0) { - return; - } - - const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; - - if (preempt_ram_used() > budget && slot.preempt_state_size() > 0) { + if (preempt_ram_used() > preempt_ram_budget() && slot.preempt_state_size() > 0) { SLT_INF(slot, "%.1f MiB of parked RAM returned: the pool is over its budget\n", slot.preempt_state_size() / (1024.0 * 1024.0)); slot.preempt_state_free(); } @@ -3431,7 +3412,7 @@ struct server_context_impl { // [TAG_PREEMPT_ASYNC] the room the pool is kept clear of, so everything still decoding has somewhere to put its tokens until a park lands; a resume candidate is charged the same runway int32_t preempt_n_margin(int32_t n_additional_running = 0) const { - if (!preempt_async_active()) { + if (!preempt_async_ok) { // [TAG_EXACT_CONCURRENCY] a margin of eight cells is no margin where a step can cost a whole page return preempt_n_cells(PREEMPT_N_MARGIN); } @@ -3493,10 +3474,6 @@ struct server_context_impl { bool res = false; for (auto & slot : slots) { - if (slot.state != SLOT_STATE_STARTED || !slot.task) { - continue; - } - const int32_t before = slot.prompt.n_tokens(); preempt_normalize_started(slot); @@ -3621,11 +3598,6 @@ struct server_context_impl { } // called once per update_slots(), before the batch is built: every slot is then at a token boundary with no draft in flight, so it can be removed whole - // [TAG_PREEMPT_ASYNC] is any slot parking or resuming through a transfer right now - bool preempt_async_active() const { - return preempt_async_ok; - } - // [TAG_PREEMPT] park a slot: a synchronous park is finished here, an asynchronous one only issued, and update_preempt_copies() counts it when its copy lands. // The notice goes with the save, not the cell release: preempt_save() has already detached the slot, so a release-time notice would leave the copy's silence unexplained. bool preempt_park(server_slot & slot, int64_t t_start) { @@ -5732,7 +5704,7 @@ std::unique_ptr server_routes::handle_completions_impl( task.params.oaicompat_model = meta->model_name; // [TAG_EXACT_CONCURRENCY] exact mode gives a page to a single sequence, so refuse an n_cmpl > 1 child here, where it becomes a 400 rather than at seq_cp - if (task.params.n_cmpl > 1 && server_exact_concurrency()) { + if (task.params.n_cmpl > 1 && common_exact_concurrency()) { throw std::runtime_error( "n > 1 is not supported while LLAMA_EXACT_CONCURRENCY is set: each " "completion needs its own sequence, and in exact mode a KV page belongs " From bb510f55612b36d4b987274d91f5889e97503974 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 09:03:41 +0000 Subject: [PATCH 80/81] server : drive a slot's two transfers through one helper The target's transfer and the draft's are always driven together, so the five places that spelled out "the target, and the draft if there is one" go through preempt_sum() for a figure and preempt_each() for a call. The parked slot that cannot fit an empty pool is found with std::find_if instead of a sentinel. --- tools/server/server-context.cpp | 81 ++++++++++++++------------------- 1 file changed, 34 insertions(+), 47 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 01a4bda7914..b71567270b6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -366,22 +366,33 @@ struct server_slot { return (bool) preempt_cpy_tgt; } - int64_t preempt_sync_us() const { + // the target's transfer and the draft's are always driven together, so a figure is the sum over both and a call is made on both + template + auto preempt_sum(F f) const -> decltype(f(preempt_cpy_tgt.get())) { if (!preempt_is_async()) { return 0; } - return llama_state_seq_copy_sync_us(preempt_cpy_tgt.get()) + - (preempt_cpy_dft ? llama_state_seq_copy_sync_us(preempt_cpy_dft.get()) : 0); + return f(preempt_cpy_tgt.get()) + (preempt_cpy_dft ? f(preempt_cpy_dft.get()) : 0); } - size_t preempt_n_copies() const { - if (!preempt_is_async()) { - return 0; + template + void preempt_each(F f) const { + if (preempt_cpy_tgt) { + f(preempt_cpy_tgt.get()); } - return llama_state_seq_copy_n_copies(preempt_cpy_tgt.get()) + - (preempt_cpy_dft ? llama_state_seq_copy_n_copies(preempt_cpy_dft.get()) : 0); + if (preempt_cpy_dft) { + f(preempt_cpy_dft.get()); + } + } + + int64_t preempt_sync_us() const { + return preempt_sum(llama_state_seq_copy_sync_us); + } + + size_t preempt_n_copies() const { + return preempt_sum(llama_state_seq_copy_n_copies); } // [TAG_PREEMPT_ASYNC] a copy is running: the slot must not be scheduled but still owns cells, so it is neither running nor parked @@ -400,24 +411,14 @@ struct server_slot { bool preempt_rotation_refused = false; // this park has logged a rotation refused for budget size_t preempt_state_size() const { - if (preempt_is_async()) { - // the capacity, not the live size: the pinned buffers are kept between parks, so --preempt-ram has to bound what is held - return llama_state_seq_copy_buf_capacity(preempt_cpy_tgt.get()) + - (preempt_cpy_dft ? llama_state_seq_copy_buf_capacity(preempt_cpy_dft.get()) : 0); - } - - return preempt_state_tgt.size() + preempt_state_dft.size(); + // for a transfer the capacity, not the live size: the pinned buffers are kept between parks, so --preempt-ram has to bound what is held + return preempt_is_async() ? preempt_sum(llama_state_seq_copy_buf_capacity) + : preempt_state_tgt.size() + preempt_state_dft.size(); } void preempt_state_free() { - // wait for anything in flight first: release() is reached with a copy possibly still using the buffer - if (preempt_cpy_tgt) { - llama_state_seq_copy_buf_free(preempt_cpy_tgt.get()); - } - - if (preempt_cpy_dft) { - llama_state_seq_copy_buf_free(preempt_cpy_dft.get()); - } + // waits for anything in flight first: release() is reached with a copy possibly still using the buffer + preempt_each(llama_state_seq_copy_buf_free); preempt_state_tgt.clear(); preempt_state_tgt.shrink_to_fit(); @@ -426,13 +427,7 @@ struct server_slot { } void preempt_copy_wait() { - if (preempt_cpy_tgt) { - llama_state_seq_copy_wait(preempt_cpy_tgt.get()); - } - - if (preempt_cpy_dft) { - llama_state_seq_copy_wait(preempt_cpy_dft.get()); - } + preempt_each(llama_state_seq_copy_wait); } size_t preempt_state_required() const { @@ -3760,23 +3755,15 @@ struct server_context_impl { server_slot * best = nullptr; // a parked slot that would not fit an empty pool can never be restored, so report it as the single-conversation overflow and rescan without it - { - server_slot * impossible = nullptr; - - for (auto * slot : parked) { - if (preempt_n_need(*slot) > n_cells) { - impossible = slot; - break; - } - } - - if (impossible) { - SLT_WRN(*impossible, "parked sequence of %d tokens cannot fit the pool of %d cells even alone, failing it\n", - preempt_n_need(*impossible), n_cells); - send_error(*impossible, "Context size has been exceeded."); - impossible->release(); - continue; - } + const auto impossible = std::find_if(parked.begin(), parked.end(), + [this, n_cells](const server_slot * slot) { return preempt_n_need(*slot) > n_cells; }); + + if (impossible != parked.end()) { + SLT_WRN(**impossible, "parked sequence of %d tokens cannot fit the pool of %d cells even alone, failing it\n", + preempt_n_need(**impossible), n_cells); + send_error(**impossible, "Context size has been exceeded."); + (*impossible)->release(); + continue; } // room for the sequence and for the next step of everything running, the candidate included, or a resume immediately preempts somebody; with nobody resident an exact fit is let in From 09cecf4f0d639bcfa235aa23e3706e0039a57284 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 09:38:28 +0000 Subject: [PATCH 81/81] server, llama, ggml : fewer comments Drop comments that restate the code they sit next to, keeping only invariants, rationale and hazards. --- common/arg.cpp | 1 - common/common.cpp | 2 - ggml/src/ggml-cuda/ggml-cuda.cu | 3 -- ggml/src/ggml-cuda/mmvq.cu | 6 +-- src/llama-batch.cpp | 3 +- src/llama-batch.h | 1 - src/llama-context.cpp | 5 -- src/llama-context.h | 1 - src/llama-impl.cpp | 1 - src/llama-kv-cache.cpp | 12 ----- src/llama-memory.h | 3 +- tests/test-state-seq-copy.cpp | 2 - tools/server/server-common.h | 1 - tools/server/server-context.cpp | 50 ++----------------- tools/server/tests/unit/test_preempt.py | 8 --- .../server/tests/unit/test_preempt_notify.py | 2 - 16 files changed, 9 insertions(+), 92 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 633337da885..af32812b930 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1305,7 +1305,6 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e } params.lr.init(); - // [TAG_EXACT_CONCURRENCY] refuse a column bound that cannot cover a decode step before anything is loaded, rather than running with the guarantee quietly off if (!common_exact_concurrency_init(ctx_arg.params)) { ctx_arg.params = params_org; return false; diff --git a/common/common.cpp b/common/common.cpp index 9d8918f4b80..18477a91ffb 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1450,7 +1450,6 @@ bool common_exact_concurrency() { return enabled; } -// [TAG_EXACT_CONCURRENCY] int common_exact_decode_width(const common_params & params) { const int64_t n_slots = std::max(1, params.n_parallel); @@ -1462,7 +1461,6 @@ int common_exact_decode_width(const common_params & params) { return n_cols > INT32_MAX ? -1 : (int) n_cols; } -// [TAG_EXACT_CONCURRENCY] bool common_exact_concurrency_init(const common_params & params) { if (!common_exact_concurrency()) { return true; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index b41b704e784..b2bc1934795 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1790,7 +1790,6 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { } static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { - // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving it on would give a solo request a different code path from a batched one if (ggml_cuda_batch_invariant()) { return false; } @@ -1962,7 +1961,6 @@ static int64_t ggml_cuda_mul_mat_invariant_width( return 1; } -// recompute dst in slices of columns so each column sees the batch-of-one configuration; false when the batched launch already gives every column that value static bool ggml_cuda_mul_mat_split_columns( ggml_backend_cuda_context & ctx, int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { @@ -2085,7 +2083,6 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor GGML_ABORT("fatal error"); } -// [TAG_BATCH_INVARIANT] true when the policy computes this MUL_MAT_ID one token at a time static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) { if (!ggml_cuda_batch_invariant()) { return false; diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 032964198a3..1fe68cebfa4 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -592,8 +592,7 @@ static __global__ void mul_mat_vec_q( ggml_cuda_pdl_sync(); sample_dst = blockIdx.z; - // [TAG_BATCH_INVARIANT] with ids, a sample is a token: every token goes on the z axis of one single-column - // launch, so each (token, expert slot) block runs the single-token configuration + // [TAG_BATCH_INVARIANT] with ids, a sample is a token: every token goes on the z axis of one single-column launch, so each (token, expert slot) block runs the single-token configuration channel_x = ncols_dst == 1 && ids ? ids[sample_dst*ids_stride + channel_dst] : fastdiv(channel_dst, channel_ratio); channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; @@ -1282,8 +1281,7 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT( nb0 == ts_dst); GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); - // [TAG_BATCH_INVARIANT] a multi-token MUL_MAT_ID becomes one launch of the single-token configuration - // with the tokens on the sample axis, so the count is not bounded by the column templates + // [TAG_BATCH_INVARIANT] a multi-token MUL_MAT_ID becomes one launch of the single-token configuration with the tokens on the sample axis, so the count is not bounded by the column templates const bool tokens_as_samples = ids && ne2 > 1 && ggml_cuda_batch_invariant(); GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE || tokens_as_samples); diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 62912596ce9..561adccd530 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -572,8 +572,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } if (add) { - // [TAG_EXACT_CONCURRENCY] a set with more tokens left than a decode step carries is a prompt and gets an - // ubatch of its own; grouped sets need equal tokens left, or the expansion below changes their sum order + // [TAG_EXACT_CONCURRENCY] a set with more tokens left than a decode step carries is a prompt and gets a ubatch of its own; grouped sets need equal tokens left, or the expansion below changes their sum order if (isolate_seqs_above > 0) { uint32_t n_left = 0; diff --git a/src/llama-batch.h b/src/llama-batch.h index ddf05843d10..edb01045b10 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -111,7 +111,6 @@ class llama_batch_allocr { // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than n_tokens left to place, i.e. what remains of the batch holds a prompt bool has_seq_wider_than(uint32_t n_tokens) const; - // [TAG_EXACT_CONCURRENCY] true if some token carries more than one sequence id bool has_shared_tokens() const; // sequence-set-wise split - each ubatch contains a single sequence-set diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 322d968708f..a62d2dbd534 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -104,7 +104,6 @@ llama_context::llama_context( // [TAG_EXACT_CONCURRENCY] the widest decode step this context can build, reported so a backend that splits columns covers it; reported at the end of the constructor if (llama_exact_concurrency()) { - // an explicit column bound below this context's width would leave decodes batched above it, so the report refuses it if (!llama_exact_check_n_seq(cparams.n_seq_max)) { throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); } @@ -1215,7 +1214,6 @@ void llama_context::set_causal_attn(bool value) { return; } - // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a context with a cache keeps causal attention rather than asserting in the next graph if (!value && memory && llama_exact_concurrency()) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so causal attention cannot be turned off; the change is refused\n", __func__); return; @@ -4812,8 +4810,6 @@ size_t llama_state_seq_set_data_ext(llama_context * ctx, const uint8_t * src, si return ctx->state_seq_set_data(seq_id, src, size, flags); } -// [TAG_STATE_ASYNC] - llama_state_seq_copy * llama_state_seq_copy_init(llama_context * ctx) { return ctx->state_seq_copy_init(); } @@ -4851,7 +4847,6 @@ void llama_state_seq_copy_buf_free(llama_state_seq_copy * cpy) { } bool llama_state_seq_copy_buf_is_pinned(llama_state_seq_copy * cpy) { - // what was allocated, not what could be: a host buffer type is free to hand back ordinary memory, as CUDA does under GGML_CUDA_NO_PINNED return cpy->pinned; } diff --git a/src/llama-context.h b/src/llama-context.h index 496ddcb05a2..f44f505a05f 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -40,7 +40,6 @@ struct llama_memory_buffer { using llama_memory_buffers = std::map; -// [TAG_STATE_ASYNC] defined in llama-context.cpp struct llama_state_seq_copy; struct llama_context { diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index ee35b914d72..50037a46b84 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -174,7 +174,6 @@ std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i) { } } -// [TAG_EXACT_CONCURRENCY] bool llama_exact_concurrency() { static const bool enabled = []() { const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index e7511dfd39e..8bbed8264b2 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -168,7 +168,6 @@ llama_kv_cache::llama_kv_cache( GGML_ASSERT(kv_size % n_pad == 0); - // [TAG_EXACT_CONCURRENCY] all of these are reachable from the command line, so name the one that failed instead of aborting on a bare assert if (exact_pages) { const char * unsupported = nullptr; @@ -321,14 +320,12 @@ llama_kv_cache::llama_kv_cache( throw std::runtime_error("exact concurrency: unsupported attention head size"); } - // [TAG_EXACT_CONCURRENCY] the paged kernel has no soft-capped variant and would assert if (exact_pages && hparams.attn_soft_cap) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but this model soft-caps its attention logits (%.1f), " "which the paged attention kernel does not apply\n", __func__, hparams.f_attn_logit_softcapping); throw std::runtime_error("exact concurrency: attention soft cap is not supported"); } - // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) { LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, " "which has no paged attention: every layer must be offloaded to the CUDA backend " @@ -492,7 +489,6 @@ llama_kv_cache::llama_kv_cache( debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0; } -// [TAG_EXACT_CONCURRENCY] void llama_kv_cache::exact_pages_rebuild() const { const auto & cells = v_cells[0]; @@ -521,7 +517,6 @@ void llama_kv_cache::exact_pages_rebuild() const { exact_page_owner_dirty = false; } -// [TAG_EXACT_CONCURRENCY] void llama_kv_cache::exact_pages_sync() const { if (exact_page_owner_dirty) { exact_pages_rebuild(); @@ -542,7 +537,6 @@ void llama_kv_cache::exact_pages_sync() const { } } -// [TAG_EXACT_CONCURRENCY] void llama_kv_cache::exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos) { if (exact_page_owner_dirty || exact_page_owner.empty()) { return; @@ -750,7 +744,6 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) { return; } - // [TAG_EXACT_CONCURRENCY] as in seq_rm, this can empty pages exact_page_owner_dirty = true; GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); @@ -838,7 +831,6 @@ void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, in return; } - // [TAG_EXACT_CONCURRENCY] as in seq_add: dividing positions breaks the position/offset identity if (exact_pages && d != 1) { LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions " "(seq %d, d %d); ignoring the division\n", @@ -1207,7 +1199,6 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, } if (exact_pages) { - // ownership is maintained as cells are placed, so this reads one entry per page rather than scanning every cell const auto & cells = v_cells[0]; exact_pages_sync(); @@ -1425,7 +1416,6 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & cells.seq_add(idx, ubatch.seq_id[i][s]); } - // [TAG_EXACT_CONCURRENCY] the page this cell belongs to is now owned by its sequence if (exact_pages) { GGML_ASSERT(ubatch.n_seq_id[i] == 1); @@ -1544,7 +1534,6 @@ ggml_tensor * llama_kv_cache::build_input_pages(ggml_context * ctx, const llama_ void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const { GGML_ASSERT(exact_pages && dst->ne[1] == ubatch->n_tokens); - // [TAG_EXACT_CONCURRENCY] one entry per physical page, not one per cell exact_pages_sync(); std::map> pages; @@ -2636,7 +2625,6 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 } else { // whole KV cache restore - // [TAG_EXACT_CONCURRENCY] refused at the top of state_read(), before anything is read GGML_ASSERT(!exact_pages); if (cell_count > cells.size()) { diff --git a/src/llama-memory.h b/src/llama-memory.h index 8f85aac4130..61cd348f2dd 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -100,8 +100,7 @@ struct llama_memory_i { // getters virtual bool get_can_shift() const = 0; - // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit: 1 unless a mode allocates in - // larger blocks, when n tokens occupy round_up(n, granularity) cells. Not pure, so old modules inherit 1. + // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit: 1 unless a mode allocates in larger blocks, when n tokens occupy round_up(n, granularity) cells. Not pure, so old modules inherit 1. virtual uint32_t alloc_granularity() const { return 1; } // diff --git a/tests/test-state-seq-copy.cpp b/tests/test-state-seq-copy.cpp index 488a268ee07..bd3cd9cc30a 100644 --- a/tests/test-state-seq-copy.cpp +++ b/tests/test-state-seq-copy.cpp @@ -85,7 +85,6 @@ int main(int argc, char ** argv) { llama_state_seq_copy_buf_is_pinned(cpy) ? "pinned" : "pageable", llama_state_seq_copy_buf_can_pin(cpy) ? "pinned" : "pageable"); - // a size beyond the buffer the transfer owns is refused, on both directions CHECK(llama_state_seq_copy_get(cpy, size + 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); CHECK(llama_state_seq_copy_set(cpy, size + 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); @@ -114,7 +113,6 @@ int main(int argc, char ** argv) { llama_memory_seq_rm(llama_get_memory(ctx), seq_id, -1, -1); - // the restore side too: a buffer claimed one byte short is refused before a copy is posted CHECK(llama_state_seq_copy_set(cpy, size - 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); CHECK(llama_state_seq_copy_n_copies(cpy) == 0); CHECK(llama_state_seq_copy_done(cpy)); diff --git a/tools/server/server-common.h b/tools/server/server-common.h index f0cf76b8c50..423e9982f1f 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -467,7 +467,6 @@ struct server_metrics { uint64_t n_decode = 0; uint64_t n_busy_slots = 0; - // [TAG_PREEMPT] slots parked to make room in the unified KV pool, and put back uint64_t n_preempt = 0; uint64_t n_resume = 0; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b71567270b6..c624c2cee57 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -71,8 +71,6 @@ constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the res constexpr int64_t PREEMPT_KEEPALIVE_MS = 2000; // SSE keepalive period while a streaming slot is parked constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected -// [TAG_PREEMPT] LLAMA_SERVER_PREEMPT_RESUME: head (the default) resumes by park time and lets nobody pass a head that does not fit; pass restores most-preempted-first -// [TAG_PREEMPT] the SSE comment for a park or resume; prompt 0 keeps the bare form single-prompt clients match on static std::string preempt_notice_comment(const server_task_result_preempt_notice & notice) { std::string res = notice.parked ? ": preempted" : ": resumed"; @@ -353,7 +351,6 @@ struct server_slot { prompt.clear(); } - // [TAG_PREEMPT] state of a slot whose cells were taken back; the task, sampler, text and stream stay, so a resume is a memcpy slot_state state_before_preempt = SLOT_STATE_IDLE; std::vector preempt_state_tgt; std::vector preempt_state_dft; @@ -366,7 +363,6 @@ struct server_slot { return (bool) preempt_cpy_tgt; } - // the target's transfer and the draft's are always driven together, so a figure is the sum over both and a call is made on both template auto preempt_sum(F f) const -> decltype(f(preempt_cpy_tgt.get())) { if (!preempt_is_async()) { @@ -450,7 +446,6 @@ struct server_slot { (!preempt_cpy_dft || llama_state_seq_copy_done(preempt_cpy_dft.get())); } - // back in the state it was parked from, with a speculative context to match: the draft went out with the cells bool preempt_resumed() { n_preempt_fail = 0; @@ -690,7 +685,6 @@ struct server_slot { n_predict_max = -1; - // [TAG_PREEMPT] preempt_state_free(); state_before_preempt = SLOT_STATE_IDLE; n_preempt = 0; @@ -850,8 +844,7 @@ struct server_slot { t_last_used = ggml_time_us(); - // [TAG_PREEMPT] a parked slot's cells are already gone, so the mirror must not outlive them or the next task prefix-matches an empty cache - // [TAG_PREEMPT_ASYNC] wait for any copy first: its host buffer and its cells are about to be handed on + // [TAG_PREEMPT] [TAG_PREEMPT_ASYNC] a parked slot's cells are already gone, so the mirror must not outlive them or the next task prefix-matches an empty cache; wait for any copy first, its buffer and its cells are about to be handed on if (preempt_is_out()) { preempt_copy_wait(); preempt_state_free(); @@ -1630,7 +1623,6 @@ struct server_context_impl { } } - // [TAG_PREEMPT_ASYNC] the slots either all park through a transfer or none do { preempt_async_ok = !slots.empty(); @@ -1671,7 +1663,6 @@ struct server_context_impl { } } - // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates rather than assume a cell per token { preempt_alloc_granularity = (int32_t) std::max(1u, llama_memory_alloc_granularity(llama_get_memory(ctx_tgt))); @@ -2428,7 +2419,6 @@ struct server_context_impl { queue_results.send(std::move(res)); } - // [TAG_PREEMPT] tell a streaming client its slot was parked or restored; the HTTP layer sends an SSE comment, which a client that does not know about preemption never sees void send_preempt_notice(server_slot & slot, bool parked) { if (!slot.task || !slot.task->params.stream) { return; @@ -3183,8 +3173,6 @@ struct server_context_impl { }; #endif - // [TAG_PREEMPT] server-side request preemption - // LLAMA_SERVER_PREEMPT_EVERY=N: preempt every generating slot every N tokens, pressure or not, so the determinism test can blame any difference on the preemption int32_t preempt_test_every = 0; @@ -3210,7 +3198,6 @@ struct server_context_impl { // LLAMA_SERVER_PREEMPT_RESUME=head or pass, read at load, per context bool preempt_resume_head = true; - // a recurrent cache has no cell pool to run out of, so preemption is off for those models; a hybrid keeps its attention cache and stays on bool preempt_recurrent = false; bool preempt_batch_abandoned = false; @@ -3222,7 +3209,6 @@ struct server_context_impl { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } - // [TAG_PREEMPT_ASYNC] whether the kind of host memory the parks got has been reported; only knowable once a buffer exists bool preempt_ram_kind_logged = false; void preempt_log_ram_kind(const server_slot & slot) { @@ -3266,7 +3252,7 @@ struct server_context_impl { return res; } - // [TAG_PREEMPT_ASYNC] whether parking this slot stays under --preempt-ram; a restored slot keeps its pinned buffer, so idle capacity is given back largest first when a park does not fit + // [TAG_PREEMPT_ASYNC] a restored slot keeps its pinned buffer, so that idle capacity is given back largest first when a park does not fit under --preempt-ram void preempt_reclaim_idle_ram(size_t budget, size_t extra, const server_slot & keep) { for (;;) { if (preempt_ram_used() + extra <= budget) { @@ -3304,7 +3290,6 @@ struct server_context_impl { } } - // the --preempt-ram ceiling in bytes; the unlimited setting is a ceiling nothing reaches size_t preempt_ram_budget() const { return params_base.preempt_ram_mib < 0 ? SIZE_MAX : (size_t) params_base.preempt_ram_mib * 1024 * 1024; } @@ -3322,7 +3307,6 @@ struct server_context_impl { return preempt_ram_used() + extra <= budget; } - // [TAG_PREEMPT_ASYNC] over budget, a buffer held by a running slot would keep every other slot from being parked at all void preempt_trim_ram(server_slot & slot) { if (preempt_ram_used() > preempt_ram_budget() && slot.preempt_state_size() > 0) { SLT_INF(slot, "%.1f MiB of parked RAM returned: the pool is over its budget\n", slot.preempt_state_size() / (1024.0 * 1024.0)); @@ -3381,8 +3365,6 @@ struct server_context_impl { // [TAG_PREEMPT_ASYNC] deliberately not skipped: a slot with a copy in flight holds cells either way, and skipping it would hand the same cells out twice - // [TAG_EXACT_CONCURRENCY] the tail page is charged in full: it cannot be given to anybody else - if (slot.state == SLOT_STATE_WAIT_OTHER) { res += preempt_n_cells(slot.prompt.n_tokens()); continue; @@ -3448,7 +3430,6 @@ struct server_context_impl { case SLOT_STATE_STARTED: case SLOT_STATE_PROCESSING_PROMPT: { - // preempt_n_retained() reads the live state, so a restoring slot is charged from what it holds const int32_t n_have = preempt_n_retained(slot); const int32_t n_left = slot.task ? slot.task->n_tokens() - n_have : 0; @@ -3463,7 +3444,6 @@ struct server_context_impl { return res + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1)); } - // keep the slot furthest along, it is the closest to giving its cells back; among the rest prefer one not preempted PREEMPT_N_STARVED times, then the smallest // [TAG_PREEMPT] trim a just-started slot to the prefix it keeps first, or it is copied out, charged and sized by the previous request's prompt bool preempt_normalize_started_all() { bool res = false; @@ -3592,9 +3572,7 @@ struct server_context_impl { return a.prompt.n_tokens() < b.prompt.n_tokens(); } - // called once per update_slots(), before the batch is built: every slot is then at a token boundary with no draft in flight, so it can be removed whole - // [TAG_PREEMPT] park a slot: a synchronous park is finished here, an asynchronous one only issued, and update_preempt_copies() counts it when its copy lands. - // The notice goes with the save, not the cell release: preempt_save() has already detached the slot, so a release-time notice would leave the copy's silence unexplained. + // [TAG_PREEMPT] park a slot: a synchronous park is finished here, an asynchronous one only issued, and update_preempt_copies() counts it when its copy lands. The notice goes with the save, not the cell release: preempt_save() has already detached the slot, so a release-time notice would leave the copy's silence unexplained. bool preempt_park(server_slot & slot, int64_t t_start) { slot.t_preempt_copy_us = t_start; @@ -3613,7 +3591,6 @@ struct server_context_impl { return true; } - // [TAG_PREEMPT_ASYNC] a park whose copy has landed; `note` says how it was waited for, if it was void preempt_parked(server_slot & slot, const char * note) { metrics.n_preempt++; @@ -3687,7 +3664,6 @@ struct server_context_impl { return false; } - // wait for one outstanding park, the last thing tried before giving up on room: the decode then waits for the copy exactly as the synchronous path did bool preempt_wait_in_flight() { for (auto & slot : slots) { if (slot.state != SLOT_STATE_PREEMPTING) { @@ -3754,7 +3730,6 @@ struct server_context_impl { server_slot * best = nullptr; - // a parked slot that would not fit an empty pool can never be restored, so report it as the single-conversation overflow and rescan without it const auto impossible = std::find_if(parked.begin(), parked.end(), [this, n_cells](const server_slot * slot) { return preempt_n_need(*slot) > n_cells; }); @@ -3882,7 +3857,6 @@ struct server_context_impl { break; } - // [TAG_PREEMPT_ASYNC] with a transfer the copy has only been issued; update_preempt_copies() counts and logs it when it lands if (best->state == SLOT_STATE_RESTORING) { SLT_WRN(*best, "resumed after %.2f s: %d tokens, restore issued in %.2f ms (%zu transfers, %.2f ms sync), kv %d/%d, preemptions %d\n", (ggml_time_us() - best->t_preempt_us) / 1e6, @@ -3960,7 +3934,6 @@ struct server_context_impl { break; // could not park it; the existing retry ladder is still behind us } - // [TAG_PREEMPT_ASYNC] the copy has only been issued and the cells are still the victim's, so nothing further can be decided about the pool this iteration if (victim->state == SLOT_STATE_PREEMPTING) { SLT_WRN(*victim, "preempted: %d cells, park issued in %.2f ms (%zu transfers, %.2f ms sync), %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, @@ -3970,8 +3943,7 @@ struct server_context_impl { preempt_kv_used(), n_cells, n_used, victim->n_preempt); - // [TAG_PREEMPT_ASYNC] short of the lookahead only, the step still fits and leaving is the point; out of room - // for it the cells are held until the copy lands, so the retry ladder ends every request instead of waiting + // [TAG_PREEMPT_ASYNC] short of the lookahead only, the step still fits and leaving is the point; out of room for it the cells are held until the copy lands, so the retry ladder ends every request instead of waiting if (n_used + preempt_n_margin() > n_cells) { continue; } @@ -3989,7 +3961,6 @@ struct server_context_impl { } // the checks a request has to pass before its prompt is processed; true when it is rejected. An empty prompt is not here: it is a final response, not an error. - // [TAG_PREEMPT] the planner asks the same question before parking a started slot: a notice opens the stream, and a rejected request would get 200 plus an in-stream error bool slot_prompt_rejected(const server_slot & slot, std::string & msg, error_type & type) const { if (!slot.task) { return false; @@ -4136,7 +4107,6 @@ struct server_context_impl { #endif if (preempt_batch_abandoned) { - // [TAG_PREEMPT] the rest of this batch never ran; the next pass rebuilds it preempt_batch_abandoned = false; break; } @@ -4170,7 +4140,6 @@ struct server_context_impl { // apply context-shift if needed // TODO: simplify and improve - // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift void pre_decode_shift() { iterate(slots, [&](server_slot & slot) { if (slot.state == SLOT_STATE_GENERATING && slot.prompt.n_tokens() + 1 >= slot.n_ctx) { @@ -4510,7 +4479,6 @@ struct server_context_impl { slot.mem.seq_rm (slot.id, head_p, head_c); slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift); - // [TAG_PREEMPT_ASYNC] applied in place inside the next llama_decode, like a context shift preempt_shift_pending = true; for (size_t i = 0; i < n_match; i++) { @@ -4886,7 +4854,6 @@ struct server_context_impl { } } - // [TAG_PREEMPT] the retry ladder ran out: give the batch up, rewind every resident to the token boundary the cache is at and park the smallest. Multimodal keeps the old path. // [TAG_PREEMPT_ASYNC] whether a park can happen at all and go asynchronously bool preempt_async_possible() const { return params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0 && @@ -4897,6 +4864,7 @@ struct server_context_impl { return params_base.kv_unified && params_base.preempt_ram_mib != 0 && !preempt_recurrent && slots.size() >= 2 && llama_get_memory(ctx_tgt); } + // [TAG_PREEMPT] the retry ladder ran out: give the batch up, rewind every resident to the token boundary the cache is at and park the smallest. Multimodal keeps the old path. bool preempt_last_resort(int32_t off) { if (!preempt_last_resort_possible()) { return false; @@ -4905,7 +4873,6 @@ struct server_context_impl { int32_t n_running = 0; for (auto & slot : slots) { - // [TAG_PREEMPT_ASYNC] a slot in transfer is not in this batch either way if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED || slot.preempt_in_flight()) { continue; } @@ -4981,7 +4948,6 @@ struct server_context_impl { return true; } - // [TAG_PREEMPT] whether a slot in the batch has its sampled token and a draft in it bool batch_has_spec_groups() const { for (const auto & slot : slots) { if (!slot.spec_i_batch.empty()) { @@ -5056,7 +5022,6 @@ struct server_context_impl { } if (n_batch == 1 && ret == 1) { - // [TAG_PREEMPT] park instead of ending everyone, when there is a budget to park into if (preempt_last_resort(off)) { preempt_batch_abandoned = true; return true; @@ -5082,7 +5047,6 @@ struct server_context_impl { SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret); for (auto & slot : slots) { - // [TAG_PREEMPT] a parked slot is not part of this failure and comes back when there is room if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED && !slot.preempt_in_flight()) { send_error(slot, err); slot.release(); @@ -5432,7 +5396,6 @@ struct server_context_impl { void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) { metrics.n_decode++; for (const auto & slot : slots) { - // [TAG_PREEMPT] a parked slot is processing but took no part in this decode if (slot.is_processing() && !slot.preempt_is_out()) { metrics.n_busy_slots++; } @@ -5756,7 +5719,6 @@ std::unique_ptr server_routes::handle_completions_impl( std::set parked_idx; // prompts of this request that are parked right now auto first_result = rd.next(req.should_stop); if (first_result != nullptr && dynamic_cast(first_result.get()) != nullptr) { - // [TAG_PREEMPT] the stream starts now, with the notice, so the parked keepalive runs through the wait instead of the client seeing nothing const auto * notice = static_cast(first_result.get()); preempt_prefix = preempt_notice_comment(*notice); if (notice->parked) { @@ -5795,7 +5757,6 @@ std::unique_ptr server_routes::handle_completions_impl( res->status = 200; res->content_type = "text/event-stream"; res->set_next([res_this = res.get(), res_type, sse_ping_interval, parked_idx](std::string & output) mutable -> bool { - // [TAG_PREEMPT] the keepalive runs while ANY prompt of the request is parked const bool parked = !parked_idx.empty(); static auto format_error = [](task_response_type res_type, const json & res_json) { @@ -5881,7 +5842,6 @@ std::unique_ptr server_routes::handle_completions_impl( SRV_DBG("%s", "error received during streaming, terminating stream\n"); return false; // terminate on error } else if (const auto * notice = dynamic_cast(result.get())) { - // [TAG_PREEMPT] an SSE comment: invisible to clients that do not know about preemption if (notice->parked) { parked_idx.insert(notice->index); } else { diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 72aeb44f693..4e7aa6d851c 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -96,7 +96,6 @@ def _assert_completed(results, n_predict: int, whole: bool = False): def test_forced_preemption_does_not_change_the_output(): - # park and restore the only running slot every 8 tokens: the batch shape is the same at every step, so any difference in the output is the preemption's fault _start(n_ctx=512) reference = _complete(64) assert reference.status_code == 200 @@ -203,7 +202,6 @@ def test_two_prompts_that_overflow_the_pool_together_both_finish(): def test_a_generating_slot_and_a_large_prompt_both_finish(): - # a long generation meets a large prompt arriving beside it: the prompt is admitted chunk by chunk, whoever is smaller is parked, and both finish log = _start(n_ctx=256) prompt_b, n_b = _prompt_of_about(150, "Charlie") @@ -228,7 +226,6 @@ def test_a_generating_slot_and_a_large_prompt_both_finish(): def test_preempt_ram_zero_disables_preemption(): - # --preempt-ram 0 switches back to the old behaviour: nothing is parked and the KV-full path ends the requests os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" log = _start(n_ctx=256) @@ -289,7 +286,6 @@ def _require_async(text: str): def test_async_preemption_does_not_change_the_output(): - # the synchronous determinism question asked of the asynchronous path: with one request the batch shape is fixed, so a continuation that is not byte-identical is the transfer's fault text = _start_async(n_ctx=512, n_gpu_layer=99) _require_async(text) @@ -389,7 +385,6 @@ def test_cancel_while_a_copy_is_in_flight_frees_the_slot(): def test_no_preempt_async_falls_back_to_the_synchronous_path(): - # The flag has to really switch it off, so that the two can be compared on one binary. os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "0" os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" log = _start(n_ctx=512, n_gpu_layer=99) @@ -471,7 +466,6 @@ def test_the_last_resort_parks_instead_of_ending_everyone(): def test_the_last_resort_works_with_an_unlimited_budget(): - # --preempt-ram -1 is the documented unlimited setting and must enable the last resort too os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" os.environ["LLAMA_ARG_PREEMPT_RAM"] = "-1" log = _start(n_ctx=256) @@ -487,7 +481,6 @@ def test_the_last_resort_works_with_an_unlimited_budget(): def test_the_last_resort_rewinds_a_prompt_in_flight(): - # the failed chunk comes back off the slot's tokens and is processed again after the resume, neither skipped nor fed twice os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" log = _start(n_ctx=256) @@ -605,7 +598,6 @@ def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_w def test_a_recurrent_model_is_served_without_preemption(): - # a recurrent cache holds one state per sequence whatever its length, so preemption is off for such a model and the forced-park knob parks nothing path = os.environ.get("LLAMA_SERVER_TEST_RECURRENT_MODEL") if path: server.model_file = path diff --git a/tools/server/tests/unit/test_preempt_notify.py b/tools/server/tests/unit/test_preempt_notify.py index 485a6bbc336..fa6a956e878 100644 --- a/tools/server/tests/unit/test_preempt_notify.py +++ b/tools/server/tests/unit/test_preempt_notify.py @@ -172,7 +172,6 @@ def test_two_overflowing_streams_both_finish_and_the_parked_one_says_so(): def test_a_stream_parked_before_its_first_token_starts_with_the_notice(): - # A request parked while still processing its prompt has no token to send yet, so the response starts with the notice instead of a silent connection. # n_batch: the whole prompt in one batch, so the planner sees its size at once _start(n_ctx=512, n_batch=512) url = f"http://{server.server_host}:{server.server_port}/completion" @@ -229,7 +228,6 @@ def test_a_resident_rotated_out_for_a_parked_head_is_told_so(): def test_an_oversized_prompt_is_errored_instead_of_parked(): # A slot just given a task has not passed the prompt checks yet, and a notice opens the stream, so parking it would turn a plain error response into 200 plus an in-stream one. os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" - # n_batch: the whole prompt in one batch, so the planner sees its size at once _start(n_ctx=512, n_batch=512) url = f"http://{server.server_host}:{server.server_port}/completion" resident = _completion_payload(390) | {"prompt": " ".join([_PROMPT_A] * 6)}