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

Cross-check 2 complete: block-aligned comp_pos + compress_rope_theta wired through

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

15
single_shot_inference.py
View File

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