token splicing in vllm

controllers/pyctrl/samples/test.py

@@ -45,10 +45,15 @@ async def test_backtrack_lang():
 
 
 async def test_hello():
-    aici.log_level = 10
+    aici.log_level = 3
     prompt = await aici.GetPrompt()
     print("prompt", prompt)
     await aici.gen_tokens(regex=r"[A-Z].*", max_tokens=5)
+    l = aici.Label()
+    await aici.FixedTokens("\n2 + 2 = ")
+    await aici.gen_tokens(regex=r"\d+", max_tokens=1)
+    await aici.FixedTokens("\n3 + 3 = ", following=l)
+    await aici.gen_tokens(regex=r"\d+", max_tokens=1)
 
 
 async def test_main():

py/pyaici/comms.py

@@ -263,6 +263,14 @@ class Branch:
     mask: Optional[int] = None
     splices: List[Splice] = field(default_factory=list)
 
+    def find_splice(self, token: Optional[int]) -> Optional[Splice]:
+        if self.mask is None:
+            token = None
+        for splice in self.splices:
+            if not splice.when_sampled or token in splice.when_sampled:
+                return splice
+        return None
+
     def is_splice(self) -> bool:
         return len(self.splices) == 1 and self.splices[0].when_sampled == []
 

scripts/vllm-server.sh

@@ -13,6 +13,8 @@ RUST_LOG=info,tokenizers=error,aicirt=trace \
 PYTHONPATH=py:py/vllm \
 python3 -m vllm.entrypoints.openai.api_server \
     --enforce-eager \
+    --use-v2-block-manager \
+    --enable-chunked-prefill \
     --aici-rt ./target/release/aicirt \
     --aici-tokenizer $AICI_TOK \
     --model $MODEL \