MB5 driver updates: PD-proxy + snapshot instrument + launcher tweaks

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

microbench/fresh_setup/instrument_kv_snapshot.py

@@ -194,24 +194,35 @@ MOONCAKE_PATCHES = [
 ]
 
 # ---------- Patch 4: vLLM 0.18.1 PD-consumer metrics counter underflow ------
-# In PromptTokenStats.update_from_output, local_cache_hit is computed as
-# (num_cached_tokens + recomputed - num_external_computed_tokens). On a
-# kv_consumer, a remote KV transfer can report more external-computed tokens
-# than the scheduler's cached count (esp. on a KV-load failure for a large
-# request), driving local_cache_hit negative. loggers.record() then calls
-# Counter.inc() with that negative value and prometheus_client raises
+# On a kv_consumer, a KV-load failure (Mooncake transfer returns -1 when the
+# D-pool is full) makes vLLM emit an iteration-stats "correction" with NEGATIVE
+# token deltas: PromptTokenStats fields (local_cache_hit, cached_tokens, ...)
+# AND iteration_stats.{num_prompt_tokens, num_generation_tokens} all go below
+# zero. Every Counter.inc() in loggers.record() then trips prometheus_client's
 # "Counters can only be incremented by non-negative amounts.", which kills the
 # EngineCore — turning one failed request into a total config collapse.
-# We clamp the per-source counts to >= 0 right before they are recorded.
-LOGGERS_ANCHOR = "        pts = iteration_stats.prompt_token_stats\n"
+#
+# Clamp every field that feeds a Counter.inc() in record() to >= 0. We anchor
+# right after the `if iteration_stats is None: return` guard so the clamp runs
+# before the first inc() (the corrupted-requests / preempted / prompt-token
+# counters at the top of the method).
+LOGGERS_ANCHOR = "        if iteration_stats is None:\n            return\n"
 LOGGERS_INSERT = (
     f"        {START_MARK}\n"
-    f"        if pts.local_cache_hit < 0:\n"
-    f"            pts.local_cache_hit = 0\n"
-    f"        if pts.computed < 0:\n"
-    f"            pts.computed = 0\n"
-    f"        if pts.external_kv_transfer < 0:\n"
-    f"            pts.external_kv_transfer = 0\n"
+    f"        _mb5_pts = iteration_stats.prompt_token_stats\n"
+    f"        for _mb5_o, _mb5_a in (\n"
+    f"            (iteration_stats, 'num_prompt_tokens'),\n"
+    f"            (iteration_stats, 'num_generation_tokens'),\n"
+    f"            (iteration_stats, 'num_preempted_reqs'),\n"
+    f"            (iteration_stats, 'num_corrupted_reqs'),\n"
+    f"            (_mb5_pts, 'computed'),\n"
+    f"            (_mb5_pts, 'local_cache_hit'),\n"
+    f"            (_mb5_pts, 'external_kv_transfer'),\n"
+    f"            (_mb5_pts, 'cached_tokens'),\n"
+    f"            (_mb5_pts, 'recomputed_tokens'),\n"
+    f"        ):\n"
+    f"            if getattr(_mb5_o, _mb5_a, 0) < 0:\n"
+    f"                setattr(_mb5_o, _mb5_a, 0)\n"
     f"        {END_MARK}\n"
 )