bugfix: fix failed unittest on spark (sm_121) #1951
Conversation
Summary of ChangesHello @yzh119, 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 addresses and resolves three specific unittest failures encountered on Spark environments (sm_121). It implements targeted adjustments to test logic, introduces conditional test skipping based on available hardware resources, and temporarily marks one test as an expected failure due to a known numerical accuracy issue. The overall goal is to enhance test suite stability and reliability on Spark without compromising the integrity of the tests. 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
|
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdded device-SM availability checks and runtime guards to green-context splitting logic and tests; tests now catch/skip on insufficient-SMs runtime errors. Added a compute-capability xfail in a JIT test for SM 12.1, and tightened softmax numeric comparison to include rtol. Changes
Sequence Diagram(s)sequenceDiagram
participant Test as Test Function
participant GreenSplit as green_ctx.split_*
participant DeviceInfo as utils.get_device_sm_count
participant PyTest as pytest.skip/xfail
Test->>GreenSplit: request split / create context (groups/min_count or sm_counts)
GreenSplit->>DeviceInfo: query device SM count
DeviceInfo-->>GreenSplit: available_sms
alt required_sms > available_sms
GreenSplit-->>Test: raise RuntimeError("Insufficient SMs: requires X, have Y")
Test->>PyTest: catch RuntimeError -> pytest.skip(...)
else required_sms <= available_sms
GreenSplit-->>Test: return split contexts / proceed
Test->>Test: run kernels / assertions
end
Note over Test,PyTest: Separate path: test_jit_example checks compute capability -> mark xfail for (12,1)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
![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 addresses three failing unit tests on Spark (sm_121) by adding a guard for SM availability in test_green_ctx.py, marking a test as xfail in test_jit_example.py due to numerical issues, and increasing the tolerance in test_sampling.py. The changes are correct and effectively fix the described issues. I've provided a couple of suggestions for test_green_ctx.py to improve code clarity and reduce duplication.
tests/utils/test_green_ctx.py
Outdated
| total = 0 | ||
| for sm_count in sm_counts: | ||
| rounded = round_up(max(sm_count, min_sm), alignment) | ||
| total += rounded | ||
| return total |
There was a problem hiding this comment.
This for-loop can be expressed more concisely using the built-in sum() function with a generator expression. This is a common Python idiom that improves readability.
| total = 0 | |
| for sm_count in sm_counts: | |
| rounded = round_up(max(sm_count, min_sm), alignment) | |
| total += rounded | |
| return total | |
| return sum(round_up(max(sm_count, min_sm), alignment) for sm_count in sm_counts) |
tests/utils/test_green_ctx.py
Outdated
| ): | ||
| required_sms = calculate_required_sms(num_groups, min_count, device) | ||
| available_sms = get_device_sm_count(torch.device(device)) | ||
| if required_sms > available_sms: |
There was a problem hiding this comment.
Should we move this check in the def split_device_green_ctx API and raise an exception?
There was a problem hiding this comment.
That would also solve gemini's concern with copy pasting the check.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
![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
♻️ Duplicate comments (2)
tests/utils/test_green_ctx.py (2)
20-24: Consider using built-insum()for improved readability.As noted in previous reviews, this for-loop can be expressed more concisely using the built-in
sum()function with a generator expression, which is a common Python idiom.Apply this diff to refactor:
- total = 0 - for sm_count in sm_counts: - rounded = round_up(max(sm_count, min_sm), alignment) - total += rounded - return total + return sum(round_up(max(sm_count, min_sm), alignment) for sm_count in sm_counts)
36-42: Address the pre-commit formatting failure.The pipeline indicates a formatting issue that needs to be resolved. Please run
pre-commit run --all-filesto apply the formatting changes.Additionally, as noted in previous reviews, this pre-check logic is duplicated across multiple tests. Consider either:
- Extracting it into a pytest fixture or helper function
- Moving the check into the
split_device_green_ctxAPI itself to raise an exception
🧹 Nitpick comments (1)
tests/utils/test_green_ctx.py (1)
43-45: Prefix unused variable with underscore.The
streamsvariable is unpacked but never used in this test function. Prefix it with an underscore to indicate it's intentionally unused.Apply this diff:
- streams, resources = green_ctx.split_device_green_ctx( + _streams, resources = green_ctx.split_device_green_ctx( dev, num_groups, min_count )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/utils/test_green_ctx.py(6 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
tests/utils/test_green_ctx.py (2)
flashinfer/utils.py (2)
get_compute_capability(251-254)get_device_sm_count(595-596)flashinfer/green_ctx.py (2)
get_sm_count_constraint(34-44)split_device_green_ctx(126-178)
🪛 GitHub Actions: pre-commit
tests/utils/test_green_ctx.py
[error] 40-40: ruff-format: 1 file reformatted by this hook. The pre-commit hook failed; please re-run with 'pre-commit run --all-files' to apply formatting changes.
[error] 40-40: Code style formatting changed by ruff-format. Updated call should be: streams, resources = green_ctx.split_device_green_ctx(dev, num_groups, min_count).
🪛 Ruff (0.14.1)
tests/utils/test_green_ctx.py
43-43: Unpacked variable streams is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Deploy Docs
🔇 Additional comments (6)
tests/utils/test_green_ctx.py (6)
5-5: LGTM!The imports are necessary for the SM calculation helpers and are correctly placed.
8-13: LGTM!The helper correctly calculates the total SM count required by rounding up the minimum count to meet alignment requirements and multiplying by the number of groups.
61-67: LGTM!The pre-check logic correctly validates SM availability before running the test.
97-103: LGTM!The pre-check correctly uses
calculate_required_sms_by_countsto validate SM availability for tests with specific SM counts.
130-136: LGTM!The pre-check correctly validates SM availability before running the kernel execution test.
165-171: LGTM!The pre-check correctly validates SM availability before running the alignment test.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
tests/utils/test_green_ctx.py (1)
39-46: Consider consistency in device object creation.Unlike
test_green_ctx_creation(line 15), this test passestorch.device(device)directly without creating adevvariable first. While both approaches work, consistent usage across all tests would improve readability.Apply this diff for consistency:
+ dev = torch.device(device) try: streams, resources = green_ctx.split_device_green_ctx( - torch.device(device), num_groups, min_count + dev, num_groups, min_count )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
flashinfer/green_ctx.py(3 hunks)tests/utils/test_green_ctx.py(5 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
flashinfer/green_ctx.py (1)
flashinfer/utils.py (3)
get_compute_capability(251-254)get_device_sm_count(595-596)round_up(589-591)
tests/utils/test_green_ctx.py (1)
flashinfer/green_ctx.py (2)
split_device_green_ctx(126-190)split_device_green_ctx_by_sm_count(193-281)
🪛 Ruff (0.14.1)
flashinfer/green_ctx.py
180-183: Avoid specifying long messages outside the exception class
(TRY003)
264-264: Avoid specifying long messages outside the exception class
(TRY003)
272-275: Avoid specifying long messages outside the exception class
(TRY003)
tests/utils/test_green_ctx.py
17-17: Unpacked variable streams is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Deploy Docs
🔇 Additional comments (5)
tests/utils/test_green_ctx.py (1)
15-23: Good error handling pattern for insufficient SMs.The try-except block properly catches and skips tests when the device lacks sufficient SMs, which addresses the spark (sm_121) test failures mentioned in the PR objectives.
flashinfer/green_ctx.py (4)
31-31: LGTM! Required import for SM count validation.The
get_device_sm_countimport is correctly added and used in both validation checks (lines 177 and 269).
173-184: Excellent early validation for SM availability.The pre-check correctly computes the required SMs and fails fast before any CUDA operations, providing a clear error message that aligns with the test expectations.
261-261: Good optimization: constraint calculation moved outside loop.Moving
get_sm_count_constraintoutside the loop avoids redundant calls, as the constraints don't change between iterations.
267-276: Proper SM validation with informative error message.The validation correctly sums the rounded SM counts and raises a clear error if insufficient. The error message helpfully includes the actual
rounded_sm_countslist to aid debugging.
There was a problem hiding this comment.
I can confirm that test_jit_example.py now passes or xfails.
test_green_ctx.py still has 7 failures:
================================================================================================================================================= short test summary info =================================================================================================================================================
FAILED tests/utils/test_green_ctx.py::test_green_ctx_creation[16-3-cuda:0] - RuntimeError: CUDA error code=914(b'CUDA_ERROR_INVALID_RESOURCE_TYPE')
FAILED tests/utils/test_green_ctx.py::test_green_ctx_kernel_execution[16-3-cuda:0] - RuntimeError: CUDA error code=914(b'CUDA_ERROR_INVALID_RESOURCE_TYPE')
FAILED tests/utils/test_green_ctx.py::test_split_device_green_ctx_by_sm_count_creation[sm_counts0-cuda:0] - RuntimeError: CUDA error code=914(b'CUDA_ERROR_INVALID_RESOURCE_TYPE')
FAILED tests/utils/test_green_ctx.py::test_split_device_green_ctx_by_sm_count_creation[sm_counts1-cuda:0] - RuntimeError: CUDA error code=914(b'CUDA_ERROR_INVALID_RESOURCE_TYPE')
FAILED tests/utils/test_green_ctx.py::test_split_device_green_ctx_by_sm_count_kernel_execution[sm_counts0-cuda:0] - RuntimeError: CUDA error code=914(b'CUDA_ERROR_INVALID_RESOURCE_TYPE')
FAILED tests/utils/test_green_ctx.py::test_split_device_green_ctx_by_sm_count_kernel_execution[sm_counts1-cuda:0] - RuntimeError: CUDA error code=914(b'CUDA_ERROR_INVALID_RESOURCE_TYPE')
FAILED tests/utils/test_green_ctx.py::test_split_device_green_ctx_by_sm_count_alignment[sm_counts1-cuda:0] - RuntimeError: CUDA error code=914(b'CUDA_ERROR_INVALID_RESOURCE_TYPE')
=================================================================================================================================== 7 failed, 10 passed, 5 skipped, 1 warning in 0.91s ====================================================================================================================================
Please see my other comment for test_sampling.py. There might be nans happening from the kernel, at least in my local env
| probs_ref = torch.softmax(logits_scaled, dim=-1) | ||
|
|
||
| assert torch.allclose(probs, probs_ref, atol=1e-5) | ||
| assert torch.allclose(probs, probs_ref, rtol=1e-5, atol=1e-5) |
There was a problem hiding this comment.
I cannot seem to repro the fix in Spark. It also seems like allclose has a default rtol=1e-5 so this may not even effectively make any change.
In fact in my local env (cu130 container), when I change the tolerance and inject print statements as
probs_ref = torch.softmax(logits_scaled, dim=-1)
print(f"{torch.isnan(probs).sum().item() = }")
print(f"{torch.isnan(probs_ref).sum().item() =}")
assert torch.allclose(probs, probs_ref, rtol=100, atol=100)
I am seeing nans.
(py312) root@c661e6d696f6:/flashinfer# pytest tests/utils/test_sampling.py -x -s
=================================================================================================================================================== test session starts ===================================================================================================================================================
platform linux -- Python 3.12.11, pytest-8.4.2, pluggy-1.6.0
rootdir: /flashinfer
configfile: pytest.ini
collected 900 items
tests/utils/test_sampling.py torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
.torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
.torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
.torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
.torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
.torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
.torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
.torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
torch.isnan(probs).sum().item() = 0
torch.isnan(probs_ref).sum().item() =0
.torch.isnan(probs).sum().item() = 4873728
torch.isnan(probs_ref).sum().item() =0
F
======================================================================================================================================================== FAILURES =========================================================================================================================================================
____________________________________________________________________________________________________________________________ test_softmax[True-True-1.0-normal_distribution(std=1)-128256-989] ____________________________________________________________________________________________________________________________
...
> assert torch.allclose(probs, probs_ref, rtol=100, atol=100)
E AssertionError: assert False
E + where False = <built-in method allclose of type object at 0x16bc850>(tensor([[0.0000e+00, 7.8481e-05, 0.0000e+00, ..., 9.0452e-06, 8.5036e-06,\n 0.0000e+00],\n [2.4505e-05, ...05],\n [0.0000e+00, 0.0000e+00, 7.0366e-06, ..., 0.0000e+00, 7.1824e-06,\n 2.0367e-06]], device='cuda:0'), tensor([[0.0000e+00, 7.8481e-05, 0.0000e+00, ..., 9.0452e-06, 8.5036e-06,\n 0.0000e+00],\n [2.4505e-05, ...05],\n [0.0000e+00, 0.0000e+00, 7.0366e-06, ..., 0.0000e+00, 7.1824e-06,\n 2.0367e-06]], device='cuda:0'), rtol=100, atol=100)
E + where <built-in method allclose of type object at 0x16bc850> = torch.allclose
tests/utils/test_sampling.py:76: AssertionError
...
3 participants
📌 Description
There are three failed unittests on spark (sm_121):
First one is because spark has small number of SMs (48) and we don't have a guard on green context splitting.
Second one is an unknown issue (logits don't match with reference) and probably related to barriers on sm_121, xfail now and will fix later.
The last one is because of the reduction size difference, and we should increase tolerance (by adding a rtol).
This PR fixes these issues.
🔍 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