# # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. # This file is a part of the vllm-ascend project. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Scheduler subclass with profiling-based dynamic chunk sizing. Compatible with vLLM v0.15.x scheduler. When the upstream ``schedule()`` method is refactored, this override should be updated accordingly. """ import inspect import time from vllm.config import VllmConfig from vllm.logger import logger from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import ( NewRequestData, SchedulerOutput, ) from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.engine import EngineCoreEventType from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request, RequestStatus from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import record_function_or_nullcontext from vllm_ascend.core.profiling_chunk_predictor import ProfilingChunkManager class ProfilingChunkScheduler(Scheduler): """Scheduler with profiling-based dynamic chunk sizing. During initialization, the scheduler profiles prefill latency at various chunk sizes by calling ``profile_prefill_latency`` on each worker via ``collective_rpc``. A quadratic latency model is then fitted, and during scheduling the model predicts the optimal chunk size for each waiting request based on its ``num_computed_tokens``. """ def __init__( self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig, structured_output_manager: StructuredOutputManager, block_size: int, # `hash_block_size` was added in vLLM #40946; keep it optional so the # subclass works on both pinned vllm and main. hash_block_size: int | None = None, mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, include_finished_set: bool = False, log_stats: bool = False, ) -> None: super().__init__( vllm_config, kv_cache_config, structured_output_manager, block_size, hash_block_size=hash_block_size, mm_registry=mm_registry, include_finished_set=include_finished_set, log_stats=log_stats, ) from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config init_ascend_config(vllm_config) profiling_cfg = get_ascend_config().profiling_chunk_config self.profiling_chunk_config = profiling_cfg base_chunk = self.max_num_scheduled_tokens self.profiling_chunk_manager = ProfilingChunkManager( base_chunk_size=base_chunk, page_size=self.cache_config.block_size, smooth_factor=profiling_cfg.smooth_factor, min_chunk=profiling_cfg.min_chunk, max_fit_chunk=profiling_cfg.max_fit_chunk, ) self._profiling_initialized = False logger.info( "[ProfilingChunk] Scheduler initialized. base_chunk=%d, page_size=%d, smooth_factor=%.2f, min_chunk=%d", base_chunk, self.cache_config.block_size, profiling_cfg.smooth_factor, profiling_cfg.min_chunk, ) # ------------------------------------------------------------------ # Profiling initialization # ------------------------------------------------------------------ def run_profiling_chunk_init(self, model_executor) -> None: """Profile prefill latency using real model forward passes. Called by EngineCore after model_executor is ready. Collects latency samples at different chunk sizes and fits the quadratic model. """ if self._profiling_initialized: return self._profiling_initialized = True if model_executor is None: logger.warning("[ProfilingChunk] No model_executor provided, skipping profiling") return logger.info("[ProfilingChunk] Running startup profiling with real model forward...") seq_lens: list[int] = [] latencies: list[float] = [] base_chunk_size = self.profiling_chunk_manager.base_chunk_size num_samples = 64 # Determine unique_reply_rank for PP setups rpc_kwargs = self._build_rpc_kwargs(model_executor) total_steps = num_samples + 1 log_interval = max(1, total_steps // 10) t_start = time.perf_counter() for i in range(total_steps): chunk_size = int(base_chunk_size - (i - 1) * (base_chunk_size / num_samples)) if chunk_size <= 0: break if i % log_interval == 0 or i == total_steps - 1: elapsed = time.perf_counter() - t_start logger.info( "[ProfilingChunk] Profiling prefill latency: %d/%d samples done (chunk=%d, elapsed=%.1fs)", max(i - 1, 0), num_samples, chunk_size, elapsed, ) try: result = model_executor.collective_rpc( "profile_prefill_latency", args=(chunk_size,), **rpc_kwargs, ) # First iteration is warm-up if i == 0: continue latency_ms = self._extract_latency(result) if latency_ms is None: continue seq_lens.append(chunk_size) latencies.append(latency_ms) except Exception as e: logger.debug( "[ProfilingChunk] Forward failed for chunk=%d: %s", chunk_size, e, ) continue if len(seq_lens) < 8: logger.warning( "[ProfilingChunk] Profiling failed: only %d/8 samples collected", len(seq_lens), ) return logger.info( "[ProfilingChunk] Collected %d samples. Latency range: [%.2f, %.2f] ms", len(seq_lens), min(latencies), max(latencies), ) predictor = self.profiling_chunk_manager.predictor if not predictor.fit(seq_lens, latencies): return predictor.set_target_latency(base_chunk_size) predictor.is_ready = True self.profiling_chunk_manager._profiling_done = True logger.info("[ProfilingChunk] Profiling completed successfully") @staticmethod def _build_rpc_kwargs(model_executor) -> dict: """Build kwargs for collective_rpc, handling PP unique_reply_rank.""" kwargs: dict = {} if not hasattr(model_executor, "collective_rpc"): return kwargs sig = inspect.signature(model_executor.collective_rpc) if "unique_reply_rank" not in sig.parameters: return kwargs try: pc = model_executor.vllm_config.parallel_config output_rank = pc.world_size - pc.tensor_parallel_size * pc.prefill_context_parallel_size kwargs["unique_reply_rank"] = output_rank except AttributeError: pass return kwargs @staticmethod def _extract_latency(result) -> float | None: """Extract latency value from collective_rpc result.""" if isinstance(result, (int, float)): return float(result) if isinstance(result, list) and len(result) > 0: return float(result[0]) return None # ------------------------------------------------------------------ # schedule() override # ------------------------------------------------------------------ # The method below is based on the upstream Scheduler.schedule() # with profiling-based chunk sizing applied to both RUNNING requests # (chunked prefill continuation) and WAITING requests (new prefill). # Modified sections are marked with ">>> PROFILING CHUNK" comments. # ------------------------------------------------------------------ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # noqa: C901 scheduled_new_reqs: list[Request] = [] scheduled_resumed_reqs: list[Request] = [] scheduled_running_reqs: list[Request] = [] preempted_reqs: list[Request] = [] req_to_new_blocks: dict[str, KVCacheBlocks] = {} num_scheduled_tokens: dict[str, int] = {} # >>> PROFILING CHUNK >>> target_latency = self.profiling_chunk_manager.predictor.target_latency time_budget = target_latency if target_latency is not None else float("inf") # <<< PROFILING CHUNK <<< token_budget = self.max_num_scheduled_tokens if self._pause_state == PauseState.PAUSED_ALL: token_budget = 0 # Encoder-related. scheduled_encoder_inputs: dict[str, list[int]] = {} encoder_compute_budget = self.max_num_encoder_input_tokens # Spec decode-related. scheduled_spec_decode_tokens: dict[str, list[int]] = {} # For logging. scheduled_timestamp = time.monotonic() self.kv_cache_manager.new_step_starts() # First, schedule the RUNNING requests. req_index = 0 # >>> PROFILING CHUNK >>> while req_index < len(self.running) and token_budget > 0 and time_budget > 0: # <<< PROFILING CHUNK <<< request = self.running[req_index] if ( request.num_output_placeholders > 0 and request.num_computed_tokens + 2 - request.num_output_placeholders >= request.num_prompt_tokens + request.max_tokens ): req_index += 1 continue num_new_tokens = ( request.num_tokens_with_spec + request.num_output_placeholders - request.num_computed_tokens ) if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: num_new_tokens = self.scheduler_config.long_prefill_token_threshold num_new_tokens = min(num_new_tokens, token_budget) # Make sure the input position does not exceed the max model len. num_new_tokens = min( num_new_tokens, self.max_model_len - 1 - request.num_computed_tokens, ) # Schedule encoder inputs. encoder_inputs_to_schedule = None external_load_encoder_input: list[int] = [] new_encoder_compute_budget = encoder_compute_budget if request.has_encoder_inputs: ( encoder_inputs_to_schedule, num_new_tokens, new_encoder_compute_budget, external_load_encoder_input, ) = self._try_schedule_encoder_inputs( request, request.num_computed_tokens, num_new_tokens, encoder_compute_budget, shift_computed_tokens=1 if self.use_eagle else 0, ) # >>> PROFILING CHUNK: dynamic chunk sizing for RUNNING >>> if ( self.profiling_chunk_manager is not None and self.profiling_chunk_manager.is_ready and request.num_computed_tokens < request.num_prompt_tokens and (request.num_computed_tokens > 0 or not self.profiling_chunk_config.need_timing) ): predicted_chunk = self.profiling_chunk_manager.predict_chunk_size( num_computed_tokens=request.num_computed_tokens, target_time=time_budget, ) if predicted_chunk is not None and predicted_chunk > 0: logger.debug( "[ProfilingChunk] Dynamic chunk for %s: %s -> %s (predicted=%s)", request.request_id, num_new_tokens, min(predicted_chunk, num_new_tokens), predicted_chunk, ) num_new_tokens = min(predicted_chunk, num_new_tokens) elif self.profiling_chunk_config.need_timing: logger.info("[Dynamic Chunk] Online calibration stage. Long requests are better") elif time_budget == target_latency: logger.warning_once( "[Dynamic Chunk] Profiling Failed. Degenerated to a fixed chunk size" "Please increase the `max_fit_chunk` to profile more data" ) else: break # <<< PROFILING CHUNK <<< if self.need_mamba_block_aligned_split: num_new_tokens = self._mamba_block_aligned_split(request, num_new_tokens) if num_new_tokens == 0: req_index += 1 continue # Schedule newly needed KV blocks for the request. with record_function_or_nullcontext("schedule: allocate_slots"): while True: new_blocks = self.kv_cache_manager.allocate_slots( request, num_new_tokens, num_lookahead_tokens=self.num_lookahead_tokens, ) if new_blocks is not None: break if self._disable_preemption: break if self.policy == SchedulingPolicy.PRIORITY: preempted_req = max( self.running, key=lambda r: (r.priority, r.arrival_time), ) self.running.remove(preempted_req) if preempted_req in scheduled_running_reqs: preempted_req_id = preempted_req.request_id scheduled_running_reqs.remove(preempted_req) token_budget += num_scheduled_tokens.pop(preempted_req_id) req_to_new_blocks.pop(preempted_req_id) scheduled_spec_decode_tokens.pop(preempted_req_id, None) preempted_encoder_inputs = scheduled_encoder_inputs.pop(preempted_req_id, None) if preempted_encoder_inputs: num_embeds_to_restore = sum( preempted_req.get_num_encoder_embeds(i) for i in preempted_encoder_inputs ) encoder_compute_budget += num_embeds_to_restore req_index -= 1 else: preempted_req = self.running.pop() self._preempt_request(preempted_req, scheduled_timestamp) preempted_reqs.append(preempted_req) logger.info( "[ProfilingChunk] Preempted request %s. running_count=%s, token_budget=%s", preempted_req.request_id, len(self.running), token_budget, ) if preempted_req == request: break if new_blocks is None: break # Schedule the request. scheduled_running_reqs.append(request) request_id = request.request_id req_to_new_blocks[request_id] = new_blocks num_scheduled_tokens[request_id] = num_new_tokens token_budget -= num_new_tokens # Decode requests (num_new_tokens == 1) have negligible latency; # skip time_budget accounting so they don't starve other requests. if request.num_computed_tokens < request.num_prompt_tokens: time_budget -= self.profiling_chunk_manager.predict_time(num_new_tokens, request.num_computed_tokens) req_index += 1 # Speculative decode related. if request.spec_token_ids: num_scheduled_spec_tokens = ( num_new_tokens + request.num_computed_tokens - request.num_tokens - request.num_output_placeholders ) if num_scheduled_spec_tokens > 0: spec_token_ids = request.spec_token_ids if len(spec_token_ids) > num_scheduled_spec_tokens: spec_token_ids = spec_token_ids[:num_scheduled_spec_tokens] scheduled_spec_decode_tokens[request_id] = spec_token_ids request.spec_token_ids = [] # Encoder-related. if encoder_inputs_to_schedule: scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule for i in encoder_inputs_to_schedule: self.encoder_cache_manager.allocate(request, i) if self.ec_connector is not None: self.ec_connector.update_state_after_alloc(request, i) encoder_compute_budget = new_encoder_compute_budget if external_load_encoder_input: for i in external_load_encoder_input: self.encoder_cache_manager.allocate(request, i) if self.ec_connector is not None: self.ec_connector.update_state_after_alloc(request, i) # Record the LoRAs in scheduled_running_reqs scheduled_loras: set[int] = set() if self.lora_config: scheduled_loras = set( req.lora_request.lora_int_id for req in scheduled_running_reqs if req.lora_request and req.lora_request.lora_int_id > 0 ) assert len(scheduled_loras) <= self.lora_config.max_loras # Next, schedule the WAITING requests. if not preempted_reqs and self._pause_state == PauseState.UNPAUSED: step_skipped_waiting = create_request_queue(self.policy) # >>> PROFILING CHUNK >>> while (self.waiting or self.skipped_waiting) and token_budget > 0 and time_budget > 0: # <<< PROFILING CHUNK <<< if len(self.running) == self.max_num_running_reqs: break request_queue = self._select_waiting_queue_for_scheduling() assert request_queue is not None request = request_queue.peek_request() request_id = request.request_id # Try to promote blocked statuses while traversing skipped queue. if self._is_blocked_waiting_status(request.status) and not self._try_promote_blocked_waiting_request( request ): if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: logger.debug( "[ProfilingChunk] %s is still in WAITING_FOR_REMOTE_KVS state.", request_id, ) request_queue.pop_request() step_skipped_waiting.prepend_request(request) continue # Check that adding the request still respects the max_loras # constraint. if ( self.lora_config and request.lora_request and ( len(scheduled_loras) == self.lora_config.max_loras and request.lora_request.lora_int_id not in scheduled_loras ) ): request_queue.pop_request() step_skipped_waiting.prepend_request(request) continue num_external_computed_tokens = 0 load_kv_async = False connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0 # Get already-cached tokens. if request.num_computed_tokens == 0: new_computed_blocks, num_new_local_computed_tokens = self.kv_cache_manager.get_computed_blocks( request ) if self.connector is not None: ext_tokens, load_kv_async = self.connector.get_num_new_matched_tokens( request, num_new_local_computed_tokens ) if ext_tokens is None: request_queue.pop_request() step_skipped_waiting.prepend_request(request) continue num_external_computed_tokens = ext_tokens connector_prefix_cache_queries = request.num_tokens - num_new_local_computed_tokens connector_prefix_cache_hits = num_external_computed_tokens num_computed_tokens = num_new_local_computed_tokens + num_external_computed_tokens if request.prefill_stats is not None: request.prefill_stats.set( num_prompt_tokens=request.num_prompt_tokens, num_local_cached_tokens=num_new_local_computed_tokens, num_external_cached_tokens=num_external_computed_tokens, ) assert num_computed_tokens <= request.num_tokens else: new_computed_blocks = self.kv_cache_manager.empty_kv_cache_blocks num_new_local_computed_tokens = 0 num_computed_tokens = request.num_computed_tokens encoder_inputs_to_schedule = None external_load_encoder_input = [] new_encoder_compute_budget = encoder_compute_budget if load_kv_async: assert num_external_computed_tokens > 0 num_new_tokens = 0 else: num_new_tokens = request.num_tokens - num_computed_tokens threshold = self.scheduler_config.long_prefill_token_threshold if 0 < threshold < num_new_tokens: num_new_tokens = threshold # >>> PROFILING CHUNK: dynamic chunk sizing >>> if ( self.profiling_chunk_manager is not None and self.profiling_chunk_manager.is_ready and request.num_computed_tokens < request.num_prompt_tokens and (request.num_computed_tokens > 0 or not self.profiling_chunk_config.need_timing) ): predicted_chunk = self.profiling_chunk_manager.predict_chunk_size( num_computed_tokens=num_computed_tokens, target_time=time_budget, ) if predicted_chunk is not None and predicted_chunk > 0: num_new_tokens = min(num_new_tokens, predicted_chunk) elif self.profiling_chunk_config.need_timing: logger.info("[Dynamic Chunk] Online calibration stage. Long requests are better") elif time_budget == target_latency: logger.warning_once( "[Dynamic Chunk] Profiling Failed. Degenerated to a fixed chunk size" "Please increase the `max_fit_chunk` to profile more data" ) else: break # <<< PROFILING CHUNK <<< if not self.scheduler_config.enable_chunked_prefill and num_new_tokens > token_budget: break num_new_tokens = min(num_new_tokens, token_budget) assert num_new_tokens > 0 # Schedule encoder inputs. if request.has_encoder_inputs: ( encoder_inputs_to_schedule, num_new_tokens, new_encoder_compute_budget, external_load_encoder_input, ) = self._try_schedule_encoder_inputs( request, num_computed_tokens, num_new_tokens, encoder_compute_budget, shift_computed_tokens=1 if self.use_eagle else 0, ) if num_new_tokens == 0: break if self.need_mamba_block_aligned_split: num_new_tokens = self._mamba_block_aligned_split( request, num_new_tokens, num_new_local_computed_tokens, num_external_computed_tokens, ) if num_new_tokens == 0: break effective_lookahead_tokens = 0 if request.num_computed_tokens == 0 else self.num_lookahead_tokens # Determine if we need to allocate cross-attention blocks. num_encoder_tokens = 0 if self.is_encoder_decoder and request.has_encoder_inputs and encoder_inputs_to_schedule: num_encoder_tokens = sum(request.get_num_encoder_embeds(i) for i in encoder_inputs_to_schedule) new_blocks = self.kv_cache_manager.allocate_slots( request, num_new_tokens, num_new_computed_tokens=num_new_local_computed_tokens, new_computed_blocks=new_computed_blocks, num_lookahead_tokens=effective_lookahead_tokens, num_external_computed_tokens=num_external_computed_tokens, delay_cache_blocks=load_kv_async, num_encoder_tokens=num_encoder_tokens, full_sequence_must_fit=self.scheduler_reserve_full_isl, ) if new_blocks is None: if request.has_encoder_inputs: self.encoder_cache_manager.free(request) break if self.connector is not None: self.connector.update_state_after_alloc( request, self.kv_cache_manager.get_blocks(request_id), num_external_computed_tokens, ) if self.connector_prefix_cache_stats is not None and connector_prefix_cache_queries != 0: self.connector_prefix_cache_stats.record( num_tokens=connector_prefix_cache_queries, num_hits=connector_prefix_cache_hits, preempted=request.num_preemptions > 0, ) request = request_queue.pop_request() if load_kv_async: request.status = RequestStatus.WAITING_FOR_REMOTE_KVS step_skipped_waiting.prepend_request(request) request.num_computed_tokens = num_computed_tokens continue self.running.append(request) if self.log_stats: request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) if request.status == RequestStatus.WAITING: scheduled_new_reqs.append(request) elif request.status == RequestStatus.PREEMPTED: scheduled_resumed_reqs.append(request) else: raise RuntimeError(f"Invalid request status: {request.status}") if self.lora_config and request.lora_request: scheduled_loras.add(request.lora_request.lora_int_id) req_to_new_blocks[request_id] = self.kv_cache_manager.get_blocks(request_id) num_scheduled_tokens[request_id] = num_new_tokens token_budget -= num_new_tokens # Decode requests (num_new_tokens == 1) have negligible latency; # skip time_budget accounting so they don't starve other requests. if request.num_computed_tokens < request.num_prompt_tokens: time_budget -= self.profiling_chunk_manager.predict_time( num_new_tokens, request.num_computed_tokens ) request.status = RequestStatus.RUNNING request.num_computed_tokens = num_computed_tokens if encoder_inputs_to_schedule: scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule for i in encoder_inputs_to_schedule: self.encoder_cache_manager.allocate(request, i) if self.ec_connector is not None: self.ec_connector.update_state_after_alloc(request, i) encoder_compute_budget = new_encoder_compute_budget if external_load_encoder_input: for i in external_load_encoder_input: self.encoder_cache_manager.allocate(request, i) if self.ec_connector is not None: self.ec_connector.update_state_after_alloc(request, i) # Re-queue requests skipped in this pass ahead of older skipped items. if step_skipped_waiting: self.skipped_waiting.prepend_requests(step_skipped_waiting) # Check if the scheduling constraints are satisfied. total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens assert token_budget >= 0 assert len(self.running) <= self.max_num_running_reqs assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len(scheduled_running_reqs) <= len(self.running) # Get the longest common prefix among all requests in the running queue. num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"): if self.running: any_request_id = self.running[0].request_id num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id) # Construct the scheduler output. if self.use_v2_model_runner: scheduled_new_reqs = scheduled_new_reqs + scheduled_resumed_reqs scheduled_resumed_reqs = [] new_reqs_data = [ NewRequestData.from_request( req, req_to_new_blocks[req.request_id].get_block_ids(), req._all_token_ids, ) for req in scheduled_new_reqs ] else: new_reqs_data = [ NewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids()) for req in scheduled_new_reqs ] with record_function_or_nullcontext("schedule: make_cached_request_data"): cached_reqs_data = self._make_cached_request_data( scheduled_running_reqs, scheduled_resumed_reqs, num_scheduled_tokens, scheduled_spec_decode_tokens, req_to_new_blocks, ) self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) new_block_ids_to_zero = ( (self.kv_cache_manager.take_new_block_ids() or None) if self.needs_kv_cache_zeroing else None ) scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, scheduled_cached_reqs=cached_reqs_data, num_scheduled_tokens=num_scheduled_tokens, total_num_scheduled_tokens=total_num_scheduled_tokens, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, scheduled_encoder_inputs=scheduled_encoder_inputs, num_common_prefix_blocks=num_common_prefix_blocks, preempted_req_ids={req.request_id for req in preempted_reqs}, finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), new_block_ids_to_zero=new_block_ids_to_zero, ) if self.connector is not None: meta = self._build_kv_connector_meta(self.connector, scheduler_output) scheduler_output.kv_connector_metadata = meta if self.ec_connector is not None: ec_meta = self.ec_connector.build_connector_meta(scheduler_output) scheduler_output.ec_connector_metadata = ec_meta with record_function_or_nullcontext("schedule: update_after_schedule"): self._update_after_schedule(scheduler_output) return scheduler_output