diff --git a/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py b/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py index 81708730e..c61037338 100644 --- a/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py +++ b/python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py @@ -96,6 +96,13 @@ class CausalSelfAttentionKVCache: if self.attention_window_size == old_cache_size: self.attention_window_size = new_cache_size + def can_direct_current_attention(self, num_new_tokens: int) -> bool: + return ( + self.sink_tokens == 0 + and self.cache_size == num_new_tokens + and self.attention_window_size == num_new_tokens + ) + def update_and_get_attention_kv( self, *, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 0366b029f..0652945ef 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py @@ -183,6 +183,9 @@ class CausalWanSelfAttention(nn.Module): block_mask=block_mask, )[:, :, :-padded_length].transpose(2, 1) else: + if kv_cache.can_direct_current_attention(roped_key.shape[1]): + return self.attn(roped_query, roped_key, v) + cache_view = kv_cache.update_and_get_attention_kv( key=roped_key, value=v, 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 766ef0c21..b455e8fdb 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py @@ -282,6 +282,13 @@ class LingBotWorldCausalSelfAttention(CausalWanSelfAttention): ) roped_query, roped_key, v = qkv.chunk(3, dim=-1) + if ( + not sequence_shard_enabled + and not update_cache_only + and kv_cache.can_direct_current_attention(roped_key.shape[1]) + ): + return self.attn(roped_query, roped_key, v) + cache_head_start = ( get_tp_rank() * roped_key.shape[2] if sequence_shard_enabled diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_causal_denoising.py b/python/sglang/multimodal_gen/test/unit/realtime/test_causal_denoising.py index fb864bba4..59689ec18 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_causal_denoising.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_causal_denoising.py @@ -399,6 +399,24 @@ def test_causal_kv_cache_allocation_sets_shapes_and_optional_int_indices(): assert cache[0].allow_growth is False +def test_causal_kv_cache_can_direct_current_attention(): + cache = CausalSelfAttentionKVCache( + k=torch.zeros(1, 3, 1, 1), + v=torch.zeros(1, 3, 1, 1), + global_end_index=torch.zeros(1, dtype=torch.long), + local_end_index=torch.zeros(1, dtype=torch.long), + cache_size=3, + sink_tokens=0, + attention_window_size=3, + ) + + assert cache.can_direct_current_attention(3) + assert not cache.can_direct_current_attention(2) + + cache.sink_tokens = 1 + assert not cache.can_direct_current_attention(3) + + def test_causal_kv_cache_update_handles_append_roll_and_recompute(): cache = CausalSelfAttentionKVCache( k=torch.zeros(1, 4, 1, 1),