[diffusion] fix: use Megatron-style tp for native encoders and dits (#28318)

python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py

@@ -117,6 +117,37 @@ class HunyuanConfig(PipelineConfig):
         self.vae_config.load_encoder = False
         self.vae_config.load_decoder = True
 
+    def get_text_encoder_pooler_output(self, outputs, encoder_index):
+        if encoder_index == 1:
+            return outputs.pooler_output
+        return None
+
+    def get_pos_prompt_embeds(self, batch):
+        return batch.prompt_embeds[0]
+
+    def get_neg_prompt_embeds(self, batch):
+        return batch.negative_prompt_embeds[0]
+
+    def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        prompt_attention_mask = batch.prompt_attention_mask[0]
+        return {
+            "encoder_attention_mask": prompt_attention_mask,
+            "encoder_hidden_states_mask": prompt_attention_mask,
+            "pooled_projections": (
+                batch.pooled_embeds[0] if batch.pooled_embeds else None
+            ),
+        }
+
+    def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        negative_attention_mask = batch.negative_attention_mask[0]
+        return {
+            "encoder_attention_mask": negative_attention_mask,
+            "encoder_hidden_states_mask": negative_attention_mask,
+            "pooled_projections": (
+                batch.neg_pooled_embeds[0] if batch.neg_pooled_embeds else None
+            ),
+        }
+
 
 @dataclass
 class FastHunyuanConfig(HunyuanConfig):

python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py

@@ -134,6 +134,10 @@ class DistributedAutograd:
 
             if scatter_dim == 2 and gather_dim == 1:
                 bs, shard_seqlen, hn, hd = input_.shape
+                assert hn % world_size == 0, (
+                    f"head dimension ({hn}) must be divisible by sequence "
+                    f"parallel world size ({world_size})"
+                )
                 seqlen = shard_seqlen * world_size
                 shard_hn = hn // world_size
 
@@ -155,6 +159,10 @@ class DistributedAutograd:
                 return output
             elif scatter_dim == 1 and gather_dim == 2:
                 bs, seqlen, shard_hn, hd = input_.shape
+                assert seqlen % world_size == 0, (
+                    f"sequence dimension ({seqlen}) must be divisible by sequence "
+                    f"parallel world size ({world_size})"
+                )
                 hn = shard_hn * world_size
                 shard_seqlen = seqlen // world_size
 

python/sglang/multimodal_gen/runtime/layers/linear.py

@@ -479,6 +479,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         prefix: str = "",
         tp_group: dist.ProcessGroup = None,
     ):
+        tp_group = tp_group or get_tp_group()
+        if get_group_size(tp_group) > 1:
+            self.output_sizes = output_sizes
         super().__init__(
             input_size=input_size,
             output_size=sum(output_sizes),
@@ -621,6 +624,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
 
         if loaded_shard_id is None:
             if isinstance(param, PerTensorScaleParameter):
+                if self.tp_size > 1 and loaded_weight.shape == param.data.shape:
+                    param.data.copy_(loaded_weight)
+                    return
                 param.load_merged_column_weight(loaded_weight=loaded_weight, shard_id=0)
                 return
             elif type(param) in (RowvLLMParameter, BasevLLMParameter):

python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py

@@ -10,7 +10,9 @@ import sglang.multimodal_gen.envs as envs
 from sglang.multimodal_gen.runtime.distributed import (
     divide,
     get_tp_group,
+    split_tensor_along_last_dim,
     tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
 )
 from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size
 from sglang.multimodal_gen.runtime.models.utils import set_weight_attrs
@@ -245,6 +247,156 @@ class WeightOnlyFP8ColumnParallelLinear(nn.Module):
         return output_parallel
 
 
+class WeightOnlyFP8MergedColumnParallelLinear(WeightOnlyFP8ColumnParallelLinear):
+    """Column-parallel storage-only FP8 packed linear."""
+
+    def __init__(
+        self,
+        in_features: int,
+        output_sizes: list[int],
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        gather_output: bool = False,
+        tp_group=None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        self.output_sizes = output_sizes
+        super().__init__(
+            in_features,
+            sum(output_sizes),
+            bias=bias,
+            compute_dtype=compute_dtype,
+            gather_output=gather_output,
+            tp_group=tp_group,
+            enable_fused_w8a8=enable_fused_w8a8,
+        )
+        assert all(output_size % self.tp_size == 0 for output_size in output_sizes)
+
+    def weight_loader(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        output_dim = getattr(param, "output_dim", None)
+        if output_dim is not None:
+            shards = []
+            current_offset = 0
+            for output_size in self.output_sizes:
+                loaded_shard = loaded_weight.narrow(
+                    output_dim, current_offset, output_size
+                )
+                shard_size = output_size // self.tp_size
+                loaded_shard = loaded_shard.narrow(
+                    output_dim, self.tp_rank * shard_size, shard_size
+                )
+                shards.append(loaded_shard)
+                current_offset += output_size
+            loaded_weight = torch.cat(shards, dim=output_dim)
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+        assert param.data.shape == loaded_weight.shape
+        param.data.copy_(loaded_weight)
+
+
+class WeightOnlyFP8RowParallelLinear(nn.Module):
+    """Row-parallel storage-only e4m3 FP8 linear."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        input_is_parallel: bool = True,
+        reduce_results: bool = True,
+        tp_group=None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.compute_dtype = compute_dtype
+        self.input_is_parallel = input_is_parallel
+        self.reduce_results = reduce_results
+        self.enable_fused_w8a8 = _resolve_enable_fused_w8a8(enable_fused_w8a8)
+        self.tp_group = tp_group or get_tp_group()
+        self.tp_size = get_group_size(self.tp_group)
+        self.tp_rank = get_group_rank(self.tp_group)
+        self.in_features_per_partition = divide(in_features, self.tp_size)
+        self.weight = nn.Parameter(
+            torch.empty(
+                out_features,
+                self.in_features_per_partition,
+                dtype=FP8_WEIGHT_DTYPE,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight,
+            {
+                "input_dim": 1,
+                "weight_loader": self.weight_loader,
+            },
+        )
+        self.weight_scale = nn.Parameter(
+            torch.empty(out_features, dtype=torch.float32),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight_scale,
+            {
+                "missing_param_init": "error",
+                "weight_loader": self.weight_loader,
+            },
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(
+                    out_features, dtype=compute_dtype or torch.get_default_dtype()
+                ),
+                requires_grad=False,
+            )
+            set_weight_attrs(self.bias, {"weight_loader": self.weight_loader})
+        else:
+            self.register_parameter("bias", None)
+
+    def weight_loader(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        input_dim = getattr(param, "input_dim", None)
+        if input_dim is not None:
+            shard_size = param.data.shape[input_dim]
+            loaded_weight = loaded_weight.narrow(
+                input_dim, self.tp_rank * shard_size, shard_size
+            )
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+        assert param.data.shape == loaded_weight.shape
+        param.data.copy_(loaded_weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.input_is_parallel:
+            input_parallel = x
+        else:
+            input_parallel = split_tensor_along_last_dim(
+                x, num_partitions=self.tp_size
+            )[self.tp_rank].contiguous()
+
+        compute_dtype = self.compute_dtype or x.dtype
+        bias = None if self.tp_rank > 0 else self.bias
+        output_parallel = _apply_weight_only_fp8_linear(
+            input_parallel,
+            self.weight,
+            self.weight_scale,
+            bias,
+            compute_dtype,
+            self.enable_fused_w8a8,
+        )
+        if self.reduce_results and self.tp_size > 1:
+            return tensor_model_parallel_all_reduce(
+                output_parallel, tp_group=self.tp_group
+            )
+        return output_parallel
+
+
 def _resolve_enable_fused_w8a8(value: bool | None) -> bool:
     if value is not None:
         return value

python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py

@@ -26,7 +26,12 @@ flex_attention = torch.compile(
 import torch.distributed as dist
 
 from sglang.multimodal_gen.configs.models.dits import WanVideoConfig
-from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size
+from sglang.multimodal_gen.runtime.distributed import (
+    divide,
+    get_sp_world_size,
+    get_tp_rank,
+    get_tp_world_size,
+)
 from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
 from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
 from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import (
@@ -38,8 +43,13 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
     LayerNormScaleShift,
     RMSNorm,
     ScaleResidualLayerNormScaleShift,
+    tensor_parallel_rms_norm,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
 )
-from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
@@ -59,6 +69,7 @@ from sglang.multimodal_gen.runtime.platforms import (
     current_platform,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.utils import add_prefix
 
 logger = init_logger(__name__)
 
@@ -73,12 +84,17 @@ class CausalWanSelfAttention(nn.Module):
         qk_norm=True,
         eps=1e-6,
         parallel_attention=False,
+        head_dim: int | None = None,
+        head_start: int = 0,
     ) -> None:
+        if head_dim is None:
             assert dim % num_heads == 0
+            head_dim = dim // num_heads
         super().__init__()
         self.dim = dim
         self.num_heads = num_heads
-        self.head_dim = dim // num_heads
+        self.head_dim = head_dim
+        self.head_start = head_start
         self.local_attn_size = local_attn_size
         self.sink_size = sink_size
         self.qk_norm = qk_norm
@@ -167,16 +183,46 @@ class CausalWanSelfAttention(nn.Module):
                 block_mask=block_mask,
             )[:, :, :-padded_length].transpose(2, 1)
         else:
+            head_slice = None
+            if kv_cache.k.shape[2] != roped_key.shape[2]:
+                head_slice = slice(self.head_start, self.head_start + self.num_heads)
+                cache_key = roped_key.new_zeros(
+                    roped_key.shape[0],
+                    roped_key.shape[1],
+                    kv_cache.k.shape[2],
+                    roped_key.shape[3],
+                )
+                cache_value = v.new_zeros(
+                    v.shape[0],
+                    v.shape[1],
+                    kv_cache.v.shape[2],
+                    v.shape[3],
+                )
+                cache_key[:, :, head_slice, :] = roped_key
+                cache_value[:, :, head_slice, :] = v
+            else:
+                cache_key = roped_key
+                cache_value = v
             cache_view = kv_cache.update_and_get_attention_kv(
-                key=roped_key,
+                key=cache_key,
-                value=v,
+                value=cache_value,
                 current_chunk_start=current_start,
                 debug_name="CausalWan KV cache",
             )
+            key = (
+                cache_view.k[:, :, head_slice, :]
+                if head_slice is not None
+                else cache_view.k
+            )
+            value = (
+                cache_view.v[:, :, head_slice, :]
+                if head_slice is not None
+                else cache_view.v
+            )
             x = self.attn(
                 roped_query,
-                cache_view.k,
+                key,
-                cache_view.v,
+                value,
             )
 
         return x
@@ -202,23 +248,71 @@ class CausalWanTransformerBlock(nn.Module):
 
         # 1. Self-attention
         self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        use_megatron_tp = getattr(
+            self, "_use_megatron_tp", type(self) is CausalWanTransformerBlock
+        )
+        if use_megatron_tp:
+            self.to_q = ColumnParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=add_prefix("to_q", prefix),
+            )
+            self.to_k = ColumnParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=add_prefix("to_k", prefix),
+            )
+            self.to_v = ColumnParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=add_prefix("to_v", prefix),
+            )
+
+            self.to_out = RowParallelLinear(
+                dim,
+                dim,
+                bias=True,
+                input_is_parallel=True,
+                quant_config=quant_config,
+                prefix=add_prefix("to_out", prefix),
+            )
+            tp_size = get_tp_world_size()
+            self.local_num_heads = divide(num_heads, tp_size)
+            head_start = get_tp_rank() * self.local_num_heads
+        else:
             self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
             self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
             self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
-
+            self.to_out = ReplicatedLinear(
-        self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config)
+                dim, dim, bias=True, quant_config=quant_config
+            )
+            tp_size = 1
+            self.local_num_heads = num_heads
+            head_start = 0
+        dim_head = dim // num_heads
         self.attn1 = CausalWanSelfAttention(
             dim,
-            num_heads,
+            self.local_num_heads,
             local_attn_size=local_attn_size,
             sink_size=sink_size,
             qk_norm=qk_norm,
             eps=eps,
+            head_dim=dim_head,
+            head_start=head_start,
         )
         self.hidden_dim = dim
         self.num_attention_heads = num_heads
         self.local_attn_size = local_attn_size
-        dim_head = dim // num_heads
+        self.dim_head = dim_head
         if qk_norm == "rms_norm":
             self.norm_q = RMSNorm(dim_head, eps=eps)
             self.norm_k = RMSNorm(dim_head, eps=eps)
@@ -229,6 +323,9 @@ class CausalWanTransformerBlock(nn.Module):
         else:
             print("QK Norm type not supported")
             raise Exception
+        self.tp_rmsnorm = (
+            use_megatron_tp and qk_norm == "rms_norm_across_heads" and tp_size > 1
+        )
         assert cross_attn_norm is True
         self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim, eps=eps, elementwise_affine=True, dtype=torch.float32
@@ -306,13 +403,19 @@ class CausalWanTransformerBlock(nn.Module):
         value, _ = self.to_v(norm_hidden_states)
 
         if self.norm_q is not None:
+            if self.tp_rmsnorm:
+                query = tensor_parallel_rms_norm(query, self.norm_q)
+            else:
                 query = self.norm_q(query)
         if self.norm_k is not None:
+            if self.tp_rmsnorm:
+                key = tensor_parallel_rms_norm(key, self.norm_k)
+            else:
                 key = self.norm_k(key)
 
-        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        query = query.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head))
-        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        key = key.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head))
-        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        value = value.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head))
 
         attn_output = self.attn1(
             query,

python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py

@@ -18,6 +18,7 @@ from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoC
 from sglang.multimodal_gen.runtime.distributed import (
     get_sp_group,
     get_sp_world_size,
+    get_tp_world_size,
     sequence_model_parallel_all_gather,
 )
 from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
@@ -337,6 +338,19 @@ class Cosmos3CausalAttention(nn.Module):
         self.num_attention_heads = num_attention_heads
         self.num_key_value_heads = num_key_value_heads
         self.head_dim = head_dim
+        self.tp_size = get_tp_world_size()
+        if num_attention_heads % self.tp_size != 0:
+            raise ValueError(
+                "Cosmos3CausalAttention requires num_attention_heads divisible "
+                f"by tp_size, got {num_attention_heads=} {self.tp_size=}."
+            )
+        if num_key_value_heads % self.tp_size != 0:
+            raise ValueError(
+                "Cosmos3CausalAttention requires num_key_value_heads divisible "
+                f"by tp_size, got {num_key_value_heads=} {self.tp_size=}."
+            )
+        self.local_num_attention_heads = num_attention_heads // self.tp_size
+        self.local_num_key_value_heads = num_key_value_heads // self.tp_size
 
         self.q_size = num_attention_heads * head_dim
         self.kv_size = num_key_value_heads * head_dim
@@ -344,16 +358,15 @@ class Cosmos3CausalAttention(nn.Module):
             hidden_size,
             [self.q_size, self.kv_size, self.kv_size],
             bias=False,
-            gather_output=True,
+            gather_output=False,
             quant_config=quant_config,
             prefix=add_prefix("to_qkv", prefix),
         )
-        # Output projection - ReplicatedLinear for quantization support
+        self.to_out = RowParallelLinear(
-        # Input is not parallel (gather_output=True on QKV)
-        self.to_out = ReplicatedLinear(
             num_attention_heads * head_dim,
             hidden_size,
             bias=False,
+            input_is_parallel=True,
             quant_config=quant_config,
             prefix=add_prefix("to_out", prefix),
         )
@@ -371,7 +384,7 @@ class Cosmos3CausalAttention(nn.Module):
         """Forward with KV cache return.
 
         Returns:
-            (output, K, V) where K/V are post-norm, post-RoPE
+            (output, K, V) where K/V are TP-local, post-norm, post-RoPE
         """
         batch_size, seq_len = hidden_states.shape[:2]
 
@@ -379,18 +392,23 @@ class Cosmos3CausalAttention(nn.Module):
         qkv = qkv.view(
             batch_size,
             seq_len,
-            self.num_attention_heads + 2 * self.num_key_value_heads,
+            self.local_num_attention_heads + 2 * self.local_num_key_value_heads,
             self.head_dim,
         )
-        q = qkv[:, :, : self.num_attention_heads, :]
+        q = qkv[:, :, : self.local_num_attention_heads, :]
         k = qkv[
             :,
             :,
-            self.num_attention_heads : self.num_attention_heads
+            self.local_num_attention_heads : self.local_num_attention_heads
-            + self.num_key_value_heads,
+            + self.local_num_key_value_heads,
+            :,
+        ]
+        v = qkv[
+            :,
+            :,
+            self.local_num_attention_heads + self.local_num_key_value_heads :,
             :,
         ]
-        v = qkv[:, :, self.num_attention_heads + self.num_key_value_heads :, :]
 
         q = F.rms_norm(
             q, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon
@@ -436,6 +454,19 @@ class Cosmos3CrossAttention(nn.Module):
         self.num_attention_heads = num_attention_heads
         self.num_key_value_heads = num_key_value_heads
         self.head_dim = head_dim
+        self.tp_size = get_tp_world_size()
+        if num_attention_heads % self.tp_size != 0:
+            raise ValueError(
+                "Cosmos3CrossAttention requires num_attention_heads divisible "
+                f"by tp_size, got {num_attention_heads=} {self.tp_size=}."
+            )
+        if num_key_value_heads % self.tp_size != 0:
+            raise ValueError(
+                "Cosmos3CrossAttention requires num_key_value_heads divisible "
+                f"by tp_size, got {num_key_value_heads=} {self.tp_size=}."
+            )
+        self.local_num_attention_heads = num_attention_heads // self.tp_size
+        self.local_num_key_value_heads = num_key_value_heads // self.tp_size
 
         self.q_size = num_attention_heads * head_dim
         self.kv_size = num_key_value_heads * head_dim
@@ -443,14 +474,15 @@ class Cosmos3CrossAttention(nn.Module):
             hidden_size,
             [self.q_size, self.kv_size, self.kv_size],
             bias=False,
-            gather_output=True,
+            gather_output=False,
             quant_config=quant_config,
             prefix=add_prefix("to_qkv", prefix),
         )
-        self.to_out = ReplicatedLinear(
+        self.to_out = RowParallelLinear(
             num_attention_heads * head_dim,
             hidden_size,
             bias=False,
+            input_is_parallel=True,
             quant_config=quant_config,
             prefix=add_prefix("to_out", prefix),
         )
@@ -459,9 +491,9 @@ class Cosmos3CrossAttention(nn.Module):
         self.norm_k = RMSNorm(head_dim, eps=1e-6)
 
         self.attn = USPAttention(
-            num_heads=num_attention_heads,
+            num_heads=self.local_num_attention_heads,
             head_size=head_dim,
-            num_kv_heads=num_key_value_heads,
+            num_kv_heads=self.local_num_key_value_heads,
             causal=False,
             supported_attention_backends=supported_attention_backends,
             prefix=add_prefix("attn", prefix),
@@ -480,8 +512,8 @@ class Cosmos3CrossAttention(nn.Module):
 
         Args:
             hidden_states: [B, S_gen_local, hidden_size] visual tokens (may be sharded)
-            k_und: [B, S_und, H_kv, D] pre-computed UND keys (always full/replicated)
+            k_und: [B, S_und, H_kv_local, D] UND keys (replicated over SP)
-            v_und: [B, S_und, H_kv, D] pre-computed UND values (always full/replicated)
+            v_und: [B, S_und, H_kv_local, D] UND values (replicated over SP)
             cos_sin_cache: [B*S_gen_local, D] local rows of [cos, sin]
             rope_cache_positions: identity row positions into cos_sin_cache
         """
@@ -491,18 +523,23 @@ class Cosmos3CrossAttention(nn.Module):
         qkv = qkv.view(
             batch_size,
             seq_len_gen,
-            self.num_attention_heads + 2 * self.num_key_value_heads,
+            self.local_num_attention_heads + 2 * self.local_num_key_value_heads,
             self.head_dim,
         )
-        q = qkv[:, :, : self.num_attention_heads, :]
+        q = qkv[:, :, : self.local_num_attention_heads, :]
         k = qkv[
             :,
             :,
-            self.num_attention_heads : self.num_attention_heads
+            self.local_num_attention_heads : self.local_num_attention_heads
-            + self.num_key_value_heads,
+            + self.local_num_key_value_heads,
+            :,
+        ]
+        v = qkv[
+            :,
+            :,
+            self.local_num_attention_heads + self.local_num_key_value_heads :,
             :,
         ]
-        v = qkv[:, :, self.num_attention_heads + self.num_key_value_heads :, :]
 
         if use_fused_qk_norm_rope:
             q, k = _apply_qwen3_qk_norm_rope(
@@ -519,7 +556,7 @@ class Cosmos3CrossAttention(nn.Module):
                 q, k, self.norm_q, self.norm_k, self.head_dim, cos_sin_cache
             )
 
-        # K/V = [text (replicated full on every SP rank) | image (sharded same as Q)].
+        # K/V = [text (replicated on every SP rank) | image (sharded same as Q)].
         # USPAttention routes through the registered attention backend (FA, sage,
         # …) and handles the Ulysses all-to-all when SP > 1.
         out = self.attn.forward_with_replicated_kv_prefix(q, k_und, v_und, k, v)
@@ -1179,7 +1216,8 @@ class Cosmos3OmniTransformer(CachableDiT):
         # pick max as the fused scale, requant each shard against the max,
         # then concat the requantized FP8 bytes. input_scale is shared across
         # shards (same activation tensor), so just take max — no requant
-        # needed.
+        # needed. The emitted fused tensors keep the full Q/K/V layout; the
+        # column-parallel loader slices each logical shard for the local TP rank.
         mapping_fn = get_param_names_mapping(self.param_names_mapping)
         pending: dict[str, dict[str, dict[int, torch.Tensor]]] = {}
         expected_count: dict[str, int] = {}

python/sglang/multimodal_gen/runtime/models/dits/glm_image.py

@@ -22,12 +22,17 @@ from sglang.multimodal_gen.configs.models.dits.glmimage import GlmImageDitConfig
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_sp_parallel_rank,
     get_sp_world_size,
+    get_tp_world_size,
 )
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
 from sglang.multimodal_gen.runtime.layers.layernorm import (
     ScaleResidualLayerNormScaleShift,
 )
-from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
 from sglang.multimodal_gen.runtime.layers.mlp import FeedForward
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
@@ -307,6 +312,74 @@ class GlmImageAdaLayerNormZero(nn.Module):
         )
 
 
+class GlmImageGELU(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        inner_dim: int,
+        bias: bool = True,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.proj = ColumnParallelLinear(
+            dim,
+            inner_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj" if prefix else "proj",
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.proj(hidden_states)
+        return F.gelu(hidden_states, approximate="tanh")
+
+
+class GlmImageFeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        inner_dim: Optional[int] = None,
+        bias: bool = True,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        if inner_dim is None:
+            inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        self.net = nn.ModuleList(
+            [
+                GlmImageGELU(
+                    dim,
+                    inner_dim,
+                    bias=bias,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.net.0" if prefix else "net.0",
+                ),
+                nn.Dropout(0.0),
+                RowParallelLinear(
+                    inner_dim,
+                    dim_out,
+                    bias=bias,
+                    input_is_parallel=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.net.2" if prefix else "net.2",
+                ),
+            ]
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.net[0](hidden_states)
+        hidden_states = self.net[1](hidden_states)
+        hidden_states, _ = self.net[2](hidden_states)
+        return hidden_states
+
+
 class GlmImageAttention(torch.nn.Module):
     def __init__(
         self,
@@ -333,23 +406,48 @@ class GlmImageAttention(torch.nn.Module):
         self.inner_kv_dim = self.inner_dim
         self.out_dim = out_dim if out_dim is not None else query_dim
 
-        self.num_kv_heads = self.dim_head // self.inner_kv_dim
+        tp_size = get_tp_world_size()
+        assert (
+            self.heads % tp_size == 0
+        ), f"heads ({self.heads}) must be divisible by tp_size ({tp_size})"
+        self.num_local_heads = self.heads // tp_size
+        self.num_local_kv_heads = self.num_local_heads
 
-        self.to_q = ReplicatedLinear(
+        self.to_q = ColumnParallelLinear(
-            query_dim, self.inner_dim, bias=bias, quant_config=quant_config
+            query_dim,
+            self.inner_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_q" if prefix else "to_q",
         )
-        self.to_k = ReplicatedLinear(
+        self.to_k = ColumnParallelLinear(
-            query_dim, self.inner_kv_dim, bias=bias, quant_config=quant_config
+            query_dim,
+            self.inner_kv_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_k" if prefix else "to_k",
         )
-        self.to_v = ReplicatedLinear(
+        self.to_v = ColumnParallelLinear(
-            query_dim, self.inner_kv_dim, bias=bias, quant_config=quant_config
+            query_dim,
+            self.inner_kv_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_v" if prefix else "to_v",
         )
 
         # (dropout omitted)
         self.to_out = nn.ModuleList(
             [
-                ReplicatedLinear(
+                RowParallelLinear(
-                    self.inner_dim, self.out_dim, bias=True, quant_config=quant_config
+                    self.inner_dim,
+                    self.out_dim,
+                    bias=True,
+                    input_is_parallel=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0" if prefix else "to_out.0",
                 )
             ]
         )
@@ -370,9 +468,9 @@ class GlmImageAttention(torch.nn.Module):
             )
 
         self.attn = USPAttention(
-            num_heads=self.heads,
+            num_heads=self.num_local_heads,
             head_size=dim_head,
-            num_kv_heads=self.num_kv_heads,
+            num_kv_heads=self.num_local_kv_heads,
             dropout_rate=0,
             softmax_scale=None,
             causal=False,
@@ -397,9 +495,9 @@ class GlmImageAttention(torch.nn.Module):
         key, _ = self.to_k(hidden_states)
         value, _ = self.to_v(hidden_states)
 
-        query = query.unflatten(2, (self.heads, -1))
+        query = query.unflatten(2, (self.num_local_heads, -1))
-        key = key.unflatten(2, (self.heads, -1))
+        key = key.unflatten(2, (self.num_local_kv_heads, -1))
-        value = value.unflatten(2, (self.heads, -1))
+        value = value.unflatten(2, (self.num_local_kv_heads, -1))
 
         # 2. QK normalization
         if self.norm_q is not None:
@@ -504,7 +602,12 @@ class GlmImageTransformerBlock(nn.Module):
         self.norm2_context = ScaleResidualLayerNormScaleShift(
             dim, eps=1e-5, elementwise_affine=False
         )
-        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        self.ff = GlmImageFeedForward(
+            dim=dim,
+            dim_out=dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ff" if prefix else "ff",
+        )
 
     def forward(
         self,

python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py

@@ -10,6 +10,7 @@ import torch.nn as nn
 
 from sglang.multimodal_gen.configs.models.dits import HunyuanVideoConfig
 from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams
+from sglang.multimodal_gen.runtime.distributed import divide, get_tp_world_size
 from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size
 from sglang.multimodal_gen.runtime.layers.attention import (
     LocalAttention,
@@ -21,7 +22,11 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
     RMSNorm,
     ScaleResidualLayerNormScaleShift,
 )
-from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
+from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
@@ -48,6 +53,33 @@ from sglang.multimodal_gen.runtime.platforms import (
 )
 
 
+class MixedRowParallelLinear(RowParallelLinear):
+    def __init__(self, input_sizes: list[int], output_size: int, **kwargs):
+        self.input_sizes = input_sizes
+        super().__init__(sum(input_sizes), output_size, **kwargs)
+
+    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+        input_dim = getattr(param, "input_dim", None)
+        if input_dim is not None:
+            shards = []
+            offset = 0
+            for input_size in self.input_sizes:
+                loaded_shard = loaded_weight.narrow(input_dim, offset, input_size)
+                shard_size = input_size // self.tp_size
+                loaded_shard = loaded_shard.narrow(
+                    input_dim, self.tp_rank * shard_size, shard_size
+                )
+                shards.append(loaded_shard)
+                offset += input_size
+            loaded_weight = torch.cat(shards, dim=input_dim)
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+        param.data.copy_(loaded_weight)
+
+    def weight_loader_v2(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+        self.weight_loader(param, loaded_weight)
+
+
 class MMDoubleStreamBlock(nn.Module):
     """
     A multimodal DiT block with separate modulation for text and image/video,
@@ -68,6 +100,8 @@ class MMDoubleStreamBlock(nn.Module):
 
         self.deterministic = False
         self.num_attention_heads = num_attention_heads
+        tp_size = get_tp_world_size()
+        self.local_num_attention_heads = divide(num_attention_heads, tp_size)
         head_dim = hidden_size // num_attention_heads
         mlp_hidden_dim = int(hidden_size * mlp_ratio)
 
@@ -90,23 +124,24 @@ class MMDoubleStreamBlock(nn.Module):
         self.img_mlp_residual = MulAdd()
 
         # Image attention components
-        self.img_attn_qkv = ReplicatedLinear(
+        self.img_attn_qkv = MergedColumnParallelLinear(
             hidden_size,
-            hidden_size * 3,
+            [hidden_size] * 3,
             bias=True,
+            gather_output=False,
             params_dtype=dtype,
             prefix=f"{prefix}.img_attn_qkv",
             quant_config=quant_config,
-            output_sizes=[hidden_size] * 3,
         )
 
         self.img_attn_q_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype)
         self.img_attn_k_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype)
 
-        self.img_attn_proj = ReplicatedLinear(
+        self.img_attn_proj = RowParallelLinear(
             hidden_size,
             hidden_size,
             bias=True,
+            input_is_parallel=True,
             params_dtype=dtype,
             prefix=f"{prefix}.img_attn_proj",
             quant_config=quant_config,
@@ -140,24 +175,25 @@ class MMDoubleStreamBlock(nn.Module):
         self.txt_mlp_residual = MulAdd()
 
         # Text attention components
-        self.txt_attn_qkv = ReplicatedLinear(
+        self.txt_attn_qkv = MergedColumnParallelLinear(
             hidden_size,
-            hidden_size * 3,
+            [hidden_size] * 3,
             bias=True,
+            gather_output=False,
             params_dtype=dtype,
             prefix=f"{prefix}.txt_attn_qkv",
             quant_config=quant_config,
-            output_sizes=[hidden_size] * 3,
         )
 
         # QK norm layers for text
         self.txt_attn_q_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype)
         self.txt_attn_k_norm = RMSNorm(head_dim, eps=1e-6, dtype=dtype)
 
-        self.txt_attn_proj = ReplicatedLinear(
+        self.txt_attn_proj = RowParallelLinear(
             hidden_size,
             hidden_size,
             bias=True,
+            input_is_parallel=True,
             params_dtype=dtype,
             prefix=f"{prefix}.txt_attn_proj",
             quant_config=quant_config,
@@ -174,7 +210,7 @@ class MMDoubleStreamBlock(nn.Module):
 
         # Use UlyssesAttention to replace Distributed attention
         self.attn = UlyssesAttention(
-            num_heads=num_attention_heads,
+            num_heads=self.local_num_attention_heads,
             head_size=head_dim,
             causal=False,
             supported_attention_backends=supported_attention_backends,
@@ -217,7 +253,7 @@ class MMDoubleStreamBlock(nn.Module):
 
         # Split QKV
         img_qkv = img_qkv.view(
-            batch_size, image_seq_len, 3, self.num_attention_heads, -1
+            batch_size, image_seq_len, 3, self.local_num_attention_heads, -1
         )
         img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2]
 
@@ -240,7 +276,7 @@ class MMDoubleStreamBlock(nn.Module):
 
         # Split QKV
         txt_qkv = txt_qkv.view(
-            batch_size, text_seq_len, 3, self.num_attention_heads, -1
+            batch_size, text_seq_len, 3, self.local_num_attention_heads, -1
         )
         txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2]
 
@@ -300,26 +336,31 @@ class MMSingleStreamBlock(nn.Module):
         self.deterministic = False
         self.hidden_size = hidden_size
         self.num_attention_heads = num_attention_heads
+        tp_size = get_tp_world_size()
+        self.local_num_attention_heads = divide(num_attention_heads, tp_size)
         head_dim = hidden_size // num_attention_heads
+        self.head_dim = head_dim
         mlp_hidden_dim = int(hidden_size * mlp_ratio)
         self.mlp_hidden_dim = mlp_hidden_dim
+        self.local_mlp_hidden_dim = divide(mlp_hidden_dim, tp_size)
 
         # Combined QKV and MLP input projection
-        self.linear1 = ReplicatedLinear(
+        self.linear1 = MergedColumnParallelLinear(
             hidden_size,
-            hidden_size * 3 + mlp_hidden_dim,
+            [hidden_size] * 3 + [mlp_hidden_dim],
             bias=True,
+            gather_output=False,
             params_dtype=dtype,
             prefix=f"{prefix}.linear1",
             quant_config=quant_config,
-            output_sizes=[hidden_size] * 3 + [mlp_hidden_dim],
         )
 
         # Combined projection and MLP output
-        self.linear2 = ReplicatedLinear(
+        self.linear2 = MixedRowParallelLinear(
-            hidden_size + mlp_hidden_dim,
+            [hidden_size, mlp_hidden_dim],
             hidden_size,
             bias=True,
+            input_is_parallel=True,
             params_dtype=dtype,
             prefix=f"{prefix}.linear2",
             quant_config=quant_config,
@@ -352,7 +393,7 @@ class MMSingleStreamBlock(nn.Module):
 
         # Use UlyssesAttention to replace Distributed attention
         self.attn = UlyssesAttention(
-            num_heads=num_attention_heads,
+            num_heads=self.local_num_attention_heads,
             head_size=head_dim,
             causal=False,
             supported_attention_backends=supported_attention_backends,
@@ -376,13 +417,16 @@ class MMSingleStreamBlock(nn.Module):
         linear1_out, _ = self.linear1(x_mod)
 
         # Split into QKV and MLP parts
+        local_qkv_dim = 3 * self.local_num_attention_heads * self.head_dim
         qkv, mlp = torch.split(
-            linear1_out, [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1
+            linear1_out,
+            [local_qkv_dim, self.local_mlp_hidden_dim],
+            dim=-1,
         )
 
         # Process QKV
         batch_size, seq_len = qkv.shape[0], qkv.shape[1]
-        qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
+        qkv = qkv.view(batch_size, seq_len, 3, self.local_num_attention_heads, -1)
         q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
 
         # Apply QK-Norm
@@ -575,6 +619,7 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi
         encoder_hidden_states: torch.Tensor | list[torch.Tensor],
         timestep: torch.LongTensor,
         encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+        pooled_projections: torch.Tensor | None = None,
         guidance=None,
         **kwargs,
     ):
@@ -604,8 +649,12 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi
 
         # Split text embeddings - first token is global, rest are per-token
         if isinstance(encoder_hidden_states, torch.Tensor):
+            if pooled_projections is None:
                 txt = encoder_hidden_states[:, 1:]
                 text_states_2 = encoder_hidden_states[:, 0, : self.text_states_dim_2]
+            else:
+                txt = encoder_hidden_states
+                text_states_2 = pooled_projections
         else:
             txt = encoder_hidden_states[0]
             text_states_2 = encoder_hidden_states[1]
@@ -883,25 +932,30 @@ class IndividualTokenRefinerBlock(nn.Module):
     ) -> None:
         super().__init__()
         self.num_attention_heads = num_attention_heads
+        tp_size = get_tp_world_size()
+        self.local_num_attention_heads = divide(num_attention_heads, tp_size)
         mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        head_dim = hidden_size // num_attention_heads
 
         # Normalization and attention
         self.norm1 = nn.LayerNorm(
             hidden_size, eps=1e-6, elementwise_affine=True, dtype=dtype
         )
 
-        self.self_attn_qkv = ReplicatedLinear(
+        self.self_attn_qkv = MergedColumnParallelLinear(
             hidden_size,
-            hidden_size * 3,
+            [hidden_size] * 3,
             bias=qkv_bias,
+            gather_output=False,
             params_dtype=dtype,
             prefix=f"{prefix}.self_attn_qkv",
         )
 
-        self.self_attn_proj = ReplicatedLinear(
+        self.self_attn_proj = RowParallelLinear(
             hidden_size,
             hidden_size,
             bias=qkv_bias,
+            input_is_parallel=True,
             params_dtype=dtype,
             prefix=f"{prefix}.self_attn_proj",
         )
@@ -930,8 +984,8 @@ class IndividualTokenRefinerBlock(nn.Module):
 
         # Scaled dot product attention
         self.attn = LocalAttention(
-            num_heads=num_attention_heads,
+            num_heads=self.local_num_attention_heads,
-            head_size=hidden_size // num_attention_heads,
+            head_size=head_dim,
             # TODO: remove hardcode; remove STA
             supported_attention_backends=(
                 AttentionBackendEnum.FA,
@@ -948,7 +1002,7 @@ class IndividualTokenRefinerBlock(nn.Module):
         qkv, _ = self.self_attn_qkv(norm_x)
 
         batch_size, seq_len = qkv.shape[0], qkv.shape[1]
-        qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
+        qkv = qkv.view(batch_size, seq_len, 3, self.local_num_attention_heads, -1)
         q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
 
         # Run scaled dot product attention

python/sglang/multimodal_gen/runtime/models/dits/ideogram.py

@@ -9,6 +9,7 @@ import torch.nn.functional as F
 
 from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
 from sglang.multimodal_gen.runtime.distributed import (
+    divide,
     get_tp_world_size,
     model_parallel_is_initialized,
 )
@@ -18,7 +19,9 @@ from sglang.multimodal_gen.runtime.layers.attention import (
 )
 from sglang.multimodal_gen.runtime.layers.linear import (
     ColumnParallelLinear,
+    MergedColumnParallelLinear,
     ReplicatedLinear,
+    RowParallelLinear,
 )
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
@@ -26,6 +29,8 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor
 from sglang.multimodal_gen.runtime.layers.quantization.weight_only_fp8 import (
     WeightOnlyFP8ColumnParallelLinear,
     WeightOnlyFP8Linear,
+    WeightOnlyFP8MergedColumnParallelLinear,
+    WeightOnlyFP8RowParallelLinear,
 )
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     Qwen3VLTextRotaryEmbedding,
@@ -57,14 +62,29 @@ class Ideogram4ColumnParallelLinear(ColumnParallelLinear):
         return super().forward(x)[0]
 
 
+class Ideogram4MergedColumnParallelLinear(MergedColumnParallelLinear):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x)[0]
+
+
+class Ideogram4RowParallelLinear(RowParallelLinear):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x)[0]
+
+
+def _tp_size() -> int:
+    return get_tp_world_size() if model_parallel_is_initialized() else 1
+
+
 def _linear(
     in_features: int,
     out_features: int,
     bias: bool = True,
     quant_config: QuantizationConfig | None = None,
     prefix: str = "",
+    gather_output: bool = True,
 ):
-    tp_size = get_tp_world_size() if model_parallel_is_initialized() else 1
+    tp_size = _tp_size()
     use_column_parallel = tp_size > 1 and out_features % tp_size == 0
     if quant_config is None:
         if use_column_parallel:
@@ -72,7 +92,7 @@ def _linear(
                 in_features,
                 out_features,
                 bias=bias,
-                gather_output=True,
+                gather_output=gather_output,
             )
         return WeightOnlyFP8Linear(in_features, out_features, bias=bias)
     if use_column_parallel:
@@ -80,7 +100,82 @@ def _linear(
             in_features,
             out_features,
             bias=bias,
-            gather_output=True,
+            gather_output=gather_output,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+    return Ideogram4QuantizedLinear(
+        in_features,
+        out_features,
+        bias=bias,
+        quant_config=quant_config,
+        prefix=prefix,
+    )
+
+
+def _merged_column_linear(
+    in_features: int,
+    output_sizes: list[int],
+    bias: bool = True,
+    quant_config: QuantizationConfig | None = None,
+    prefix: str = "",
+):
+    tp_size = _tp_size()
+    use_column_parallel = tp_size > 1 and all(
+        output_size % tp_size == 0 for output_size in output_sizes
+    )
+    out_features = sum(output_sizes)
+    if quant_config is None:
+        if use_column_parallel:
+            return WeightOnlyFP8MergedColumnParallelLinear(
+                in_features,
+                output_sizes,
+                bias=bias,
+                gather_output=False,
+            )
+        return WeightOnlyFP8Linear(in_features, out_features, bias=bias)
+    if use_column_parallel:
+        return Ideogram4MergedColumnParallelLinear(
+            in_features,
+            output_sizes,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+    return Ideogram4QuantizedLinear(
+        in_features,
+        out_features,
+        bias=bias,
+        quant_config=quant_config,
+        prefix=prefix,
+    )
+
+
+def _row_linear(
+    in_features: int,
+    out_features: int,
+    bias: bool = True,
+    quant_config: QuantizationConfig | None = None,
+    prefix: str = "",
+):
+    tp_size = _tp_size()
+    use_row_parallel = tp_size > 1 and in_features % tp_size == 0
+    if quant_config is None:
+        if use_row_parallel:
+            return WeightOnlyFP8RowParallelLinear(
+                in_features,
+                out_features,
+                bias=bias,
+                input_is_parallel=True,
+            )
+        return WeightOnlyFP8Linear(in_features, out_features, bias=bias)
+    if use_row_parallel:
+        return Ideogram4RowParallelLinear(
+            in_features,
+            out_features,
+            bias=bias,
+            input_is_parallel=True,
             quant_config=quant_config,
             prefix=prefix,
         )
@@ -107,9 +202,12 @@ class Ideogram4Attention(nn.Module):
         self.hidden_size = hidden_size
         self.num_heads = num_heads
         self.head_dim = hidden_size // num_heads
-        self.qkv = _linear(
+        tp_size = _tp_size()
+        assert num_heads % tp_size == 0
+        self.local_num_heads = divide(num_heads, tp_size)
+        self.qkv = _merged_column_linear(
             hidden_size,
-            hidden_size * 3,
+            [hidden_size, hidden_size, hidden_size],
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.qkv",
@@ -117,14 +215,14 @@ class Ideogram4Attention(nn.Module):
         self.norm_q = Ideogram4RMSNorm(self.head_dim, eps=eps)
         self.norm_k = Ideogram4RMSNorm(self.head_dim, eps=eps)
         self.attn = USPAttention(
-            num_heads=num_heads,
+            num_heads=self.local_num_heads,
             head_size=self.head_dim,
             dropout_rate=0,
             softmax_scale=None,
             causal=False,
             supported_attention_backends=supported_attention_backends,
         )
-        self.o = _linear(
+        self.o = _row_linear(
             hidden_size,
             hidden_size,
             bias=False,
@@ -134,13 +232,15 @@ class Ideogram4Attention(nn.Module):
 
     def forward(self, x, cos, sin, attn_mask, attn_mask_meta):
         batch_size, seq_len, _ = x.shape
-        qkv = self.qkv(x).view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
+        qkv = self.qkv(x).view(
+            batch_size, seq_len, 3, self.local_num_heads, self.head_dim
+        )
         q, k, v = qkv.unbind(dim=2)
         q = self.norm_q(q)
         k = self.norm_k(k)
         q, k = qwen3_apply_rotary_pos_emb(q, k, cos, sin)
         out = self.attn(q, k, v, attn_mask=attn_mask, attn_mask_meta=attn_mask_meta)
-        out = out.reshape(batch_size, seq_len, self.hidden_size)
+        out = out.reshape(batch_size, seq_len, self.local_num_heads * self.head_dim)
         return self.o(out)
 
 
@@ -159,8 +259,9 @@ class Ideogram4MLP(nn.Module):
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.w1",
+            gather_output=False,
         )
-        self.w2 = _linear(
+        self.w2 = _row_linear(
             hidden_dim,
             dim,
             bias=False,
@@ -173,6 +274,7 @@ class Ideogram4MLP(nn.Module):
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.w3",
+            gather_output=False,
         )
 
     def forward(self, x):

python/sglang/multimodal_gen/runtime/models/dits/joy_image.py

@@ -10,8 +10,10 @@ from einops import rearrange
 
 from sglang.multimodal_gen.configs.models.dits.joy_image import JoyImageDiTConfig
 from sglang.multimodal_gen.runtime.distributed import (
+    divide,
     get_sp_group,
     get_sp_world_size,
+    get_tp_world_size,
     sequence_model_parallel_all_gather,
 )
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
@@ -20,7 +22,11 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
     RMSNorm,
     apply_qk_norm_with_optional_rope,
 )
-from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
+from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
@@ -104,6 +110,8 @@ class MMDoubleStreamBlock(nn.Module):
         super().__init__()
         self.heads_num = heads_num
         self.hidden_size = hidden_size
+        self.tp_size = get_tp_world_size()
+        self.local_heads_num = divide(self.heads_num, self.tp_size)
         self.head_dim = self.hidden_size // self.heads_num
         self.mlp_hidden_dim = int(self.hidden_size * mlp_width_ratio)
 
@@ -114,10 +122,11 @@ class MMDoubleStreamBlock(nn.Module):
             elementwise_affine=False,
         )
 
-        self.img_attn_qkv = ReplicatedLinear(
+        self.img_attn_qkv = MergedColumnParallelLinear(
             self.hidden_size,
-            hidden_size * 3,
+            [hidden_size, hidden_size, hidden_size],
             bias=True,
+            gather_output=False,
             quant_config=quant_config,
             prefix=f"{prefix}.img_attn_qkv",
         )
@@ -129,10 +138,11 @@ class MMDoubleStreamBlock(nn.Module):
             self.head_dim,
             eps=1e-6,
         )
-        self.img_attn_proj = ReplicatedLinear(
+        self.img_attn_proj = RowParallelLinear(
             self.hidden_size,
             hidden_size,
             bias=True,
+            input_is_parallel=True,
             quant_config=quant_config,
             prefix=f"{prefix}.img_attn_proj",
         )
@@ -157,10 +167,11 @@ class MMDoubleStreamBlock(nn.Module):
             eps=1e-6,
             elementwise_affine=False,
         )
-        self.txt_attn_qkv = ReplicatedLinear(
+        self.txt_attn_qkv = MergedColumnParallelLinear(
             self.hidden_size,
-            self.hidden_size * 3,
+            [self.hidden_size, self.hidden_size, self.hidden_size],
             bias=True,
+            gather_output=False,
             quant_config=quant_config,
             prefix=f"{prefix}.txt_attn_qkv",
         )
@@ -172,10 +183,11 @@ class MMDoubleStreamBlock(nn.Module):
             self.head_dim,
             eps=1e-6,
         )
-        self.txt_attn_proj = ReplicatedLinear(
+        self.txt_attn_proj = RowParallelLinear(
             self.hidden_size,
             self.hidden_size,
             bias=True,
+            input_is_parallel=True,
             quant_config=quant_config,
             prefix=f"{prefix}.txt_attn_proj",
         )
@@ -193,7 +205,7 @@ class MMDoubleStreamBlock(nn.Module):
             prefix=f"{prefix}.txt_mlp",
         )
         self.attn = USPAttention(
-            num_heads=self.heads_num,
+            num_heads=self.local_heads_num,
             head_size=self.head_dim,
             causal=False,
             supported_attention_backends=supported_attention_backends,
@@ -233,7 +245,7 @@ class MMDoubleStreamBlock(nn.Module):
         )
         img_qkv, _ = self.img_attn_qkv(img_modulated)
         img_q, img_k, img_v = rearrange(
-            img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num
+            img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.local_heads_num
         )
 
         if vis_freqs_cis is None:
@@ -266,7 +278,7 @@ class MMDoubleStreamBlock(nn.Module):
         )
         txt_qkv, _ = self.txt_attn_qkv(txt_modulated)
         txt_q, txt_k, txt_v = rearrange(
-            txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num
+            txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.local_heads_num
         )
 
         if txt_freqs_cis is not None and not (

python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py

@@ -18,6 +18,7 @@ from sglang.multimodal_gen.runtime.distributed import (
     get_sp_group,
     get_sp_parallel_rank,
     get_sp_world_size,
+    get_tp_rank,
     get_tp_world_size,
     sequence_model_parallel_all_gather,
 )
@@ -281,19 +282,54 @@ class LingBotWorldCausalSelfAttention(CausalWanSelfAttention):
             )
             roped_query, roped_key, v = qkv.chunk(3, dim=-1)
 
+        head_slice = None
+        if kv_cache.k.shape[2] != roped_key.shape[2]:
+            if sequence_shard_enabled:
+                head_start = get_tp_rank() * roped_key.shape[2]
+            else:
+                head_start = self.head_start
+            head_slice = slice(head_start, head_start + roped_key.shape[2])
+            cache_key = roped_key.new_zeros(
+                roped_key.shape[0],
+                roped_key.shape[1],
+                kv_cache.k.shape[2],
+                roped_key.shape[3],
+            )
+            cache_value = v.new_zeros(
+                v.shape[0],
+                v.shape[1],
+                kv_cache.v.shape[2],
+                v.shape[3],
+            )
+            cache_key[:, :, head_slice, :] = roped_key
+            cache_value[:, :, head_slice, :] = v
+        else:
+            cache_key = roped_key
+            cache_value = v
+
         cache_view = kv_cache.update_and_get_attention_kv(
-            key=roped_key,
+            key=cache_key,
-            value=v,
+            value=cache_value,
             current_chunk_start=current_start,
             debug_name="LingBot KV cache",
         )
         if update_cache_only:
             return v
+        key = (
+            cache_view.k[:, :, head_slice, :]
+            if head_slice is not None
+            else cache_view.k
+        )
+        value = (
+            cache_view.v[:, :, head_slice, :]
+            if head_slice is not None
+            else cache_view.v
+        )
         attn_impl = self.ulysses_attn if sequence_shard_enabled else self.attn
         x = attn_impl(
             roped_query,
-            cache_view.k,
+            key,
-            cache_view.v,
+            value,
         )
         if sequence_shard_enabled:
             assert seq_splits is not None
@@ -862,15 +898,20 @@ class LingBotWorldTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi
 
 
 class CausalLingBotWorldTransformerBlock(CausalWanTransformerBlock):
+    _use_megatron_tp = True
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        head_start = self.attn1.head_start
         self.attn1 = LingBotWorldCausalSelfAttention(
             dim=self.hidden_dim,
-            num_heads=self.num_attention_heads,
+            num_heads=self.local_num_heads,
             local_attn_size=self.local_attn_size,
             sink_size=self.attn1.sink_size,
             qk_norm=self.attn1.qk_norm,
             eps=self.attn1.eps,
+            head_dim=self.dim_head,
+            head_start=head_start,
         )
         self.cam_conditioner = LingBotWorldCamConditioner(self.hidden_dim)
         self._fused_qkv_weight = None
@@ -1042,11 +1083,15 @@ class CausalLingBotWorldTransformerBlock(CausalWanTransformerBlock):
             .to(orig_dtype)
         )
         query, key, value = self._project_qkv(norm_hidden_states)
+        if self.tp_rmsnorm:
+            query = tensor_parallel_rms_norm(query, self.norm_q)
+            key = tensor_parallel_rms_norm(key, self.norm_k)
+        else:
             query = self.norm_q(query)
             key = self.norm_k(key)
-        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        query = query.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head))
-        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        key = key.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head))
-        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        value = value.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head))
 
         attn_output = self.attn1(
             query,

python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py

@@ -21,6 +21,7 @@ import torch
 from torch import nn
 from torch.nn.attention import SDPBackend, sdpa_kernel
 from transformers import Cache, DynamicCache, LlavaConfig, Mistral3Config, MistralConfig
+from transformers.activations import ACT2FN
 from transformers.masking_utils import (
     create_causal_mask,
     create_sliding_window_causal_mask,
@@ -32,7 +33,6 @@ from transformers.models.mistral3.modeling_mistral3 import (
     Mistral3ModelOutputWithPast,
 )
 from transformers.models.mistral.modeling_mistral import (
-    MistralMLP,
     MistralPreTrainedModel,
     MistralRMSNorm,
     MistralRotaryEmbedding,
@@ -40,6 +40,14 @@ from transformers.models.mistral.modeling_mistral import (
     eager_attention_forward,
 )
 
+from sglang.multimodal_gen.runtime.distributed import (
+    get_tp_world_size,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
 from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
 from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
     LayerwiseOffloadableModuleMixin,
@@ -55,6 +63,50 @@ _CREATE_CAUSAL_MASK_ARG = (
 )
 
 
+def _tp_world_size() -> int:
+    if not model_parallel_is_initialized():
+        return 1
+    return get_tp_world_size()
+
+
+def _linear_output(linear: nn.Module, x: torch.Tensor) -> torch.Tensor:
+    output = linear(x)
+    return output[0] if isinstance(output, tuple) else output
+
+
+def _make_column_linear(
+    in_features: int,
+    out_features: int,
+    *,
+    bias: bool,
+    use_tensor_parallel: bool,
+):
+    if use_tensor_parallel:
+        return ColumnParallelLinear(
+            in_features,
+            out_features,
+            bias=bias,
+            gather_output=False,
+        )
+    return nn.Linear(in_features, out_features, bias=bias)
+
+
+def _make_row_linear(
+    in_features: int,
+    out_features: int,
+    *,
+    bias: bool,
+    use_tensor_parallel: bool,
+):
+    if use_tensor_parallel:
+        return RowParallelLinear(
+            in_features,
+            out_features,
+            bias=bias,
+        )
+    return nn.Linear(in_features, out_features, bias=bias)
+
+
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     """
     This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
@@ -81,26 +133,56 @@ class MistralAttention(nn.Module):
             getattr(config, "head_dim", None)
             or config.hidden_size // config.num_attention_heads
         )
-        self.num_key_value_groups = (
-            config.num_attention_heads // config.num_key_value_heads
-        )
         self.scaling = self.head_dim**-0.5
         self.attention_dropout = config.attention_dropout
-        self.q_proj = nn.Linear(
+        self.total_num_heads = config.num_attention_heads
-            config.hidden_size, config.num_attention_heads * self.head_dim, bias=False
+        self.total_num_key_value_heads = config.num_key_value_heads
+        tp_size = _tp_world_size()
+        q_size = self.total_num_heads * self.head_dim
+        kv_size = self.total_num_key_value_heads * self.head_dim
+        self.use_tensor_parallel = (
+            tp_size > 1
+            and self.total_num_heads % tp_size == 0
+            and self.total_num_key_value_heads % tp_size == 0
+            and q_size % tp_size == 0
+            and kv_size % tp_size == 0
         )
-        self.k_proj = nn.Linear(
+        self.num_heads = (
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False
+            self.total_num_heads // tp_size
+            if self.use_tensor_parallel
+            else self.total_num_heads
         )
-        self.v_proj = nn.Linear(
+        self.num_key_value_heads = (
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False
+            self.total_num_key_value_heads // tp_size
+            if self.use_tensor_parallel
+            else self.total_num_key_value_heads
         )
-        self.o_proj = nn.Linear(
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
-            config.num_attention_heads * self.head_dim, config.hidden_size, bias=False
+        self.q_proj = _make_column_linear(
+            config.hidden_size,
+            q_size,
+            bias=False,
+            use_tensor_parallel=self.use_tensor_parallel,
+        )
+        self.k_proj = _make_column_linear(
+            config.hidden_size,
+            kv_size,
+            bias=False,
+            use_tensor_parallel=self.use_tensor_parallel,
+        )
+        self.v_proj = _make_column_linear(
+            config.hidden_size,
+            kv_size,
+            bias=False,
+            use_tensor_parallel=self.use_tensor_parallel,
+        )
+        self.o_proj = _make_row_linear(
+            q_size,
+            config.hidden_size,
+            bias=False,
+            use_tensor_parallel=self.use_tensor_parallel,
         )
         self.is_causal = True
-        self.num_heads = config.num_attention_heads
-        self.num_key_value_heads = config.num_key_value_heads
 
     def forward(
         self,
@@ -114,9 +196,21 @@ class MistralAttention(nn.Module):
         input_shape = hidden_states.shape[:-1]
         hidden_shape = (*input_shape, -1, self.head_dim)
 
-        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        query_states = (
-        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            _linear_output(self.q_proj, hidden_states)
-        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            .view(hidden_shape)
+            .transpose(1, 2)
+        )
+        key_states = (
+            _linear_output(self.k_proj, hidden_states)
+            .view(hidden_shape)
+            .transpose(1, 2)
+        )
+        value_states = (
+            _linear_output(self.v_proj, hidden_states)
+            .view(hidden_shape)
+            .transpose(1, 2)
+        )
 
         cos, sin = position_embeddings
         query_states, key_states = apply_rotary_pos_emb(
@@ -154,16 +248,48 @@ class MistralAttention(nn.Module):
         )
 
         attn_output = attn_output.reshape(*input_shape, -1).contiguous()
-        attn_output = self.o_proj(attn_output)
+        attn_output = _linear_output(self.o_proj, attn_output)
         return attn_output, attn_weights
 
 
+class MistralTPMLP(nn.Module):
+    def __init__(self, config: MistralConfig):
+        super().__init__()
+        tp_size = _tp_world_size()
+        use_tensor_parallel = tp_size > 1 and config.intermediate_size % tp_size == 0
+        self.gate_proj = _make_column_linear(
+            config.hidden_size,
+            config.intermediate_size,
+            bias=False,
+            use_tensor_parallel=use_tensor_parallel,
+        )
+        self.up_proj = _make_column_linear(
+            config.hidden_size,
+            config.intermediate_size,
+            bias=False,
+            use_tensor_parallel=use_tensor_parallel,
+        )
+        self.down_proj = _make_row_linear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=False,
+            use_tensor_parallel=use_tensor_parallel,
+        )
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.act_fn(_linear_output(self.gate_proj, x)) * _linear_output(
+            self.up_proj, x
+        )
+        return _linear_output(self.down_proj, x)
+
+
 class MistralDecoderLayer(nn.Module):
     def __init__(self, config: MistralConfig, layer_idx: int):
         super().__init__()
         self.hidden_size = config.hidden_size
         self.self_attn = MistralAttention(config=config, layer_idx=layer_idx)
-        self.mlp = MistralMLP(config)
+        self.mlp = MistralTPMLP(config)
         self.input_layernorm = MistralRMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )

python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py

@@ -16,8 +16,13 @@ from transformers.modeling_outputs import BaseModelOutputWithPast
 from transformers.utils import TransformersKwargs, is_torchdynamo_compiling
 
 from sglang.multimodal_gen.configs.models.encoders.qwen_image import Qwen2_5VLConfig
+from sglang.multimodal_gen.runtime.distributed import (
+    get_tp_world_size,
+    model_parallel_is_initialized,
+)
 from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
 from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
     MergedColumnParallelLinear,
     RowParallelLinear,
 )
@@ -70,7 +75,6 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
     Qwen2_5_VLCausalLMOutputWithPast,
     Qwen2_5_VLModelOutputWithPast,
     Qwen2_5_VLRotaryEmbedding,
-    Qwen2MLP,
     apply_multimodal_rotary_pos_emb,
     eager_attention_forward,
 )
@@ -78,6 +82,50 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
 logger = logging.getLogger(__name__)
 
 
+def _tp_world_size() -> int:
+    if not model_parallel_is_initialized():
+        return 1
+    return get_tp_world_size()
+
+
+def _linear_output(linear: nn.Module, x: torch.Tensor) -> torch.Tensor:
+    output = linear(x)
+    return output[0] if isinstance(output, tuple) else output
+
+
+def _make_column_linear(
+    in_features: int,
+    out_features: int,
+    *,
+    bias: bool,
+    use_tensor_parallel: bool,
+):
+    if use_tensor_parallel:
+        return ColumnParallelLinear(
+            in_features,
+            out_features,
+            bias=bias,
+            gather_output=False,
+        )
+    return nn.Linear(in_features, out_features, bias=bias)
+
+
+def _make_row_linear(
+    in_features: int,
+    out_features: int,
+    *,
+    bias: bool,
+    use_tensor_parallel: bool,
+):
+    if use_tensor_parallel:
+        return RowParallelLinear(
+            in_features,
+            out_features,
+            bias=bias,
+        )
+    return nn.Linear(in_features, out_features, bias=bias)
+
+
 class Qwen2_5_VLAttention(nn.Module):
     """
     Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
@@ -96,31 +144,63 @@ class Qwen2_5_VLAttention(nn.Module):
             )
 
         self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
+        self.total_num_heads = config.num_attention_heads
-        self.head_dim = self.hidden_size // self.num_heads
+        self.head_dim = self.hidden_size // self.total_num_heads
-        self.num_key_value_heads = config.num_key_value_heads
+        self.total_num_key_value_heads = config.num_key_value_heads
+        if (self.head_dim * self.total_num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.total_num_heads})."
+            )
+        tp_size = _tp_world_size()
+        q_size = self.total_num_heads * self.head_dim
+        kv_size = self.total_num_key_value_heads * self.head_dim
+        self.use_tensor_parallel = (
+            tp_size > 1
+            and self.total_num_heads % tp_size == 0
+            and self.total_num_key_value_heads % tp_size == 0
+            and q_size % tp_size == 0
+            and kv_size % tp_size == 0
+        )
+        self.num_heads = (
+            self.total_num_heads // tp_size
+            if self.use_tensor_parallel
+            else self.total_num_heads
+        )
+        self.num_key_value_heads = (
+            self.total_num_key_value_heads // tp_size
+            if self.use_tensor_parallel
+            else self.total_num_key_value_heads
+        )
         self.num_key_value_groups = self.num_heads // self.num_key_value_heads
         self.is_causal = True
         self.attention_dropout = config.attention_dropout
         self.rope_scaling = config.rope_scaling
         self.scaling = self.head_dim**-0.5
 
-        if (self.head_dim * self.num_heads) != self.hidden_size:
+        self.q_proj = _make_column_linear(
-            raise ValueError(
+            self.hidden_size,
-                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+            q_size,
-                f" and `num_heads`: {self.num_heads})."
+            bias=True,
+            use_tensor_parallel=self.use_tensor_parallel,
         )
-        self.q_proj = nn.Linear(
+        self.k_proj = _make_column_linear(
-            self.hidden_size, self.num_heads * self.head_dim, bias=True
+            self.hidden_size,
+            kv_size,
+            bias=True,
+            use_tensor_parallel=self.use_tensor_parallel,
         )
-        self.k_proj = nn.Linear(
+        self.v_proj = _make_column_linear(
-            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
+            self.hidden_size,
+            kv_size,
+            bias=True,
+            use_tensor_parallel=self.use_tensor_parallel,
         )
-        self.v_proj = nn.Linear(
+        self.o_proj = _make_row_linear(
-            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
+            q_size,
-        )
+            self.hidden_size,
-        self.o_proj = nn.Linear(
+            bias=False,
-            self.num_heads * self.head_dim, self.hidden_size, bias=False
+            use_tensor_parallel=self.use_tensor_parallel,
         )
         self.sliding_window = (
             config.sliding_window
@@ -157,9 +237,9 @@ class Qwen2_5_VLAttention(nn.Module):
     ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
 
-        query_states = self.q_proj(hidden_states)
+        query_states = _linear_output(self.q_proj, hidden_states)
-        key_states = self.k_proj(hidden_states)
+        key_states = _linear_output(self.k_proj, hidden_states)
-        value_states = self.v_proj(hidden_states)
+        value_states = _linear_output(self.v_proj, hidden_states)
 
         query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
         key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
@@ -202,10 +282,42 @@ class Qwen2_5_VLAttention(nn.Module):
         # )
 
         attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
-        attn_output = self.o_proj(attn_output)
+        attn_output = _linear_output(self.o_proj, attn_output)
         return attn_output
 
 
+class Qwen2_5_VLTextMLP(nn.Module):
+    def __init__(self, config: Qwen2_5_VLTextConfig):
+        super().__init__()
+        tp_size = _tp_world_size()
+        use_tensor_parallel = tp_size > 1 and config.intermediate_size % tp_size == 0
+        self.gate_proj = _make_column_linear(
+            config.hidden_size,
+            config.intermediate_size,
+            bias=False,
+            use_tensor_parallel=use_tensor_parallel,
+        )
+        self.up_proj = _make_column_linear(
+            config.hidden_size,
+            config.intermediate_size,
+            bias=False,
+            use_tensor_parallel=use_tensor_parallel,
+        )
+        self.down_proj = _make_row_linear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=False,
+            use_tensor_parallel=use_tensor_parallel,
+        )
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.act_fn(_linear_output(self.gate_proj, x)) * _linear_output(
+            self.up_proj, x
+        )
+        return _linear_output(self.down_proj, x)
+
+
 class Qwen2_5_VLDecoderLayer(nn.Module):
     def __init__(self, config: Qwen2_5_VLTextConfig, layer_idx: int):
         super().__init__()
@@ -221,7 +333,7 @@ class Qwen2_5_VLDecoderLayer(nn.Module):
             )
         self.self_attn = Qwen2_5_VLAttention(config, layer_idx)
 
-        self.mlp = Qwen2MLP(config)
+        self.mlp = Qwen2_5_VLTextMLP(config)
         self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = Qwen2RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps

python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py

@@ -18,6 +18,7 @@ from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
 from sglang.multimodal_gen.runtime.layers.linear import (
     ColumnParallelLinear,
     ReplicatedLinear,
+    RowParallelLinear,
 )
 from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
@@ -25,6 +26,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor
 from sglang.multimodal_gen.runtime.layers.quantization.weight_only_fp8 import (
     WeightOnlyFP8ColumnParallelLinear,
     WeightOnlyFP8Linear,
+    WeightOnlyFP8RowParallelLinear,
 )
 from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
 from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
@@ -70,6 +72,11 @@ class Qwen3VLColumnParallelLinear(ColumnParallelLinear):
         return super().forward(x)[0]
 
 
+class Qwen3VLRowParallelLinear(RowParallelLinear):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x)[0]
+
+
 def _tp_world_size() -> int:
     if not model_parallel_is_initialized():
         return 1
@@ -132,6 +139,51 @@ def _make_text_linear(
     return nn.Linear(in_features, out_features, bias=bias)
 
 
+def _make_text_row_linear(
+    in_features: int,
+    out_features: int,
+    *,
+    bias: bool,
+    quant_config: QuantizationConfig | None,
+    use_weight_only_fp8: bool,
+    use_tensor_parallel: bool,
+    prefix: str,
+):
+    tp_size = _tp_world_size()
+    use_row_parallel = (
+        use_tensor_parallel and tp_size > 1 and in_features % tp_size == 0
+    )
+    if use_weight_only_fp8:
+        if use_row_parallel:
+            return WeightOnlyFP8RowParallelLinear(
+                in_features,
+                out_features,
+                bias=bias,
+                input_is_parallel=True,
+                enable_fused_w8a8=False,
+            )
+        return WeightOnlyFP8Linear(
+            in_features, out_features, bias=bias, enable_fused_w8a8=False
+        )
+    if use_row_parallel:
+        return Qwen3VLRowParallelLinear(
+            input_size=in_features,
+            output_size=out_features,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+    if quant_config is not None:
+        return Qwen3VLQuantizedLinear(
+            in_features,
+            out_features,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+    return nn.Linear(in_features, out_features, bias=bias)
+
+
 def _gather_tensor_parallel_activation(
     x: torch.Tensor, linear: nn.Module
 ) -> torch.Tensor:
@@ -159,10 +211,14 @@ class Qwen3VLTextAttention(nn.Module):
         self.head_dim = config.hidden_size // config.num_attention_heads
         self.total_num_heads = config.num_attention_heads
         self.total_num_key_value_heads = config.num_key_value_heads
-        self.tp_size = _tp_world_size() if use_tensor_parallel else 1
+        tp_size = _tp_world_size() if use_tensor_parallel else 1
-        if self.tp_size > 1:
+        use_tensor_parallel = (
-            assert self.total_num_heads % self.tp_size == 0
+            use_tensor_parallel
-            assert self.total_num_key_value_heads % self.tp_size == 0
+            and tp_size > 1
+            and self.total_num_heads % tp_size == 0
+            and self.total_num_key_value_heads % tp_size == 0
+        )
+        self.tp_size = tp_size if use_tensor_parallel else 1
         self.num_heads = self.total_num_heads // self.tp_size
         self.num_key_value_heads = self.total_num_key_value_heads // self.tp_size
         self.num_key_value_groups = self.num_heads // self.num_key_value_heads
@@ -200,14 +256,13 @@ class Qwen3VLTextAttention(nn.Module):
             gather_output=False,
             prefix=f"{prefix}.v_proj",
         )
-        self.o_proj = _make_text_linear(
+        self.o_proj = _make_text_row_linear(
             config.num_attention_heads * self.head_dim,
             config.hidden_size,
             bias=config.attention_bias,
             quant_config=quant_config,
             use_weight_only_fp8=use_weight_only_fp8,
             use_tensor_parallel=use_tensor_parallel,
-            gather_output=True,
             prefix=f"{prefix}.o_proj",
         )
         self.q_norm = Qwen3VLTextRMSNorm(
@@ -279,6 +334,9 @@ class Qwen3VLTextAttention(nn.Module):
         attn_output = self.attn(query_states, key_states, value_states)
 
         attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        if not isinstance(
+            self.o_proj, (Qwen3VLRowParallelLinear, WeightOnlyFP8RowParallelLinear)
+        ):
             attn_output = _gather_tensor_parallel_activation(attn_output, self.q_proj)
         attn_output = self.o_proj(attn_output)
         return attn_output
@@ -317,20 +375,22 @@ class Qwen3VLTextMLP(nn.Module):
             gather_output=False,
             prefix=f"{prefix}.up_proj",
         )
-        self.down_proj = _make_text_linear(
+        self.down_proj = _make_text_row_linear(
             self.intermediate_size,
             self.hidden_size,
             bias=False,
             quant_config=quant_config,
             use_weight_only_fp8=use_weight_only_fp8,
             use_tensor_parallel=use_tensor_parallel,
-            gather_output=True,
             prefix=f"{prefix}.down_proj",
         )
         self.act_fn = ACT2FN[config.hidden_act]
 
     def forward(self, x):
         hidden_states = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
+        if not isinstance(
+            self.down_proj, (Qwen3VLRowParallelLinear, WeightOnlyFP8RowParallelLinear)
+        ):
             hidden_states = _gather_tensor_parallel_activation(
                 hidden_states, self.gate_proj
             )

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -114,6 +114,13 @@ from sglang.srt.utils.common import get_compiler_backend
 logger = init_logger(__name__)
 
 
+def _ensure_tensor_model_output(model_output):
+    sample = getattr(model_output, "sample", None)
+    if isinstance(sample, torch.Tensor):
+        return sample
+    return model_output
+
+
 @dataclass(slots=True)
 class DenoisingContext:
     """Loop-scoped state shared across the denoising skeleton and its hooks."""
@@ -1810,12 +1817,13 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             getattr(current_model, "forward", current_model),
             {"guidance": guidance},
         )
-        return current_model(
+        model_output = current_model(
             hidden_states=latent_model_input,
             timestep=timestep,
             **guidance_kwargs,
             **kwargs,
         )
+        return _ensure_tensor_model_output(model_output)
 
     def prepare_sta_param(self, batch: Req, server_args: ServerArgs):
         """

python/sglang/multimodal_gen/test/unit/test_ideogram4.py

@@ -37,6 +37,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.weight_only_fp8 import (
     W8A8_FP8_GEMM_ENV,
     WeightOnlyFP8ColumnParallelLinear,
     WeightOnlyFP8Linear,
+    WeightOnlyFP8RowParallelLinear,
     dequantize_rowwise_fp8_weight,
 )
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
@@ -52,6 +53,7 @@ from sglang.multimodal_gen.runtime.loader.fsdp_load import (
 )
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.models.dits.ideogram import (
+    Ideogram4RowParallelLinear,
     Ideogram4Transformer2DModel,
 )
 from sglang.multimodal_gen.runtime.models.encoders.ideogram import (
@@ -762,7 +764,7 @@ class TestIdeogram4(unittest.TestCase):
             (1,),
         )
 
-    def test_ideogram_dit_tp_nvfp4_uses_column_parallel_quant_linears(self):
+    def test_ideogram_dit_tp_nvfp4_uses_megatron_parallel_quant_linears(self):
         import sglang.multimodal_gen.runtime.server_args as server_args_module
 
         fake_tp_group = SimpleNamespace(world_size=2, rank_in_group=1)
@@ -802,7 +804,7 @@ class TestIdeogram4(unittest.TestCase):
         finally:
             set_global_server_args(prev_args)
 
-        self.assertTrue(model.layers[0].attention.qkv.gather_output)
+        self.assertFalse(model.layers[0].attention.qkv.gather_output)
         self.assertEqual(
             tuple(model.layers[0].attention.qkv.weight.shape), (6912, 2304)
         )
@@ -810,6 +812,14 @@ class TestIdeogram4(unittest.TestCase):
             model.layers[0].attention.qkv.quant_method,
             ModelOptFp4LinearMethod,
         )
+        self.assertIsInstance(model.layers[0].attention.o, Ideogram4RowParallelLinear)
+        self.assertTrue(model.layers[0].attention.o.input_is_parallel)
+        self.assertFalse(model.layers[0].feed_forward.w1.gather_output)
+        self.assertFalse(model.layers[0].feed_forward.w3.gather_output)
+        self.assertIsInstance(
+            model.layers[0].feed_forward.w2, Ideogram4RowParallelLinear
+        )
+        self.assertTrue(model.layers[0].feed_forward.w2.input_is_parallel)
 
     def test_bitsandbytes_tp_quant_state_uses_local_output_shard(self):
         param = torch.nn.Parameter(
@@ -998,7 +1008,7 @@ class TestIdeogram4(unittest.TestCase):
             any(isinstance(module, torch.nn.Linear) for module in encoder.modules())
         )
 
-    def test_ideogram_text_encoder_tp_fp8_uses_column_parallel_linears(self):
+    def test_ideogram_text_encoder_tp_fp8_uses_megatron_parallel_linears(self):
         config = Ideogram4TextEncoderConfig()
         config.post_diffusers_config_update()
         config.arch_config.text_config = Qwen3VLTextConfig(
@@ -1044,10 +1054,16 @@ class TestIdeogram4(unittest.TestCase):
         self.assertEqual(layer.self_attn.num_key_value_heads, 2)
         self.assertIsInstance(layer.self_attn.q_proj, WeightOnlyFP8ColumnParallelLinear)
         self.assertFalse(layer.self_attn.q_proj.gather_output)
-        self.assertTrue(layer.self_attn.o_proj.gather_output)
+        self.assertIsInstance(layer.self_attn.o_proj, WeightOnlyFP8RowParallelLinear)
+        self.assertTrue(layer.self_attn.o_proj.input_is_parallel)
+        self.assertTrue(layer.self_attn.o_proj.reduce_results)
         self.assertIsInstance(layer.mlp.gate_proj, WeightOnlyFP8ColumnParallelLinear)
         self.assertFalse(layer.mlp.gate_proj.gather_output)
-        self.assertTrue(layer.mlp.down_proj.gather_output)
+        self.assertIsInstance(layer.mlp.up_proj, WeightOnlyFP8ColumnParallelLinear)
+        self.assertFalse(layer.mlp.up_proj.gather_output)
+        self.assertIsInstance(layer.mlp.down_proj, WeightOnlyFP8RowParallelLinear)
+        self.assertTrue(layer.mlp.down_proj.input_is_parallel)
+        self.assertTrue(layer.mlp.down_proj.reduce_results)
 
     def test_denoise_and_decode_shape_smoke(self):
         import sglang.multimodal_gen.runtime.server_args as server_args_module