cuda: opt-in exact concurrency, byte-identical output regardless of neighbours or parks - #194
cuda: opt-in exact concurrency, byte-identical output regardless of neighbours or parks#194danielhanchen wants to merge 47 commits into
Conversation
…atch 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.
…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.
…ENCY 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.
|
Status on a mixture-of-experts model, Qwen3.6-35B-A3B, same harness, same settings, seed 0, temperature 0, 512 tokens. The mode does not hold there yet. Mode on diverges from the solo run at token 54 in every round, with and without 98 forced parks (mode off diverges at 47). The node probe over one decode step, sequence 0 alone against with three neighbours, is byte-identical for nodes 0 to 3249, which covers every layer through 33 and layer 34's routing, gate and up projections, and first differs at node 3250, Cost on that model is the same as on the 4B, about 9 percent solo decode, within 1 percent on prompt, aggregate and wall. A follow-up commit extending the batch-invariant policy to |
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_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.
|
Two more commits and the mode now holds on the mixture-of-experts model as well. Qwen3.6-35B-A3B, same harness, seed 0, temperature 0, 512 tokens, P0 against its solo run, three rounds per cell: speculation off identical, speculation off with 98 forced parks identical, draft-mtp with two drafts identical, and the three solo references are byte-identical to each other. At the library level 512 greedy decode steps of sequence 0 alone against inside a four-wide batch never differ in a single logit, on the 35B and on the 4B. What the earlier divergence was, corrected. The three differing
A second defect only became visible once that was fixed, and no node probe could have found it: registering an eval callback disables CUDA op fusion, so both sides of every probe took the unfused path. The fused top-k routing kernel is only admitted at one row (its memory-range exception is sound only there), so a request decoding alone computed its routing weights with the fused kernel and the same request next to three others computed them with the unfused chain. Cost on the 35B, five interleaved pairs on an uncontended GPU: solo decode 198 to 167 tok/s (15.6 percent, of which about two thirds is the unfused routing, which also applies to a solo decode), four-chat aggregate 440 to 360 (18.3 percent, the four-way split of 123 expert matmuls per step), prompt within 3 percent. On the dense 4B nothing changed. Limitations added: |
…y 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.
|
One more commit, Measured on the 35B, mode on, P0 against its solo run, three rounds per cell: draft-mtp with two drafts identical at 16, the same with 98 forced parks identical, and as a stress a four-draft verify batch of up to 20 columns identical both at a bound of 8 (left unsplit) and at 32 (split). On the 4B the forced-park MTP cell is identical at 16. Cost of 16 against 8, interleaved pairs: within 0.6 percent on every metric with speculation on and flat with it off, inside the run-to-run spread. The default is not self-adjusting: more than four slots, or more than two drafts, exceeds 16 and the bound has to be raised by hand, since the CUDA backend has no view of |
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)
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.
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.
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.
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
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.
…ve to A llama-server built with LLAMA_EXACT_CONCURRENCY set (unslothai/llama.cpp#194) gives a chat the same generated tokens whether it decodes alone, beside three other chats in one unified KV cache, or across a park and restore that moved its cells. There is no command-line flag for it, no --help entry and nothing in /props: the whole interface is an environment variable. So the only way to ask for it today is to set a llama.cpp variable on the Studio process, which is not discoverable, is not per load, and hands the mode to every child whether or not that load wanted it. UNSLOTH_LLAMA_EXACT_CONCURRENCY=auto|off|on is the switch, shaped like UNSLOTH_LLAMA_PREEMPT_MODE next to it, with the same three values also on the load request (exact_concurrency) and in a stored setting the GUI can write through GET/PUT /api/settings/exact-concurrency. Highest wins: the environment, then the request, then the store. Default off, because the mode costs about 9 percent of solo decode on a dense model and more on a mixture of experts, and buys nothing for a chat that never shares its cache. An inherited LLAMA_EXACT_CONCURRENCY is read as the default rather than ignored, so whoever set it before this existed keeps it; an explicit off takes it back out of the child environment, since off is the one answer that has to be obeyed exactly. The launch sets the variable on the child rather than relying on inheritance, and squares Studio's own line with the mode: --parallel 1 skips --kv-unified, which the paged KV pool needs whether or not anything else is decoding, so exact mode adds it. Studio never emits --cache-reuse and already emits --no-context-shift; a contradiction in the user's extra args (--cache-reuse, --context-shift, --no-kv-offload, a quantized KV cache, flash attention off) is named in a load warning instead of arriving as an error about a launch line the user did not compose. Detection is by trying, because there is nothing to ask. Under auto, a child that dies naming the mode is relaunched once without the variable, ahead of every other rung in the ladder since the server said which of its own requirements it could not meet; the attempt bound goes from three to four so the ROCm correction and the fit recovery do not lose their slot to it. Under on the refusal is the answer, and it is classified into a message that names the setting and what the mode requires. The --flash-attn off respawn drops the variable under auto for the same reason. What is reported is read off the argv and environment that actually launched, so a respawn that took away flash attention or the unified cache reads as unavailable rather than still claiming the guarantee. The load and status responses gain exact_concurrency (on/off/unavailable) and requested_exact_concurrency, the duplicate-load check treats a different setting as a reload since the mode cannot change without a new child, the preemption snapshot carries it, and the "llama preemption armed" line gains exact=. The one thing none of this can do is tell a build that predates #194, which ignores the variable and starts perfectly, from one that granted the mode: that reports on, and the tests record it as the limit it is. 103 new tests over the resolution order, the child environment, the launch line, the refusal detection and the auto fallback, the reported state and the armed line; 2278 pass across the preemption, start-failure and load-mode suites.
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.
…he 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.
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.
|
Measured the prompt serialisation limitation, since the description names it as the one known cost of the isolation rule. Setup: Qwen3.5-4B UD-Q4_K_XL, Four prompts of P tokens arriving at once, median of 3 repeats:
The cost is per extra ubatch, not per prompt token. A prompt of at least Mixed case, three chats mid-decode when a 3500-token prompt arrives: the decoders emit the same 5, 6, 6 tokens during the prefill in both modes, so decoders are not split apart, and the longer stall (p50 gap 151 ms off, 184 ms on) is entirely the slower prefill. The GPU was shared with another job, so the absolute numbers are not clean; the paired ratios are. The crossover is |
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.
|
Nine more commits, and the review items are closed. The 37 to 42 percent concurrency cost measured on
The last two commits are the deferred Medium row of the review, the one that said exact mode and #184's preemption planner cannot both be right about how full the pool is. They are not. Four chats with 1000-token prompts generating 2048 tokens each at
After, same configuration, three rounds:
Exact mode off on the same binary is 4 of 4 in all three rounds with 3, 2 and 2 parks, and its planner still logs the token figures it always did ( Server harness Two things this does not answer. The live cell above is not a cost measurement: the three configurations ran back to back rather than interleaved on a GPU shared with another tenant, and mode-off aggregate swung from 94 to 151 tok/s across its own three rounds, so the only cost figures worth quoting are the interleaved ones above. And, found while checking whether the harness model could reach the page case: on a model whose head size the CUDA paged kernel does not support, |
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.
|
Pushed c6c3cb6, which takes the speculative case from 0.43 to 0.74 of the mode-off rate. The 0.45 measured earlier was not the column policy. Slicing the column split into batch-of-one-equivalent slices (four columns on this GPU's MMVQ table, so three launches for a twelve-column batch instead of twelve) changed nothing on its own, because no matmul was ever twelve wide: the batch splitter isolated every sequence set with more than one token left to place as a prompt, and a slot's three-token verify batch is such a set, so with two MTP drafts each decode step ran the whole graph once per slot. The splitter now isolates by width. Exactness, three rounds per cell, P0 against its own solo run, all byte-identical:
Full pool, four chats of 937 prompt plus 2048 generated tokens at
Speculation off, the same shape: 0.92, unchanged from before the commit. The |
…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.
|
Pushed 379ca5d, which removes the per-token cost of the mode on mixture-of-experts models. Under the knob a Single-op probe, token 0's slice for 2 to 12 tokens against its own single-token run, knob on, neighbours routed to token 0's experts: 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 before the knob): 0 differing elements in all 99 comparisons. Qwen3.6-35B-A3B, four chats, P0 against its own solo run, three rounds per cell, exact on and off interleaved on a shared GPU:
Exact mode on the 35B now runs at the mode-off rate within the noise of the shared GPU, with and without speculation; the per-token re-entry had left it at about 0.7 with two drafts. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
… 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.
…o HEAD # Conflicts: # tools/server/server-context.cpp
…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.
…e 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.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…ion, 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.
|
@codex review |
|
Codex Review: Didn't find any major issues. Another round soon, please! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
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.
…o HEAD # Conflicts: # tools/server/server-context.cpp
1 participant
Stacked on #184; the commits after it are new. Opt-in, default off.
Summary
LLAMA_EXACT_CONCURRENCY=1makes a sequence's generated tokens byte-identical whether it decodes alone, next to other sequences in one unified KV cache, or across a park and restore that moved its cells. Today's output under concurrency differs from a solo run of the same prompt even with no preemption, for two reasons this branch removes: kernels change algorithm with batch width, and a sequence's attention rounds differently depending on where its cells sit next to its neighbours.Policy
GGML_CUDA_BATCH_INVARIANT: a CUDA row's result does not depend on how many other columns share the launch. The dispatcher pins the single-column configuration ofMUL_MAT(MMVF against cuBLAS on F32, MMVQ, MMQ) andFLASH_ATTN_EXT(vector against MMA kernel,parallel_blocks), bounded byGGML_CUDA_BATCH_INVARIANT_MAX_COLS.n_parallel * (1 + n_draft)(16 for the default four slots with MTP drafts), so a speculative verify batch is covered by construction. An explicitGGML_CUDA_BATCH_INVARIANT_MAX_COLSbelow the derived bound is refused at start rather than silently narrowing the guarantee.MUL_MAT_IDtoken is computed with the single-token configuration, in one launch that puts the tokens on the MMVQ kernel's sample axis and reads the expert index per sample, so each (token, expert slot) block runs the single-token instructions on the same data. The fused top-k routing kernel is left unfused, since it is admitted only at one row and would otherwise route a solo decode with a different algorithm than a concurrent one.n_seq_maxtimes the per-sequence width, never lowered), so a program using the library directly is covered without knowing there is a bound; a caller building wider steps reports the width itself throughllama_set_exact_decode_width(). What a context reports is its sequence count; the width is recomputed from the widest count seen whenever the per-sequence token figure changes, so a context created under a narrower figure keeps a bound that covers it. Both setters return false and change nothing when an explicitGGML_CUDA_BATCH_INVARIANT_MAX_COLScannot cover the width, since that bound wins in the backend; the context constructor and the server setup treat that as an error. Reports are serialised and the backend keeps the widest figure whatever order they arrive in; the per-sequence token figure is never lowered either, so a narrower context set up later cannot turn an existing context's verify steps into prompts. The three figures change under one lock, and the setup runs at the front of the init result's constructor, before any context exists.n > 1completions are refused with a 400 rather than aborting in the cache. A context created non-causal over a KV cache is refused at creation rather than asserting on its first graph. A model that soft-caps its attention logits is refused at load too, since the paged kernel has no soft-capped variant.split_equalon whenever the mode is on, not only while a prompt is in the batch: 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. The plain KV cache keeps prompt-only isolation, since its kernels are column invariant. An isolated ubatch takes only sets that finish in it, so the expansion never cuts them part way either.llama_set_exact_decode_tokens()carries one plus the draft length, so a speculative verify batch stays grouped with the other slots' and is not run once per sequence, while a prompt still gets a ubatch of its own.llama_memory_alloc_granularity()tells the server how many cells one allocation takes (256 under exact mode, 1 otherwise), so the preemption planner from server: preempt a slot instead of ending every conversation when the KV pool fills #184 counts pages instead of tokens when deciding whether a sequence fits. The planner charges every slot's tail page in full, reserves a page only when a step crosses a page boundary, and rounds the margin; at a granularity of 1 every figure is the parent's, held bystatic_assert. Without it the planner saw room in tokens thatfind_slotcould not find in pages, never parked anybody, and the retry ladder ended every request withContext size has been exceeded.Results
The shape that used to fail, four chats of 937 prompt plus 2048 generated tokens at
-c 8192with MTP drafts, three rounds each, on one binary with the planner switched between token and page arithmetic through the test knob:Context size has been exceededQwen3.5-4B,
--parallel 4 --kv-unified -c 8192 --flash-attn on, seed 0, temperature 0, top_k 1, 512 generated tokens, P0 against its own solo run, three rounds per cell, independently rebuilt and reproduced:-c 16384Mode off on the same binary diverges at token 118, so the harness sees. The same three cells (spec off, spec off with 98 parks, MTP) were re-run after the last commit against the strictest reference, an unparked non-speculative solo run, and read identical in every round.
Qwen3.6-35B-A3B (mixture of experts, Unsloth dynamic quant), same settings, three rounds per cell:
The three solo references are byte-identical to each other, and 512 lockstep decode steps of one sequence alone against inside a four-wide batch never differ in a single logit, on either model. At the library level a probe over all 1831 graph nodes of a decode step finds 0 differing for a sequence alone against inside a four-sequence batch, and 0 for a parked and restored sequence.
Exactness
Every solo reference across spec off, MTP, forced parks and 16384 context is byte-identical to every other, and every concurrent run matches them. Exact mode's output differs from mode-off output at one near-tied greedy step, as any bit-level change to a canonical computation would; the claim is neighbour independence, not equality with the unpinned kernels.
Cost when it does not fire
None. Off by default: no page tables, no reserved pages, no extra kernels, no ubatch restriction. The op tests and the four-chat and solo rates with the mode off reproduce the parent commit's.
Cost when it fires
Three interleaved pairs per row on one GPU, medians, ratios to the mode off:
-c 8192, pool full and preemption running-c 65536, no preemption-c 8192, speculation off, three interleaved pairsOf the 8.5 percent four-chat loss on the 4B, the kernel policy on its own (
GGML_CUDA_BATCH_INVARIANT=2with no paged pool) costs more than the whole mode, 0.791 of off, because the paged pool lets attention use the single-column kernels without the split-K combine.Prompt serialisation, the one structural cost of the ubatch rule, was measured directly. Four prompts of P tokens arriving at once on the 4B, default
n_batch 2048 / n_ubatch 512:The cost is per extra ubatch, not per prompt token: a prompt of at least
n_ubatchtokens fills a ubatch by itself, so above that the two modes need the same ubatch count and exact mode is not slower to first token. With three chats mid-decode and a 3500-token prompt arriving, the decoders emit the same number of tokens during the prefill in both modes, so decoders are not split apart; the stall is longer only by the prefill's own slowdown.Tests
tools/server/tests/unit/test_preempt.py14 passed, including the page-planner case and two prompts within a page of the pool size.test-backend-ops -b CUDA0 -o MUL_MAT_ID,MUL_MAT,FLASH_ATTN_EXT,SET_ROWS5234 of 5234 with the mode on and off;MUL_MAT_IDalone 919 of 919 with the knob on and off, including new gate, up and down cases over Q4_K, Q5_K, Q6_K, Q8_0 and F16 at 1 to 17 tokens. New op cases: nonadjacent pages in reversed order, partial tail, query widths 1, 4 and 12; shared-weight projections F32, Q4_K, Q6_K, Q8_0 at widths 1, 17 and 307 with sequence axes 2 and 3. Server-level matrix above. The evidence harness refuses to certify a run that completed fewer rounds than requested.Limitations
FLASH_ATTN_EXT(CANN, Metal, OpenCL, SYCL, Vulkan, WebGPU, ExecuTorch, Hexagon, OpenVINO, RPC) refuses the op when the page table is present, and those rejections are compile-time only here.n > 1are refused in exact mode, and a refusal leaves the cache as it was: a whole-context restore is refused before a byte is read, sollama_state_set_datareturns 0 with every sequence intact. Per-sequence park and restore is supported and tested.static_assert. Without it the planner saw room in tokens thatfind_slotcould not find in pages, never parked anybody, and the retry ladder ended every request withContext size has been exceeded.n_ubatchpay up to 1.4x on time to first token at four chats; the crossover moves with-ub.Follow-up commits after review
A whole-context restore under the mode was refused inside the cache's metadata reader, after which the generic restore path cleared the live cache; it is refused at the top of the context's and the cache's readers now, before anything is read, and a probe on the 4B shows
llama_state_set_datareturning 0 with the sequence still at its position and decoding afterwards. The context reported its sequence count at the front of its constructor, so a construction that failed later (an unsupported cache layout, say) left a width behind that no context needed; the count is checked against the explicit bound early and reported at the end. DSpark runs on the DFlash implementation and turns causal attention off on its draft, so it is refused alongside DFlash. The decode width is computed in 64 bits and refused aboveINT32_MAX. The five backends above that still claimed the paged op refuse it. The probe'sPROBE_A_PERMkeeps prompt 0 on sequence 0, which its second phase compares against. One more from the next pass: the cache accepted a layer on any CUDA, ROCm or MUSA device by the registry name alone, and a build without the flash attention kernels, or a device and head shape they do not cover, would have had the scheduler hand the paged op to the CPU, which accepts the page table as the reference for the backend test and ignores it. 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 K/V types.