[Feat] support CFG Parallel for FLUX kontext by RuixiangMa · Pull Request #2281 · vllm-project/vllm-omni

RuixiangMa · 2026-03-28T01:43:49Z

Purpose

support CFG Parallel for FLUX kontext

Test Plan

Test Result

#!/bin/bash

# Offline image-to-image test matrix.
# Compare eager/non-eager and cfg-parallel on/off.

export CUDA_VISIBLE_DEVICES=1,2,3,4

PROMPT="Add a hat to the dog"
IMAGE="dog.png"
NEGATIVE_PROMPT="ugly, unclear, blurry, gray"
CFG_SCALE=4.0

# Arrays to track success and failure
declare -a SUCCESS_TASKS
declare -a FAILED_TASKS

# Format: "model_name|model_path"
declare -a MODELS=(
  "FLUX.1-Kontext-dev|black-forest-labs/FLUX.1-Kontext-dev"
)

# Eager mode configurations
declare -a EAGER_CONFIGS=(
  "no_eager|"
  "with_eager|--enforce-eager"
)

# CFG parallel configurations
declare -a CFG_CONFIGS=(
  "no_cfg_parallel|"
  "with_cfg_parallel|--cfg-parallel-size 2"
)

echo "=========================================="
echo "Starting offline image edit inference tests"
echo "Script: examples/offline_inference/image_to_image/image_edit.py"
echo "Testing combinations of eager mode and CFG parallel"
echo "=========================================="
echo ""

if [[ ! -f "$IMAGE" ]]; then
  echo "Input image not found: $IMAGE"
  exit 1
fi

TASK_NUM=0
TOTAL_TASKS=$((${#MODELS[@]} * ${#EAGER_CONFIGS[@]} * ${#CFG_CONFIGS[@]}))

for model_info in "${MODELS[@]}"; do
  IFS='|' read -r model_name model_path <<< "$model_info"

  model_dir="${model_name// /_}"
  mkdir -p "$model_dir"

  for eager_info in "${EAGER_CONFIGS[@]}"; do
    IFS='|' read -r eager_label eager_args <<< "$eager_info"

    for cfg_info in "${CFG_CONFIGS[@]}"; do
      IFS='|' read -r cfg_label cfg_args <<< "$cfg_info"
      TASK_NUM=$((TASK_NUM + 1))

      base_name="${model_name,,}"
      base_name="${base_name// /_}"
      output_file="$model_dir/${base_name}_output_${eager_label}_${cfg_label}.png"
      log_file="$model_dir/${base_name}_${eager_label}_${cfg_label}.log"
      task_label="$model_name ($eager_label + $cfg_label)"

      # no_cfg_parallel: TP=2 (2 GPUs)
      # with_cfg_parallel: TP=2 + CFG=2 (4 GPUs)
      tp_size=2

      echo "=========================================="
      echo "$TASK_NUM/$TOTAL_TASKS: Running $task_label..."
      echo "=========================================="
      echo "Using CFG scale: $CFG_SCALE"

      if python examples/offline_inference/image_to_image/image_edit.py \
        --model "$model_path" \
        --cfg-scale "$CFG_SCALE" \
        --prompt "$PROMPT" \
        --negative-prompt "$NEGATIVE_PROMPT" \
        --tensor-parallel-size "$tp_size" \
        --enable-cpu-offload \
        --image "$IMAGE" \
        --output "$output_file" \
        $eager_args \
        $cfg_args \
        2>&1 | tee "$log_file"; then
        echo "PASS $task_label completed."
        SUCCESS_TASKS+=("$task_label")
      else
        echo "FAIL $task_label failed."
        FAILED_TASKS+=("$task_label")
      fi
      echo ""
    done
  done
done

echo "=========================================="
echo "All tasks completed!"
echo "=========================================="
echo "Summary: ${#SUCCESS_TASKS[@]}/$TOTAL_TASKS successful, ${#FAILED_TASKS[@]}/$TOTAL_TASKS failed"
echo ""

if [ ${#SUCCESS_TASKS[@]} -gt 0 ]; then
  echo "PASS Successful tasks:"
  for task in "${SUCCESS_TASKS[@]}"; do
    echo "  - $task"
  done
  echo ""
fi

if [ ${#FAILED_TASKS[@]} -gt 0 ]; then
  echo "FAIL Failed tasks:"
  for task in "${FAILED_TASKS[@]}"; do
    echo "  - $task"
  done
  echo ""
  echo "Check model directories for error logs."
  echo ""
fi

echo "Output directories:"
for model_info in "${MODELS[@]}"; do
  IFS='|' read -r model_name _ <<< "$model_info"
  model_dir="${model_name// /_}"
  echo "  - $model_dir/"
done

if [ ${#FAILED_TASKS[@]} -gt 0 ]; then
  exit 1
fi

Configuration	Total Time
no_eager_no_cfg_parallel	114.822s
no_eager_with_cfg_parallel	67.569s
with_eager_no_cfg_parallel	117.827s
with_eager_with_cfg_parallel	68.837s

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
The test results. Please paste the results comparison before and after, or the e2e results.
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
(Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

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

lishunyang12

left a couple comments

lishunyang12 · 2026-04-02T15:38:20Z

-                img_ids=latent_ids,
-                joint_attention_kwargs=self.joint_attention_kwargs,
-                return_dict=False,
-            )[0]


check_cfg_parallel_validity returns False on invalid config but the return value is discarded here. Should this raise or at least skip the CFG path? Right now an invalid config silently runs and produces bad output.

Now, the CFG path now respects check_cfg_parallel_validity(...) instead of discarding the return value, so invalid CFG-parallel configurations fall back cleanly instead of continuing with bad output.

lishunyang12 · 2026-04-02T15:38:20Z

        raise NotImplementedError("Subclasses must implement diffuse")

+    def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool) -> bool:
+        """


Nit: the mixin already imports from vllm_omni.diffusion.distributed.parallel_state — is there a project-wide logger convention (e.g. from vllm) you could use instead of import logging + getLogger? Minor.

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

RuixiangMa · 2026-04-24T09:11:18Z

@hsliuustc0106 verified it and ready to merge

wtomin · 2026-05-14T02:34:31Z

        )
        do_true_cfg = true_cfg_scale > 1.0 and has_neg_prompt

+        do_true_cfg = do_true_cfg and self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)


About the CFG scale, I notice in FLUX.1-Kontext-dev HF webpage, diffusers use guidance_scale not true_cfg_scale. Please double check. Better to stay consistent with diffusers' key name.

FLUX.1-dev HF webpage uses guidance_scale too. I am indeed very confused about the two key names: guidance_scale and true_cfg_scale.

wtomin · 2026-05-14T03:10:21Z

Please update docs/user_guide/diffusion_features.md. Besides, please increase the CFG parallel feature in tests/e2e/online_serving/test_flux_kontext_expansion.py, taking tests/e2e/online_serving/test_flux_2_dev_expansion.py as an example.

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

RuixiangMa · 2026-05-19T14:18:14Z

Please update docs/user_guide/diffusion_features.md. Besides, please increase the CFG parallel feature in tests/e2e/online_serving/test_flux_kontext_expansion.py, taking tests/e2e/online_serving/test_flux_2_dev_expansion.py as an example.

updated it

wtomin · 2026-06-22T02:41:47Z

@RuixiangMa Pls resolve the conflicts. When it is ready, I think we can launch a nightly CI test for this PR.

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

wtomin · 2026-06-24T03:50:35Z

@hsliuustc0106 Could you please add the label related diffusion text to image CI test? The npu ci test failure is not relevant to this PR, I think.

I think the nightly CI label is too large, so I removed it. @RuixiangMa

RuixiangMa · 2026-06-24T04:44:01Z

is too large, so I remo

Thx, get it

wtomin · 2026-07-08T02:05:30Z

I notice that the nightly CI failed:

[2026-06-25T03:11:09Z] =========================== short test summary info ============================
--
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux2_expansion.py::test_flux2_klein[layerwise_hsdp] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_text_to_image[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_image_edit[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_image_edit_no_negative[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_high_resolution[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_multiple_outputs[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_text_to_image[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_image_edit[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_image_edit_no_negative[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_high_resolution[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_multiple_outputs[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] = 17 passed, 4 skipped, 398 deselected, 22 warnings, 11 errors in 3168.67s (0:52:48) =

Could you pls run this test tests/e2e/online_serving/test_flux_kontext_expansion.py locally to check if it is correct?

And please resolve the conflicts.

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

RuixiangMa · 2026-07-10T01:55:06Z

I notice that the nightly CI failed:

[2026-06-25T03:11:09Z] =========================== short test summary info ============================
--
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux2_expansion.py::test_flux2_klein[layerwise_hsdp] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_text_to_image[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_image_edit[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_image_edit_no_negative[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_high_resolution[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_multiple_outputs[parallel_tp_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_text_to_image[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_image_edit[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_image_edit_no_negative[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_high_resolution[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] ERROR tests/e2e/online_serving/test_flux_kontext_expansion.py::test_flux_kontext_multiple_outputs[parallel_cfg_2] - RuntimeError: Server processes exited with code 1 before becoming ready.
  | [2026-06-25T03:11:09Z] = 17 passed, 4 skipped, 398 deselected, 22 warnings, 11 errors in 3168.67s (0:52:48) =

Could you pls run this test tests/e2e/online_serving/test_flux_kontext_expansion.py locally to check if it is correct?

And please resolve the conflicts.

Yes, local testing works fine;

I suspect that FLUX.1-Kontext-dev et al are all gated models on HuggingFace. If the CI environment doesn't have an HF_TOKEN with approved access to these repos, the model download would fail and the tests would error. I added skip_if_gated_repo_inaccessible()

wtomin

LGTM.

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

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

RuixiangMa requested a review from hsliuustc0106 as a code owner March 28, 2026 01:43

[Feat] support CFG Parallel for FLUX kontext

4cf10d0

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

RuixiangMa force-pushed the knotextcfg branch from 46f80ef to 4cf10d0 Compare March 28, 2026 01:55

RuixiangMa mentioned this pull request Mar 28, 2026

[RFC]: Continuous Diffusion Model Acceleration Support #1217

Open

1 task done

lishunyang12 reviewed Apr 2, 2026

View reviewed changes

RuixiangMa added 5 commits April 5, 2026 15:11

upd

31dd7ce

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

Merge branch 'main' into knotextcfg

cc4d7f6

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

upd

60f1055

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

Merge branch 'main' into knotextcfg

a2f0c51

upd

21e493f

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

RuixiangMa added 2 commits May 7, 2026 05:39

Merge branch 'main' into knotextcfg

8eda48d

Merge branch 'main' into knotextcfg

6021012

RuixiangMa requested review from Isotr0py, SamitHuang, ZJY0516, david6666666, princepride and wtomin as code owners May 13, 2026 16:46

wtomin reviewed May 14, 2026

View reviewed changes

upd

54738de

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

RuixiangMa requested review from Gaohan123 and yenuo26 as code owners May 19, 2026 13:56

RuixiangMa added 2 commits May 19, 2026 13:59

Merge remote-tracking branch origin/main into knotextcfg

432e2eb

upd

fe880da

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

Merge branch 'main' into knotextcfg

9630151

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

RuixiangMa added the nightly-test label to trigger buildkite nightly test CI label Jun 22, 2026

Merge branch 'main' into knotextcfg

69901fa

wtomin removed the nightly-test label to trigger buildkite nightly test CI label Jun 24, 2026

wtomin added this to the v0.24.0 milestone Jun 24, 2026

RuixiangMa added the diffusion-x2iat-test label Jun 24, 2026

Merge branch 'main' into knotextcfg

1370d2d

wtomin added diffusion-x2iat-test and removed diffusion-x2iat-test labels Jun 25, 2026

Gaohan123 modified the milestones: v0.24.0, v0.26.0 Jul 7, 2026

Merge branch main into knotextcfg: resolve pipeline_flux conflict

1328ec7

RuixiangMa requested a review from xuechendi as a code owner July 10, 2026 01:38

upd

353013b

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

wtomin added diffusion-x2iat-test and removed diffusion-x2iat-test labels Jul 13, 2026

hsliuustc0106 added diffusion codes related to diffusion models enhancement New feature or request labels Jul 17, 2026 — with ChatGPT Codex Connector

wtomin approved these changes Jul 20, 2026

View reviewed changes

wtomin merged commit eef3bd5 into vllm-project:main Jul 20, 2026
6 checks passed

yenuo26 mentioned this pull request Jul 21, 2026

[Bug]: Intel CI, tests/e2e/online_serving/test_flux_kontext_expansion.py, pytest.skip outside of a test will skip the entire module #5250

Closed

1 task done

linyueqian mentioned this pull request Jul 21, 2026

[Bug] Fix Ming feature extractor/processor registration crash under transformers >= 5.x #5243

Merged

y-null pushed a commit to y-null/vllm-omni that referenced this pull request Jul 21, 2026

[Feat] support CFG Parallel for FLUX kontext (vllm-project#2281)

24acf4e

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

bowieshi pushed a commit to bowieshi/vllm-omni that referenced this pull request Jul 22, 2026

[Feat] support CFG Parallel for FLUX kontext (vllm-project#2281)

7045821

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

linyueqian mentioned this pull request Jul 22, 2026

[Bug]: weekly CI, tests/e2e/online_serving/test_flux_kontext_expansion.py [parallel_cfg_2], CUDA OOM on L4 during dummy run (5 setup errors) #5307

Closed

1 task done

nagisa-kunhah pushed a commit to nagisa-kunhah/vllm-omni that referenced this pull request Jul 30, 2026

[Feat] support CFG Parallel for FLUX kontext (vllm-project#2281)

ce15307

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

khairulkabir1661 pushed a commit to khairulkabir1661/vllm-omni that referenced this pull request Sep 25, 2026

[Feat] support CFG Parallel for FLUX kontext (vllm-project#2281)

584c2b1

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

Conversation

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

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

lishunyang12 left a comment

Choose a reason for hiding this comment

Uh oh!

lishunyang12 Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

RuixiangMa Apr 5, 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!

lishunyang12 Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

RuixiangMa Apr 5, 2026

Choose a reason for hiding this comment

Uh oh!

RuixiangMa commented Apr 24, 2026

Uh oh!

wtomin May 14, 2026

Choose a reason for hiding this comment

Uh oh!

wtomin May 14, 2026

Choose a reason for hiding this comment

Uh oh!

wtomin commented May 14, 2026

Uh oh!

RuixiangMa commented May 19, 2026

Uh oh!

wtomin commented Jun 22, 2026

Uh oh!

wtomin commented Jun 24, 2026

Uh oh!

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

Uh oh!

Uh oh!

wtomin commented Jul 8, 2026

Uh oh!

RuixiangMa commented Jul 10, 2026

Uh oh!

wtomin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

RuixiangMa commented Mar 28, 2026 •

edited

Loading

RuixiangMa Apr 5, 2026 •

edited

Loading

RuixiangMa commented Jun 24, 2026 •

edited

Loading