MB5 proxy: session-affinity P routing (MB5_P_ROUTING=session)

The upstream mooncake_connector_proxy round-robins both P and D
selection. For agentic multi-turn sessions this destroys prefix-cache
reuse on the producer side — every turn of a session lands on a
different P, so the prefix-cache hit ratio collapses to 0 (observed in
the 6P+2D round-robin baseline) and every turn re-prefills from
scratch, piling extra load on the P pool.

Add an env-gated routing mode so the same proxy serves both arms of a
clean A/B:
  MB5_P_ROUTING=rr       round-robin (default, = upstream behavior)
  MB5_P_ROUTING=session  consistent md5 hash on X-Session-Id -> same
                         producer for all turns of a session

Decode side stays round-robin (load balance) in both modes — decode
KV is freshly transferred per turn, so D gains nothing from affinity
but everything from even load spreading.

mb5_launch.sh threads MB5_P_ROUTING through to the proxy and logs the
active mode. Default path is byte-for-byte the old behavior, so an
in-flight round-robin sweep is unaffected if this is redeployed.

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

microbench/fresh_setup/mb5_launch.sh

@@ -190,7 +190,9 @@ for port in "${all_ports[@]}"; do
 done
 
 if [ "${ROLES}" = "pd" ]; then
-    echo "[mb5] launching mooncake_connector_proxy on ${PROXY_PORT}"
+    P_ROUTING="${MB5_P_ROUTING:-rr}"
+    echo "[mb5] launching mooncake_connector_proxy on ${PROXY_PORT} (P routing=${P_ROUTING})"
+    MB5_P_ROUTING="${P_ROUTING}" \
     nohup python "${PROXY_SRC}" "${proxy_args[@]}" --port "${PROXY_PORT}" --host 0.0.0.0 \
         > "${LOGS_DIR}/proxy.log" 2>&1 &
     disown

microbench/fresh_setup/mb5_pd_proxy.py

@@ -3,6 +3,7 @@
 
 import argparse
 import asyncio
+import hashlib
 import ipaddress
 import itertools
 import os
@@ -227,6 +228,28 @@ def _parse_decode_urls(decode_list):
     return [url[0] for url in decode_list]
 
 
+# MB5: routing mode for the prefill (producer) side.
+#   "rr"      — round-robin (official upstream behavior)
+#   "session" — consistent hash on X-Session-Id, so all turns of a session
+#               land on the same producer and reuse its prefix cache.
+# Decode side stays round-robin (load balance) regardless.
+MB5_P_ROUTING = os.environ.get("MB5_P_ROUTING", "rr").lower()
+
+
+def get_prefill_by_session(app, session_id: str):
+    """Pick a (prefill_client, dp_rank) deterministically from session_id.
+
+    Uses a stable (non-PYTHONHASHSEED-dependent) hash so the mapping is
+    reproducible across processes. dp_size is usually 1 here (TP=1, no DP),
+    but we hash into the flat (client, dp_rank) slot space to stay correct
+    if a producer ever reports dp_size > 1.
+    """
+    clients = app.state.prefill_clients
+    slots = [(c, r) for c in clients for r in range(max(1, c.get("dp_size", 1)))]
+    h = int(hashlib.md5(session_id.encode()).hexdigest()[:8], 16)
+    return slots[h % len(slots)]
+
+
 def get_next_client(app, service_type: str):
     """
     Get the next client in round-robin fashion.
@@ -325,8 +348,17 @@ async def _handle_completions(api: str, request: Request):
         req_data = await request.json()
         request_id = str(uuid.uuid4())
 
-        # Get the next prefill client in round-robin fashion
-        prefill_client_info, prefill_dp_rank = get_next_client(request.app, "prefill")
+        # Select the prefill (producer) client.
+        if MB5_P_ROUTING == "session":
+            session_id = request.headers.get("X-Session-Id") or request_id
+            prefill_client_info, prefill_dp_rank = get_prefill_by_session(
+                request.app, session_id
+            )
+        else:
+            # Round-robin (official upstream behavior).
+            prefill_client_info, prefill_dp_rank = get_next_client(
+                request.app, "prefill"
+            )
 
         # Send request to prefill service
         asyncio.create_task(