feat(kvc): Option D - delegate seed/reseed admission to D worker

v4 (cap=16) saw 35% session-cap fallback because the local soft_cap
min(16, usable / target) evaluates to 1-2 for large agentic inputs.
The cap was hit not because D was full but because replay's heuristic
underestimated capacity.

This change makes worker admission_mode authoritative for ALL paths:

SGLang side:
- io_struct.py: DirectAppendAdmissionReqInput gains a `mode` field
  ("direct_append" | "seed", default "direct_append" preserves prior
  behavior).
- scheduler.py:admit_direct_append: when mode == "seed", skip the
  resident-on-D requirement and run the same capacity check + LRU
  eviction (maybe_trim_decode_session_cache) that direct_append uses.
  This lets D atomically decide if a new session can be admitted based
  on actual token_to_kv_pool_allocator state.

Replay side (replay.py):
- _query_decode_direct_admission gains a `mode` parameter.
- _reserve_decode_session_capacity: in worker admission_mode, the
  seed/reseed branch now queries D with mode="seed" and trusts the
  result, instead of estimating capacity from the residency snapshot.
- _should_admit_new_decode_session: in worker mode, skip the local
  soft_cap pre-check and let D decide. Same-D session fast-path is
  preserved.

Effects:
- Local hardcoded cap of 16 is bypassed under worker mode; D's real
  KV pool size is the only constraint.
- LRU eviction runs in D's process atomically with admission, so
  starvation (the v3 bimodal "lucky vs starved sessions" pattern)
  should resolve.

scripts/sweep_tp1_v5_optD.sh added to run the same 1P7D / 2P6D
configs as v4 with the new admission path.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
kzlin
2026-04-28 23:40:03 +08:00
parent 74194e660a
commit 6e5ed8da80
4 changed files with 173 additions and 2 deletions

View File

@@ -651,6 +651,7 @@ async def _query_decode_direct_admission(
session_id: str,
uncached_input_tokens: int,
output_tokens: int,
mode: str = "direct_append",
) -> dict[str, Any]:
try:
response = await client.post(
@@ -659,6 +660,7 @@ async def _query_decode_direct_admission(
"session_id": session_id,
"uncached_input_tokens": max(0, uncached_input_tokens),
"output_tokens": max(0, output_tokens),
"mode": mode,
},
timeout=_ADMISSION_PROBE_TIMEOUT_S,
)
@@ -913,6 +915,7 @@ def _should_admit_new_decode_session(
session: DirectSessionState,
direct_sessions: dict[str, DirectSessionState],
treat_as_fresh_session: bool,
admission_mode: KvCacheAdmissionMode = "router",
) -> bool:
if (
not treat_as_fresh_session
@@ -920,6 +923,11 @@ def _should_admit_new_decode_session(
and session.server_url == server_url
):
return True
if admission_mode == "worker":
# Defer the capacity decision to D's admit_direct_append (mode=seed),
# which checks real KV pool availability and runs LRU eviction. The
# local soft cap is router-mode only.
return True
open_sessions = sum(
1
for candidate in direct_sessions.values()
@@ -1331,6 +1339,7 @@ async def _reserve_decode_session_capacity(
session_id=session.session_id,
uncached_input_tokens=max(0, request.input_length - current_tokens),
output_tokens=request.output_length,
mode="direct_append",
)
if not bool(admission.get("resident")):
return False, 0, 0, 0, str(admission.get("reason") or "d-session-not-resident")
@@ -1355,6 +1364,41 @@ async def _reserve_decode_session_capacity(
None,
)
# Seed / reseed path: ask D itself via the seed-mode admission endpoint
# instead of estimating capacity from a stale router-state snapshot. D
# will run LRU eviction internally to make room. Falls through to the
# legacy router-state logic below if the endpoint is unavailable.
seed_admission = await _query_decode_direct_admission(
client=client,
server_url=server_url,
session_id=session.session_id,
uncached_input_tokens=max(0, request.input_length - current_tokens),
output_tokens=request.output_length,
mode="seed",
)
seed_reason = seed_admission.get("reason")
if seed_reason != "admission-query-failed":
if not bool(seed_admission.get("can_admit")):
return (
False,
0,
int(seed_admission.get("evicted_session_count", 0) or 0),
0,
str(seed_reason or "d-no-space"),
)
reserved_tokens = int(
seed_admission.get("required_tokens", required_extra_tokens)
or required_extra_tokens
)
_add_reserved_tokens(residency, server_url, reserved_tokens)
return (
True,
reserved_tokens,
int(seed_admission.get("evicted_session_count", 0) or 0),
0,
None,
)
session_cache, max_total_num_tokens, reserved_decode_tokens = (
await _fetch_decode_server_state(
client=client,
@@ -1906,6 +1950,7 @@ async def _execute_request(
session=decode_session,
direct_sessions=direct_sessions,
treat_as_fresh_session=True,
admission_mode=config.kvcache_admission_mode,
)
if not admit_new_decode_session:
can_seed = False
@@ -2060,6 +2105,7 @@ async def _execute_request(
session=decode_session,
direct_sessions=direct_sessions,
treat_as_fresh_session=True,
admission_mode=config.kvcache_admission_mode,
)
if not admit_new_decode_session:
can_seed = False