Skip to content

fix(scheduler): reserve paged KV at allocation granularity - #367

Open
taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:fix/prefill-page-admission
Open

fix(scheduler): reserve paged KV at allocation granularity#367
taking-lying-flat wants to merge 1 commit into
FlashML-org:mainfrom
taking-lying-flat:fix/prefill-page-admission

Conversation

@taking-lying-flat

Copy link
Copy Markdown
Contributor

Summary

  • Account prefill KV reservations in whole pages per request.
  • Use the same page-ceiling span as CacheManager.allocate_paged() for both admission checks and cumulative reservations.
  • Keep requests queued when the remaining KV pool cannot provide a distinct page for each request.

Problem

CacheManager.available_size reports token-equivalent capacity, but the allocator assigns whole pages independently to each request. PrefillAdder previously charged raw token counts, so multiple short requests could pass admission even when they required more physical pages than the pool contained. The subsequent allocate_paged() call could then fail during eviction or hit the "Eviction did not free enough space" assertion.

For example, with a 128-token page and one free page, two requests that each need two tokens were charged as four tokens total even though they require two separate pages.

Validation

  • 51 relevant scheduler tests passed across radix, chunked prefill, hybrid, SWA page-size, DSV4, and abort paths.
  • Ruff check passed for python/freetoken/scheduler/prefill.py.
  • git diff --check passed.

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.

1 participant