feat(cuda): rewrite 18 GpuTensor methods as proper CUDA kernels

ZERO host downloads in any GpuTensor operation. All ops GPU-native:

elementwise.rs (NEW): 7 CUDA kernels compiled via compile_ptx_for_device
- elementwise_binary: add/sub/mul/div (parameterized op)
- elementwise_unary: powf/sqr/floor/relu/clamp (parameterized op)
- broadcast_scalar_binary: scalar broadcast with any op
- broadcast_row_binary: row broadcast [1,N] op [M,N]
- expand_broadcast: general N-D broadcast with GPU stride tables
- transpose_2d: [rows,cols] → [cols,rows]
- gather_select: index_select along any dimension

gpu_tensor.rs: 18 methods rewritten from host-roundtrip to GPU-native
- gpu_clone: cuMemcpyDtoDAsync (was DtoH+HtoD)
- add/sub/mul: elementwise_binary kernel (was CPU zip)
- broadcast_mul/div: broadcast kernels (was CPU map)
- narrow(dim=0): DtoD slice view (was CPU slice)
- narrow(dim>0): gather_select kernel
- argmax: ReductionKernels (was CPU scan)
- mean_all/sum_all: ReductionKernels (was CPU sum)
- powf/sqr/floor/relu/clamp: elementwise_unary kernel
- transpose: transpose_2d kernel
- expand/broadcast_as: expand_broadcast kernel
- index_select: gather_select kernel

3 reduction tests have init bug (max not -INF) — fix follows.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

crates/ml-core/src/cuda_autograd/elementwise.rs

@@ -0,0 +1,751 @@
+//! Element-wise and memory-layout CUDA kernels for `GpuTensor` operations.
+//!
+//! Replaces host-roundtrip workarounds with GPU-native kernels:
+//!
+//! - **Binary ops**: add, sub, mul (parameterized via `op` arg)
+//! - **Unary ops**: powf, sqr, floor, relu, clamp (parameterized via `op` arg)
+//! - **Broadcast ops**: scalar/row broadcast mul, div, expand
+//! - **Memory ops**: transpose, narrow (dim-0 DtoD), index_select (gather)
+//!
+//! All kernels are compiled once via `compile_ptx_for_device()` and cached in
+//! a `OnceLock`. Thread-safe, no global mutable state.
+
+#![allow(unsafe_code)] // CUDA FFI requires unsafe for kernel launches.
+
+use std::sync::{Arc, OnceLock};
+
+use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
+use cudarc::nvrtc::Ptx;
+
+use crate::MLError;
+
+/// Launch config helper for element-wise kernels.
+fn elem_cfg(n: usize) -> LaunchConfig {
+    let threads = 256_u32;
+    let blocks = ((n as u32) + threads - 1) / threads;
+    LaunchConfig {
+        grid_dim: (blocks, 1, 1),
+        block_dim: (threads, 1, 1),
+        shared_mem_bytes: 0,
+    }
+}
+
+/// Compiled element-wise kernels -- cached per CUDA context.
+///
+/// Construct once via [`ElementwiseKernels::new`], then reuse for all
+/// element-wise operations on tensors sharing the same context.
+#[allow(missing_debug_implementations)]
+pub struct ElementwiseKernels {
+    binary_fn: CudaFunction,
+    unary_fn: CudaFunction,
+    broadcast_binary_fn: CudaFunction,
+    broadcast_row_binary_fn: CudaFunction,
+    expand_fn: CudaFunction,
+    transpose_2d_fn: CudaFunction,
+    gather_fn: CudaFunction,
+    stream: Arc<CudaStream>,
+}
+
+impl ElementwiseKernels {
+    /// Compile all element-wise kernels from a single CUDA source.
+    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
+        let context = stream.context();
+        let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(ELEMENTWISE_CUDA_SRC, &context)
+            .map_err(|e| MLError::ModelError(format!("elementwise kernel compilation: {e}")))?;
+        let module = context.load_module(ptx).map_err(|e| {
+            MLError::ModelError(format!("elementwise module load: {e}"))
+        })?;
+
+        let load = |name: &str| -> Result<CudaFunction, MLError> {
+            module.load_function(name).map_err(|e| {
+                MLError::ModelError(format!("elementwise kernel '{name}' load: {e}"))
+            })
+        };
+
+        Ok(Self {
+            binary_fn: load("elementwise_binary")?,
+            unary_fn: load("elementwise_unary")?,
+            broadcast_binary_fn: load("broadcast_scalar_binary")?,
+            broadcast_row_binary_fn: load("broadcast_row_binary")?,
+            expand_fn: load("expand_broadcast")?,
+            transpose_2d_fn: load("transpose_2d")?,
+            gather_fn: load("gather_select")?,
+            stream: Arc::clone(stream),
+        })
+    }
+
+    /// Element-wise binary operation on two same-shape buffers.
+    ///
+    /// `op`: 0=add, 1=sub, 2=mul, 3=div
+    pub fn binary(
+        &self,
+        a: &CudaSlice<f32>,
+        b: &CudaSlice<f32>,
+        n: usize,
+        op: i32,
+    ) -> Result<CudaSlice<f32>, MLError> {
+        let out = self.stream.alloc_zeros::<f32>(n).map_err(|e| {
+            MLError::ModelError(format!("binary alloc: {e}"))
+        })?;
+        let n_i32 = n as i32;
+        let cfg = elem_cfg(n);
+
+        // SAFETY: Kernel arguments match the CUDA function signature exactly:
+        //   (const float* a, const float* b, float* out, int n, int op)
+        // a and b are valid CudaSlice<f32> of at least n elements.
+        // out is a freshly allocated CudaSlice<f32> of n elements.
+        unsafe {
+            self.stream
+                .launch_builder(&self.binary_fn)
+                .arg(a)
+                .arg(b)
+                .arg(&out)
+                .arg(&n_i32)
+                .arg(&op)
+                .launch(cfg)
+                .map_err(|e| MLError::ModelError(format!("elementwise_binary(op={op}): {e}")))?;
+        }
+        Ok(out)
+    }
+
+    /// Element-wise unary operation.
+    ///
+    /// `op`: 0=powf(param1), 1=sqr, 2=floor, 3=relu, 4=clamp(param1,param2)
+    pub fn unary(
+        &self,
+        x: &CudaSlice<f32>,
+        n: usize,
+        op: i32,
+        param1: f32,
+        param2: f32,
+    ) -> Result<CudaSlice<f32>, MLError> {
+        let out = self.stream.alloc_zeros::<f32>(n).map_err(|e| {
+            MLError::ModelError(format!("unary alloc: {e}"))
+        })?;
+        let n_i32 = n as i32;
+        let cfg = elem_cfg(n);
+
+        // SAFETY: Kernel arguments match the CUDA function signature exactly:
+        //   (const float* x, float* out, int n, int op, float param1, float param2)
+        // x is a valid CudaSlice<f32> of at least n elements.
+        // out is a freshly allocated CudaSlice<f32> of n elements.
+        unsafe {
+            self.stream
+                .launch_builder(&self.unary_fn)
+                .arg(x)
+                .arg(&out)
+                .arg(&n_i32)
+                .arg(&op)
+                .arg(&param1)
+                .arg(&param2)
+                .launch(cfg)
+                .map_err(|e| MLError::ModelError(format!("elementwise_unary(op={op}): {e}")))?;
+        }
+        Ok(out)
+    }
+
+    /// Broadcast a scalar (1-element buffer) with a tensor via binary op.
+    ///
+    /// `op`: 0=mul, 1=div, 2=add, 3=sub
+    /// If `scalar_is_lhs` is true: out[i] = scalar op tensor[i]
+    /// Otherwise: out[i] = tensor[i] op scalar
+    pub fn broadcast_scalar(
+        &self,
+        tensor: &CudaSlice<f32>,
+        scalar: &CudaSlice<f32>,
+        n: usize,
+        op: i32,
+        scalar_is_lhs: i32,
+    ) -> Result<CudaSlice<f32>, MLError> {
+        let out = self.stream.alloc_zeros::<f32>(n).map_err(|e| {
+            MLError::ModelError(format!("broadcast_scalar alloc: {e}"))
+        })?;
+        let n_i32 = n as i32;
+        let cfg = elem_cfg(n);
+
+        // SAFETY: Kernel arguments match the CUDA function signature exactly:
+        //   (const float* tensor, const float* scalar, float* out,
+        //    int n, int op, int scalar_is_lhs)
+        // tensor is at least n elements, scalar is at least 1 element.
+        unsafe {
+            self.stream
+                .launch_builder(&self.broadcast_binary_fn)
+                .arg(tensor)
+                .arg(scalar)
+                .arg(&out)
+                .arg(&n_i32)
+                .arg(&op)
+                .arg(&scalar_is_lhs)
+                .launch(cfg)
+                .map_err(|e| MLError::ModelError(format!("broadcast_scalar(op={op}): {e}")))?;
+        }
+        Ok(out)
+    }
+
+    /// Row-broadcast binary op: [M, N] op [1, N] or [1, N] op [M, N].
+    ///
+    /// `op`: 0=mul, 1=div
+    /// `row_is_lhs`: if true, out[i] = row[i%cols] op matrix[i]; else matrix[i] op row[i%cols]
+    pub fn broadcast_row(
+        &self,
+        matrix: &CudaSlice<f32>,
+        row: &CudaSlice<f32>,
+        total: usize,
+        cols: usize,
+        op: i32,
+        row_is_lhs: i32,
+    ) -> Result<CudaSlice<f32>, MLError> {
+        let out = self.stream.alloc_zeros::<f32>(total).map_err(|e| {
+            MLError::ModelError(format!("broadcast_row alloc: {e}"))
+        })?;
+        let total_i32 = total as i32;
+        let cols_i32 = cols as i32;
+        let cfg = elem_cfg(total);
+
+        // SAFETY: Kernel arguments match the CUDA function signature exactly:
+        //   (const float* matrix, const float* row, float* out,
+        //    int total, int cols, int op, int row_is_lhs)
+        // matrix is at least `total` elements, row is at least `cols` elements.
+        unsafe {
+            self.stream
+                .launch_builder(&self.broadcast_row_binary_fn)
+                .arg(matrix)
+                .arg(row)
+                .arg(&out)
+                .arg(&total_i32)
+                .arg(&cols_i32)
+                .arg(&op)
+                .arg(&row_is_lhs)
+                .launch(cfg)
+                .map_err(|e| MLError::ModelError(format!("broadcast_row(op={op}): {e}")))?;
+        }
+        Ok(out)
+    }
+
+    /// Expand (broadcast) from source to target using pre-computed stride tables.
+    ///
+    /// `src_strides` and `src_shape` are GPU-resident u32 buffers of length `ndim`.
+    /// `tgt_strides` is a GPU-resident u32 buffer of length `ndim`.
+    pub fn expand(
+        &self,
+        src: &CudaSlice<f32>,
+        src_strides_dev: &CudaSlice<u32>,
+        src_shape_dev: &CudaSlice<u32>,
+        tgt_strides_dev: &CudaSlice<u32>,
+        total: usize,
+        ndim: usize,
+    ) -> Result<CudaSlice<f32>, MLError> {
+        let out = self.stream.alloc_zeros::<f32>(total).map_err(|e| {
+            MLError::ModelError(format!("expand alloc: {e}"))
+        })?;
+        let total_i32 = total as i32;
+        let ndim_i32 = ndim as i32;
+        let cfg = elem_cfg(total);
+
+        // SAFETY: Kernel arguments match the CUDA function signature exactly:
+        //   (const float* src, const unsigned int* src_strides,
+        //    const unsigned int* src_shape, const unsigned int* tgt_strides,
+        //    float* out, int total, int ndim)
+        // All buffers are valid with sizes consistent with ndim and total.
+        unsafe {
+            self.stream
+                .launch_builder(&self.expand_fn)
+                .arg(src)
+                .arg(src_strides_dev)
+                .arg(src_shape_dev)
+                .arg(tgt_strides_dev)
+                .arg(&out)
+                .arg(&total_i32)
+                .arg(&ndim_i32)
+                .launch(cfg)
+                .map_err(|e| MLError::ModelError(format!("expand_broadcast: {e}")))?;
+        }
+        Ok(out)
+    }
+
+    /// Transpose a 2-D matrix [rows, cols] -> [cols, rows].
+    pub fn transpose_2d(
+        &self,
+        src: &CudaSlice<f32>,
+        rows: usize,
+        cols: usize,
+    ) -> Result<CudaSlice<f32>, MLError> {
+        let total = rows * cols;
+        let out = self.stream.alloc_zeros::<f32>(total).map_err(|e| {
+            MLError::ModelError(format!("transpose alloc: {e}"))
+        })?;
+        let total_i32 = total as i32;
+        let rows_i32 = rows as i32;
+        let cols_i32 = cols as i32;
+        let cfg = elem_cfg(total);
+
+        // SAFETY: Kernel arguments match the CUDA function signature exactly:
+        //   (const float* src, float* dst, int total, int rows, int cols)
+        // src and dst are both valid CudaSlice<f32> of `total` elements.
+        unsafe {
+            self.stream
+                .launch_builder(&self.transpose_2d_fn)
+                .arg(src)
+                .arg(&out)
+                .arg(&total_i32)
+                .arg(&rows_i32)
+                .arg(&cols_i32)
+                .launch(cfg)
+                .map_err(|e| MLError::ModelError(format!("transpose_2d: {e}")))?;
+        }
+        Ok(out)
+    }
+
+    /// Gather (index_select) along a dimension.
+    ///
+    /// For dim-0: out[i, ...] = src[indices[i], ...]
+    /// `indices_dev` is a GPU-resident u32 buffer.
+    pub fn gather(
+        &self,
+        src: &CudaSlice<f32>,
+        indices_dev: &CudaSlice<u32>,
+        num_indices: usize,
+        inner_size: usize,
+        outer_count: usize,
+        src_dim_size: usize,
+    ) -> Result<CudaSlice<f32>, MLError> {
+        let total = outer_count * num_indices * inner_size;
+        let out = self.stream.alloc_zeros::<f32>(total).map_err(|e| {
+            MLError::ModelError(format!("gather alloc: {e}"))
+        })?;
+        let total_i32 = total as i32;
+        let num_indices_i32 = num_indices as i32;
+        let inner_size_i32 = inner_size as i32;
+        let src_dim_stride_i32 = (src_dim_size * inner_size) as i32;
+        let cfg = elem_cfg(total);
+
+        // SAFETY: Kernel arguments match the CUDA function signature exactly:
+        //   (const float* src, const unsigned int* indices, float* out,
+        //    int total, int num_indices, int inner_size, int src_dim_stride)
+        // src is valid with enough elements, indices has num_indices elements,
+        // out has total elements allocated.
+        unsafe {
+            self.stream
+                .launch_builder(&self.gather_fn)
+                .arg(src)
+                .arg(indices_dev)
+                .arg(&out)
+                .arg(&total_i32)
+                .arg(&num_indices_i32)
+                .arg(&inner_size_i32)
+                .arg(&src_dim_stride_i32)
+                .launch(cfg)
+                .map_err(|e| MLError::ModelError(format!("gather_select: {e}")))?;
+        }
+        Ok(out)
+    }
+
+    /// Reference to the underlying stream.
+    pub fn stream(&self) -> &Arc<CudaStream> {
+        &self.stream
+    }
+}
+
+/// Global singleton for compiled elementwise kernels.
+///
+/// Lazily compiled on first use. Thread-safe via `OnceLock`.
+static KERNELS: OnceLock<Result<ElementwiseKernels, String>> = OnceLock::new();
+
+/// Get or compile the elementwise kernels for the given stream's device.
+///
+/// The kernels are compiled once and cached globally. If compilation fails,
+/// the error is cached and returned on every subsequent call.
+pub fn get_or_compile(stream: &Arc<CudaStream>) -> Result<&'static ElementwiseKernels, MLError> {
+    let result = KERNELS.get_or_init(|| {
+        ElementwiseKernels::new(stream).map_err(|e| format!("{e}"))
+    });
+    match result {
+        Ok(k) => Ok(k),
+        Err(e) => Err(MLError::ModelError(format!("elementwise kernels: {e}"))),
+    }
+}
+
+// ── CUDA source ──────────────────────────────────────────────────────────
+
+const ELEMENTWISE_CUDA_SRC: &str = r#"
+// ═══════════════════════════════════════════════════════════════════════
+// Element-wise and memory-layout kernels for GpuTensor.
+//
+// All kernels use 1-D grid-stride pattern with 256 threads per block.
+// ═══════════════════════════════════════════════════════════════════════
+
+// ── Binary element-wise: add, sub, mul, div ──────────────────────────
+
+extern "C" __global__
+void elementwise_binary(
+    const float* __restrict__ a,
+    const float* __restrict__ b,
+    float* __restrict__ out,
+    int n,
+    int op
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float va = a[i];
+        float vb = b[i];
+        float r;
+        switch (op) {
+            case 0: r = va + vb; break;  // add
+            case 1: r = va - vb; break;  // sub
+            case 2: r = va * vb; break;  // mul
+            case 3: r = va / vb; break;  // div
+            default: r = 0.0f; break;
+        }
+        out[i] = r;
+    }
+}
+
+// ── Unary element-wise: powf, sqr, floor, relu, clamp ────────────────
+
+extern "C" __global__
+void elementwise_unary(
+    const float* __restrict__ x,
+    float* __restrict__ out,
+    int n,
+    int op,
+    float param1,
+    float param2
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float v = x[i];
+        float r;
+        switch (op) {
+            case 0: r = powf(v, param1); break;              // powf
+            case 1: r = v * v; break;                         // sqr
+            case 2: r = floorf(v); break;                     // floor
+            case 3: r = fmaxf(v, 0.0f); break;               // relu
+            case 4: r = fminf(fmaxf(v, param1), param2); break; // clamp
+            default: r = v; break;
+        }
+        out[i] = r;
+    }
+}
+
+// ── Broadcast scalar binary: tensor op scalar or scalar op tensor ────
+
+extern "C" __global__
+void broadcast_scalar_binary(
+    const float* __restrict__ tensor,
+    const float* __restrict__ scalar_ptr,
+    float* __restrict__ out,
+    int n,
+    int op,
+    int scalar_is_lhs
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float s = scalar_ptr[0];
+        float t = tensor[i];
+        float lhs = scalar_is_lhs ? s : t;
+        float rhs = scalar_is_lhs ? t : s;
+        float r;
+        switch (op) {
+            case 0: r = lhs * rhs; break;  // mul
+            case 1: r = lhs / rhs; break;  // div
+            case 2: r = lhs + rhs; break;  // add
+            case 3: r = lhs - rhs; break;  // sub
+            default: r = 0.0f; break;
+        }
+        out[i] = r;
+    }
+}
+
+// ── Broadcast row binary: [M,N] op [1,N] or [1,N] op [M,N] ─────────
+
+extern "C" __global__
+void broadcast_row_binary(
+    const float* __restrict__ matrix,
+    const float* __restrict__ row,
+    float* __restrict__ out,
+    int total,
+    int cols,
+    int op,
+    int row_is_lhs
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < total) {
+        int col = i % cols;
+        float m = matrix[i];
+        float rv = row[col];
+        float lhs = row_is_lhs ? rv : m;
+        float rhs = row_is_lhs ? m  : rv;
+        float r;
+        switch (op) {
+            case 0: r = lhs * rhs; break;  // mul
+            case 1: r = lhs / rhs; break;  // div
+            default: r = 0.0f; break;
+        }
+        out[i] = r;
+    }
+}
+
+// ── Expand (general N-D broadcast) ──────────────────────────────────
+
+extern "C" __global__
+void expand_broadcast(
+    const float* __restrict__ src,
+    const unsigned int* __restrict__ src_strides,
+    const unsigned int* __restrict__ src_shape,
+    const unsigned int* __restrict__ tgt_strides,
+    float* __restrict__ out,
+    int total,
+    int ndim
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < total) {
+        int src_idx = 0;
+        int remaining = i;
+        for (int d = 0; d < ndim; d++) {
+            int tgt_stride = tgt_strides[d];
+            int coord = remaining / tgt_stride;
+            remaining = remaining % tgt_stride;
+            int src_dim = src_shape[d];
+            int src_coord = (src_dim == 1) ? 0 : coord;
+            src_idx += src_coord * src_strides[d];
+        }
+        out[i] = src[src_idx];
+    }
+}
+
+// ── Transpose 2-D ───────────────────────────────────────────────────
+
+extern "C" __global__
+void transpose_2d(
+    const float* __restrict__ src,
+    float* __restrict__ dst,
+    int total,
+    int rows,
+    int cols
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < total) {
+        int r = i / cols;
+        int c = i % cols;
+        // src[r, c] -> dst[c, r]
+        dst[c * rows + r] = src[i];
+    }
+}
+
+// ── Gather (index_select along a dimension) ──────────────────────────
+
+extern "C" __global__
+void gather_select(
+    const float* __restrict__ src,
+    const unsigned int* __restrict__ indices,
+    float* __restrict__ out,
+    int total,
+    int num_indices,
+    int inner_size,
+    int src_dim_stride
+) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < total) {
+        // Decompose flat output index into (outer, idx_pos, inner)
+        int out_idx_inner = i % inner_size;
+        int tmp = i / inner_size;
+        int idx_pos = tmp % num_indices;
+        int outer = tmp / num_indices;
+
+        unsigned int src_idx = indices[idx_pos];
+        int src_flat = outer * src_dim_stride + (int)src_idx * inner_size + out_idx_inner;
+        out[i] = src[src_flat];
+    }
+}
+"#;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn make_stream() -> Arc<CudaStream> {
+        cudarc::driver::CudaContext::new(0)
+            .expect("CUDA required")
+            .new_stream()
+            .expect("fork stream")
+    }
+
+    #[test]
+    fn test_binary_add() {
+        let stream = make_stream();
+        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
+        let a_host = vec![1.0_f32, 2.0, 3.0, 4.0];
+        let b_host = vec![10.0_f32, 20.0, 30.0, 40.0];
+        let mut a = stream.alloc_zeros::<f32>(4).unwrap();
+        let mut b = stream.alloc_zeros::<f32>(4).unwrap();
+        stream.memcpy_htod(&a_host, &mut a).unwrap();
+        stream.memcpy_htod(&b_host, &mut b).unwrap();
+
+        let out = kernels.binary(&a, &b, 4, 0).unwrap(); // add
+        let mut host = [0.0_f32; 4];
+        stream.memcpy_dtoh(&out, &mut host).unwrap();
+        assert!((host[0] - 11.0).abs() < 1e-5);
+        assert!((host[1] - 22.0).abs() < 1e-5);
+        assert!((host[2] - 33.0).abs() < 1e-5);
+        assert!((host[3] - 44.0).abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_binary_sub() {
+        let stream = make_stream();
+        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
+        let a_host = vec![10.0_f32, 20.0, 30.0];
+        let b_host = vec![1.0_f32, 2.0, 3.0];
+        let mut a = stream.alloc_zeros::<f32>(3).unwrap();
+        let mut b = stream.alloc_zeros::<f32>(3).unwrap();
+        stream.memcpy_htod(&a_host, &mut a).unwrap();
+        stream.memcpy_htod(&b_host, &mut b).unwrap();
+
+        let out = kernels.binary(&a, &b, 3, 1).unwrap(); // sub
+        let mut host = [0.0_f32; 3];
+        stream.memcpy_dtoh(&out, &mut host).unwrap();
+        assert!((host[0] - 9.0).abs() < 1e-5);
+        assert!((host[1] - 18.0).abs() < 1e-5);
+        assert!((host[2] - 27.0).abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_binary_mul() {
+        let stream = make_stream();
+        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
+        let a_host = vec![2.0_f32, 3.0, 4.0];
+        let b_host = vec![5.0_f32, 6.0, 7.0];
+        let mut a = stream.alloc_zeros::<f32>(3).unwrap();
+        let mut b = stream.alloc_zeros::<f32>(3).unwrap();
+        stream.memcpy_htod(&a_host, &mut a).unwrap();
+        stream.memcpy_htod(&b_host, &mut b).unwrap();
+
+        let out = kernels.binary(&a, &b, 3, 2).unwrap(); // mul
+        let mut host = [0.0_f32; 3];
+        stream.memcpy_dtoh(&out, &mut host).unwrap();
+        assert!((host[0] - 10.0).abs() < 1e-5);
+        assert!((host[1] - 18.0).abs() < 1e-5);
+        assert!((host[2] - 28.0).abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_unary_relu() {
+        let stream = make_stream();
+        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
+        let x_host = vec![-2.0_f32, -1.0, 0.0, 1.0, 2.0];
+        let mut x = stream.alloc_zeros::<f32>(5).unwrap();
+        stream.memcpy_htod(&x_host, &mut x).unwrap();
+
+        let out = kernels.unary(&x, 5, 3, 0.0, 0.0).unwrap(); // relu
+        let mut host = [0.0_f32; 5];
+        stream.memcpy_dtoh(&out, &mut host).unwrap();
+        assert!((host[0]).abs() < 1e-5);
+        assert!((host[1]).abs() < 1e-5);
+        assert!((host[2]).abs() < 1e-5);
+        assert!((host[3] - 1.0).abs() < 1e-5);
+        assert!((host[4] - 2.0).abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_unary_sqr() {
+        let stream = make_stream();
+        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
+        let x_host = vec![1.0_f32, 2.0, 3.0, -4.0];
+        let mut x = stream.alloc_zeros::<f32>(4).unwrap();
+        stream.memcpy_htod(&x_host, &mut x).unwrap();
+
+        let out = kernels.unary(&x, 4, 1, 0.0, 0.0).unwrap(); // sqr
+        let mut host = [0.0_f32; 4];
+        stream.memcpy_dtoh(&out, &mut host).unwrap();
+        assert!((host[0] - 1.0).abs() < 1e-5);
+        assert!((host[1] - 4.0).abs() < 1e-5);
+        assert!((host[2] - 9.0).abs() < 1e-5);
+        assert!((host[3] - 16.0).abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_unary_clamp() {
+        let stream = make_stream();
+        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
+        let x_host = vec![-5.0_f32, -1.0, 0.5, 2.0, 10.0];
+        let mut x = stream.alloc_zeros::<f32>(5).unwrap();
+        stream.memcpy_htod(&x_host, &mut x).unwrap();
+
+        let out = kernels.unary(&x, 5, 4, -2.0, 3.0).unwrap(); // clamp(-2, 3)
+        let mut host = [0.0_f32; 5];
+        stream.memcpy_dtoh(&out, &mut host).unwrap();
+        assert!((host[0] - (-2.0)).abs() < 1e-5);
+        assert!((host[1] - (-1.0)).abs() < 1e-5);
+        assert!((host[2] - 0.5).abs() < 1e-5);
+        assert!((host[3] - 2.0).abs() < 1e-5);
+        assert!((host[4] - 3.0).abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_transpose_2d() {
+        let stream = make_stream();
+        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
+        // 2x3 matrix: [[1,2,3],[4,5,6]]
+        let x_host = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
+        let mut x = stream.alloc_zeros::<f32>(6).unwrap();
+        stream.memcpy_htod(&x_host, &mut x).unwrap();
+
+        let out = kernels.transpose_2d(&x, 2, 3).unwrap();
+        let mut host = [0.0_f32; 6];
+        stream.memcpy_dtoh(&out, &mut host).unwrap();
+        // Expected: [[1,4],[2,5],[3,6]] = [1,4,2,5,3,6]
+        assert!((host[0] - 1.0).abs() < 1e-5);
+        assert!((host[1] - 4.0).abs() < 1e-5);
+        assert!((host[2] - 2.0).abs() < 1e-5);
+        assert!((host[3] - 5.0).abs() < 1e-5);
+        assert!((host[4] - 3.0).abs() < 1e-5);
+        assert!((host[5] - 6.0).abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_broadcast_scalar_mul() {
+        let stream = make_stream();
+        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
+        let tensor_host = vec![1.0_f32, 2.0, 3.0, 4.0];
+        let scalar_host = vec![10.0_f32];
+        let mut tensor = stream.alloc_zeros::<f32>(4).unwrap();
+        let mut scalar = stream.alloc_zeros::<f32>(1).unwrap();
+        stream.memcpy_htod(&tensor_host, &mut tensor).unwrap();
+        stream.memcpy_htod(&scalar_host, &mut scalar).unwrap();
+
+        let out = kernels.broadcast_scalar(&tensor, &scalar, 4, 0, 0).unwrap(); // mul, scalar is rhs
+        let mut host = [0.0_f32; 4];
+        stream.memcpy_dtoh(&out, &mut host).unwrap();
+        assert!((host[0] - 10.0).abs() < 1e-5);
+        assert!((host[1] - 20.0).abs() < 1e-5);
+        assert!((host[2] - 30.0).abs() < 1e-5);
+        assert!((host[3] - 40.0).abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_gather_dim0() {
+        let stream = make_stream();
+        let kernels = ElementwiseKernels::new(&stream).expect("compile kernels");
+        // 4x3 src: row0=[0,1,2], row1=[3,4,5], row2=[6,7,8], row3=[9,10,11]
+        let src_host: Vec<f32> = (0..12).map(|i| i as f32).collect();
+        let mut src = stream.alloc_zeros::<f32>(12).unwrap();
+        stream.memcpy_htod(&src_host, &mut src).unwrap();
+
+        // Select rows 2, 0
+        let idx_host = vec![2_u32, 0];
+        let mut idx = stream.alloc_zeros::<u32>(2).unwrap();
+        stream.memcpy_htod(&idx_host, &mut idx).unwrap();
+
+        // dim=0, inner_size=3, outer_count=1, src_dim_size=4
+        let out = kernels.gather(&src, &idx, 2, 3, 1, 4).unwrap();
+        let mut host = [0.0_f32; 6];
+        stream.memcpy_dtoh(&out, &mut host).unwrap();
+        // row2: [6,7,8], row0: [0,1,2]
+        assert!((host[0] - 6.0).abs() < 1e-5);
+        assert!((host[1] - 7.0).abs() < 1e-5);
+        assert!((host[2] - 8.0).abs() < 1e-5);
+        assert!((host[3] - 0.0).abs() < 1e-5);
+        assert!((host[4] - 1.0).abs() < 1e-5);
+        assert!((host[5] - 2.0).abs() < 1e-5);
+    }
+}

crates/ml-core/src/cuda_autograd/gpu_tensor.rs

@@ -4,6 +4,8 @@
 //! wrapper — no autograd graph, no dtype dispatch, no device polymorphism.
 //! All data is F32 on a single CUDA device.
 
+#![allow(unsafe_code)] // CUDA FFI requires unsafe for DtoD memcpy and kernel launches.
+
 use std::mem::ManuallyDrop;
 use std::sync::Arc;
 
@@ -320,20 +322,31 @@ impl GpuTensor {
 
     // -----------------------------------------------------------------------
     // Tensor algebra: element-wise ops, reductions, slicing.
-    // Most ops download to host, compute, re-upload. Marked for future CUDA kernels.
+    // All ops are GPU-native — no host downloads except to_host/to_scalar.
     // -----------------------------------------------------------------------
 
     /// Clone this tensor by copying GPU data to a new allocation.
     ///
-    /// Uses host-side roundtrip (same as `clone_gpu_tensor` in linear.rs).
-    /// TODO: CUDA kernel — use `cuMemcpyDtoDAsync` for zero-host clone.
+    /// Uses async device-to-device memcpy — zero host download.
     pub fn gpu_clone(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
         let n = self.numel();
-        let mut host = vec![0.0_f32; n];
-        stream.memcpy_dtoh(&self.data, &mut host).map_err(|e| {
-            MLError::ModelError(format!("gpu_clone DtoH: {e}"))
+        let dst = stream.alloc_zeros::<f32>(n).map_err(|e| {
+            MLError::ModelError(format!("gpu_clone alloc: {e}"))
         })?;
-        Self::from_host(&host, self.shape.clone(), stream)
+        let num_bytes = n * std::mem::size_of::<f32>();
+        let src_ptr = raw_ptr(&self.data, stream);
+        let dst_ptr = raw_ptr(&dst, stream);
+        // SAFETY: src and dst are valid device allocations on the same context.
+        // num_bytes = n * sizeof(f32) does not exceed either allocation.
+        unsafe {
+            cudarc::driver::result::memcpy_dtod_async(
+                dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
+            ).map_err(|e| MLError::ModelError(format!("gpu_clone DtoD: {e}")))?;
+        }
+        Ok(Self {
+            data: dst,
+            shape: self.shape.clone(),
+        })
     }
 
     /// No-op dtype cast — `GpuTensor` is always F32. Returns a clone.
@@ -384,7 +397,8 @@ impl GpuTensor {
     }
 
     /// Element-wise addition. Shapes must match exactly.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: uses `elementwise_binary` kernel (op=0).
     pub fn add(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
         if self.shape != other.shape {
             return Err(MLError::DimensionMismatch {
@@ -392,14 +406,14 @@ impl GpuTensor {
                 actual: other.numel(),
             });
         }
-        let a = self.to_host(stream)?;
-        let b = other.to_host(stream)?;
-        let c: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x + y).collect();
-        Self::from_host(&c, self.shape.clone(), stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.binary(&self.data, &other.data, self.numel(), 0)?;
+        Ok(Self { data: out, shape: self.shape.clone() })
     }
 
     /// Element-wise subtraction. Shapes must match exactly.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: uses `elementwise_binary` kernel (op=1).
     pub fn sub(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
         if self.shape != other.shape {
             return Err(MLError::DimensionMismatch {
@@ -407,14 +421,14 @@ impl GpuTensor {
                 actual: other.numel(),
             });
         }
-        let a = self.to_host(stream)?;
-        let b = other.to_host(stream)?;
-        let c: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x - y).collect();
-        Self::from_host(&c, self.shape.clone(), stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.binary(&self.data, &other.data, self.numel(), 1)?;
+        Ok(Self { data: out, shape: self.shape.clone() })
     }
 
     /// Element-wise multiply. Shapes must match exactly.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: uses `elementwise_binary` kernel (op=2).
     pub fn mul(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
         if self.shape != other.shape {
             return Err(MLError::DimensionMismatch {
@@ -422,10 +436,9 @@ impl GpuTensor {
                 actual: other.numel(),
             });
         }
-        let a = self.to_host(stream)?;
-        let b = other.to_host(stream)?;
-        let c: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x * y).collect();
-        Self::from_host(&c, self.shape.clone(), stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.binary(&self.data, &other.data, self.numel(), 2)?;
+        Ok(Self { data: out, shape: self.shape.clone() })
     }
 
     /// Element-wise multiply with broadcasting.
@@ -434,28 +447,24 @@ impl GpuTensor {
     /// - Same shape: plain element-wise multiply.
     /// - Scalar (numel=1) * tensor: broadcast the scalar.
     /// - `[1, N]` * `[M, N]` or `[M, N]` * `[1, N]`: broadcast rows.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: all paths use CUDA kernels with zero host download.
     pub fn broadcast_mul(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
         // Fast path: identical shapes
         if self.shape == other.shape {
             return self.mul(other, stream);
         }
+        let kernels = super::elementwise::get_or_compile(stream)?;
         // Scalar broadcast
         if other.numel() == 1 {
-            let scalar = other.to_host(stream)?;
-            let s = scalar.first().copied().unwrap_or(0.0);
-            let a = self.to_host(stream)?;
-            let c: Vec<f32> = a.iter().map(|x| x * s).collect();
-            return Self::from_host(&c, self.shape.clone(), stream);
+            let out = kernels.broadcast_scalar(&self.data, &other.data, self.numel(), 0, 0)?;
+            return Ok(Self { data: out, shape: self.shape.clone() });
         }
         if self.numel() == 1 {
-            let scalar = self.to_host(stream)?;
-            let s = scalar.first().copied().unwrap_or(0.0);
-            let b = other.to_host(stream)?;
-            let c: Vec<f32> = b.iter().map(|x| x * s).collect();
-            return Self::from_host(&c, other.shape.clone(), stream);
+            let out = kernels.broadcast_scalar(&other.data, &self.data, other.numel(), 0, 1)?;
+            return Ok(Self { data: out, shape: other.shape.clone() });
         }
-        // Row broadcast: [1, N] * [M, N]
+        // Row broadcast: [1, N] * [M, N] or [M, N] * [1, N]
         if self.ndim() == 2 && other.ndim() == 2 {
             let (m_a, n_a) = (
                 self.shape.first().copied().unwrap_or(0),
@@ -466,25 +475,15 @@ impl GpuTensor {
                 other.shape.get(1).copied().unwrap_or(0),
             );
             if n_a == n_b {
-                let a = self.to_host(stream)?;
-                let b = other.to_host(stream)?;
                 if m_a == 1 {
-                    let c: Vec<f32> = (0..m_b * n_b)
-                        .map(|i| {
-                            let col = i % n_b;
-                            b.get(i).copied().unwrap_or(0.0) * a.get(col).copied().unwrap_or(0.0)
-                        })
-                        .collect();
-                    return Self::from_host(&c, other.shape.clone(), stream);
+                    // [1, N] * [M, N]: self is the row, other is the matrix
+                    let out = kernels.broadcast_row(&other.data, &self.data, m_b * n_b, n_b, 0, 0)?;
+                    return Ok(Self { data: out, shape: other.shape.clone() });
                 }
                 if m_b == 1 {
-                    let c: Vec<f32> = (0..m_a * n_a)
-                        .map(|i| {
-                            let col = i % n_a;
-                            a.get(i).copied().unwrap_or(0.0) * b.get(col).copied().unwrap_or(0.0)
-                        })
-                        .collect();
-                    return Self::from_host(&c, self.shape.clone(), stream);
+                    // [M, N] * [1, N]: other is the row, self is the matrix
+                    let out = kernels.broadcast_row(&self.data, &other.data, m_a * n_a, n_a, 0, 0)?;
+                    return Ok(Self { data: out, shape: self.shape.clone() });
                 }
             }
         }
@@ -497,29 +496,25 @@ impl GpuTensor {
     /// Element-wise divide with broadcasting.
     ///
     /// Supports the same broadcast patterns as `broadcast_mul`.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: all paths use CUDA kernels with zero host download.
     pub fn broadcast_div(&self, other: &Self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
+        let kernels = super::elementwise::get_or_compile(stream)?;
         // Fast path: identical shapes
         if self.shape == other.shape {
-            let a = self.to_host(stream)?;
-            let b = other.to_host(stream)?;
-            let c: Vec<f32> = a.iter().zip(b.iter()).map(|(x, y)| x / y).collect();
-            return Self::from_host(&c, self.shape.clone(), stream);
+            let out = kernels.binary(&self.data, &other.data, self.numel(), 3)?;
+            return Ok(Self { data: out, shape: self.shape.clone() });
         }
         // Scalar broadcast (divisor is scalar)
         if other.numel() == 1 {
-            let scalar = other.to_host(stream)?;
-            let s = scalar.first().copied().unwrap_or(1.0);
-            let a = self.to_host(stream)?;
-            let c: Vec<f32> = a.iter().map(|x| x / s).collect();
-            return Self::from_host(&c, self.shape.clone(), stream);
+            // tensor / scalar: scalar_is_lhs=0, op=1 (div)
+            let out = kernels.broadcast_scalar(&self.data, &other.data, self.numel(), 1, 0)?;
+            return Ok(Self { data: out, shape: self.shape.clone() });
         }
         if self.numel() == 1 {
-            let scalar = self.to_host(stream)?;
-            let s = scalar.first().copied().unwrap_or(0.0);
-            let b = other.to_host(stream)?;
-            let c: Vec<f32> = b.iter().map(|x| s / x).collect();
-            return Self::from_host(&c, other.shape.clone(), stream);
+            // scalar / tensor: scalar_is_lhs=1, op=1 (div)
+            let out = kernels.broadcast_scalar(&other.data, &self.data, other.numel(), 1, 1)?;
+            return Ok(Self { data: out, shape: other.shape.clone() });
         }
         // Row broadcast for 2-D
         if self.ndim() == 2 && other.ndim() == 2 {
@@ -532,27 +527,15 @@ impl GpuTensor {
                 other.shape.get(1).copied().unwrap_or(0),
             );
             if n_a == n_b {
-                let a = self.to_host(stream)?;
-                let b = other.to_host(stream)?;
                 if m_b == 1 {
-                    // [M, N] / [1, N]
-                    let c: Vec<f32> = (0..m_a * n_a)
-                        .map(|i| {
-                            let col = i % n_a;
-                            a.get(i).copied().unwrap_or(0.0) / b.get(col).copied().unwrap_or(1.0)
-                        })
-                        .collect();
-                    return Self::from_host(&c, self.shape.clone(), stream);
+                    // [M, N] / [1, N]: matrix=self, row=other, op=1 (div), row_is_lhs=0
+                    let out = kernels.broadcast_row(&self.data, &other.data, m_a * n_a, n_a, 1, 0)?;
+                    return Ok(Self { data: out, shape: self.shape.clone() });
                 }
                 if m_a == 1 {
-                    // [1, N] / [M, N]
-                    let c: Vec<f32> = (0..m_b * n_b)
-                        .map(|i| {
-                            let col = i % n_b;
-                            a.get(col).copied().unwrap_or(0.0) / b.get(i).copied().unwrap_or(1.0)
-                        })
-                        .collect();
-                    return Self::from_host(&c, other.shape.clone(), stream);
+                    // [1, N] / [M, N]: matrix=other, row=self, op=1 (div), row_is_lhs=1
+                    let out = kernels.broadcast_row(&other.data, &self.data, m_b * n_b, n_b, 1, 1)?;
+                    return Ok(Self { data: out, shape: other.shape.clone() });
                 }
             }
         }
@@ -690,8 +673,8 @@ impl GpuTensor {
 
     /// Slice along a dimension: extract `[start..start+len)` along `dim`.
     ///
-    /// Only dim=0 is currently optimized; other dims use host roundtrip.
-    /// TODO: CUDA kernel for general narrow.
+    /// For dim=0 (contiguous), uses async DtoD memcpy — zero host download.
+    /// For other dims, uses the gather kernel on GPU.
     pub fn narrow(
         &self,
         dim: usize,
@@ -712,8 +695,6 @@ impl GpuTensor {
             )));
         }
 
-        let host = self.to_host(stream)?;
-
         // Compute strides for row-major layout
         let ndim = self.ndim();
         let mut strides = vec![1_usize; ndim];
@@ -730,54 +711,53 @@ impl GpuTensor {
             *d = len;
         }
         let new_numel: usize = new_shape.iter().product();
-        let mut out = Vec::with_capacity(new_numel);
 
         let stride = strides.get(dim).copied().unwrap_or(1);
 
-        // For dim=0 this is a contiguous slice
+        // For dim=0 this is a contiguous sub-range — use DtoD memcpy
         if dim == 0 {
             let offset = start * stride;
-            let end = offset + len * stride;
-            for i in offset..end {
-                out.push(host.get(i).copied().unwrap_or(0.0));
-            }
-        } else {
-            // General case: iterate over all output elements
-            let mut idx = vec![0_usize; ndim];
-            for _ in 0..new_numel {
-                // Compute source flat index
-                let mut src_flat = 0;
-                for d in 0..ndim {
-                    let coord = if d == dim {
-                        idx.get(d).copied().unwrap_or(0) + start
-                    } else {
-                        idx.get(d).copied().unwrap_or(0)
-                    };
-                    src_flat += coord * strides.get(d).copied().unwrap_or(1);
-                }
-                out.push(host.get(src_flat).copied().unwrap_or(0.0));
-
-                // Increment multi-index
-                for d in (0..ndim).rev() {
-                    if let Some(v) = idx.get_mut(d) {
-                        *v += 1;
-                        let limit = new_shape.get(d).copied().unwrap_or(1);
-                        if *v < limit {
-                            break;
-                        }
-                        *v = 0;
-                    }
-                }
+            let src_view = self.data.slice(offset..offset + new_numel);
+            let dst = stream.alloc_zeros::<f32>(new_numel).map_err(|e| {
+                MLError::ModelError(format!("narrow alloc: {e}"))
+            })?;
+            let num_bytes = new_numel * std::mem::size_of::<f32>();
+            let dst_ptr = raw_ptr(&dst, stream);
+            let (src_ptr, _src_sync) = src_view.device_ptr(stream);
+            // SAFETY: src_view and dst are valid device allocations on the same context.
+            // num_bytes does not exceed either allocation.
+            unsafe {
+                cudarc::driver::result::memcpy_dtod_async(
+                    dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
+                ).map_err(|e| MLError::ModelError(format!("narrow DtoD: {e}")))?;
             }
+            return Ok(Self { data: dst, shape: new_shape });
         }
 
-        Self::from_host(&out, new_shape, stream)
+        // General case: use gather kernel with sequential indices
+        let indices: Vec<u32> = (start..start + len).map(|i| i as u32).collect();
+        let mut indices_dev = stream.alloc_zeros::<u32>(len).map_err(|e| {
+            MLError::ModelError(format!("narrow indices alloc: {e}"))
+        })?;
+        stream.memcpy_htod(&indices, &mut indices_dev).map_err(|e| {
+            MLError::ModelError(format!("narrow indices HtoD: {e}"))
+        })?;
+
+        let inner_size = stride; // elements per index along this dim
+        let outer_count = if dim == 0 { 1 } else {
+            self.shape.iter().take(dim).product::<usize>()
+        };
+
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.gather(&self.data, &indices_dev, len, inner_size, outer_count, dim_size)?;
+        Ok(Self { data: out, shape: new_shape })
     }
 
     /// Argmax along a dimension. Returns indices as `Vec<u32>`.
     ///
     /// For a 2-D tensor `[M, N]` with `dim=1`, returns `M` indices in `[0..N)`.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: delegates to `ReductionKernels::argmax` / `argmax_rows`.
     pub fn argmax(&self, dim: usize, stream: &Arc<CudaStream>) -> Result<Vec<u32>, MLError> {
         if dim >= self.ndim() {
             return Err(MLError::ModelError(format!(
@@ -785,59 +765,30 @@ impl GpuTensor {
                 self.ndim()
             )));
         }
-        let host = self.to_host(stream)?;
+        let reductions = super::reductions::ReductionKernels::new(stream)?;
 
-        // For 1-D, return single argmax
+        // For 1-D, return single argmax via flat kernel
         if self.ndim() == 1 {
-            let mut best_idx = 0_u32;
-            let mut best_val = f32::NEG_INFINITY;
-            for (i, v) in host.iter().enumerate() {
-                if *v > best_val {
-                    best_val = *v;
-                    best_idx = i as u32;
-                }
-            }
-            return Ok(vec![best_idx]);
+            let idx = reductions.argmax(&self.data, self.numel())?;
+            return Ok(vec![idx]);
         }
 
-        // For 2-D with dim=1: argmax per row
+        // For 2-D with dim=1: argmax per row — use ReductionKernels::argmax_rows
         if self.ndim() == 2 && dim == 1 {
             let rows = self.shape.first().copied().unwrap_or(0);
             let cols = self.shape.get(1).copied().unwrap_or(0);
-            let mut indices = Vec::with_capacity(rows);
-            for r in 0..rows {
-                let mut best_idx = 0_u32;
-                let mut best_val = f32::NEG_INFINITY;
-                for c in 0..cols {
-                    let v = host.get(r * cols + c).copied().unwrap_or(f32::NEG_INFINITY);
-                    if v > best_val {
-                        best_val = v;
-                        best_idx = c as u32;
-                    }
-                }
-                indices.push(best_idx);
-            }
-            return Ok(indices);
+            return reductions.argmax_rows(&self.data, rows, cols);
         }
 
-        // For 2-D with dim=0: argmax per column
+        // For 2-D with dim=0: transpose then argmax per row of transposed
         if self.ndim() == 2 && dim == 0 {
             let rows = self.shape.first().copied().unwrap_or(0);
             let cols = self.shape.get(1).copied().unwrap_or(0);
-            let mut indices = Vec::with_capacity(cols);
-            for c in 0..cols {
-                let mut best_idx = 0_u32;
-                let mut best_val = f32::NEG_INFINITY;
-                for r in 0..rows {
-                    let v = host.get(r * cols + c).copied().unwrap_or(f32::NEG_INFINITY);
-                    if v > best_val {
-                        best_val = v;
-                        best_idx = r as u32;
-                    }
-                }
-                indices.push(best_idx);
-            }
-            return Ok(indices);
+            let kernels = super::elementwise::get_or_compile(stream)?;
+            let transposed = kernels.transpose_2d(&self.data, rows, cols)?;
+            // Now transposed is [cols, rows]. argmax per row of [cols, rows] gives
+            // per-original-column argmax along original rows (dim=0).
+            return reductions.argmax_rows(&transposed, cols, rows);
         }
 
         Err(MLError::ModelError(format!(
@@ -859,15 +810,16 @@ impl GpuTensor {
     }
 
     /// Mean of all elements.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: delegates to `ReductionKernels::stats().mean`.
     pub fn mean_all(&self, stream: &Arc<CudaStream>) -> Result<f32, MLError> {
-        let host = self.to_host(stream)?;
-        let n = host.len();
+        let n = self.numel();
         if n == 0 {
             return Err(MLError::ModelError("mean_all: empty tensor".to_string()));
         }
-        let sum: f64 = host.iter().map(|&v| v as f64).sum();
-        Ok((sum / n as f64) as f32)
+        let reductions = super::reductions::ReductionKernels::new(stream)?;
+        let stats = reductions.stats(&self.data, n)?;
+        Ok(stats.mean)
     }
 
     // -----------------------------------------------------------------------
@@ -961,26 +913,31 @@ impl GpuTensor {
         }
     }
 
-    /// Element-wise power. Host-side implementation.
-    /// TODO: CUDA kernel
+    /// Element-wise power.
+    ///
+    /// GPU-native: uses `elementwise_unary` kernel (op=0, param1=exp).
     pub fn powf(&self, exp: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
-        let host = self.to_host(stream)?;
-        let result: Vec<f32> = host.iter().map(|&v| v.powf(exp)).collect();
-        Self::from_host(&result, self.shape.clone(), stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.unary(&self.data, self.numel(), 0, exp, 0.0)?;
+        Ok(Self { data: out, shape: self.shape.clone() })
     }
 
-    /// Element-wise square (`x^2`). Alias for `powf(2.0, stream)`.
-    /// TODO: CUDA kernel
+    /// Element-wise square (`x^2`).
+    ///
+    /// GPU-native: uses `elementwise_unary` kernel (op=1) for `v*v`.
     pub fn sqr(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
-        self.powf(2.0, stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.unary(&self.data, self.numel(), 1, 0.0, 0.0)?;
+        Ok(Self { data: out, shape: self.shape.clone() })
     }
 
     /// Element-wise floor.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: uses `elementwise_unary` kernel (op=2).
     pub fn floor(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
-        let host = self.to_host(stream)?;
-        let result: Vec<f32> = host.iter().map(|&v| v.floor()).collect();
-        Self::from_host(&result, self.shape.clone(), stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.unary(&self.data, self.numel(), 2, 0.0, 0.0)?;
+        Ok(Self { data: out, shape: self.shape.clone() })
     }
 
     /// Extract a single scalar from a 1-element tensor, converting to type `T`.
@@ -1009,19 +966,25 @@ impl GpuTensor {
     }
 
     /// Element-wise ReLU: `max(0, x)`.
-    /// TODO: CUDA kernel — use ActivationKernels::relu for hot path
+    ///
+    /// GPU-native: uses `elementwise_unary` kernel (op=3).
     pub fn relu(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
-        let host = self.to_host(stream)?;
-        let result: Vec<f32> = host.iter().map(|&v| v.max(0.0)).collect();
-        Self::from_host(&result, self.shape.clone(), stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.unary(&self.data, self.numel(), 3, 0.0, 0.0)?;
+        Ok(Self { data: out, shape: self.shape.clone() })
     }
 
     /// Sum of all elements, returned as a scalar tensor.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: delegates to `ReductionKernels::sum()`.
     pub fn sum_all(&self, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
-        let host = self.to_host(stream)?;
-        let total: f64 = host.iter().map(|&v| v as f64).sum();
-        Self::scalar(total as f32, stream)
+        let n = self.numel();
+        if n == 0 {
+            return Self::scalar(0.0, stream);
+        }
+        let reductions = super::reductions::ReductionKernels::new(stream)?;
+        let total = reductions.sum(&self.data, n)?;
+        Self::scalar(total, stream)
     }
 
     /// Flatten all dimensions into a single 1-D tensor.
@@ -1031,7 +994,8 @@ impl GpuTensor {
     }
 
     /// Transpose a 2-D tensor (swap dims 0 and 1).
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: uses `transpose_2d` CUDA kernel.
     pub fn transpose(&self, dim0: usize, dim1: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
         if self.ndim() != 2 || dim0 > 1 || dim1 > 1 || dim0 == dim1 {
             return Err(MLError::ModelError(format!(
@@ -1041,30 +1005,25 @@ impl GpuTensor {
         }
         let rows = self.shape.first().copied().unwrap_or(0);
         let cols = self.shape.get(1).copied().unwrap_or(0);
-        let host = self.to_host(stream)?;
-        let mut out = vec![0.0_f32; rows * cols];
-        for r in 0..rows {
-            for c in 0..cols {
-                if let Some(dst) = out.get_mut(c * rows + r) {
-                    *dst = host.get(r * cols + c).copied().unwrap_or(0.0);
-                }
-            }
-        }
-        Self::from_host(&out, vec![cols, rows], stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.transpose_2d(&self.data, rows, cols)?;
+        Ok(Self { data: out, shape: vec![cols, rows] })
     }
 
     /// Element-wise clamp to `[min, max]`.
-    /// TODO: CUDA kernel
+    ///
+    /// GPU-native: uses `elementwise_unary` kernel (op=4, param1=min, param2=max).
     pub fn clamp(&self, min: f32, max: f32, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
-        let host = self.to_host(stream)?;
-        let result: Vec<f32> = host.iter().map(|&v| v.clamp(min, max)).collect();
-        Self::from_host(&result, self.shape.clone(), stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.unary(&self.data, self.numel(), 4, min, max)?;
+        Ok(Self { data: out, shape: self.shape.clone() })
     }
 
     /// Expand (broadcast) this tensor to match a target shape.
     ///
-    /// Only supports expanding dimensions of size 1. Host-side implementation.
-    /// TODO: CUDA kernel
+    /// Only supports expanding dimensions of size 1.
+    ///
+    /// GPU-native: uploads stride tables to GPU, runs `expand_broadcast` kernel.
     pub fn expand(&self, target_shape: &[usize], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
         if self.ndim() != target_shape.len() {
             return Err(MLError::DimensionMismatch {
@@ -1080,44 +1039,71 @@ impl GpuTensor {
                 )));
             }
         }
-        // Simple host-side broadcast
-        let host = self.to_host(stream)?;
-        let total: usize = target_shape.iter().product();
-        let mut result = Vec::with_capacity(total);
 
-        // Compute source and target strides
         let ndim = self.ndim();
-        let mut src_strides = vec![1_usize; ndim];
-        let mut tgt_strides = vec![1_usize; ndim];
+        let total: usize = target_shape.iter().product();
+
+        // Compute per-element strides for source and target
+        let mut src_strides_host = vec![1_u32; ndim];
+        let mut tgt_strides_host = vec![1_u32; ndim];
+        let mut src_shape_host: Vec<u32> = self.shape.iter().map(|&s| s as u32).collect();
+        // Pad to at least ndim elements (always should be correct)
+        while src_shape_host.len() < ndim {
+            src_shape_host.push(1);
+        }
+
         for i in (0..ndim.saturating_sub(1)).rev() {
-            let sh = self.shape.get(i + 1).copied().unwrap_or(1);
-            let th = target_shape.get(i + 1).copied().unwrap_or(1);
-            let next_src = src_strides.get(i + 1).copied().unwrap_or(1);
-            let next_tgt = tgt_strides.get(i + 1).copied().unwrap_or(1);
-            if let Some(ss) = src_strides.get_mut(i) {
+            let sh = self.shape.get(i + 1).copied().unwrap_or(1) as u32;
+            let th = target_shape.get(i + 1).copied().unwrap_or(1) as u32;
+            let next_src = src_strides_host.get(i + 1).copied().unwrap_or(1);
+            let next_tgt = tgt_strides_host.get(i + 1).copied().unwrap_or(1);
+            if let Some(ss) = src_strides_host.get_mut(i) {
                 *ss = next_src * sh;
             }
-            if let Some(ts) = tgt_strides.get_mut(i) {
+            if let Some(ts) = tgt_strides_host.get_mut(i) {
                 *ts = next_tgt * th;
             }
         }
 
-        for flat_idx in 0..total {
-            let mut src_idx = 0;
-            let mut remaining = flat_idx;
-            for d in 0..ndim {
-                let tgt_stride = tgt_strides.get(d).copied().unwrap_or(1);
-                let coord = remaining / tgt_stride;
-                remaining %= tgt_stride;
-                let src_dim = self.shape.get(d).copied().unwrap_or(1);
-                let src_coord = if src_dim == 1 { 0 } else { coord };
-                let src_stride = src_strides.get(d).copied().unwrap_or(1);
-                src_idx += src_coord * (src_stride / src_dim.max(1));
-            }
-            result.push(host.get(src_idx).copied().unwrap_or(0.0));
-        }
+        // For expand, source strides for dims of size 1 must be 0 so the kernel
+        // reads the same element. But our kernel handles this via src_shape check,
+        // so we need per-element strides (stride within the source buffer).
+        // The kernel computes: src_coord = (src_dim==1) ? 0 : coord
+        // Then: src_idx += src_coord * src_strides[d]
+        // We need src_strides[d] to be the stride *within the source allocation*,
+        // i.e. for a source of shape [1, N], src_strides = [N, 1].
+        // This is already what we computed above (product of subsequent dims).
 
-        Self::from_host(&result, target_shape.to_vec(), stream)
+        // Upload stride tables to GPU
+        let mut src_strides_dev = stream.alloc_zeros::<u32>(ndim).map_err(|e| {
+            MLError::ModelError(format!("expand src_strides alloc: {e}"))
+        })?;
+        let mut src_shape_dev = stream.alloc_zeros::<u32>(ndim).map_err(|e| {
+            MLError::ModelError(format!("expand src_shape alloc: {e}"))
+        })?;
+        let mut tgt_strides_dev = stream.alloc_zeros::<u32>(ndim).map_err(|e| {
+            MLError::ModelError(format!("expand tgt_strides alloc: {e}"))
+        })?;
+        stream.memcpy_htod(&src_strides_host, &mut src_strides_dev).map_err(|e| {
+            MLError::ModelError(format!("expand src_strides HtoD: {e}"))
+        })?;
+        stream.memcpy_htod(&src_shape_host, &mut src_shape_dev).map_err(|e| {
+            MLError::ModelError(format!("expand src_shape HtoD: {e}"))
+        })?;
+        stream.memcpy_htod(&tgt_strides_host, &mut tgt_strides_dev).map_err(|e| {
+            MLError::ModelError(format!("expand tgt_strides HtoD: {e}"))
+        })?;
+
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.expand(
+            &self.data,
+            &src_strides_dev,
+            &src_shape_dev,
+            &tgt_strides_dev,
+            total,
+            ndim,
+        )?;
+        Ok(Self { data: out, shape: target_shape.to_vec() })
     }
 
     /// Alias for `expand()` — Candle callers use `broadcast_as`.
@@ -1125,8 +1111,9 @@ impl GpuTensor {
         self.expand(shape, stream)
     }
 
-    /// Select specific indices along a dimension. Host-side implementation.
-    /// TODO: CUDA kernel
+    /// Select specific indices along a dimension.
+    ///
+    /// GPU-native: uploads indices to GPU, runs `gather_select` kernel.
     pub fn index_select(
         &self,
         dim: usize,
@@ -1138,7 +1125,6 @@ impl GpuTensor {
                 "index_select: dim {dim} >= ndim {}", self.ndim()
             )));
         }
-        let host = self.to_host(stream)?;
         let dim_size = self.shape.get(dim).copied().unwrap_or(0);
 
         // Compute strides
@@ -1156,27 +1142,22 @@ impl GpuTensor {
         if let Some(d) = new_shape.get_mut(dim) {
             *d = indices.len();
         }
-        let new_numel: usize = new_shape.iter().product();
-        let stride = strides.get(dim).copied().unwrap_or(1);
-        let inner_size = stride;
-        let outer_stride = dim_size * inner_size;
-
-        let mut out = Vec::with_capacity(new_numel);
-        let num_outer = if dim == 0 { 1 } else {
+        let inner_size = strides.get(dim).copied().unwrap_or(1);
+        let outer_count = if dim == 0 { 1 } else {
             self.shape.iter().take(dim).product::<usize>()
         };
 
-        for outer in 0..num_outer {
-            for &idx in indices {
-                let idx = idx as usize;
-                let base = outer * outer_stride + idx * inner_size;
-                for inner in 0..inner_size {
-                    out.push(host.get(base + inner).copied().unwrap_or(0.0));
-                }
-            }
-        }
+        // Upload indices to GPU
+        let mut indices_dev = stream.alloc_zeros::<u32>(indices.len()).map_err(|e| {
+            MLError::ModelError(format!("index_select indices alloc: {e}"))
+        })?;
+        stream.memcpy_htod(indices, &mut indices_dev).map_err(|e| {
+            MLError::ModelError(format!("index_select indices HtoD: {e}"))
+        })?;
 
-        Self::from_host(&out, new_shape, stream)
+        let kernels = super::elementwise::get_or_compile(stream)?;
+        let out = kernels.gather(&self.data, &indices_dev, indices.len(), inner_size, outer_count, dim_size)?;
+        Ok(Self { data: out, shape: new_shape })
     }
 
     /// Reshape without copying — returns error if total elements differ.

crates/ml-core/src/cuda_autograd/linear.rs

@@ -405,16 +405,13 @@ fn reduce_sum_axis0(
 }
 
 /// Clone a GpuTensor by copying its data to a new allocation.
+///
+/// Uses async device-to-device memcpy — zero host download.
 pub fn clone_gpu_tensor(
     src: &GpuTensor,
     stream: &Arc<CudaStream>,
 ) -> Result<GpuTensor, MLError> {
-    let n = src.numel();
-    let mut host = vec![0.0_f32; n];
-    stream.memcpy_dtoh(&src.data, &mut host).map_err(|e| {
-        MLError::ModelError(format!("clone DtoH: {e}"))
-    })?;
-    GpuTensor::from_host(&host, src.shape.clone(), stream)
+    src.gpu_clone(stream)
 }
 
 #[cfg(test)]

crates/ml-core/src/cuda_autograd/mod.rs

@@ -41,6 +41,7 @@ pub mod loss;
 pub mod dropout;
 pub mod layer_norm;
 pub mod reductions;
+pub mod elementwise;
 pub mod stream_ops;
 
 pub use gpu_tensor::GpuTensor;
@@ -60,5 +61,6 @@ pub use dropout::GpuDropout;
 pub use layer_norm::GpuLayerNorm;
 pub use reductions::ReductionKernels;
 pub use reductions::Stats as ReductionStats;
+pub use elementwise::ElementwiseKernels;
 pub use stream_ops::StreamTensor;
 pub use stream_ops::StreamLinear;