From 3c295f225a577a0294c90eb432bf22c9383374b0 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Tue, 2 Jun 2026 09:08:07 +0000
Subject: [PATCH] =?UTF-8?q?P3:=20integrate=20CUDA=20RoPE=20kernel=20into?=
 =?UTF-8?q?=20single=5Fshot=20=E2=80=94=20732=20launches/token=20eliminate?=
 =?UTF-8?q?d?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_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.
---
 single_shot_inference.py | 37 ++++++++++++++++++++-----------------
 1 file changed, 20 insertions(+), 17 deletions(-)

diff --git a/single_shot_inference.py b/single_shot_inference.py
index 82afad7f..dff625ca 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -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