[NPU][diffusion]: support parallel decoding of qwen-image

[gxxx-hum]

Motivation

Parallel VAE decoding reduces peak memory for Qwen-Image and improves memory stability on large-resolution image generation workloads.

In our comparison against main, resolutions above 1024x1024 have a high chance of running into OOM on the baseline path. With this change, peak memory is reduced by about 5% in our tests, while end-to-end performance remains effectively unchanged.

Modifications

• Qwen-Image VAE config
  • Enable parallel VAE decoding by default for the Qwen-Image VAE configuration so that autoencoder_kl_qwenimage uses the ParallelTiledVAE path when SP world size > 1.
• Qwen-Image VAE runtime
  • Wire Qwen-Image VAE decoding to the existing parallel_tiled_decode implementation when use_parallel_decode is enabled and sequence-parallel world size is greater than 1.
  • Keep the non-parallel _decode path unchanged and still available as a fallback for environments without SP or when parallel decode is explicitly disabled.
• Gather decoded tile payloads on rank 0 only, merge once, and broadcast the final decoded output to all ranks.
• Vectorize blend_v, blend_h to replace Python loops with tensor operations and reduce host-side overhead.

Accuracy Tests

startup script:

nohup sglang serve \
--model-path Qwen-Image-2512 \
--num-gpus 2 \
--sp-degree 2 \
--port 8080 \
--host 0.0.0.0 > sglang.log 2>&1 & 

test script:

import requests

response=requests.post(
   "http://localhost:8080/v1/images/generations",
   json={
      "prompt": "a black and white cat wearing a princess tiara",
      "size": "1024x1024",
      "num_inference_steps": 50,
      "response_format": "b64_json",
      "seed": 42,
      "n":1
   }
)

The current version would have caused an OOM (Out of Memory) problem，the error log:

[DecodingStage] Error during execution after 11612.2197 ms: The Inner error is reported as above. The process exits for this inner error, and the current working operator name is Conv2D.
Since the operator is called asynchronously, the stacktrace may be inaccurate. If you want to get the accurate stacktrace, please set the environment variable ASCEND_LAUNCH_BLOCKING=1.
Note: ASCEND_LAUNCH_BLOCKING=1 will force ops to run in synchronous mode, resulting in performance degradation. Please unset ASCEND_LAUNCH_BLOCKING in time after debugging.
[ERROR] 2026-03-16-03:20:06 (PID:97494, Device:0, RankID:0) ERR00100 PTA call acl api failed.
[PID: 97494] 2026-03-16-03:20:04.278.899 Memory_Allocation_Failure(EL0004): Failed to allocate memory requested by GE module.
        Possible Cause: Available memory is insufficient.
        Solution: Close applications not in use.
        TraceBack (most recent call last):
        Call rtMalloc fail, purpose: page caching, type = 2, size:805371904, device_id:0[FUNC:Alloc][FILE:device_allocator.cc][LINE:72]
        Call rtMalloc fail, purpose: page caching, type = 2, size:402718720, device_id:0[FUNC:Alloc][FILE:device_allocator.cc][LINE:72]
        [Exec][Op]Execute op failed. op type = Conv2D, ge result = 4294967295[FUNC:ReportCallError][FILE:acl_op_executor_log_inner.cpp][LINE:124]]

after vae parallel decoding adaptation:

Benchmarking and Profiling

Since 1024x1024 images will cause OOM, 768x768 images were chosen for performance testing and comparison.

	decoding time	peak gpu memory(GB)	peak allocated memory(GB)	memory pool overhead(GB)	remaining gpu memory at peak(GB)
main	0.1101s	58.39	43.22	15.17	2.57
pr	0.2196	54.57	42.62	11.95	6.39

Parallel VAE decoding cuts peak GPU memory by 6.5% (with remaining headroom up by about 2.5× and memory‑pool overhead down by 21.2%), while increasing decode latency modestly from 0.1101 s to 0.2196 s.

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, 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 addresses critical NPU Out-Of-Memory (OOM) issues encountered during high-resolution Qwen-Image generation by introducing and enabling parallel VAE decoding. While this change slightly increases VAE decode latency for single requests, it significantly enhances memory stability and improves performance under higher concurrency, ensuring more robust image generation.

Highlights

• NPU OOM Mitigation: Enabled parallel VAE decoding by default for Qwen-Image to mitigate NPU Out-Of-Memory (OOM) issues during 1024x1024 image generation.
• Parallel Decoding Integration: Integrated the Qwen-Image VAE decoding path with the existing parallel_tiled_decode implementation when parallel decoding is enabled and sequence-parallel world size is greater than 1.
• Fallback Mechanism: Maintained the original non-parallel decoding path as a fallback for environments without sequence parallelism or when parallel decoding is explicitly disabled.

🧠 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/configs/models/vaes/qwenimage.py
  • Added use_parallel_decode configuration option and set its default to True.
• python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py
  • Imported get_sp_world_size for distributed environment checks.
  • Initialized self.use_parallel_decode based on the VAE configuration.
  • Introduced _decode_with_parallel_dispatch to conditionally route decoding to ParallelTiledVAE.parallel_tiled_decode if parallel decoding is enabled and sp_world_size > 1.
  • Refactored _decode and tiled_decode methods to consistently return torch.Tensor directly, simplifying their return types.
  • Updated the main decode method to utilize the new _decode_with_parallel_dispatch for handling both sliced and non-sliced inputs.

Activity

• Provided a motivation detailing NPU OOMs on 1024x1024 Qwen-Image generation and the acceptable latency trade-off for improved memory stability.
• Included accuracy tests with a startup script and a test script demonstrating successful image generation after the changes, resolving a previously observed OOM error.
• Presented benchmarking results comparing decoding time and peak memory usage before and after the changes, showing a reduction in peak memory with a slight increase in decoding time.

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 support for parallel VAE decoding for Qwen-Image models to mitigate out-of-memory errors on NPUs. The changes are well-implemented, enabling parallel decoding by default when sequence parallelism is active. A new dispatch method, _decode_with_parallel_dispatch, cleanly routes decoding to either the new parallel path or the existing standard path.

A significant part of this PR is the refactoring of the _decode and tiled_decode methods to always return a tensor, removing the return_dict parameter. This not only simplifies the code but also fixes an existing issue where the method signature was not compatible with the base class, improving correctness and maintainability.

I have one suggestion to further improve the code by using super() for calling the base class method, which is a more standard and robust approach.

[gxxx-hum]

I can go in two directions here and would appreciate your preference:

A) Tensor-only API (align with ParallelTiledVAE / wanvae.py)

• Make AutoencoderKLQwenImage.encode/decode/forward return only the posterior / torch.Tensor
• Remove return_dict, DecoderOutput, and tuple-style returns from the internal tiled/parallel helpers

B) Keep the legacy return_dict signature for compatibility

• Keep return_dict (even if unused) and/or keep DecoderOutput wrapping

Repo-wide search didn’t find any internal callers passing return_dict to the QwenImage VAE, so A should be safe inside this codebase, and I’ll also run thorough tests to validate the change. But I’m happy to follow the project’s compatibility preference.
@zhuyijie88

[gxxx-hum]

/tag-and-rerun-ci

[gxxx-hum]

I have already fixed the lint and addressed the review comments. Could a maintainer help trigger the CI ？🫶🏼 @ping1jing2

[ping1jing2]

/tag-and-rerun-ci

[gxxx-hum]

/rerun-failed-ci

[gxxx-hum]

/rerun-failed-ci

[mickqian]

looks quite the same with ParallelTiledVAE.parallel_tiled_decode, AutoencoderKLQwenImage is already a subclass of ParallelTiledVAE, why do we need a new method here?

[gxxx-hum]

It is not the same as the base，the base path uses all-gather and merges on every rank, while qwen-Image gathers to rank 0, merges once on rank 0, and then broadcasts the final result.
I scoped it to the Qwen-Image subclass to avoid changing the shared base behavior for other subclasses, especially since Wan may be affected.Also, on NPU I did observe a noticeable peak-memory reduction.
I can first move it back to the base class and test the impact. If it does affect other subclasses, I’ll refactor the shared front half and abstract the collect/merge part to reduce duplicated code.

[mickqian]

we should make the new approach a method of ParallelTiledVAE, providing an alternative option.

How does these two approaches differ, in terms of applicable scenarios, performance?

[gxxx-hum]

I see, thanks. You mean this should be a selectable strategy in ParallelTiledVAE, rather than a Qwen-Image-specific override. I can further make it a base-level alternative option while keeping the default behavior unchanged.

This new path seems to be a better fit for image-generation scenarios. I’ll also run more comparative experiments to better characterize the differences and refine this part.

[gxxx-hum]

/rerun-failed-ci

[gxxx-hum]

/rerun-failed-ci

[ping1jing2]

/rerun-failed-ci

[ping1jing2]

/rerun-failed-ci

[ping1jing2]

/rerun-failed-ci

[gxxx-hum]

@mickqian I've addressed the comments and updated the pr, PTAL, thanks !

all the comments have been resolved and CI passed

[mickqian]

if x.is_cuda and not current_platform.is_hip(): this pattern is appearing repeatedly recently, please move the implementation to platform.py, including all the changes recently

[gxxx-hum]

I’ll move it under runtime/platforms and clean up the recent similar cases there as well.