Fix fused RoPE backward pass computing wrong gradients (#2)

_FusedRoPE.backward reused the forward kernel verbatim with the same
sin, but RoPE's forward is a rotation matrix [[c,-s],[s,c]] applied to
each (x0,x1) pair, so the correct backward is that matrix's transpose
— equivalent to the same kernel with sin negated, not reapplied
unchanged. Every gradient through the fused path was wrong whenever
sin != 0 (max abs error 3.64 in the reported repro vs 0.0 after the
fix), silently corrupting training for any model using
RotaryEmbedding(use_fused=True).

test_gradient_flows only checked grad is not None, so this shipped
undetected. Added test_gradient_correctness, which compares the fused
path's gradient against the vanilla autograd-correct path via
assert_close, and fails against the old code (85% of elements
mismatched) while passing with the fix.

Reported by @tohtana during review of deepspeedai/DeepSpeed#8052.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: py-ai-dev

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "torchembed"
-version = "0.3.0"
+version = "0.3.1"
 description = "A PyTorch library of modern embedding strategies missing from torch.nn"
 readme = "README.md"
 license = { file = "LICENSE" }

tests/test_positional.py

@@ -140,6 +140,42 @@ def test_gradient_flows(self, dim):
         assert q.grad is not None
         assert k.grad is not None
 
+    def test_gradient_correctness(self, dim):
+        """Fused kernel gradients must numerically match the vanilla path's.
+
+        RoPE's forward is a per-position rotation matrix [[c, -s], [s, c]] applied
+        to each (x0, x1) pair; its backward is that matrix's transpose, not the
+        forward transform reapplied unchanged. A prior bug reused the forward
+        kernel verbatim for backward, which silently produced wrong gradients
+        for every dim with sin != 0. `torch.testing.assert_close` on the grads
+        (not just shape/is-not-None) is what catches a regression here.
+        """
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA not available")
+        rope = RotaryEmbedding(dim=dim, device="cuda")
+        cos, sin = rope.cos_cache[:8], rope.sin_cache[:8]
+
+        torch.manual_seed(0)
+        q_base = torch.randn(2, 4, 8, dim, device="cuda")
+        k_base = torch.randn(2, 4, 8, dim, device="cuda")
+        grad_q_up = torch.randn_like(q_base)
+        grad_k_up = torch.randn_like(k_base)
+
+        q_ref = q_base.clone().requires_grad_(True)
+        k_ref = k_base.clone().requires_grad_(True)
+        q_rot_ref, k_rot_ref = rope(q_ref, k_ref)
+        q_rot_ref.backward(grad_q_up, retain_graph=True)
+        k_rot_ref.backward(grad_k_up)
+
+        q_fused = q_base.clone().requires_grad_(True)
+        k_fused = k_base.clone().requires_grad_(True)
+        q_rot_fused, k_rot_fused = rope._fused_forward(q_fused, k_fused, cos, sin)
+        q_rot_fused.backward(grad_q_up, retain_graph=True)
+        k_rot_fused.backward(grad_k_up)
+
+        torch.testing.assert_close(q_fused.grad, q_ref.grad, atol=1e-4, rtol=1e-4)
+        torch.testing.assert_close(k_fused.grad, k_ref.grad, atol=1e-4, rtol=1e-4)
+
     def test_fused_forward_function(self, dim):
         """Direct call to fused_rope_forward should match."""
         if not torch.cuda.is_available():

torchembed/_triton.py

@@ -91,7 +91,10 @@ def forward(ctx, x, cos, sin):
         @staticmethod
         def backward(ctx, grad_output):
             cos, sin = ctx.saved_tensors
-            grad_input = _fused_rope_forward_core(grad_output, cos, sin)
+            # Forward applies rotation matrix [[c, -s], [s, c]] to x. Its gradient
+            # w.r.t. x is the transpose [[c, s], [-s, c]], which is the same rotation
+            # kernel run with sin negated (a rotation matrix's transpose is its inverse).
+            grad_input = _fused_rope_forward_core(grad_output, cos, -sin)
             return grad_input, None, None
 
     def _fused_rope_forward_core(x, cos, sin):