[diffusion] refactor: refactor causal KV local head cache updates (#28888)

python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py

@@ -102,21 +102,36 @@ class CausalSelfAttentionKVCache:
         key: torch.Tensor,
         value: torch.Tensor,
         current_chunk_start: int,
+        cache_head_start: int | None = None,
         debug_name: str = "causal KV cache",
     ) -> CausalAttentionKVView:
-        """write kv into the cache, returns the part visible to the current chunk
+        """write fresh kv into the cache, returns the part of view visible to the current chunk
 
         Args:
             current_chunk_start: the global position of the start of the chunk
+            cache_head_start: first cache head for key/value when they only
+                carry a local slice of the cache heads; other heads are left untouched
 
         """
         num_new_tokens = key.shape[1]
+        num_input_heads = key.shape[2]
+        num_cache_heads = self.k.shape[2]
+        cache_head_slice = None
+        if num_cache_heads != num_input_heads:
+            if cache_head_start is None:
+                raise ValueError(
+                    f"{debug_name} requires cache_head_start when cache heads "
+                    f"({num_cache_heads}) differ from input heads ({num_input_heads})."
+                )
+            cache_head_slice = slice(
+                cache_head_start, cache_head_start + num_input_heads
+            )
         current_chunk_end = current_chunk_start + num_new_tokens
         kv_cache_size = self.cache_size
         sink_tokens = self.sink_tokens
         global_end_index, local_end_index_prev = self._read_indices()
 
-        # local_end_index: the local position of the end of current chunk
+        # local_start(/end)_index: the local position of the start/end of current chunk
         # updated_local_end: the updated local end
         # updated_global_end: the updated global end
 
@@ -159,20 +174,54 @@ class CausalSelfAttentionKVCache:
                     local_end_index_prev - num_evicted_tokens - sink_tokens,
                 )
                 if num_rolled_tokens > 0:
-                    self.k[:, sink_tokens : sink_tokens + num_rolled_tokens] = self.k[
+                    if cache_head_slice is None:
+                        self.k[:, sink_tokens : sink_tokens + num_rolled_tokens] = (
+                            self.k[
                                 :,
                                 sink_tokens
                                 + num_evicted_tokens : sink_tokens
                                 + num_evicted_tokens
                                 + num_rolled_tokens,
                             ].clone()
-                    self.v[:, sink_tokens : sink_tokens + num_rolled_tokens] = self.v[
+                        )
+                        self.v[:, sink_tokens : sink_tokens + num_rolled_tokens] = (
+                            self.v[
                                 :,
                                 sink_tokens
                                 + num_evicted_tokens : sink_tokens
                                 + num_evicted_tokens
                                 + num_rolled_tokens,
                             ].clone()
+                        )
+                    else:
+                        self.k[
+                            :,
+                            sink_tokens : sink_tokens + num_rolled_tokens,
+                            cache_head_slice,
+                            :,
+                        ] = self.k[
+                            :,
+                            sink_tokens
+                            + num_evicted_tokens : sink_tokens
+                            + num_evicted_tokens
+                            + num_rolled_tokens,
+                            cache_head_slice,
+                            :,
+                        ].clone()
+                        self.v[
+                            :,
+                            sink_tokens : sink_tokens + num_rolled_tokens,
+                            cache_head_slice,
+                            :,
+                        ] = self.v[
+                            :,
+                            sink_tokens
+                            + num_evicted_tokens : sink_tokens
+                            + num_evicted_tokens
+                            + num_rolled_tokens,
+                            cache_head_slice,
+                            :,
+                        ].clone()
 
                 # if we move the minimum number of tokens, the right bound of the append token would be aligned with end of the buffer
                 local_end_index = kv_cache_size
@@ -203,17 +252,31 @@ class CausalSelfAttentionKVCache:
             self.k = self.k.detach()
         if self.v.requires_grad:
             self.v = self.v.detach()
+        attn_start_index = max(0, updated_local_end - self.attention_window_size)
+
+        # write fresh kv and return visible view
+        if cache_head_slice is None:
             self.k[:, local_start_index:local_end_index] = key
             self.v[:, local_start_index:local_end_index] = value
+            visible_k = self.k[:, attn_start_index:updated_local_end]
+            visible_v = self.v[:, attn_start_index:updated_local_end]
+        else:
+            self.k[:, local_start_index:local_end_index, cache_head_slice, :] = key
+            self.v[:, local_start_index:local_end_index, cache_head_slice, :] = value
+            visible_k = self.k[
+                :, attn_start_index:updated_local_end, cache_head_slice, :
+            ]
+            visible_v = self.v[
+                :, attn_start_index:updated_local_end, cache_head_slice, :
+            ]
 
-        attn_start_index = max(0, updated_local_end - self.attention_window_size)
         self._write_indices(
             global_end_index=updated_global_end,
             local_end_index=updated_local_end,
         )
         return CausalAttentionKVView(
-            k=self.k[:, attn_start_index:updated_local_end],
+            k=visible_k,
-            v=self.v[:, attn_start_index:updated_local_end],
+            v=visible_v,
             local_start_index=local_start_index,
             local_end_index=local_end_index,
             visible_local_end=updated_local_end,

python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py

@@ -183,46 +183,17 @@ class CausalWanSelfAttention(nn.Module):
                 block_mask=block_mask,
             )[:, :, :-padded_length].transpose(2, 1)
         else:
-            head_slice = None
-            if kv_cache.k.shape[2] != roped_key.shape[2]:
-                head_slice = slice(self.head_start, self.head_start + self.num_heads)
-                cache_key = roped_key.new_zeros(
-                    roped_key.shape[0],
-                    roped_key.shape[1],
-                    kv_cache.k.shape[2],
-                    roped_key.shape[3],
-                )
-                cache_value = v.new_zeros(
-                    v.shape[0],
-                    v.shape[1],
-                    kv_cache.v.shape[2],
-                    v.shape[3],
-                )
-                cache_key[:, :, head_slice, :] = roped_key
-                cache_value[:, :, head_slice, :] = v
-            else:
-                cache_key = roped_key
-                cache_value = v
             cache_view = kv_cache.update_and_get_attention_kv(
-                key=cache_key,
+                key=roped_key,
-                value=cache_value,
+                value=v,
                 current_chunk_start=current_start,
+                cache_head_start=self.head_start,
                 debug_name="CausalWan KV cache",
             )
-            key = (
-                cache_view.k[:, :, head_slice, :]
-                if head_slice is not None
-                else cache_view.k
-            )
-            value = (
-                cache_view.v[:, :, head_slice, :]
-                if head_slice is not None
-                else cache_view.v
-            )
             x = self.attn(
                 roped_query,
-                key,
+                cache_view.k,
-                value,
+                cache_view.v,
             )
 
         return x
@@ -285,6 +256,7 @@ class CausalWanTransformerBlock(nn.Module):
                 quant_config=quant_config,
                 prefix=add_prefix("to_out", prefix),
             )
+            # megatron-style tp shards the weight (qkv) column-wise, effectively splitting the attention heads
             tp_size = get_tp_world_size()
             self.local_num_heads = divide(num_heads, tp_size)
             head_start = get_tp_rank() * self.local_num_heads

python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py

@@ -282,54 +282,25 @@ class LingBotWorldCausalSelfAttention(CausalWanSelfAttention):
             )
             roped_query, roped_key, v = qkv.chunk(3, dim=-1)
 
-        head_slice = None
+        cache_head_start = (
-        if kv_cache.k.shape[2] != roped_key.shape[2]:
+            get_tp_rank() * roped_key.shape[2]
-            if sequence_shard_enabled:
+            if sequence_shard_enabled
-                head_start = get_tp_rank() * roped_key.shape[2]
+            else self.head_start
-            else:
-                head_start = self.head_start
-            head_slice = slice(head_start, head_start + roped_key.shape[2])
-            cache_key = roped_key.new_zeros(
-                roped_key.shape[0],
-                roped_key.shape[1],
-                kv_cache.k.shape[2],
-                roped_key.shape[3],
         )
-            cache_value = v.new_zeros(
-                v.shape[0],
-                v.shape[1],
-                kv_cache.v.shape[2],
-                v.shape[3],
-            )
-            cache_key[:, :, head_slice, :] = roped_key
-            cache_value[:, :, head_slice, :] = v
-        else:
-            cache_key = roped_key
-            cache_value = v
-
         cache_view = kv_cache.update_and_get_attention_kv(
-            key=cache_key,
+            key=roped_key,
-            value=cache_value,
+            value=v,
             current_chunk_start=current_start,
+            cache_head_start=cache_head_start,
             debug_name="LingBot KV cache",
         )
         if update_cache_only:
             return v
-        key = (
-            cache_view.k[:, :, head_slice, :]
-            if head_slice is not None
-            else cache_view.k
-        )
-        value = (
-            cache_view.v[:, :, head_slice, :]
-            if head_slice is not None
-            else cache_view.v
-        )
         attn_impl = self.ulysses_attn if sequence_shard_enabled else self.attn
         x = attn_impl(
             roped_query,
-            key,
+            cache_view.k,
-            value,
+            cache_view.v,
         )
         if sequence_shard_enabled:
             assert seq_splits is not None

python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py

@@ -789,6 +789,7 @@ class CausalDMDDenoisingStage(DenoisingStage):
             prepare_model_input=prepare_model_input,
             progress_bar=progress_bar,
         )
+        # after clean latent is generated, fill the causal kv cache with a forward with clean latent as input
         self._update_causal_context_cache(
             batch,
             server_args,

python/sglang/multimodal_gen/test/unit/realtime/test_causal_denoising.py

@@ -445,6 +445,63 @@ def test_causal_kv_cache_update_handles_append_roll_and_recompute():
     assert recompute_view.k.flatten().tolist() == [50.0, 6.0]
 
 
+def test_causal_kv_cache_update_for_cache_head_slice_returns_local_view():
+    cache = CausalSelfAttentionKVCache(
+        k=torch.full((1, 4, 4, 1), 99.0),
+        v=torch.full((1, 4, 4, 1), 999.0),
+        global_end_index=torch.zeros(1, dtype=torch.long),
+        local_end_index=torch.zeros(1, dtype=torch.long),
+        global_end_index_int=0,
+        local_end_index_int=0,
+    )
+
+    view = cache.update_and_get_attention_kv(
+        key=torch.tensor([[[[1.0], [2.0]], [[3.0], [4.0]]]]),
+        value=torch.tensor([[[[10.0], [20.0]], [[30.0], [40.0]]]]),
+        current_chunk_start=0,
+        cache_head_start=1,
+    )
+
+    assert view.k.shape == (1, 2, 2, 1)
+    assert view.v.shape == (1, 2, 2, 1)
+    assert view.k.flatten().tolist() == [1.0, 2.0, 3.0, 4.0]
+    assert view.v.flatten().tolist() == [10.0, 20.0, 30.0, 40.0]
+    assert cache.k[0, :2, :, 0].tolist() == [
+        [99.0, 1.0, 2.0, 99.0],
+        [99.0, 3.0, 4.0, 99.0],
+    ]
+
+    cache.update_and_get_attention_kv(
+        key=torch.tensor([[[[5.0], [6.0]], [[7.0], [8.0]]]]),
+        value=torch.tensor([[[[50.0], [60.0]], [[70.0], [80.0]]]]),
+        current_chunk_start=2,
+        cache_head_start=1,
+    )
+    rolled_view = cache.update_and_get_attention_kv(
+        key=torch.tensor([[[[9.0], [10.0]], [[11.0], [12.0]]]]),
+        value=torch.tensor([[[[90.0], [100.0]], [[110.0], [120.0]]]]),
+        current_chunk_start=4,
+        cache_head_start=1,
+    )
+
+    assert rolled_view.k.flatten().tolist() == [
+        5.0,
+        6.0,
+        7.0,
+        8.0,
+        9.0,
+        10.0,
+        11.0,
+        12.0,
+    ]
+    assert cache.k[0, :, :, 0].tolist() == [
+        [99.0, 5.0, 6.0, 99.0],
+        [99.0, 7.0, 8.0, 99.0],
+        [99.0, 9.0, 10.0, 99.0],
+        [99.0, 11.0, 12.0, 99.0],
+    ]
+
+
 def test_causal_kv_cache_update_grows_without_rolling_when_enabled():
     cache = CausalSelfAttentionKVCache(
         k=torch.zeros(1, 2, 1, 1),