[CI] Fix vLLM integration CI (fp8 MoE, continue decode, routed-experts tests)

.buildkite/vllm_lkg.version

@@ -1 +1 @@
-e7bb7e2f20752ad38f8d44a6e7258cf002e29bea
+2fb320d8304381e224dd27fbf8c13c30b4e1020f

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

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()
 

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,25 +47,27 @@ 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,
                                                           new_token_ids,
                                                           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
 

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)