Add gqa_max_head_dim feature flag for configurable GQA limit - #241
Closed
justinchuby wants to merge 5 commits into
Closed
Add gqa_max_head_dim feature flag for configurable GQA limit#241justinchuby wants to merge 5 commits into
justinchuby wants to merge 5 commits into
Conversation
Make the Attention→GroupQueryAttention rewrite rule's head_dim limit configurable via the gqa_max_head_dim flag (default: 512). Previously the limit was hardcoded as _MAX_GQA_HEAD_DIM = 256, which prevented GQA fusion for Gemma4 global attention layers (head_dim=512). The new default of 512 matches newer ORT builds. Users targeting older ORT builds can restore the limit with: export MOBIUS_GQA_MAX_HEAD_DIM=256 or programmatically: from mobius import override_flags with override_flags(gqa_max_head_dim=256): ... Changes: - _flags.py: Add _env_int helper, gqa_max_head_dim flag, update override_flags type hint to accept int - _group_query_attention.py: Replace _MAX_GQA_HEAD_DIM constant with flags.gqa_max_head_dim - _flags_test.py: 6 new tests covering default, env var, override, and integration with the GQA rewrite rule 2674 tests pass, lint clean. Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR makes the Attention→GroupQueryAttention rewrite rule’s head_dim cutoff configurable via a runtime feature flag (flags.gqa_max_head_dim) instead of using a hardcoded constant, enabling GQA fusion for models with larger head_dim (e.g., 512) while still allowing users to opt back into the legacy 256 limit via env var or override_flags().
Changes:
- Add
_env_int()and introducegqa_max_head_dimflag (default 512), plus widenoverride_flags()typing to allow integer overrides. - Update the GQA rewrite rule to use
flags.gqa_max_head_dimfor the head-dim guard and for the corresponding failure message. - Add unit tests covering default/env/override behavior and verifying the GQA head-dim check respects the flag.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
src/mobius/rewrite_rules/_group_query_attention.py |
Switches the GQA head-dim limit check to use flags.gqa_max_head_dim instead of a constant. |
src/mobius/_flags.py |
Adds integer env parsing and introduces the gqa_max_head_dim feature flag (default 512). |
src/mobius/_flags_test.py |
Adds tests for the new integer flag and verifies the GQA rewrite guard respects it. |
- Update inline comments from hardcoded '(256)' to reference gqa_max_head_dim flag - Update docs/_generate_flags_docs.py to recognize _env_int() calls alongside _env_bool() so generated docs show correct defaults - Regenerate docs/feature-flags.md with gqa_max_head_dim entry Signed-off-by: Justin Chu <justinchu@microsoft.com>
docs/feature-flags.md is auto-generated by docs/_generate_flags_docs.py and should not be committed. Add it to .gitignore. Signed-off-by: Justin Chu <justinchu@microsoft.com>
Keep the conservative default (256) matching older ORT builds. Users targeting newer ORT can set MOBIUS_GQA_MAX_HEAD_DIM=512. Signed-off-by: Justin Chu <justinchu@microsoft.com>
| """ | ||
|
|
||
| gqa_max_head_dim: int = dataclasses.field( | ||
| default_factory=lambda: _env_int("MOBIUS_GQA_MAX_HEAD_DIM", 256) |
| (ORT ≤1.24.x workaround). | ||
| * - ``gqa_max_head_dim`` | ||
| - ``MOBIUS_GQA_MAX_HEAD_DIM`` | ||
| - ``256`` |
Comment on lines
+116
to
+120
| def test_default_is_512(self, monkeypatch): | ||
| monkeypatch.delenv("MOBIUS_GQA_MAX_HEAD_DIM", raising=False) | ||
| f = _flags._Flags() | ||
| assert f.gqa_max_head_dim == 256 | ||
|
|
Comment on lines
+171
to
+176
| # Skip when head_dim exceeds gqa_max_head_dim flag (default 256). | ||
| hd = _head_dim_exceeds_gqa_limit(past_key) | ||
| if hd is not None: | ||
| return result.fail(f"head_dim={hd} exceeds GQA MAX_HEAD_SIZE={_MAX_GQA_HEAD_DIM}") | ||
| return result.fail( | ||
| f"head_dim={hd} exceeds GQA MAX_HEAD_SIZE={flags.gqa_max_head_dim}" | ||
| ) |
Comment on lines
+55
to
+57
| # The limit is now configurable via the ``gqa_max_head_dim`` flag (default 256 | ||
| # for newer ORT builds). Set ``MOBIUS_GQA_MAX_HEAD_DIM=256`` to restore the | ||
| # old limit for ORT ≤1.24. |
titaiwangms
reviewed
May 4, 2026
| """Tests for the gqa_max_head_dim integer flag.""" | ||
|
|
||
| def test_default_is_512(self, monkeypatch): | ||
| monkeypatch.delenv("MOBIUS_GQA_MAX_HEAD_DIM", raising=False) |
Contributor
There was a problem hiding this comment.
We set it as ENV? Why not just hard code it to 512 since it's a "max" value anyway.
The feature flag is unnecessary — old ORT versions have other CUDA issues regardless, and CPU EP is unaffected by head_dim limits. Simply update the hardcoded limit from 256 to 512 to support Gemma4 global attention (head_dim=512) with newer ORT builds. Removes: gqa_max_head_dim flag, _env_int helper, associated tests. Signed-off-by: Justin Chu <justinchu@microsoft.com>
Member
Author
|
Closing this PR. After further discussion, the feature flag is not needed — we will hardcode MAX_HEAD_DIM=512 since older ORT versions have other CUDA issues that prevent usage anyway, and CPU EP is unaffected by head_dim. The 512 limit is already set in #240. |
37 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Make the Attention→GroupQueryAttention rewrite rule head_dim limit configurable via a feature flag instead of a hardcoded constant.
Problem
The GQA rewrite rule had a hardcoded
_MAX_GQA_HEAD_DIM = 256that prevented GQA fusion for Gemma4 global attention layers (head_dim=512). PR #240 intended to change this to 512 but the value was not actually updated.Solution
Replace the constant with
flags.gqa_max_head_dim(default: 512), configurable via:MOBIUS_GQA_MAX_HEAD_DIM=256override_flags(gqa_max_head_dim=256)Changes
_flags.py: Add_env_inthelper,gqa_max_head_dimflag (default 512), updateoverride_flagstype hint_group_query_attention.py: Replace_MAX_GQA_HEAD_DIMconstant withflags.gqa_max_head_dim_flags_test.py: 6 new tests (default value, env var, override, integration with GQA rewrite)Testing
2674 tests pass, lint clean.