SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

[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 -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](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