[Test][Triton] Add RMS norm kernel unit tests - #14700
Conversation
Summary of ChangesHello, 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 introduces unit test coverage for the triton_q_rms operator wrapper. By utilizing a mock kernel approach, the changes validate that the operator correctly interfaces with the underlying Triton kernel, properly calculates launch grids based on hardware properties, and enforces necessary constraints on input dimensions. Highlights
New Features🧠 You can now enable Memory (public preview) 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 the 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 counterproductive. 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. 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
Suggested PR Title:
[Ops][Test][Feature] Add unit tests for Triton RMS Norm operatorSuggested PR Summary:
### What this PR does / why we need it?
This pull request adds unit tests for the Triton RMS Norm operator (`triton_q_rms`) in `tests/ut/ops/test_triton_rms_norm.py`. It verifies that the kernel is launched with the expected arguments across various shapes and device configurations, and checks that unsupported dimensions (greater than 2048) are correctly rejected.
Feedback:
- An issue was identified where directly monkeypatching `rms_norm.triton.runtime.driver.active.utils` can raise an `AttributeError` in environments where Triton's active driver is not initialized (such as CPU-only test environments). It is recommended to mock the `active` attribute on `triton.runtime.driver` instead.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
This PR introduces new unit tests (`test_triton_q_rms_launches_kernel_with_expected_args` and `test_triton_q_rms_rejects_unsupported_dim`) to verify the correctness of the Triton RMS Norm operator.| monkeypatch.setattr( | ||
| rms_norm.triton.runtime.driver.active.utils, | ||
| "get_device_properties", | ||
| lambda device: {"num_vectorcore": num_vectorcore}, | ||
| ) |
There was a problem hiding this comment.
Attempting to monkeypatch rms_norm.triton.runtime.driver.active.utils directly will raise an AttributeError on environments where Triton's active driver is not initialized (such as CPU-only test environments or local development machines), because triton.runtime.driver.active is None. To make the unit tests robust and runnable on CPU-only environments, mock the active attribute itself on triton.runtime.driver instead.
| monkeypatch.setattr( | |
| rms_norm.triton.runtime.driver.active.utils, | |
| "get_device_properties", | |
| lambda device: {"num_vectorcore": num_vectorcore}, | |
| ) | |
| class MockActive: | |
| class utils: | |
| @staticmethod | |
| def get_device_properties(device): | |
| return {"num_vectorcore": num_vectorcore} | |
| monkeypatch.setattr(rms_norm.triton.runtime.driver, "active", MockActive) |
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
Add unit coverage for triton_q_rms wrapper launch behavior, block size selection, output shape preservation, and unsupported hidden dimension handling. Signed-off-by: GDzhu01 <116337067+GDzhu01@users.noreply.github.com>
f2a09a1 to
63d82fd
Compare
|
CI gate is blocked because selected tests require a readiness label. Could a maintainer please add |
1 participant
What this PR does / why we need it?
Adds unit test coverage for
triton_q_rms, the Python wrapper aroundtriton_rms_kernel.The new tests verify:
total_batch, hidden dimension, row stride, epsilon, and output tensor shape are passed correctlyBLOCK_Mselection covers representative batch-per-core casestriton_q_rmsis used by the DSA Q RMS norm paths throughDeviceOperator.apply_dsa_q_rms, and is also warmed up byrms_triton_warmup.Does this PR introduce any user-facing change?
No.
How was this patch tested?
tests/ut/ops/test_triton_rms_norm.pypython -m py_compile tests/ut/ops/test_triton_rms_norm.pyCould not run pytest locally because this Windows PC Python environment does not have
pytestinstalled, and this workspace targets Ascend/NPU execution for hardware-specific tests.