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

[diffusion] optimize: optimize flux1 tensor parallel sharding (#27826)

Browse Source
This commit is contained in:
Mick
2026-06-12 13:13:18 +08:00
committed by GitHub
parent 0a1fb0da86
commit 8038806557
5 changed files with 313 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+12 -4
View File
@@ -359,20 +359,28 @@ class LocalAttention(nn.Module):
mask = mask[:, None, :, :]
mask = (mask - 1.0) * torch.finfo(q_.dtype).max
if q_.shape[1] != k_.shape[1]:
repeat_factor = q_.shape[1] // k_.shape[1]
k_ = k_.repeat_interleave(repeat_factor, dim=1)
v_ = v_.repeat_interleave(repeat_factor, dim=1)
sdpa_context = (
sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
if self.allow_cudnn_sdp and q_.device.type == "cuda"
else nullcontext()
)
attn_kwargs = {
"attn_mask": mask,
"dropout_p": 0.0,
"is_causal": False,
"scale": self.softmax_scale,
}
with sdpa_context:
return torch.nn.functional.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=self.softmax_scale,
**attn_kwargs,
).transpose(1, 2)
output = self.attn_impl.forward(q, k, v, attn_metadata=ctx_attn_metadata)
python/sglang/multimodal_gen/runtime/models/dits/flux.py
+182 -28
View File
@@ -28,6 +28,7 @@ from diffusers.models.normalization import (
from torch.nn import LayerNorm as LayerNorm
from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
from sglang.multimodal_gen.runtime.distributed import divide, get_tp_world_size
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.layernorm import (
RMSNorm,
@@ -36,6 +37,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
from sglang.multimodal_gen.runtime.layers.linear import (
ColumnParallelLinear,
MergedColumnParallelLinear,
RowParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.mlp import FeedForward
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
@@ -209,6 +211,75 @@ def _get_qkv_projections(
return query, key, value, encoder_query, encoder_key, encoder_value
class FluxGELU(nn.Module):
def __init__(
self,
dim: int,
inner_dim: int,
bias: bool = True,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
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",
)
self.gelu = nn.GELU(approximate="tanh")
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states, _ = self.proj(hidden_states)
return self.gelu(hidden_states)
class FluxParallelFeedForward(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: Optional[QuantizationConfig] = None,
prefix: str = "",
):
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(
[
FluxGELU(
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 FluxAttention(torch.nn.Module, AttentionModuleMixin):
def __init__(
self,
@@ -238,6 +309,11 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
self.context_pre_only = context_pre_only
self.pre_only = pre_only
self.heads = out_dim // dim_head if out_dim is not None else num_heads
self.tp_size = get_tp_world_size()
self.shard_qkv = self.tp_size > 1 and not isinstance(
quant_config, NunchakuConfig
)
self.local_heads = divide(self.heads, self.tp_size)
self.added_kv_proj_dim = added_kv_proj_dim
self.added_proj_bias = added_proj_bias
@@ -252,7 +328,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
query_dim,
[self.inner_dim] * 3,
bias=bias,
gather_output=True,
gather_output=not self.shard_qkv,
quant_config=quant_config,
prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
)
@@ -261,7 +337,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
query_dim,
self.inner_dim,
bias=bias,
gather_output=True,
gather_output=not self.shard_qkv,
quant_config=quant_config,
prefix=f"{prefix}.to_q" if prefix else "to_q",
)
@@ -269,7 +345,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
query_dim,
self.inner_dim,
bias=bias,
gather_output=True,
gather_output=not self.shard_qkv,
quant_config=quant_config,
prefix=f"{prefix}.to_k" if prefix else "to_k",
)
@@ -277,18 +353,24 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
query_dim,
self.inner_dim,
bias=bias,
gather_output=True,
gather_output=not self.shard_qkv,
quant_config=quant_config,
prefix=f"{prefix}.to_v" if prefix else "to_v",
)
if not self.pre_only:
self.to_out = torch.nn.ModuleList([])
out_proj_cls = RowParallelLinear if self.shard_qkv else ColumnParallelLinear
out_proj_kwargs = (
{"input_is_parallel": True}
if self.shard_qkv
else {"gather_output": True}
)
self.to_out.append(
ColumnParallelLinear(
out_proj_cls(
self.inner_dim,
self.out_dim,
bias=out_bias,
gather_output=True,
**out_proj_kwargs,
quant_config=quant_config,
prefix=f"{prefix}.to_out.0" if prefix else "",
)
@@ -304,7 +386,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
added_kv_proj_dim,
[self.inner_dim] * 3,
bias=added_proj_bias,
gather_output=True,
gather_output=not self.shard_qkv,
quant_config=quant_config,
prefix=f"{prefix}.to_added_qkv" if prefix else "to_added_qkv",
)
@@ -313,7 +395,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
gather_output=not self.shard_qkv,
quant_config=quant_config,
prefix=f"{prefix}.add_q_proj" if prefix else "add_q_proj",
)
@@ -321,7 +403,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
gather_output=not self.shard_qkv,
quant_config=quant_config,
prefix=f"{prefix}.add_k_proj" if prefix else "add_k_proj",
)
@@ -329,21 +411,29 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
added_kv_proj_dim,
self.inner_dim,
bias=added_proj_bias,
gather_output=True,
gather_output=not self.shard_qkv,
quant_config=quant_config,
prefix=f"{prefix}.add_v_proj" if prefix else "add_v_proj",
)
self.to_add_out = ColumnParallelLinear(
add_out_proj_cls = (
RowParallelLinear if self.shard_qkv else ColumnParallelLinear
)
add_out_proj_kwargs = (
{"input_is_parallel": True}
if self.shard_qkv
else {"gather_output": True}
)
self.to_add_out = add_out_proj_cls(
self.inner_dim,
query_dim,
bias=out_bias,
gather_output=True,
**add_out_proj_kwargs,
quant_config=quant_config,
prefix=f"{prefix}.to_add_out" if prefix else "",
)
self.attn = USPAttention(
num_heads=num_heads,
num_heads=self.local_heads if self.shard_qkv else num_heads,
head_size=self.head_dim,
dropout_rate=0,
softmax_scale=None,
@@ -366,9 +456,10 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
encoder_value,
) = _get_qkv_projections(self, x, encoder_hidden_states)
query = query.unflatten(-1, (self.heads, -1))
key = key.unflatten(-1, (self.heads, -1))
value = value.unflatten(-1, (self.heads, -1))
num_heads = self.local_heads if self.shard_qkv else self.heads
query = query.unflatten(-1, (num_heads, -1))
key = key.unflatten(-1, (num_heads, -1))
value = value.unflatten(-1, (num_heads, -1))
cos_sin_cache = None
if freqs_cis is not None:
cos, sin = freqs_cis
@@ -381,9 +472,9 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
)
if self.added_kv_proj_dim is not None:
encoder_query = encoder_query.unflatten(-1, (self.heads, -1))
encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
encoder_query = encoder_query.unflatten(-1, (num_heads, -1))
encoder_key = encoder_key.unflatten(-1, (num_heads, -1))
encoder_value = encoder_value.unflatten(-1, (num_heads, -1))
text_seq_len = encoder_query.shape[1]
encoder_query, encoder_key = apply_qk_norm_with_optional_rope(
@@ -466,6 +557,9 @@ class FluxSingleTransformerBlock(nn.Module):
super().__init__()
self.mlp_hidden_dim = int(dim * mlp_ratio)
self.use_nunchaku_structure = isinstance(quant_config, NunchakuConfig)
self.tp_size = get_tp_world_size()
self.local_mlp_hidden_dim = divide(self.mlp_hidden_dim, self.tp_size)
self.local_dim = divide(dim, self.tp_size)
self.norm = AdaLayerNormZeroSingle(dim)
@@ -502,23 +596,34 @@ class FluxSingleTransformerBlock(nn.Module):
if is_nunchaku_available():
self.norm = NunchakuAdaLayerNormZeroSingle(self.norm, scale_shift=0)
else:
shard_single_block = self.tp_size > 1
self.proj_mlp = ColumnParallelLinear(
dim,
self.mlp_hidden_dim,
bias=True,
gather_output=True,
gather_output=not shard_single_block,
quant_config=quant_config,
prefix=f"{prefix}.proj_mlp" if prefix else "proj_mlp",
)
self.act_mlp = nn.GELU(approximate="tanh")
self.proj_out = ColumnParallelLinear(
proj_out_cls = (
RowParallelLinear if shard_single_block else ColumnParallelLinear
)
proj_out_kwargs = (
{"input_is_parallel": True}
if shard_single_block
else {"gather_output": True}
)
self.proj_out = proj_out_cls(
dim + self.mlp_hidden_dim,
dim,
bias=True,
gather_output=True,
**proj_out_kwargs,
quant_config=quant_config,
prefix=f"{prefix}.proj_out" if prefix else "proj_out",
)
if shard_single_block:
self._patch_proj_out_weight_loader()
self.attn = FluxAttention(
query_dim=dim,
dim_head=attention_head_dim,
@@ -531,6 +636,30 @@ class FluxSingleTransformerBlock(nn.Module):
prefix=f"{prefix}.attn" if prefix else "attn",
)
def _patch_proj_out_weight_loader(self) -> None:
dim, mlp_dim = self.local_dim, self.local_mlp_hidden_dim
tp_rank = self.proj_out.tp_rank
def _loader(param, loaded_weight):
input_dim = getattr(param, "input_dim", None)
if input_dim is not None:
# checkpoint columns are [attn_full | mlp_full], while TP consumes [attn_shard | mlp_shard]
attn_cols = loaded_weight.narrow(input_dim, tp_rank * dim, dim)
mlp_cols = loaded_weight.narrow(
input_dim,
self.tp_size * dim + tp_rank * mlp_dim,
mlp_dim,
)
param.data.copy_(torch.cat([attn_cols, mlp_cols], dim=input_dim))
else:
param.data.copy_(loaded_weight)
self.proj_out.weight_loader = _loader
if hasattr(self.proj_out.weight, "_weight_loader"):
self.proj_out.weight._weight_loader = _loader
else:
self.proj_out.weight.weight_loader = _loader
def forward(
self,
hidden_states: torch.Tensor,
@@ -634,13 +763,16 @@ class FluxTransformerBlock(nn.Module):
and is_nunchaku_available()
)
self.use_nunchaku_structure = nunchaku_enabled
self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
self.ff_context = FeedForward(
dim=dim,
dim_out=dim,
activation_fn="gelu-approximate",
)
self.tp_size = get_tp_world_size()
if nunchaku_enabled:
self.ff = FeedForward(
dim=dim, dim_out=dim, activation_fn="gelu-approximate"
)
self.ff_context = FeedForward(
dim=dim,
dim_out=dim,
activation_fn="gelu-approximate",
)
nunchaku_kwargs = {
"precision": quant_config.precision,
"rank": quant_config.rank,
@@ -652,6 +784,28 @@ class FluxTransformerBlock(nn.Module):
self.norm1_context = NunchakuAdaLayerNormZero(
self.norm1_context, scale_shift=0
)
elif self.tp_size > 1:
self.ff = FluxParallelFeedForward(
dim=dim,
dim_out=dim,
quant_config=quant_config,
prefix=f"{prefix}.ff" if prefix else "ff",
)
self.ff_context = FluxParallelFeedForward(
dim=dim,
dim_out=dim,
quant_config=quant_config,
prefix=f"{prefix}.ff_context" if prefix else "ff_context",
)
else:
self.ff = FeedForward(
dim=dim, dim_out=dim, activation_fn="gelu-approximate"
)
self.ff_context = FeedForward(
dim=dim,
dim_out=dim,
activation_fn="gelu-approximate",
)
def forward(
self,
python/sglang/multimodal_gen/runtime/utils/model_overlay.py
+61 -3
View File
@@ -51,6 +51,16 @@ MODEL_OVERLAY_METADATA_PATTERNS = [
"**/*.txt",
]
_MATERIALIZED_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pth", ".pt")
_MATERIALIZED_CONFIG_ONLY_COMPONENTS = {
"feature_extractor",
"image_processor",
"processor",
"scheduler",
"tokenizer",
"tokenizer_2",
}
_MODEL_OVERLAY_REGISTRY_CACHE: dict[str, dict[str, Any]] | None = None
@@ -196,6 +206,52 @@ def load_model_index_from_dir(model_dir: str) -> dict[str, Any]:
return config
def _component_has_weight_file(component_dir: str) -> bool:
for root, _, file_names in os.walk(component_dir):
if any(
file_name.endswith(_MATERIALIZED_WEIGHT_SUFFIXES)
and os.path.isfile(os.path.join(root, file_name))
for file_name in file_names
):
return True
return False
def _materialized_overlay_has_component_weights(model_dir: str) -> bool:
model_index = load_model_index_from_dir(model_dir)
for component_name, entry in model_index.items():
if (
component_name.startswith("_")
or component_name == "pipeline_name"
or component_name in _MATERIALIZED_CONFIG_ONLY_COMPONENTS
or not isinstance(entry, list)
):
continue
component_dir = os.path.join(model_dir, component_name)
if not os.path.isdir(component_dir) or not _component_has_weight_file(
component_dir
):
logger.warning(
"Materialized overlay cache for %s is missing weights for component %s",
model_dir,
component_name,
)
return False
return True
def _materialized_overlay_cache_complete(
final_dir: str,
marker_path: str,
verify_diffusers_model_complete_fn: Callable[[str], bool],
) -> bool:
return (
verify_diffusers_model_complete_fn(final_dir)
and os.path.exists(marker_path)
and _materialized_overlay_has_component_weights(final_dir)
)
def _ensure_dir(path: str) -> None:
os.makedirs(path, exist_ok=True)
@@ -503,15 +559,17 @@ def materialize_overlay_model(
safe_name = source_model_id.replace("/", "__")
final_dir = os.path.join(cache_root, f"{safe_name}-{cache_key}")
marker_path = os.path.join(final_dir, ".sglang_overlay_materialized.json")
if verify_diffusers_model_complete_fn(final_dir) and os.path.exists(marker_path):
if _materialized_overlay_cache_complete(
final_dir, marker_path, verify_diffusers_model_complete_fn
):
return final_dir
lock_name = (
f"overlay-materialize::{source_model_id}::{overlay_repo_id}::{overlay_revision}"
)
with get_lock(lock_name).acquire(poll_interval=2):
if verify_diffusers_model_complete_fn(final_dir) and os.path.exists(
marker_path
if _materialized_overlay_cache_complete(
final_dir, marker_path, verify_diffusers_model_complete_fn
):
return final_dir
python/sglang/multimodal_gen/test/test_utils.py
+53 -24
View File
@@ -995,29 +995,56 @@ def _remote_consistency_gt_candidates(
return [(filename, f"{base_url}/{filename}") for filename in filenames]
def _is_ascend_consistency_case(case_id: str) -> bool:
return "npu" in case_id
def _remote_file_exists(url: str) -> bool:
for method in ("head", "get"):
try:
if method == "head":
resp = requests.head(url, timeout=10, allow_redirects=True)
else:
resp = requests.get(
url,
timeout=10,
allow_redirects=True,
headers={"Range": "bytes=0-0"},
stream=True,
)
for _ in range(3):
for method in ("head", "get"):
try:
if resp.status_code in (200, 206):
return True
if resp.status_code not in (403, 405, 429) and resp.status_code < 500:
return False
if method == "head":
resp = requests.head(url, timeout=30, allow_redirects=True)
else:
resp = requests.get(
url,
timeout=30,
allow_redirects=True,
headers={"Range": "bytes=0-0"},
stream=True,
)
try:
if resp.status_code in (200, 206):
return True
if (
resp.status_code not in (403, 405, 429)
and resp.status_code < 500
):
return False
finally:
resp.close()
except requests.RequestException:
pass
return False
def _load_remote_gt_image(url: str) -> np.ndarray:
last_error: Exception | None = None
for _ in range(3):
try:
resp = requests.get(url, timeout=60)
try:
if resp.status_code == 200:
image = Image.open(io.BytesIO(resp.content)).convert("RGB")
return np.array(image)
last_error = FileNotFoundError(f"GT image not found: {url}")
if resp.status_code not in (403, 429) and resp.status_code < 500:
break
finally:
resp.close()
except requests.RequestException:
pass
return False
except requests.RequestException as exc:
last_error = exc
raise FileNotFoundError(f"GT image not found: {url}") from last_error
def _find_remote_consistency_gt_files(
@@ -1026,7 +1053,12 @@ def _find_remote_consistency_gt_files(
is_video: bool,
output_format: str | None = None,
) -> list[tuple[str, str]]:
if case_id in SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_CASES:
if _is_ascend_consistency_case(case_id):
bases = (
SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE_ASCEND,
SGL_TEST_FILES_CONSISTENCY_GT_BASE,
)
elif case_id in SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_CASES:
bases = SGL_TEST_FILES_CONSISTENCY_GT_BASES
else:
# Avoid accidentally comparing non-comparable CI cases against official GT.
@@ -1115,10 +1147,7 @@ def load_consistency_gt(
f"GT image not found for {case_id}. Tried: {', '.join(filenames)}"
)
for _, url in remote_files:
resp = requests.get(url, timeout=30)
if resp.status_code != 200:
raise FileNotFoundError(f"GT image not found: {url}")
images.append(np.array(Image.open(io.BytesIO(resp.content)).convert("RGB")))
images.append(_load_remote_gt_image(url))
source_dir = remote_files[0][1].rsplit("/", 1)[0]
logger.info(f"Loaded {len(images)} GT images for {case_id} from {source_dir}")
python/sglang/multimodal_gen/test/unit/test_transformer_quant.py
+5
View File
@@ -390,8 +390,13 @@ class TestTransformerQuantHelpers(unittest.TestCase):
"sglang.multimodal_gen.runtime.layers.attention.selector.get_global_server_args",
return_value=SimpleNamespace(attention_backend=None),
)
@patch(
"sglang.multimodal_gen.runtime.models.dits.flux.get_tp_world_size",
return_value=1,
)
def test_flux_single_transformer_block_modelopt_excludes_use_full_prefix(
self,
_mock_tp_world_size,
_mock_server_args,
_mock_ring_world_size,
_mock_tp_group,