diff --git a/rust/proto/control.proto b/rust/proto/control.proto index 428ad64604..1903405e27 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -38,6 +38,7 @@ message ServerInfo { uint64 total_kv_blocks = 7; uint64 max_running_requests = 8; uint64 max_batched_tokens = 9; + uint32 max_loras = 10; RlCapabilities rl_capabilities = 11; } @@ -65,6 +66,7 @@ message ModelInfo { bool supports_text_input = 20; bool supports_token_ids_input = 21; + bool supports_lora = 22; bool supports_multimodal = 23; string reasoning_parser = 24; string tool_call_parser = 25; diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 7f8663f852..1efca89621 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -286,6 +286,7 @@ impl EngineCoreClient { config: EngineCoreClientConfig, connected: transport::ConnectedTransport, ) -> Result { + validate_lora_capabilities(&connected.engines)?; let (output_tx, output_rx) = mpsc::channel(64); let (abort_tx, abort_rx) = mpsc::unbounded_channel(); let engines = connected.engines; @@ -476,6 +477,32 @@ impl EngineCoreClient { } } +fn validate_lora_capabilities(engines: &[ConnectedEngine]) -> Result<()> { + let first = engines.first().expect("engine core client requires at least one engine"); + for engine in engines { + let ready = &engine.ready_response; + if ready.supports_lora != (ready.max_loras > 0) { + return Err(Error::UnexpectedHandshakeMessage { + message: format!( + "engine {:?} reported inconsistent LoRA capability (supports_lora={}, max_loras={})", + engine.engine_id, ready.supports_lora, ready.max_loras + ), + }); + } + if ready.supports_lora != first.ready_response.supports_lora + || ready.max_loras != first.ready_response.max_loras + { + return Err(Error::UnexpectedHandshakeMessage { + message: format!( + "engine {:?} reported LoRA capability inconsistent with engine {:?}", + engine.engine_id, first.engine_id + ), + }); + } + } + Ok(()) +} + // Client API implementation. impl EngineCoreClient { /// Add a new request to the engine and return a per-request raw output diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index 28fb4d84fe..8bd6d77360 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -64,6 +64,8 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { max_num_seqs: 256, max_num_batched_tokens: 8192, instance_id: "test-instance".to_string(), + supports_lora: false, + max_loras: 0, kv_cache_size_tokens: None, kv_cache_max_concurrency: None, kv_events_config: None, diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index b4214a06f2..5cb4c98046 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -80,6 +80,10 @@ pub struct EngineCoreReadyResponse { pub max_num_batched_tokens: u64, /// Unique identifier for this server instance. pub instance_id: String, + /// Whether the engine was started with LoRA support enabled. + pub supports_lora: bool, + /// Maximum number of LoRA adapters the engine may keep active. + pub max_loras: u32, /// Total KV cache capacity in tokens, if reported. pub kv_cache_size_tokens: Option, /// Maximum achievable request concurrency given the KV cache, if reported. diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 8f07a8059e..913b894ada 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2686,6 +2686,8 @@ fn python_msgpack_fixtures_match_rust_encoding() { let ready_response: EngineCoreReadyResponse = rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap(); + assert!(ready_response.supports_lora); + assert_eq!(ready_response.max_loras, 8); assert_eq!( ready_response.weight_transfer_backend.as_deref(), Some("nccl") diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 389e97cd99..ac033914b0 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -409,6 +409,8 @@ class EngineCoreReadyResponse: max_num_seqs: int max_num_batched_tokens: int instance_id: str + supports_lora: bool + max_loras: int kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None kv_events_config: KVEventsConfig | None = None @@ -433,6 +435,8 @@ ready_response = EngineCoreReadyResponse( max_num_seqs=256, max_num_batched_tokens=8192, instance_id="test-instance", + supports_lora=True, + max_loras=8, weight_transfer_backend="nccl", enable_sleep_mode=True, supports_draft_weight_updates=True, diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index c7d9af88c8..c08803306f 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -150,6 +150,7 @@ impl pb::control_server::Control for ControlServiceImpl { total_kv_blocks: self.state.engine_core_client().total_num_gpu_blocks(), max_running_requests: ready.max_num_seqs, max_batched_tokens: ready.max_num_batched_tokens, + max_loras: ready.max_loras, rl_capabilities: Some(self.rl_capabilities()), })) } @@ -166,6 +167,7 @@ impl pb::control_server::Control for ControlServiceImpl { // GenerateRequest accepts both prompt representations. supports_text_input: true, supports_token_ids_input: true, + supports_lora: self.ready().supports_lora, supports_multimodal: self.state.chat.supports_multimodal(), reasoning_parser: self .state diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 74f7b803b8..1727c51b19 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -335,6 +335,8 @@ def test_apply_ready_response_syncs_block_size(): max_num_seqs=256, max_num_batched_tokens=8192, instance_id="test-instance", + supports_lora=False, + max_loras=0, ) ) client._apply_ready_response(payload) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index d70778eb51..e07df08f3d 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -88,6 +88,8 @@ class EngineCoreReadyResponse: max_num_seqs: int max_num_batched_tokens: int instance_id: str + supports_lora: bool + max_loras: int # KV cache capacity (None for encoder-only/attention-free models). kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 295c66b360..641c5ecc4e 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1636,6 +1636,12 @@ class EngineCoreProc(EngineCore): max_num_seqs=scheduler_config.max_num_seqs, max_num_batched_tokens=scheduler_config.max_num_batched_tokens, instance_id=self.vllm_config.instance_id, + supports_lora=self.vllm_config.lora_config is not None, + max_loras=( + self.vllm_config.lora_config.max_loras + if self.vllm_config.lora_config is not None + else 0 + ), kv_events_config=self.scheduler.get_kv_event_publisher_config(), weight_transfer_backend=( self.vllm_config.weight_transfer_config.backend