Share the partner label and partner starts in the join minimum

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,29 +600,67 @@ 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` |
 
-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
-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,
-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.
-
-Input tokens depend on which prompts the plan evaluates. Fresh tokens measure
-model computation. Two engines can therefore have the same input tokens and
+#### 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
+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.
+
+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.
+
+#### 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.
+
+#### 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
 
 ```text

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
 
 

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