Skip to content

fix(exact-prefix): adapt LingBot binding to TeleFuser int KV indices #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yJader wants to merge 1 commit into
base: main
Choose a base branch
from
+20 −15
Open

fix(exact-prefix): adapt LingBot binding to TeleFuser int KV indices #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions cacheseek/reuse/exact_prefix/telefuser_lingbot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,8 +222,8 @@ def set_resume_depth(self, depth: int) -> None:
rt = self._rt
global_end = (depth + 1) * rt.chunk_size * rt.frame_tokens
for kv in rt.self_kv_cache:
kv["global_end_index"][0] = global_end
kv["local_end_index"][0] = self._local_end_tokens
kv["global_end_index"] = global_end
kv["local_end_index"] = self._local_end_tokens


class LingBotWorldKVBinding:
Expand Down Expand Up @@ -376,7 +376,7 @@ def on_chunk_finalized(
# In rolling mode this chunk's clean KV sits at the physical tail
# [local_end-ct : local_end] (logical position idx*ct holds only in
# full-length mode); local_end was just advanced by the clean rewrite.
e = int(kv["local_end_index"][0])
e = int(kv["local_end_index"])
s = e - ct
payload.append(
(
Expand Down
9 changes: 7 additions & 2 deletions examples/exact_prefix_reuse/e2e_telefuser_lingbot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@
from cacheseek.reuse.exact_prefix import NamespaceForest, WorldKVManager
from cacheseek.reuse.exact_prefix.telefuser_lingbot import (
LingBotWorldKVBinding,
make_full_kv_config,
make_rolling_config,
)
from cacheseek.stores import InMemoryTierStore
Expand Down Expand Up @@ -244,9 +245,13 @@ def make_store():

def fresh_stack():
forest = NamespaceForest()
world_kv_cfg = (
make_full_kv_config()
if int(cfg.local_attn_size) == -1
else make_rolling_config(local_attn_size=cfg.local_attn_size, sink_size=cfg.sink_size, chunk_size=3)
)
mgr = WorldKVManager(
forest, shared_store,
make_rolling_config(local_attn_size=cfg.local_attn_size, sink_size=cfg.sink_size, chunk_size=3),
forest, shared_store, world_kv_cfg,
)
return forest, mgr

Expand Down
20 changes: 10 additions & 10 deletions tests/test_world_kv_telefuser_binding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,8 +69,8 @@ def make_runtime(seed: int, actions: list[int]) -> SimpleNamespace:
{
"k": torch.zeros((1, KV_TOKENS, HEADS, HEAD_DIM)),
"v": torch.zeros((1, KV_TOKENS, HEADS, HEAD_DIM)),
"global_end_index": torch.tensor([0]),
"local_end_index": torch.tensor([0]),
"global_end_index": 0,
"local_end_index": 0,
}
for _ in range(N_LAYERS)
],
Expand All @@ -89,8 +89,8 @@ def kv_write(kv: dict, k_new: torch.Tensor, v_new: torch.Tensor, current_start:
(lingbot_world_fast_dit.py:108-145)."""
num_new = k_new.shape[1]
current_end = current_start + num_new
global_end = int(kv["global_end_index"][0])
local_end = int(kv["local_end_index"][0])
global_end = int(kv["global_end_index"])
local_end = int(kv["local_end_index"])
if current_end > global_end and num_new + local_end > KV_TOKENS:
evicted = num_new + local_end - KV_TOKENS
rolled = local_end - evicted - SINK_TOKENS
Expand All @@ -106,8 +106,8 @@ def kv_write(kv: dict, k_new: torch.Tensor, v_new: torch.Tensor, current_start:
local_start = local_end - num_new
kv["k"][:, local_start:local_end] = k_new
kv["v"][:, local_start:local_end] = v_new
kv["global_end_index"][0] = current_end
kv["local_end_index"][0] = local_end
kv["global_end_index"] = current_end
kv["local_end_index"] = local_end


def fake_denoise(runtime: SimpleNamespace, idx: int) -> torch.Tensor:
Expand All @@ -134,7 +134,7 @@ def snapshot(runtime: SimpleNamespace) -> list[dict]:
return [
{
"k": kv["k"].clone(), "v": kv["v"].clone(),
"ge": int(kv["global_end_index"][0]), "le": int(kv["local_end_index"][0]),
"ge": int(kv["global_end_index"]), "le": int(kv["local_end_index"]),
}
for kv in runtime.self_kv_cache
]
Expand Down Expand Up @@ -182,8 +182,8 @@ def test_warm_ring_equals_cold_ring_all_resume_points():
assert b.last_fast_forward == K, f"K={K}: matched {b.last_fast_forward}"
for layer, (kv, r) in enumerate(zip(rt.self_kv_cache, snaps[K], strict=False)):
le = r["le"]
assert int(kv["local_end_index"][0]) == le, f"K={K} layer{layer} local_end"
assert int(kv["global_end_index"][0]) == r["ge"], f"K={K} layer{layer} global_end"
assert int(kv["local_end_index"]) == le, f"K={K} layer{layer} local_end"
assert int(kv["global_end_index"]) == r["ge"], f"K={K} layer{layer} global_end"
assert torch.equal(kv["k"][:, :le], r["k"][:, :le]), f"K={K} layer{layer} k ring"
assert torch.equal(kv["v"][:, :le], r["v"][:, :le]), f"K={K} layer{layer} v ring"

Expand All @@ -206,7 +206,7 @@ def test_branch_resume_rng_aligned():
for i in (2, 3):
assert torch.equal(out_warm[i], out_cold[i]), f"chunk {i} diverged: RNG misaligned"
for lw, lc in zip(rt_warm.self_kv_cache, rt_cold.self_kv_cache, strict=False):
le = int(lc["local_end_index"][0])
le = int(lc["local_end_index"])
assert torch.equal(lw["k"][:, :le], lc["k"][:, :le]) # ring also matches after continuation

# After branch write-back, both paths hit the full chain.
Expand Down