feat: load config support hybrid parallel (#777)

* load config support hybrid parallel

* load config support hybrid parallel

* load config support hybrid parallel

* load config support hybrid parallel

* load config support hybrid parallel

* load config support hybrid parallel

* load config support hybrid parallel

Author: DefTruth

README.md

@@ -34,7 +34,9 @@
 
 ## 🔥Latest News 
 
-- [2026/01] **[🎉v1.2.0 Major Release](https://github.com/vipshop/cache-dit)** is ready: New Models Support(Z-Image, FLUX.2, LTX-2, etc), Request level Cache Context, HTTP Serving, [Ulysses Anything](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/#uaa-ulysses-anything-attention), TE-P, VAE-P, CN-P and [Ascend NPUs](https://cache-dit.readthedocs.io/en/latest/user_guide/ASCEND_NPU/) Support.
+- [2026/02] **[🎉v1.2.1](https://github.com/vipshop/cache-dit)** release is ready, the major updates including: [Ring](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL) Attention w/ [batched P2P](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL), [USP](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/) (Hybrid Ring and Ulysses), Hybrid 2D and 3D Parallelism (💥[USP + TP](https://cache-dit.readthedocs.io/en/latest/user_guide/HYBRID_PARALLEL/)),  VAE-P Comm overhead reduce.
+- [2026/01] **[🎉v1.2.0](https://github.com/vipshop/cache-dit)** stable release is ready: New Models Support(Z-Image, FLUX.2, LTX-2, etc), Request level Cache Context, HTTP Serving, [Ulysses Anything](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/#uaa-ulysses-anything-attention), TE-P, VAE-P, CN-P and [Ascend NPUs](https://cache-dit.readthedocs.io/en/latest/user_guide/ASCEND_NPU/) support.
+
 
 ## 🚀Quick Start 
 
@@ -55,14 +57,13 @@ Then accelerate your DiTs with just **♥️one line♥️** of code ~
 >>> cache_dit.enable_cache(
 ...   pipe, cache_config=DBCacheConfig(), # w/ default
 ...   parallelism_config=ParallelismConfig(ulysses_size=2))
+>>> # Or, Use Distributed Inference without Cache Acceleration.
+>>> cache_dit.enable_cache(
+...   pipe, parallelism_config=ParallelismConfig(ulysses_size=2))
 >>> # Or, Hybrid Cache Acceleration + 2D Parallelism.
 >>> cache_dit.enable_cache(
 ...   pipe, cache_config=DBCacheConfig(), # w/ default
 ...   parallelism_config=ParallelismConfig(ulysses_size=2, tp_size=2))
->>> # Or, Use Distributed Inference without Cache Acceleration.
->>> cache_dit.enable_cache(
-...   pipe, cache_config=None, # Set cache_config as None.
-...   parallelism_config=ParallelismConfig(ulysses_size=2))
 >>> from cache_dit import load_configs
 >>> # Or, Load Acceleration config from a custom yaml file.
 >>> cache_dit.enable_cache(pipe, **load_configs("config.yaml"))

docs/README.md

@@ -40,7 +40,8 @@
 
 ## 🔥Latest News 
 
-- [2026/01] **[🎉v1.2.0 Major Release](https://github.com/vipshop/cache-dit)** is ready: New Models Support(Z-Image, FLUX.2, LTX-2, etc), Request level Cache Context, HTTP Serving, [Ulysses Anything](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/#uaa-ulysses-anything-attention), TE-P, VAE-P, CN-P and [Ascend NPUs](https://cache-dit.readthedocs.io/en/latest/user_guide/ASCEND_NPU/) Support.
+- [2026/02] **[🎉v1.2.1](https://github.com/vipshop/cache-dit)** release is ready, the major updates including: [Ring](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL) Attention w/ [batched P2P](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL), [USP](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/) (Hybrid Ring and Ulysses), Hybrid 2D and 3D Parallelism (💥[USP + TP](https://cache-dit.readthedocs.io/en/latest/user_guide/HYBRID_PARALLEL/)),   VAE-P Comm overhead reduce.
+- [2026/01] **[🎉v1.2.0](https://github.com/vipshop/cache-dit)** stable release is ready: New Models Support(Z-Image, FLUX.2, LTX-2, etc), Request level Cache Context, HTTP Serving, [Ulysses Anything](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/#uaa-ulysses-anything-attention), TE-P, VAE-P, CN-P and [Ascend NPUs](https://cache-dit.readthedocs.io/en/latest/user_guide/ASCEND_NPU/) Support.
 
 ## 🚀Quick Start 
 
@@ -61,14 +62,13 @@ Then accelerate your DiTs with just **♥️one line♥️** of code ~
 >>> cache_dit.enable_cache(
 ...   pipe, cache_config=DBCacheConfig(), # w/ default
 ...   parallelism_config=ParallelismConfig(ulysses_size=2))
+>>> # Or, Use Distributed Inference without Cache Acceleration.
+>>> cache_dit.enable_cache(
+...   pipe, parallelism_config=ParallelismConfig(ulysses_size=2))
 >>> # Or, Hybrid Cache Acceleration + 2D Parallelism.
 >>> cache_dit.enable_cache(
 ...   pipe, cache_config=DBCacheConfig(), # w/ default
 ...   parallelism_config=ParallelismConfig(ulysses_size=2, tp_size=2))
->>> # Or, Use Distributed Inference without Cache Acceleration.
->>> cache_dit.enable_cache(
-...   pipe, cache_config=None, # Set cache_config as None.
-...   parallelism_config=ParallelismConfig(ulysses_size=2))
 >>> from cache_dit import load_configs
 >>> # Or, Load Acceleration config from a custom yaml file.
 >>> cache_dit.enable_cache(pipe, **load_configs("config.yaml"))

docs/user_guide/LOAD_CONFIGS.md

@@ -27,6 +27,8 @@ Then, apply the acceleration config from yaml.
 
 ## Distributed inference  
 
+- 1D Parallelism
+
 Define a parallelism only config yaml `parallel.yaml` file that contains:
 
 ```yaml
@@ -42,6 +44,43 @@ Then, apply the distributed inference acceleration config from yaml. `ulysses_si
 >>> cache_dit.enable_cache(pipe, **cache_dit.load_configs("parallel.yaml"))
 ```
 
+- 2D Parallelism
+
+You can also define a 2D parallelism config yaml `parallel_2d.yaml` file that contains:
+
+```yaml
+parallelism_config:
+  ulysses_size: auto
+  tp_size: 2
+  parallel_kwargs:
+    attention_backend: native
+    extra_parallel_modules: ["text_encoder", "vae"]
+```
+Then, apply the 2D parallelism config from yaml. Here `tp_size: 2` means using tensor parallelism with size 2. The `ulysses_size: auto` means that cache-dit will auto detect the `world_size // tp_size` as the ulysses_size.
+```python
+>>> import cache_dit
+>>> cache_dit.enable_cache(pipe, **cache_dit.load_configs("parallel_2d.yaml"))
+```
+
+- 3D Parallelism
+
+You can also define a 3D parallelism config yaml `parallel_3d.yaml` file that contains:
+
+```yaml
+parallelism_config:
+  ulysses_size: 2
+  ring_size: 2
+  tp_size: 2
+  parallel_kwargs:
+    attention_backend: native
+    extra_parallel_modules: ["text_encoder", "vae"]
+```
+Then, apply the 3D parallelism config from yaml. Here `ulysses_size: 2`, `ring_size: 2`, `tp_size: 2` means using ulysses parallelism with size 2, ring parallelism with size 2 and tensor parallelism with size 2.
+```python
+>>> import cache_dit
+>>> cache_dit.enable_cache(pipe, **cache_dit.load_configs("parallel_3d.yaml"))
+```
+
 ## Hybrid Cache and Parallelism
 
 Define a hybrid cache and parallel acceleration config yaml `hybrid.yaml` file that contains:
@@ -81,6 +120,8 @@ pip3 install git+https://github.com/huggingface/diffusers.git # latest or >= 0.3
 pip3 install git+https://github.com/vipshop/cache-dit.git # latest
 
 python3 -m cache_dit.generate flux --config cache.yaml
-torchrun --nproc_per_node=4 -m cache_dit.generate flux --config parallel.yaml
 torchrun --nproc_per_node=4 -m cache_dit.generate flux --config hybrid.yaml
+torchrun --nproc_per_node=4 -m cache_dit.generate flux --config parallel.yaml
+torchrun --nproc_per_node=4 -m cache_dit.generate flux --config parallel_2d.yaml
+torchrun --nproc_per_node=8 -m cache_dit.generate flux --config parallel.yaml
 ```

examples/configs/parallel_2d.yaml

@@ -0,0 +1,6 @@
+parallelism_config:
+  ulysses_size: auto
+  tp_size: 2
+  parallel_kwargs:
+    attention_backend: native
+    extra_parallel_modules: ["text_encoder", "vae"]

examples/configs/parallel_3d.yaml

@@ -0,0 +1,7 @@
+parallelism_config:
+  ulysses_size: 2
+  ring_size: 2
+  tp_size: 2
+  parallel_kwargs:
+    attention_backend: native
+    extra_parallel_modules: ["text_encoder", "vae"]

src/cache_dit/caching/utils.py

@@ -225,7 +225,9 @@ def load_parallelism_config(
         backend_str = parallelism_config_kwargs["backend"]
         parallelism_config_kwargs["backend"] = ParallelismBackend.from_str(backend_str)
 
-    def _maybe_auto_parallel_size(size: str | int | None) -> Optional[int]:
+    def _maybe_auto_parallel_size(
+        size: str | int | None, partial_max_size: Optional[int] = None
+    ) -> Optional[int]:
         if size is None:
             return None
         if isinstance(size, int):
@@ -236,7 +238,11 @@ def _maybe_auto_parallel_size(size: str | int | None) -> Optional[int]:
             size = 1
             if dist.is_initialized():
                 # Assume world size is the parallel size
-                size = dist.get_world_size()
+                world_size = dist.get_world_size()
+                if partial_max_size is not None:
+                    size = world_size // partial_max_size
+                else:
+                    size = world_size
             if size == 1:
                 logger.warning(
                     "Auto parallel size selected as 1. Make sure to run with torch.distributed "
@@ -247,20 +253,40 @@ def _maybe_auto_parallel_size(size: str | int | None) -> Optional[int]:
             return size
         raise ValueError(f"Invalid parallel size value: {size}. Must be int or 'auto'.")
 
-    if kwargs.get("auto_parallel_size", True):
-        if "ulysses_size" in parallelism_config_kwargs:
-            parallelism_config_kwargs["ulysses_size"] = _maybe_auto_parallel_size(
-                parallelism_config_kwargs["ulysses_size"]
-            )
-        if "ring_size" in parallelism_config_kwargs:
-            parallelism_config_kwargs["ring_size"] = _maybe_auto_parallel_size(
-                parallelism_config_kwargs["ring_size"]
-            )
-        if "tp_size" in parallelism_config_kwargs:
-            parallelism_config_kwargs["tp_size"] = _maybe_auto_parallel_size(
-                parallelism_config_kwargs["tp_size"]
+    def _maybe_auto_parallel_sizes(parallelism_config_kwargs: dict) -> dict:
+        # Only allow one of the parallel size to be auto for simplicity
+        auto_count = sum(
+            1
+            for key in ["ulysses_size", "ring_size", "tp_size"]
+            if key in parallelism_config_kwargs and parallelism_config_kwargs[key] == "auto"
+        )
+        if auto_count > 1:
+            raise ValueError(
+                "Only one of 'ulysses_size', 'ring_size', or 'tp_size' can be set to 'auto'."
             )
 
+        ulysses_size = parallelism_config_kwargs.get("ulysses_size", 1)
+        ring_size = parallelism_config_kwargs.get("ring_size", 1)
+        tp_size = parallelism_config_kwargs.get("tp_size", 1)
+        partial_max_size = None
+        if isinstance(ulysses_size, str) and ulysses_size.lower() == "auto":
+            partial_max_size = ring_size * tp_size
+        elif isinstance(ring_size, str) and ring_size.lower() == "auto":
+            partial_max_size = ulysses_size * tp_size
+        elif isinstance(tp_size, str) and tp_size.lower() == "auto":
+            partial_max_size = ulysses_size * ring_size
+
+        for key in ["ulysses_size", "ring_size", "tp_size"]:
+            if key in parallelism_config_kwargs:
+                parallelism_config_kwargs[key] = _maybe_auto_parallel_size(
+                    parallelism_config_kwargs[key], partial_max_size=partial_max_size
+                )
+        return parallelism_config_kwargs
+
+    if kwargs.get("auto_parallel_size", True):
+
+        parallelism_config_kwargs = _maybe_auto_parallel_sizes(parallelism_config_kwargs)
+
     parallelism_config = ParallelismConfig(**parallelism_config_kwargs)
     return parallelism_config