Fix filtered topk unified kernel#2319
Conversation
murphymatt
requested review from
IwakuraRein,
jiahanc,
kahyunnam and
yzh119
as code owners
|
Warning Rate limit exceeded@murphymatt has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 2 minutes and 23 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds explicit per-element length checks to a top-k CUDA kernel to prevent out-of-bounds reads/writes for tail elements and adds a ragged-row test exercising unaligned per-row lengths to validate index bounds. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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 @murphymatt, 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 resolves a critical bug in the filtered topk unified kernel that caused out-of-bounds memory access during vectorized processing. The issue arose when the input data length was not perfectly divisible by the vector size, leading to incorrect indexing and Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. 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
|
![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 aims to fix a bug in the filtered_topk_unified kernel by adding bounds checking to vectorized loops. This is an important fix for edge cases where the input length is not a multiple of the vector size. However, the current implementation introduces two critical issues. First, a logic error in the histogram update will cause double-counting of elements. Second, a typo in a conditional check uses a variable before it's declared, which will lead to a compilation error. I've provided suggestions to correct these issues. The other bounds check added in the PR appears to be correct.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In @include/flashinfer/topk.cuh:
- Around line 2044-2051: The loop in topk.cuh uses an undefined variable `bin`
in the bounds check; change the condition to check the computed index using
`base + j` before computing `bin`. Specifically, in the for-loop that iterates
VEC_SIZE, replace `if (bin + j < length)` with a bounds check like `if (base + j
< length)` so you only compute `const auto bin =
static_cast<int>(Traits::ToCoarseKey(score_vec[j]));` afterwards, and then
proceed to compare `bin` to `threshold_bin` and use `atomicAdd(&s_counter, 1)` /
`s_indices[pos] = base + j` as before.
- Around line 1996-2002: The loop currently increments s_histogram twice for
in-bounds elements because there is an unconditional atomicAdd followed by a
guarded atomicAdd; remove the unconditional increment (the first call to
Traits::ToCoarseKey(score_vec[j]) and atomicAdd(&s_histogram[bin], 1)) and keep
only the bounds-checked version so each element contributes at most one
histogram increment; ensure you compute bin once per j using
Traits::ToCoarseKey(score_vec[j]) and call atomicAdd(&s_histogram[bin], 1) only
inside the if (base + j < length) branch to avoid double-counting.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
include/flashinfer/topk.cuh
🧰 Additional context used
📓 Path-based instructions (1)
include/**/*.cuh
📄 CodeRabbit inference engine (CLAUDE.md)
include/**/*.cuh: Torch headers MUST NOT be included in files within theinclude/directory - keep framework-agnostic CUDA kernels that accept raw pointers
Kernel code ininclude/flashinfer/is automatically picked up by JIT compilation on changes - no pip reinstall needed
For performance-critical hot paths, leave comments explaining special algorithmic choices and potential alternatives for future reviewers
Consult the PTX ISA documentation (https://docs.nvidia.com/cuda/parallel-thread-execution/) for low-level instruction details and new GPU architecture features when writing inline PTX assembly
Files:
include/flashinfer/topk.cuh
🧠 Learnings (1)
📚 Learning: 2025-12-30T09:34:39.900Z
Learnt from: CR
Repo: flashinfer-ai/flashinfer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-30T09:34:39.900Z
Learning: Applies to include/**/*.cuh : For performance-critical hot paths, leave comments explaining special algorithmic choices and potential alternatives for future reviewers
Applied to files:
include/flashinfer/topk.cuh
🔇 Additional comments (1)
include/flashinfer/topk.cuh (1)
2066-2081: LGTM: Correct bounds checking implementation.This segment properly guards all operations on vector elements with
if (base + j < length), preventing out-of-bounds reads fromscore_vec[j]and invalid writes to shared memory arrays when the vector chunk extends beyond the input length.This is the correct pattern for vectorized boundary handling in CUDA kernels.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
tests/utils/test_topk.py (1)
382-382: Good addition to test non-aligned lengths.The
length_offsetparameter effectively creates test cases whereactual_length = max_len - length_offsetmay not be aligned toVEC_SIZE, which exercises the bounds checking fix described in the PR.📝 Optional: Consider adding a comment explaining the purpose
Adding a brief comment would help future maintainers understand why these specific offset values are tested:
+# Test both VEC_SIZE-aligned (0) and non-aligned (1) lengths to ensure bounds checking @pytest.mark.parametrize("length_offset", [0, 1])Alternatively, you could test additional offset values (2, 3, etc.) to cover different VEC_SIZE alignments (e.g., VEC_SIZE=2, 4, 8), though
[0, 1]is sufficient for catching the out-of-bounds issue.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/utils/test_topk.py
🧰 Additional context used
📓 Path-based instructions (1)
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Test implementations should useflashinfer.utilsfunctions (get_compute_capability,is_sm90a_supported,is_sm100a_supported, etc.) to skip tests on unsupported GPU architectures
For testing withmpirunon multi-GPU systems, use the pattern:mpirun -np <num_gpus> pytest tests/path/to/test.py::test_function
Avoid OOM (out-of-memory) errors in tests by using appropriate problem sizes -tests/conftest.pyprovides auto-skipping for OOM tests as a safety net but should not be relied upon
Files:
tests/utils/test_topk.py
🔇 Additional comments (4)
tests/utils/test_topk.py (4)
386-391: LGTM! Correctly implements non-aligned length testing.The function signature update and
actual_lengthcomputation properly introduce variable row lengths into the test. The skip condition correctly usesactual_lengthto prevent invalid test cases wherek > actual_length.
405-406: Correctly applies actual_length to test data.The lengths tensor now uses
actual_lengthinstead ofmax_len, ensuring the kernel processes non-VEC_SIZE-aligned row lengths. This setup properly exercises the bounds checking fix.
420-430: Excellent bounds validation to verify the kernel fix.This explicit validation loop ensures that all output indices fall within the valid range
[offset, offset + length)for each row. This is the critical check that would catch the bug described in the PR, where vectorized loads without bounds checking could produce out-of-range indices whenlength % VEC_SIZE != 0.The validation logic is sound:
- Filters padding values (
>= 0)- Converts to relative indices for range checking
- Provides clear error messages with actual values
1209-1211: Good coverage of aligned and non-aligned test cases.The test invocations now exercise both
length_offset=0(VEC_SIZE-aligned) andlength_offset=1(potentially non-aligned) scenarios. This ensures the kernel fix handles both common and edge cases correctly.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tests/utils/test_topk.py (2)
786-788: Consider addingtorch.bfloat16to dtype parameters.Other tests in this file include
bfloat16(e.g., lines 343, 384), and the function supports it per the relevant code snippets. Sincebfloat16would have the same VEC_SIZE asfloat16(8), it would provide additional validation of the bounds-checking fix for unaligned lengths.♻️ Suggested enhancement
@pytest.mark.parametrize("length_offset", [1, 2, 3, 5, 6, 7]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) def test_ragged_transform_unaligned_length_bounds(length_offset, dtype):
788-819: Consider testingtop_k_page_table_transformwith unaligned lengths.The enriched summary indicates that bounds checks were added to multiple locations in
topk.cuhfor both histogram increments and bin computations. While this test coverstop_k_ragged_transform, it may be worth adding a similar test fortop_k_page_table_transformto ensure the bounds-checking fix applies to both code paths.Additionally, since the PR specifically mentions the "filtered topk unified kernel," you might consider explicitly testing with the filtered algorithm:
def test_ragged_transform_unaligned_length_bounds_filtered(length_offset, dtype): """Test with filtered algorithm explicitly.""" from flashinfer.topk import can_implement_filtered_topk if not can_implement_filtered_topk(): pytest.skip("GPU does not support filtered topk") os.environ["FLASHINFER_TOPK_ALGO"] = "filtered" try: # ... existing test logic ... finally: os.environ.pop("FLASHINFER_TOPK_ALGO", None)However, the existing test should catch the bug regardless of algorithm selection, so this is optional.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/utils/test_topk.py
🧰 Additional context used
📓 Path-based instructions (1)
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Test implementations should useflashinfer.utilsfunctions (get_compute_capability,is_sm90a_supported,is_sm100a_supported, etc.) to skip tests on unsupported GPU architectures
For testing withmpirunon multi-GPU systems, use the pattern:mpirun -np <num_gpus> pytest tests/path/to/test.py::test_function
Avoid OOM (out-of-memory) errors in tests by using appropriate problem sizes -tests/conftest.pyprovides auto-skipping for OOM tests as a safety net but should not be relied upon
Files:
tests/utils/test_topk.py
🧬 Code graph analysis (1)
tests/utils/test_topk.py (1)
flashinfer/topk.py (1)
top_k_ragged_transform(348-421)
🔇 Additional comments (1)
tests/utils/test_topk.py (1)
794-819: Well-designed test for the bounds-checking bug.The test implementation effectively validates the fix:
- Creates unaligned lengths (
max_len - length_offset) to trigger the vectorized boundary condition- Properly validates that all returned indices fall within
[offset, offset + length)- Clear assertions with helpful error messages
|
Hi @murphymatt thanks for working on this, looking like a duplicate of #2308 ? |
awesome, didn't see this before |
2 participants
📌 Description
There's currently a minor bug in the filtered topk unified kernel affecting some edge cases (missing bounds checking in vectorized loops).
We currently loop over the length axis and perform vectorized processing over chunks. However, when length % VEC_SIZE != 0, the vectorized load loops can process elements beyond the valid range, and it is possible for out-of-bound indices to appear in the output.
Simple fix is to apply bounds checking in the affected locations.
🔍 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
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.