diff --git a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py index 15def69..f66553f 100644 --- a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py +++ b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py @@ -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: @@ -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( ( diff --git a/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py b/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py index 6cfbe8c..0a4f2e2 100644 --- a/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py +++ b/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py @@ -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 @@ -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 diff --git a/tests/test_world_kv_telefuser_binding.py b/tests/test_world_kv_telefuser_binding.py index b9bd011..bf5310d 100644 --- a/tests/test_world_kv_telefuser_binding.py +++ b/tests/test_world_kv_telefuser_binding.py @@ -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) ], @@ -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 @@ -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: @@ -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 ] @@ -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" @@ -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.