[Feat] support SP for FLUX.2-klein by RuixiangMa · Pull Request #1250 · vllm-project/vllm-omni

RuixiangMa · 2026-02-06T09:47:26Z

Purpose

support SP (Ulysses & Ring) for FLUX.2-klein

Test Plan

Test Result

Target image:

tp = 1 + 4 * nvidia 4090(24G)

curl -s -X POST "http://localhost:8004/v1/images/edits" -F "image=@test.jpg" -F "prompt=Change the sky to orange sunset." -F "guidance_scale=1.0" -F "num_inference_steps=50" -F "n=1" -F "size=1024x1024" -F "output_format=png" | jq -r '.data[0].b64_json' | base64 --decode > output.png

Configuration	Ulysses degree	Ring degree	Generation Time	Speedup
Baseline	1	1	25.503s	1.00x
Ulysses	4	1	13.173s	1.94x
Ring	1	4	16.866s	1.51x
Hybrid Ulysses + Ring	2	2	14.812s	1.72x

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f3436b8532

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

Signed-off-by: Lancer <maruixiang6688@gmail.com>

hsliuustc0106 · 2026-02-06T14:20:46Z

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py


-        image_rotary_emb = self.pos_embed(img_ids)
-        text_rotary_emb = self.pos_embed(txt_ids)
+        if current_omni_platform.is_npu():


@gcanlin do we have better ways to handle this difference? this is so awkward

@wtomin PTAL

I find that there exists npu and mps hardcode in Flux2PosEmbed, which may be from diffusers library I guess. I will take a micro-refactoring PR for removing it. We could delete this npu branch temporarily if we'd like to merge this PR first.

I find that there exists npu and mps hardcode in Flux2PosEmbed, which may be from diffusers library I guess. I will take a micro-refactoring PR for removing it. We could delete this npu branch temporarily if we'd like to merge this PR first.

Got it! I'll just remove this branch right now

I remember some of these is because torch_npu doesn't support complex number

Signed-off-by: Lancer <maruixiang6688@gmail.com>

hsliuustc0106 · 2026-02-07T23:13:14Z

update the docs as well

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa · 2026-02-08T00:56:28Z

update the docs as well

Done

wtomin · 2026-02-09T03:31:19Z

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

        num_txt_tokens = encoder_hidden_states.shape[1]

+        sp_size = self.parallel_config.sequence_parallel_size
+        get_forward_context().sequence_parallel_size = sp_size


To fetch the sp_size, using self.parallel_config.sequence_parallel_size would be sufficient. I don't see why you need to set get_forward_context().sequence_parallel_size here.

I find that LongCatImageTransformer2DModel also edits get_forward_context...

od_config = get_forward_context().omni_diffusion_config
parallel_config = od_config.parallel_config
sequence_parallel_size = parallel_config.sequence_parallel_size

This would be my recommendation.

To fetch the sp_size, using self.parallel_config.sequence_parallel_size would be sufficient. I don't see why you need to set get_forward_context().sequence_parallel_size here.

ths, I keep only sp_size = self.parallel_config.sequence_parallel_size—that matches what you suggested

wtomin · 2026-02-09T03:32:02Z

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

+                    f"rank={sp_rank}, original_shape={original_shape}, chunked_shape={hidden_states.shape}"
+                )
+        else:
+            if not hasattr(self, "_sp_forward_logged"):


Is self._sp_forward_logged used for debugging only?

Is self._sp_forward_logged used for debugging only

yes, I'm removing it now.

Signed-off-by: Lancer <maruixiang6688@gmail.com>

wtomin · 2026-02-10T07:03:16Z

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

+        sp_size = self.parallel_config.sequence_parallel_size
+        sp_pad_size = 0
+        if sp_size > 1:
+            hidden_states, sp_pad_size = sp_shard_with_padding(hidden_states, dim=1)


I am a little concerned about the padding behavior here. It looks like flux2_klein does not support hidden_states_mask (like qwen_image), therefore, the padded tokens will participate in the attention computation, which is wrong.

If using Non-intrusive _sp_plan, sp_padding_size will be automatically set in get_forward_context(), however, hidden_states_mask is required to exclude the padded tokens.

# vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py hidden_states_mask = None # default if self.parallel_config is not None and self.parallel_config.sequence_parallel_size > 1: ctx = get_forward_context() if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: # Create mask for the full (padded) sequence # valid positions = True, padding positions = False batch_size = hidden_states.shape[0] padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size hidden_states_mask = torch.ones( batch_size, padded_seq_len, dtype=torch.bool, device=hidden_states.device, ) hidden_states_mask[:, ctx.sp_original_seq_len :] = False

@RuixiangMa I would like to invite you to join the discussion here #1324.

I have two suggestions:

refer to [Feature] Support Wan2.2 output with irregular shapes #1279 , you may use auto_pad=True and use hidden_states_mask to exclude the padded tokens.

You may set SequenceParallelInput(auto_pad=False), this will raise an error when seq_len not divisible by sp_size, which I think might be rare for Flux.2-klein. We can take care of rare cases in the future, after we make our discussion clear in [RFC]: Ulysses-SP Constraints Solution #1324.

hsliuustc0106 · 2026-02-24T08:02:37Z

@vllm-omni-reviewer

Signed-off-by: Lancer <maruixiang6688@gmail.com>

wtomin · 2026-03-09T12:10:58Z

Do you have plan to support SP with _sp_plan hooks?

RuixiangMa · 2026-03-09T12:39:12Z

Do you have plan to support SP with _sp_plan hooks?

Sorry, just submitted a revision

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa · 2026-03-09T13:19:50Z

Do you have plan to support SP with _sp_plan hooks?

Sorry, just submitted a revision

@wtomin Can you take a look?

wtomin · 2026-03-09T14:31:54Z

There is a bug related to sp in #1556, I have a bugfix in #1704. Please check if it affects your PR.

hsliuustc0106 · 2026-03-09T14:33:44Z

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

            sin = sin.to(query.dtype)
-            query = self.rope(query, cos, sin)
-            key = self.rope(key, cos, sin)
+            txt_cos, img_cos = cos[:text_seq_len], cos[text_seq_len:]


are these added lines from 397 to 410 independent of any diffusion model and can be extracted?

I've extracted the RoPE logic into a reusable helper function(only verified flux1, flux2, and z-image so far, but the others (e.g., ovis-image) should be reusable as well.), and SP-specific logic in flux2_klein (lines 393-410) remains unchanged as it handles model-specific sequence parallelism splitting.

wtomin · 2026-03-09T14:35:18Z

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py


+        hidden_states_mask = None
+        ctx = get_forward_context()
+        if ctx.sp_active:


I think ctx.sp_active is should not be exposed to users or developers.

In _sp_plan hooks design, ctx.sp_active is determined by _sp_shard_depth. _sp_shard_depth will be automatically handled by hook function. Therefore, no need to include ctx.sp_active here.

I think ctx.sp_active is should not be exposed to users or developers.

In _sp_plan hooks design, ctx.sp_active is determined by _sp_shard_depth. _sp_shard_depth will be automatically handled by hook function. Therefore, no need to include ctx.sp_active here.

It was redundant, I removed it

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa · 2026-03-10T06:37:44Z

There is a bug related to sp in #1556, I have a bugfix in #1704. Please check if it affects your PR.

Flux2's hook-based impl seems unaffected. I'll merge and verify

wtomin · 2026-03-10T08:01:29Z

I would recommend you to use _sp_plan hooks for sp implementation instead of intrusive modification.

Taking PR #1772 as an example, one major reason that intrusive modification is not good, it affects teacache extractor function, making it over complicated.

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa · 2026-03-10T12:31:11Z

There is a bug related to sp in #1556, I have a bugfix in #1704. Please check if it affects your PR.

merge the pr and test, it work well.

wtomin · 2026-03-11T03:08:43Z

Have you run offline inference with irreguar shape that involves auto padding and hidden_states_mask?

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa · 2026-03-11T11:23:28Z

Have you run offline inference with irreguar shape that involves auto padding and hidden_states_mask?

Both irreguar and offload tests passed.

wtomin

LGTM.

RuixiangMa requested a review from hsliuustc0106 as a code owner February 6, 2026 09:47

chatgpt-codex-connector bot reviewed Feb 6, 2026

View reviewed changes

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py Outdated Show resolved Hide resolved

[Feat] support SP for FLUX.2-klein

e8fb739

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa force-pushed the spforflux2klein branch from f3436b8 to e8fb739 Compare February 6, 2026 09:56

hsliuustc0106 reviewed Feb 6, 2026

View reviewed changes

RuixiangMa added 2 commits February 8, 2026 01:05

upd

f2daa52

Signed-off-by: Lancer <maruixiang6688@gmail.com>

upd

225ec8a

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa force-pushed the spforflux2klein branch from 482940d to 225ec8a Compare February 7, 2026 17:18

upd

c57c866

Signed-off-by: Lancer <maruixiang6688@gmail.com>

wtomin reviewed Feb 9, 2026

View reviewed changes

upd

b33ce10

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa force-pushed the spforflux2klein branch from 0b50742 to b33ce10 Compare February 9, 2026 14:57

wtomin reviewed Feb 10, 2026

View reviewed changes

This was referenced Feb 10, 2026

[RFC]: Continuous Diffusion Model Acceleration Support #1217

Open

[RFC]: Ulysses-SP Constraints Solution #1324

Closed

gcanlin mentioned this pull request Feb 26, 2026

[RFC]: vLLM-Omni NPU 2026 Q1 Roadmap #886

Open

22 tasks

RuixiangMa added 2 commits March 3, 2026 11:44

Merge branch 'main' into spforflux2klein

7166799

Signed-off-by: Lancer <maruixiang6688@gmail.com>

upd

780ed85

Signed-off-by: Lancer <maruixiang6688@gmail.com>

upd

8a7e687

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa force-pushed the spforflux2klein branch from 823dfcf to 8a7e687 Compare March 9, 2026 12:57

hsliuustc0106 reviewed Mar 9, 2026

View reviewed changes

wtomin reviewed Mar 9, 2026

View reviewed changes

RuixiangMa added 2 commits March 10, 2026 10:23

upd

63c4658

Signed-off-by: Lancer <maruixiang6688@gmail.com>

Merge branch 'main' into spforflux2klein

219d4f7

upd

876d492

Signed-off-by: Lancer <maruixiang6688@gmail.com>

RuixiangMa added 2 commits March 11, 2026 14:27

Merge branch 'main' into spforflux2klein

ec5f9f2

upd

6a1e5f2

Signed-off-by: Lancer <maruixiang6688@gmail.com>

wtomin reviewed Mar 11, 2026

View reviewed changes

wtomin approved these changes Mar 11, 2026

View reviewed changes

Conversation

RuixiangMa commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ZJY0516 Feb 10, 2026 • 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!

hsliuustc0106 commented Feb 7, 2026

Uh oh!

RuixiangMa commented Feb 8, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 commented Feb 24, 2026

Uh oh!

wtomin commented Mar 9, 2026

Uh oh!

RuixiangMa commented Mar 9, 2026

Uh oh!

RuixiangMa commented Mar 9, 2026

Uh oh!

wtomin commented Mar 9, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

RuixiangMa Mar 10, 2026 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

RuixiangMa commented Mar 10, 2026

Uh oh!

wtomin commented Mar 10, 2026

Uh oh!

RuixiangMa commented Mar 10, 2026

Uh oh!

wtomin commented Mar 11, 2026

Uh oh!

RuixiangMa commented Mar 11, 2026

Uh oh!

wtomin left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

RuixiangMa commented Feb 6, 2026 •

edited

Loading

ZJY0516 Feb 10, 2026 •

edited

Loading

RuixiangMa Mar 10, 2026 •

edited

Loading