alpha_sys_1.py: shared-prefix inference for multi-question requests

alpha_sys_1.py

@@ -12,6 +12,8 @@ on them and reads the answer distribution from one forward pass.
     m.system_one({"state": ..., "images": [...], "questions": {"q1": {...}, "q2": {...}}})
     # -> {"model": ..., "answers": {"q1": {...}, "q2": {...}}}   (TypeSafe's System One shape)
 
+The questions of one request share their images and state, so that prefix is computed once
+and only the questions run against its cache (`shared_prefix=False` turns this off).
 Question types: choice (criteria = {option: description or None} or a list of options, up to
 26), noul (a statement; criteria = {"true": ..., "false": ...} optional), score (criteria = the
 levels, lowest first; the score is the expected level index). Images: a PIL image, a path, or
@@ -88,9 +90,13 @@ def confidence(p: list[float]) -> float:
 
 
 class SystemOne:
-    def __init__(self, repo: str, revision: str | None = None, device: str | None = None, dtype=torch.bfloat16, temperature: float = 1.0):
-        """temperature: the label logits are divided by it (1.0 = the model as released; see fit_temperature)."""
-        self.repo, self.revision, self.temperature = repo, revision, temperature
+    def __init__(self, repo: str, revision: str | None = None, device: str | None = None, dtype=torch.bfloat16, temperature: float = 1.0,
+                 shared_prefix: bool = True):
+        """temperature: the label logits are divided by it (1.0 = the model as released; see fit_temperature).
+        shared_prefix: when several questions share their images and state, run that prefix once and
+        only the questions against its cache (same probabilities up to bfloat16 noise, several times
+        faster with an image). False runs every question as its own full sequence."""
+        self.repo, self.revision, self.temperature, self.shared_prefix = repo, revision, temperature, shared_prefix
         self.processor = AutoProcessor.from_pretrained(repo, revision=revision)
         self.processor.tokenizer.padding_side = "left"
         self.model = AutoModelForImageTextToText.from_pretrained(repo, revision=revision, dtype=dtype)
@@ -105,25 +111,68 @@ class SystemOne:
             self._ids[label] = ids[0]
         return self._ids[label]
 
-    @torch.inference_mode()
-    def distributions(self, items: list[tuple[Any, dict, list | None]]) -> list[list[float]]:
-        """items: (state, question, images or None) -> probabilities in the answer-space order."""
+    def _messages(self, items: list[tuple[Any, dict, list | None]]) -> tuple[list, list[list[str]]]:
         msgs, labels_per = [], []
         for state, q, images in items:
             text, labels, _ = render(state, q)
             content = [{"type": "image", "image": load_image(im)} for im in (images or [])] + [{"type": "text", "text": text}]
             msgs.append([{"role": "user", "content": content}])
             labels_per.append(labels)
-        inputs = self.processor.apply_chat_template(
-            msgs, add_generation_prompt=True, tokenize=True, return_dict=True,
-            processor_kwargs={"return_tensors": "pt", "padding": True}).to(self.device)
-        logits = self.model(**inputs, logits_to_keep=1).logits[:, -1].float()
+        return msgs, labels_per
+
+    def _probs(self, logits: torch.Tensor, labels_per: list[list[str]]) -> list[list[float]]:
         out = []
         for i, labels in enumerate(labels_per):
             ids = torch.tensor([self._label_id(lab) for lab in labels], device=logits.device)
             out.append(torch.softmax(logits[i, ids] / self.temperature, -1).tolist())
         return out
 
+    @torch.inference_mode()
+    def distributions(self, items: list[tuple[Any, dict, list | None]]) -> list[list[float]]:
+        """items: (state, question, images or None) -> probabilities in the answer-space order.
+        Every item is its own sequence, read at its last position; items that share images and
+        state share the prefix's computation when `shared_prefix` is on."""
+        msgs, labels_per = self._messages(items)
+        if self.shared_prefix and len(items) > 1 and all(it[0] == items[0][0] and (it[2] or None) == (items[0][2] or None) for it in items):
+            logits = self._shared_prefix_logits(msgs)
+            if logits is not None:
+                return self._probs(logits, labels_per)
+        inputs = self.processor.apply_chat_template(
+            msgs, add_generation_prompt=True, tokenize=True, return_dict=True,
+            processor_kwargs={"return_tensors": "pt", "padding": True}).to(self.device)
+        return self._probs(self.model(**inputs, logits_to_keep=1).logits[:, -1].float(), labels_per)
+
+    def _shared_prefix_logits(self, msgs: list) -> torch.Tensor | None:
+        """One pass over the longest common token prefix (images, state), then the question suffixes,
+        right-padded, against that cache repeated across the batch. None when there is too little to share."""
+        encs = [self.processor.apply_chat_template([m], add_generation_prompt=True, tokenize=True, return_dict=True,
+                                                   processor_kwargs={"return_tensors": "pt"}) for m in msgs]
+        ids = [e["input_ids"][0] for e in encs]
+        L = min(len(x) for x in ids) - 1  # at least one token per suffix
+        for x in ids[1:]:
+            diff = (x[:L] != ids[0][:L]).nonzero()
+            if len(diff):
+                L = min(L, int(diff[0]))
+        image_id = getattr(self.model.config, "image_token_id", None)
+        if L < 64 or (image_id is not None and any((x[L:] == image_id).any() for x in ids)):
+            return None
+        n = len(ids)
+        image_kw = {k: v.to(self.device) for k, v in encs[0].items() if k in ("pixel_values", "spatial_shapes", "pixel_attention_mask")}
+        cache = self.model(input_ids=ids[0][:L][None].to(self.device), **image_kw, use_cache=True).past_key_values
+        cache.reorder_cache(torch.zeros(n, dtype=torch.long, device=self.device))
+        sufs = [x[L:] for x in ids]
+        lens = torch.tensor([len(s) for s in sufs])
+        width = int(lens.max())
+        suffix = torch.full((n, width), self.processor.tokenizer.pad_token_id, dtype=torch.long)
+        attention = torch.zeros((n, L + width), dtype=torch.long)
+        attention[:, :L] = 1
+        for i, s in enumerate(sufs):
+            suffix[i, : len(s)] = s
+            attention[i, L : L + len(s)] = 1
+        out = self.model(input_ids=suffix.to(self.device), attention_mask=attention.to(self.device), past_key_values=cache,
+                         cache_position=torch.arange(L, L + width, device=self.device), use_cache=True)
+        return out.logits[torch.arange(n), (lens - 1).to(self.device)].float()
+
     def answer(self, q: dict, p: list[float]) -> dict:
         _, _, keys = render(None, q)
         if q["type"] == "choice":