Skip to content

Support page size > 1#4356

Merged
merrymercy merged 21 commits intomainfrom
lianmin/page-size
Mar 13, 2025
Merged

Support page size > 1#4356
merrymercy merged 21 commits intomainfrom
lianmin/page-size

Conversation

@merrymercy
Copy link
Copy Markdown
Contributor

@merrymercy merrymercy commented Mar 12, 2025

  • This PR implements PagedTokenToKVPoolAllocator. It has almost the same interface as the old TokenToKVPoolAllocator, but it guarantees that the returned kv cache locations of one request are always page aligned. It is compatible with prefix cache.
  • Some other cleanup (e.g., move batch_is_full to ScheduleBatch)

Page size = 1, 4, 8, 16 shows no perf diff, because the page table management is all overlapped with GPU computation.
This PR will enable the integration of more attention kernels (e.g., FlashMLA and CuDNN attention)

Todos for the next PRs:

  • Change the attention kernel calls
  • Support speculative decoding

Comment thread python/sglang/srt/managers/schedule_policy.py
@merrymercy merrymercy mentioned this pull request Mar 13, 2025
67 tasks
@merrymercy merrymercy merged commit c76040e into main Mar 13, 2025
@merrymercy merrymercy deleted the lianmin/page-size branch March 13, 2025 05:22
hebiao064 pushed a commit to hebiao064/sglang that referenced this pull request Mar 13, 2025
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.

2 participants