diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index 9aab2badd..884c47bd3 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py @@ -530,10 +530,6 @@ class Flux2PipelineConfig(FluxPipelineConfig): txt_ids = _prepare_text_ids(prompt_embeds).to(device=device) img_ids = batch.latent_ids - if batch.image_latent is not None: - image_latent_ids = batch.condition_image_latent_ids - img_ids = torch.cat([img_ids, image_latent_ids], dim=1).to(device=device) - if img_ids.ndim == 3: img_ids = img_ids[0] if txt_ids.ndim == 3: @@ -544,6 +540,16 @@ class Flux2PipelineConfig(FluxPipelineConfig): img_cos = shard_rotary_emb_for_sp(img_cos) img_sin = shard_rotary_emb_for_sp(img_sin) + if batch.image_latent is not None: + cond_ids = batch.condition_image_latent_ids + if cond_ids.ndim == 3: + cond_ids = cond_ids[0] + cond_cos, cond_sin = rotary_emb.forward(cond_ids) + cond_cos = shard_rotary_emb_for_sp(cond_cos) + cond_sin = shard_rotary_emb_for_sp(cond_sin) + img_cos = torch.cat([img_cos, cond_cos], dim=0) + img_sin = torch.cat([img_sin, cond_sin], dim=0) + txt_cos, txt_sin = rotary_emb.forward(txt_ids) cos = torch.cat([txt_cos, img_cos], dim=0).to(device=device) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 7550514d5..4bfe9f2c0 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -13,6 +13,7 @@ from sglang.multimodal_gen.runtime.distributed.communication_op import ( from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_ring_parallel_world_size, get_sequence_parallel_world_size, + get_sp_group, get_sp_parallel_rank, get_sp_world_size, get_ulysses_parallel_world_size, @@ -352,23 +353,23 @@ class USPAttention(nn.Module): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - replicated_q: torch.Tensor | None = None, - replicated_k: torch.Tensor | None = None, - replicated_v: torch.Tensor | None = None, + num_replicated_prefix: int = 0, ) -> torch.Tensor: """ Forward pass for USPAttention. q, k, v: [B, S_local, H, D] + num_replicated_prefix: number of leading tokens in q/k/v that are + replicated (identical) across all SP ranks, e.g. text tokens + in FLUX joint attention. These tokens are excluded from the + Ulysses all-to-all so they appear exactly once in the gathered + sequence, preserving correct attention weights. Note: Replicated tensors are not supported in this implementation. When skip_sequence_parallel=True (set at construction time), all SP communication is bypassed — use this for cross-attention where KV content is replicated across ranks (distinct from replicated_k/v args). """ - assert ( - replicated_q is None and replicated_k is None and replicated_v is None - ), "USPAttention does not support replicated_qkv." forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata if self.skip_sequence_parallel or get_sequence_parallel_world_size() == 1: @@ -376,8 +377,14 @@ class USPAttention(nn.Module): out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) return out + sp_size = get_ulysses_parallel_world_size() + if sp_size > 1 and num_replicated_prefix > 0: + return self._forward_with_replicated_prefix( + q, k, v, ctx_attn_metadata, num_replicated_prefix + ) + # Ulysses-style All-to-All for sequence/head sharding - if get_ulysses_parallel_world_size() > 1: + if sp_size > 1: # -> [B, S, H_local, D] q = _usp_input_all_to_all(q, head_dim=2) k = _usp_input_all_to_all(k, head_dim=2) @@ -398,8 +405,66 @@ class USPAttention(nn.Module): out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) # Ulysses-style All-to-All to restore original sharding - if get_ulysses_parallel_world_size() > 1: + if sp_size > 1: # -> [B, S_local, H, D] out = _usp_output_all_to_all(out, head_dim=2) return out + + def _forward_with_replicated_prefix( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ctx_attn_metadata, + num_rep: int, + ) -> torch.Tensor: + """Ulysses attention where the first *num_rep* tokens are replicated + across SP ranks (e.g. text tokens) and should NOT be duplicated by the + all-to-all. + + Strategy: + 1. Split q/k/v into replicated prefix and SP-sharded suffix. + 2. All-to-all only the sharded suffix (gathers sequence, shards heads). + 3. Locally slice the replicated prefix to the same head shard. + 4. Concatenate [prefix_h_local, gathered_suffix] and run attention. + 5. Split output, all-to-all back the suffix, all-gather prefix heads. + """ + sp_size = get_ulysses_parallel_world_size() + sp_rank = get_sp_parallel_rank() + + q_rep, q_shard = q[:, :num_rep], q[:, num_rep:] + k_rep, k_shard = k[:, :num_rep], k[:, num_rep:] + v_rep, v_shard = v[:, :num_rep], v[:, num_rep:] + + q_shard = _usp_input_all_to_all(q_shard, head_dim=2) + k_shard = _usp_input_all_to_all(k_shard, head_dim=2) + v_shard = _usp_input_all_to_all(v_shard, head_dim=2) + + h_local = q_shard.shape[2] + h_start = sp_rank * h_local + h_end = h_start + h_local + q_rep = q_rep[:, :, h_start:h_end, :].contiguous() + k_rep = k_rep[:, :, h_start:h_end, :].contiguous() + v_rep = v_rep[:, :, h_start:h_end, :].contiguous() + + q = torch.cat([q_rep, q_shard], dim=1) + k = torch.cat([k_rep, k_shard], dim=1) + v = torch.cat([v_rep, v_shard], dim=1) + + out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + + out_rep = out[:, :num_rep] + out_shard = out[:, num_rep:] + + out_shard = _usp_output_all_to_all(out_shard, head_dim=2) + + gathered = [torch.empty_like(out_rep) for _ in range(sp_size)] + torch.distributed.all_gather( + gathered, + out_rep.contiguous(), + group=get_sp_group().ulysses_group, + ) + out_rep = torch.cat(gathered, dim=2) + + return torch.cat([out_rep, out_shard], dim=1) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index bbf7cf489..5cd447f35 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -264,7 +264,10 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): query, key, cos_sin_cache, is_neox=False ) - hidden_states = self.attn(query, key, value) + num_rep = ( + encoder_hidden_states.shape[1] if encoder_hidden_states is not None else 0 + ) + hidden_states = self.attn(query, key, value, num_replicated_prefix=num_rep) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) @@ -366,6 +369,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, freqs_cis: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + num_replicated_prefix: int = 0, **kwargs, ) -> torch.Tensor: # Parallel in (QKV + MLP in) projection @@ -398,7 +402,9 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) - hidden_states = self.attn(query, key, value) + hidden_states = self.attn( + query, key, value, num_replicated_prefix=num_replicated_prefix + ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) @@ -468,6 +474,7 @@ class Flux2SingleTransformerBlock(nn.Module): attn_output = self.attn( hidden_states=norm_hidden_states, freqs_cis=freqs_cis, + num_replicated_prefix=text_seq_len or 0, **joint_attention_kwargs, ) @@ -879,6 +886,7 @@ class Flux2Transformer2DModel(CachableDiT, OffloadableDiTMixin): temb_mod_params=single_stream_mod, freqs_cis=freqs_cis, joint_attention_kwargs=joint_attention_kwargs, + text_seq_len=num_txt_tokens, ) # Remove text tokens from concatenated stream hidden_states = hidden_states[:, num_txt_tokens:, ...] 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 844999554..48646bb1a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -23,7 +23,6 @@ from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode from sglang.multimodal_gen.configs.pipeline_configs.wan import ( Wan2_2_TI2V_5B_Config, - WanI2V480PConfig, ) from sglang.multimodal_gen.runtime.cache.cache_dit_integration import ( CacheDitConfig, @@ -794,13 +793,9 @@ class DenoisingStage(PipelineStage): else: batch.did_sp_shard_latents = False - # For I2I tasks like QwenImageEdit, where the image latents is provided as condition, the image_latent (input image) should be - # replicated on all SP ranks, not sharded, as it provides global context. - # For Wan2_2_TI2V_5B_Config, it has very special settings - if ( - isinstance(server_args.pipeline_config, WanI2V480PConfig) - and batch.image_latent is not None - ): + # image_latent must be sharded consistently with latents when it is + # concatenated along the sequence dimension in the denoising loop. + if batch.image_latent is not None: batch.image_latent, _ = server_args.pipeline_config.shard_latents_for_sp( batch, batch.image_latent ) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index cdd9f6700..f4fdd93d9 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -2051,6 +2051,26 @@ "expected_e2e_ms": 6359.87, "expected_avg_denoise_ms": 99.47, "expected_median_denoise_ms": 100.27 + }, + "flux_2_klein_ti2i_2_gpus": { + "stages_ms": { + "InputValidationStage": 40.19, + "TextEncodingStage": 88.84, + "ImageVAEEncodingStage": 80.81, + "LatentPreparationStage": 1.05, + "TimestepPreparationStage": 28.64, + "DenoisingStage": 354.04, + "DecodingStage": 11.11 + }, + "denoise_step_ms": { + "0": 33.54, + "1": 61.3, + "2": 86.9, + "3": 87.55 + }, + "expected_e2e_ms": 716.81, + "expected_avg_denoise_ms": 67.32, + "expected_median_denoise_ms": 74.1 } } } diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 3ba45f0db..f72de0a88 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -932,8 +932,9 @@ def get_generate_fn( if is_image_url(image_path): new_image_paths.append(download_image_from_url(str(image_path))) else: - new_image_paths.append(Path(image_path)) - if not image_path.exists(): + local_path = Path(image_path) + new_image_paths.append(local_path) + if not local_path.exists(): pytest.skip(f"{case_id}: file missing: {image_path}") image_paths = new_image_paths diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 1a4e32d59..4b14a6752 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -790,6 +790,15 @@ TWO_GPU_CASES_B = [ ), T2I_sampling_params, ), + DiffusionTestCase( + "flux_2_klein_ti2i_2_gpus", + DiffusionServerArgs( + model_path="black-forest-labs/FLUX.2-klein-4B", + modality="image", + num_gpus=2, + ), + TI2I_sampling_params, + ), ] # Load global configuration