Add scheduler config support

src/cli/groups/add.py

@@ -4,6 +4,9 @@
 import json
 
 import click
+from pydantic import ValidationError
+
+from src.server.schemas.modeling.contract_ovgenai_llm_and_vlm import SchedulerConfigSchema
 
 from ..main import cli, console
 from ..utils import validate_model_path
@@ -14,7 +17,7 @@
     required=True,
     help='Public facing name of the model.')
 @click.option('--model-path', '--m',
-    required=True, 
+    required=True,
     help='Path to OpenVINO IR converted model.')
 @click.option('--engine', '--en',
     type=click.Choice(['ovgenai', 'openvino', 'optimum']),
@@ -34,6 +37,9 @@
 @click.option("--runtime-config", "--rtc",
     default=None,
     help='OpenVINO runtime configuration as JSON string (e.g., \'{"MODEL_DISTRIBUTION_POLICY": "PIPELINE_PARALLEL"}\').')
+@click.option("--scheduler-config", "-sc",
+    default=None,
+    help='OpenVINO runtime scheduler configuration as JSON string (e.g., \'{"use_sparse_attention": true}\').')
 @click.option('--cache-dir', '--cd',
     required=False,
     default=None,
@@ -57,14 +63,14 @@
     type=float,
     help='Confidence threshold for accepting draft tokens.')
 @click.pass_context
-def add(ctx, model_path, model_name, engine, model_type, device, runtime_config, cache_dir, draft_model_path, draft_device, num_assistant_tokens, assistant_confidence_threshold):
+def add(ctx, model_path, model_name, engine, model_type, device, runtime_config, scheduler_config, cache_dir, draft_model_path, draft_device, num_assistant_tokens, assistant_confidence_threshold):
     """- Add a model configuration to the config file."""
-    
+
     # Validate model path
     if not validate_model_path(model_path):
         console.print(f"[red]Model file check failed! {model_path} does not contain openvino model files OR your chosen path is malformed. Verify chosen path is correct and acquired model files match source on the hub, or the destination of converted model.[/red]")
         ctx.exit(1)
-    
+
     # Parse runtime_config if provided
     parsed_runtime_config = {}
     if runtime_config:
@@ -78,17 +84,34 @@ def add(ctx, model_path, model_name, engine, model_type, device, runtime_config,
             console.print(f"[red]Error parsing runtime_config JSON:[/red] {e}")
             console.print('[yellow]Example format: \'{"MODEL_DISTRIBUTION_POLICY": "PIPELINE_PARALLEL"}\'[/yellow]')
             ctx.exit(1)
-
+    parsed_scheduler_config = {}
+    if scheduler_config:
+        # Let the model validate the JSON itself. If it validates, assume we can safely load the JSON.
+        try:
+            parsed_scheduler_config = json.loads(scheduler_config)
+            if not isinstance(parsed_scheduler_config, dict):
+                console.print(f"[red]Error: scheduler_config must be a JSON object (dictionary), got {type(scheduler_config).__name__}[/red]")
+                console.print('[yellow]Example format: \'{"max_num_batched_tokens": 256, "enable_prefix_caching": true}\'[/yellow]')
+            SchedulerConfigSchema.model_validate_json(scheduler_config)
+        except ValidationError as e:
+                console.print("[red]Error: Failed validating scheduler_config:[/red]")
+                console.print('[yellow]Example format: \'{"max_num_batched_tokens": 256, "enable_prefix_caching": true}\'[/yellow]')
+                console.print('')
+                console.print('[yellow]Error:[/yellow]')
+                console.print(e)
+                ctx.exit(1)
+
     # Legacy configs may still contain vlm_type, but new configs resolve VLM tokens from config.json.
     load_config = {
         "model_name": model_name,
-        "model_path": model_path,  
-        "model_type": model_type,  
-        "engine": engine,    
+        "model_path": model_path,
+        "model_type": model_type,
+        "engine": engine,
         "device": device,
         "runtime_config": parsed_runtime_config,
+        "scheduler_config": parsed_scheduler_config,
     }
-    
+
     # Store the cache directory (resolved relative to the config file at load time)
     if cache_dir:
         load_config["cache_dir"] = cache_dir
@@ -105,7 +128,7 @@ def add(ctx, model_path, model_name, engine, model_type, device, runtime_config,
         load_config["num_assistant_tokens"] = num_assistant_tokens
     if assistant_confidence_threshold is not None:
         load_config["assistant_confidence_threshold"] = assistant_confidence_threshold
-    
+
     ctx.obj.server_config.save_model_config(model_name, load_config)
     console.print(f"[green]Model configuration saved:[/green] {model_name}")
     console.print(f"[dim]Use 'openarc load {model_name}' to load this model.[/dim]")

src/engine/ov_genai/llm.py

@@ -1,3 +1,4 @@
+from src.engine.ov_genai.utils import extract_scheduler_config_from_loader
 import asyncio
 import gc
 import logging
@@ -265,6 +266,7 @@ def load_model(self, loader: ModelLoadConfig):
             self.model_assistant_confidence_threshold = None
 
         pipeline_kwargs = {**(loader.runtime_config or {})}
+        scheduler_config = extract_scheduler_config_from_loader(loader)
         if loader.cache_dir:
             pipeline_kwargs['CACHE_DIR'] = loader.cache_dir
         if draft_model is not None:
@@ -273,7 +275,8 @@ def load_model(self, loader: ModelLoadConfig):
         self.model = LLMPipeline(
             loader.model_path,
             loader.device,
-            **pipeline_kwargs
+            **scheduler_config,
+            **pipeline_kwargs,
         )
 
         self.encoder_tokenizer = AutoTokenizer.from_pretrained(loader.model_path)

src/engine/ov_genai/utils.py

@@ -0,0 +1,56 @@
+"""OV GenAI engine utilities."""
+import logging
+from typing import Literal
+
+from openvino_genai import SchedulerConfig
+
+from src.server.schemas.modeling.contract_ovgenai_llm_and_vlm import SchedulerConfigSchema
+from src.server.schemas.registration import ModelLoadConfig
+
+logger = logging.getLogger(__name__)
+
+def generate_ov_scheduler_config(scheduler_config: SchedulerConfigSchema) -> dict:
+  """Generates a SchedulerConfig object from the scheduler config model.
+
+     Note: `scheduler_config` cannot be passed to SDPA pipelines without raising
+     an error. Ensure you test if the pipeline configuration is set to SDPA first
+     by using methods such as `extract_scheduler_config_from_loader`.
+  """
+  sched_config = SchedulerConfig()
+  if scheduler_config.max_num_batched_tokens:
+    sched_config.max_num_batched_tokens = scheduler_config.max_num_batched_tokens
+  if scheduler_config.num_kv_blocks:
+    sched_config.num_kv_blocks = scheduler_config.num_kv_blocks
+  if scheduler_config.cache_size:
+    sched_config.cache_size = scheduler_config.cache_size
+  if scheduler_config.num_linear_attention_blocks:
+    sched_config.num_linear_attention_blocks = scheduler_config.num_linear_attention_blocks
+  if scheduler_config.cache_interval_multiplier:
+    sched_config.cache_interval_multiplier = scheduler_config.cache_interval_multiplier
+  if scheduler_config.dynamic_split_fuse:
+    sched_config.dynamic_split_fuse = scheduler_config.dynamic_split_fuse
+  if scheduler_config.max_num_seqs:
+    sched_config.max_num_seqs = scheduler_config.max_num_seqs
+  if scheduler_config.enable_prefix_caching:
+    sched_config.enable_prefix_caching = scheduler_config.enable_prefix_caching
+  if scheduler_config.use_cache_eviction:
+    sched_config.use_cache_eviction = scheduler_config.use_cache_eviction
+  if scheduler_config.use_sparse_attention:
+    sched_config.use_sparse_attention = scheduler_config.use_sparse_attention
+  return {"scheduler_config": sched_config}
+
+def extract_scheduler_config_from_loader(loader: ModelLoadConfig) -> dict[Literal["scheduler_config"], SchedulerConfig]:
+  """Extract the scheduler configuration from the loader config and return as a dict to be piped to the pipeline.
+
+     If pipeline is SDPA, returns an empty dictionary and raises an error to the user.
+     Otherwise, returns a dictonary with the SchedulerConfig object
+  """
+  pipeline_kwargs = loader.runtime_config or {}
+  sched_config = loader.scheduler_config or SchedulerConfigSchema()
+  sched_config_dict = sched_config.model_dump(exclude_unset=True)
+  if pipeline_kwargs.get("ATTENTION_BACKEND") == "SDPA" and sched_config_dict:
+    logger.error("Cannot set scheduler_config for model: scheduler config is unsupported for SDPA backends")
+    return {}
+  if not sched_config_dict:
+    return {}
+  return generate_ov_scheduler_config(sched_config)

src/engine/ov_genai/vlm.py

@@ -1,3 +1,4 @@
+from src.engine.ov_genai.utils import extract_scheduler_config_from_loader
 import asyncio
 import base64
 import gc
@@ -271,14 +272,16 @@ def load_model(self, loader: ModelLoadConfig):
         """
         try:
             logger.info(f"{loader.model_type} on {loader.device} with {loader.runtime_config}")
-
+
+            scheduler_config = extract_scheduler_config_from_loader(loader)
             pipeline_kwargs = {**(loader.runtime_config or {})}
             if loader.cache_dir:
                 pipeline_kwargs['CACHE_DIR'] = loader.cache_dir
 
             self.model_path = VLMPipeline(
                 loader.model_path,
                 loader.device,
+                **scheduler_config,
                 **pipeline_kwargs
             )
 

src/server/schemas/modeling/contract_ovgenai_llm_and_vlm.py

@@ -52,7 +52,7 @@ class OVGenAI_GenConfig(BaseModel):
         default=None,
         description="Confidence threshold for accepting draft tokens (typically 0.3-0.5)"
     )
-    
+
     stream: bool = Field(
         default=False,
         description="Stream output in chunks of tokens."
@@ -91,3 +91,63 @@ def text_messages(self) -> List[Dict[str, Any]]:
         """Messages with their `content` coerced to plain strings for text models."""
 
         return flatten_messages(self.messages)
+
+
+class SchedulerConfigSchema(BaseModel):
+    """Model for OV scheduler config."""
+
+    max_num_batched_tokens: Optional[int] = Field(
+        default=None,
+        description=(
+            "Maximum number of tokens to batch (in contrast to max_batch_size which "
+        "combines independent sequences, we consider total amount of tokens in a batch)."
+    ))
+    num_kv_blocks: Optional[int] = Field(
+        default=None,
+        description="Total number of KV blocks available to scheduler logic.",
+    )
+    cache_size: Optional[int] = Field(
+        default=None,
+        description="Total size of cache in GB."
+    )
+    num_linear_attention_blocks: Optional[int] = Field(
+        default=None,
+        description="Total number of linear attention blocks available to scheduler logic. Only applicable for models with linear attention cache inputs."
+    )
+    cache_interval_multiplier: Optional[int] = Field(
+        default=None,
+        description="""
+        Optional multiplier used to derive the linear-attention checkpoint interval for prefix caching.
+        The internal interval is KV cache block size * cache_interval_multiplier.
+        When unset, the default value 8 is used for hybrid models with prefix caching.
+        Explicit values are supported only for models with linear attention cache inputs.
+        0 is valid only when prefix caching is disabled.
+        """
+    )
+    dynamic_split_fuse: Optional[bool] = Field(
+        default=None,
+        description="Whether to split prompt / generate to different scheduling phases."
+    )
+    max_num_seqs: Optional[int] = Field(
+        default=None,
+        description="Max number of scheduled sequences (you can think of it as \"max batch size\")."
+    )
+    enable_prefix_caching: Optional[bool] = Field(
+        default=None,
+        description="""
+        Enable caching of KV-blocks.
+        When turned on all previously calculated KV-caches are kept in memory for future usages.
+        KV-caches can be overridden if KV-cache limit is reached, but blocks are not released.
+        This results in more RAM usage, maximum RAM usage is determined by cache_size or num_kv_blocks parameters.
+        When turned off only KV-cache required for batch calculation is kept in memory and
+        when a sequence has finished generation its cache is released.
+        """
+    )
+    use_cache_eviction: Optional[bool] = Field(
+        default=None,
+        description="Whether to use cache eviction during generation."
+    )
+    use_sparse_attention: Optional[bool] = Field(
+        default=None,
+        description="Whether to use sparse attention during prefill."
+    )

src/server/schemas/registration.py

@@ -1,8 +1,11 @@
+
 from enum import Enum
 from typing import Any, Dict, Optional
 
 from pydantic import BaseModel, Field
 
+from src.server.schemas.modeling.contract_ovgenai_llm_and_vlm import SchedulerConfigSchema
+
 
 class ModelStatus(str, Enum):
     """loading status.
@@ -111,7 +114,11 @@ class ModelLoadConfig(BaseModel):
         default=None,
         description="Default assistant_confidence_threshold for speculative decoding with this model"
     )
-
+    scheduler_config: Optional[SchedulerConfigSchema] = Field(
+        default=None,
+        description="Optional OpenVINO scheduler properties.",
+    )
+
 
 class ModelUnloadConfig(BaseModel):
     model_name: str = Field(..., description="Name of the model to unload")