Permute page table in benchmarking#2194
Conversation
WalkthroughChanged benchmark tests to generate page tables by applying a random permutation per batch instead of sequential block indices; updated Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello @jhjpark, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the benchmarking of paged KV cache operations by introducing random permutation of page tables. This change aims to provide a more robust and realistic evaluation of different backends by ensuring they correctly handle non-sequential page allocations, which is a common scenario in real-world memory management. The modification ensures that the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
jhjpark
requested review from
bkryu,
cyx-6,
jiahanc and
nvmbreughe
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request modifies the benchmarks/routines/attention.py file, specifically within the testBatchDecodeWithPagedKVCacheWrapper and testBatchPrefillWithPagedKVCacheWrapper functions. The changes introduce randomization to the initialization of block_tables by using torch.randperm for page assignments, moving away from a sequential range. Additionally, the assignment of kv_indices is simplified by directly slicing from the block_tables instead of generating values with torch.arange.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
benchmarks/routines/attention.py (1)
397-405: Block table permutation per sequence is correct; consider a slightly cleaner constructionThe new
block_tablesdefinition correctly:
- Keeps each sequence
iwithin its own page range[i * num_pages_per_seq, (i + 1) * num_pages_per_seq).- Applies an independent random permutation of pages per sequence, which matches the PR goal of testing non‑trivial page layouts.
One minor style nit: iterating over
torch.randperm(num_pages_per_seq)in a list comprehension yields 0‑D tensors, whichtorch.tensor(...)then converts. Functionally this is fine, but you could make it a bit clearer and avoid 0‑D tensors by constructing directly in Torch, e.g.:perm = torch.stack( [torch.randperm(num_pages_per_seq) for _ in range(batch_size)], dim=0, ) block_tables = ( perm + torch.arange(batch_size).unsqueeze(1) * num_pages_per_seq ).to(device=device, dtype=torch.int)Not critical at all, just a readability/clarity option for benchmark setup code.
Also applies to: 833-841
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
benchmarks/routines/attention.py(4 hunks)
🔇 Additional comments (1)
benchmarks/routines/attention.py (1)
421-425: kv_indices now correctly respect the per‑sequence page permutationUsing
kv_indices[start_idx:end_idx] = block_tables[i, : end_idx - start_idx]ties the linearkv_indicesview directly to the permutedblock_tablesrows, whileend_idx - start_idxmatches the page count derived fromkv_indptr. This ensures:
- Each sequence uses exactly its required number of pages.
- The actual page IDs seen by the backends via
kv_indicesstay consistent with the permutedblock_tables, avoiding mismatches between kernels that consult one vs the other.This aligns with the PR objective of making the paged‑KV benchmarks exercise non‑trivial page tables while still passing a coherent mapping to all backends.
Also applies to: 883-887
|
Thanks @jhjpark. Let's confirm that we get expected results when we run a larger testlist and leave a comment confirming it is good to go |
Head branch was pushed to by a user without write access
jhjpark
force-pushed
the
jopark/fix_benchmark_page_table
branch
from
ee0f239 to
1b5503c
Compare
<!-- .github/pull_request_template.md --> ## 📌 Description This PR will permute page tables during benchmarking of different backends and ensure the correct page table gets passed to the backends. <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Improved batch operation tests by introducing randomized page ordering and updated indexing to validate behavior under non-sequential page arrangements, increasing robustness and coverage. * Confirmed no changes to public interfaces or external behavior; these updates only strengthen internal test validation. <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub> <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Joshua Park <jopark@nvidia.com>
3 participants
📌 Description
This PR will permute page tables during benchmarking of different backends and ensure the correct page table gets passed to the backends.
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.