diff --git a/single_shot_inference.py b/single_shot_inference.py
index 0040f402..232379a8 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -305,7 +305,7 @@ class Indexer:
 # KV Cache
 # =====================================================================
 class KVCache:
-    def __init__(self, head_dim, window_size=128, max_comp=32768, device='cuda:0'):
+    def __init__(self, head_dim, window_size=128, max_comp=65536, device='cuda:0'):
         self.hd, self.ws, self.dev = head_dim, window_size, device
         self.swa = torch.zeros(window_size, head_dim, dtype=torch.bfloat16, device=device)
         self.swa_pos = torch.zeros(window_size, dtype=torch.long, device=device)
@@ -886,7 +886,7 @@ def main():
     rt = rp.get("type", rp.get("rope_type", "yarn")); rf = rp.get("factor", 16.0)
     rtheta = cfg.get("rope_theta", 10000.); romax = rp.get("original_max_position_embeddings", 65536)
     rbfast, rbslow = rp.get("beta_fast", 32), rp.get("beta_slow", 1)
-    rope_caches = {g: build_rope_cache(8192, rd, f"cuda:{g}", rtheta, rt, rf, romax, rbfast, rbslow) for g in range(NUM_GPUS)}
+    rope_caches = {g: build_rope_cache(romax, rd, f"cuda:{g}", rtheta, rt, rf, romax, rbfast, rbslow) for g in range(NUM_GPUS)}
 
     # KV caches, compressors, indexers
     kv_caches, compressors, indexers = {}, {}, {}