From 6a3374da1874321799024e4b164a1f71dae7dec7 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 09:19:11 +0000 Subject: [PATCH] Cross-check 2 complete: block-aligned comp_pos + compress_rope_theta wired through - Fixed comp_pos: (bi*r) block-aligned instead of ((bi+1)*r-1) last-position - compress_rope_theta: separate rope cache for compressed KV entries - comp_rope_cos/comp_rope_sin wired to all forward_layer call sites (prefill chunk loop, decode loop, CUDAGraphDecoder capture) - forward_layer uses comp_rope caches for compressed RoPE, falls back to normal - Only single_shot_inference.py modified, no kernel code touched --- single_shot_inference.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index 29d841a0..6f3d34fa 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -158,7 +158,7 @@ class CUDAGraphDecoder: def pre_allocate(self, cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms, kv_caches, compressors, indexers, moe_runners, se_runners, routers, prod_lins, layer_w, rope_caches, hc_head, - final_norm_w, lm_head_lin): + final_norm_w, lm_head_lin, comp_rope_caches=None): """Pre-allocate all I/O buffers with fixed addresses.""" for li in range(self.n_layers): dev = self.devices[li % self.num_gpus] @@ -170,7 +170,7 @@ class CUDAGraphDecoder: def capture(self, cfg, attn_mhcs, ffn_mhcs, attn_norms, ffn_norms, kv_caches, compressors, indexers, moe_runners, se_runners, routers, prod_lins, layer_w, rope_caches, hc_head, - final_norm_w, lm_head_lin, positions, token_id): + final_norm_w, lm_head_lin, positions, token_id, comp_rope_caches=None): """Capture CUDA graphs for all layers + lm_head. Must be called after one warmup step so that: @@ -199,7 +199,9 @@ class CUDAGraphDecoder: compressors.get(li), indexers.get(li), moe_runners.get(li), se_runners.get(li), routers.get(li), prod_lin=prod_lins.get(li), - _use_fused_rmsnorm_quantize=True + _use_fused_rmsnorm_quantize=True, + comp_rope_cos=comp_rope_caches[gpu][0] if comp_rope_caches else None, + comp_rope_sin=comp_rope_caches[gpu][1] if comp_rope_caches else None, ) # Copy output to fixed buffer self.x_out_bufs[li].copy_(X_out) @@ -855,7 +857,10 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, rope_bf16 = comp_kv_fp32[:, nope_dim:].bfloat16().contiguous() # (n_comp, 64) BF16 # Apply RoPE on BF16 rope dims (existing BF16 RoPE kernel) rope_3d = rope_bf16.unsqueeze(1) # (n_comp, 1, 64) - rope_3d = _apply_rope(rope_3d, comp_pos, rope_cos, rope_sin, rd) + # Use compress_rope_theta cache for compressed entries if available + crc = comp_rope_cos if comp_rope_cos is not None else rope_cos + crs = comp_rope_sin if comp_rope_sin is not None else rope_sin + rope_3d = _apply_rope(rope_3d, comp_pos, crc, crs, rd) rope_bf16 = rope_3d.squeeze(1) # (n_comp, 64) BF16 # Quantize non-RoPE part FP32 → FP8_E4M3 nope_fp8, nope_scale = kv_mod.quantize_fp8_e4m3_from_fp32(nope_fp32) @@ -1527,6 +1532,7 @@ def main(): moe_runners.get(li), se_runners.get(li), routers.get(li), prod_lin=prod_lins.get(li), _use_fused_rmsnorm_quantize=not _args.no_fused_rmsnorm, + comp_rope_cos=comp_rope_caches[gpu][0], comp_rope_sin=comp_rope_caches[gpu][1], ) except Exception as e: torch.cuda.synchronize() @@ -1607,6 +1613,7 @@ def main(): _profile_detail=(profile and step == 1), _profile_times=cuda_layer_events if (profile and step == 1) else None, _use_fused_rmsnorm_quantize=not _args.no_fused_rmsnorm, + comp_rope_cos=comp_rope_caches[gpu][0], comp_rope_sin=comp_rope_caches[gpu][1], ) X = X.to('cuda:0'); torch.cuda.set_device(0) t_layers = time.perf_counter()