[Diffusion] Fix Sana corrupted output by removing spurious QK norm layers by ChangyiYang · Pull Request #20656 · sgl-project/sglang

ChangyiYang · 2026-03-16T04:41:47Z

Problem

The native SGLang Sana pipeline produced corrupted/garbled images. The output looked like a mess of random shapes with no semantic content.

Root Cause

SanaLinearAttention and SanaCrossAttention both had norm_q / norm_k RMSNorm layers that do not exist in the Sana checkpoint. These layers were initialized with random weights and applied to query/key tensors before attention computation, completely corrupting the attention outputs.

Verified by inspecting the checkpoint: no norm_q or norm_k keys exist in diffusion_pytorch_model.safetensors for Efficient-Large-Model/Sana_600M_1024px_diffusers.

Fix

Remove norm_q and norm_k from SanaLinearAttention and SanaCrossAttention to match the actual Sana model architecture.

Test

Ran pytest python/sglang/multimodal_gen/test/server/test_server_a.py -k sana_image_t2i — PASSED.

🤖 Generated with Claude Code

…yers SanaLinearAttention and SanaCrossAttention had norm_q/norm_k RMSNorm layers that are absent from the Sana checkpoint weights. These layers ran with random initialization, corrupting attention outputs and producing garbled images. Remove norm_q/norm_k from both attention classes to match the actual Sana model architecture. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

gemini-code-assist · 2026-03-16T04:42:03Z

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 resolves a critical bug in the SGLang Sana pipeline that caused corrupted image generation. The problem was traced to the erroneous inclusion of RMSNorm layers (norm_q and norm_k) within the SanaLinearAttention and SanaCrossAttention modules, which were not part of the official Sana model checkpoint. By removing these superfluous normalization layers, the model now correctly produces high-quality images, ensuring the proper functionality of the Sana pipeline.

Highlights

Corrupted Sana Output: The native SGLang Sana pipeline produced garbled images, appearing as a mess of random shapes with no semantic content.
Root Cause Identified: The issue stemmed from SanaLinearAttention and SanaCrossAttention incorrectly including norm_q and norm_k RMSNorm layers that do not exist in the actual Sana checkpoint, leading to corrupted attention outputs.
Fix Implemented: The norm_q and norm_k layers were removed from SanaLinearAttention and SanaCrossAttention to align with the correct Sana model architecture.
Verification and Testing: The fix was verified by running pytest python/sglang/multimodal_gen/test/server/test_server_a.py -k sana_image_t2i, which now passes with correct, high-quality image output (e.g., Doraemon eating dorayaki), contrasting with the garbled images produced before the fix.

🧠 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/models/dits/sana.py
- Removed the qk_norm_dim parameter from the SanaLinearAttention constructor.
- Removed the initialization of self.norm_q and self.norm_k in SanaLinearAttention.
- Removed the application of self.norm_q and self.norm_k within the forward method of SanaLinearAttention.
- Removed the initialization of self.norm_q and self.norm_k in SanaCrossAttention.
- Removed the application of self.norm_q and self.norm_k within the forward method of SanaCrossAttention.
- Removed the qk_norm_dim argument during the instantiation of SanaLinearAttention.

Activity

No human activity 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.

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 correctly addresses the issue of corrupted image output from the Sana model by removing the spurious norm_q and norm_k RMSNorm layers from SanaLinearAttention and SanaCrossAttention. The changes are consistent with the pull request description, which states these layers do not exist in the model checkpoint. The modifications are confined to deleting the incorrect layer initializations and their applications, which is the right approach to fix this bug. The code is now aligned with the actual model architecture, and I find no further issues with the changes.

mickqian

...amazed by these modern art

mickqian · 2026-03-19T02:42:33Z

/tag-and-rerun-ci

yhyang201 · 2026-03-19T16:10:13Z

/rerun-failed-ci

yhyang201 · 2026-03-19T23:47:36Z

/rerun-failed-ci

yhyang201 · 2026-03-20T04:43:53Z

/rerun-failed-ci

yhyang201 · 2026-03-20T14:43:47Z

/rerun-failed-ci

yhyang201 · 2026-03-20T23:16:56Z

/rerun-failed-ci

yhyang201 · 2026-03-21T04:29:19Z

/rerun-failed-ci

…rm layers (sgl-project#20656) Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Mick <mickjagger19@icloud.com>

…rm layers (sgl-project#20656) Co-authored-by: Mick <mickjagger19@icloud.com>

ChangyiYang requested review from BBuf, mickqian, ping1jing2, yhyang201 and yingluosanqian as code owners March 16, 2026 04:41

github-actions Bot added the diffusion SGLang Diffusion label Mar 16, 2026

gemini-code-assist Bot reviewed Mar 16, 2026

View reviewed changes

ChangyiYang added 2 commits March 15, 2026 21:45

add sana before/after images for PR

de29631

add sana before/after images for PR

6a0a6ab

github-actions Bot added the documentation Improvements or additions to documentation label Mar 16, 2026

mickqian reviewed Mar 16, 2026

View reviewed changes

Comment thread docs/sana-after.jpg Outdated

ChangyiYang added 3 commits March 18, 2026 16:05

Merge branch 'main' into sana-fix-remove-spurious-qk-norm

35910a2

remove sana before/after images

eb80634

remove sana before/after images

e8a0c94

github-actions Bot added the run-ci label Mar 19, 2026

Merge branch 'main' into sana-fix-remove-spurious-qk-norm

261d3a9

mickqian merged commit c1794e2 into sgl-project:main Mar 22, 2026
70 of 73 checks passed

yhyang201 pushed a commit to yhyang201/sglang that referenced this pull request Apr 22, 2026

[diffusion] fix: fix Sana corrupted output by removing spurious QK no…

98742dd

…rm layers (sgl-project#20656) Co-authored-by: Mick <mickjagger19@icloud.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Diffusion] Fix Sana corrupted output by removing spurious QK norm layers#20656

[Diffusion] Fix Sana corrupted output by removing spurious QK norm layers#20656
mickqian merged 7 commits intosgl-project:mainfrom
ChangyiYang:sana-fix-remove-spurious-qk-norm

ChangyiYang commented Mar 16, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Mar 16, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

mickqian left a comment •

edited

Loading

Uh oh!

Uh oh!

mickqian commented Mar 19, 2026

Uh oh!

yhyang201 commented Mar 19, 2026

Uh oh!

yhyang201 commented Mar 19, 2026

Uh oh!

yhyang201 commented Mar 20, 2026

Uh oh!

yhyang201 commented Mar 20, 2026

Uh oh!

yhyang201 commented Mar 20, 2026

Uh oh!

yhyang201 commented Mar 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

ChangyiYang commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Root Cause

Fix

Test

Uh oh!

gemini-code-assist Bot commented Mar 16, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

mickqian left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mickqian commented Mar 19, 2026

Uh oh!

yhyang201 commented Mar 19, 2026

Uh oh!

yhyang201 commented Mar 19, 2026

Uh oh!

yhyang201 commented Mar 20, 2026

Uh oh!

yhyang201 commented Mar 20, 2026

Uh oh!

yhyang201 commented Mar 20, 2026

Uh oh!

yhyang201 commented Mar 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ChangyiYang commented Mar 16, 2026 •

edited

Loading

mickqian left a comment •

edited

Loading