tpu-inference 0.11.1__py3-none-any.whl

tpu_inference/interfaces/cache.py

@@ -0,0 +1,31 @@
+"""
+Defines the abstract contracts for cache managers.
+"""
+from typing import Protocol
+
+
+class IKVCacheManager(Protocol):
+    """
+    Abstract contract for a KVCacheManager.
+    """
+    # Add methods and properties from vllm.v1.core.kv_cache_manager.KVCacheManager
+    # that tpu_inference actually uses.
+    ...
+
+
+class IEncoderCacheManager(Protocol):
+    """
+    Abstract contract for an EncoderCacheManager.
+    """
+    # Add methods and properties from vllm.v1.core.encoder_cache_manager.EncoderCacheManager
+    # that tpu_inference actually uses.
+    ...
+
+
+class IMirroredProcessingCache(Protocol):
+    """
+    Abstract contract for a MirroredProcessingCache.
+    """
+    # Add methods and properties from vllm.v1.engine.mm_input_cache.MirroredProcessingCache
+    # that tpu_inference actually uses.
+    ...

tpu_inference/interfaces/config.py

@@ -0,0 +1,47 @@
+"""
+Defines the abstract contract for a configuration object.
+"""
+from typing import Any, Optional, Protocol
+
+from .config_parts import (ICacheConfig, ICompilationConfig, IModelConfig,
+                           IParallelConfig, ISchedulerConfig,
+                           ISpeculativeConfig)
+
+
+class IConfig(Protocol):
+    """
+    A minimal, abstract interface for a configuration object.
+
+    This protocol defines only the methods and properties that tpu_inference
+    requires to operate. Client libraries (like vLLM) will provide concrete
+    implementations that satisfy this contract.
+    """
+
+    @property
+    def cache_config(self) -> ICacheConfig:
+        ...
+
+    @property
+    def compilation_config(self) -> ICompilationConfig:
+        ...
+
+    @property
+    def model_config(self) -> Optional[IModelConfig]:
+        ...
+
+    @property
+    def parallel_config(self) -> IParallelConfig:
+        ...
+
+    @property
+    def scheduler_config(self) -> ISchedulerConfig:
+        ...
+
+    @property
+    def speculative_config(self) -> Optional[ISpeculativeConfig]:
+        ...
+
+    # Escape hatch for direct access when needed by the adapter.
+    @property
+    def vllm_config(self) -> Any:
+        ...

tpu_inference/interfaces/config_parts.py

@@ -0,0 +1,117 @@
+"""
+Defines the abstract contracts for the component parts of an IConfig.
+"""
+from typing import Any, Optional, Protocol
+
+import torch
+
+
+class IModelConfig(Protocol):
+
+    @property
+    def dtype(self) -> torch.dtype:
+        ...
+
+    @dtype.setter
+    def dtype(self, value: torch.dtype) -> None:
+        ...
+
+    @property
+    def use_mla(self) -> bool:
+        ...
+
+
+class ICacheConfig(Protocol):
+
+    @property
+    def block_size(self) -> Optional[int]:
+        ...
+
+    @block_size.setter
+    def block_size(self, value: Optional[int]) -> None:
+        ...
+
+
+class IParallelConfig(Protocol):
+
+    @property
+    def worker_cls(self) -> str:
+        ...
+
+    @worker_cls.setter
+    def worker_cls(self, value: str) -> None:
+        ...
+
+
+class ISchedulerConfig(Protocol):
+
+    @property
+    def max_num_seqs(self) -> int:
+        ...
+
+    @property
+    def is_multi_step(self) -> bool:
+        ...
+
+    @property
+    def is_multimodal_model(self) -> bool:
+        ...
+
+    @property
+    def disable_chunked_mm_input(self) -> bool:
+        ...
+
+    @disable_chunked_mm_input.setter
+    def disable_chunked_mm_input(self, value: bool) -> None:
+        ...
+
+    @property
+    def enable_chunked_prefill(self) -> bool:
+        ...
+
+    @enable_chunked_prefill.setter
+    def enable_chunked_prefill(self, value: bool) -> None:
+        ...
+
+    @property
+    def chunked_prefill_enabled(self) -> bool:
+        ...
+
+    @chunked_prefill_enabled.setter
+    def chunked_prefill_enabled(self, value: bool) -> None:
+        ...
+
+    @property
+    def max_model_len(self) -> int:
+        ...
+
+    @property
+    def max_num_batched_tokens(self) -> int:
+        ...
+
+    @max_num_batched_tokens.setter
+    def max_num_batched_tokens(self, value: int) -> None:
+        ...
+
+
+class ICompilationConfig(Protocol):
+
+    @property
+    def level(self) -> Any:
+        ...
+
+    @level.setter
+    def level(self, value: Any) -> None:
+        ...
+
+    @property
+    def backend(self) -> str:
+        ...
+
+    @backend.setter
+    def backend(self, value: str) -> None:
+        ...
+
+
+class ISpeculativeConfig(Protocol):
+    ...

tpu_inference/interfaces/engine.py

@@ -0,0 +1,51 @@
+"""
+This module defines the engine interface contracts required by tpu_inference.
+"""
+
+from typing import TYPE_CHECKING, Any, Protocol
+
+# tpu_inference now depends on its own, locally defined interfaces.
+from .cache import IMirroredProcessingCache
+from .outputs import IStructuredOutputManager
+from .scheduler import IScheduler
+
+# This block is only processed by type checkers, not at runtime.
+if TYPE_CHECKING:
+    from .outputs import IModelRunnerOutput
+
+
+class IEngineProc(Protocol):
+    """
+    A high-level interface for any process that can be launched by a client.
+    It defines the single entry point for starting the process's main loop.
+    """
+
+    def run_busy_loop(self) -> None:
+        ...
+
+
+class IDisaggEngineCoreProc(IEngineProc):
+    """
+    An interface for the disaggregated engine process. It inherits the common
+    IEngineProc contract.
+    """
+    pass
+
+
+class IEngineCore(Protocol):
+    """
+    An interface defining the contract for an Engine Core building block.
+    This mirrors the public API of a vLLM Engine Core that is used by the
+    DisaggEngineCoreProc.
+    """
+    scheduler: IScheduler
+    mm_input_cache_server: IMirroredProcessingCache
+    structured_output_manager: IStructuredOutputManager
+    model_executor: Any
+
+    def execute_model_with_error_logging(self, *args,
+                                         **kwargs) -> "IModelRunnerOutput":
+        ...
+
+    def shutdown(self) -> None:
+        ...

tpu_inference/interfaces/outputs.py

@@ -0,0 +1,22 @@
+"""
+Defines the abstract contracts for model and structured outputs.
+"""
+from typing import Protocol
+
+
+class IModelRunnerOutput(Protocol):
+    """
+    Abstract contract for the output of a model runner.
+    """
+    # Add methods and properties from vllm.v1.outputs.ModelRunnerOutput
+    # that tpu_inference actually uses.
+    ...
+
+
+class IStructuredOutputManager(Protocol):
+    """
+    Abstract contract for a StructuredOutputManager.
+    """
+    # Add methods and properties from vllm.v1.structured_output.StructuredOutputManager
+    # that tpu_inference actually uses.
+    ...

tpu_inference/interfaces/params.py

@@ -0,0 +1,21 @@
+"""
+Defines the abstract contracts for sampling and pooling parameters.
+"""
+from typing import Any, Protocol
+
+
+class IPoolingParams(Protocol):
+    """
+    Abstract contract for PoolingParams.
+    """
+    ...
+
+
+class ISamplingParams(Protocol):
+    """
+    Abstract contract for SamplingParams.
+    """
+
+    @property
+    def sampling_type(self) -> Any:
+        ...

tpu_inference/interfaces/platform.py

@@ -0,0 +1,74 @@
+"""
+Defines the abstract contract for a hardware platform.
+"""
+from typing import Any, Optional, Protocol, Union
+
+import torch
+
+from .config import IConfig
+from .params import IPoolingParams, ISamplingParams
+
+
+class IPlatform(Protocol):
+    """
+    A minimal, abstract interface for a hardware platform.
+    """
+
+    def can_update_inplace(self) -> bool:
+        ...
+
+    def check_and_update_config(self, vllm_config: IConfig) -> None:
+        ...
+
+    def get_attn_backend_cls(self, selected_backend: Any, head_size: int,
+                             dtype: torch.dtype, kv_cache_dtype: Optional[str],
+                             block_size: int, use_v1: bool, use_mla: bool,
+                             has_sink: bool, use_spare: bool) -> str:
+        ...
+
+    def get_device_communicator_cls(self) -> str:
+        ...
+
+    def get_device_name(self, device_id: int = 0) -> str:
+        ...
+
+    def get_device_total_memory(self, device_id: int = 0) -> int:
+        ...
+
+    def get_infinity_values(self, dtype: torch.dtype) -> tuple[float, float]:
+        ...
+
+    def get_lora_vocab_padding_size(self) -> int:
+        ...
+
+    def get_punica_wrapper(self) -> str:
+        ...
+
+    def inference_mode(self) -> Any:
+        ...
+
+    def is_async_output_supported(self, enforce_eager: Optional[bool]) -> bool:
+        ...
+
+    def is_kv_cache_dtype_supported(self, kv_cache_dtype: str) -> bool:
+        ...
+
+    def is_pin_memory_available(self) -> bool:
+        ...
+
+    def set_device(self, device: torch.device) -> None:
+        ...
+
+    def supports_v1(self, model_config: Any) -> bool:
+        ...
+
+    def use_all_gather(self) -> bool:
+        ...
+
+    def validate_request(
+        self,
+        prompt: Any,
+        params: Union[ISamplingParams, IPoolingParams],
+        processed_inputs: Any,
+    ) -> None:
+        ...

tpu_inference/interfaces/request.py

@@ -0,0 +1,39 @@
+"""
+Defines the abstract contract for a Request.
+"""
+from typing import Any, Protocol
+
+
+class IRequest(Protocol):
+    """
+    A minimal, abstract interface for a request.
+
+    This protocol defines only the methods and properties that tpu_inference
+    requires to operate. Client libraries (like vLLM) will provide concrete
+    implementations that satisfy this contract.
+    """
+
+    @property
+    def vllm_request(self) -> Any:
+        ...
+
+    def is_finished(self) -> bool:
+        ...
+
+    def get_request_id(self) -> str:
+        ...
+
+    # Add mm_hashes. it's used by `if request.mm_hashes is not None:`.
+
+    # Add other methods and properties from vllm.v1.request.Request that are
+    # actually used by the orchestration logic.
+    # For example:
+    # @property
+    # def prompt(self) -> str: ...
+    #
+    # @property
+    # def prompt_token_ids(self) -> list[int]: ...
+    #
+    # def is_finished(self) -> bool: ...
+    #
+    # ... etc.

tpu_inference/interfaces/scheduler.py

@@ -0,0 +1,31 @@
+"""
+This module defines the scheduler interface contract required by tpu_inference.
+"""
+
+from typing import Dict, Protocol
+
+# tpu_inference now depends on its own, locally defined interfaces.
+from .cache import IEncoderCacheManager, IKVCacheManager
+from .request import IRequest
+
+
+class IScheduler(Protocol):
+    """
+    An extended interface for a scheduler, tailored to the needs
+    of advanced orchestration engines.
+
+    This contract is defined by tpu_inference and must be implemented by
+    any client library (like vLLM) that wishes to use this orchestrator.
+    """
+
+    @property
+    def requests(self) -> Dict[str, IRequest]:
+        ...
+
+    @property
+    def kv_cache_manager(self) -> IKVCacheManager:
+        ...
+
+    @property
+    def encoder_cache_manager(self) -> IEncoderCacheManager:
+        ...