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

Fix: thread comp_rope_cos/comp_rope_sin through forward_attention

Browse Source 
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.
This commit is contained in:
biondizzle
2026-06-03 09:30:57 +00:00
parent 6a3374da18
commit 284fc9ca86
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
single_shot_inference.py
View File

@@ -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)