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: