From 73ae9393da9d12af599cbbe089645101ac90390a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 23:18:37 +0000 Subject: [PATCH] =?UTF-8?q?FIX:=20RoPE=20cache=208192=E2=86=9265536=20(ori?= =?UTF-8?q?ginal=5Fmax=5Fposition=5Fembeddings),=20KVCache=20max=5Fcomp=20?= =?UTF-8?q?32768=E2=86=9265536?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- single_shot_inference.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 = {}, {}, {}