minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-06-17 13:07:44 +08:00
committed by GitHub
parent 9371062ef3
commit 0f5e14e1d9
16 changed files with 1151 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py
+31
View File
@@ -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
+8
View File
@@ -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
+6
View File
@@ -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
+152
View File
@@ -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
+117 -14
View File
@@ -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,
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)
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(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:
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))
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,
python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py
+61 -23
View File
@@ -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] = {}
python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
+119 -16
View File
@@ -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,
python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py
+80 -26
View File
@@ -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):
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
python/sglang/multimodal_gen/runtime/models/dits/ideogram.py
+112 -10
View File
@@ -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
+22 -10
View File
@@ -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
+53 -8
View File
@@ -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)
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))
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,
python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py
+145 -19
View File
@@ -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
)
python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py
+134 -22
View File
@@ -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 = _make_column_linear(
self.hidden_size,
q_size,
bias=True,
use_tensor_parallel=self.use_tensor_parallel,
)
self.q_proj = nn.Linear(
self.hidden_size, self.num_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.k_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.v_proj = nn.Linear(
self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
)
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
python/sglang/multimodal_gen/runtime/models/encoders/qwen3vl.py
+68 -8
View File
@@ -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,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
+9 -1
View File
@@ -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
+21 -5
View File
@@ -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