add tg-mxfp4-moe-test #22540
add tg-mxfp4-moe-test #22540
Conversation
Signed-off-by: siyuanf <siyuanf@nvidia.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
![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 adds a new test for trtllm-gen mxfp4 fused MoE kernels. The changes include reference implementations for dequantization and MoE, and a new test function that compares the kernel's output against the reference. My review focuses on correctness and maintainability. I've identified two high-severity issues: one related to potential side effects from in-place modification of function arguments, which could make tests fragile, and another concerning the inclusion of tests that are known to fail, which can disrupt CI pipelines. I've provided suggestions to address these points.
tests/kernels/moe/test_mxfp4_moe.py
Outdated
| router_logits, | ||
| topk, | ||
| num_experts, | ||
| intermediate_size, | ||
| hidden_size, | ||
| hidden_states, | ||
| hidden_states_scale, | ||
| w13_weight, | ||
| w13_weight_scale, | ||
| w2_weight, | ||
| w2_weight_scale, | ||
| act_type, | ||
| ) -> torch.Tensor: |
There was a problem hiding this comment.
This function modifies its input tensors (w13_weight, w13_weight_scale, w2_weight, w2_weight_scale) in-place. This is a side effect that can lead to unexpected behavior and makes the code fragile. For instance, in test_trtllm_gen_mxfp4_fused_moe, the test passes only because the reference implementation is called before this function. If the call order were reversed, the test would fail.
To prevent these side effects and improve the function's robustness, it's best to work on copies of the input tensors by cloning them at the beginning of the function. For example:
def tg_mxfp4_moe(...):
w13_weight = w13_weight.clone()
w13_weight_scale = w13_weight_scale.clone()
w2_weight = w2_weight.clone()
w2_weight_scale = w2_weight_scale.clone()
...
tests/kernels/moe/test_mxfp4_moe.py
Outdated
| check_accuracy(ref_result, tg_result, atol=0, rtol=0.35, percent=0.9) |
There was a problem hiding this comment.
The comment indicates that some test configurations are known to fail due to accuracy issues. Tests that are committed to the repository should pass. If some configurations are expected to fail, they should be marked with pytest.mark.xfail with a clear reason. This prevents CI from being blocked and clearly documents known issues. Please either fix the accuracy issues or mark the failing test cases appropriately.
|
Thanks for the kernel tests! LGTM just a few comments |
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Essential Elements of an Effective PR Description Checklist
Add unit tests for mxfp4 x mxfp8 and mxfp4 x bf16 trtllm-gen moe.
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating
supported_models.mdandexamplesfor a new model.Purpose
Test Plan
Test Result
(Optional) Documentation Update