From 2ea6d7e24862f27c72225f6a44c6d1b17ec34ddc Mon Sep 17 00:00:00 2001 From: Michael Carroll Date: Thu, 11 Jun 2026 09:45:12 -0400 Subject: [PATCH 1/6] Create utility function to load scheduler config --- src/engine/ov_genai/utils.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 src/engine/ov_genai/utils.py diff --git a/src/engine/ov_genai/utils.py b/src/engine/ov_genai/utils.py new file mode 100644 index 0000000..1594c1b --- /dev/null +++ b/src/engine/ov_genai/utils.py @@ -0,0 +1,30 @@ +"""OV GenAI engine utilities.""" +from src.server.models.ov_genai import SchedulerConfigSchema + +def generate_ov_scheduler_config(scheduler_config: SchedulerConfigSchema) -> dict: + """Generates a SchedulerConfig object from the scheduler config model.""" + # SchedulerConfig cannot be used as SDPA pipelines refuse to accept a scheduler config + # which includes any invalid properties, including those set to `None` by default. + # Fortunately, the pipelines also accept a dictionary. + sched_config = {} + 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 sched_config \ No newline at end of file From 00ea0690a8a8649b79396e45ac48875afc94efbc Mon Sep 17 00:00:00 2001 From: Michael Carroll Date: Wed, 5 Aug 2026 09:46:48 -0400 Subject: [PATCH 2/6] Use scheduler config object --- src/engine/ov_genai/utils.py | 34 ++++++++++++++++++---------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/src/engine/ov_genai/utils.py b/src/engine/ov_genai/utils.py index 1594c1b..eec5113 100644 --- a/src/engine/ov_genai/utils.py +++ b/src/engine/ov_genai/utils.py @@ -2,29 +2,31 @@ from src.server.models.ov_genai import SchedulerConfigSchema def generate_ov_scheduler_config(scheduler_config: SchedulerConfigSchema) -> dict: - """Generates a SchedulerConfig object from the scheduler config model.""" - # SchedulerConfig cannot be used as SDPA pipelines refuse to accept a scheduler config - # which includes any invalid properties, including those set to `None` by default. - # Fortunately, the pipelines also accept a dictionary. - sched_config = {} + """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 + 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 + sched_config.num_kv_blocks = scheduler_config.num_kv_blocks if scheduler_config.cache_size: - sched_config["cache_size"] = 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 + 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 + 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 + 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 + 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 + 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 + 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 sched_config \ No newline at end of file + sched_config.use_sparse_attention = scheduler_config.use_sparse_attention + return {"scheduler_config": sched_config} From 6e2e59438e369bdddce6f3ed671f1cf64d3b801b Mon Sep 17 00:00:00 2001 From: Michael Carroll Date: Wed, 5 Aug 2026 09:48:56 -0400 Subject: [PATCH 3/6] Create function to extract scheduler from load config --- src/engine/ov_genai/utils.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/src/engine/ov_genai/utils.py b/src/engine/ov_genai/utils.py index eec5113..b437f0b 100644 --- a/src/engine/ov_genai/utils.py +++ b/src/engine/ov_genai/utils.py @@ -1,5 +1,13 @@ """OV GenAI engine utilities.""" +import logging +from typing import Literal + +from openvino_genai import SchedulerConfig + from src.server.models.ov_genai import SchedulerConfigSchema +from src.server.models.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. @@ -30,3 +38,19 @@ def generate_ov_scheduler_config(scheduler_config: SchedulerConfigSchema) -> dic 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) From 9c0c841b42fb7ec4db1b37f9337c333ed2ec3b41 Mon Sep 17 00:00:00 2001 From: Michael Carroll Date: Wed, 5 Aug 2026 10:01:42 -0400 Subject: [PATCH 4/6] Create scheduler config model and apply to model load config --- src/cli/groups/add.py | 43 ++++++++++++++++------ src/server/models/ov_genai.py | 60 +++++++++++++++++++++++++++++++ src/server/models/registration.py | 7 +++- 3 files changed, 99 insertions(+), 11 deletions(-) diff --git a/src/cli/groups/add.py b/src/cli/groups/add.py index 7bf89d9..b2320c2 100644 --- a/src/cli/groups/add.py +++ b/src/cli/groups/add.py @@ -4,6 +4,9 @@ import json import click +from pydantic import ValidationError + +from src.server.models.ov_genai 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]") diff --git a/src/server/models/ov_genai.py b/src/server/models/ov_genai.py index d924902..6663c68 100644 --- a/src/server/models/ov_genai.py +++ b/src/server/models/ov_genai.py @@ -96,3 +96,63 @@ def text_messages(self) -> List[Dict[str, Any]]: class OVGenAI_WhisperGenConfig(BaseModel): audio_base64: str = Field(..., description="Base64 encoded audio") + + +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." + ) diff --git a/src/server/models/registration.py b/src/server/models/registration.py index 2257f37..2250d84 100644 --- a/src/server/models/registration.py +++ b/src/server/models/registration.py @@ -1,3 +1,4 @@ +from src.server.models.ov_genai import SchedulerConfigSchema from enum import Enum from typing import Any, Dict, Optional @@ -111,7 +112,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") From b85e9aef795700dc8673f779b036e9d39f500ad2 Mon Sep 17 00:00:00 2001 From: Michael Carroll Date: Wed, 5 Aug 2026 10:03:19 -0400 Subject: [PATCH 5/6] Apply scheduler config when loading the models --- src/engine/ov_genai/llm.py | 5 ++++- src/engine/ov_genai/vlm.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/engine/ov_genai/llm.py b/src/engine/ov_genai/llm.py index 00267e2..8653291 100755 --- a/src/engine/ov_genai/llm.py +++ b/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 @@ -269,6 +270,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: @@ -277,7 +279,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) diff --git a/src/engine/ov_genai/vlm.py b/src/engine/ov_genai/vlm.py index 1799caf..69c7830 100644 --- a/src/engine/ov_genai/vlm.py +++ b/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 @@ -270,7 +271,8 @@ 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 @@ -278,6 +280,7 @@ def load_model(self, loader: ModelLoadConfig): self.model_path = VLMPipeline( loader.model_path, loader.device, + **scheduler_config, **pipeline_kwargs ) From 4ff7a39cf9efa4650170498ffb3611f111a59da5 Mon Sep 17 00:00:00 2001 From: Michael Carroll Date: Wed, 5 Aug 2026 10:20:35 -0400 Subject: [PATCH 6/6] Correct for changes from merging back --- src/cli/groups/add.py | 2 +- src/engine/ov_genai/utils.py | 4 +- .../modeling/contract_ovgenai_llm_and_vlm.py | 60 +++++++++++++++++++ src/server/schemas/registration.py | 4 +- 4 files changed, 66 insertions(+), 4 deletions(-) diff --git a/src/cli/groups/add.py b/src/cli/groups/add.py index b2320c2..7de4413 100644 --- a/src/cli/groups/add.py +++ b/src/cli/groups/add.py @@ -6,7 +6,7 @@ import click from pydantic import ValidationError -from src.server.models.ov_genai import SchedulerConfigSchema +from src.server.schemas.modeling.contract_ovgenai_llm_and_vlm import SchedulerConfigSchema from ..main import cli, console from ..utils import validate_model_path diff --git a/src/engine/ov_genai/utils.py b/src/engine/ov_genai/utils.py index b437f0b..973f3a1 100644 --- a/src/engine/ov_genai/utils.py +++ b/src/engine/ov_genai/utils.py @@ -4,8 +4,8 @@ from openvino_genai import SchedulerConfig -from src.server.models.ov_genai import SchedulerConfigSchema -from src.server.models.registration import ModelLoadConfig +from src.server.schemas.modeling.contract_ovgenai_llm_and_vlm import SchedulerConfigSchema +from src.server.schemas.registration import ModelLoadConfig logger = logging.getLogger(__name__) diff --git a/src/server/schemas/modeling/contract_ovgenai_llm_and_vlm.py b/src/server/schemas/modeling/contract_ovgenai_llm_and_vlm.py index 61adbce..7233603 100644 --- a/src/server/schemas/modeling/contract_ovgenai_llm_and_vlm.py +++ b/src/server/schemas/modeling/contract_ovgenai_llm_and_vlm.py @@ -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." + ) diff --git a/src/server/schemas/registration.py b/src/server/schemas/registration.py index 2250d84..825ad54 100644 --- a/src/server/schemas/registration.py +++ b/src/server/schemas/registration.py @@ -1,9 +1,11 @@ -from src.server.models.ov_genai import SchedulerConfigSchema + 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.