wip: T10 batched forward (validation)

crates/xtrain-tensor/src/tensor.rs

@@ -454,13 +454,20 @@ impl Tensor {
         dx
     }
 
-    /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; the position
-    /// of each token is its row index. Returns the rotated tensor.
+    /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; each token's
+    /// position is `row % period`. `period` = sequence length, so a flattened
+    /// batch `[B*S,heads,head_dim]` gets per-sequence positions (pass `period=S`);
+    /// pass `period=tokens` for a single sequence (position = row). Returns the
+    /// rotated tensor.
     #[cfg(not(no_cuda))]
-    pub fn rope(&self, theta: f32) -> Self {
+    pub fn rope(&self, theta: f32, period: usize) -> Self {
         assert_eq!(self.ndim(), 3, "rope requires [tokens,heads,head_dim]");
         let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
         assert_eq!(head_dim % 2, 0, "head_dim must be even");
+        assert!(
+            period > 0 && tokens % period == 0,
+            "tokens must be a multiple of period"
+        );
         let out = Tensor::zeros(&self.shape, DType::F32, self.device());
         unsafe {
             xtrain_cuda::ffi::launch_rope_f32(
@@ -470,6 +477,7 @@ impl Tensor {
                 heads as i32,
                 head_dim as i32,
                 theta,
+                period as i32,
                 std::ptr::null_mut(),
             );
         }
@@ -477,9 +485,9 @@ impl Tensor {
     }
 
     /// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
-    /// orthogonal map, so it needs no cached forward values, only `theta`.
+    /// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
     #[cfg(not(no_cuda))]
-    pub fn rope_backward(dy: &Tensor, theta: f32) -> Self {
+    pub fn rope_backward(dy: &Tensor, theta: f32, period: usize) -> Self {
         let (tokens, heads, head_dim) = (dy.shape[0], dy.shape[1], dy.shape[2]);
         let dx = Tensor::zeros(&dy.shape, DType::F32, dy.device());
         unsafe {
@@ -490,6 +498,7 @@ impl Tensor {
                 heads as i32,
                 head_dim as i32,
                 theta,
+                period as i32,
                 std::ptr::null_mut(),
             );
         }