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 823c28722..81708730e 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 @@ -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[ - :, - 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[ - :, - sink_tokens - + num_evicted_tokens : sink_tokens - + num_evicted_tokens - + num_rolled_tokens, - ].clone() + 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[ + :, + 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() - self.k[:, local_start_index:local_end_index] = key - self.v[:, local_start_index:local_end_index] = value - 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, : + ] + 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], - v=self.v[:, attn_start_index:updated_local_end], + k=visible_k, + v=visible_v, local_start_index=local_start_index, local_end_index=local_end_index, visible_local_end=updated_local_end, 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 132a35fb0..0366b029f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/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, - value=cache_value, + key=roped_key, + 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, - value, + cache_view.k, + 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 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 248f8ae6e..766ef0c21 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py +++ b/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 - if kv_cache.k.shape[2] != roped_key.shape[2]: - if sequence_shard_enabled: - head_start = get_tp_rank() * roped_key.shape[2] - 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_head_start = ( + get_tp_rank() * roped_key.shape[2] + if sequence_shard_enabled + else self.head_start + ) cache_view = kv_cache.update_and_get_attention_kv( - key=cache_key, - value=cache_value, + key=roped_key, + 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, - value, + cache_view.k, + cache_view.v, ) if sequence_shard_enabled: assert seq_splits is not None diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index 382d95d42..1f3458797 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/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, 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 ca0e0d678..fb864bba4 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 @@ -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),