diff --git a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py index 15def69..9a38205 100644 --- a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py +++ b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py @@ -50,27 +50,36 @@ from .types import ActionKey -def make_full_kv_config(*, break_even_k: int = 1) -> WorldKVConfig: - """Full-length KV (local_attn_size=-1) => materialization needs the entire - prefix.""" - return WorldKVConfig( - window_chunks=1_000_000, sink_chunks=1, break_even_k=break_even_k - ) - - -def make_rolling_config( +def make_world_kv_config( *, local_attn_size: int = 7, # latent frames, incl. sink (matches TeleFuser pipeline config) sink_size: int = 3, # latent frames chunk_size: int = 3, # latent frames per chunk break_even_k: int = 1, ) -> WorldKVConfig: - """Rolling window: materialization needs only the sink chunk plus the chunks - covering the most recent (L-S) frames -- O(W) rather than O(K).""" - recent_frames = max(local_attn_size - sink_size, 1) + """Build WorldKVConfig from TeleFuser LingBot KV geometry. + + Uses local_attn_size=-1 for full-length KV. Positive values for + rolling KV. + + Rolling window: materialization needs only the sink chunk plus the chunks + covering the most recent (L-S) frames -- O(W) rather than O(K). + """ + if local_attn_size == -1: + window_chunks = 1_000_000 + sink_chunks = 1 + elif local_attn_size > 0: + recent_frames = max(local_attn_size - sink_size, 1) + + window_chunks=-(-recent_frames // chunk_size) # ceil + sink_chunks=-(-sink_size // chunk_size) + else: + raise ValueError(f"local_attn_size must be -1 for full KV and positive " + f"for rolling, got {local_attn_size}") + return WorldKVConfig( - window_chunks=-(-recent_frames // chunk_size), # ceil - sink_chunks=-(-sink_size // chunk_size), + window_chunks=window_chunks, + sink_chunks=sink_chunks, break_even_k=break_even_k, ) @@ -222,8 +231,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 +385,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..c84b8b6 100644 --- a/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py +++ b/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py @@ -69,7 +69,7 @@ from cacheseek.reuse.exact_prefix import NamespaceForest, WorldKVManager from cacheseek.reuse.exact_prefix.telefuser_lingbot import ( LingBotWorldKVBinding, - make_rolling_config, + make_world_kv_config, ) from cacheseek.stores import InMemoryTierStore @@ -147,6 +147,7 @@ def run_request(pipeline, binding, frame_num: int, poses, intrinsics, *, image_p control_mode="cam", frame_num=frame_num, seed=SEED, + sammple_shift=10, # better generation quality poses=poses, intrinsics=intrinsics, world_kv_binding=binding, @@ -246,7 +247,11 @@ def fresh_stack(): forest = NamespaceForest() mgr = WorldKVManager( forest, shared_store, - make_rolling_config(local_attn_size=cfg.local_attn_size, sink_size=cfg.sink_size, chunk_size=3), + make_world_kv_config( + local_attn_size=cfg.local_attn_size, + sink_size=cfg.sink_size, + chunk_size=3 + ), ) return forest, mgr diff --git a/tests/test_world_kv_telefuser_binding.py b/tests/test_world_kv_telefuser_binding.py index b9bd011..aa3e84a 100644 --- a/tests/test_world_kv_telefuser_binding.py +++ b/tests/test_world_kv_telefuser_binding.py @@ -29,7 +29,7 @@ ) from cacheseek.reuse.exact_prefix.telefuser_lingbot import ( # noqa: E402 LingBotWorldKVBinding, - make_rolling_config, + make_world_kv_config, ) # Tiny geometry: chunk=3 frames, frame_tokens=2; window L=7 frames including @@ -159,7 +159,7 @@ def run_session(binding, runtime, session, *, snapshots: list | None = None): def make_stack(): forest = NamespaceForest() - cfg = make_rolling_config(local_attn_size=LOCAL_ATTN, sink_size=SINK, chunk_size=CHUNK_FRAMES) + cfg = make_world_kv_config(local_attn_size=LOCAL_ATTN, sink_size=SINK, chunk_size=CHUNK_FRAMES) return forest, WorldKVManager(forest, InMemoryTierStore(), cfg)