Compare commits
10 Commits
9aac36fd89
...
f09562123b
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f09562123b | ||
|
|
9cca2c60c9 | ||
|
|
5c09a3a0cb | ||
|
|
19612ff3a3 | ||
|
|
a953346a0c | ||
|
|
2dfe22ab20 | ||
|
|
6be5f9b57e | ||
|
|
f926a7b87d | ||
|
|
552f3f564e | ||
|
|
051d9220f4 |
5
.gitignore
vendored
5
.gitignore
vendored
@@ -13,6 +13,11 @@ src/*.egg-info
|
||||
outputs/
|
||||
|
||||
# Vendored dependencies. Track only the maintained SGLang fork/snapshot.
|
||||
# third_party/traces/ holds the replay trace files used by the benchmark
|
||||
# (~56 MB each) for convenient transfer between hosts; they would otherwise
|
||||
# live under outputs/ but outputs/ is gitignored.
|
||||
third_party/*
|
||||
!third_party/sglang/
|
||||
!third_party/agentic-kvcache/
|
||||
!third_party/traces/
|
||||
*.log
|
||||
|
||||
3
.gitmodules
vendored
Normal file
3
.gitmodules
vendored
Normal file
@@ -0,0 +1,3 @@
|
||||
[submodule "third_party/agentic-kvcache"]
|
||||
path = third_party/agentic-kvcache
|
||||
url = git@ipads.se.sjtu.edu.cn:scaleaisys/projects/agentic-kvcache.git
|
||||
202
docs/E4_V8_RESULTS_ZH.md
Normal file
202
docs/E4_V8_RESULTS_ZH.md
Normal file
@@ -0,0 +1,202 @@
|
||||
# E4-v8 完整结果 — KVC 在真实节奏 trace 上的表现
|
||||
|
||||
**日期**:2026-05-13
|
||||
**Status**:实验跑完
|
||||
**Run**:`outputs/e4p_kvc_v2_d_to_p_sync_pressured_50sess/...20260513T075500Z/`
|
||||
**前置**:`docs/SNAPSHOT_STORE_REFACTOR_ZH.md`、`docs/E4_VS_E1_RESULTS_ZH.md`
|
||||
|
||||
---
|
||||
|
||||
## 0. TL;DR
|
||||
|
||||
V8 跑 **真实节奏 trace**(`third_party/traces/qwen35-swebench-50sess.jsonl`,4449 reqs × 52 sessions,原始 5.44h 时间线)在 TIME_SCALE=2 压缩到 ~2.7h wall clock:
|
||||
|
||||
| 指标 | V8 实测 |
|
||||
|---|---:|
|
||||
| 总请求 | 4449 |
|
||||
| Failure / Error / Abort | **0 / 0 / 0** |
|
||||
| Success rate | **100%** |
|
||||
| Latency mean / p50 / p90 / p99 | 1.28s / 0.51s / 3.17s / **7.44s** |
|
||||
| **TTFT mean / p50 / p90 / p99** | **49ms / 40ms / 68ms / 167ms** |
|
||||
| Direct-to-D fast path | **96.4%** (4291/4449) |
|
||||
| Reseed paths | 51 (1.1%) |
|
||||
| D→P sync OK | **0** (architecturally wired but no successful pushes — see §3) |
|
||||
|
||||
**关键结论**:先前 E1 和 E4-v3 上 TTFT 上百秒的"灾难数字"是**burst trace 排队累积的人为产物**。在真实节奏 SWE-Bench trace 上,**KVC 表现为亚秒到个位数秒的正常生产 serving 性能**。
|
||||
|
||||
---
|
||||
|
||||
## 1. 实验配置
|
||||
|
||||
```
|
||||
Workload: third_party/traces/qwen35-swebench-50sess.jsonl
|
||||
4449 reqs / 52 sessions / 5.44h original wall-clock span
|
||||
per-session inter-turn p50: 2.53s (real SWE-agent timing)
|
||||
input length p50: 27K, p99: 92K, max: 104K
|
||||
|
||||
Compression: TIME_SCALE=2 → 2.72h actual run-time
|
||||
Topology: 1P + 3D, 4× H200 80GB single-node
|
||||
RDMA: mlx5_60 NDR 400Gb / mooncake
|
||||
Model: Qwen3-30B-A3B-Instruct-2507 (TP=1)
|
||||
Concurrency: 32
|
||||
|
||||
Memory: PREFILL_MEM_FRAC=0.7 / DECODE_MEM_FRAC=0.8
|
||||
snapshot_buf=16 GB on each worker (alloc succeeded)
|
||||
|
||||
KVC config: --kvcache-load-floor-bonus 200
|
||||
--kvcache-migration-reject-threshold 1
|
||||
--kvcache-direct-max-uncached-tokens 8192
|
||||
--enable-d-to-p-sync (with SnapshotStore refactor)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. 完整 v8 数据
|
||||
|
||||
### 2.1 Headline
|
||||
|
||||
```
|
||||
request_count : 4449
|
||||
abort_count : 0
|
||||
error_count : 0
|
||||
failure_count : 0
|
||||
cache_hit_request_count : 4446 / 4449 = 99.9%
|
||||
mean cached_tokens : 30,513 / req (out of avg 32K input)
|
||||
```
|
||||
|
||||
### 2.2 Latency / TTFT
|
||||
|
||||
```
|
||||
count mean p50 p90 p99
|
||||
latency_stats_s 4449 1.28 0.51 3.17 7.44 s
|
||||
ttft_stats_s 4449 0.049 0.040 0.068 0.167 s ← p99 = 167ms
|
||||
```
|
||||
|
||||
### 2.3 Execution_mode 分布
|
||||
|
||||
```
|
||||
kvcache-direct-to-d-session 4291 (96.4%) ← KVC 独特 fast path
|
||||
pd-router-turn1-seed 52 ( 1.2%) ← 每个 session 第一个 turn
|
||||
pd-router-fallback-session-not-resident-seed-filter 52 ( 1.2%) ← seed-filter 早 turn fallback
|
||||
pd-router-d-session-reseed 47 ( 1.1%) ← 真正的 reseed (session 曾在 D)
|
||||
pd-router-fallback-real-large-append-session-cap 3
|
||||
pd-router-fallback-session-not-resident-session-cap 1
|
||||
pd-router-policy-no-bypass-reseed 1
|
||||
pd-router-real-large-append-reseed 1
|
||||
pd-router-session-not-resident-reseed 1
|
||||
-----
|
||||
4449
|
||||
```
|
||||
|
||||
### 2.4 Per-decode load
|
||||
|
||||
```
|
||||
decode-0: 1505 bindings (33.8%)
|
||||
decode-1: 1497 bindings (33.6%)
|
||||
decode-2: 1447 bindings (32.5%)
|
||||
```
|
||||
|
||||
负载完美均衡(load-floor bonus K=200 起作用)。
|
||||
|
||||
---
|
||||
|
||||
## 3. D→P snapshot link 状态(重构验证)
|
||||
|
||||
**SnapshotStore 重构(commit 2dfe22a)成功**:
|
||||
- 旧设计 prepare_receive 用 `token_to_kv_pool_allocator.alloc(N)` 抢 P 的 KV pool slot → 90%+ alloc-failed
|
||||
- 新设计 prepare_receive 从独立 16 GB GPU `snapshot_buf` 分配 slab → **0 alloc-failed**
|
||||
|
||||
```
|
||||
sync events total: 102
|
||||
by (stage, reason):
|
||||
('dump', 'session-not-resident'): 96 (D 端 session 已 evict 或从未 resident)
|
||||
('prepare', 'snapshot-buf-full'): 6 (snapshot_buf 偶尔满)
|
||||
('ok', None): 0 (无成功 push)
|
||||
```
|
||||
|
||||
**为什么 0 OK?**
|
||||
|
||||
mem_fraction=0.8 让 D 的 trim 机制总是成功 → admission 不拒绝 → reseed path 不通过"D 曾持有 session"分支触发,而是通过 first-turn-fallback 等路径触发,那些路径下 D 端**从未持有** session,dump 必然失败。
|
||||
|
||||
102 个 sync 事件中:
|
||||
- 96 个 dump session-not-resident:包含 52 个 turn-1 first-seed-fallback(session 从未 resident)+ 44 个其他 fallback
|
||||
- 6 个 snapshot-buf-full:偶尔出现,证明 buffer 在 working
|
||||
|
||||
D→P **底层链路 + agentic orchestration 都已就位**——只是 agentic 触发的 reseed 场景里 D 端 session 不存在。要让 D→P 真正 fire OK,需要:
|
||||
1. 给 D-side SessionAwareCache 加 "pending-snapshot pinning" 保护,让 evict 不打掉等 sync 的 session
|
||||
2. **或者** 加 D-side push-on-eviction:D 端在 evict 一个 session 前先 push 给 P(D-driven 主动模式)
|
||||
3. **或者** 调小 mem_fraction 让 admission 真正拒绝("还有 session 时就拒"),让 reseed 命中真正"session 仍在 D"的场景
|
||||
|
||||
---
|
||||
|
||||
## 4. 跟之前几次实验对比
|
||||
|
||||
| Run | Trace | failures | TTFT p99 | Latency p99 | D→P OK |
|
||||
|---|---|---:|---:|---:|---:|
|
||||
| E1 (naive PD) | inferact 1285 burst | 6.6% | **207s** | 219s | n/a |
|
||||
| E4-v3 (KVC + load-floor, no D→P fix) | inferact 1285 burst | 0% | 225s | 234s | n/a |
|
||||
| E4-v4/v5 (KVC + D→P, bug) | inferact 1285 burst | 0% / 12% | similar | similar | 0 (logger NameError or alloc-fail) |
|
||||
| **E4-v8 (refactor + real trace)** | **swebench 4449 real-time** | **0%** | **167ms** | **7.4s** | 0 (D-side eviction timing) |
|
||||
|
||||
E1 vs v8 的数字差距巨大但**不直接可比**——因为 trace 完全不同:
|
||||
- E1 burst trace:所有 1285 req 在 t=0 全部到达 → 队列累积 → TTFT 上百秒
|
||||
- v8 real-time trace:req 按 2.53s p50 inter-turn 真实节奏到达 → 系统不饱和 → TTFT 几十 ms
|
||||
|
||||
**To be fair**: 要跟 v8 真实对比 KVC vs naive PD,需要也用 swebench trace 跑一遍 naive PD。这是下一步。
|
||||
|
||||
---
|
||||
|
||||
## 5. 给 D→P sync 真正生效的下一步
|
||||
|
||||
按重要性排序:
|
||||
|
||||
### P1:让 sync 能在 reseed 时 fire OK
|
||||
|
||||
**最直接的方法**:在 agentic 监测到 admission 拒绝时**立即**触发 dump(**在 D evict 之前**)。当前实现是 reseed 决策做完才 dump,已经太晚。
|
||||
|
||||
**方案**:
|
||||
1. 改 agentic `admit_direct_append` 调用之后,如果返回 reason=`no-space`,**立即 invoke sync** 到 source D,把 session KV 推给 P → 然后 retry admit 或转 fallback
|
||||
2. 在 D-side SessionAwareCache 加 "pending-snapshot pinning",让 eviction 暂时 skip 这个 session
|
||||
|
||||
### P2:D-driven 主动模式
|
||||
|
||||
每次 D 完成 `cache_finished_req` 后,**异步**推 incremental KV 给所有注册的 P。这是设计 doc §2.5 提到的方向。开销显著(每次 turn 都推流量)但确保 sync 一直有数据。
|
||||
|
||||
### P3:mem-fraction tuning
|
||||
|
||||
把 decode mem-fraction 调到 0.5-0.55,让 admission 自然拒绝更多,从而 reseed 路径命中真正的"session-resident-on-some-D"分支。但这降低 throughput。
|
||||
|
||||
---
|
||||
|
||||
## 6. 对 ProjectGoal 的回答
|
||||
|
||||
> 寻找 KVC 如何才能在保持自身独特性的情况下胜过 naive PD Disagg
|
||||
|
||||
**V8 数据回答**:在真实节奏 SWE-Bench workload 下:
|
||||
- **96.4% 请求走 direct-to-D fast path**(KVC 独特价值)
|
||||
- TTFT p99 = 167ms,latency p99 = 7.44s
|
||||
- **0% failure**
|
||||
- D→P snapshot 底层架构 ready,但 trigger 的时机问题导致目前 OK rate=0
|
||||
|
||||
**要全面证明 KVC > naive PD**,需要补:
|
||||
- 用 swebench trace 跑一次 naive PD baseline → 直接对比
|
||||
- 修 P1(agentic admission-rejection 时立即 sync)→ 让 D→P 真起作用
|
||||
|
||||
---
|
||||
|
||||
## 7. 当前 branch HEAD
|
||||
|
||||
```
|
||||
git log --oneline -5
|
||||
9cca2c6 feat(experiments): expose PREFILL_MEM_FRAC + plumb --prefill-mem-fraction-static
|
||||
5c09a3a feat(experiments): per-second GPU util sampler in E4-pressured sweep
|
||||
19612ff feat(experiments): parameterize TIME_SCALE in E4-pressured sweep
|
||||
a953346 feat(experiments): E4-pressured points at third_party/traces SWE-Bench trace
|
||||
2dfe22a refactor(snapshot): dedicated GPU snapshot_buf replaces kv_pool alloc
|
||||
```
|
||||
|
||||
`outputs/e4p_kvc_v2_d_to_p_sync_pressured_50sess/` 包含完整 metrics + structural logs + GPU util CSV,会另外做对比图(与 swebench-on-naive-PD 一旦跑出)。
|
||||
|
||||
---
|
||||
|
||||
**核心句**:V8 数据把 KVC TTFT 数字从 100+s(burst trace 假象)拉回 167ms(真实 workload),证明 KVC 在真实在线 serving 节奏下表现优异。D→P snapshot link 架构全栈 deploy 完毕但 trigger 时机仍需调整才能真正 fire。
|
||||
174
docs/SNAPSHOT_STORE_REFACTOR_ZH.md
Normal file
174
docs/SNAPSHOT_STORE_REFACTOR_ZH.md
Normal file
@@ -0,0 +1,174 @@
|
||||
# SnapshotStore 重构(解决 P-side alloc-failed 死局)
|
||||
|
||||
**日期**:2026-05-13
|
||||
**Status**:设计阶段,开始实施
|
||||
**根因**:`docs/E4_VS_E1_RESULTS_ZH.md` §3 + E4-v4/v5 forensic 显示 D→P sync 167 次尝试 0 OK,全部因 `prepare_receive` 试图从 `token_to_kv_pool_allocator.alloc(N)` 拿 N 个 slot 而 P 的池被自己 prefill 工作占满
|
||||
|
||||
---
|
||||
|
||||
## 0. TL;DR
|
||||
|
||||
- 当前 P-side `prepare_receive` 用 `token_to_kv_pool_allocator.alloc(N)` 抢 kv_pool slot —— 跟 P 自己的 prefill 工作直接争抢资源 → 90%+ 时间 alloc-failed
|
||||
- 重构方向:**P-side 用独立 GPU buffer 接收 snapshot**,与 kv_pool 解耦
|
||||
- 在 finalize_ingest 时才把 snapshot bytes copy 进 kv_pool slots(此时可以等更优的时机)
|
||||
- ~250 LOC 新代码,主要在 `disaggregation/snapshot/controller.py`
|
||||
|
||||
---
|
||||
|
||||
## 1. 当前实现的死局
|
||||
|
||||
```
|
||||
prepare_receive(sid, num_tokens=50000):
|
||||
indices = self.token_to_kv_pool_allocator.alloc(50000)
|
||||
if indices is None:
|
||||
return ok=False, reason="alloc-failed" ← 90%+ 时间走这里
|
||||
return slot_indices = indices.tolist()
|
||||
```
|
||||
|
||||
`alloc(50000)` 在 P 池中找 50000 个 contiguous 空 slot。当 P 正在 prefill 自己的 request 时(这是 P 的常态),池里大部分 slot 被锁定 → 找不出 50K 个空闲的 → fail.
|
||||
|
||||
E4-v5 167 次 sync 尝试统计:
|
||||
- 148 个 alloc-failed(**88%**)
|
||||
- 19 个 session-not-resident(D 端已 evict)
|
||||
- 0 个 OK
|
||||
|
||||
---
|
||||
|
||||
## 2. 新设计:PrefillSnapshotStore 侧表
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────────┐
|
||||
│ P worker scheduler │
|
||||
│ │
|
||||
│ kv_pool (existing, owned by P's prefill work) │
|
||||
│ ┌────────────────────────────────────────────────┐ │
|
||||
│ │ k_buffer[0..L]: (max_tokens, head, dim) │ │
|
||||
│ │ v_buffer[0..L]: (max_tokens, head, dim) │ │
|
||||
│ └────────────────────────────────────────────────┘ │
|
||||
│ │
|
||||
│ snapshot_buf (NEW, dedicated for D→P snapshot reception) │
|
||||
│ ┌────────────────────────────────────────────────┐ │
|
||||
│ │ pinned GPU tensor of size SNAPSHOT_BUF_BYTES │ │
|
||||
│ │ (default 8 GB) │ │
|
||||
│ │ • registered with mooncake (one-time at init) │ │
|
||||
│ │ • slab-allocator manages free space │ │
|
||||
│ └────────────────────────────────────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────────┘
|
||||
|
||||
Flow:
|
||||
1. prepare_receive(sid, N):
|
||||
slab = snapshot_buf_allocator.alloc(N * per_token_bytes_total)
|
||||
record = (sid, slab_offset, N)
|
||||
return (snapshot_buf_base + slab_offset for K_L, V_L per layer)
|
||||
← never blocks on kv_pool
|
||||
|
||||
2. (out-of-band) D pushes KV bytes into the slab via mooncake RDMA
|
||||
|
||||
3. finalize_ingest(sid, token_ids):
|
||||
record = pop ingest_record[sid]
|
||||
slots = token_to_kv_pool_allocator.alloc(N) ← can fail here
|
||||
if alloc-failed:
|
||||
snapshot_buf_allocator.free(record.slab)
|
||||
return ok=False, reason=alloc-failed-on-finalize
|
||||
# copy snapshot_buf[layer L][token range] → kv_pool.k_buffer[L][slots]
|
||||
for L in range(layer_num):
|
||||
kv_pool.k_buffer[L][slots] = snapshot_buf[K_L_offset : K_L_offset + N * K_stride].view(N, head, dim)
|
||||
kv_pool.v_buffer[L][slots] = snapshot_buf[V_L_offset : V_L_offset + N * V_stride].view(N, head, dim)
|
||||
tree_cache.insert(InsertParams(key=token_ids, value=slots))
|
||||
snapshot_buf_allocator.free(record.slab)
|
||||
return ok=True
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. 关键 design choices
|
||||
|
||||
| 决策 | 选择 | 原因 |
|
||||
|---|---|---|
|
||||
| Snapshot buffer 存哪 | GPU memory | 与 D RDMA 目标对称(D 端 KV 也在 GPU),避免 host↔device 拷贝 |
|
||||
| 默认大小 | **8 GB** | Qwen3-30B 一个 ~50K-token session 的 KV ~5 GB;8 GB 让我们至少 hold 一个 + 部分备份 |
|
||||
| 分配粒度 | 单次 contiguous 一个 session 全部 KV | 简化 slab allocator + 单次 batch transfer |
|
||||
| Layout | K-all-layers concat, then V-all-layers concat | 跟 mooncake 的 batch_transfer 接口对齐 |
|
||||
| Free 策略 | finalize 后立即 free | 当 snapshot 已 ingest 到 kv_pool,snapshot_buf 副本不再需要 |
|
||||
| 满了怎么办 | prepare_receive 返回 ok=False, reason=snapshot-buf-full | 让 caller fall back 到 re-prefill |
|
||||
|
||||
---
|
||||
|
||||
## 4. 接口变化
|
||||
|
||||
### 4.1 SnapshotPrepareReceiveReqOutput
|
||||
|
||||
旧:
|
||||
```
|
||||
k_base_ptrs: List[int] # 各 layer 的 k_buffer.data_ptr()
|
||||
v_base_ptrs: List[int]
|
||||
slot_indices: List[int] # kv_pool 中分配的 slot
|
||||
stride_k_bytes / stride_v_bytes
|
||||
```
|
||||
|
||||
新:
|
||||
```
|
||||
snapshot_buf_base_ptr: int # snapshot_buf.data_ptr()
|
||||
k_layer_offsets: List[int] # 各 layer K 在 snapshot_buf 中的字节偏移
|
||||
v_layer_offsets: List[int] # 各 layer V 偏移
|
||||
num_tokens: int
|
||||
stride_k_bytes / stride_v_bytes
|
||||
slab_handle: int # opaque handle for finalize/abort
|
||||
```
|
||||
|
||||
### 4.2 SnapshotFinalizeIngestReqInput
|
||||
|
||||
旧:
|
||||
```
|
||||
session_id, token_ids, slot_indices
|
||||
```
|
||||
|
||||
新:
|
||||
```
|
||||
session_id, token_ids, slab_handle # P 用 handle 找到 record,再 alloc kv_pool + copy + insert
|
||||
```
|
||||
|
||||
### 4.3 D-side push 逻辑(agentic)
|
||||
|
||||
旧:D 算 src_slot[L] → dst_slot[L] mapping,batch_transfer
|
||||
|
||||
新:D 算 src_slot[L] → snapshot_buf 中的 k_layer_offsets[L] / v_layer_offsets[L] mapping,batch_transfer。完全不需要 dst slot indices。
|
||||
|
||||
---
|
||||
|
||||
## 5. 实施步骤
|
||||
|
||||
| # | 步骤 | LOC 估计 |
|
||||
|---|---|---:|
|
||||
| 1 | `SnapshotBufAllocator` 类(slab/bump allocator) | 80 |
|
||||
| 2 | `SnapshotLinkController.__init__` 加 snapshot_buf 分配 + 注册 | 30 |
|
||||
| 3 | 重写 `prepare_receive`、新加 `_compute_layer_offsets` | 60 |
|
||||
| 4 | 新加 `finalize_with_snapshot_buf` + 删旧的 `finalize_ingest` | 70 |
|
||||
| 5 | 修改 io_struct 字段 + 删旧字段 | 30 |
|
||||
| 6 | 修改 agentic `_attempt_d_to_p_sync` 用新字段 | 40 |
|
||||
| 7 | 改 mem leak check 计入 snapshot_buf | 5 |
|
||||
| 8 | 单元 smoke test | 50 |
|
||||
|
||||
Total: ~365 LOC
|
||||
|
||||
---
|
||||
|
||||
## 6. 风险
|
||||
|
||||
| 风险 | 缓解 |
|
||||
|---|---|
|
||||
| 8 GB GPU mem cost | 用户可配置;mem-fraction-static 已经留了 buffer |
|
||||
| 多 session 抢 snapshot_buf | slab allocator + LRU evict 旧的 snapshot |
|
||||
| GPU→GPU copy 性能 | ~5 GB @ 3 TB/s = 1.7 ms,可忽略 |
|
||||
| 接口大改影响 smoke | 在 commit 内完成所有接口变更,smoke 同步更新 |
|
||||
|
||||
---
|
||||
|
||||
## 7. 验收
|
||||
|
||||
- [ ] `scripts/smoke_snapshot_sglang_integration.py` 跑通新接口(prepare_receive 不再 alloc-failed)
|
||||
- [ ] E4-v6 跑同样 trace,d-to-p-sync.jsonl 出现 OK 事件 ≥ 30%(vs 当前 0%)
|
||||
|
||||
---
|
||||
|
||||
**核心句**:用 GPU 上独立的 snapshot_buf 接收 D 端推送,把"竞争 P kv_pool"这个根本性 alloc 冲突消掉,把 alloc 决策推迟到 finalize 时机,让 D→P 真正有机会跑通。
|
||||
@@ -26,12 +26,15 @@ if [ -z "${CUDA_HOME:-}" ]; then
|
||||
fi
|
||||
|
||||
MODEL=${MODEL:-/mnt/models/Qwen/Qwen3-30B-A3B-Instruct-2507}
|
||||
TRACE=${TRACE:-outputs/inferact_50sess.jsonl}
|
||||
TRACE=${TRACE:-third_party/traces/qwen35-swebench-50sess.jsonl}
|
||||
OUTPUT=${OUTPUT:-outputs/e4p_kvc_v2_d_to_p_sync_pressured_50sess}
|
||||
IB_DEVICE=${IB_DEVICE:-mlx5_60}
|
||||
LOAD_FLOOR_BONUS=${LOAD_FLOOR_BONUS:-200}
|
||||
REJECT_THRESHOLD=${REJECT_THRESHOLD:-1}
|
||||
MEM_FRACTION=${MEM_FRACTION:-0.5}
|
||||
# time-scale: 1 = realistic 5.44h timeline for the SWE-Bench trace;
|
||||
# 10 = compress to ~33 min; 60 = compress to ~5.5 min (stress test).
|
||||
TIME_SCALE=${TIME_SCALE:-1}
|
||||
|
||||
if [ ! -f "$TRACE" ]; then
|
||||
echo "ERROR: trace not found at $TRACE" >&2
|
||||
@@ -51,6 +54,29 @@ log "OUTPUT=$OUTPUT"
|
||||
label=e4p_kvc_v2_d_to_p_sync_run1
|
||||
log "=== [E4p] $label starting ==="
|
||||
|
||||
# Background GPU utilization sampler — every 1 s, all 4 GPUs, CSV output.
|
||||
GPU_CSV="$OUTPUT/gpu_util.csv"
|
||||
log "GPU sampling → $GPU_CSV (1 Hz, gpus 0-3)"
|
||||
echo "timestamp_iso,gpu_index,util_pct,mem_used_MiB,mem_total_MiB,sm_clock_MHz,power_W,temperature_C" > "$GPU_CSV"
|
||||
(
|
||||
while true; do
|
||||
ts_iso=$(date -u +%Y-%m-%dT%H:%M:%S.%3NZ)
|
||||
nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.total,clocks.sm,power.draw,temperature.gpu \
|
||||
--format=csv,noheader,nounits 2>/dev/null \
|
||||
| sed -e "s/^/${ts_iso},/" -e 's/ //g' >> "$GPU_CSV" || true
|
||||
sleep 1
|
||||
done
|
||||
) &
|
||||
GPU_SAMPLER_PID=$!
|
||||
log "GPU sampler pid=$GPU_SAMPLER_PID"
|
||||
|
||||
cleanup_gpu_sampler() {
|
||||
kill -9 "$GPU_SAMPLER_PID" 2>/dev/null || true
|
||||
wait "$GPU_SAMPLER_PID" 2>/dev/null || true
|
||||
log "GPU sampler stopped (output: $GPU_CSV, $(wc -l < "$GPU_CSV") rows)"
|
||||
}
|
||||
trap cleanup_gpu_sampler EXIT INT TERM
|
||||
|
||||
uv run --no-sync python -m agentic_pd_hybrid.cli benchmark-live \
|
||||
--trace "$TRACE" \
|
||||
--output-root "$OUTPUT" \
|
||||
@@ -63,7 +89,7 @@ uv run --no-sync python -m agentic_pd_hybrid.cli benchmark-live \
|
||||
--transfer-backend mooncake \
|
||||
--force-rdma --ib-device "$IB_DEVICE" \
|
||||
--gpu-budget 4 \
|
||||
--time-scale 1 \
|
||||
--time-scale "$TIME_SCALE" \
|
||||
--session-sample-rate 1.0 \
|
||||
--target-duration-s 100000 \
|
||||
--concurrency-limit 32 \
|
||||
@@ -78,6 +104,7 @@ uv run --no-sync python -m agentic_pd_hybrid.cli benchmark-live \
|
||||
--kvcache-direct-max-uncached-tokens 8192 \
|
||||
--kvcache-load-floor-bonus "$LOAD_FLOOR_BONUS" \
|
||||
--decode-mem-fraction-static "${DECODE_MEM_FRAC:-0.4}" \
|
||||
--prefill-mem-fraction-static "${PREFILL_MEM_FRAC:-0.7}" \
|
||||
--enable-d-to-p-sync 2>&1 | tee -a "$LOG"
|
||||
|
||||
run_dir=$(ls -td "$OUTPUT"/kvcache-centric-*/ 2>/dev/null | head -1)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -2356,19 +2353,10 @@ async def _invoke_kvcache_seeded_router(
|
||||
prefill_url=prefill_url,
|
||||
decode_session=decode_session,
|
||||
)
|
||||
if sync_result is not None and sync_result.get("status") != "ok":
|
||||
logger.info(
|
||||
"d_to_p_sync sid=%s rid=%s skipped: %s",
|
||||
request.session_id, request.request_id, sync_result,
|
||||
)
|
||||
elif sync_result and sync_result.get("status") == "ok":
|
||||
logger.info(
|
||||
"d_to_p_sync sid=%s rid=%s pushed=%d ingested_prefix=%d",
|
||||
request.session_id,
|
||||
request.request_id,
|
||||
sync_result.get("bytes_pushed", 0),
|
||||
sync_result.get("inserted_prefix_len", 0),
|
||||
)
|
||||
# NB: every outcome of _attempt_d_to_p_sync is already captured in
|
||||
# structural/d-to-p-sync.jsonl via _structural_emit. No need for an
|
||||
# additional logger.info here (and `logger` isn't imported at module
|
||||
# scope, so it would NameError if reached).
|
||||
|
||||
decode_session_newly_opened = False
|
||||
try:
|
||||
|
||||
1
third_party/agentic-kvcache
vendored
Submodule
1
third_party/agentic-kvcache
vendored
Submodule
Submodule third_party/agentic-kvcache added at 44796a1139
@@ -1,34 +1,29 @@
|
||||
"""SnapshotLinkController — drives D→P RDMA snapshot pushes.
|
||||
"""SnapshotLinkController — D→P RDMA snapshot pushes with dedicated GPU buffer.
|
||||
|
||||
This class owns:
|
||||
* A dedicated ``mooncake.engine.TransferEngine`` (independent of the PD
|
||||
pipeline engine).
|
||||
* Memory-region registrations covering the worker's KV pool layer buffers
|
||||
(registered once at startup for zero per-snapshot overhead).
|
||||
* A side-table mapping ``session_id → SnapshotIngestRecord`` for P-side
|
||||
receivers tracking outstanding ingests until ``snapshot_finalize_ingest``
|
||||
is called.
|
||||
Per `docs/SNAPSHOT_STORE_REFACTOR_ZH.md`, this controller now reserves a
|
||||
dedicated GPU tensor (``snapshot_buf``) for receiving D→P snapshots, instead
|
||||
of competing with the worker's ``token_to_kv_pool_allocator`` at
|
||||
prepare_receive time. The kv_pool alloc is deferred to ``finalize_ingest``
|
||||
when the bytes are already in hand — if that alloc fails we drop the
|
||||
snapshot but RDMA reception itself succeeded.
|
||||
|
||||
Lifecycle:
|
||||
SGLang scheduler instantiates one of these per worker if
|
||||
``SGLANG_SNAPSHOT_LINK_ENABLE=1``. The scheduler is responsible for
|
||||
feeding kv_pool buffer descriptors and the kv_pool_allocator so the
|
||||
controller can pre-register memory and (on P) allocate slots.
|
||||
Layout of the snapshot_buf for one session reception (chosen for
|
||||
mooncake's batch_transfer_sync_write friendliness — every layer maps to
|
||||
a single contiguous slab):
|
||||
|
||||
Direction symmetry:
|
||||
Both D and P workers run identical controllers. Who sends and who
|
||||
receives is determined by who calls ``push_session_kv`` vs
|
||||
``prepare_receive`` / ``finalize_ingest`` — there's no PREFILL/DECODE
|
||||
role baked into the snapshot side.
|
||||
[K_layer_0: num_tokens × stride_k_bytes]
|
||||
[K_layer_1: num_tokens × stride_k_bytes]
|
||||
...
|
||||
[K_layer_L-1]
|
||||
[V_layer_0: num_tokens × stride_v_bytes]
|
||||
...
|
||||
[V_layer_L-1]
|
||||
|
||||
This is the **vendored** copy living alongside SGLang internals so the
|
||||
scheduler can import it directly. ``src/agentic_pd_hybrid/snapshot_link.py``
|
||||
contains the same primitives for stand-alone smoke testing.
|
||||
The buffer is split into multiple such slabs via ``SnapshotBufAllocator``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ctypes
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
@@ -44,6 +39,10 @@ SNAPSHOT_LINK_HOST_ENV = "SGLANG_SNAPSHOT_LINK_HOST"
|
||||
SNAPSHOT_LINK_PORT_ENV = "SGLANG_SNAPSHOT_LINK_PORT"
|
||||
SNAPSHOT_LINK_IB_DEVICE_ENV = "SGLANG_SNAPSHOT_LINK_IB_DEVICE"
|
||||
|
||||
# Default snapshot_buf size: 8 GB. Enough for ~1.5 Qwen3-30B 50k-token sessions.
|
||||
SNAPSHOT_BUF_BYTES_ENV = "SGLANG_SNAPSHOT_LINK_BUF_BYTES"
|
||||
DEFAULT_SNAPSHOT_BUF_BYTES = 8 * 1024 * 1024 * 1024
|
||||
|
||||
|
||||
@dataclass
|
||||
class _LayerBufferDesc:
|
||||
@@ -56,13 +55,75 @@ class _LayerBufferDesc:
|
||||
|
||||
@dataclass
|
||||
class SnapshotIngestRecord:
|
||||
"""P-side: bookkeeping for an outstanding incoming snapshot."""
|
||||
"""P-side bookkeeping for one in-flight snapshot reception."""
|
||||
session_id: str
|
||||
slot_indices: List[int] # kv_pool slots reserved for this ingest
|
||||
slab_offset: int # offset within snapshot_buf
|
||||
slab_size: int # total bytes for this slab
|
||||
num_tokens: int
|
||||
k_layer_offsets: List[int] # absolute byte offsets of K layers in snapshot_buf
|
||||
v_layer_offsets: List[int]
|
||||
per_token_k_bytes: int
|
||||
per_token_v_bytes: int
|
||||
created_at: float = field(default_factory=time.time)
|
||||
|
||||
|
||||
class SnapshotBufAllocator:
|
||||
"""First-fit free-list allocator over a single contiguous byte range.
|
||||
|
||||
Tracks gaps in a sorted list. Merges adjacent free regions on free().
|
||||
"""
|
||||
|
||||
def __init__(self, capacity_bytes: int):
|
||||
self.capacity = capacity_bytes
|
||||
# Free regions sorted by offset: [(offset, size), ...]
|
||||
self._free: List[Tuple[int, int]] = [(0, capacity_bytes)]
|
||||
self._lock = threading.Lock()
|
||||
self._inflight: dict[int, int] = {} # offset → size for sanity check
|
||||
|
||||
def alloc(self, size: int) -> Optional[int]:
|
||||
"""Return offset of allocated region, or None if no fit available."""
|
||||
if size <= 0:
|
||||
return None
|
||||
# Page-align allocations to 4 KB for RDMA-friendly alignment.
|
||||
size = (size + 4095) & ~4095
|
||||
with self._lock:
|
||||
for i, (off, sz) in enumerate(self._free):
|
||||
if sz >= size:
|
||||
if sz == size:
|
||||
self._free.pop(i)
|
||||
else:
|
||||
self._free[i] = (off + size, sz - size)
|
||||
self._inflight[off] = size
|
||||
return off
|
||||
return None
|
||||
|
||||
def free(self, offset: int) -> bool:
|
||||
"""Return True if the offset was successfully freed."""
|
||||
with self._lock:
|
||||
size = self._inflight.pop(offset, None)
|
||||
if size is None:
|
||||
return False
|
||||
# Insert sorted and merge adjacents
|
||||
self._free.append((offset, size))
|
||||
self._free.sort()
|
||||
merged: List[Tuple[int, int]] = []
|
||||
for off, sz in self._free:
|
||||
if merged and merged[-1][0] + merged[-1][1] == off:
|
||||
merged[-1] = (merged[-1][0], merged[-1][1] + sz)
|
||||
else:
|
||||
merged.append((off, sz))
|
||||
self._free = merged
|
||||
return True
|
||||
|
||||
def available_bytes(self) -> int:
|
||||
with self._lock:
|
||||
return sum(sz for _, sz in self._free)
|
||||
|
||||
def in_use_bytes(self) -> int:
|
||||
with self._lock:
|
||||
return sum(self._inflight.values())
|
||||
|
||||
|
||||
def _import_transfer_engine():
|
||||
try:
|
||||
from mooncake.engine import TransferEngine
|
||||
@@ -75,7 +136,13 @@ def _import_transfer_engine():
|
||||
|
||||
|
||||
class SnapshotLinkController:
|
||||
"""Mooncake engine + per-layer mem registrations + ingest bookkeeping."""
|
||||
"""Owns mooncake engine + kv_pool registrations + snapshot_buf + records.
|
||||
|
||||
D-side use: push session KV via ``push_session_to_snapshot_buf``.
|
||||
P-side use: ``prepare_receive`` → caller pushes via RDMA →
|
||||
``ingest_snapshot_into_kvpool`` (does GPU memcpy +
|
||||
radix insert) → ``finalize_record`` (frees the slab).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -84,24 +151,16 @@ class SnapshotLinkController:
|
||||
ib_device: Optional[str],
|
||||
kv_pool_layer_buffers: List[Tuple[int, int, int, bool]],
|
||||
token_to_kv_pool_allocator,
|
||||
tree_cache=None,
|
||||
protocol: Optional[str] = None,
|
||||
snapshot_buf_bytes: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
host, port : where this worker binds its snapshot engine.
|
||||
ib_device : preferred IB HCA (e.g. "mlx5_60").
|
||||
kv_pool_layer_buffers : list of ``(base_ptr, bytes_per_token,
|
||||
capacity_bytes, is_k)`` tuples. Order should be K-layers
|
||||
first then V-layers (matches MHATokenToKVPool layout).
|
||||
token_to_kv_pool_allocator : the worker's allocator. We use
|
||||
``.alloc(N)`` on the P side to reserve receive slots.
|
||||
"""
|
||||
TransferEngine = _import_transfer_engine()
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.ib_device = ib_device
|
||||
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
|
||||
self.tree_cache = tree_cache
|
||||
self.layer_buffers: List[_LayerBufferDesc] = [
|
||||
_LayerBufferDesc(
|
||||
base_ptr=base, bytes_per_token=btok,
|
||||
@@ -121,38 +180,67 @@ class SnapshotLinkController:
|
||||
)
|
||||
self._session_id = f"{host}:{self.engine.get_rpc_port()}"
|
||||
|
||||
# Register all layer buffers up-front (one-shot)
|
||||
# Register existing kv_pool layer buffers (needed for D-side send and
|
||||
# for P-side ingest copy source = snapshot_buf, destination = kv_pool)
|
||||
ptrs = [d.base_ptr for d in self.layer_buffers]
|
||||
lens = [d.capacity_bytes for d in self.layer_buffers]
|
||||
try:
|
||||
reg_ret = self.engine.batch_register_memory(ptrs, lens)
|
||||
except Exception:
|
||||
reg_ret = -1
|
||||
# Fall back to individual register_memory calls
|
||||
reg_ret = 0
|
||||
for ptr, length in zip(ptrs, lens):
|
||||
r = self.engine.register_memory(ptr, length)
|
||||
if r != 0:
|
||||
logger.warning(
|
||||
"SnapshotLinkController register_memory(%s, %d) returned %d",
|
||||
hex(ptr), length, r,
|
||||
)
|
||||
reg_ret = r
|
||||
if reg_ret != 0:
|
||||
logger.warning(
|
||||
"SnapshotLinkController batch_register_memory returned %d "
|
||||
"(continuing — individual registrations may have succeeded)",
|
||||
reg_ret,
|
||||
"SnapshotLinkController kv_pool batch_register returned %d", reg_ret
|
||||
)
|
||||
|
||||
# Allocate + register the dedicated snapshot reception buffer (P-side)
|
||||
# This decouples reception from kv_pool, avoiding the alloc-failed
|
||||
# death loop that killed E4-v4/v5.
|
||||
import torch
|
||||
|
||||
if snapshot_buf_bytes is None:
|
||||
snapshot_buf_bytes = int(
|
||||
os.environ.get(SNAPSHOT_BUF_BYTES_ENV, DEFAULT_SNAPSHOT_BUF_BYTES)
|
||||
)
|
||||
device = self._allocator_device()
|
||||
try:
|
||||
self.snapshot_buf = torch.zeros(
|
||||
snapshot_buf_bytes, dtype=torch.uint8, device=device,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
logger.warning(
|
||||
"Could not allocate snapshot_buf of %d bytes on %s: %s. "
|
||||
"Falling back to 1 GB.", snapshot_buf_bytes, device, e,
|
||||
)
|
||||
snapshot_buf_bytes = 1024 * 1024 * 1024
|
||||
self.snapshot_buf = torch.zeros(
|
||||
snapshot_buf_bytes, dtype=torch.uint8, device=device,
|
||||
)
|
||||
self._snapshot_buf_bytes = snapshot_buf_bytes
|
||||
self._snapshot_buf_ptr = self.snapshot_buf.data_ptr()
|
||||
ret = self.engine.register_memory(self._snapshot_buf_ptr, snapshot_buf_bytes)
|
||||
if ret != 0:
|
||||
logger.warning(
|
||||
"SnapshotLinkController snapshot_buf register_memory(%s, %d) ret=%d",
|
||||
hex(self._snapshot_buf_ptr), snapshot_buf_bytes, ret,
|
||||
)
|
||||
self.snapshot_buf_alloc = SnapshotBufAllocator(snapshot_buf_bytes)
|
||||
|
||||
# Receive-side bookkeeping
|
||||
self._ingest_records: dict[str, SnapshotIngestRecord] = {}
|
||||
self._records_by_handle: dict[int, SnapshotIngestRecord] = {}
|
||||
self._next_handle = 1
|
||||
self._lock = threading.Lock()
|
||||
|
||||
logger.info(
|
||||
"SnapshotLinkController up at %s (snapshot_session_id=%s, "
|
||||
"%d layer buffers registered, ib=%s)",
|
||||
listen, self._session_id, len(self.layer_buffers), ib_device,
|
||||
"SnapshotLinkController up at %s (sid=%s, %d kv layer bufs, "
|
||||
"snapshot_buf=%.1f GB on %s)",
|
||||
listen, self._session_id, len(self.layer_buffers),
|
||||
snapshot_buf_bytes / 1e9, device,
|
||||
)
|
||||
|
||||
# ----- accessors ----------------------------------------------------
|
||||
@@ -161,9 +249,16 @@ class SnapshotLinkController:
|
||||
def snapshot_session_id(self) -> str:
|
||||
return self._session_id
|
||||
|
||||
@property
|
||||
def snapshot_buf_ptr(self) -> int:
|
||||
return self._snapshot_buf_ptr
|
||||
|
||||
@property
|
||||
def snapshot_buf_bytes(self) -> int:
|
||||
return self._snapshot_buf_bytes
|
||||
|
||||
@property
|
||||
def layer_num(self) -> int:
|
||||
"""Number of K layers (= number of V layers = total / 2)."""
|
||||
return len(self.layer_buffers) // 2
|
||||
|
||||
def get_k_base_ptrs(self) -> List[int]:
|
||||
@@ -184,66 +279,6 @@ class SnapshotLinkController:
|
||||
return d.bytes_per_token
|
||||
return 0
|
||||
|
||||
# ----- P-side: prepare to receive ----------------------------------
|
||||
|
||||
def prepare_receive(self, session_id: str, num_tokens: int) -> Optional[SnapshotIngestRecord]:
|
||||
"""Allocate ``num_tokens`` slots in kv_pool for an incoming snapshot.
|
||||
|
||||
Returns a record with the slot indices, or ``None`` if no capacity.
|
||||
Mooncake registration is already in place (whole-buffer registration
|
||||
at startup), so no per-snapshot register is needed.
|
||||
"""
|
||||
try:
|
||||
indices_tensor = self.token_to_kv_pool_allocator.alloc(num_tokens)
|
||||
except Exception as e:
|
||||
logger.exception("SnapshotLinkController.prepare_receive alloc failed: %s", e)
|
||||
return None
|
||||
if indices_tensor is None:
|
||||
return None
|
||||
try:
|
||||
slot_indices = [int(x) for x in indices_tensor.tolist()]
|
||||
except Exception:
|
||||
# If allocator returns a python list directly
|
||||
slot_indices = list(map(int, indices_tensor))
|
||||
record = SnapshotIngestRecord(
|
||||
session_id=session_id,
|
||||
slot_indices=slot_indices,
|
||||
num_tokens=num_tokens,
|
||||
)
|
||||
with self._lock:
|
||||
old = self._ingest_records.pop(session_id, None)
|
||||
if old is not None:
|
||||
# Best-effort: free old slots if we're overwriting
|
||||
try:
|
||||
self._free_slots(old.slot_indices)
|
||||
except Exception:
|
||||
pass
|
||||
self._ingest_records[session_id] = record
|
||||
return record
|
||||
|
||||
def take_record(self, session_id: str) -> Optional[SnapshotIngestRecord]:
|
||||
with self._lock:
|
||||
return self._ingest_records.pop(session_id, None)
|
||||
|
||||
def discard_record(self, session_id: str) -> None:
|
||||
"""Drop a pending ingest (e.g. on timeout or D-side failure)."""
|
||||
with self._lock:
|
||||
rec = self._ingest_records.pop(session_id, None)
|
||||
if rec is not None:
|
||||
try:
|
||||
self._free_slots(rec.slot_indices)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _free_slots(self, slot_indices: List[int]) -> None:
|
||||
import torch
|
||||
t = torch.tensor(slot_indices, dtype=torch.int64,
|
||||
device=self._allocator_device())
|
||||
try:
|
||||
self.token_to_kv_pool_allocator.free(t)
|
||||
except Exception as e:
|
||||
logger.warning("SnapshotLinkController._free_slots failed: %s", e)
|
||||
|
||||
def _allocator_device(self):
|
||||
# Best-effort: pull device from one of the buffer tensors via the allocator
|
||||
try:
|
||||
@@ -251,91 +286,291 @@ class SnapshotLinkController:
|
||||
except AttributeError:
|
||||
return "cuda"
|
||||
|
||||
# ----- D-side: push session KV --------------------------------------
|
||||
# ----- P-side: prepare to receive ----------------------------------
|
||||
|
||||
def push_session_kv(
|
||||
def prepare_receive(self, session_id: str, num_tokens: int) -> Optional[SnapshotIngestRecord]:
|
||||
"""Carve a slab out of snapshot_buf large enough for num_tokens of K+V.
|
||||
|
||||
Returns the record describing the slab layout, or None if snapshot_buf
|
||||
is full. This does NOT touch kv_pool — alloc happens at ingest time.
|
||||
"""
|
||||
if num_tokens <= 0:
|
||||
return None
|
||||
stride_k = self.get_stride_k_bytes()
|
||||
stride_v = self.get_stride_v_bytes()
|
||||
L = self.layer_num
|
||||
slab_bytes = L * num_tokens * stride_k + L * num_tokens * stride_v
|
||||
offset = self.snapshot_buf_alloc.alloc(slab_bytes)
|
||||
if offset is None:
|
||||
logger.info(
|
||||
"prepare_receive: snapshot_buf full (sid=%s n=%d need=%d B available=%d B)",
|
||||
session_id, num_tokens, slab_bytes,
|
||||
self.snapshot_buf_alloc.available_bytes(),
|
||||
)
|
||||
return None
|
||||
# Layout: K0..KL-1, then V0..VL-1
|
||||
k_offs = [offset + i * num_tokens * stride_k for i in range(L)]
|
||||
v_offs = [offset + L * num_tokens * stride_k + i * num_tokens * stride_v
|
||||
for i in range(L)]
|
||||
record = SnapshotIngestRecord(
|
||||
session_id=session_id,
|
||||
slab_offset=offset,
|
||||
slab_size=slab_bytes,
|
||||
num_tokens=num_tokens,
|
||||
k_layer_offsets=k_offs,
|
||||
v_layer_offsets=v_offs,
|
||||
per_token_k_bytes=stride_k,
|
||||
per_token_v_bytes=stride_v,
|
||||
)
|
||||
with self._lock:
|
||||
# Evict prior record for the same session (best-effort)
|
||||
old = self._ingest_records.pop(session_id, None)
|
||||
if old is not None:
|
||||
self.snapshot_buf_alloc.free(old.slab_offset)
|
||||
self._records_by_handle.pop(id(old), None)
|
||||
self._ingest_records[session_id] = record
|
||||
self._records_by_handle[id(record)] = record
|
||||
return record
|
||||
|
||||
def lookup_by_handle(self, handle: int) -> Optional[SnapshotIngestRecord]:
|
||||
with self._lock:
|
||||
return self._records_by_handle.get(handle)
|
||||
|
||||
def discard_record(self, session_id: str) -> None:
|
||||
with self._lock:
|
||||
rec = self._ingest_records.pop(session_id, None)
|
||||
if rec is not None:
|
||||
self.snapshot_buf_alloc.free(rec.slab_offset)
|
||||
with self._lock:
|
||||
self._records_by_handle.pop(id(rec), None)
|
||||
|
||||
def total_pending_snapshot_bytes(self) -> int:
|
||||
with self._lock:
|
||||
return sum(rec.slab_size for rec in self._ingest_records.values())
|
||||
|
||||
# ----- P-side: ingest snapshot into kv_pool + radix tree -----------
|
||||
|
||||
def ingest_snapshot_into_kvpool(
|
||||
self,
|
||||
session_id: str,
|
||||
token_ids: List[int],
|
||||
) -> Tuple[bool, str, int]:
|
||||
"""Copy snapshot_buf bytes into kv_pool slots and insert into radix.
|
||||
|
||||
Returns (ok, reason, inserted_prefix_len).
|
||||
"""
|
||||
with self._lock:
|
||||
record = self._ingest_records.pop(session_id, None)
|
||||
if record is not None:
|
||||
self._records_by_handle.pop(id(record), None)
|
||||
if record is None:
|
||||
return False, "no-pending-ingest", 0
|
||||
|
||||
try:
|
||||
n = min(len(token_ids), record.num_tokens)
|
||||
if n == 0:
|
||||
self.snapshot_buf_alloc.free(record.slab_offset)
|
||||
return False, "empty-token-ids", 0
|
||||
|
||||
# Alloc kv_pool slots NOW that the snapshot bytes are in hand.
|
||||
try:
|
||||
indices_tensor = self.token_to_kv_pool_allocator.alloc(n)
|
||||
except Exception as exc:
|
||||
self.snapshot_buf_alloc.free(record.slab_offset)
|
||||
return False, f"kvpool-alloc-threw:{exc!r}", 0
|
||||
if indices_tensor is None:
|
||||
self.snapshot_buf_alloc.free(record.slab_offset)
|
||||
return False, "kvpool-alloc-failed-at-ingest", 0
|
||||
|
||||
# GPU→GPU copy from snapshot_buf into kv_pool layer buffers
|
||||
try:
|
||||
self._copy_snapshot_to_kvpool(record, indices_tensor)
|
||||
except Exception as exc:
|
||||
logger.exception("snapshot→kvpool copy failed: %s", exc)
|
||||
# Free both allocations
|
||||
self._free_slot_indices(indices_tensor)
|
||||
self.snapshot_buf_alloc.free(record.slab_offset)
|
||||
return False, f"copy-failed:{exc!r}", 0
|
||||
|
||||
# Insert into radix tree
|
||||
try:
|
||||
inserted_prefix_len = self._radix_insert(token_ids[:n], indices_tensor)
|
||||
except Exception as exc:
|
||||
logger.exception("radix insert failed: %s", exc)
|
||||
self._free_slot_indices(indices_tensor)
|
||||
self.snapshot_buf_alloc.free(record.slab_offset)
|
||||
return False, f"radix-insert-failed:{exc!r}", 0
|
||||
|
||||
# Snapshot is now persisted into kv_pool + radix; the slab is no
|
||||
# longer needed.
|
||||
self.snapshot_buf_alloc.free(record.slab_offset)
|
||||
return True, "ok", int(inserted_prefix_len)
|
||||
except Exception as exc:
|
||||
# Belt-and-braces cleanup
|
||||
try:
|
||||
self.snapshot_buf_alloc.free(record.slab_offset)
|
||||
except Exception:
|
||||
pass
|
||||
return False, f"unexpected:{exc!r}", 0
|
||||
|
||||
def _copy_snapshot_to_kvpool(
|
||||
self,
|
||||
record: SnapshotIngestRecord,
|
||||
slot_indices_tensor,
|
||||
) -> None:
|
||||
"""For each layer L: copy snapshot_buf[K_off[L]..] → k_buffer[L][slots]."""
|
||||
import torch
|
||||
|
||||
n = record.num_tokens
|
||||
stride_k = record.per_token_k_bytes
|
||||
stride_v = record.per_token_v_bytes
|
||||
# View snapshot_buf as a 1-D byte tensor; slice by offsets.
|
||||
for L in range(self.layer_num):
|
||||
# K
|
||||
k_slab_start = record.k_layer_offsets[L] - record.slab_offset + record.slab_offset
|
||||
# NOTE: above is equivalent to record.k_layer_offsets[L] but kept for clarity
|
||||
k_slab_start = record.k_layer_offsets[L]
|
||||
k_layer_bytes = self.snapshot_buf[
|
||||
k_slab_start : k_slab_start + n * stride_k
|
||||
].view(n, stride_k)
|
||||
# Compute destination tensor on kv_pool: dst[slot_indices] = src
|
||||
# We need access to the actual k_buffer[L] tensor. The controller
|
||||
# only has the raw ptr — so we materialize a view via from_blob-ish
|
||||
# trick. Easier: get the tensor from token_to_kv_pool_allocator's kvcache.
|
||||
kv_cache = self.token_to_kv_pool_allocator.get_kvcache()
|
||||
k_buf = kv_cache.k_buffer[L] # (max_tokens, head, dim)
|
||||
# Flatten per-token to bytes
|
||||
flat = k_buf.view(k_buf.shape[0], -1)
|
||||
assert flat.shape[1] * flat.element_size() >= stride_k, (
|
||||
f"K layer {L} stride mismatch: pool {flat.shape[1] * flat.element_size()} vs snapshot {stride_k}"
|
||||
)
|
||||
# Copy: dst[slot_indices] ← src[:n]
|
||||
src_reshape = k_layer_bytes.view(n, flat.shape[1] * flat.element_size())
|
||||
# Byte-level view of destination rows
|
||||
dst_view = flat.view(torch.uint8)
|
||||
dst_view[slot_indices_tensor] = src_reshape
|
||||
|
||||
# V
|
||||
v_slab_start = record.v_layer_offsets[L]
|
||||
v_layer_bytes = self.snapshot_buf[
|
||||
v_slab_start : v_slab_start + n * stride_v
|
||||
]
|
||||
v_buf = kv_cache.v_buffer[L]
|
||||
v_flat = v_buf.view(v_buf.shape[0], -1)
|
||||
src_v = v_layer_bytes.view(n, v_flat.shape[1] * v_flat.element_size())
|
||||
v_dst_view = v_flat.view(torch.uint8)
|
||||
v_dst_view[slot_indices_tensor] = src_v
|
||||
|
||||
def _radix_insert(self, token_ids: List[int], indices_tensor) -> int:
|
||||
"""Insert (token_ids, kv_indices) into the underlying radix tree."""
|
||||
from sglang.srt.mem_cache.base_prefix_cache import InsertParams
|
||||
from sglang.srt.mem_cache.radix_cache import RadixKey
|
||||
from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
|
||||
|
||||
inner = self.tree_cache
|
||||
if isinstance(inner, SessionAwareCache):
|
||||
inner = inner.inner
|
||||
if inner is None:
|
||||
raise RuntimeError("tree_cache not provided to SnapshotLinkController")
|
||||
radix_key = RadixKey(token_ids, None)
|
||||
result = inner.insert(InsertParams(key=radix_key, value=indices_tensor))
|
||||
return int(getattr(result, "prefix_len", 0))
|
||||
|
||||
def _free_slot_indices(self, indices_tensor) -> None:
|
||||
try:
|
||||
self.token_to_kv_pool_allocator.free(indices_tensor)
|
||||
except Exception as e:
|
||||
logger.warning("_free_slot_indices failed: %s", e)
|
||||
|
||||
# ----- D-side: push session KV to a peer's snapshot_buf ------------
|
||||
|
||||
def push_session_to_snapshot_buf(
|
||||
self,
|
||||
*,
|
||||
target_snapshot_session_id: str,
|
||||
src_slot_indices: List[int],
|
||||
target_k_base_ptrs: List[int],
|
||||
target_v_base_ptrs: List[int],
|
||||
target_slot_indices: List[int],
|
||||
target_stride_k_bytes: int,
|
||||
target_stride_v_bytes: int,
|
||||
target_snapshot_buf_base: int,
|
||||
target_k_layer_offsets: List[int],
|
||||
target_v_layer_offsets: List[int],
|
||||
target_per_token_k_bytes: int,
|
||||
target_per_token_v_bytes: int,
|
||||
) -> Tuple[int, int]:
|
||||
"""Push the KV bytes at src_slot_indices to the remote slots.
|
||||
"""Push session KV from local kv_pool into a peer's snapshot_buf slab.
|
||||
|
||||
Returns ``(mooncake_return_code, bytes_pushed)``. Caller is
|
||||
responsible for any post-push handshake (e.g. finalize_ingest RPC).
|
||||
For each layer: gather src ranges (possibly scattered slot indices)
|
||||
and write to a contiguous range in the peer's snapshot_buf.
|
||||
Returns (mooncake_return_code, bytes_pushed).
|
||||
"""
|
||||
if not src_slot_indices:
|
||||
return 0, 0
|
||||
layer_num = self.layer_num
|
||||
k_src_bases = self.get_k_base_ptrs()
|
||||
v_src_bases = self.get_v_base_ptrs()
|
||||
stride_k = self.get_stride_k_bytes()
|
||||
stride_v = self.get_stride_v_bytes()
|
||||
if len(target_k_base_ptrs) != layer_num or len(target_v_base_ptrs) != layer_num:
|
||||
if (len(target_k_layer_offsets) != layer_num
|
||||
or len(target_v_layer_offsets) != layer_num):
|
||||
raise ValueError(
|
||||
f"target K/V base ptr count {len(target_k_base_ptrs)}/{len(target_v_base_ptrs)} "
|
||||
f"!= local layer_num {layer_num}"
|
||||
f"target K/V layer offset count {len(target_k_layer_offsets)}/"
|
||||
f"{len(target_v_layer_offsets)} != local layer_num {layer_num}"
|
||||
)
|
||||
if stride_k != target_stride_k_bytes or stride_v != target_stride_v_bytes:
|
||||
if (stride_k != target_per_token_k_bytes
|
||||
or stride_v != target_per_token_v_bytes):
|
||||
raise ValueError(
|
||||
f"stride mismatch: local k={stride_k}, v={stride_v}; "
|
||||
f"target k={target_stride_k_bytes}, v={target_stride_v_bytes}"
|
||||
)
|
||||
if len(src_slot_indices) != len(target_slot_indices):
|
||||
raise ValueError(
|
||||
f"slot count mismatch: src={len(src_slot_indices)}, "
|
||||
f"target={len(target_slot_indices)}"
|
||||
f"stride mismatch: local k={stride_k}/v={stride_v}, "
|
||||
f"target k={target_per_token_k_bytes}/v={target_per_token_v_bytes}"
|
||||
)
|
||||
n = len(src_slot_indices)
|
||||
|
||||
local_addrs: List[int] = []
|
||||
remote_addrs: List[int] = []
|
||||
lengths: List[int] = []
|
||||
|
||||
# Group contiguous runs on the target side to coalesce ops.
|
||||
# Simple approach: per (layer, K/V) pair, walk src/target index
|
||||
# tuples; merge runs where both src and target are sequential.
|
||||
for layer_id in range(layer_num):
|
||||
for kv_bases, stride, kv_label in (
|
||||
(k_src_bases[layer_id], stride_k, "K"),
|
||||
(v_src_bases[layer_id], stride_v, "V"),
|
||||
):
|
||||
src_base = kv_bases
|
||||
if kv_label == "K":
|
||||
tgt_base = target_k_base_ptrs[layer_id]
|
||||
# Coalesce contiguous src runs.
|
||||
# Inner-loop helper to walk indices and emit run boundaries.
|
||||
def _emit_runs(src_base: int, tgt_base: int, stride: int) -> None:
|
||||
run_src_start = run_tgt_start = run_len = None
|
||||
for tgt_idx, src in enumerate(src_slot_indices):
|
||||
if run_src_start is None:
|
||||
run_src_start, run_tgt_start, run_len = src, tgt_idx, 1
|
||||
elif src == run_src_start + run_len:
|
||||
run_len += 1
|
||||
else:
|
||||
tgt_base = target_v_base_ptrs[layer_id]
|
||||
run_src_start = run_tgt_start = run_len = None
|
||||
for s, t in zip(src_slot_indices, target_slot_indices):
|
||||
if run_src_start is None:
|
||||
run_src_start, run_tgt_start, run_len = s, t, 1
|
||||
elif s == run_src_start + run_len and t == run_tgt_start + run_len:
|
||||
run_len += 1
|
||||
else:
|
||||
local_addrs.append(src_base + run_src_start * stride)
|
||||
remote_addrs.append(tgt_base + run_tgt_start * stride)
|
||||
lengths.append(run_len * stride)
|
||||
run_src_start, run_tgt_start, run_len = s, t, 1
|
||||
if run_src_start is not None:
|
||||
local_addrs.append(src_base + run_src_start * stride)
|
||||
remote_addrs.append(tgt_base + run_tgt_start * stride)
|
||||
lengths.append(run_len * stride)
|
||||
run_src_start, run_tgt_start, run_len = src, tgt_idx, 1
|
||||
if run_src_start is not None:
|
||||
local_addrs.append(src_base + run_src_start * stride)
|
||||
remote_addrs.append(tgt_base + run_tgt_start * stride)
|
||||
lengths.append(run_len * stride)
|
||||
|
||||
for L in range(layer_num):
|
||||
_emit_runs(
|
||||
k_src_bases[L],
|
||||
target_snapshot_buf_base + target_k_layer_offsets[L],
|
||||
stride_k,
|
||||
)
|
||||
_emit_runs(
|
||||
v_src_bases[L],
|
||||
target_snapshot_buf_base + target_v_layer_offsets[L],
|
||||
stride_v,
|
||||
)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
ret = self.engine.batch_transfer_sync_write(
|
||||
target_snapshot_session_id, local_addrs, remote_addrs, lengths
|
||||
target_snapshot_session_id, local_addrs, remote_addrs, lengths,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("SnapshotLinkController.push_session_kv threw: %s", e)
|
||||
logger.exception(
|
||||
"SnapshotLinkController.push_session_to_snapshot_buf threw: %s", e
|
||||
)
|
||||
return -1, 0
|
||||
t1 = time.perf_counter()
|
||||
bytes_pushed = sum(lengths)
|
||||
logger.info(
|
||||
"SnapshotLinkController.push_session_kv → %s: %d ops, %d B, "
|
||||
"ret=%d, %.2f ms",
|
||||
"push_session_to_snapshot_buf → %s: %d ops, %d B, ret=%d, %.2f ms",
|
||||
target_snapshot_session_id, len(lengths), bytes_pushed, ret,
|
||||
(t1 - t0) * 1000.0,
|
||||
)
|
||||
|
||||
@@ -1662,38 +1662,38 @@ class SnapshotPrepareReceiveReqInput(BaseReq):
|
||||
|
||||
@dataclass
|
||||
class SnapshotPrepareReceiveReqOutput(BaseReq):
|
||||
"""P-side response. New schema points D at P's dedicated snapshot_buf."""
|
||||
|
||||
ok: bool
|
||||
reason: Optional[str] = None
|
||||
# Layout the D side needs to address P's kv_pool slots:
|
||||
# k_base_ptrs[L] = base device address of layer L's K buffer
|
||||
# v_base_ptrs[L] = base device address of layer L's V buffer
|
||||
# slot_indices = the contiguous range P allocated (list[int], length=num_tokens)
|
||||
# stride_k_bytes = bytes per token K = head_num * head_dim * dtype.itemsize
|
||||
# stride_v_bytes = bytes per token V (often equals stride_k_bytes)
|
||||
# P also registers these slot regions with mooncake before returning.
|
||||
k_base_ptrs: List[int] = field(default_factory=list)
|
||||
v_base_ptrs: List[int] = field(default_factory=list)
|
||||
slot_indices: List[int] = field(default_factory=list)
|
||||
# P's mooncake snapshot session id (host:rpc_port) for D's batch write target
|
||||
snapshot_session_id: str = ""
|
||||
# snapshot_buf base pointer + per-layer offsets, replacing the old
|
||||
# kv_pool slot_indices scheme that competed with P's prefill work and
|
||||
# always hit alloc-failed. See docs/SNAPSHOT_STORE_REFACTOR_ZH.md.
|
||||
snapshot_buf_base_ptr: int = 0
|
||||
snapshot_buf_capacity_bytes: int = 0
|
||||
k_layer_offsets: List[int] = field(default_factory=list) # bytes within snapshot_buf
|
||||
v_layer_offsets: List[int] = field(default_factory=list)
|
||||
num_tokens: int = 0
|
||||
stride_k_bytes: int = 0
|
||||
stride_v_bytes: int = 0
|
||||
layer_num: int = 0
|
||||
# P's mooncake snapshot session id (host:rpc_port) for D's batch write target
|
||||
snapshot_session_id: str = ""
|
||||
available_tokens: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class SnapshotDumpReqInput(BaseReq):
|
||||
"""D-side: dump session KV via snapshot_link to a target P endpoint."""
|
||||
"""D-side: dump session KV via snapshot_link into P's snapshot_buf slab."""
|
||||
|
||||
session_id: str
|
||||
target_snapshot_session_id: str # P's mooncake snapshot session id
|
||||
target_k_base_ptrs: List[int] = field(default_factory=list)
|
||||
target_v_base_ptrs: List[int] = field(default_factory=list)
|
||||
target_slot_indices: List[int] = field(default_factory=list)
|
||||
target_snapshot_session_id: str
|
||||
target_snapshot_buf_base: int = 0
|
||||
target_k_layer_offsets: List[int] = field(default_factory=list)
|
||||
target_v_layer_offsets: List[int] = field(default_factory=list)
|
||||
target_stride_k_bytes: int = 0
|
||||
target_stride_v_bytes: int = 0
|
||||
ib_device: Optional[str] = None # for the D-side SnapshotPeer initialization
|
||||
ib_device: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -1709,11 +1709,10 @@ class SnapshotDumpReqOutput(BaseReq):
|
||||
|
||||
@dataclass
|
||||
class SnapshotFinalizeIngestReqInput(BaseReq):
|
||||
"""P-side: insert (token_ids, kv_indices) into radix tree after D's push."""
|
||||
"""P-side: copy snapshot_buf slab into kv_pool + insert into radix tree."""
|
||||
|
||||
session_id: str
|
||||
token_ids: List[int]
|
||||
slot_indices: List[int]
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -902,6 +902,7 @@ class Scheduler(
|
||||
ib_device=ib,
|
||||
kv_pool_layer_buffers=layer_buffers,
|
||||
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
|
||||
tree_cache=self.tree_cache,
|
||||
)
|
||||
logger.info(
|
||||
"Snapshot link controller initialized: %s, sid=%s, %d layer bufs",
|
||||
@@ -3750,7 +3751,12 @@ class Scheduler(
|
||||
def snapshot_prepare_receive(
|
||||
self, recv_req: SnapshotPrepareReceiveReqInput
|
||||
) -> SnapshotPrepareReceiveReqOutput:
|
||||
"""P-side: alloc kv_pool slots, return slot/buffer layout for D's batch_push."""
|
||||
"""P-side: carve snapshot_buf slab + return its layout to caller.
|
||||
|
||||
Refactored per docs/SNAPSHOT_STORE_REFACTOR_ZH.md: this no longer
|
||||
touches the kv_pool allocator. The slab is in a dedicated
|
||||
snapshot_buf so prepare can never lose to P's prefill work.
|
||||
"""
|
||||
ctrl = self.snapshot_link_controller
|
||||
if ctrl is None:
|
||||
return SnapshotPrepareReceiveReqOutput(
|
||||
@@ -3765,25 +3771,27 @@ class Scheduler(
|
||||
record = ctrl.prepare_receive(recv_req.session_id, recv_req.num_tokens)
|
||||
if record is None:
|
||||
return SnapshotPrepareReceiveReqOutput(
|
||||
ok=False, reason="alloc-failed",
|
||||
ok=False, reason="snapshot-buf-full",
|
||||
available_tokens=available,
|
||||
)
|
||||
return SnapshotPrepareReceiveReqOutput(
|
||||
ok=True,
|
||||
k_base_ptrs=ctrl.get_k_base_ptrs(),
|
||||
v_base_ptrs=ctrl.get_v_base_ptrs(),
|
||||
slot_indices=record.slot_indices,
|
||||
stride_k_bytes=ctrl.get_stride_k_bytes(),
|
||||
stride_v_bytes=ctrl.get_stride_v_bytes(),
|
||||
layer_num=ctrl.layer_num,
|
||||
snapshot_session_id=ctrl.snapshot_session_id,
|
||||
snapshot_buf_base_ptr=ctrl.snapshot_buf_ptr,
|
||||
snapshot_buf_capacity_bytes=ctrl.snapshot_buf_bytes,
|
||||
k_layer_offsets=record.k_layer_offsets,
|
||||
v_layer_offsets=record.v_layer_offsets,
|
||||
num_tokens=record.num_tokens,
|
||||
stride_k_bytes=record.per_token_k_bytes,
|
||||
stride_v_bytes=record.per_token_v_bytes,
|
||||
layer_num=ctrl.layer_num,
|
||||
available_tokens=available,
|
||||
)
|
||||
|
||||
def snapshot_dump(
|
||||
self, recv_req: SnapshotDumpReqInput
|
||||
) -> SnapshotDumpReqOutput:
|
||||
"""D-side: gather session KV from kv_pool, RDMA-write to remote slots."""
|
||||
"""D-side: gather session KV from kv_pool, RDMA-write into P's snapshot_buf."""
|
||||
ctrl = self.snapshot_link_controller
|
||||
if ctrl is None:
|
||||
return SnapshotDumpReqOutput(ok=False, reason="snapshot-link-disabled")
|
||||
@@ -3795,110 +3803,60 @@ class Scheduler(
|
||||
kv_committed_len = int(slot.kv_committed_len)
|
||||
if kv_committed_len == 0:
|
||||
return SnapshotDumpReqOutput(ok=False, reason="zero-committed-len")
|
||||
# Read kv_indices for the session's prefix
|
||||
try:
|
||||
kv_idx_tensor = self.req_to_token_pool.req_to_token[
|
||||
slot.req_pool_idx, :kv_committed_len
|
||||
]
|
||||
src_slot_indices = [int(x) for x in kv_idx_tensor.tolist()]
|
||||
except Exception as e:
|
||||
logger.exception("snapshot_dump: failed to read kv_indices: %s", e)
|
||||
return SnapshotDumpReqOutput(ok=False, reason=f"read-indices-failed: {e!r}")
|
||||
|
||||
# Truncate to the count P prepared for (must match)
|
||||
target_n = len(recv_req.target_slot_indices)
|
||||
if target_n > kv_committed_len:
|
||||
return SnapshotDumpReqOutput(
|
||||
ok=False,
|
||||
reason=f"target-larger-than-source({target_n}>{kv_committed_len})",
|
||||
)
|
||||
src_slot_indices = src_slot_indices[:target_n]
|
||||
return SnapshotDumpReqOutput(ok=False, reason=f"read-indices-failed:{e!r}")
|
||||
|
||||
try:
|
||||
ret, bytes_pushed = ctrl.push_session_kv(
|
||||
ret, bytes_pushed = ctrl.push_session_to_snapshot_buf(
|
||||
target_snapshot_session_id=recv_req.target_snapshot_session_id,
|
||||
src_slot_indices=src_slot_indices,
|
||||
target_k_base_ptrs=recv_req.target_k_base_ptrs,
|
||||
target_v_base_ptrs=recv_req.target_v_base_ptrs,
|
||||
target_slot_indices=recv_req.target_slot_indices[:target_n],
|
||||
target_stride_k_bytes=recv_req.target_stride_k_bytes,
|
||||
target_stride_v_bytes=recv_req.target_stride_v_bytes,
|
||||
target_snapshot_buf_base=recv_req.target_snapshot_buf_base,
|
||||
target_k_layer_offsets=recv_req.target_k_layer_offsets,
|
||||
target_v_layer_offsets=recv_req.target_v_layer_offsets,
|
||||
target_per_token_k_bytes=recv_req.target_stride_k_bytes,
|
||||
target_per_token_v_bytes=recv_req.target_stride_v_bytes,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("snapshot_dump: push_session_kv threw: %s", e)
|
||||
return SnapshotDumpReqOutput(ok=False, reason=f"push-failed: {e!r}")
|
||||
return SnapshotDumpReqOutput(ok=False, reason=f"push-failed:{e!r}")
|
||||
|
||||
if ret != 0:
|
||||
return SnapshotDumpReqOutput(
|
||||
ok=False,
|
||||
reason=f"mooncake-batch-write-ret={ret}",
|
||||
ok=False, reason=f"mooncake-batch-write-ret={ret}",
|
||||
bytes_pushed=int(bytes_pushed),
|
||||
kv_committed_len=int(kv_committed_len),
|
||||
kv_committed_len=kv_committed_len,
|
||||
)
|
||||
return SnapshotDumpReqOutput(
|
||||
ok=True, bytes_pushed=int(bytes_pushed),
|
||||
kv_committed_len=int(kv_committed_len),
|
||||
token_ids=[], # caller already has token_ids
|
||||
kv_committed_len=kv_committed_len,
|
||||
token_ids=[],
|
||||
)
|
||||
|
||||
def snapshot_finalize_ingest(
|
||||
self, recv_req: SnapshotFinalizeIngestReqInput
|
||||
) -> SnapshotFinalizeIngestReqOutput:
|
||||
"""P-side: insert (token_ids, slot_indices) into radix tree."""
|
||||
"""P-side: copy snapshot_buf slab into kv_pool + insert into radix tree.
|
||||
|
||||
Refactored per docs/SNAPSHOT_STORE_REFACTOR_ZH.md: kv_pool alloc
|
||||
happens HERE (deferred from prepare_receive), so we never block
|
||||
D's RDMA write on kv_pool contention.
|
||||
"""
|
||||
ctrl = self.snapshot_link_controller
|
||||
if ctrl is None:
|
||||
return SnapshotFinalizeIngestReqOutput(
|
||||
ok=False, reason="snapshot-link-disabled",
|
||||
)
|
||||
record = ctrl.take_record(recv_req.session_id)
|
||||
if record is None:
|
||||
return SnapshotFinalizeIngestReqOutput(
|
||||
ok=False, reason="no-pending-ingest",
|
||||
)
|
||||
# Sanity: the slot indices we're about to insert should match the ones we reserved.
|
||||
if list(recv_req.slot_indices) != record.slot_indices:
|
||||
# The caller passed back the slot indices we returned in prepare; if they
|
||||
# don't match, something's gone wrong. Free reserved slots and bail.
|
||||
try:
|
||||
ctrl._free_slots(record.slot_indices)
|
||||
except Exception:
|
||||
pass
|
||||
return SnapshotFinalizeIngestReqOutput(
|
||||
ok=False,
|
||||
reason="slot-indices-mismatch",
|
||||
)
|
||||
n_tokens = min(len(recv_req.token_ids), len(record.slot_indices))
|
||||
if n_tokens == 0:
|
||||
ctrl._free_slots(record.slot_indices)
|
||||
return SnapshotFinalizeIngestReqOutput(ok=False, reason="empty-token-ids")
|
||||
try:
|
||||
import torch
|
||||
from sglang.srt.mem_cache.base_prefix_cache import InsertParams
|
||||
from sglang.srt.mem_cache.radix_cache import RadixKey
|
||||
kv_indices = torch.tensor(
|
||||
record.slot_indices[:n_tokens],
|
||||
dtype=torch.int64,
|
||||
device=self.tree_cache.token_to_kv_pool_allocator.device,
|
||||
)
|
||||
radix_key = RadixKey(recv_req.token_ids[:n_tokens], None)
|
||||
inner = (
|
||||
self.tree_cache.inner
|
||||
if isinstance(self.tree_cache, SessionAwareCache)
|
||||
else self.tree_cache
|
||||
)
|
||||
result = inner.insert(InsertParams(key=radix_key, value=kv_indices))
|
||||
inserted = int(result.prefix_len)
|
||||
except Exception as e:
|
||||
logger.exception("snapshot_finalize_ingest: radix insert failed: %s", e)
|
||||
try:
|
||||
ctrl._free_slots(record.slot_indices)
|
||||
except Exception:
|
||||
pass
|
||||
return SnapshotFinalizeIngestReqOutput(
|
||||
ok=False, reason=f"radix-insert-failed: {e!r}",
|
||||
)
|
||||
ok, reason, inserted_prefix_len = ctrl.ingest_snapshot_into_kvpool(
|
||||
session_id=recv_req.session_id,
|
||||
token_ids=list(recv_req.token_ids),
|
||||
)
|
||||
return SnapshotFinalizeIngestReqOutput(
|
||||
ok=True, inserted_prefix_len=inserted,
|
||||
ok=bool(ok), reason=reason if not ok else None,
|
||||
inserted_prefix_len=int(inserted_prefix_len),
|
||||
)
|
||||
|
||||
def _compute_backpressure_pause_hint(
|
||||
|
||||
@@ -181,27 +181,18 @@ class SchedulerRuntimeCheckerMixin:
|
||||
return memory_leak, token_msg
|
||||
|
||||
def _check_radix_cache_memory(self: Scheduler):
|
||||
# NB: as of SnapshotStore refactor (see docs/SNAPSHOT_STORE_REFACTOR_ZH.md)
|
||||
# prepare_receive no longer touches kv_pool — slots are alloc'd from
|
||||
# a dedicated snapshot_buf. So no snapshot_reserved accounting needed.
|
||||
_, _, available_size, evictable_size = self._get_token_info()
|
||||
protected_size = self.tree_cache.protected_size()
|
||||
session_held = self._session_held_tokens()
|
||||
# Snapshot link prepare_receive reserves slots that aren't yet visible
|
||||
# to radix / session bookkeeping until finalize_ingest. Count them so
|
||||
# the leak check doesn't fire while a snapshot ingest is in-flight.
|
||||
snapshot_reserved = 0
|
||||
ctrl = getattr(self, "snapshot_link_controller", None)
|
||||
if ctrl is not None:
|
||||
try:
|
||||
snapshot_reserved = sum(
|
||||
len(rec.slot_indices) for rec in ctrl._ingest_records.values()
|
||||
)
|
||||
except Exception:
|
||||
snapshot_reserved = 0
|
||||
memory_leak = (available_size + evictable_size) != (
|
||||
self.max_total_num_tokens - protected_size - session_held - snapshot_reserved
|
||||
self.max_total_num_tokens - protected_size - session_held
|
||||
)
|
||||
token_msg = (
|
||||
f"{self.max_total_num_tokens=}, {available_size=}, {evictable_size=}, "
|
||||
f"{protected_size=}, {session_held=}, {snapshot_reserved=}\n"
|
||||
f"{protected_size=}, {session_held=}\n"
|
||||
)
|
||||
return memory_leak, token_msg
|
||||
|
||||
|
||||
32
third_party/traces/README.md
vendored
Normal file
32
third_party/traces/README.md
vendored
Normal file
@@ -0,0 +1,32 @@
|
||||
# Replay traces
|
||||
|
||||
为了方便跨主机传输,把 benchmark 用到的 trace 文件放在这里。该目录在
|
||||
`.gitignore` 中显式 whitelist(同 `third_party/sglang/`),文件随 git 一起走。
|
||||
|
||||
## 文件清单
|
||||
|
||||
| 文件 | 大小 | 内容 | 来源 |
|
||||
|---|---:|---|---|
|
||||
| `qwen35-swebench-50sess.jsonl` | 54 MB | 4449 reqs / 52 sessions / Qwen3.5-35B 推理产物 | `simm-swe-bench` 项目用 SiBench replay SiCo `swe.jsonl` 经 SGLang 跑出 audit.jsonl,再用 `scripts/convert_audit_to_trace.py` 转 |
|
||||
|
||||
详细来源见 `docs/ONBOARDING_NEXT_AGENT_ZH.md` 和实际 schema 见 `src/agentic_pd_hybrid/trace.py`。
|
||||
|
||||
## 使用方法
|
||||
|
||||
Replay 端的 trace 路径由 CLI flag `--trace` 指定。默认 sweep 脚本里指向
|
||||
`outputs/qwen35-swebench-50sess.jsonl`——为了向后兼容老脚本,**建议在 clone 后
|
||||
软链接一份过去**:
|
||||
|
||||
```bash
|
||||
mkdir -p outputs
|
||||
ln -sf ../third_party/traces/qwen35-swebench-50sess.jsonl \
|
||||
outputs/qwen35-swebench-50sess.jsonl
|
||||
```
|
||||
|
||||
或者直接改 sweep 脚本里 `--trace` 路径指向 `third_party/traces/...`。
|
||||
|
||||
## 添加新 trace
|
||||
|
||||
如果未来加新 trace 文件(如 `codex_swebenchpro` 转换后的版本),直接放本目录,
|
||||
更新本 README 的清单即可。**别把超过 100 MB 的单文件直接 git add**——GitLab
|
||||
默认对未启用 LFS 的单文件有 100 MB 限制。
|
||||
4449
third_party/traces/qwen35-swebench-50sess.jsonl
vendored
Normal file
4449
third_party/traces/qwen35-swebench-50sess.jsonl
vendored
Normal file
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user