diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py index 7eba6e057..66ec770fc 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py +++ b/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): diff --git a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py b/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py index b474b5d5f..228c5c1cf 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index 89f232592..7b4cce559 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/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): diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py index b7bd5491e..2d540eb5a 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 370968609..132a35fb0 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/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: - assert dim % num_heads == 0 + 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, - value=v, + key=cache_key, + 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, - cache_view.v, + key, + value, ) return x @@ -202,23 +248,71 @@ class CausalWanTransformerBlock(nn.Module): # 1. Self-attention self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) - 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) + 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 = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config) + 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( + 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: - query = self.norm_q(query) + 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: - key = self.norm_k(key) + 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)) - key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - value = value.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.local_num_heads, self.dim_head)) + value = value.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) attn_output = self.attn1( query, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index 0f3f28333..9cb0fa060 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/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 - # Input is not parallel (gather_output=True on QKV) - 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), ) @@ -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.num_key_value_heads, + self.local_num_attention_heads : self.local_num_attention_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) - v_und: [B, S_und, H_kv, D] pre-computed UND values (always full/replicated) + k_und: [B, S_und, H_kv_local, D] UND keys (replicated over SP) + 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.num_key_value_heads, + self.local_num_attention_heads : self.local_num_attention_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] = {} diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index 0cb937df8..ca227ac5f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/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( - query_dim, self.inner_dim, bias=bias, quant_config=quant_config + self.to_q = ColumnParallelLinear( + 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( - query_dim, self.inner_kv_dim, bias=bias, quant_config=quant_config + self.to_k = ColumnParallelLinear( + 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( - query_dim, self.inner_kv_dim, bias=bias, quant_config=quant_config + self.to_v = ColumnParallelLinear( + 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( - self.inner_dim, self.out_dim, bias=True, quant_config=quant_config + RowParallelLinear( + 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)) - key = key.unflatten(2, (self.heads, -1)) - value = value.unflatten(2, (self.heads, -1)) + query = query.unflatten(2, (self.num_local_heads, -1)) + key = key.unflatten(2, (self.num_local_kv_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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index 6313d1e7f..c173e54d5 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/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( - hidden_size + mlp_hidden_dim, + self.linear2 = MixedRowParallelLinear( + [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): - txt = encoder_hidden_states[:, 1:] - text_states_2 = encoder_hidden_states[:, 0, : self.text_states_dim_2] + 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, - head_size=hidden_size // num_attention_heads, + num_heads=self.local_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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py index 169f784f5..805c2a486 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py +++ b/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): diff --git a/python/sglang/multimodal_gen/runtime/models/dits/joy_image.py b/python/sglang/multimodal_gen/runtime/models/dits/joy_image.py index c988fd925..14bf66faf 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/joy_image.py +++ b/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 ( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py index 25a2a858f..248f8ae6e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py +++ b/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, - value=v, + key=cache_key, + 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, - cache_view.v, + key, + 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) - query = self.norm_q(query) - key = self.norm_k(key) - query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + 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.local_num_heads, self.dim_head)) + key = key.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) + value = value.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) attn_output = self.attn1( query, diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py b/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py index 72a7d9587..95a8cd592 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py +++ b/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( - config.hidden_size, config.num_attention_heads * self.head_dim, bias=False + self.total_num_heads = config.num_attention_heads + 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( - config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False + self.num_heads = ( + self.total_num_heads // tp_size + if self.use_tensor_parallel + else self.total_num_heads ) - self.v_proj = nn.Linear( - config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False + 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.o_proj = nn.Linear( - config.num_attention_heads * self.head_dim, config.hidden_size, bias=False + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + 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) - key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) - value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + query_states = ( + _linear_output(self.q_proj, hidden_states) + .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 ) diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py index 491bbbe66..cdf5ccc05 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py +++ b/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.head_dim = self.hidden_size // self.num_heads - self.num_key_value_heads = config.num_key_value_heads + self.total_num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.total_num_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: - raise ValueError( - f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" - f" and `num_heads`: {self.num_heads})." - ) - self.q_proj = nn.Linear( - self.hidden_size, self.num_heads * self.head_dim, bias=True + self.q_proj = _make_column_linear( + self.hidden_size, + q_size, + bias=True, + use_tensor_parallel=self.use_tensor_parallel, ) - self.k_proj = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True + self.k_proj = _make_column_linear( + self.hidden_size, + kv_size, + bias=True, + use_tensor_parallel=self.use_tensor_parallel, ) - self.v_proj = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True + self.v_proj = _make_column_linear( + self.hidden_size, + kv_size, + bias=True, + use_tensor_parallel=self.use_tensor_parallel, ) - self.o_proj = nn.Linear( - self.num_heads * self.head_dim, self.hidden_size, bias=False + self.o_proj = _make_row_linear( + q_size, + 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) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + query_states = _linear_output(self.q_proj, hidden_states) + key_states = _linear_output(self.k_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 diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py index 1e68c5ba7..440742c13 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py +++ b/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 - if self.tp_size > 1: - assert self.total_num_heads % self.tp_size == 0 - assert self.total_num_key_value_heads % self.tp_size == 0 + tp_size = _tp_world_size() if use_tensor_parallel else 1 + use_tensor_parallel = ( + use_tensor_parallel + 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,7 +334,10 @@ class Qwen3VLTextAttention(nn.Module): attn_output = self.attn(query_states, key_states, value_states) attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = _gather_tensor_parallel_activation(attn_output, self.q_proj) + 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,23 +375,25 @@ 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) - hidden_states = _gather_tensor_parallel_activation( - hidden_states, self.gate_proj - ) + if not isinstance( + self.down_proj, (Qwen3VLRowParallelLinear, WeightOnlyFP8RowParallelLinear) + ): + hidden_states = _gather_tensor_parallel_activation( + hidden_states, self.gate_proj + ) down_proj = self.down_proj(hidden_states) return down_proj diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 10f739084..e7fa7402a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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): """ diff --git a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py b/python/sglang/multimodal_gen/test/unit/test_ideogram4.py index 9e86cf0e5..967cbe75c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py +++ b/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