From e1e79b1ff5c341343f6170cfcc6a918a916f035a Mon Sep 17 00:00:00 2001 From: Procr <193802945+procr1337@users.noreply.github.com> Date: Tue, 28 Jul 2026 17:00:47 +0000 Subject: [PATCH] Deduplicate lockstep MLA pending job cleanup Lockstep MLA groups share physical block IDs, so a transfer job may contain the same block ID once per logical group. Iterate the deduplicated set of block IDs in _remove_pending_job() so each (block, job) pair is only removed once from the set-backed pending-job index, instead of crashing with KeyError on the second occurrence. Fixes a KeyError in _remove_pending_job() observed when KV cache offloading (--kv-offloading-size) is enabled while serving GLM-5.2 (MLA + decode context parallelism), reproduced by offload-fix-glm52/probe_gpu_cache.py. --- .../kv_transfer/kv_connector/v1/offloading/scheduler.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 5e98c1266e20..3ae65dc523e9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -438,7 +438,10 @@ def _generate_job_id(self) -> int: return job_id def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None: - for bid in block_ids or (): + # Lockstep MLA groups share physical block IDs, so a transfer job may + # contain the same ID once per logical group. The pending-job index is + # set-backed and therefore registers each (block, job) pair only once. + for bid in set(block_ids or ()): pending = self._block_id_to_pending_jobs[bid] pending.remove(job_id) if not pending: