lmcache-cli 0.5.7rc2.dev0__py3-none-any.whl → 0.5.7rc4.dev0__py3-none-any.whl

lmcache/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.7rc2.dev0'
-__version_tuple__ = version_tuple = (0, 5, 7, 'rc2', 'dev0')
+__version__ = version = '0.5.7rc4.dev0'
+__version_tuple__ = version_tuple = (0, 5, 7, 'rc4', 'dev0')
 
-__commit_id__ = commit_id = 'gaa12d34e3'
+__commit_id__ = commit_id = 'g45a0bb420'

lmcache/cli/commands/bench/engine_bench/quality/alignment.py

@@ -5,11 +5,27 @@ A passage cached on its own matches inside a longer prompt only if it starts
 on a cache-chunk boundary in both places: chunks are content-addressed, so a
 passage that lands off-phase shares no chunk with its cached copy and nothing
 is reused, silently.  :class:`ChunkAligner` pads the shared system block and
-each passage so every block occupies whole chunks.
+each passage so every block occupies whole chunks, and
+:meth:`ChunkAligner.place_blocks` checks, on the rendered chat prompt, that
+they really do.
+
+What fills the padding is the :class:`Filler`:
+
+* ``words`` (legacy, and ``rag-qa-quality``'s) appends random single-token
+  words.
+* ``dots`` / ``dashes`` wrap each passage as
+  ``Passage <id>:\\n<text>\\n\\n . . … .\\nend of passage <id>`` -- a repeated
+  low-information token the model can read as a divider.  The block starts
+  with a letter and ends with a digit, so no tokenizer merges it with its
+  neighbours, and no chat template trims it (see :func:`block_id`).
 """
 
 # Standard
+from collections.abc import Mapping
+from dataclasses import dataclass
 from typing import Any
+import enum
+import hashlib
 import math
 import random
 
@@ -31,14 +47,79 @@ _TEMPLATE_SENTINEL = "██SENTINEL██"
 # Re-encode attempts per block; a merged boundary token shifts the estimate.
 _PAD_ATTEMPTS = 8
 
+# Chunk targets tried per block by the unit fillers before giving up: a
+# boundary merge can make a target unreachable, never the next one too.
+_PAD_TARGETS = 3
+
+
+class Filler(enum.Enum):
+    """What pads a block to a whole number of chunks."""
+
+    WORDS = "words"
+    """Random single-token words (legacy; ``rag-qa-quality`` uses it).
+    Noisy for the model to read."""
+
+    DOTS = "dots"
+    """Repeated ``" ."`` inside a ``Passage <id>: … end of passage <id>``
+    frame."""
+
+    DASHES = "dashes"
+    """Repeated ``" -"`` inside the same frame as ``dots``."""
+
+
+_PAD_UNITS = {Filler.DOTS: " .", Filler.DASHES: " -"}
+
+# Ends the padded system block in the framed layouts: a fixed digit, so the
+# first passage's "Passage" never merges into it.
+_SYSTEM_TAIL = "\nend of instructions 0"
+
+
+def block_id(text: str) -> str:
+    """Return the id a framed block is labelled with.
+
+    Six hex characters of the text's digest plus one decimal digit, so the
+    id -- and the block, which ends with it -- always ends in a digit. Every
+    tokenizer checked splits a digit from a following letter or newline, while
+    a final letter merges into the next block's first word on some, a final
+    punctuation mark glues to the newlines after it on others, and final
+    whitespace is merged with its neighbours or stripped by chat templates.
+
+    Args:
+        text: The raw passage text.
+
+    Returns:
+        The id, e.g. ``"4f7a1c9"``. Identical for identical text.
+    """
+    digest = hashlib.sha256(text.encode()).hexdigest()
+    return digest[:6] + str(int(digest, 16) % 10)
+
+
+@dataclass(frozen=True)
+class BlockPlacement:
+    """Where one block sits in a rendered prompt.
+
+    Attributes:
+        block: The block text.
+        start_token: Its first token's index in the tokenized prompt, or
+            ``-1`` when the block was not found verbatim.
+        tokens: Its token ids in this prompt.
+        problem: Why it cannot be reused, or ``""`` when it starts and ends
+            on chunk boundaries with no token straddling either.
+    """
+
+    block: str
+    start_token: int
+    tokens: tuple[int, ...]
+    problem: str
+
 
 class ChunkAligner:
     """Pads prompt blocks so each occupies a whole number of cache chunks.
 
-    Padding is deterministic: the same tokenizer, chunk size, seed and text
-    always yield the same padded block, so two runs build byte-identical
-    prompts.  Passage padding is seeded from the passage text, so a passage
-    shared by several samples is padded once, identically.
+    Padding is deterministic: the same tokenizer, chunk size, filler and text
+    always yield the same padded block (and, for ``words``, the same seed),
+    so two runs build byte-identical prompts.  A passage shared by several
+    samples is padded once, identically.
     """
 
     def __init__(
@@ -47,19 +128,25 @@ class ChunkAligner:
         model_name: str,
         align_tokens: int,
         seed: int,
+        filler: Filler = Filler.WORDS,
+        template_kwargs: Mapping[str, Any] | None = None,
     ) -> None:
-        """Build the filler pool for *tokenizer*.
+        """Build the aligner (and, for ``words``, the filler pool).
 
         Args:
             tokenizer: A HuggingFace tokenizer for the served model.
             model_name: The model's name, used only in log messages.
             align_tokens: Chunk size to align to.  Should equal the
                 deployment's LMCache chunk size.
-            seed: Seed for filler selection.
+            seed: Seed for filler selection (``words`` only).
+            filler: What pads each block.
+            template_kwargs: Chat-template variables the engine renders
+                with; they can change the template's length ahead of the
+                content.
 
         Raises:
             ValueError: If *align_tokens* is not positive, or the tokenizer
-                cannot supply a single-token filler pool.
+                cannot supply a single-token filler pool for ``words``.
         """
         if align_tokens < 1:
             raise ValueError(f"align_tokens must be >= 1, got {align_tokens}")
@@ -67,8 +154,25 @@ class ChunkAligner:
         self._model_name = model_name
         self._align_tokens = align_tokens
         self._seed = seed
-        self._pool = build_single_token_pool(tokenizer, _FILLER_VOCAB_SIZE, seed=seed)
-        self._passage_blocks: dict[str, str] = {}
+        self._filler = filler
+        self._template_kwargs = dict(template_kwargs or {})
+        self._pool = (
+            build_single_token_pool(tokenizer, _FILLER_VOCAB_SIZE, seed=seed)
+            if filler is Filler.WORDS
+            else None
+        )
+        self._passage_blocks: dict[str, tuple[str, int]] = {}
+        self._system_pad_tokens = 0
+
+    @property
+    def filler(self) -> Filler:
+        """What pads each block."""
+        return self._filler
+
+    @property
+    def system_pad_tokens(self) -> int:
+        """Padding tokens in the last block :meth:`system_block` returned."""
+        return self._system_pad_tokens
 
     def token_length(self, text: str) -> int:
         """Return the token length of *text*, without special tokens.
@@ -95,9 +199,16 @@ class ChunkAligner:
             The padded system block.
         """
         prefix_tokens = self._chat_prefix_tokens()
-        block = self._pad_to_multiple(
-            system_prompt, prefix_tokens, random.Random(self._seed)
-        )
+        if self._filler is Filler.WORDS:
+            block = self._pad_to_multiple(
+                system_prompt, prefix_tokens, random.Random(self._seed)
+            )
+            pad = self.token_length(block) - self.token_length(system_prompt)
+        else:
+            block, pad = self._pad_framed(
+                "", system_prompt, _SYSTEM_TAIL, prefix_tokens
+            )
+        self._system_pad_tokens = pad
         logger.info(
             "System block: %d tokens after a %d-token chat-template prefix",
             self.token_length(block),
@@ -114,12 +225,161 @@ class ChunkAligner:
         Returns:
             The padded passage; identical across calls for the same text.
         """
-        block = self._passage_blocks.get(passage)
-        if block is None:
-            rng = random.Random(f"{self._seed}:{passage}")
-            block = self._pad_to_multiple(passage, 0, rng)
-            self._passage_blocks[passage] = block
-        return block
+        return self._padded_passage(passage)[0]
+
+    def pad_tokens(self, passage: str) -> int:
+        """Return how many padding tokens :meth:`passage_block` adds.
+
+        Args:
+            passage: Raw passage text.
+
+        Returns:
+            Tokens of filler (``words``: the words; framed fillers: the
+            repeated unit), excluding the frame around the text.
+        """
+        return self._padded_passage(passage)[1]
+
+    def framed(self, passage: str) -> str:
+        """Return *passage* in the framed layout without padding.
+
+        ``Passage <id>:\n<text>\n\nend of passage <id>``: the same label
+        and end marker as the ``dots`` / ``dashes`` blocks, for prompts that
+        are not chunk-aligned (``bench quality``'s ``concatenated``
+        scenario).
+
+        Args:
+            passage: Raw passage text.
+
+        Returns:
+            The framed passage.
+        """
+        ident = block_id(passage)
+        return f"Passage {ident}:\n{passage}\n\nend of passage {ident}"
+
+    def render_and_encode(
+        self, messages: list[dict[str, str]]
+    ) -> tuple[str, list[int]]:
+        """Render *messages* with the chat template and tokenize the result.
+
+        Args:
+            messages: Chat messages, as sent to the engine.
+
+        Returns:
+            The rendered prompt and its token ids, as the engine sees them.
+        """
+        rendered = self._render(messages)
+        return rendered, self._encode(rendered)
+
+    def token_boundary(
+        self, rendered: str, char_index: int, ids: list[int]
+    ) -> tuple[int, bool]:
+        """Map a character position in a rendered prompt to a token index.
+
+        Args:
+            rendered: A prompt from :meth:`render_and_encode`.
+            char_index: A character position in it.
+            ids: Its token ids.
+
+        Returns:
+            The number of tokens before *char_index*, and whether a token
+            boundary falls exactly there (the text before it tokenizes to a
+            prefix of the whole prompt's tokens).
+        """
+        head = self._encode(rendered[:char_index])
+        return len(head), ids[: len(head)] == head
+
+    def place_blocks(
+        self, messages: list[dict[str, str]], blocks: list[str]
+    ) -> tuple[int, list[BlockPlacement]]:
+        """Locate *blocks* in the prompt the engine will tokenize.
+
+        Renders *messages* with the chat template (as the engine does),
+        tokenizes the result, and checks each block: it must appear
+        verbatim, in order, start and end on chunk boundaries, and no token
+        may straddle either boundary. Only then do its chunks hash the same
+        here as in any other prompt holding it -- which callers check by
+        comparing ``tokens`` across prompts.
+
+        Args:
+            messages: The chat messages, as sent to the engine.
+            blocks: The padded blocks the messages contain, in order.
+
+        Returns:
+            The rendered prompt's token count, and one placement per block.
+        """
+        rendered, ids = self.render_and_encode(messages)
+        placements: list[BlockPlacement] = []
+        cursor = 0
+        for block in blocks:
+            start = rendered.find(block, cursor)
+            if start < 0:
+                placements.append(
+                    BlockPlacement(
+                        block,
+                        -1,
+                        (),
+                        "not found verbatim in the rendered prompt (the chat "
+                        "template changed it, e.g. trimmed whitespace)",
+                    )
+                )
+                continue
+            end = start + len(block)
+            cursor = end
+            first, clean_start = self.token_boundary(rendered, start, ids)
+            last, clean_end = self.token_boundary(rendered, end, ids)
+            problems = []
+            if not clean_start:
+                problems.append("a token straddles its start")
+            if not clean_end:
+                problems.append("a token straddles its end")
+            if clean_start and first % self._align_tokens:
+                problems.append(
+                    f"starts at token {first}, not on a "
+                    f"{self._align_tokens}-token chunk boundary"
+                )
+            if clean_start and clean_end and (last - first) % self._align_tokens:
+                problems.append(
+                    f"spans {last - first} tokens, not a whole number of chunks"
+                )
+            placements.append(
+                BlockPlacement(
+                    block, first, tuple(ids[first:last]), "; ".join(problems)
+                )
+            )
+        return len(ids), placements
+
+    def _padded_passage(self, passage: str) -> tuple[str, int]:
+        """Pad *passage* once and cache the block with its padding count."""
+        cached = self._passage_blocks.get(passage)
+        if cached is None:
+            if self._filler is Filler.WORDS:
+                rng = random.Random(f"{self._seed}:{passage}")
+                block = self._pad_to_multiple(passage, 0, rng)
+                cached = (block, self.token_length(block) - self.token_length(passage))
+            else:
+                ident = block_id(passage)
+                cached = self._pad_framed(
+                    f"Passage {ident}:\n", passage, f"\nend of passage {ident}", 0
+                )
+            self._passage_blocks[passage] = cached
+        return cached
+
+    def _render(self, messages: list[dict[str, str]]) -> str:
+        """Render *messages* as the engine will, or their content without one."""
+        try:
+            return str(
+                self._tokenizer.apply_chat_template(
+                    messages,
+                    tokenize=False,
+                    add_generation_prompt=True,
+                    **self._template_kwargs,
+                )
+            )
+        except Exception:  # noqa: BLE001 - no template: content only, as below
+            return "".join(m["content"] for m in messages)
+
+    def _encode(self, text: str) -> list[int]:
+        return list(self._tokenizer.encode(text, add_special_tokens=False))
 
     def _chat_prefix_tokens(self) -> int:
         """Count the tokens the chat template inserts before the content.
@@ -133,6 +393,7 @@ class ChunkAligner:
                 [{"role": "user", "content": _TEMPLATE_SENTINEL}],
                 tokenize=False,
                 add_generation_prompt=True,
+                **self._template_kwargs,
             )
         except Exception as e:  # noqa: BLE001 - a template failure is non-fatal
             logger.warning(
@@ -146,8 +407,54 @@ class ChunkAligner:
         head = str(rendered).split(_TEMPLATE_SENTINEL)[0]
         return self.token_length(head)
 
+    def _pad_framed(
+        self, head: str, text: str, tail: str, offset: int
+    ) -> tuple[str, int]:
+        """Pad ``head + text + "\\n\\n" + unit*n + tail`` to a chunk boundary.
+
+        Args:
+            head: Text before *text* (the passage label, or ``""``).
+            text: The passage or system prompt.
+            tail: Text after the padding (the end marker).
+            offset: Tokens already ahead of the block.
+
+        Returns:
+            The padded block and its padding token count. When no count
+            lands on a boundary the closest attempt is returned; the
+            in-context check (:meth:`place_blocks`) then reports it.
+        """
+        unit = _PAD_UNITS[self._filler]
+        align = self._align_tokens
+
+        def build(n: int) -> str:
+            return head + text + "\n\n" + unit * n + tail
+
+        base = self.token_length(build(0))
+        target = math.ceil((offset + base) / align) * align
+        block = build(0)
+        for _ in range(_PAD_TARGETS):
+            n = target - offset - base
+            for _ in range(_PAD_ATTEMPTS):
+                n = max(n, 0)
+                block = build(n)
+                actual = offset + self.token_length(block)
+                if actual == target:
+                    return block, actual - offset - base
+                n += target - actual
+            target += align
+        logger.warning(
+            "Could not pad a block to a %d-token boundary with %r; cache "
+            "reuse will be partial",
+            align,
+            unit,
+        )
+        return block, self.token_length(block) - base
+
     def _pad_to_multiple(self, text: str, offset: int, rng: random.Random) -> str:
         """Pad *text* so ``offset + len(text)`` is a whole chunk count."""
+        pool = self._pool
+        if pool is None:
+            raise RuntimeError("word padding needs the words filler")
         align = self._align_tokens
         total = offset + self.token_length(text)
         target = math.ceil(total / align) * align
@@ -159,8 +466,8 @@ class ChunkAligner:
         # Correct against a re-encode: a merged boundary token shifts the
         # estimate, and an off-by-one phase error costs a whole chunk.
         for _ in range(_PAD_ATTEMPTS):
-            words = [rng.choice(self._pool.words) for _ in range(max(num_words, 0))]
-            padded = text + "\n" + self._pool.join(words)
+            words = [rng.choice(pool.words) for _ in range(max(num_words, 0))]
+            padded = text + "\n" + pool.join(words)
             actual = offset + self.token_length(padded)
             if actual == target:
                 return padded

lmcache/cli/commands/bench/engine_bench/quality/prompts.py

@@ -4,8 +4,19 @@
 Kept in one module so every benchmark that scores answers over the same
 passages asks byte-identical questions: a prompt difference between two runs
 would read as a quality difference.
+
+Two versions exist. :data:`ANSWER_TAG_PROMPT` (``bench quality``) asks for
+``<answer>`` tags. :data:`QA_SYSTEM_PROMPT` / :data:`QA_QUESTION_TEMPLATE`
+(``rag-qa-quality``) ask for ``<final_answer>`` tags, which gpt-oss trips
+over: its response format names the answer channel ``final``, and a prompt
+full of ``final_answer`` leads it to write ``<|channel|>final_answer>…``,
+dropping the opening tag. ``extract_final_answer`` accepts both tags and that
+form.
 """
 
+# Standard
+from dataclasses import dataclass
+
 QA_SYSTEM_PROMPT = (
     "Answer the question using only the given passages. You may reason "
     "through the problem, but put only the concise final answer between "
@@ -20,16 +31,50 @@ QA_QUESTION_TEMPLATE = (
 )
 
 
+@dataclass(frozen=True)
+class QAPrompt:
+    """The instructions around a question's passages.
+
+    Attributes:
+        system: Placed before the passages (then chunk-aligned).
+        question_template: Placed after them; ``{question}`` is filled in.
+    """
+
+    system: str
+    question_template: str
+
+
+ANSWER_TAG_PROMPT = QAPrompt(
+    system=(
+        "Answer the question using only the given passages. You may reason "
+        "through the problem, but put only the concise answer between "
+        "<answer> and </answer>. Always emit both tags.\n\n"
+        "The following are the given passages.\n"
+    ),
+    question_template=(
+        "\n\nAnswer the question using only the passages above. End with "
+        "exactly one concise answer in this form: <answer>answer</answer>.\n\n"
+        "Question: {question}\nAnswer:"
+    ),
+)
+"""The prompt with ``<answer>`` tags, which no checked chat format uses."""
+
+
 def compose_qa_messages(
-    system_block: str, passage_blocks: list[str], question: str
+    system_block: str,
+    passage_blocks: list[str],
+    question: str,
+    question_template: str = QA_QUESTION_TEMPLATE,
 ) -> list[dict[str, str]]:
     """Build the chat messages for one question over its passages.
 
     Args:
-        system_block: :data:`QA_SYSTEM_PROMPT`, chunk-aligned by
+        system_block: A :class:`QAPrompt` system text, chunk-aligned by
             :meth:`ChunkAligner.system_block`.
         passage_blocks: Chunk-aligned passages, in the order to present them.
         question: The question text.
+        question_template: The text after the passages, with a
+            ``{question}`` field; the system block's prompt version.
 
     Returns:
         A single user message:
@@ -38,7 +83,7 @@ def compose_qa_messages(
     content = (
         system_block
         + "".join(passage_blocks)
-        + QA_QUESTION_TEMPLATE.format(question=question)
+        + question_template.format(question=question)
     )
     return [{"role": "user", "content": content}]
 

lmcache/cli/commands/bench/engine_bench/quality/scoring.py

@@ -17,8 +17,12 @@ _PUNCTUATION = str.maketrans("", "", string.punctuation)
 
 # Lazy match, last region wins: reasoning models may echo an example answer
 # before their own. Include an unfinished last region to avoid scoring the example.
+# Opening tags: <final_answer> or <answer>, or a bare "final_answer>" -- gpt-oss
+# opens its answer channel with "<|channel|>final" and, primed by the prompt,
+# continues "_answer>" instead of "<|message|>", so the "<" never appears.
 _FINAL_ANSWER = re.compile(
-    r"<final_answer>\s*(.*?)\s*(</final_answer>|$)",
+    r"(?:<(?:final_answer|answer)>|(?<![<\w/])final_answer>)"
+    r"\s*(.*?)\s*(</(?:final_answer|answer)>|$)",
     re.IGNORECASE | re.DOTALL,
 )
 
@@ -106,8 +110,11 @@ _ABSTENTION_ANSWERS = frozenset({"none", "n/a", "na", "null", "unknown"})
 def extract_final_answer(output: str) -> str:
     """Extract the model's delimited final answer from a response.
 
-    An unterminated region counts as no answer: generation was cut off, so
-    scoring the reasoning before it would report an answer never produced.
+    Accepts ``<final_answer>…</final_answer>`` and ``<answer>…</answer>``
+    (either prompt version), and gpt-oss's ``final_answer>…</final_answer>``
+    with the opening ``<`` lost to its channel header. An unterminated region
+    counts as no answer: generation was cut off, so scoring the reasoning
+    before it would report an answer never produced.
 
     Args:
         output: The model's full response text.

lmcache/cli/commands/bench/quality_bench/__init__.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``lmcache bench quality`` command group.
+
+Measures whether a serving engine change -- e.g. quantization, KV cache
+compression, post-training, or another engine or model -- changes the
+answers compared with a reference, in two stages:
+
+* ``run`` sends byte-identical prompts to one or more arms (a reference, a
+  candidate under test, and optionally a control) and
+  saves their answers to a results directory.
+* ``analyze`` scores the saved answers and reports per-question rows and a
+  verdict.
+
+Arms can be run one at a time into the same directory, so they never need
+to be served at once. The subcommands are auto-discovered from the
+``*_command`` modules in this package; the runtime lives in ``core`` so that
+discovery imports only these lightweight wrappers.
+"""
+
+# First Party
+from lmcache.cli.commands.base import CompositeCommand
+
+
+class QualityBenchCommand(CompositeCommand):
+    """Compare a candidate serving stack's answers with a reference."""
+
+    def name(self) -> str:
+        return "quality"
+
+    def help(self) -> str:
+        return (
+            "Compare a candidate serving stack's answers with a reference "
+            "(e.g. a quantized model vs the original)."
+        )

lmcache/cli/commands/bench/quality_bench/analyze_command.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``lmcache bench quality analyze`` subcommand."""
+
+# Standard
+import argparse
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+
+
+class QualityAnalyzeCommand(BaseCommand):
+    """Score saved arm answers and reach the quality verdict."""
+
+    def name(self) -> str:
+        return "analyze"
+
+    def help(self) -> str:
+        return "Score the saved answers in a results directory and print a verdict."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        # First Party
+        from lmcache.cli.commands.bench.quality_bench.core.analyze_stage import (
+            add_analyze_arguments,
+        )
+
+        add_analyze_arguments(parser)
+
+    def register(self, subparsers: argparse._SubParsersAction) -> None:
+        """Register without the common output args (--format/--output).
+
+        Results are written into the results directory, so the generic
+        output flags do not apply.
+        """
+        parser = subparsers.add_parser(self.name(), help=self.help())
+        self.add_arguments(parser)
+        parser.set_defaults(func=self.execute)
+
+    def execute(self, args: argparse.Namespace) -> None:
+        # First Party
+        from lmcache.cli.commands.bench.quality_bench.core.analyze_stage import (
+            run_quality_analyze,
+        )
+
+        run_quality_analyze(args)

lmcache/cli/commands/bench/quality_bench/core/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Runtime for ``lmcache bench quality``: plans, arm runs, scoring, verdict.
+
+Kept in a sub-package so that command discovery, which imports every
+module directly under ``quality_bench``, imports only the lightweight
+command wrappers.
+"""