--- a/python/src/captchakraken/pointer/service.py +++ b/python/src/captchakraken/pointer/service.py @@ -29,6 +29,8 @@ from . import vllm_plugin +VLLM_VIDEO = "<|vision_start|><|video_pad|><|vision_end|>" + class PointerService: def __init__(self, checkpoint: Union[str, Path], detector: Union[str, Path, None] = None, gpu_memory_utilization: float = 0.45, max_model_len: int = 8192, **llm_kwargs): @@ -41,7 +43,10 @@ self.llm = LLM(model=str(checkpoint), runner="pooling", max_model_len=max_model_len, gpu_memory_utilization=gpu_memory_utilization, limit_mm_per_prompt={"image": 1 + L.MAX_HISTORY, "video": 1}, - mm_processor_kwargs={"min_pixels": 64 * 64, "max_pixels": L.MAX_PIXELS}, **llm_kwargs) + mm_processor_kwargs={"min_pixels": 64 * 64, "max_pixels": L.MAX_PIXELS}, + # The heads read every token's hidden state; vLLM returns none for a + # prefix it served from cache, so a repeated prompt came back short. + enable_prefix_caching=False, **llm_kwargs) self.detector = None if detector is not None: from .detector import Detector @@ -59,7 +64,7 @@ marked: bool = True, history: Sequence[Image.Image] = ()) -> Dict: """The request vLLM gets, and what decoding its answer needs. `history`: the puzzle's earlier screens, oldest first, as they were sent (marked the - same way), so the prompt up to them is the one already cached.""" + same way). The prefix is recomputed each time: the heads need its hidden states.""" if (image is None) == (frames is None): raise ValueError("give exactly one of image or frames") if frames is not None and history: @@ -78,7 +83,12 @@ if tracks: fr = np.stack([np.asarray(M.draw(im, E.on_frame(tracks, k))) for k, im in enumerate(ims)]) - return {"request": {"prompt": L.chat(self.processor, True), + # vLLM (0.25.1) expands a video only where it finds the template's whole + # <|vision_start|><|video_pad|><|vision_end|>; `chat` sends transformers + # the bare pad. Both then give the same tokens but one: training merges the + # instruction's last "." with the first timestamp's "<", vLLM splits them. + prompt = L.chat(self.processor, True).replace(L.VIDEO_PAD, VLLM_VIDEO) + return {"request": {"prompt": prompt, "multi_modal_data": {"video": (fr, meta)}, "mm_processor_kwargs": {"do_sample_frames": False}}, "size": (w, h), "elements": tracks, "shown": None, --- a/python/src/captchakraken/pointer/vllm_plugin.py +++ b/python/src/captchakraken/pointer/vllm_plugin.py @@ -65,6 +65,9 @@ # A real prompt without its markers still fails in Positions. out.append(torch.zeros(L.PACKED + TRAILER, device=h.device)) continue + if h.shape[0] != ids.shape[0]: + raise RuntimeError(f"{h.shape[0]} hidden states for a {ids.shape[0]}-token prompt: " + "the heads need them all (is prefix caching on?)") pos = Positions(ids, m.pointer_token_ids) v = m.pointer_heads.decide(h, pos, m.language_model.lm_head.weight) out.append(torch.cat([v, torch.tensor([float(pos.per_moment), float(pos.n_moments)],