rtx-backend-cuda: transpose of tall tensors failed to launch, and the generic permute was an identity copy

Found by running dg-gnn's HetGAT training on an RTX 5060 Ti: the first backward
pass panicked in transpose_2d with CUDA_ERROR_INVALID_VALUE. The tensor was a
[4,211,136 x 1] gradient. The tiled kernel maps rows onto grid.y, and CUDA caps
grid.y at 65,535 blocks, so anything past ~2.1M rows was rejected at launch.
Inference never hit it — the forward pass has no transposes.

Three fixes, each tested on the GPU:

  - A row or column vector transposes to itself in memory ([N,1] and [1,N] are
    the same N floats). swap_dims now shares the buffer and swaps the shape; no
    kernel. This is the case that actually failed.
  - transpose_2d_gpu hands matrices with rows/32 > 65,535 to the 1-D generic
    permute, whose grid.x allows 2^31-1 blocks.
  - transpose_generic_gpu passed the INPUT's strides (swapped) as the kernel's
    output_strides — but the kernel DECODES each flat output index with those,
    so they must be the contiguous strides of the output shape. For a 2-D input
    that was [1, cols], which decodes every index to itself: the "permute" was
    a plain copy. Verified against the original code — a [2,3,4] swap_dims(0,2)
    returned b[1][0][0] = 12.0 where 1.0 is correct. Every higher-D dim swap
    went through this path; nothing had checked it numerically.

Also stamps kernel outputs with the contiguous strides of their new shape
rather than the input's strides swapped. Nothing outside ops/shape.rs reads
strides today, so that was latent, but is_contiguous() now tells the truth.

rtx-backend-cuda --features cuda: 55 + 16 passed, 0 failed (was 51 + 16; four
new tests, one of which fails on the original code for each bug above).

Co-Authored-By: Claude Opus 5 <[email protected]>

crates/core/rtx-backend-cuda/src/ops/shape.rs

@@ -66,6 +66,20 @@ pub fn swap_dims<const D: usize>(
 
     // For 2D tensors swapping dims 0 and 1, use optimized transpose kernel
     if D == 2 && dim1 == 0 && dim2 == 1 {
+        // A row or column vector transposes to itself in memory: [N,1] and
+        // [1,N] are the same N contiguous floats. Share the buffer and swap
+        // the shape — no kernel. This is not just an optimisation: a [4.2M,1]
+        // gradient reaching the tiled kernel needs grid.y = 131,598, and CUDA
+        // caps grid.y at 65,535, so the launch failed with INVALID_VALUE.
+        if tensor.shape[0] == 1 || tensor.shape[1] == 1 {
+            return CudaTensorPrimitive {
+                data: tensor.data.clone(),
+                shape: new_shape,
+                strides: CudaTensorPrimitive::<D>::compute_strides(&new_shape),
+                device: tensor.device.clone(),
+                offset: tensor.offset,
+            };
+        }
         return transpose_2d_gpu(tensor);
     }
 
@@ -73,6 +87,9 @@ pub fn swap_dims<const D: usize>(
     transpose_generic_gpu(tensor, dim1, dim2, new_shape, new_strides)
 }
 
+/// Hardware limit on `gridDim.y` (and `.z`); `gridDim.x` allows 2^31-1.
+const MAX_GRID_Y: usize = 65_535;
+
 /// Optimized 2D matrix transpose using shared memory tiled kernel.
 fn transpose_2d_gpu<const D: usize>(tensor: &CudaTensorPrimitive<D>) -> CudaTensorPrimitive<D> {
     let device = &tensor.device;
@@ -80,6 +97,16 @@ fn transpose_2d_gpu<const D: usize>(tensor: &CudaTensorPrimitive<D>) -> CudaTens
     let cols = tensor.shape[1];
     let numel = tensor.numel();
 
+    // The tiled kernel maps rows onto grid.y. Past 65,535 * 32 rows the launch
+    // is rejected outright, so hand tall matrices to the 1-D generic permute.
+    if rows.div_ceil(32) > MAX_GRID_Y {
+        let mut new_shape = tensor.shape;
+        let mut new_strides = tensor.strides;
+        new_shape.swap(0, 1);
+        new_strides.swap(0, 1);
+        return transpose_generic_gpu(tensor, 0, 1, new_shape, new_strides);
+    }
+
     // Allocate output
     let output = device
         .stream()
@@ -111,16 +138,18 @@ fn transpose_2d_gpu<const D: usize>(tensor: &CudaTensorPrimitive<D>) -> CudaTens
             .expect("Failed to launch transpose_2d kernel");
     }
 
-    // Swapped shape and strides for 2D
+    // The kernel wrote a fresh buffer that is contiguous in the transposed
+    // shape, so its strides are the contiguous ones for that shape — NOT the
+    // input's strides swapped, which would describe a view of the *original*
+    // buffer. Nothing outside this module reads strides today, so the old
+    // stamp was latent rather than live; keep it honest anyway.
     let mut new_shape = tensor.shape;
-    let mut new_strides = tensor.strides;
     new_shape.swap(0, 1);
-    new_strides.swap(0, 1);
 
     CudaTensorPrimitive {
         data: Arc::new(output),
         shape: new_shape,
-        strides: new_strides,
+        strides: CudaTensorPrimitive::<D>::compute_strides(&new_shape),
         device: device.clone(),
         offset: 0,
     }
@@ -148,13 +177,23 @@ fn transpose_generic_gpu<const D: usize>(
     let mut permuted_input_strides = tensor.strides;
     permuted_input_strides.swap(dim1, dim2);
 
+    // The kernel DECODES each flat output index with `output_strides` — so they
+    // must be the contiguous strides of the output shape, which is what the
+    // kernel writes. `new_strides` (the input's strides swapped) was being
+    // passed here instead; for a contiguous 2-D input that is [1, cols], which
+    // decodes every index to itself and turns the transpose into a plain copy.
+    // Nothing exercised this path with a numeric check until the tall-matrix
+    // fallback did.
+    let _ = new_strides;
+    let out_strides = CudaTensorPrimitive::<D>::compute_strides(&new_shape);
+
     // Upload strides and shape to GPU
     let stream = device.stream();
     let input_strides_gpu = stream
         .clone_htod(&permuted_input_strides[..])
         .expect("Failed to upload input strides");
     let output_strides_gpu = stream
-        .clone_htod(&new_strides[..])
+        .clone_htod(&out_strides[..])
         .expect("Failed to upload output strides");
     let output_shape_gpu = stream
         .clone_htod(&new_shape[..])
@@ -197,7 +236,7 @@ fn transpose_generic_gpu<const D: usize>(
     CudaTensorPrimitive {
         data: Arc::new(output),
         shape: new_shape,
-        strides: new_strides,
+        strides: out_strides,
         device: device.clone(),
         offset: 0,
     }
@@ -225,6 +264,77 @@ mod tests {
         }
     }
 
+    /// grid.y is capped at 65,535 blocks of 32 rows. A column vector past that
+    /// used to fail the launch with CUDA_ERROR_INVALID_VALUE; it is a reshape.
+    #[test]
+    fn transpose_of_a_tall_column_vector_is_a_reshape_not_a_launch() {
+        if let Ok(device) = CudaDevice::new(0) {
+            let n = 4_211_136; // the exact size that failed in a dg-gnn backward pass
+            let host: Vec<f32> = (0..n).map(|i| i as f32).collect();
+            let a = creation::from_data(&host, [n, 1], &device);
+            let at = transpose(&a);
+            assert_eq!(at.shape, [1, n]);
+            assert!(at.is_contiguous());
+            let back = dev_ops::copy_to_host(&at);
+            assert_eq!(back[0], 0.0);
+            assert_eq!(back[n - 1], (n - 1) as f32);
+            assert_eq!(back[123_456], 123_456.0);
+        }
+    }
+
+    /// A tall matrix with more than one column cannot be a reshape and must
+    /// take the 1-D generic path once rows/32 exceeds the grid.y limit.
+    #[test]
+    fn transpose_of_a_tall_matrix_falls_back_past_the_grid_limit() {
+        if let Ok(device) = CudaDevice::new(0) {
+            let rows = MAX_GRID_Y * 32 + 32; // one tile past the cap
+            let cols = 3;
+            let host: Vec<f32> = (0..rows * cols).map(|i| i as f32).collect();
+            let a = creation::from_data(&host, [rows, cols], &device);
+            let at = transpose(&a);
+            assert_eq!(at.shape, [cols, rows]);
+            assert!(at.is_contiguous());
+            let back = dev_ops::copy_to_host(&at);
+            // at[c][r] == a[r][c] == r*cols + c
+            for (r, c) in [(0, 0), (1, 2), (rows - 1, 0), (rows - 1, 2), (1_000_000, 1)] {
+                assert_eq!(back[c * rows + r], (r * cols + c) as f32, "at[{c}][{r}]");
+            }
+        }
+    }
+
+    /// The tiled kernel's output is contiguous in the transposed shape; the
+    /// strides must say so rather than echo the input's strides swapped.
+    #[test]
+    fn tiled_transpose_output_is_stamped_contiguous() {
+        if let Ok(device) = CudaDevice::new(0) {
+            let a = creation::from_data(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3], &device);
+            let at = transpose(&a);
+            assert_eq!(at.shape, [3, 2]);
+            assert_eq!(at.strides, [2, 1]);
+            assert!(at.is_contiguous());
+        }
+    }
+
+    /// The generic permute decoded output indices with the wrong strides and
+    /// degenerated into a copy. A 3-D dim swap goes through it unconditionally.
+    #[test]
+    fn generic_permute_actually_permutes() {
+        if let Ok(device) = CudaDevice::new(0) {
+            // a[i][j][k] = 100*i + 10*j + k, shape [2,3,4]
+            let host: Vec<f32> = (0..2 * 3 * 4)
+                .map(|n| { let (i, j, k) = (n / 12, (n / 4) % 3, n % 4); (100 * i + 10 * j + k) as f32 })
+                .collect();
+            let a = creation::from_data(&host, [2, 3, 4], &device);
+            let b = swap_dims(&a, 0, 2); // b[k][j][i] == a[i][j][k], shape [4,3,2]
+            assert_eq!(b.shape, [4, 3, 2]);
+            assert!(b.is_contiguous());
+            let back = dev_ops::copy_to_host(&b);
+            for i in 0..2 { for j in 0..3 { for k in 0..4 {
+                assert_eq!(back[k * 6 + j * 2 + i], (100 * i + 10 * j + k) as f32, "b[{k}][{j}][{i}]");
+            }}}
+        }
+    }
+
     #[test]
     fn test_reshape() {
         if let Ok(device) = CudaDevice::new(0) {