Skip to content

fix(scheduler): stop the swa prefill commit from freeing pages the request still reads - #651

Open
ischencheng wants to merge 1 commit into
FlashML-org:mainfrom
ischencheng:fix/swa-commit-locked-tombstone
Open

ischencheng wants to merge 1 commit into
FlashML-org:mainfrom
ischencheng:fix/swa-commit-locked-tombstone

Conversation

@ischencheng

Copy link
Copy Markdown
Contributor

Fixes #204.

_cache_req_swa can free pages that the committing request is still reading. It happens when the request's freshly prefilled span runs through a tombstone that another running request still full-locks:

  1. C is decoding on a shared prefix S. It full-locks S, but its swa lock only covers its own last window.
  2. S becomes a tombstone while C still holds it. A request A that also starts with S finishes and its finish-time trim_head_swa tombstones the head (or evict_swa does it under pool pressure).
  3. B = S + a tail shorter than the sliding window arrives. match_prefix stops at the tombstone, so B prefills S itself into its own pages.
  4. B's prefill commit calls insert. For a tombstone with ref_count > 0, insert keeps the tree's slots and returns B's copy of S in freed, which is then freed from both pools. The re-point after it only covers m.cached_len, and the re-match truncates to 0 again because the live run after the tombstone (B's tail) is shorter than the window. So B's row still names pages on the free list, and its in-window positions map to swa slot 0.
  5. B decodes on those pages, and its finish commit frees them a second time. That leaves the duplicate pages in free_slots that Concurrent tool-calling load kills the worker: duplicate pages in the full KV free list (root cause), surfacing as SWA-slot leak/double-free #204 measured, and the idle check_integrity fails with SWA-slot leak/double-free.

The radix and hybrid paths re-point the deduped span at the tree's pages. That doesn't work here: the tree's copy has no swa, and B's window still covers the end of S. So this skips the unfinished commit when insert would take that branch, like the hybrid path does for a chunk with no tracked boundary. B keeps its own pages and its admission handle, and the finish commit inserts and frees the duplicates once. The finish path is unchanged. SWARadixCache.has_locked_tombstone is a read-only walk with the same condition as that branch of insert.

The cost is that in this case B holds its own copy of S until it finishes instead of sharing the tree's. Those are pages it already allocated for its prefill.

test_swa_unfinished_commit_keeps_its_pages_under_a_locked_tombstone builds the C/A/B sequence on a real CacheManager and SWARadixCache (window 16), next to the radix and hybrid versions of the same test. On main, 64 positions of B's row are on the free list and 8 of its 16 in-window positions map to swa slot 0. After both requests finish, free_slots holds 64 duplicates and check_integrity raises SWA-slot leak/double-free: free(191) + tree(72) != capacity(255).

Tested on an NVIDIA L4 (sm_89), driver 580.95.05, CUDA 13.0, torch 2.11.0+cu130, on top of 0781324:

command main this PR
pytest tests/scheduler/test_commit_repoints_page_table.py (this PR's file) 1 failed, 4 passed 5 passed
pytest tests/scheduler -m "not slow" 102 passed 103 passed
pytest tests/kvcache -m "not slow" 270 passed, 1 skipped 270 passed, 1 skipped

End to end on an L40S (sm_89), same driver and torch, with openai/gpt-oss-20b (sliding window 128): ft serve --model <snapshot> --max-seq-len-override 8192, which resolves to triton attention, offload MoE, swa_radix and page size 1. The script below sends, at temperature 0: C = S + 261 tokens decoding 1500 tokens, A = S + 514 tokens with max_tokens=1 once C is decoding, then B = S + 14 tokens. S is 2558 tokens. The scheduler log shows A reusing all 2558 tokens of S and B reusing none.

main this PR main, --cache-type naive
B's output degenerates into repeated We / The lines (differs run to run) answers G-84, same text in all 4 runs same text as this PR
20 s after going idle check_integrity: SWA-slot leak/double-free: free(192409) + tree(781) != capacity(193175), then Backend worker is gone and cannot be restarted; stopping the API server still serving still serving

main was run in 2 containers and this PR in 3, with the same outcome each time (the free/tree numbers above are from the first main run). B served alone on a fresh server also answers G-84, but its text drifts from the runs above after a few tokens. The naive run drifts from it in the same way, so that comes from B being served next to C, not from the prefix cache.

repro
# Against a running `ft serve --model openai/gpt-oss-20b --max-seq-len-override 8192`.
import json
import threading
import time
import urllib.request

URL = "http://127.0.0.1:1919"


def complete(prompt, max_tokens, **kw):
    body = {"model": "m", "prompt": prompt, "max_tokens": max_tokens, "temperature": 0, **kw}
    req = urllib.request.Request(URL + "/v1/completions", json.dumps(body).encode(),
                                 {"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=900) as r:
        return json.loads(r.read())["choices"][0]["text"]


s = "Plant operations manual.\n" + "".join(
    f"Rule {i}: when the operator asks about subsystem {i}, check the gauge labelled "
    f"G-{i * 7 % 97}, write the reading in the log, and only then open valve V-{i * 13 % 89}. "
    for i in range(1, 70)) + "\n"
c = s + ("Task: write a very long, detailed shift report that walks through every rule above one "
         "by one, explaining the reasoning behind each step, the failure it prevents, and an "
         "example of a past incident where skipping it caused trouble. Be thorough and do not "
         "summarize. " * 5)
a = s + "Separate note from the night shift: " + " ".join(
    f"the pump on line {i} was serviced and its filter replaced;" for i in range(1, 40))
b = s + "Question: which gauge should be checked for subsystem 12? Answer:"

complete("The quick brown fox jumps over the lazy dog. " * 40, 8)  # warm up
t = threading.Thread(target=complete, args=(c, 1500), kwargs={"ignore_eos": True})
t.start()
time.sleep(5)                        # C is decoding and full-locks S
complete(a, 1)                       # A finishes; its trim tombstones S
print("B:", repr(complete(b, 96)))   # B reuses nothing, prefills S itself, commits
t.join()
time.sleep(20)                       # idle -> check_integrity
try:
    urllib.request.urlopen(URL + "/v1/models", timeout=5)
    print("server still up")
except OSError as e:
    print("server gone:", e)

#488 would remove the trim trigger, but not the evict_swa one. Until this lands, --cache-type naive avoids it, as the last column shows.

KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 10, 2026
…ll commit from freeing pages the request still reads
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 10, 2026
…d tombstone through swa pool pressure

next keeps the finished prompt head's swa (FlashML-org#488), so the finish no longer
tombstones it; ensure_swa_slots does, and the test fails without the fix.

Assisted-by: Claude Opus 5.5

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Concurrent tool-calling load kills the worker: duplicate pages in the full KV free list (root cause), surfacing as SWA-slot leak/double-free

1 participant