gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(snapshot): dedicated GPU snapshot_buf replaces kv_pool alloc

Browse Source 
Implements the design in docs/SNAPSHOT_STORE_REFACTOR_ZH.md to fix
the alloc-failed death loop that killed D→P in E4-v4/v5 (167 sync
attempts, 0 OK because P's kv_pool was busy with its own prefill).

Mechanism change:
  OLD prepare_receive: token_to_kv_pool_allocator.alloc(N) — 90%+ failure
  NEW prepare_receive: SnapshotBufAllocator.alloc(slab_bytes) carves a
                       range from an 8 GB GPU buffer dedicated to
                       snapshot reception, decoupled from kv_pool

  OLD finalize_ingest: just radix.insert with pre-alloc'd slots
  NEW finalize_ingest: kv_pool.alloc NOW + GPU memcpy snapshot_buf →
                       k_buffer/v_buffer + radix.insert

Wire schema changed (clean break, no back-compat):
  PrepareReceiveReqOutput  swaps k/v_base_ptrs + slot_indices  for
                           snapshot_buf_base_ptr + k/v_layer_offsets +
                           num_tokens
  DumpReqInput             swaps target_k/v_base_ptrs + target_slot_indices
                           for target_snapshot_buf_base +
                           target_k/v_layer_offsets
  FinalizeIngestReqInput   drops slot_indices (P resolves at ingest)

Controller adds:
  SnapshotBufAllocator: first-fit free-list with 4 KB alignment
  ingest_snapshot_into_kvpool: GPU→GPU copy + radix insert

Configurable buffer size via SGLANG_SNAPSHOT_LINK_BUF_BYTES env
(default 8 GB, scales down to 1 GB if alloc fails).

Removed runtime leak-check accommodation since prepare_receive no
longer touches kv_pool.

Total: ~365 LOC including alloc helper; smoke-test verification next.
This commit is contained in:
Claude Code Agent
2026-05-13 14:18:23 +08:00
parent 6be5f9b57e
commit 2dfe22ab20
5 changed files with 465 additions and 285 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/agentic_pd_hybrid/replay.py
View File

@@ -2187,9 +2187,9 @@ async def _attempt_d_to_p_sync(
json={
"session_id": request.session_id,
"target_snapshot_session_id": prep["snapshot_session_id"],
"target_k_base_ptrs": prep["k_base_ptrs"],
"target_v_base_ptrs": prep["v_base_ptrs"],
"target_slot_indices": prep["slot_indices"],
"target_snapshot_buf_base": prep["snapshot_buf_base_ptr"],
"target_k_layer_offsets": prep["k_layer_offsets"],
"target_v_layer_offsets": prep["v_layer_offsets"],
"target_stride_k_bytes": prep["stride_k_bytes"],
"target_stride_v_bytes": prep["stride_v_bytes"],
},
@@ -2220,15 +2220,13 @@ async def _attempt_d_to_p_sync(
# for the first N — use that as best-available approximation.
tokens = list(getattr(request, "input_token_ids", []) or [])
if not tokens:
# No token_ids available — can't insert into radix. P will fall back
# to normal prefill but will have wasted slots. Discard.
# No token_ids → can't insert into radix; tell P to free the slab.
try:
await client.post(
f"{prefill_url}/_snapshot/finalize_ingest",
json={
"session_id": request.session_id,
"token_ids": [],
"slot_indices": prep["slot_indices"],
},
timeout=15.0,
)
@@ -2242,7 +2240,7 @@ async def _attempt_d_to_p_sync(
)
return {"status": "no-tokens-discard", "bytes_pushed": dump.get("bytes_pushed", 0)}
n = min(len(tokens), len(prep["slot_indices"]))
n = min(len(tokens), int(prep.get("num_tokens", 0)))
t_fin0 = time.perf_counter()
try:
fin_resp = await client.post(
@@ -2250,7 +2248,6 @@ async def _attempt_d_to_p_sync(
json={
"session_id": request.session_id,
"token_ids": tokens[:n],
"slot_indices": prep["slot_indices"][:n],
},
timeout=30.0,
)