[RFC]: Bagel deployment

Motivation

offline: https://github.com/vllm-project/vllm-omni/pull/1669 online: https://github.com/vllm-project/vllm-omni/pull/1804 Training and serving the Bagel omni-modal model presents distinct challenges focused on high-quality text and image understanding and generation. Bagel processes text and images as both input and output, prioritizing generation quality and complex cross-modal reasoning over ultra-low latency.

(1) Cross-modal heterogeneity: Bagel must effectively fuse text and image representations to ensure disjoint modalities do not degrade each other. The model needs to handle complex instructions involving both modalities and generate high-fidelity images or text responses accordingly.

(2) Throughput and Efficiency: The heavy computational cost of image generation (via DiT) and large-scale language modeling requires significant optimization to ensure reasonable inference times and high throughput for concurrent requests.

(3) Complex Pipeline Integration: Bagel's architecture involves orchestrating multiple heavy components—Multimodal LLM (Thinker), Diffusion Transformer (DiT), and VAEs. Integrating these efficiently, especially with features like cross-stage VAE and TeaCache, is critical for making deployment viable in production environments.

For Bagel, we aim to implement a deployment strategy that addresses these challenges while supporting specific features like DiT batch inference, TeaCache, and advanced parallelism to maximize hardware utilization and reduce cost.

Proposed Change

For detailed designs, please refer to the specific feature tracking below.

P0: Core Architecture & Foundation

• Add Bagel E2E Smoked Test [Feature] https://github.com/vllm-project/vllm-omni/pull/1074

  • Need test Bagel text2img multi-stage use case, using mooncake or shared-memory.
  • Related: https://github.com/vllm-project/vllm-omni/issues/400

• Think Model Support [Feature]

  • Integrate support for the "Think" model component within the Bagel deployment pipeline.

• DiT Offline Batch Inference [Feature]

  • Support offline batch inference for the Diffusion Transformer (DiT) component.
  • Related: https://github.com/vllm-project/vllm-omni/pull/797

• AR Model VAE Support (Img2Img) [Feature] https://github.com/vllm-project/vllm-omni/pull/1669

  • Enable VAE module support within the Auto-Regressive (AR) model.
  • Allow Img2Img tasks to span across different stages (cross-stage support).

• KV Cache Transfer Encapsulation [Refactor] https://github.com/vllm-project/vllm-omni/pull/979

  • Encapsulate the KV cache transfer logic between Diffusion Runner and AR Runner into a dedicated utility object to facilitate future rebases and maintenance.
  • Related: https://github.com/vllm-project/vllm-omni/issues/944

• Diffusion KV Manager [Optimization]

  • Implement a Diffusion KV Manager to decouple KV cache lifecycle from the DiT stage transfer.
  • Prevent immediate release of KV cache after transfer to better support multi-turn conversations and reduce memory fragmentation.
  • Related: https://github.com/vllm-project/vllm-omni/issues/800

• KV Cache Extraction & Transfer Decoupling [Architecture]

  • Decouple KV cache format extraction from the transfer mechanism.
  • Design a generic interface to support other models that require cross-stage KV cache transfer dependencies.
  • Related: https://github.com/vllm-project/vllm-omni/issues/800

P1: Performance & Production Readiness

• TeaCache Support [Optimization]https://github.com/vllm-project/vllm-omni/pull/848

  • Implement TeaCache integration for efficient caching and acceleration.

• Cross-node RDMA KV Cache Transfer [Feature] @natureofnature https://github.com/vllm-project/vllm-omni/pull/1019

  • Support Remote Direct Memory Access (RDMA) for transferring Key-Value (KV) caches in cross-node deployment scenarios to minimize latency.
  • Related: https://github.com/vllm-project/vllm-omni/issues/955

• Multi-turn Conversation Support [Feature]

  • Enable robust management of multi-turn interaction contexts (Dependent on Diffusion KV Manager).

• Bagel Documentation Update [Docs] @nussejzz https://github.com/vllm-project/vllm-omni/pull/987

  • Update Bagel documentation with comprehensive online and offline deployment strategies and usage guides.

P2: Advanced Scalability

• Ulysses-SP Implementation [Feature] https://github.com/vllm-project/vllm-omni/pull/1903

  • Implement Ulysses Sequence Parallelism (SP) for improved efficiency on long sequences.

• Ring-Attention Implementation [Feature] https://github.com/vllm-project/vllm-omni/pull/1903

  • Implement Ring-Attention to support longer context windows and efficient distributed training/inference.

• CFG-Parallel Implementation [Feature] https://github.com/vllm-project/vllm-omni/pull/1578

  • Implement parallelism for Classifier-Free Guidance (CFG) to speed up generation steps.
  • Related: https://github.com/vllm-project/vllm-omni/issues/1047