# CP HiCache L1/L2 Free-Room and Contiguous Allocation Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Reduce repeated L1/L2 eviction churn under CP HiCache pressure and make later allocations more likely to receive contiguous page runs.

**Architecture:** Add trigger/target free-room watermarks for L2 host and L1 device owner lanes. Trigger controls when extra eviction starts; target controls how much room remains after a triggered eviction. Preserve CP owner-lane correctness and avoid new collectives. Add contiguous-preferred allocation only as best effort: it may improve descriptor coalescing, but must never choose pages with the wrong owner.

**Tech Stack:** Python, PyTorch tensors, SGLang CP shared-KV allocator, HiRadixCache, HiCacheController, unit tests under `test/registered/unit`.

---

## Findings from code reading

1. `HostKVCache.alloc()` currently slices `free_slots[:need_size]`; `HostKVCache.free()` appends freed token indices. Over-evicting alone cannot guarantee contiguous later host allocations because older fragmented `free_slots` remain at the front.
2. CP host write admission lives in `HiRadixCache._cp_build_write_admission()`. It currently computes exact per-owner deficits: `max(0, required - available)` for target and draft host pools.
3. CP L1 owner-lane prefill allocation uses `alloc_paged_token_slots_extend()` -> `CPSharedPagedTokenToKVPoolAllocator.alloc_extend_compute_owner()`. If first allocation fails, `_evict_for_compute_owner_lanes()` evicts exact owner deficits.
4. CP L1 load-back allocation uses `HiRadixCache._build_cp_load_back_plan()` -> `HiCacheController.load_cp()` -> `alloc_pages_with_owners()`. The load-back plan also computes exact owner-lane deficits.
5. `CPSharedPagedTokenToKVPoolAllocator._select_compute_owner_pages()` preserves owner correctness but is not contiguity-aware beyond consuming the current `free_pages`/`release_pages` order.
6. Free-room must be trigger/target based, not always-on target maintenance. The accepted formula is:
   ```text
   if available >= required + trigger_room:
       deficit = 0
   else:
       deficit = max(0, required + target_room - available)
   ```
7. Implementation correction: L1 owner-lane free-room deficits are page-unit deficits, not token-unit deficits. The owner-lane eviction API already consumes page counts in `owner_lane_deficits`, while `num_tokens` remains the scalar coarse eviction amount.
8. Implementation correction: L1 contiguous-preferred allocation must preserve owner order in the final request page sequence but may change which logical page IDs satisfy each owner. Existing tests that asserted exact FIFO pages were updated to assert the new contiguous preference rather than the old incidental page choice.

## Progress as of 2026-06-02

- Done: Task 1-3 host/L2 ratio knobs, trigger-target helper, and CP host write admission.
- Done: Task 4 host `alloc_contiguous_preferred()` and CP target/draft reservation use.
- Done: Task 5 L1 owner-lane free-room in prefill eviction and CP load-back planning.
- Done: Task 6 L1 owner-lane contiguous-preferred free-page selection without `torch.isin`, sort-merge, or CUDA `.tolist()` path.
- Verification: remote py_compile passed for touched runtime/test files. Remote targeted tests passed with an sgl_kernel stub runner for allocator/load-back tests because importing the installed `sgl_kernel` directly can abort this container during pytest collection.

## Files

- Modify: `python/sglang/srt/server_args.py`
  - Add four runtime knobs: host target/trigger ratios and L1 target/trigger ratios.
  - Validate each ratio and `trigger <= target`.
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
  - Add page-aligned free-room helper.
  - Apply helper to CP host write admission.
  - Apply helper to CP L1 load-back owner-lane plan.
  - Store configured ratios on `HiRadixCache`.
- Modify: `python/sglang/srt/mem_cache/common.py`
  - Apply L1 owner-lane trigger/target deficits before `_evict_for_compute_owner_lanes()` calls `tree_cache.evict()`.
- Modify: `python/sglang/srt/mem_cache/allocator.py`
  - Add owner-lane capacity helper for CP shared paged allocator.
  - Add optional owner-lane contiguous-preferred page selection without changing owner correctness.
- Modify: `python/sglang/srt/mem_cache/memory_pool_host.py`
  - Add `alloc_contiguous_preferred()` that finds page-contiguous host token runs before falling back to FIFO.
- Modify: `python/sglang/srt/managers/cache_controller.py`
  - Use host `alloc_contiguous_preferred()` for CP target/draft host reservations when available.
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`
- Test: `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
- Test: `test/registered/unit/managers/test_hicache_controller_cp.py`

---

### Task 1: Add ratio knobs and validation

**Files:**
- Modify: `python/sglang/srt/server_args.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`

- [ ] **Step 1: Write failing validation tests**

Add tests that instantiate `ServerArgs` with invalid ratios and expect `ValueError` from `__post_init__()` or the existing validation path. Cover negative ratio and trigger greater than target.

- [ ] **Step 2: Run validation tests and verify RED**

Run:
```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -k free_room_args
```
Expected: FAIL because the new attributes do not exist or validation is missing.

- [ ] **Step 3: Add args**

Add dataclass fields with default `0.0`:
```python
hicache_host_free_room_ratio: float = 0.0
hicache_host_free_room_trigger_ratio: float = 0.0
hicache_l1_free_room_ratio: float = 0.0
hicache_l1_free_room_trigger_ratio: float = 0.0
```

Add parser args near existing HiCache args:
```python
--hicache-host-free-room-ratio
--hicache-host-free-room-trigger-ratio
--hicache-l1-free-room-ratio
--hicache-l1-free-room-trigger-ratio
```

- [ ] **Step 4: Add validation**

In `_handle_hicache()` or adjacent HiCache validation, reject values outside `[0, 1)` and reject `trigger > target` for host and L1 pairs.

- [ ] **Step 5: Run validation tests and verify GREEN**

Run the same targeted pytest. Expected: PASS.

---

### Task 2: Implement reusable trigger/target free-room deficit helper

**Files:**
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`

- [ ] **Step 1: Write failing helper tests**

Add tests for a helper with page-size rounding:

- required fits above trigger -> deficit `0`.
- required fits but below target only -> still deficit `0` if trigger is `0`.
- required does not fit trigger -> deficit includes target room.
- ratio target is rounded up to page size.

- [ ] **Step 2: Run helper tests and verify RED**

Run:
```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -k free_room_deficit
```
Expected: FAIL because the helper does not exist.

- [ ] **Step 3: Add helper**

Add a pure helper in `hiradix_cache.py`:
```python
def _page_aligned_room_from_ratio(capacity: int, ratio: float, page_size: int) -> int:
    if ratio <= 0 or capacity <= 0:
        return 0
    raw = int(math.ceil(float(capacity) * float(ratio)))
    return ((raw + page_size - 1) // page_size) * page_size


def _free_room_deficit(
    *,
    required: int,
    available: int,
    capacity: int,
    page_size: int,
    target_ratio: float,
    trigger_ratio: float,
) -> int:
    target_room = _page_aligned_room_from_ratio(capacity, target_ratio, page_size)
    trigger_room = _page_aligned_room_from_ratio(capacity, trigger_ratio, page_size)
    if available >= required + trigger_room:
        return 0
    return max(0, required + target_room - available)
```

- [ ] **Step 4: Run helper tests and verify GREEN**

Run targeted pytest. Expected: PASS.

---

### Task 3: Apply free-room to L2 CP host write admission

**Files:**
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`

- [ ] **Step 1: Write failing admission tests**

Add tests around `_cp_build_write_admission()` using `HiRadixCache.__new__()`:

1. `host_free_room_ratio=0.0` preserves current exact deficits.
2. If `available >= required + trigger_room`, deficit is zero even if `available < required + target_room`.
3. If `available < required + trigger_room`, deficit is `required + target_room - available`.
4. Draft host lower availability drives the max deficit.

- [ ] **Step 2: Run tests and verify RED**

Run:
```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -k cp_host_free_room
```
Expected: FAIL because admission still uses exact deficits.

- [ ] **Step 3: Store ratios and update admission**

In `HiRadixCache.__init__`, store:
```python
self.hicache_host_free_room_ratio = server_args.hicache_host_free_room_ratio
self.hicache_host_free_room_trigger_ratio = server_args.hicache_host_free_room_trigger_ratio
self.hicache_l1_free_room_ratio = server_args.hicache_l1_free_room_ratio
self.hicache_l1_free_room_trigger_ratio = server_args.hicache_l1_free_room_trigger_ratio
```

In `_cp_build_write_admission()`, compute target and draft deficits with `_free_room_deficit()` using snapshot capacity per owner.

- [ ] **Step 4: Run admission tests and verify GREEN**

Run targeted pytest. Expected: PASS.

---

### Task 4: Add contiguous-preferred host allocation and use it for CP HiCache reservations

**Files:**
- Modify: `python/sglang/srt/mem_cache/memory_pool_host.py`
- Modify: `python/sglang/srt/managers/cache_controller.py`
- Test: `test/registered/unit/managers/test_hicache_controller_cp.py`

- [ ] **Step 1: Write failing host allocator test**

Create a fake or direct `HostKVCache.__new__()` instance with `page_size=4` and fragmented `free_slots` such as:
```python
[100,101,102,103, 8,9,10,11, 12,13,14,15]
```
Request `8` tokens. Expected contiguous-preferred allocation returns `[8..15]`, not the fragmented FIFO prefix if the prefix does not satisfy the requested run.

- [ ] **Step 2: Write failing controller test**

Extend `FakeHostPool` with `alloc_contiguous_preferred()` and verify `reserve_write_cp()` calls it instead of `alloc()` for CP target and draft host pools.

- [ ] **Step 3: Verify RED**

Run:
```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py -k contiguous_preferred
```
Expected: FAIL because the method/call does not exist.

- [ ] **Step 4: Implement host method**

Add `HostKVCache.alloc_contiguous_preferred(need_size)`:

- assert page-aligned size;
- return `None` if capacity is insufficient;
- if FIFO prefix is already a contiguous token run, reuse `alloc()` behavior;
- otherwise scan page starts in `free_slots` and select the first run of `need_size // page_size` consecutive pages;
- remove selected indices by boolean mask;
- fallback to `alloc()` if no run exists.

- [ ] **Step 5: Use it in CP reservation**

In `HiCacheController.reserve_write_cp()`, replace direct host allocation with:
```python
alloc_host = getattr(self.mem_pool_host, "alloc_contiguous_preferred", self.mem_pool_host.alloc)
host_indices = alloc_host(len(physical_device_indices))
```
Do the same for `draft_mem_pool_host`.

- [ ] **Step 6: Verify GREEN**

Run the targeted controller tests. Expected: PASS.

---

### Task 5: Apply L1 owner-lane free-room to prefill allocation and load-back

**Files:**
- Modify: `python/sglang/srt/mem_cache/allocator.py`
- Modify: `python/sglang/srt/mem_cache/common.py`
- Modify: `python/sglang/srt/mem_cache/hiradix_cache.py`
- Test: `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
- Test: `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`

- [ ] **Step 1: Write failing allocator capacity tests**

Add a test for `CPSharedPagedTokenToKVPoolAllocator.compute_owner_lane_capacity_pages()` expecting each lane capacity to equal `num_pages // cp_size`.

- [ ] **Step 2: Write failing prefill eviction tests**

Extend `_evict_for_compute_owner_lanes()` tests:

- with target/trigger ratio zero, current exact deficits are preserved;
- with trigger zero and target room two pages, an exhausted lane evicts `required + target_room - available` pages for that owner;
- with available satisfying trigger, no extra eviction runs.

- [ ] **Step 3: Write failing load-back plan tests**

Add a `HiRadixCache.__new__()` test proving `_build_cp_load_back_plan()` uses target/trigger L1 room in `deficit_by_owner`.

- [ ] **Step 4: Verify RED**

Run:
```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -k "owner_lane and free_room" \
  test/registered/unit/mem_cache/test_cp_hicache_metadata.py -k "load_back and free_room"
```
Expected: FAIL because L1 free-room is not implemented.

- [ ] **Step 5: Implement allocator capacity helper**

Add to `CPSharedPagedTokenToKVPoolAllocator`:
```python
def compute_owner_lane_capacity_pages(self) -> list[int]:
    capacity = self.num_pages // self.cp_size
    return [int(capacity) for _ in range(self.cp_size)]
```

- [ ] **Step 6: Update `_evict_for_compute_owner_lanes()`**

Give it optional `target_ratio` and `trigger_ratio`. Compute deficits from required/available/capacity using page units. Preserve current behavior when both ratios are zero.

- [ ] **Step 7: Pass ratios from `alloc_paged_token_slots_extend()`**

When CP shared KV server args are enabled, pass `server_args.hicache_l1_free_room_ratio` and `server_args.hicache_l1_free_room_trigger_ratio` into `_evict_for_compute_owner_lanes()`.

- [ ] **Step 8: Update load-back plan**

In `_build_cp_load_back_plan()` and `_refresh_cp_load_back_plan()`, compute L1 deficits with the same trigger/target helper. Use allocator capacity pages converted to token counts or compute directly in pages consistently with existing `compute_owner_lane_stats()` output.

- [ ] **Step 9: Verify GREEN**

Run targeted tests. Expected: PASS.

---

### Task 6: Add L1 owner-lane contiguous-preferred page selection

**Files:**
- Modify: `python/sglang/srt/mem_cache/allocator.py`
- Test: `test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`

- [ ] **Step 1: Write failing contiguity preference tests**

Use CPU allocator with fragmented owner-lane pages. For one owner requiring three pages, set candidates so the first owner-correct pages are fragmented but a later three-page physical run exists. Verify selected logical pages map to consecutive physical pages for that owner.

- [ ] **Step 2: Verify RED**

Run:
```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -k contiguous_owner_lane
```
Expected: FAIL because current selection takes first owner pages.

- [ ] **Step 3: Implement best-effort selection**

In `_select_compute_owner_pages()`:

- keep existing owner correctness;
- prefer a contiguous run in each owner lane when one can satisfy that owner's required count;
- otherwise fall back to the current first-available behavior;
- do not introduce CPU `.tolist()` or host sync in the CUDA hot path;
- preserve the existing no-`torch.isin` and no forced sort-merge tests.

- [ ] **Step 4: Verify GREEN**

Run targeted allocator tests. Expected: PASS.

---

### Task 7: Full verification and remote sync

**Files:** all modified files.

- [ ] **Step 1: Run local CPU tests**

Run:
```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
  test/registered/unit/mem_cache/test_cp_hicache_metadata.py \
  test/registered/unit/managers/test_hicache_controller_cp.py
```
Expected: PASS or dependency-related skip only; real failures must be fixed.

- [ ] **Step 2: Compile touched files**

Run:
```bash
cd /root/sglang-work/sglang-dev
python -m py_compile \
  python/sglang/srt/server_args.py \
  python/sglang/srt/mem_cache/hiradix_cache.py \
  python/sglang/srt/mem_cache/common.py \
  python/sglang/srt/mem_cache/allocator.py \
  python/sglang/srt/mem_cache/memory_pool_host.py \
  python/sglang/srt/managers/cache_controller.py
```
Expected: no output.

- [ ] **Step 3: Sync code to remote only after local CPU tests are clean**

Use `scp` to `g0034:/mnt/beegfs/cjy/sglang-dev`. Do not chown inside container.

- [ ] **Step 4: Run remote targeted tests in container**

Run via ssh/docker:
```bash
cd /sgl-workspace/sglang-tai
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
  test/registered/unit/mem_cache/test_cp_hicache_metadata.py \
  test/registered/unit/managers/test_hicache_controller_cp.py
```
Expected: PASS.

- [ ] **Step 5: Runtime validation knobs**

Start conservatively:
```bash
--hicache-host-free-room-trigger-ratio 0.0 \
--hicache-host-free-room-ratio 0.03 \
--hicache-l1-free-room-trigger-ratio 0.0 \
--hicache-l1-free-room-ratio 0.02
```
Observe eviction frequency, allocation fallback logs, cache hit, accept length, and Mooncake transfer block statistics.

---

## Self-review

- The plan keeps default behavior unchanged with all ratios `0.0`.
- The plan separates trigger and target room to avoid wasting cache capacity when current allocation fits.
- The plan avoids new collectives.
- The plan preserves CP owner-lane correctness.
- The plan records that over-eviction alone does not guarantee contiguous allocation; allocator policy is required.
- Risk retained: L1 contiguous-preferred selection must avoid CUDA host sync. If a correct no-sync implementation is too invasive, Task 6 should stop at tests and defer implementation rather than adding a slow hot-path algorithm.

---

## Microbenchmark follow-up: allocator overhead guardrail

### Added benchmark

File:
- `benchmark/hicache/bench_cp_hicache_allocator_overhead.py`
- unit smoke: `test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py`

Purpose:
- measure metadata/control-path cost without allocating real 220GB HiCache buffers;
- compare FIFO allocation vs contiguous-preferred allocation;
- cover Host/L2 and L1 CP owner-lane allocation;
- keep CUDA validation remote-only.

Important default behavior:
- `--host-sizes-gb 220` maps to page count via `floor(size_gb * 1e9 / (bytes_per_token * page_size))`.
- With default `bytes_per_token=100000` and `page_size=64`, 220GB maps to 34375 pages.
- `--host-pages` defaults to `8192,16384,32768` only when neither `--host-pages` nor `--host-sizes-gb` is supplied.

### Remote verification already run

Remote container path:
- `/sgl-workspace/sglang-tai`

Commands:
```bash
python -m py_compile \
  benchmark/hicache/bench_cp_hicache_allocator_overhead.py \
  test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py
PYTHONPATH=. python -m pytest -q \
  test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py
```
Result:
- `4 passed, 1 warning`

Host 220GB-equivalent smoke:
```bash
PYTHONPATH=. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \
  --bench host --host-sizes-gb 220 \
  --request-pages 1,8,64,512 --repeat 5 --warmup 2 \
  --patterns contiguous_fifo,fragmented_prefix_later_run,random_fragmented \
  --host-methods fifo,contiguous \
  --json-output /tmp/cp_hicache_allocator_host_220.json
```
Representative result:
- contiguous free list: contiguous-preferred remains cheap (`req=8 p50≈22us`, `req=512 p50≈86us`).
- fragmented/random free list: contiguous-preferred is multi-ms (`req=8/64 p50≈4.6ms`, `req=512 p50≈3.5ms`, random p95 can exceed 8ms).
- FIFO stays low-latency but loses contiguity in fragmented cases.

L1 CPU proxy smoke:
```bash
PYTHONPATH=python:. python benchmark/hicache/bench_cp_hicache_allocator_overhead.py \
  --bench l1 --device cpu --stub-sgl-kernel --physical-pages 8192 \
  --request-pages 64,512 --repeat 5 --warmup 2 \
  --l1-free-patterns owner_fragmented_later_run,random \
  --l1-owner-patterns round_robin,single_owner,zigzag \
  --l1-impl fifo,current \
  --json-output /tmp/cp_hicache_allocator_l1_cpu.json
```
Representative result:
- current contiguous-preferred L1 CPU proxy is slower than FIFO by ~0.4ms for single-owner and ~2.5-3.6ms for round-robin/zigzag at 8192 physical pages.
- it restores contiguity for `owner_fragmented_later_run`, but cannot restore contiguity for fully random free lists.

CUDA note:
- CUDA L1 smoke was not run in this pass because a production prefill server was active on `g0034` using the GPUs. Do not run CUDA allocator benchmark concurrently with active ETE unless explicitly scheduled.

### Finding / next implementation implication

The current contiguous-preferred scan is acceptable only when free lists are already mostly contiguous. Under fragmented 220GB-equivalent metadata it can add multi-ms CPU/control-path overhead. This means the free-room/evict strategy cannot rely on repeated full-list scans to recover contiguity in the hot path.

Next preferred direction:
1. Keep this benchmark as regression evidence.
2. Optimize allocator metadata before enabling aggressive contiguous-preferred allocation broadly:
   - Host/L2: maintain page-run or interval metadata instead of scanning/sorting the whole free-slot tensor on each allocation.
   - L1: avoid per-owner full free-list masks for small requests, or gate contiguous-preferred to fragmented-risk / large-request paths.
3. Use trigger/target free-room eviction to reduce eviction frequency, but do not assume it alone fixes allocator overhead.
