minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-06-22 21:00:38 +08:00
committed by GitHub
parent 1adb53f147
commit ead39d38fc
5 changed files with 158 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py
+84 -21
View File
@@ -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,
python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
+6 -34
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py
+10 -39
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py
+1
View File
@@ -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
+57
View File
@@ -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),