mostlyai-engine 1.5.6__tar.gz → 1.5.7__tar.gz

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.5.6
+Version: 1.5.7
 Summary: Synthetic Data Engine
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-engine
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-engine
@@ -37,15 +37,15 @@ Requires-Dist: psutil<6,>=5.9.5
 Requires-Dist: pyarrow>=16.0.0
 Requires-Dist: setuptools>=77.0.3
 Requires-Dist: tokenizers>=0.21.0
-Requires-Dist: torch<2.7.1,>=2.7.0
-Requires-Dist: torchaudio<2.7.1,>=2.7.0
-Requires-Dist: torchvision<0.22.1,>=0.22.0
-Requires-Dist: transformers<4.54.0,>=4.51.0
-Requires-Dist: xgrammar>=0.1.19
+Requires-Dist: torch<2.9.0,>=2.8.0
+Requires-Dist: torchaudio<2.9.0,>=2.8.0
+Requires-Dist: torchvision<0.24.0,>=0.23.0
+Requires-Dist: transformers>=4.55.0
+Requires-Dist: xgrammar>=0.1.21
 Provides-Extra: gpu
 Requires-Dist: bitsandbytes==0.42.0; (sys_platform == 'darwin') and extra == 'gpu'
 Requires-Dist: bitsandbytes>=0.45.5; (sys_platform == 'linux') and extra == 'gpu'
-Requires-Dist: vllm==0.9.1; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'gpu'
+Requires-Dist: vllm==0.10.2; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'gpu'
 Description-Content-Type: text/markdown
 
 # Synthetic Data Engine 💎

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/mostlyai/engine/__init__.py

@@ -22,7 +22,7 @@ from mostlyai.engine.splitting import split
 from mostlyai.engine.training import train
 
 __all__ = ["split", "analyze", "encode", "train", "generate", "init_logging", "set_random_state"]
-__version__ = "1.5.6"
+__version__ = "1.5.7"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/mostlyai/engine/_language/engine/base.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from collections.abc import Generator
 from dataclasses import dataclass
 
 from pydantic import BaseModel
@@ -26,10 +25,6 @@ class EngineMetrics:
 
 
 class LanguageEngine(ABC):
-    @abstractmethod
-    def initialize_logits_processors(self, schemas: Generator[BaseModel]):
-        pass
-
     @abstractmethod
     def generate(
         self, text: list[str], sampling_temperature: float, sampling_top_p: float
@@ -47,3 +42,24 @@ class LanguageEngine(ABC):
     @abstractmethod
     def cleanup(self):
         pass
+
+    def update_json_constraints(self, schemas: list[BaseModel] | None) -> None:
+        """Update JSON schema constraints for the next generation call.
+
+        Args:
+            schemas: Schema constraints to apply to the next generate() call.
+                    None to clear any existing constraints.
+
+        Default implementation does nothing.
+        Engines that support JSON constraints should override this method.
+        """
+        pass
+
+    def can_reuse_schemas(self) -> bool:
+        """Whether the engine can reuse JSON schema constraints across batches with different sizes.
+
+        Returns:
+            True if the engine can handle variable batch sizes with reused schema constraints,
+            False if schema constraints need to be recreated for each batch.
+        """
+        return True

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/mostlyai/engine/_language/engine/hf_engine.py

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 import time
-from collections.abc import Generator
 from os import PathLike
 from pathlib import Path
 
@@ -75,15 +74,6 @@ class HuggingFaceEngine(LanguageEngine):
     def supports_json_enforcing(self) -> bool:
         return self._json_enforcing_possible
 
-    def initialize_logits_processors(self, schemas: Generator[BaseModel]):
-        compiled_grammars = create_compiled_grammars(
-            schemas=schemas,
-            tokenizer=self.tokenizer,
-            vocab_size=self._model_config.vocab_size,
-            is_peft_adapter=self.is_peft_adapter,
-        )
-        self._logits_processors = [LogitsProcessor(list(compiled_grammars))]
-
     def generate(
         self, text: list[str], sampling_temperature: float, sampling_top_p: float
     ) -> tuple[list[int], EngineMetrics]:
@@ -124,3 +114,20 @@ class HuggingFaceEngine(LanguageEngine):
 
     def cleanup(self):
         pass
+
+    def update_json_constraints(self, schemas: list[BaseModel] | None) -> None:
+        """Update JSON schema constraints for the next generation call."""
+        if schemas:
+            compiled_grammars = create_compiled_grammars(
+                schemas=schemas,
+                tokenizer=self.tokenizer,
+                vocab_size=self._model_config.vocab_size,
+                is_peft_adapter=self.is_peft_adapter,
+            )
+            self._logits_processors = [LogitsProcessor(list(compiled_grammars))]
+        else:
+            self._logits_processors = None
+
+    def can_reuse_schemas(self) -> bool:
+        """HuggingFaceEngine cannot reuse LogitsProcessor across different batch sizes."""
+        return False

mostlyai_engine-1.5.7/mostlyai/engine/_language/engine/vllm_engine.py

@@ -0,0 +1,174 @@
+# Copyright 2025 MOSTLY AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+
+os.environ["VLLM_USE_V1"] = "1"
+
+
+import time
+from os import PathLike
+
+import torch
+from peft import PeftConfig
+from pydantic import BaseModel
+from transformers import AutoConfig, AutoTokenizer
+from vllm import LLM, SamplingParams
+from vllm.config import _get_and_verify_max_len
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.inputs.data import TokensPrompt
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import GuidedDecodingParams
+
+from mostlyai.engine._language.common import is_bf16_supported
+from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
+from mostlyai.engine._language.tokenizer_utils import tokenize_fn
+
+
+def get_dynamic_gpu_memory_utilization(utilization_ratio: float = 0.9) -> float:
+    """
+    Calculate dynamic GPU memory utilization based on available memory.
+
+    Args:
+        utilization_ratio: Fraction of available GPU memory to use (default: 0.9)
+
+    Returns:
+        GPU memory utilization as a fraction.
+    """
+    if not torch.cuda.is_available():
+        return utilization_ratio  # fallback for non-GPU environments
+
+    try:
+        # Get free and total memory from CUDA
+        free_memory, total_memory = torch.cuda.mem_get_info()
+
+        # Use specified ratio of free memory
+        target_memory = free_memory * utilization_ratio
+        utilization = target_memory / total_memory
+
+        # Ensure utilization is within reasonable bounds (0.1 to 0.95)
+        return max(0.1, min(0.95, utilization))
+
+    except Exception:
+        # Fallback to provided ratio if anything goes wrong
+        return utilization_ratio
+
+
+class VLLMEngine(LanguageEngine):
+    def __init__(
+        self, model_path: PathLike | str, device: torch.device, max_new_tokens: int, tokenizer_max_length: int
+    ):
+        self.device = device
+        self.tokenizer_max_length = tokenizer_max_length
+        self.max_new_tokens = max_new_tokens
+
+        peft_config = PeftConfig.from_pretrained(model_path)
+        base_config = AutoConfig.from_pretrained(peft_config.base_model_name_or_path)
+
+        model_path = str(model_path)
+        self._lora_request = LoRARequest("adapter", 1, model_path)
+        config_max_model_len = _get_and_verify_max_len(
+            base_config, tokenizer_config=None, max_model_len=None, disable_sliding_window=False, sliding_window=None
+        )
+
+        self.llm = LLM(
+            model=peft_config.base_model_name_or_path,
+            tokenizer=model_path,
+            max_model_len=min(config_max_model_len, self.tokenizer_max_length + max_new_tokens),
+            enable_lora=True,
+            dtype=torch.bfloat16 if is_bf16_supported(device) else torch.float16,
+            # enforce_eager=True,  # results in big slowdown, but is needed when running pytest locally
+            swap_space=0,
+            disable_log_stats=True,
+            tensor_parallel_size=torch.cuda.device_count(),
+            gpu_memory_utilization=get_dynamic_gpu_memory_utilization(),
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            padding_side="left",
+            truncation_side="left",
+            legacy=True,
+            # these must be False at initialization, as we manually add them later in tokenize_fn
+            add_bos_token=False,
+            add_eos_token=False,
+        )
+        self._prepared_schemas = None
+
+    def get_default_batch_size(self) -> int:
+        return 192
+
+    def supports_json_enforcing(self) -> bool:
+        return True
+
+    def generate(
+        self, text: list[str], sampling_temperature: float, sampling_top_p: float
+    ) -> tuple[list[int], EngineMetrics]:
+        tokenize_kwargs = dict(
+            tokenizer=self.tokenizer,
+            return_tensors=None,
+            add_bos_token=True,
+            add_eos_token=False,
+            padding=False,
+            truncation=True,
+            max_length=self.tokenizer_max_length,  # truncates input
+        )
+        t_tokenize = time.time()
+        inputs = tokenize_fn(text=text, **tokenize_kwargs)
+        tokenize_time = time.time() - t_tokenize
+
+        actual_batch_size = len(inputs["input_ids"])
+
+        # Create sampling params with guided decoding if schemas are prepared
+        effective_schemas = self._prepared_schemas
+
+        sampling_params = []
+        for i in range(actual_batch_size):
+            guided_decoding = None
+            if effective_schemas and i < len(effective_schemas):
+                # Convert Pydantic model to JSON schema for guided decoding
+                schema_dict = effective_schemas[i].model_json_schema()
+                guided_decoding = GuidedDecodingParams(json=schema_dict)
+
+            sampling_params.append(
+                SamplingParams(
+                    max_tokens=self.max_new_tokens,
+                    temperature=sampling_temperature,
+                    top_p=sampling_top_p,
+                    guided_decoding=guided_decoding,
+                )
+            )
+        t_generate = time.time()
+        outputs = self.llm.generate(
+            prompts=[TokensPrompt(prompt_token_ids=token_ids) for token_ids in inputs["input_ids"]],
+            sampling_params=sampling_params,
+            use_tqdm=False,
+            lora_request=self._lora_request,
+        )
+        generate_time = time.time() - t_generate
+        metrics = EngineMetrics(tokenize_time=tokenize_time, generate_time=generate_time)
+        return [r.outputs[0].token_ids for r in outputs], metrics
+
+    def cleanup(self):
+        del self.llm
+        cleanup_dist_env_and_memory()
+
+    def update_json_constraints(self, schemas: list[BaseModel] | None) -> None:
+        """Update JSON schema constraints for the next generation call."""
+        self._prepared_schemas = list(schemas) if schemas else None
+
+    def can_reuse_schemas(self) -> bool:
+        """VLLMEngine can handle variable batch sizes since it creates sampling params per sample."""
+        return True

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/mostlyai/engine/_language/generation.py

@@ -287,15 +287,18 @@ def generate(
         enforce_json_output = engine.supports_json_enforcing()
         _LOG.info(f"{enforce_json_output=}")
 
-        initialize_logits_processors_once = len(seeded_tgt_columns) == 0 and engine.__class__.__name__ == "VLLMEngine"
-        if enforce_json_output and initialize_logits_processors_once:
+        # Check if we can optimize by reusing schemas/constraints across batches
+        can_reuse_schemas = len(seeded_tgt_columns) == 0 and engine.can_reuse_schemas()
+
+        # Prepare schemas once if optimization is possible
+        if enforce_json_output and can_reuse_schemas:
             t0 = time.time()
-            schemas = create_schemas(
+            schemas_for_optimization = create_schemas(
                 size=batch_size,
                 stats=tgt_stats,
                 rare_category_replacement_method=rare_category_replacement_method,
             )
-            engine.initialize_logits_processors(schemas=schemas)
+            engine.update_json_constraints(schemas_for_optimization)
             total_logits_processor_build_time += time.time() - t0
 
         # keep at most 500k samples in memory before decoding and writing to disk
@@ -309,16 +312,21 @@ def generate(
             ctx_batch = ctx_data.iloc[samples_processed : samples_processed + batch_size]
             ctx_keys = ctx_batch[ctx_primary_key]
 
-            if enforce_json_output and not initialize_logits_processors_once:
+            # Update JSON constraints if needed per-batch (when schema reuse is not possible)
+            if enforce_json_output and not can_reuse_schemas:
                 t0 = time.time()
                 schemas = create_schemas(
                     seed_df=seed_data_batch,
                     stats=tgt_stats,
                     rare_category_replacement_method=rare_category_replacement_method,
                 )
-                engine.initialize_logits_processors(schemas=schemas)
+                engine.update_json_constraints(schemas)
                 total_logits_processor_build_time += time.time() - t0
+            elif not enforce_json_output:
+                # Clear any existing constraints if JSON enforcement is disabled
+                engine.update_json_constraints(None)
 
+            # Generate outputs using single generate method
             outputs, metrics = engine.generate(
                 encoded_ctx_batch["ctx"].tolist(),
                 sampling_temperature=sampling_temperature,

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/mostlyai/engine/_language/lstm.py

@@ -25,6 +25,11 @@ _LOG = logging.getLogger(__name__)
 class LSTMFromScratchConfig(PretrainedConfig):
     model_type = model_id = "MOSTLY_AI/LSTMFromScratch-3m"
 
+    # Map standard transformer attributes to our custom LSTM attributes
+    attribute_map = {
+        "num_hidden_layers": "num_layers",
+    }
+
     def __init__(
         self,
         vocab_size: int | None = None,

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/mostlyai/engine/_language/training.py

@@ -63,6 +63,7 @@ from mostlyai.engine._training_utils import (
     ModelCheckpoint,
     ProgressMessage,
     check_early_training_exit,
+    gpu_memory_cleanup,
 )
 from mostlyai.engine._workspace import Workspace, ensure_workspace_dir
 from mostlyai.engine.domain import DifferentialPrivacyConfig, ModelStateStrategy
@@ -277,6 +278,7 @@ def _gpu_estimate_max_batch_size(
     return batch_size
 
 
+@gpu_memory_cleanup
 def train(
     *,
     model: str = "MOSTLY_AI/LSTMFromScratch-3m",
@@ -865,4 +867,5 @@ def train(
             )
             # ensure everything gets uploaded
             upload_model_data_callback()
+
     _LOG.info(f"TRAIN_LANGUAGE finished in {time.time() - t0_:.2f}s")

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/mostlyai/engine/_tabular/training.py

@@ -63,6 +63,7 @@ from mostlyai.engine._training_utils import (
     ModelCheckpoint,
     ProgressMessage,
     check_early_training_exit,
+    gpu_memory_cleanup,
 )
 from mostlyai.engine._workspace import Workspace, ensure_workspace_dir
 from mostlyai.engine.domain import DifferentialPrivacyConfig, ModelStateStrategy
@@ -337,6 +338,7 @@ def _calculate_average_trn_loss(trn_sample_losses: list[torch.Tensor], n: int |
 ################
 
 
+@gpu_memory_cleanup
 def train(
     *,
     model: str = "MOSTLY_AI/Medium",
@@ -867,4 +869,5 @@ def train(
             progress.update(completed=steps, total=steps, message=progress_message)
             # ensure everything gets uploaded
             upload_model_data_callback()
+
     _LOG.info(f"TRAIN_TABULAR finished in {time.time() - t0:.2f}s")

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/mostlyai/engine/_training_utils.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import abc
+import functools
+import gc
 import logging
 import time
 
@@ -157,3 +159,18 @@ def check_early_training_exit(workspace: Workspace, trn_cnt: int, val_cnt: int)
     trn_files = workspace.encoded_data_trn.fetch_all()
     val_files = workspace.encoded_data_val.fetch_all()
     return any((len(trn_files) == 0, len(val_files) == 0, trn_cnt == 0, val_cnt == 0))
+
+
+def gpu_memory_cleanup(func):
+    """Decorator to clean up GPU memory after function execution."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        finally:
+            for _ in range(5):
+                gc.collect()
+            torch.cuda.empty_cache()
+
+    return wrapper

{mostlyai_engine-1.5.6 → mostlyai_engine-1.5.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "1.5.6"
+version = "1.5.7"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"
@@ -25,31 +25,31 @@ classifiers = [
 ]
 
 dependencies = [
-    "setuptools>=77.0.3",  # similar to vllm 0.9.1
+    "setuptools>=77.0.3",  # similar to vllm 0.10.1.1
     "numpy>=2.0.0",
     "pandas~=2.2.0",
     "pyarrow>=16.0.0",
     "joblib>=1.4.2",
     "psutil>=5.9.5,<6",  # upgrade when colab psutil is updated
     "tokenizers>=0.21.0",
-    "transformers>=4.51.0,<4.54.0",  # update when vllm is updated to >0.10.1
+    "transformers>=4.55.0",  # updated (removed ,<4.54.0), similar to vllm 0.10.1.1
     "datasets>=3.0.0",
     "accelerate>=1.5.0",
     "peft>=0.12.0",
     "huggingface-hub[hf-xet]>=0.30.2",
     "opacus>=1.5.4",
-    "xgrammar>=0.1.19",  # for vllm 0.9.1 compatibility
+    "xgrammar>=0.1.21",  # for vllm 0.10.1.1 compatibility (=0.1.21)
     "json-repair>=0.47.0",
-    "torch>=2.7.0,<2.7.1",  # for vllm 0.9.1 compatibility
-    "torchaudio>=2.7.0,<2.7.1",  # for vllm 0.9.1 compatibility
-    "torchvision>=0.22.0,<0.22.1"  # for vllm 0.9.1 compatibility
+    "torch>=2.8.0,<2.9.0",  # for vllm 0.10.2 compatibility
+    "torchaudio>=2.8.0,<2.9.0",  # for vllm 0.10.2 compatibility
+    "torchvision>=0.23.0,<0.24.0"  # for vllm 0.10.2 compatibility
 ]
 
 [project.optional-dependencies]
 gpu = [
     "bitsandbytes==0.42.0; sys_platform == 'darwin'",
     "bitsandbytes>=0.45.5; sys_platform == 'linux'",
-    "vllm==0.9.1; sys_platform == 'linux' or sys_platform == 'darwin'",
+    "vllm==0.10.2; sys_platform == 'linux' or sys_platform == 'darwin'",
 ]
 
 [dependency-groups]

mostlyai_engine-1.5.6/mostlyai/engine/_language/engine/vllm_engine.py

@@ -1,223 +0,0 @@
-# Copyright 2025 MOSTLY AI
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import os
-
-# Per-Request Logits Processors are not supported in V1 (https://docs.vllm.ai/en/latest/getting_started/v1_user_guide.html#feature-model)
-# Global Logits Processors will be new mechanism to use, but it is not yet released
-# Here is PR for it: https://github.com/vllm-project/vllm/pull/13360
-os.environ["VLLM_USE_V1"] = "0"
-
-import contextlib
-import gc
-import time
-from collections.abc import Generator
-from os import PathLike
-
-import torch
-import xgrammar as xgr
-from peft import PeftConfig
-from pydantic import BaseModel
-from transformers import AutoConfig, AutoTokenizer
-from vllm import LLM, SamplingParams
-from vllm.config import _get_and_verify_max_len
-from vllm.distributed import destroy_distributed_environment, destroy_model_parallel
-from vllm.lora.request import LoRARequest
-from vllm.platforms import current_platform
-
-from mostlyai.engine._language.common import is_bf16_supported
-from mostlyai.engine._language.engine.base import EngineMetrics, LanguageEngine
-from mostlyai.engine._language.tokenizer_utils import tokenize_fn
-from mostlyai.engine._language.xgrammar_utils import create_compiled_grammars
-
-
-def cleanup_dist_env_and_memory():
-    """Copy from current main of vllm replace by import when possible"""
-    destroy_model_parallel()
-    destroy_distributed_environment()
-    with contextlib.suppress(AssertionError):
-        torch.distributed.destroy_process_group()
-    gc.collect()
-    if not current_platform.is_cpu():
-        torch.cuda.empty_cache()
-
-
-class XGrammarLogitsProcessor:
-    """
-    Inspired by [XGrammarLogitsProcessor](https://github.com/vllm-project/vllm/blob/a43aa183dc0cb2639044c15d272e0ce1941392b0/vllm/model_executor/guided_decoding/xgrammar_decoding.py#L280).
-    VLLM's XGrammarLogitsProcessor can be reused.
-    """
-
-    def __init__(self, compiled_grammar: xgr.CompiledGrammar):
-        self.compiled_grammar = compiled_grammar
-        self.tokenizer_info = compiled_grammar.tokenizer_info
-        self.batch_size = 1
-
-        self.matchers: list[xgr.GrammarMatcher] | None = None
-        self.token_bitmask: torch.Tensor | None = None
-        self.prefilled = False
-
-    def __call__(self, input_ids: tuple[int], scores: torch.Tensor) -> torch.Tensor:
-        # lazily initialize GrammarMatchers and bitmask
-        if self.matchers is None:
-            self.matchers = [xgr.GrammarMatcher(self.compiled_grammar) for _ in range(self.batch_size)]
-            self.token_bitmask = xgr.allocate_token_bitmask(self.batch_size, self.tokenizer_info.vocab_size)
-
-        if not self.prefilled:
-            # have not sampled a token yet
-            self.prefilled = True
-        else:
-            for i, matcher in enumerate(self.matchers):
-                if not matcher.is_terminated():
-                    sampled_token = input_ids[-1]
-                    assert self.matchers[i].accept_token(sampled_token)
-
-        for i, matcher in enumerate(self.matchers):
-            if not matcher.is_terminated():
-                matcher.fill_next_token_bitmask(self.token_bitmask, i)
-
-        # token_bitmask is a CPU tensor for use with accept_token and
-        # fill_next_token_bitmask so we move it to the device of scores
-        device_type = scores.device.type
-        dtype = scores.dtype
-        if device_type != "cuda":
-            # xgrammar on cpu only supports float32 scores
-            # see: https://github.com/mlc-ai/xgrammar/blob/c1b64920cad24f44f235778c1c00bb52d57da01a/python/xgrammar/kernels/apply_token_bitmask_inplace_cpu.py#L22
-            scores = scores.to("cpu").float().unsqueeze(0)
-
-        # Note: In this method, if the tensors have different dimensions
-        # on CPU device fails, but on GPU it runs without error. Hence the
-        # unsqueeze above for scores, to match the token bitmask shape
-        xgr.apply_token_bitmask_inplace(scores, self.token_bitmask.to(scores.device, non_blocking=True))
-        if device_type != "cuda":
-            scores = scores.to(dtype).to(device_type).squeeze()
-        return scores
-
-    def clone(self) -> XGrammarLogitsProcessor:
-        """Create a new instance with shared compiled grammar but separate state"""
-        new_processor = XGrammarLogitsProcessor(self.compiled_grammar)
-
-        # create fresh matchers for the new sequence
-        if self.matchers is not None:
-            new_processor.matchers = [xgr.GrammarMatcher(self.compiled_grammar) for _ in range(self.batch_size)]
-
-        # create a new token bitmask with the same size
-        if self.token_bitmask is not None:
-            new_processor.token_bitmask = self.token_bitmask
-
-        new_processor.batch_size = self.batch_size
-
-        # reset prefilled state for new sequence
-        new_processor.prefilled = False
-
-        return new_processor
-
-
-class VLLMEngine(LanguageEngine):
-    def __init__(
-        self, model_path: PathLike | str, device: torch.device, max_new_tokens: int, tokenizer_max_length: int
-    ):
-        self.device = device
-        self.tokenizer_max_length = tokenizer_max_length
-        self.max_new_tokens = max_new_tokens
-
-        peft_config = PeftConfig.from_pretrained(model_path)
-        base_config = AutoConfig.from_pretrained(peft_config.base_model_name_or_path)
-
-        model_path = str(model_path)
-        self._lora_request = LoRARequest("adapter", 1, model_path)
-        config_max_model_len = _get_and_verify_max_len(
-            base_config, max_model_len=None, disable_sliding_window=None, sliding_window_len=None
-        )
-        self.llm = LLM(
-            model=peft_config.base_model_name_or_path,
-            tokenizer=model_path,
-            device=device.type,
-            max_model_len=min(config_max_model_len, self.tokenizer_max_length + max_new_tokens),
-            enable_lora=True,
-            dtype=torch.bfloat16 if is_bf16_supported(device) else torch.float16,
-            # enforce_eager=True,  # results in big slowdown, but is needed when running pytest locally
-            swap_space=0,
-            disable_log_stats=True,
-            tensor_parallel_size=torch.cuda.device_count(),
-        )
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            model_path,
-            padding_side="left",
-            truncation_side="left",
-            legacy=True,
-            # these must be False at initialization, as we manually add them later in tokenize_fn
-            add_bos_token=False,
-            add_eos_token=False,
-        )
-        self._logits_processors = None
-
-    def get_default_batch_size(self) -> int:
-        return 192
-
-    def supports_json_enforcing(self) -> bool:
-        return True
-
-    def initialize_logits_processors(self, schemas: Generator[BaseModel]):
-        compiled_grammars = create_compiled_grammars(
-            schemas=schemas,
-            tokenizer=self.llm.get_tokenizer(),
-            vocab_size=self.llm.llm_engine.get_model_config().get_vocab_size(),
-            is_peft_adapter=True,
-        )
-        self._logits_processors = [XGrammarLogitsProcessor(compiled_grammar) for compiled_grammar in compiled_grammars]
-
-    def generate(
-        self, text: list[str], sampling_temperature: float, sampling_top_p: float
-    ) -> tuple[list[int], EngineMetrics]:
-        tokenize_kwargs = dict(
-            tokenizer=self.tokenizer,
-            return_tensors=None,
-            add_bos_token=True,
-            add_eos_token=False,
-            padding=False,
-            truncation=True,
-            max_length=self.tokenizer_max_length,  # truncates input
-        )
-        t_tokenize = time.time()
-        inputs = tokenize_fn(text=text, **tokenize_kwargs)
-        tokenize_time = time.time() - t_tokenize
-
-        actual_batch_size = len(inputs["input_ids"])
-        sampling_params = [
-            SamplingParams(
-                max_tokens=self.max_new_tokens,
-                temperature=sampling_temperature,
-                top_p=sampling_top_p,
-                logits_processors=[lp],
-            )
-            for lp in self._logits_processors[:actual_batch_size]
-        ]
-        t_generate = time.time()
-        outputs = self.llm.generate(
-            prompts=None,
-            prompt_token_ids=inputs["input_ids"],
-            sampling_params=sampling_params,
-            use_tqdm=False,
-            lora_request=self._lora_request,
-        )
-        generate_time = time.time() - t_generate
-        metrics = EngineMetrics(tokenize_time=tokenize_time, generate_time=generate_time)
-        return [r.outputs[0].token_ids for r in outputs], metrics
-
-    def cleanup(self):
-        del self.llm
-        cleanup_dist_env_and_memory()