diff --git a/.buildkite/vllm_lkg.version b/.buildkite/vllm_lkg.version index f5401896d4..73c1062ca1 100644 --- a/.buildkite/vllm_lkg.version +++ b/.buildkite/vllm_lkg.version @@ -1 +1 @@ -e7bb7e2f20752ad38f8d44a6e7258cf002e29bea +2fb320d8304381e224dd27fbf8c13c30b4e1020f diff --git a/tests/core/test_sched_utils.py b/tests/core/test_sched_utils.py index 1f31c04ffe..c0d00a26bc 100644 --- a/tests/core/test_sched_utils.py +++ b/tests/core/test_sched_utils.py @@ -30,11 +30,15 @@ def restore_scheduler_patch(): from vllm.v1.core.sched.async_scheduler import AsyncScheduler from vllm.v1.core.sched.scheduler import Scheduler + async_had_own = "_update_request_with_output" in AsyncScheduler.__dict__ originals = (Scheduler._update_request_with_output, Scheduler.__init__, AsyncScheduler._update_request_with_output) yield - (Scheduler._update_request_with_output, Scheduler.__init__, - AsyncScheduler._update_request_with_output) = originals + Scheduler._update_request_with_output, Scheduler.__init__ = originals[:2] + if async_had_own: + AsyncScheduler._update_request_with_output = originals[2] + elif "_update_request_with_output" in AsyncScheduler.__dict__: + del AsyncScheduler._update_request_with_output for cls in (Scheduler, AsyncScheduler): if "_continue_decode_patched" in cls.__dict__: del cls._continue_decode_patched diff --git a/tests/runner/test_routed_experts_slots.py b/tests/runner/test_routed_experts_slots.py index 1984ece02f..30ea646a9b 100644 --- a/tests/runner/test_routed_experts_slots.py +++ b/tests/runner/test_routed_experts_slots.py @@ -379,7 +379,16 @@ def _worker(num_blocks=16, dp_size=1): def _connector(hash_block_size=_BLOCK_SIZE): - if inspect.signature(AuxOutputSchedulerConnector).parameters: + params = inspect.signature(AuxOutputSchedulerConnector).parameters + if "kv_cache_manager" in params: + manager = SimpleNamespace(block_size=_BLOCK_SIZE, block_pool=None) + kv_cache_manager = SimpleNamespace( + coordinator=SimpleNamespace(full_attention_group_id=0, + single_type_managers=[manager]), + block_pool=SimpleNamespace(hash_block_size=hash_block_size), + kv_cache_config=SimpleNamespace(num_blocks=16)) + return AuxOutputSchedulerConnector(kv_cache_manager) + if params: return AuxOutputSchedulerConnector(hash_block_size) return AuxOutputSchedulerConnector() diff --git a/tpu_inference/core/sched/utils.py b/tpu_inference/core/sched/utils.py index e766d9c6f8..a9db241856 100644 --- a/tpu_inference/core/sched/utils.py +++ b/tpu_inference/core/sched/utils.py @@ -27,9 +27,9 @@ def patch_vllm_scheduler_for_continue_decode(): or 2 * (max_decode_steps - 1) under async_scheduling). 2. patched_update_base: Reconciles request.num_computed_tokens on host by adding the extra (N - 1) tokens generated on-device once model output returns. - 3. patched_async_update_request_with_output: Pre-compensates in-flight - num_output_placeholders in AsyncScheduler by adding (N - 1) before - subtracting N, preventing placeholder underflow in async mode. + 3. Placeholder pre-compensation in patched_update_base: adds (N - 1) to + num_output_placeholders before vLLM subtracts N, preventing placeholder + underflow in async mode. """ from vllm.v1.core.sched.scheduler import Scheduler @@ -47,6 +47,11 @@ def patched_update_base(scheduler_self, stale = is_stale or getattr(scheduler_self, "_cd_stale_in_flight", False) computed_before = request.num_computed_tokens + num_new_tokens = len(new_token_ids) + if not stale and num_new_tokens > 1 and \ + request.num_output_placeholders > 0: + request.num_output_placeholders += num_new_tokens - 1 + # Original update appends new_token_ids to request output and trims on stop token. res_token_ids, stopped = original_update_base(scheduler_self, request, @@ -54,18 +59,15 @@ def patched_update_base(scheduler_self, is_stale=stale, **kwargs) - # schedule() only incremented num_computed_tokens (and in-flight - # placeholders under AsyncScheduler) by 1. Reconcile against the - # trimmed output count, accounting for any stop-trimmed subtraction - # performed by the base scheduler. A stale delivery predates - # preemption rollback and must not advance either counter. + # schedule() only incremented num_computed_tokens by 1. Reconcile + # against the trimmed output count, accounting for any stop-trimmed + # subtraction performed by the base scheduler. A stale delivery + # predates preemption rollback and must not advance the counter. if res_token_ids and not stale: delta = (computed_before + len(res_token_ids) - 1) - request.num_computed_tokens if delta > 0: request.num_computed_tokens += delta - if getattr(scheduler_self, "_cd_in_async_update", False): - request.num_output_placeholders += delta return res_token_ids, stopped diff --git a/tpu_inference/layers/vllm/quantization/fp8.py b/tpu_inference/layers/vllm/quantization/fp8.py index 9630445a15..d2d49ff495 100644 --- a/tpu_inference/layers/vllm/quantization/fp8.py +++ b/tpu_inference/layers/vllm/quantization/fp8.py @@ -388,6 +388,7 @@ def __init__(self, )) self.fp8_backend = None + self.per_act_token_quant = False self.mesh = mesh self.moe_backend = select_moe_backend_from_fused_moe_config(self.moe)