From f3ad79657aa2aec1545c179122dc60f9544c75af Mon Sep 17 00:00:00 2001 From: Shreya Shankar Date: Fri, 2 Oct 2026 22:55:45 +0000 Subject: [PATCH 1/5] Share the partner label and partner starts in the join minimum The minimum counted each pair's partner label, partner document, and answer cue in full. The label is the same for every pair of one anchor, and partners of one anchor that begin with the same tokens share those tokens in a prefix cache, so the minimum was above what unlimited KV needs. On LEP-6 the stock vLLM run computed fewer fresh tokens than the minimum. The label now counts once per anchor, after the frame, and the partner documents of one anchor count as a prefix trie. --- quail_b/minimum.py | 85 +++++++++++++++++++++++++++++++++++------- tests/test_minimum.py | 86 +++++++++++++++++++++++++++++++++++-------- 2 files changed, 142 insertions(+), 29 deletions(-) diff --git a/quail_b/minimum.py b/quail_b/minimum.py index 3508625..550ac47 100644 --- a/quail_b/minimum.py +++ b/quail_b/minimum.py @@ -8,10 +8,11 @@ sequences is computed once: each document's text once, and the questions and frames after one document once each, sharing the lead they have in common (an engine that rewinds KV to where two questions -diverge computes that lead once). A pair's suffix after its anchor -(label, partner document, answer cue) is computed once per pair, with -nothing shared between the pairs of one anchor: those tokens are one -request's own, so they are never regret. What an engine computed +diverge computes that lead once). The partner label after a frame is +the same for every pair of one anchor, so it is computed once per +anchor. Each pair's partner document and answer cue follow it; the +pairs of one anchor share the tokens their partner documents begin +with, as a prefix cache over whole requests does. What an engine computed beyond the minimum is its regret, whatever the cause: an evicted anchor computed again, a set scanned twice under two aliases, or a prompt prefix the documents share computed once per document. @@ -234,6 +235,61 @@ def __getitem__(self, document) -> np.ndarray: return self._tokens[document] +class _PartnerTrie: + """Prefix trie sizes of partner documents that one tail follows. + + The documents are sorted once. Two sorted neighbours of a subset + share the shortest common prefix of the sorted run between them, so + a subset's trie size is its length minus a range minimum per + neighbour pair. + + Args: + documents: The document tokens, fetched. + keys: The ((table, column), id) partner documents. + tail: Token ids after every partner document. + """ + + def __init__(self, documents, keys, tail): + keys = list(keys) + tail = _tokens(tail) + packed = [np.concatenate((documents[(*table_set, row_id)], tail)) + .astype(">u4").tobytes() for table_set, row_id in keys] + order = sorted(range(len(keys)), key=packed.__getitem__) + self.rank = {keys[index]: rank for rank, index in enumerate(order)} + self.length = np.array([len(packed[index]) // 4 for index in order], + dtype=np.int64) + shared = np.zeros(max(len(order) - 1, 0), dtype=np.int64) + for position, (earlier, later) in enumerate(zip(order, order[1:])): + left, right = packed[earlier], packed[later] + length = min(len(left), len(right)) + differs = (np.frombuffer(left, np.uint8, length) + != np.frombuffer(right, np.uint8, length)) + first = int(differs.argmax()) if differs.any() else length + shared[position] = first // 4 + self.levels = [shared] + width = 1 + while 2 * width <= len(shared): + last = self.levels[-1] + self.levels.append(np.minimum(last[:-width], last[width:])) + width *= 2 + + def size(self, members) -> int: + """Return the prefix trie size of the member documents with the tail.""" + ranks = np.sort(np.fromiter((self.rank[key] for key in members), + dtype=np.int64, count=len(members))) + total = int(self.length[ranks].sum()) + if len(ranks) < 2: + return total + low, high = ranks[:-1], ranks[1:] + level = np.floor(np.log2(high - low)).astype(np.int64) + for step in np.unique(level).tolist(): + pick = level == step + table = self.levels[step] + total -= int(np.minimum( + table[low[pick]], table[high[pick] - (1 << step)]).sum()) + return total + + def _encoded(column) -> tuple[np.ndarray, pa.Array]: """Return a string id column as dictionary indices and the dictionary.""" column = pc.cast(column, pa.string()) @@ -333,7 +389,7 @@ def record(alias, row_id) -> _Document: (partner_set, row_id) for row_id in pc.take( partner_ids, pa.array(members)).to_pylist()) document = record(anchor, anchor_ids[int(anchors[start])].as_py()) - document.suffixes.add(group[0]) + document.suffixes.add(group[0] + group[1]) document.groups.setdefault(group, set()).add(key) documents.fetch([(*table_set, row_id) for table_set, row_id in records] @@ -345,20 +401,21 @@ def record(alias, row_id) -> _Document: for table_set, row_id in records) suffix_sizes: dict = {} partner_sizes: dict = {} + partner_tries: dict = {} + partners = frozenset().union(*members_of.values()) for document in records.values(): suffixes = frozenset(document.suffixes) if suffixes not in suffix_sizes: suffix_sizes[suffixes] = prefix_trie_size(suffixes) total += suffix_sizes[suffixes] - for (_, label, tail), keys in document.groups.items(): - keys = frozenset(keys) - if keys not in partner_sizes: - members = frozenset().union(*(members_of[key] for key in keys)) - partner_sizes[keys] = (len(members), sum( - len(documents[(*member_set, row_id)]) - for member_set, row_id in members)) - count, size = partner_sizes[keys] - total += (len(label) + len(tail)) * count + size + for (_, _, tail), keys in document.groups.items(): + key = (tail, frozenset(keys)) + if key not in partner_sizes: + if tail not in partner_tries: + partner_tries[tail] = _PartnerTrie(documents, partners, tail) + partner_sizes[key] = partner_tries[tail].size( + frozenset().union(*(members_of[member] for member in keys))) + total += partner_sizes[key] return total diff --git a/tests/test_minimum.py b/tests/test_minimum.py index 9751b13..00f7605 100644 --- a/tests/test_minimum.py +++ b/tests/test_minimum.py @@ -1,10 +1,14 @@ """CPU checks for the prefix trie and the minimum input tokens of a run.""" +import random + +import numpy as np import pyarrow as pa import pytest from quail_b.minimum import ( DocumentTokens, + _PartnerTrie, input_tokens, minimum_input_tokens, prefix_trie_size, @@ -83,13 +87,13 @@ def test_minimum_input_tokens_counts_each_document_prefix_once_and_pairs_apart() # the three documents share the preamble, and the first two share # "same start " (11 tokens) beyond it; every document gets the - # question once, the two anchors get the frame once, sharing the - # lead it has in common with the question; each pair gets its own - # label, partner, and tail + # question once, the two anchors get the frame and label once, + # sharing the lead they have in common with the question; each + # anchor's partners "aa" and "ab" share their first token pre = len(PRE) documents = 3 * pre + 14 + 14 + 5 - (pre + pre + 11) - anchored = len(QUESTION) + len(FRAME) - _lcp(QUESTION, FRAME) - pairs = 2 * len(LABEL) + 4 + 2 * len(TAIL) + anchored = prefix_trie_size([QUESTION, FRAME + LABEL]) + pairs = 2 * (2 + len(TAIL)) - 1 assert minimum == documents + len(QUESTION) + 2 * anchored + 2 * pairs full_inputs = ( 3 * (len(PRE) + len(QUESTION)) + 14 + 14 + 5 @@ -129,10 +133,9 @@ def test_minimum_input_tokens_counts_a_document_once_across_uses(): DocumentTokens(corpus, _encode)) documents = 2 * len(PRE) + 5 + 4 - len(PRE) - alpha = len(first) + len(second) + len(FRAME) - sum(( - _lcp(first, second), max(_lcp(FRAME, first), _lcp(FRAME, second)))) - beta = len(first) + len(FRAME) - _lcp(first, FRAME) - pairs = 2 * len(LABEL) + 9 + 2 * len(TAIL) + alpha = prefix_trie_size([first, second, FRAME + LABEL]) + beta = prefix_trie_size([first, FRAME + LABEL]) + pairs = 9 + 2 * len(TAIL) assert minimum == documents + alpha + beta + 2 * pairs full_inputs = ( 2 * (len(PRE) + len(first)) + 9 @@ -142,6 +145,59 @@ def test_minimum_input_tokens_counts_a_document_once_across_uses(): DocumentTokens(corpus, _encode)) == full_inputs +def test_minimum_input_tokens_of_a_join_is_the_trie_of_its_requests(): + spec = _spec( + "TEST-5", + "one join", + Join(Scan("docs", "d", "body"), Scan("notes", "n", "text"), + ("d", "n"), "{0} cites {1}"), + ) + corpus = { + "docs": pa.table({"id": ["r", "g"], "body": ["red apple", "green pear"]}), + "notes": pa.table({"id": ["1", "2", "3", "4"], + "text": ["the cat sat", "the cat ran", "a dog", + "the cat sat"]}), + } + pieces = validate_prompt_pieces(spec, { + "tokenizer": "test", "preamble": PRE, + "joins": [{"id": "join-1", "anchor": "d", "frame": FRAME, + "label": LABEL, "tail": TAIL}], + }) + pairs = [("r", "1"), ("r", "2"), ("r", "3"), ("r", "4"), + ("g", "1"), ("g", "3")] + join_answers = {"join-1": pa.table({ + "d": [anchor for anchor, _ in pairs], + "n": [partner for _, partner in pairs], + "answer": [True] * len(pairs)})} + minimum = minimum_input_tokens( + spec, pieces, {}, join_answers, DocumentTokens(corpus, _encode)) + + # with unlimited KV every distinct prefix of the requests is computed + # once: the label once per anchor, and the partners of one anchor + # sharing "the cat " and the whole of a repeated text + bodies = {"r": "red apple", "g": "green pear"} + texts = {"1": "the cat sat", "2": "the cat ran", "3": "a dog", + "4": "the cat sat"} + requests = [PRE + _ids(bodies[anchor]) + FRAME + LABEL + + _ids(texts[partner]) + TAIL for anchor, partner in pairs] + assert minimum == prefix_trie_size(requests) + + +def test_partner_trie_size_matches_the_trie_of_any_subset(): + rng = random.Random(0) + table = ("notes", "text") + texts = {str(i): [rng.choice((1, 2, 3)) for _ in range(rng.randint(0, 6))] + for i in range(60)} + documents = {(*table, i): np.asarray(t, dtype=np.uint32) + for i, t in texts.items()} + keys = [(table, i) for i in texts] + trie = _PartnerTrie(documents, keys, TAIL) + for _ in range(200): + members = rng.sample(keys, rng.randint(1, len(keys))) + assert trie.size(frozenset(members)) == prefix_trie_size( + [texts[i] + TAIL for _, i in members]) + + def _filter_spec(): return _spec( "TEST-0", @@ -313,12 +369,12 @@ def test_minimum_input_tokens_count_joined_rows_as_pairs_after_the_anchor(): minimum = minimum_input_tokens( spec, pieces, {}, join_answers, tokens, classify_answers) - # the anchor is computed once with both frames after it; the join's - # two pairs and the classified pair each get their own suffix - anchored = (len(FRAME) + len(classify_frame) - - _lcp(FRAME, classify_frame)) - joined = 2 * (len(LABEL) + len(TAIL)) + 4 - classified = len(LABEL) + len(CLASSIFY_TAIL) + 2 + # the anchor is computed once with both frames and their label after + # it; the join's partners "aa" and "ab" share their first token, and + # the classified pair gets its own partner and tail + anchored = prefix_trie_size([FRAME + LABEL, classify_frame + LABEL]) + joined = 2 * (2 + len(TAIL)) - 1 + classified = 2 + len(CLASSIFY_TAIL) assert minimum == len(PRE) + 5 + anchored + joined + classified From c40e99955a3a794f9853e47f6bd7daac90070153 Mon Sep 17 00:00:00 2001 From: Shreya Shankar Date: Fri, 2 Oct 2026 23:28:19 +0000 Subject: [PATCH 2/5] Describe the join minimum with the shared label and partner starts --- docs/reference.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/reference.md b/docs/reference.md index 1ef3b38..d7c172e 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -600,8 +600,10 @@ is evaluated document pairs, summed across all joins, divided by `runtime_s`. | Input token throughput | Input tokens divided by `runtime_s` | The minimum counts each distinct prompt prefix once, so a document's questions -share the document and any leading tokens they have in common. For joins, each -pair's label, partner document, and answer cue count once per pair. Recomputed +share the document and any leading tokens they have in common. For joins, the +partner label counts once per anchor, and each pair's partner document and +answer cue count once per pair, less the leading tokens shared with another +partner of the same anchor. Recomputed tokens are therefore the work a perfect prefix KV cache would have avoided. A classification's tail counts like a filter question: once per document, From bc854496d4aa4d29a9757caef355cb7ef326f92c Mon Sep 17 00:00:00 2001 From: Shreya Shankar Date: Fri, 2 Oct 2026 23:44:15 +0000 Subject: [PATCH 3/5] List what the minimum counts, one rule per bullet --- docs/reference.md | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/docs/reference.md b/docs/reference.md index d7c172e..6e14e94 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -599,12 +599,21 @@ is evaluated document pairs, summed across all joins, divided by `runtime_s`. | KV regret | Recomputed tokens divided by fresh tokens, as a percentage | | Input token throughput | Input tokens divided by `runtime_s` | -The minimum counts each distinct prompt prefix once, so a document's questions -share the document and any leading tokens they have in common. For joins, the -partner label counts once per anchor, and each pair's partner document and -answer cue count once per pair, less the leading tokens shared with another -partner of the same anchor. Recomputed -tokens are therefore the work a perfect prefix KV cache would have avoided. +The minimum is the number of input tokens that the requests need when the KV of +every prompt prefix stays in memory. It counts each distinct prompt prefix one +time: + +- A document counts one time, however many questions ask about it. +- If two questions about one document start with the same tokens, these tokens + count one time. +- In a join, the anchor document, the frame, and the partner label count one + time for each anchor. The partner label is the same for every pair of that + anchor. +- Each pair adds its partner document and its answer cue. If two partners of + the same anchor start with the same tokens, these tokens count one time. + +Recomputed tokens are fresh tokens minus minimum tokens. They are the work that +a prefix KV cache with unlimited memory does not do. A classification's tail counts like a filter question: once per document, sharing its lead with the document's other questions. A classification of From b4c6fb31ee11e93ce23a1a30b8bba239bdfae93f Mon Sep 17 00:00:00 2001 From: Shreya Shankar Date: Fri, 2 Oct 2026 23:45:24 +0000 Subject: [PATCH 4/5] List how the minimum counts classifications, one rule per bullet --- docs/reference.md | 33 ++++++++++++++++++++++----------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/docs/reference.md b/docs/reference.md index 6e14e94..ae269a7 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -615,17 +615,28 @@ time: Recomputed tokens are fresh tokens minus minimum tokens. They are the work that a prefix KV cache with unlimited memory does not do. -A classification's tail counts like a filter question: once per document, -sharing its lead with the document's other questions. A classification of -joined rows counts like a join. Reading a label takes no position after the -answer cue, so the minimum does not depend on how an engine reads labels. - -For example, suppose an engine classifies 100 reviews after asking one filter -question of each. Each review's tokens count once, and the filter question -and the classification tail count once per review, minus the tokens they -share at their start. Computing a review twice, sending a longer lettered -category list, or feeding label tokens after the cue to score each label is -work beyond the minimum, and counts as recomputed tokens. +A classification counts as follows: + +- The classification tail (the question, the labels, and the answer cue) counts + like a filter question: one time for each document. If it starts with the + same tokens as another question about the document, these tokens count one + time. +- A classification of joined rows counts like a join. +- The minimum stops at the answer cue. It does not depend on how an engine + reads the label. + +For example, an engine asks one filter question about each of 100 reviews, and +then classifies each review. The minimum counts: + +- the tokens of each review, one time; +- for each review, the filter question and the classification tail, less the + tokens that they share at their start. + +These are recomputed tokens: + +- a review that the engine computes two times; +- a lettered list of categories that is longer than the reference tail; +- label tokens that the engine feeds after the answer cue to score each label. Input tokens depend on which prompts the plan evaluates. Fresh tokens measure model computation. Two engines can therefore have the same input tokens and From fbbe454d9fd13900df2f4fbd064c022ab4427fc9 Mon Sep 17 00:00:00 2001 From: Shreya Shankar Date: Fri, 2 Oct 2026 23:52:30 +0000 Subject: [PATCH 5/5] Order the token metrics as in their table, one heading per metric --- docs/reference.md | 39 +++++++++++++++++++++++++++++---------- 1 file changed, 29 insertions(+), 10 deletions(-) diff --git a/docs/reference.md b/docs/reference.md index ae269a7..ee173b4 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -585,8 +585,11 @@ made. ### Throughput -Document throughput is input documents divided by `runtime_s`. Join throughput -is evaluated document pairs, summed across all joins, divided by `runtime_s`. +| Metric | Definition | +| --- | --- | +| Document throughput | Input documents divided by `runtime_s` | +| Join throughput | Evaluated document pairs, summed across all joins, divided by `runtime_s` | +| Input token throughput | Input tokens divided by `runtime_s` | ### Tokens and KV @@ -597,7 +600,20 @@ is evaluated document pairs, summed across all joins, divided by `runtime_s`. | Minimum tokens | Input positions required with an unlimited prefix KV cache | | Recomputed tokens | Fresh tokens minus minimum tokens | | KV regret | Recomputed tokens divided by fresh tokens, as a percentage | -| Input token throughput | Input tokens divided by `runtime_s` | + +#### Input tokens + +Input tokens are the full length of every prompt that the plan evaluates. They +include tokens that the engine serves from KV. Thus input tokens depend on +which prompts the plan evaluates, not on how the engine computes them. + +#### Fresh tokens + +Fresh tokens are the input positions that model forward passes process. They +measure model computation. Two engines can have the same input tokens and +different fresh tokens. + +#### Minimum tokens The minimum is the number of input tokens that the requests need when the KV of every prompt prefix stays in memory. It counts each distinct prompt prefix one @@ -612,9 +628,6 @@ time: - Each pair adds its partner document and its answer cue. If two partners of the same anchor start with the same tokens, these tokens count one time. -Recomputed tokens are fresh tokens minus minimum tokens. They are the work that -a prefix KV cache with unlimited memory does not do. - A classification counts as follows: - The classification tail (the question, the labels, and the answer cue) counts @@ -632,15 +645,21 @@ then classifies each review. The minimum counts: - for each review, the filter question and the classification tail, less the tokens that they share at their start. -These are recomputed tokens: +#### Recomputed tokens + +Recomputed tokens are fresh tokens minus minimum tokens. They are the work that +a prefix KV cache with unlimited memory does not do. In the example above, these +are recomputed tokens: - a review that the engine computes two times; - a lettered list of categories that is longer than the reference tail; - label tokens that the engine feeds after the answer cue to score each label. -Input tokens depend on which prompts the plan evaluates. Fresh tokens measure -model computation. Two engines can therefore have the same input tokens and -different fresh tokens. +#### KV regret + +KV regret is recomputed tokens divided by fresh tokens, as a percentage. For +example, an engine that computes 100 fresh tokens for requests with a minimum +of 80 tokens has 20 recomputed tokens and a KV regret of 20%. ### Cost