From 284fc9ca867618a137031147c7ebb455612a0754 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 09:30:57 +0000 Subject: [PATCH] Fix: thread comp_rope_cos/comp_rope_sin through forward_attention MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previous commit added params to forward_layer but forward_attention (where compressed RoPE is applied) didn't receive them, causing NameError. Also confirmed from B200 test output: compress_rope_theta=160000 vs rope_theta=10000 — a 16x difference. The separate cache is essential. --- single_shot_inference.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index 6f3d34fa..fc6afeb1 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -789,7 +789,8 @@ def _run_production_fmha_mixed(q_heads, kv_nope_fp8, kv_nope_scale, kv_rope_bf16 def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, kv_cache, positions, compressor, indexer, prod_lin, x_quant=None, - _profile_detail=False, _profile_times=None): + _profile_detail=False, _profile_times=None, + comp_rope_cos=None, comp_rope_sin=None): dev = x_normed.device; T = x_normed.shape[0] n_h = cfg["num_attention_heads"]; hd = cfg["head_dim"]; rd = cfg.get("qk_rope_head_dim", 64) o_groups = cfg.get("o_groups", 16); o_rank = cfg.get("o_lora_rank", 1024) @@ -1030,7 +1031,8 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, F_attn, _ = forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, kv_cache, positions, compressor, indexer, prod_lin, x_quant=x_quant_attn, - _profile_detail=_profile_detail, _profile_times=_profile_times) + _profile_detail=_profile_detail, _profile_times=_profile_times, + comp_rope_cos=comp_rope_cos, comp_rope_sin=comp_rope_sin) if _profile_detail: torch.cuda.synchronize(); t_attn1 = time.perf_counter() X_mid = attn_mhc.post_block(X_l, F_attn, ctx_a)