[diffusion] optimize: optimize realtime causal attention fastpath (#28760)

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,
         *,

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,

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

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),