[diffusion]: Enable torch.compile for UlyssesAttention

[Ratish1]

Motivation

I noticed that the `UlyssesAttention.forward` and `UlyssesAttention_VSA.forward` methods were decorated with `@torch.compiler.disable`. This prevents `torch.compile` from optimizing the attention layer, which is a significant part of the model's computation.

Modifications

• Removed @torch.compiler.disable from UlyssesAttention.forward inpython/sglang/multimodal_gen/runtime/layers/attention/layer.py.
• Removed @torch.compiler.disable from UlyssesAttention_VSA.forward in python/sglang/multimodal_gen/runtime/layers/attention/layer.py.

Accuracy Tests

I tested this change on HunyuanVideo
sglang generate --model-path hunyuanvideo-community/HunyuanVideo --log-level info --prompt "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." --negative-prompt ' ' --720p --num-inference-steps 40 --num-frames 81 --guidance-scale 5.0 --seed 42 --save-output --num-gpus 4 --ulysses-degree 4 --dit-layerwise-offload true --dit-cpu-offload false --vae-cpu-offload false --text-encoder-cpu-offload true --warmup --enable-torch-compile true

Performance Comparison Report

1. High-level Summary

Metric	Baseline	New	Diff
E2E Latency	251454.00 ms	239544.09 ms	-11909.91 ms (-4.7%)
Throughput	0.00 req/s	0.00 req/s	-

2. Stage Breakdown

Stage Name	Baseline (ms)	New (ms)	Diff (ms)	Diff (%)
InputValidationStage	0.07	0.06	-0.01	-11.1%
TextEncodingStage	1110.32	1121.54	+11.22	+1.0%
ConditioningStage	0.01	0.02	+0.00	+11.6%
TimestepPreparationStage	0.43	0.41	-0.01	-3.1%
LatentPreparationStage	0.18	0.15	-0.02	-14.0%
DenoisingStage	235593.86	223559.39	-12034.47	-5.1%
DecodingStage	14746.33	14860.19	+113.86	+0.8%

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello @Ratish1, 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 aims to improve the performance of the attention layers within the model by allowing PyTorch's torch.compile feature to optimize these critical components. By removing specific decorators that previously disabled compilation, the change enables significant latency reductions, particularly in the denoising stage, thereby enhancing overall model efficiency.

Highlights

• Performance Optimization: Enabled torch.compile for UlyssesAttention.forward and UlyssesAttention_VSA.forward methods by removing the @torch.compiler.disable decorator.
• Latency Reduction: Achieved a 4.7% reduction in End-to-End Latency and a 5.1% reduction in DenoisingStage latency, as demonstrated by performance comparison reports.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• python/sglang/multimodal_gen/runtime/layers/attention/layer.py
  • Removed @torch.compiler.disable from UlyssesAttention.forward.
  • Removed @torch.compiler.disable from UlyssesAttention_VSA.forward.

Activity

• The pull request was opened by Ratish1.
• No human activity (comments, reviews, etc.) has occurred on this pull request yet.

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 enables torch.compile for UlyssesAttention and UlyssesAttention_VSA by removing the @torch.compiler.disable decorators. This is a valuable optimization that allows PyTorch's compiler to improve the performance of these attention layers. The included benchmark results clearly show a reduction in end-to-end latency, validating the effectiveness of this change. The code modification is straightforward and well-justified. This is a great improvement.

[mickqian]

/tag-and-rerun-ci

[BBuf]

Can you give a benchmark on wan2.2 model, thanks.

[Ratish1]

@BBuf, the performance didnt change for Wan-AI/Wan2.2-T2V-A14B-Diffusers, I tested it.

[BBuf]

/tag-and-rerun-ci