From e8668508d123b1419b06a5d065527240cd589b4c Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 6 Jun 2026 09:37:29 +0800 Subject: [PATCH] [diffusion] optimize: optimize LingBot realtime sp cache path (#27383) --- .../runtime/models/dits/lingbot_world.py | 28 +++++++++++++++---- .../realtime/test_lingbot_causal_denoising.py | 11 +++++--- 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py index 6a2bc5b38..c3a28638b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py @@ -52,7 +52,9 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( get_rotary_pos_embed, ) from sglang.multimodal_gen.runtime.layers.usp import ( + _usp_input_all_to_all, _usp_input_all_to_all_varlen, + _usp_output_all_to_all, _usp_output_all_to_all_varlen, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import ( @@ -97,6 +99,12 @@ def _compute_sequence_splits(total_len: int, world_size: int) -> list[int]: return [base + (1 if rank < remainder else 0) for rank in range(world_size)] +def _sequence_splits_are_uniform(seq_splits: list[int]) -> bool: + return len(seq_splits) <= 1 or all( + seq_len == seq_splits[0] for seq_len in seq_splits + ) + + def _sequence_shard_tensor( x: torch.Tensor, seq_splits: list[int], rank: int ) -> torch.Tensor: @@ -111,6 +119,9 @@ def _sequence_all_gather_varlen( group: dist.ProcessGroup, ) -> torch.Tensor: rank = get_sp_parallel_rank() + if _sequence_splits_are_uniform(seq_splits): + return sequence_model_parallel_all_gather(x.contiguous(), dim=1) + max_seq = max(seq_splits) local_seq = seq_splits[rank] if local_seq < max_seq: @@ -215,6 +226,7 @@ class LingBotWorldCausalSelfAttention(CausalWanSelfAttention): roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) forward_batch = get_forward_context().forward_batch seq_splits = None + uniform_seq_splits = False sequence_shard_enabled = ( kv_cache is not None and forward_batch is not None @@ -245,9 +257,14 @@ class LingBotWorldCausalSelfAttention(CausalWanSelfAttention): "LingBot causal sequence sharding requires forward_batch.sequence_shard_splits." ) seq_splits = list(seq_splits) + uniform_seq_splits = _sequence_splits_are_uniform(seq_splits) # Pack Q/K/V to avoid launching three Ulysses all-to-all collectives. qkv = torch.cat([roped_query, roped_key, v], dim=-1) - qkv = _usp_input_all_to_all_varlen(qkv, seq_splits, head_dim=2) + qkv = ( + _usp_input_all_to_all(qkv, head_dim=2) + if uniform_seq_splits + else _usp_input_all_to_all_varlen(qkv, seq_splits, head_dim=2) + ) roped_query, roped_key, v = qkv.chunk(3, dim=-1) cache_view = kv_cache.update_and_get_attention_kv( @@ -266,7 +283,11 @@ class LingBotWorldCausalSelfAttention(CausalWanSelfAttention): ) if sequence_shard_enabled: assert seq_splits is not None - x = _usp_output_all_to_all_varlen(x, seq_splits, head_dim=2) + x = ( + _usp_output_all_to_all(x, head_dim=2) + if uniform_seq_splits + else _usp_output_all_to_all_varlen(x, seq_splits, head_dim=2) + ) return x @@ -937,9 +958,6 @@ class CausalLingBotWorldTransformerBlock(CausalWanTransformerBlock): return None forward_context = get_forward_context() - if forward_context.current_timestep < 0: - return self.cam_conditioner.compute_scale_shift(c2ws_plucker_emb) - forward_batch = forward_context.forward_batch if not CausalLingBotWorldTransformer3DModel._should_cache_cam_conditioner( forward_batch diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py b/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py index 71441094b..cdc745914 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_lingbot_causal_denoising.py @@ -322,7 +322,7 @@ def test_lingbot_cam_conditioner_cache_skips_single_ulysses_world(monkeypatch): assert "lingbot_cam_conditioner" not in forward_batch.extra -def test_lingbot_cam_conditioner_cache_skips_context_update(monkeypatch): +def test_lingbot_cam_conditioner_cache_reuses_context_update(monkeypatch): class _CamConditioner: def __init__(self): self.calls = 0 @@ -336,6 +336,9 @@ def test_lingbot_cam_conditioner_cache_skips_context_update(monkeypatch): ) block.cam_conditioner = _CamConditioner() forward_batch = SimpleNamespace(extra={}, enable_sequence_shard=True) + monkeypatch.setattr( + lingbot_world_module, "get_ulysses_parallel_world_size", lambda: 2 + ) monkeypatch.setattr( lingbot_world_module, "get_forward_context", @@ -346,9 +349,9 @@ def test_lingbot_cam_conditioner_cache_skips_context_update(monkeypatch): first = block._cam_conditioner_scale_shift(c2ws_plucker_emb) second = block._cam_conditioner_scale_shift(c2ws_plucker_emb) - assert first is not second - assert block.cam_conditioner.calls == 2 - assert "lingbot_cam_conditioner" not in forward_batch.extra + assert first is second + assert block.cam_conditioner.calls == 1 + assert "lingbot_cam_conditioner" in forward_batch.extra def test_lingbot_model_prepares_cam_conditioner_scale_shifts(monkeypatch):