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

P3: integrate CUDA RoPE kernel into single_shot — 732 launches/token eliminated

Browse Source 
_apply_rope now uses dsv4.ops.rope_cuda (1 CUDA kernel per call)
instead of PyTorch ops (5-6 kernels per call).
Total: 183 RoPE calls × (5-1) = 732 launches saved per token.
With fallback to PyTorch if CUDA kernel fails.
This commit is contained in:
biondizzle
2026-06-02 09:08:07 +00:00
parent 54a9b6961b
commit 3c295f225a
1 changed files with 20 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
single_shot_inference.py
View File

@@ -69,25 +69,28 @@ def build_rope_cache(max_pos, rope_dim, device, theta=10000., rope_type="default
return torch.cos(angles).to(device), torch.sin(angles).to(device)
def _apply_rope(x, pos, cos, sin, rope_dim, inverse=False):
"""In-place RoPE — mutates x, no full clone, no empty_like allocation.
"""In-place RoPE — uses CUDA kernel (1 launch) instead of PyTorch ops (5-6 launches).
P5: Eliminates x.clone() + empty_like per RoPE call.
Old: 183 calls/token × 128KB clone = 23MB pointless memcpy + 183 kernel launches.
New: Operates on the rope dims in-place, one slice copy back.
P3: Eliminates ~732 kernel launches per token across 61 layers.
"""
T, nh, hd = x.shape; nope = hd - rope_dim
if pos.device != cos.device: pos = pos.to(cos.device)
c, s = cos[pos].unsqueeze(1), sin[pos].unsqueeze(1)
xr = x[:, :, nope:] # view, not copy
ev = xr[..., 0::2].clone() # need original ev for the mix
od = xr[..., 1::2] # view; will be overwritten below
if inverse:
xr[..., 0::2] = (ev * c + od * s).bfloat16()
xr[..., 1::2] = (-ev * s + od * c).bfloat16()
else:
xr[..., 0::2] = (ev * c - od * s).bfloat16()
xr[..., 1::2] = (ev * s + od * c).bfloat16()
return x # mutated in place
try:
from dsv4.ops.rope_cuda import apply_rope
return apply_rope(x, pos, cos, sin, rope_dim, inverse=inverse)
except Exception:
# Fallback to PyTorch (should never happen in production)
T, nh, hd = x.shape; nope = hd - rope_dim
if pos.device != cos.device: pos = pos.to(cos.device)
c, s = cos[pos].unsqueeze(1), sin[pos].unsqueeze(1)
xr = x[:, :, nope:]
ev = xr[..., 0::2].clone()
od = xr[..., 1::2]
if inverse:
xr[..., 0::2] = (ev * c + od * s).bfloat16()
xr[..., 1::2] = (-ev * s + od * c).bfloat16()
else:
xr[..., 0::2] = (ev * c - od * s).bfloat16()
xr[..., 1::2] = (ev * s + od * c).bfloat16()
return x
# =====================================================================
# Weight loading