Skip to content

feat: support separate DiT/VAE service instances for Flux. - #2256

Draft
z-jun03 wants to merge 1 commit into
xLLM-AI:mainfrom
z-jun03:vae_edit
Draft

feat: support separate DiT/VAE service instances for Flux.#2256
z-jun03 wants to merge 1 commit into
xLLM-AI:mainfrom
z-jun03:vae_edit

Conversation

@z-jun03

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

Copy link
Copy Markdown
Collaborator

Description

Add role-based DiT/VAE separation for Flux inference services.

This change:

  • Adds all, dit, and vae instance roles for Flux pipelines.
  • Allows DiT instances to generate latent outputs without loading VAE weights.
  • Allows VAE instances to load only VAE weights and serve latent decoding requests.
  • Adds configurable VAE service endpoints through dit_vae_service_addresses.
  • Adds DiT-to-VAE RPC communication through CommChannel and RemoteWorker.
  • Supports round-robin VAE selection, health checks, request timeout configuration, and failover.
  • Preserves the existing all mode for backward compatibility.
  • Adds configurable DiT Worker service ports for multi-instance deployment.

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>.

Recommended title:

feat: support separate DiT and VAE service instances for Flux.

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: Flux DiT/VAE separated deployment and runtime validation have been completed on an NPU machine.
  • MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

  • The implementation supports multiple DiT instances sharing one or more VAE service instances through configurable VAE endpoints.
  • DiT instances send latent tensors to VAE instances through the existing Worker RPC interface.
  • VAE instances return decoded image tensors through the RPC response.
  • VAE selection uses local round-robin routing with health checks and failover when a VAE request fails.
  • The existing dit_instance_role=all mode is preserved for backward compatibility.
  • The current implementation uses static dit_vae_service_addresses; dynamic service discovery, a centralized router, and global load balancing are not included.
  • DiT and VAE currently reuse the existing DiTWorkerImpl and WorkerService infrastructure, while model component loading and forward behavior are selected by dit_instance_role.
  • Automated unit tests for the separated DiT/VAE deployment path are not included in this change; validation is currently based on NPU build and runtime testing.

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