biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

FIX: RoPE cache 8192→65536 (original_max_position_embeddings), KVCache max_comp 32768→65536

Browse Source
This commit is contained in:
biondizzle
2026-06-01 23:18:37 +00:00
parent 36f9782bad
commit 73ae9393da
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
single_shot_inference.py
View File

@@ -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 = {}, {}, {}