From d613864dbfca0669c91a289a82f41bde6a871c47 Mon Sep 17 00:00:00 2001
From: fywc <hanzheli@kuaishou.com>
Date: Mon, 20 Apr 2026 18:51:16 +0800
Subject: [PATCH] [Model] Add HSDP support for LTX-2 (#2899)

Signed-off-by: hanzheli <hanzheli@kuaishou.com>
Signed-off-by: fywc <hanzheli@kuaishou.com>
---
 docs/user_guide/diffusion_features.md         |  2 +-
 .../examples/online_serving/text_to_video.md  |  8 ++++++
 tests/diffusion/models/ltx2/test_ltx2_hsdp.py | 25 +++++++++++++++++++
 .../diffusion/models/ltx2/ltx2_transformer.py |  2 ++
 4 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 tests/diffusion/models/ltx2/test_ltx2_hsdp.py

diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md
index 7bdeede44..be5b2d10c 100644
--- a/docs/user_guide/diffusion_features.md
+++ b/docs/user_guide/diffusion_features.md
@@ -140,7 +140,7 @@ The following tables show which models support each feature:
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
 | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
 | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md
index 00a9c1672..b918aac19 100644
--- a/docs/user_guide/examples/online_serving/text_to_video.md
+++ b/docs/user_guide/examples/online_serving/text_to_video.md
@@ -288,6 +288,14 @@ vllm serve Lightricks/LTX-2 --omni --port 8098 \
     --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0
 ```
 
+For multi-GPU memory reduction, you can enable HSDP:
+
+```bash
+vllm serve Lightricks/LTX-2 --omni --port 8098 \
+    --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0 \
+    --use-hsdp --hsdp-shard-size 2
+```
+
 #### Start with Optimization Presets
 
 Use the LTX-2 startup script with built-in optimization presets:
diff --git a/tests/diffusion/models/ltx2/test_ltx2_hsdp.py b/tests/diffusion/models/ltx2/test_ltx2_hsdp.py
new file mode 100644
index 000000000..4dd07e1bf
--- /dev/null
+++ b/tests/diffusion/models/ltx2/test_ltx2_hsdp.py
@@ -0,0 +1,25 @@
+import pytest
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_ltx2_exposes_hsdp_shard_conditions_for_transformer_blocks():
+    model = object.__new__(LTX2VideoTransformer3DModel)
+    nn.Module.__init__(model)
+    model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
+    model.norm_out = nn.LayerNorm(4)
+
+    conditions = getattr(model, "_hsdp_shard_conditions", None)
+
+    assert conditions is not None
+    assert len(conditions) == 1
+
+    matched = []
+    for name, module in model.named_modules():
+        if any(cond(name, module) for cond in conditions):
+            matched.append(name)
+
+    assert matched == ["transformer_blocks.0", "transformer_blocks.1"]
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
index 95ef919c2..0ae71aa2d 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
@@ -41,6 +41,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 
@@ -1265,6 +1266,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
     _skip_layerwise_casting_patterns = ["norm"]
     _repeated_blocks = ["LTX2VideoTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
+    _hsdp_shard_conditions = [is_transformer_block_module]
     _sp_plan: dict[str, Any] | None = None
 
     @staticmethod