[feat] Add gemma RMS AR fusion (#3322)
<!-- .github/pull_request_template.md --> ## ๐ Description Verified from framework side. See perf at https://github.com/vllm-project/vllm/pull/42646 Add weight bias to RMS norm AR fusion to support gemma and qwen3.5 RMS ## ๐ 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 - [x] Tests have been added or updated as needed. - [x] 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 * **New Features** * Added a GPU benchmark to compare fused vs unfused AllReduce+RMSNorm performance. * Introduced an optional weight_bias parameter to AllReduce-fusion APIs to support alternate RMSNorm scaling. * **Tests** * Added a distributed correctness test for Gemma/Qwen3.5-style RMSNorm AllReduce fusion. * Extended AllReduce fusion tests to validate weight_bias variants (e.g., 0.0 and 1.0). * **Documentation** * Updated reference traces and docstrings to describe weight_bias behavior. <!-- review_stack_entry_start --> [](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3322?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack) <!-- review_stack_entry_end --> <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
J
Jiahan Chang (Cyrus) committed
7b9e054db2dd825a460b6c2ea63f3bc57cbf41b1
Parent: 9c3eb00
Committed by GitHub <noreply@github.com>
on 5/19/2026, 4:13:18 PM