From 0d69012ef891be43db479fbe1247fad203a14679 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Tue, 21 Apr 2026 16:29:23 +0800 Subject: [PATCH] Optimize LTX2 feed-forward tensor parallelism (#23221) --- .../sglang/multimodal_gen/runtime/models/dits/ltx_2.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 84bd3cd81..ec255d63b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -764,11 +764,15 @@ class LTX2FeedForward(nn.Module): inner_dim = int(dim * mult) self.proj_in = ColumnParallelLinear( - dim, inner_dim, bias=True, gather_output=True, quant_config=quant_config + dim, inner_dim, bias=True, gather_output=False, quant_config=quant_config ) self.act = nn.GELU(approximate="tanh") - self.proj_out = ColumnParallelLinear( - inner_dim, dim_out, bias=True, gather_output=True, quant_config=quant_config + self.proj_out = RowParallelLinear( + inner_dim, + dim_out, + bias=True, + input_is_parallel=True, + quant_config=quant_config, ) def forward(self, x: torch.Tensor) -> torch.Tensor: