llguidance 1.1.2__cp39-abi3-macosx_10_12_x86_64.whl → 1.2.0__cp39-abi3-macosx_10_12_x86_64.whl

llguidance/_lib.pyi

@@ -526,18 +526,49 @@ class LLExecutor:
         self,
         interpreters: List[Tuple[LLMatcher, int]],
         trg_pointer: int,
-        one_mask_byte_size: int,
+        one_mask_bytes: int,
         trg_batch_size: int,
     ) -> None:
         """
         Compute the token mask directly into memory at the specified pointer.
         For each matcher, provide the index of the target mask.
-        If index is K, the memory will be written at trg_pointer + K * one_mask_byte_size,
+        If index is K, the memory will be written at trg_pointer + K * one_mask_bytes,
         where K < trg_batch_size.
-        Memory has to have size trg_batch_size * one_mask_byte_size.
+        Memory has to have size trg_batch_size * one_mask_bytes.
         Prefer to use fill_next_token_bitmask_par(), which wraps this.
         """
 
+    def unsafe_compute_mask_ptr_with_draft_token(
+        self,
+        interpreters: List[Tuple[LLMatcher, int, List[int]]],
+        trg_pointer: int,
+        one_mask_bytes: int,
+        trg_batch_size: int,
+    ) -> None:
+        """
+        Compute the token mask directly into memory at the specified pointer, including draft tokens.
+
+        This function extends unsafe_compute_mask_ptr() to handle draft tokens in speculative decoding.
+        For each matcher in the batch, it computes masks for both the current position and all draft tokens.
+
+        Args:
+            interpreters: List of tuples containing:
+                - LLMatcher: The matcher object for constrained generation
+                - int: Index K indicating the target mask position (K < trg_batch_size)  
+                - List[int]: Draft tokens to be processed for speculative decoding
+            trg_pointer: Memory address where mask data will be written
+            one_mask_bytes: Size in bytes of a single token mask
+            trg_batch_size: Total batch size for memory allocation validation
+
+        Memory Layout:
+            - Main mask written at: trg_pointer + K * one_mask_bytes
+            - Draft token i mask written at: trg_pointer + (K + i + 1) * one_mask_bytes
+            - Total memory required: trg_batch_size * one_mask_bytes
+
+        The function processes each matcher's draft tokens sequentially, advancing the matcher state
+        for each valid token until encountering an invalid token or termination condition.
+        State rollback is performed to maintain matcher consistency.
+        """
 
 class JsonCompileOptions(TypedDict, total=False):
     # defaults to ","

llguidance/numpy.py

@@ -66,3 +66,17 @@ def fill_next_token_bitmask_par(executor: LLExecutor,
     batch, vocab = bitmask.shape
     assert bitmask.flags["C_CONTIGUOUS"], "Mask must be contiguous"
     executor.unsafe_compute_mask_ptr(matchers, bitmask.ctypes.data, vocab * 4, batch)
+
+
+def fill_next_token_bitmask_par_with_draft_tokens(executor: LLExecutor,
+                                matchers: List[Tuple[LLMatcher, int, List[int]]],
+                                bitmask: NDArray[np.int32]) -> None:
+    """
+    Compute the token mask directly into the specified array.
+    For each matcher, provide the index of the target mask.
+    """
+    assert bitmask.dtype == np.int32, "Mask must be int32"
+    assert bitmask.ndim == 2, "Mask must be 2D"
+    batch, vocab = bitmask.shape
+    assert bitmask.flags["C_CONTIGUOUS"], "Mask must be contiguous"
+    executor.unsafe_compute_mask_ptr_with_draft_token(matchers, bitmask.ctypes.data, vocab * 4, batch)

llguidance/torch.py

@@ -66,3 +66,14 @@ def fill_next_token_bitmask_par(executor: LLExecutor,
     assert bitmask.is_contiguous(), "Mask must be contiguous"
     executor.unsafe_compute_mask_ptr(matchers, bitmask.data_ptr(), vocab * 4,
                                      batch)
+
+
+def fill_next_token_bitmask_par_with_draft_tokens(executor: LLExecutor,
+                                matchers: List[Tuple[LLMatcher, int, List[int]]],
+                                bitmask: torch.Tensor) -> None:
+    assert bitmask.dtype == torch.int32, "Mask must be int32"
+    assert bitmask.is_cpu, "Mask must be on CPU"
+    assert bitmask.dim() == 2, "Mask must be 2D"
+    batch, vocab = bitmask.shape
+    assert bitmask.is_contiguous(), "Mask must be contiguous"
+    executor.unsafe_compute_mask_ptr_with_draft_token(matchers, bitmask.data_ptr(), vocab * 4, batch)

{llguidance-1.1.2.dist-info → llguidance-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llguidance
-Version: 1.1.2
+Version: 1.2.0
 License-File: LICENSE
 Summary: Bindings for the Low-level Guidance (llguidance) Rust library for use within Guidance
 Author: Michal Moskal

{llguidance-1.1.2.dist-info → llguidance-1.2.0.dist-info}/RECORD

@@ -1,10 +1,10 @@
-llguidance-1.1.2.dist-info/METADATA,sha256=CcBfni8m-8SJTKFhZXxNKOrUbamYuOqXFIyvmpceptE,10300
-llguidance-1.1.2.dist-info/WHEEL,sha256=TiMJekJwYXi-5FCpHPqncJXv9UVKDzSHt4YRv5UDSSg,104
-llguidance-1.1.2.dist-info/licenses/LICENSE,sha256=ws_MuBL-SCEBqPBFl9_FqZkaaydIJmxHrJG2parhU4M,1141
+llguidance-1.2.0.dist-info/METADATA,sha256=0zDkSmKb7n-5DphoLqb3-upZlzpwWBQUHUUBW57F-Eg,10300
+llguidance-1.2.0.dist-info/WHEEL,sha256=TiMJekJwYXi-5FCpHPqncJXv9UVKDzSHt4YRv5UDSSg,104
+llguidance-1.2.0.dist-info/licenses/LICENSE,sha256=ws_MuBL-SCEBqPBFl9_FqZkaaydIJmxHrJG2parhU4M,1141
 llguidance/__init__.py,sha256=F9svXvm6oafbuUf_eq34PHJV4c7-yN133vmbWN6nIkc,590
 llguidance/_grammar_from.py,sha256=-vHqkPqJe6t0JKKuQhlUu08kYpPIVknMh8tZlh8FYeQ,2384
-llguidance/_lib.abi3.so,sha256=4ogPfVCHQd_3ION9E6lBgRX22f6q9ZxTNumGjeklsK4,8587852
-llguidance/_lib.pyi,sha256=WIiIRVO5GKfGpowqy7wBmpgXMQ9Z7dXs4WVcOTT78cs,22617
+llguidance/_lib.abi3.so,sha256=1H38cQvfwzSB7DgGPgBjVEcG24iBcjm7TvORFmVhrlA,8625700
+llguidance/_lib.pyi,sha256=1m69VQ-7XaQLJKf6Vw5qf-UEh1cEigVMZM0Hj-XgSiw,24213
 llguidance/_struct_tag.py,sha256=83okmGWShxZud7S2vHjPRiInhFw0QVTHkeN8wtR8hR8,4430
 llguidance/_tokenizer.py,sha256=yC-RcgyMZN-olV-PnN4XkjlH-fOU8E9jrwO8VkXLv4M,1084
 llguidance/_util.py,sha256=6JV5SxjoH7hZPaSHhPRD_G6JzIhbKFFTqWTpp88VIiU,260
@@ -13,8 +13,8 @@ llguidance/gbnf_to_lark.py,sha256=32XJ5Dzq-iSySnkV_rLaNZ888JjHBIr_QkSYdhtMAME,16
 llguidance/hf.py,sha256=sLJKZxGpftuAY5eSCYpogzim3WZA7-vs9SOVgnzf3xI,1933
 llguidance/llamacpp.py,sha256=d_LjNbomBhj7uTo90h0muvPkOMso8NLe7H_YfPCqA8U,2284
 llguidance/mlx.py,sha256=ydMNSjnI2db8F685waOSWQLIoQTB9oehGBXnp9T3A6Q,2427
-llguidance/numpy.py,sha256=Z4888IfdT5oAfq4m2gn6ARfkaJ9d7nxMS79g_t8aXkg,2653
+llguidance/numpy.py,sha256=ZxHoQ_OmAgUHmxa_6KA8ruuKzXZLjmanPM9BYlKI0bo,3327
 llguidance/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 llguidance/tiktoken.py,sha256=-mYDPxq4LM2FI9K8kLL0Us0qN-fZpNK0qIwboDPjfSk,1101
-llguidance/torch.py,sha256=7LmVtRX8-fdUULcpPh4KYuDwVrHFJEvyVxzhyQflpgM,2795
-llguidance-1.1.2.dist-info/RECORD,,
+llguidance/torch.py,sha256=qLwQjaxZbvaVss7Qy_XP-SEPGYRqJRkvd6KE4LmzZJ8,3372
+llguidance-1.2.0.dist-info/RECORD,,