diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 9ecdf003e6ee..7a427a3afb43 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -16,6 +16,7 @@ from typing import Any, AsyncGenerator, Dict, Final import torch +from vllm.config import VllmConfig from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt from vllm.lora.request import LoRARequest from vllm.outputs import RequestOutput @@ -260,6 +261,32 @@ def build_sampling_params_openai( return sampling_params +def get_dp_range_for_worker(vllm_config: VllmConfig) -> range: + """ + Get the global DP rank range that this worker is responsible for based on vLLM config. + Note that the 'vllm_config' is normalized so the load balancing flags are set properly. + The return value is in the format of (start_dp_rank, managed_dp_size).""" + if vllm_config.parallel_config.data_parallel_external_lb: + # external load balancing, each worker is responsible for exactly 1 rank + return (vllm_config.parallel_config.data_parallel_rank, 1) + elif vllm_config.parallel_config.data_parallel_hybrid_lb: + # hybrid load balancing, each worker is responsible for a subset of local ranks + return ( + vllm_config.parallel_config.data_parallel_rank, + vllm_config.parallel_config.data_parallel_size_local, + ) + else: + # internal load balancing, the worker is responsible for all DP ranks + logger.warning( + "vLLM selects internal DP load balancing. If you are launching multiple workers for DP deployment," + " hybrid or external load balancing is recommended." + ) + return ( + vllm_config.parallel_config.data_parallel_rank, + vllm_config.parallel_config.data_parallel_size, + ) + + class BaseWorkerHandler(ABC): """ Request handler for the generate and clear_kv_blocks endpoints. @@ -302,6 +329,8 @@ def __init__( self.use_vllm_tokenizer = use_vllm_tokenizer + self.dp_range = get_dp_range_for_worker(self.engine_client.vllm_config) + # Initialize InputParamManager for text-in-text-out mode tokenizer = None if use_vllm_tokenizer and hasattr(engine, "tokenizer"): @@ -463,6 +492,21 @@ def add_temp_dir(self, temp_dir: tempfile.TemporaryDirectory) -> None: if temp_dir is not None: self.temp_dirs.append(temp_dir) + def _to_local_dp_rank(self, dp_rank: int | None) -> int | None: + """Convert global DP rank to local DP rank based on engine config.""" + if dp_rank is None: + return None + if dp_rank < self.dp_range[0] or dp_rank >= self.dp_range[0] + self.dp_range[1]: + logger.warning( + f"Received DP rank {dp_rank} is out of range [{self.dp_range[0]} - {self.dp_range[0] + self.dp_range[1]}), fallback to vLLM internal DP selection" + ) + return None + local_dp_rank = (dp_rank - self.dp_range[0]) % self.dp_range[1] + logger.debug( + f"Converted global DP rank {dp_rank} to local DP rank {local_dp_rank}" + ) + return local_dp_rank + def _resolve_lora_request(self, model_name: str | None) -> LoRARequest | None: """Return a LoRARequest if model_name is a loaded adapter, else None.""" if model_name and (lora := self.loaded_loras.get(model_name)): @@ -1317,7 +1361,7 @@ async def _generate_token_mode(self, request, context, request_id): f"Decode request {request_id} has no LoRA specified (model: {model_name})" ) routing = request.get("routing") or {} - dp_rank = routing.get("dp_rank") + dp_rank = self._to_local_dp_rank(routing.get("dp_rank")) priority = routing.get("priority", 0) trace_headers = build_trace_headers(context) @@ -1364,7 +1408,7 @@ async def _generate_text_mode(self, request, context, request_id): ) routing = request.get("routing") or {} - dp_rank = routing.get("dp_rank") + dp_rank = self._to_local_dp_rank(routing.get("dp_rank")) priority = routing.get("priority", 0) openai_request_id = request.get("id") or request.get("request_id", request_id) previous_text = "" @@ -1525,7 +1569,7 @@ async def _generate_token_mode(self, request, context, request_id): ) routing = request.get("routing") or {} - dp_rank = routing.get("dp_rank") + dp_rank = self._to_local_dp_rank(routing.get("dp_rank")) priority = routing.get("priority", 0) trace_headers = build_trace_headers(context) diff --git a/components/src/dynamo/vllm/main.py b/components/src/dynamo/vllm/main.py index 5935acf41ac2..c8d32ddba08b 100644 --- a/components/src/dynamo/vllm/main.py +++ b/components/src/dynamo/vllm/main.py @@ -55,7 +55,7 @@ from .args import Config, _uses_dynamo_connector, parse_args from .checkpoint_restore import get_checkpoint_config from .constants import DisaggregationMode -from .handlers import DecodeWorkerHandler, PrefillWorkerHandler +from .handlers import DecodeWorkerHandler, PrefillWorkerHandler, get_dp_range_for_worker from .health_check import ( VllmHealthCheckPayload, VllmOmniHealthCheckPayload, @@ -69,19 +69,6 @@ CHECKPOINT_SLEEP_MODE_LEVEL = 1 -async def _handle_non_leader_node(dp_rank: int) -> None: - """ - Handle non-leader node (data_parallel_rank >= 1) in multi-node deployments. - Non-leader nodes run vLLM workers but don't serve Dynamo endpoints. - """ - logger.info( - f"Non-leader node detected (data_parallel_rank={dp_rank}). " - "Skipping endpoint serving." - ) - # Wait indefinitely - process terminated via signal handlers - await asyncio.Event().wait() - - def build_headless_namespace(config: Config) -> argparse.Namespace: """Build an argparse Namespace from engine_args for vLLM's run_headless(). @@ -339,11 +326,12 @@ def setup_kv_event_publisher( ) return None - # Get data_parallel_size to create publishers for all dp_ranks - data_parallel_size = getattr(vllm_config.parallel_config, "data_parallel_size", 1) + # Get DP rank range managed by this worker to create publishers for corresponding dp_ranks, + # all served workers should cover all ranks. + dp_start, dp_size = get_dp_range_for_worker(vllm_config) kv_publishers = [] - for dp_rank in range(data_parallel_size): + for dp_rank in range(dp_start, dp_start + dp_size): if consolidator_enabled: # TODO: Use different port for each dp_rank once KVBM supports DP zmq_endpoint = f"tcp://127.0.0.1:{consolidator_port}" @@ -561,8 +549,9 @@ async def register_vllm_model( runtime_config.reasoning_parser = config.dyn_reasoning_parser # Get data_parallel_size from vllm_config (defaults to 1) - data_parallel_size = getattr(vllm_config.parallel_config, "data_parallel_size", 1) - runtime_config.data_parallel_size = data_parallel_size + dp_range = get_dp_range_for_worker(vllm_config) + runtime_config.data_parallel_start_rank = dp_range[0] + runtime_config.data_parallel_size = dp_range[1] # Configure media decoder for frontend image decoding when enabled # This enables frontend to decode images and transfer via NIXL RDMA @@ -675,10 +664,6 @@ async def init_prefill( runtime.register_engine_route("wake_up", handler.wake_up) logger.info("Registered engine routes: /engine/sleep, /engine/wake_up") - # Handle non-leader nodes - don't serve endpoints - if config.engine_args.data_parallel_rank: - await _handle_non_leader_node(config.engine_args.data_parallel_rank) - return shutdown_endpoints[:] = [generate_endpoint, clear_endpoint] # Register prefill model with ModelType.Prefill @@ -791,7 +776,6 @@ async def init( factory = StatLoggerFactory( endpoint=generate_endpoint, component_gauges=component_gauges, - dp_rank=config.engine_args.data_parallel_rank or 0, ) else: # Factory is created without component_gauges; setup_vllm_engine() will @@ -799,7 +783,6 @@ async def init( # on the factory before vLLM calls create_stat_logger(). factory = StatLoggerFactory( endpoint=generate_endpoint, - dp_rank=config.engine_args.data_parallel_rank or 0, ) ( engine_client, @@ -858,11 +841,6 @@ async def init( runtime.register_engine_route("wake_up", handler.wake_up) logger.info("Registered engine routes: /engine/sleep, /engine/wake_up") - # Handle non-leader nodes - don't serve endpoints - if config.engine_args.data_parallel_rank: - await _handle_non_leader_node(config.engine_args.data_parallel_rank) - return - # Parse endpoint types from --endpoint-types flag model_type = parse_endpoint_types(config.endpoint_types) logger.info(f"Registering model with endpoint types: {config.endpoint_types}") @@ -1011,11 +989,6 @@ async def init_omni( # Set up metrics collection for vLLM and LMCache metrics setup_metrics_collection(config, generate_endpoint, logger) - # Handle non-leader nodes - don't serve endpoints - if config.engine_args.data_parallel_rank: - await _handle_non_leader_node(config.engine_args.data_parallel_rank) - return - # TODO: extend for multi-stage pipelines model_type = get_output_modalities(config.output_modalities, config.model) if model_type is None: diff --git a/components/src/dynamo/vllm/publisher.py b/components/src/dynamo/vllm/publisher.py index 6ceba22ab58f..9477f4d28365 100644 --- a/components/src/dynamo/vllm/publisher.py +++ b/components/src/dynamo/vllm/publisher.py @@ -19,24 +19,6 @@ DYNAMO_COMPONENT_REGISTRY = CollectorRegistry() -class NullStatLogger(StatLoggerBase): - def __init__(self): - pass - - def record( - self, - scheduler_stats: Optional[SchedulerStats], - iteration_stats: Optional[IterationStats], - engine_idx: int = 0, - *args, - **kwargs, - ): - pass - - def log_engine_initialized(self): - pass - - class DynamoStatLoggerPublisher(StatLoggerBase): """Stat logger publisher. Wrapper for the WorkerMetricsPublisher to match the StatLoggerBase interface.""" @@ -106,22 +88,17 @@ def __init__( self, endpoint: Endpoint, component_gauges: Optional[LLMBackendMetrics] = None, - dp_rank: int = 0, ) -> None: self.endpoint = endpoint self.component_gauges = component_gauges self.created_logger: Optional[DynamoStatLoggerPublisher] = None - self.dp_rank = dp_rank def create_stat_logger(self, dp_rank: int) -> StatLoggerBase: - if self.dp_rank != dp_rank: - return NullStatLogger() # component_gauges must be set by setup_vllm_engine() before vLLM # calls create_stat_logger() during engine initialization. assert ( self.component_gauges is not None ), "component_gauges must be set before creating stat loggers" - logger = DynamoStatLoggerPublisher( endpoint=self.endpoint, dp_rank=dp_rank, diff --git a/examples/backends/vllm/launch/dep.sh b/examples/backends/vllm/launch/dep.sh index 0c1d0d850757..3f76a0189750 100755 --- a/examples/backends/vllm/launch/dep.sh +++ b/examples/backends/vllm/launch/dep.sh @@ -35,16 +35,16 @@ python -m dynamo.frontend --router-mode kv & # Routing to DP workers managed by Dynamo # Chose Qwen3-30B because its a small MOE that can fit on smaller GPUs (L40S for example) # --enforce-eager is added for quick deployment. for production use, need to remove this flag -for i in {0..3}; do - VLLM_NIXL_SIDE_CHANNEL_PORT=$((20096 + i)) \ - CUDA_VISIBLE_DEVICES=$i python3 -m dynamo.vllm \ - --model "$MODEL" \ - --data-parallel-rank $i \ - --data-parallel-size 4 \ - --enable-expert-parallel \ - --enforce-eager \ - --kv-events-config "{\"publisher\":\"zmq\",\"topic\":\"kv-events\",\"endpoint\":\"tcp://*:$((20080 + i))\",\"enable_kv_cache_events\":true}" & -done +VLLM_NIXL_SIDE_CHANNEL_PORT=20096 \ +python3 -m dynamo.vllm \ +--model Qwen/Qwen3-30B-A3B \ +--data-parallel-hybrid-lb \ +--data-parallel-size 4 \ +--data-parallel-size-local 4 \ +--data-parallel-start-rank 0 \ +--enable-expert-parallel \ +--enforce-eager \ +--kv-events-config "{\"publisher\":\"zmq\",\"topic\":\"kv-events\",\"endpoint\":\"tcp://*:20080\",\"enable_kv_cache_events\":true}" & echo "All workers starting. (press Ctrl+C to stop)..." wait diff --git a/examples/backends/vllm/launch/dsr1_dep.sh b/examples/backends/vllm/launch/dsr1_dep.sh index de0f26c51967..dd225aa059ff 100755 --- a/examples/backends/vllm/launch/dsr1_dep.sh +++ b/examples/backends/vllm/launch/dsr1_dep.sh @@ -116,25 +116,24 @@ mkdir -p $LOG_DIR # the GPU memory requires for vLLM reservation and runtime spike (not # reserved by vLLM) can be different and cause model fails to start, # adjust '--gpu-memory-utilization' as needed -for ((i=0; i&1 | tee $LOG_DIR/dsr1_dep_${dp_rank}.log & -done +dp_start_rank=$((NODE_RANK * GPUS_PER_NODE)) +VLLM_NIXL_SIDE_CHANNEL_PORT=20096 \ +VLLM_ALL2ALL_BACKEND="deepep_low_latency" \ +VLLM_USE_DEEP_GEMM=1 \ +VLLM_RANDOMIZE_DP_DUMMY_INPUTS=1 \ +python3 -m dynamo.vllm \ +--model $MODEL \ +--data-parallel-hybrid-lb \ +--data-parallel-size $DATA_PARALLEL_SIZE \ +--data-parallel-size-local $GPUS_PER_NODE \ +--data-parallel-start-rank $dp_start_rank \ +--enable-expert-parallel \ +--max-model-len 4096 \ +--data-parallel-address $MASTER_ADDR \ +--data-parallel-rpc-port 13345 \ +--gpu-memory-utilization 0.91 \ +--enforce-eager \ +--kv-events-config "{\"publisher\":\"zmq\",\"topic\":\"kv-events\",\"endpoint\":\"tcp://*:20080\",\"enable_kv_cache_events\":true}" 2>&1 | tee $LOG_DIR/dsr1_dep_${dp_start_rank}.log & echo "All workers starting. (press Ctrl+C to stop)..." wait diff --git a/lib/bindings/python/rust/llm/local_model.rs b/lib/bindings/python/rust/llm/local_model.rs index 49dc15cd32d3..075e71af2817 100644 --- a/lib/bindings/python/rust/llm/local_model.rs +++ b/lib/bindings/python/rust/llm/local_model.rs @@ -45,6 +45,11 @@ impl ModelRuntimeConfig { self.inner.reasoning_parser = reasoning_parser; } + #[setter] + fn set_data_parallel_start_rank(&mut self, data_parallel_start_rank: u32) { + self.inner.data_parallel_start_rank = data_parallel_start_rank; + } + #[setter] fn set_data_parallel_size(&mut self, data_parallel_size: u32) { self.inner.data_parallel_size = data_parallel_size; diff --git a/lib/kv-router/benches/active_sequences_bench.rs b/lib/kv-router/benches/active_sequences_bench.rs index 214f8df3cbde..40acea2fd2e7 100644 --- a/lib/kv-router/benches/active_sequences_bench.rs +++ b/lib/kv-router/benches/active_sequences_bench.rs @@ -289,11 +289,12 @@ async fn run_benchmark( // Total bench workers = trace workers × duplication factor. // Each gets a unique WorkerWithDpRank in the shared multi-worker. let total_workers = num_trace_workers * inference_worker_duplication_factor; - let dp_sizes: HashMap = (0..total_workers as u64).map(|id| (id, 1)).collect(); + let dp_range: HashMap = + (0..total_workers as u64).map(|id| (id, (0, 1))).collect(); let multi = Arc::new(ActiveSequencesMultiWorker::new( NoopSequencePublisher, block_size as usize, - dp_sizes, + dp_range, false, 0, "bench", diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 54ffe15e475d..eb0d099909e8 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -96,6 +96,7 @@ pub fn compute_seq_hash_for_block(block_hashes: &[LocalBlockHash]) -> Vec u32; fn data_parallel_size(&self) -> u32; fn max_num_batched_tokens(&self) -> Option; fn total_kv_blocks(&self) -> Option; diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index a19fd1e01797..93fbfcec415b 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -209,11 +209,12 @@ impl SchedulerQueue { for (&worker_id, config) in configs.iter() { let dp_size = config.data_parallel_size(); + let dp_start_rank = config.data_parallel_start_rank(); let max_batched = config .max_num_batched_tokens() .unwrap_or(DEFAULT_MAX_BATCHED_TOKENS); - for dp_rank in 0..dp_size { + for dp_rank in dp_start_rank..dp_start_rank + dp_size { let worker = WorkerWithDpRank::new(worker_id, dp_rank); let tokens = active_tokens.get(&worker).copied().unwrap_or(0); if (tokens as f64) <= threshold * (max_batched as f64) { @@ -247,11 +248,12 @@ mod tests { Arc>, Arc>, ) { - let dp_sizes: HashMap = (0..num_workers as u64).map(|id| (id, 1)).collect(); + let dp_range: HashMap = + (0..num_workers as u64).map(|id| (id, (0, 1))).collect(); let slots = Arc::new(ActiveSequencesMultiWorker::new( NoopSequencePublisher, block_size as usize, - dp_sizes, + dp_range, false, 0, "test", diff --git a/lib/kv-router/src/scheduling/selector.rs b/lib/kv-router/src/scheduling/selector.rs index 68cc89c1e5d4..6d1a20674f77 100644 --- a/lib/kv-router/src/scheduling/selector.rs +++ b/lib/kv-router/src/scheduling/selector.rs @@ -129,8 +129,10 @@ impl WorkerSelector for DefaultWorkerSelector { .filter(|(wid, _)| allowed_ids.is_none_or(|ids| ids.contains(wid))) { let data_parallel_size = config.data_parallel_size(); + let data_parallel_start_rank = config.data_parallel_start_rank(); - for dp_rank in 0..data_parallel_size { + for dp_rank in data_parallel_start_rank..(data_parallel_start_rank + data_parallel_size) + { let worker = WorkerWithDpRank::new(*worker_id, dp_rank); let overlap = *overlaps.get(&worker).unwrap_or(&0); diff --git a/lib/kv-router/src/sequences/multi_worker.rs b/lib/kv-router/src/sequences/multi_worker.rs index 638e0171efc3..1e3cffc54f4d 100644 --- a/lib/kv-router/src/sequences/multi_worker.rs +++ b/lib/kv-router/src/sequences/multi_worker.rs @@ -103,11 +103,11 @@ struct WorkerTable { } impl WorkerTable { - fn new(block_size: usize, dp_sizes: &HashMap) -> Self { + fn new(block_size: usize, dp_range: &HashMap) -> Self { let mut slots = Vec::new(); let mut index = HashMap::new(); - for (&worker_id, &dp_size) in dp_sizes { - for dp_rank in 0..dp_size { + for (&worker_id, &(dp_start, dp_size)) in dp_range { + for dp_rank in dp_start..dp_start + dp_size { let worker = WorkerWithDpRank::new(worker_id, dp_rank); let idx = slots.len(); slots.push((worker, RwLock::new(ActiveSequences::new(block_size)))); @@ -149,7 +149,7 @@ impl ActiveSequencesMultiWorker

{ pub fn new( publisher: P, block_size: usize, - dp_sizes: HashMap, + dp_range: HashMap, replica_sync: bool, router_id: u64, worker_type: &'static str, @@ -157,7 +157,7 @@ impl ActiveSequencesMultiWorker

{ assert!(block_size > 1, "block_size must be greater than 1"); Self { - workers: RwLock::new(WorkerTable::new(block_size, &dp_sizes)), + workers: RwLock::new(WorkerTable::new(block_size, &dp_range)), request_to_worker: DashMap::new(), request_to_lora: DashMap::new(), block_size, @@ -276,13 +276,13 @@ impl ActiveSequencesMultiWorker

{ /// Update the set of workers, adding and removing as needed. /// - /// `new_dp_sizes` maps worker IDs to their data-parallel size. - pub fn update_workers(&self, new_dp_sizes: &HashMap) { + /// `new_dp_range` maps worker IDs to their data-parallel range (start, size). + pub fn update_workers(&self, new_dp_range: &HashMap) { let mut table = self.workers.write(); let mut target_workers: HashSet = HashSet::new(); - for (&worker_id, &dp_size) in new_dp_sizes { - for dp_rank in 0..dp_size { + for (&worker_id, &(dp_start, dp_size)) in new_dp_range { + for dp_rank in dp_start..(dp_start + dp_size) { target_workers.insert(WorkerWithDpRank::new(worker_id, dp_rank)); } } diff --git a/lib/kv-router/src/test_utils.rs b/lib/kv-router/src/test_utils.rs index a97217cc2f6d..0bac542a006a 100644 --- a/lib/kv-router/src/test_utils.rs +++ b/lib/kv-router/src/test_utils.rs @@ -85,6 +85,7 @@ impl SequencePublisher for NoopSequencePublisher { /// Minimal [`WorkerConfigLike`] for scheduler/queue tests and benchmarks. #[derive(Debug, Clone)] pub struct SimpleWorkerConfig { + pub data_parallel_start_rank: u32, pub data_parallel_size: u32, pub max_num_batched_tokens: Option, pub total_kv_blocks: Option, @@ -93,6 +94,7 @@ pub struct SimpleWorkerConfig { impl Default for SimpleWorkerConfig { fn default() -> Self { Self { + data_parallel_start_rank: 0, data_parallel_size: 1, max_num_batched_tokens: None, total_kv_blocks: None, @@ -101,6 +103,10 @@ impl Default for SimpleWorkerConfig { } impl WorkerConfigLike for SimpleWorkerConfig { + fn data_parallel_start_rank(&self) -> u32 { + self.data_parallel_start_rank + } + fn data_parallel_size(&self) -> u32 { self.data_parallel_size } diff --git a/lib/llm/src/discovery/worker_monitor.rs b/lib/llm/src/discovery/worker_monitor.rs index 73a2f9997ed6..1f09740c4243 100644 --- a/lib/llm/src/discovery/worker_monitor.rs +++ b/lib/llm/src/discovery/worker_monitor.rs @@ -456,22 +456,25 @@ impl WorkerLoadMonitor for KvWorkerMonitor { for (lease_id, runtime_config) in runtime_configs.iter() { let mut state = worker_load_states.entry(*lease_id).or_default(); + let dp_start = runtime_config.data_parallel_start_rank; + let dp_end = dp_start + runtime_config.data_parallel_size; + // Track dp_ranks for this worker (for cleanup when worker disappears) let dp_ranks_set = known_worker_dp_ranks.entry(*lease_id).or_default(); - for dp_rank in 0..runtime_config.data_parallel_size { + for dp_rank in dp_start..dp_end { dp_ranks_set.insert(dp_rank); } // Populate total_blocks for all dp_ranks (they share the same total) if let Some(total_blocks) = runtime_config.total_kv_blocks { - for dp_rank in 0..runtime_config.data_parallel_size { + for dp_rank in dp_start..dp_end { state.kv_total_blocks.insert(dp_rank, total_blocks); } } // Populate max_num_batched_tokens for all dp_ranks if let Some(max_batched) = runtime_config.max_num_batched_tokens { - for dp_rank in 0..runtime_config.data_parallel_size { + for dp_rank in dp_start..dp_end { state.max_num_batched_tokens.insert(dp_rank, max_batched); } } diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs index 3b6f719940fc..5d28102a8854 100644 --- a/lib/llm/src/kv_router/scheduler.rs +++ b/lib/llm/src/kv_router/scheduler.rs @@ -87,11 +87,11 @@ impl KvScheduler { let current_workers = monitor_rx.borrow_and_update().clone(); if current_workers != last_workers { - let dp_sizes: HashMap = current_workers + let dp_range: HashMap = current_workers .iter() - .map(|(&id, c)| (id, c.data_parallel_size)) + .map(|(&id, c)| (id, (c.data_parallel_start_rank, c.data_parallel_size))) .collect(); - slots_monitor.update_workers(&dp_sizes); + slots_monitor.update_workers(&dp_range); last_workers = current_workers; } } diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index b6e86ef5cec0..de6493c6456f 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -103,15 +103,20 @@ pub async fn create_multi_worker_sequences( metrics_publisher, }; - let dp_sizes: HashMap = workers_with_configs + let dp_range: HashMap = workers_with_configs .into_iter() - .map(|(id, config)| (id, config.data_parallel_size)) + .map(|(id, config)| { + ( + id, + (config.data_parallel_start_rank, config.data_parallel_size), + ) + }) .collect(); let multi_worker = ActiveSequencesMultiWorker::new( publisher, block_size, - dp_sizes, + dp_range, replica_sync, router_id, worker_type, diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index 477a6285a6f9..1ef0077528a3 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -28,6 +28,10 @@ pub struct ModelRuntimeConfig { pub reasoning_parser: Option, + /// Starting rank of data parallel ranks for this worker (0 if DP not enabled) + #[serde(default = "default_data_parallel_start_rank")] + pub data_parallel_start_rank: u32, + /// Total number of data parallel ranks for this worker (1 if DP not enabled) #[serde(default = "default_data_parallel_size")] pub data_parallel_size: u32, @@ -55,6 +59,10 @@ pub struct ModelRuntimeConfig { pub disaggregated_endpoint: Option, } +const fn default_data_parallel_start_rank() -> u32 { + 0 +} + const fn default_data_parallel_size() -> u32 { 1 } @@ -71,6 +79,7 @@ impl Default for ModelRuntimeConfig { max_num_batched_tokens: None, tool_call_parser: None, reasoning_parser: None, + data_parallel_start_rank: default_data_parallel_start_rank(), data_parallel_size: default_data_parallel_size(), enable_local_indexer: true, runtime_data: HashMap::new(), @@ -81,6 +90,10 @@ impl Default for ModelRuntimeConfig { } impl dynamo_kv_router::WorkerConfigLike for ModelRuntimeConfig { + fn data_parallel_start_rank(&self) -> u32 { + self.data_parallel_start_rank + } + fn data_parallel_size(&self) -> u32 { self.data_parallel_size }