Skip to content

feat: add sequence parallel and cache support for Flux DiT. - #2236

Open
z-jun03 wants to merge 1 commit into
xLLM-AI:mainfrom
z-jun03:flux_sp
Open

feat: add sequence parallel and cache support for Flux DiT.#2236
z-jun03 wants to merge 1 commit into
xLLM-AI:mainfrom
z-jun03:flux_sp

Conversation

@z-jun03

@z-jun03 z-jun03 commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Description

Add sequence parallelism and DiT cache support for Flux DiT models.

This change:

  • Adds Ulysses-style sequence parallel execution for Flux attention.
  • Adds sequence padding, splitting, all-to-all communication, gathering, and unpadding.
  • Integrates Flux with the existing DiT cache framework.
  • Supports independent cache states for positive and negative CFG passes.
  • Preserves existing LongCat-Image behavior through backward-compatible constructors.
  • Applies the changes to the Flux, Flux Control, and Flux Fill pipelines.

Related Issues

Change Type

  • Bug fix
  • New feature
  • Performance improvement
  • Refactor
  • Documentation
  • Test
  • Build or CI

Pull Request Checklist

Thank you for contributing to xLLM. Before requesting review, please make sure the following items are complete.

PR Title and Commit Messages

  • The PR title and each commit message follow the xLLM commit format: <type>: <subject>.

Pre-commit Checks

  • I have installed pre-commit by running pip install pre-commit or an equivalent command.
  • I have installed the hooks with pre-commit install.
  • I have run pre-commit run --all-files and fixed any reported issues.

Self Review

  • I have self-reviewed the code according to .agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.
  • I have rebased this PR onto the latest main branch.

Build and Test Coverage

  • Tests have been added or updated as needed.
  • CUDA: python setup.py build test has passed on a CUDA machine.
  • NPU: python setup.py build test has passed on an NPU machine.
  • MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

  • Flux build validation passed successfully.
  • Flux generation was validated on Ascend NPU with both SP=1 and SP=2.
  • SP=2 output matched the single-card baseline byte-for-byte for the same prompt, seed, resolution, and inference steps.
  • FBCache output matched the no-cache baseline in the tested configuration.
  • Flux, Flux Control, and Flux Fill pipelines were updated.
  • Existing LongCat-Image usage was preserved through backward-compatible constructors.
  • The current implementation supports asynchronous communication submission but does not implement true computation-communication overlap.
  • Flux-specific automated SP and DiT cache tests were not added; validation was performed through NPU build and runtime integration tests.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant