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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user