Sage v2 improvements

[ksikiric]

When the first version of Sage v2 was merged to AITER, there was an assumption that the input tensors would be contiguous. This was not the case when called from xDiT, hence the tensors needed to be made contiguous every time the kernel was called, causing a small bottleneck.

I further optimization is the Hardamard matrix, which previously was created every time the kernel was called (on the AITER side). In this change this matrix is now created once, and reused for all kernel calls.

These two changes above improve E2E performance by up to 0.5s or close to 1% on MI355X.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 significant performance enhancements to the Sage v2 attention mechanism within the xfuser framework. By addressing inefficiencies related to tensor contiguity and Hadamard matrix generation, the changes aim to reduce computational overhead and improve overall execution speed, particularly for specific hardware configurations.

Highlights

• Tensor Contiguity Optimization: Eliminated the need for input tensors to be explicitly made contiguous before calling Sage v2, addressing a bottleneck when called from xDiT.
• Hadamard Matrix Pre-computation: Optimized the creation of the Hadamard matrix, which is now generated once and reused across kernel calls, improving efficiency.
• Performance Improvement: Achieved end-to-end performance improvements of up to 0.5 seconds or nearly 1% on MI355X hardware.

Changelog

• xfuser/core/distributed/attention_backend.py
  • Introduced logic to read and validate the FAV3_SAGE_MXFP4_BLOCK_R environment variable.
  • Implemented pre-computation and device replication of the Hadamard matrix for reuse.
  • Added a conditional check to bypass explicit tensor contiguity operations if the fav3_sage_mxfp4_wrapper supports the BLOCK_R parameter, indicating an internal fix for contiguity.
  • Modified the _aiter_sage_v2_attn_call function to pass the pre-computed Hadamard matrix to the attention wrapper.
• xfuser/envs.py
  • Added XFUSER_FAV3_SAGE_MXFP4_BLOCK_R to the list of managed environment variables, with a default value of "128".

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize 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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. 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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces performance improvements for Sage v2 by pre-calculating the Hadamard matrix and making input tensors contiguous only when necessary. The changes are logical and well-aligned with the goal of reducing bottlenecks. My review includes a few suggestions to enhance code clarity and robustness, such as refactoring an environment variable parsing logic and using more specific exception handling.

[avjves]

I'm not sure if this sort of gating is the best way to go 🤔 They work for now, but we should aim to replace these ASAP. Latest AITER release is from Dec 2025, so these are not clearly in the latest release for now. (Maybe some other features are there already and could then be gated through versions instead of signature inspection etc?)

[ksikiric]

Yes I agree