diff --git a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs index 42f983f74..44e2f325b 100644 --- a/crates/ml-core/src/cuda_autograd/gpu_tensor.rs +++ b/crates/ml-core/src/cuda_autograd/gpu_tensor.rs @@ -229,63 +229,97 @@ impl GpuTensor { "cat: dim {dim} >= ndim {ndim}" ))); } - // Simple host-side implementation - let mut all_data = Vec::new(); + let mut total_dim = 0_usize; for t in tensors { - let host = t.to_host(stream)?; - all_data.push(host); - total_dim = total_dim.checked_add( - t.shape().get(dim).copied().unwrap_or(0) - ).unwrap_or(total_dim); + total_dim += t.shape().get(dim).copied().unwrap_or(0); } - // For dim=0 concat (most common), just flatten - if dim == 0 { - let flat: Vec = all_data.into_iter().flatten().collect(); - let mut new_shape = tensors.first().map_or_else(Vec::new, |t| t.shape().to_vec()); - if let Some(d) = new_shape.get_mut(0) { - *d = total_dim; - } - return Self::from_host(&flat, new_shape, stream); - } - // General dim>0: interleave rows along the concat dimension + let ref_shape = tensors.first().map_or(&[] as &[usize], |t| t.shape()); - // outer = product of dims before `dim`, inner = product of dims after `dim` - let outer: usize = ref_shape[..dim].iter().product(); - let inner: usize = if dim + 1 < ref_shape.len() { ref_shape[dim + 1..].iter().product() } else { 1 }; - let mut flat = Vec::with_capacity(outer * total_dim * inner); - for o in 0..outer { - for (ti, t) in tensors.iter().enumerate() { - let d = t.shape().get(dim).copied().unwrap_or(0); - let src = &all_data[ti]; - let src_stride = d * inner; - let start = o * src_stride; - let end = start + src_stride; - if end <= src.len() { - flat.extend_from_slice(&src[start..end]); + let mut new_shape = ref_shape.to_vec(); + new_shape[dim] = total_dim; + let total_elems: usize = new_shape.iter().product(); + + // GPU-native: allocate output and DtoD copy each tensor's data + let mut output = stream.alloc_zeros::(total_elems).map_err(|e| { + MLError::ModelError(format!("cat: alloc {total_elems} f32: {e}")) + })?; + + if dim == 0 { + // dim=0: tensors are contiguous blocks — sequential DtoD copies + let mut offset = 0_usize; + for t in tensors { + let n = t.numel(); + if n > 0 { + let src_view = t.data.slice(..n); + let mut dst_view = output.slice_mut(offset..offset + n); + stream.memcpy_dtod(&src_view, &mut dst_view).map_err(|e| { + MLError::ModelError(format!("cat dim=0 DtoD: {e}")) + })?; + } + offset += n; + } + } else { + // dim>0: interleaved copy via per-row DtoD + let inner: usize = if dim + 1 < ref_shape.len() { + ref_shape[dim + 1..].iter().product() + } else { + 1 + }; + let outer: usize = ref_shape[..dim].iter().product(); + let dst_stride = total_dim * inner; + + for o in 0..outer { + let mut col_offset = 0_usize; + for t in tensors.iter() { + let d = t.shape().get(dim).copied().unwrap_or(0); + let src_stride = d * inner; + let src_start = o * src_stride; + let dst_start = o * dst_stride + col_offset; + let n = src_stride; + if n > 0 && src_start + n <= t.numel() { + let src_view = t.data.slice(src_start..src_start + n); + let mut dst_view = output.slice_mut(dst_start..dst_start + n); + stream.memcpy_dtod(&src_view, &mut dst_view).map_err(|e| { + MLError::ModelError(format!("cat dim={dim} DtoD: {e}")) + })?; + } + col_offset += src_stride; } } } - let mut new_shape = ref_shape.to_vec(); - new_shape[dim] = total_dim; - Self::from_host(&flat, new_shape, stream) + + Self::new(output, new_shape) } - /// Stack tensors along a new dimension 0. + /// Stack tensors along a new dimension 0 (GPU-native DtoD). pub fn stack(tensors: &[&Self], _dim: usize, stream: &Arc) -> Result { if tensors.is_empty() { return Err(MLError::ModelError("stack: empty tensor list".to_string())); } let inner_shape = tensors.first().map_or(&[][..], |t| t.shape()); let batch = tensors.len(); - let mut flat = Vec::with_capacity(batch * inner_shape.iter().product::()); - for t in tensors { - let host = t.to_host(stream)?; - flat.extend(host); + let elems_per: usize = inner_shape.iter().product(); + let total = batch * elems_per; + + let mut output = stream.alloc_zeros::(total).map_err(|e| { + MLError::ModelError(format!("stack: alloc {total} f32: {e}")) + })?; + + for (i, t) in tensors.iter().enumerate() { + let n = t.numel(); + if n > 0 { + let src_view = t.data.slice(..n); + let mut dst_view = output.slice_mut(i * elems_per..(i + 1) * elems_per); + stream.memcpy_dtod(&src_view, &mut dst_view).map_err(|e| { + MLError::ModelError(format!("stack DtoD [{i}]: {e}")) + })?; + } } + let mut new_shape = vec![batch]; new_shape.extend_from_slice(inner_shape); - Self::from_host(&flat, new_shape, stream) + Self::new(output, new_shape) } // ----------------------------------------------------------------------- diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 354143638..a6f6d3726 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -865,42 +865,16 @@ impl GpuDqnTrainer { } // ── Synchronize stream before readback ────────────────────── - // CUDA Graph replay and direct kernel launches are async. The stream - // must be synchronized before DtoH to ensure kernel writes are visible. - // Also required because graph capture disables event tracking, so - // cudarc's device_ptr() can't rely on write events for ordering. self.stream.synchronize() .map_err(|e| MLError::ModelError(format!("stream sync before readback: {e}")))?; // ── Gather 2 scalars to host: total_loss + grad_norm ──────── - // After CUDA Graph, CudaSlice events are stale (graph capture disables - // event tracking). Use raw_device_ptr (ManuallyDrop guard) to get the - // device pointer without triggering event wait, then raw DtoH. - // Stream is already synchronized above, so this is safe. let mut loss_host = [0.0_f32; 1]; + self.stream.memcpy_dtoh(&self.total_loss_buf, &mut loss_host) + .map_err(|e| MLError::ModelError(format!("DtoH total_loss: {e}")))?; let mut norm_host = [0.0_f32; 1]; - let loss_ptr = raw_device_ptr(&self.total_loss_buf, &self.stream); - let norm_ptr = raw_device_ptr(&self.grad_norm_buf, &self.stream); - let result_loss = unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - loss_host.as_mut_ptr().cast(), - loss_ptr, - std::mem::size_of::(), - ) - }; - if result_loss != cudarc::driver::sys::CUresult::CUDA_SUCCESS { - return Err(MLError::ModelError(format!("DtoH total_loss: {result_loss:?}"))); - } - let result_norm = unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - norm_host.as_mut_ptr().cast(), - norm_ptr, - std::mem::size_of::(), - ) - }; - if result_norm != cudarc::driver::sys::CUresult::CUDA_SUCCESS { - return Err(MLError::ModelError(format!("DtoH grad_norm: {result_norm:?}"))); - } + self.stream.memcpy_dtoh(&self.grad_norm_buf, &mut norm_host) + .map_err(|e| MLError::ModelError(format!("DtoH grad_norm: {e}")))?; self.scalars_readback_host = [loss_host[0], norm_host[0]]; Ok(FusedTrainScalars { @@ -1209,12 +1183,8 @@ impl GpuDqnTrainer { self.stream.synchronize() .map_err(|e| MLError::ModelError(format!("stream sync before capture: {e}")))?; - // Disable event tracking during capture — cudarc's device_ptr/device_ptr_mut - // record CudaEvents which are DISALLOWED inside CUDA Graph capture - // (cause CUDA_ERROR_STREAM_CAPTURE_INVALIDATED). - // SAFETY: all CudaSlices created after disable won't track events. This is - // safe because we're on a single stream with explicit ordering via capture. - unsafe { self.stream.context().disable_event_tracking(); } + // Event tracking is already disabled at DQNTrainer construction (single-stream + // context — events are unnecessary with explicit sync barriers). // Begin stream capture — only work submitted from this thread on this // stream is captured (THREAD_LOCAL mode, safe for single-stream use). @@ -1235,8 +1205,10 @@ impl GpuDqnTrainer { cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, ); - // Re-enable event tracking after capture - unsafe { self.stream.context().enable_event_tracking(); } + // Keep event tracking disabled — graph capture wrote stale events to + // CudaSlices. All subsequent operations use explicit stream sync + + // ManuallyDrop guards (raw_device_ptr), so event tracking is unnecessary. + // Re-enabling would leave stale events that cause CUDA_ERROR_INVALID_VALUE. // Propagate submission error first submit_result?; @@ -1870,8 +1842,6 @@ impl GpuDqnTrainer { ) -> Result<(), MLError> { let sizes = compute_param_sizes(&self.config); - // Synchronize stream: ensures CUDA Graph execution is fully complete - // before EMA kernels modify target weights. self.stream.synchronize() .map_err(|e| MLError::ModelError(format!("EMA pre-sync: {e}")))?; @@ -2214,29 +2184,24 @@ pub(crate) fn query_max_shmem_bytes() -> usize { /// Returns true if CUDA Graph capture should be used for training. /// -/// Disabled on consumer GPUs (< 164KB max shmem opt-in) — CUDA Graph -/// replay triggers CUDA_ERROR_ILLEGAL_ADDRESS in the driver's local -/// memory management on RTX 3050/4090 (confirmed with compute-sanitizer: -/// crash at libcuda.so+0x28d50 during cuGraphLaunch, not in user kernel). -/// This is a driver-level issue with graph replay on Ampere consumer GPUs. +/// CUDA Graph capture disables cudarc event tracking, leaving stale events +/// on CudaSlice buffers. This causes CUDA_ERROR_INVALID_VALUE on subsequent +/// cudarc API calls (memcpy_dtoh, device_ptr, synchronize). Raw driver calls +/// bypass this, but downstream code (EMA kernel, validation) also uses cudarc. /// -/// A100/H100 (>= 164KB shmem) use CUDA Graph for zero-overhead replay. -/// Consumer GPUs fall back to direct kernel launches (slightly higher -/// dispatch latency but correct execution). +/// Enable only on H100/A100 where the full pipeline has been validated. +/// Consumer GPUs use direct kernel launches (slightly higher dispatch overhead +/// but avoids the stale event problem). pub(crate) fn should_use_cuda_graph(_shmem_bytes: usize) -> bool { use cudarc::driver::sys::{cuDeviceGetAttribute, CUdevice_attribute}; let mut max_shmem: i32 = 0; - let result = unsafe { + let _ = unsafe { cuDeviceGetAttribute( &mut max_shmem, CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, 0, ) }; - if result != cudarc::driver::sys::CUresult::CUDA_SUCCESS { - return false; - } - // A100 = 164KB, H100 = 228KB — these support reliable CUDA Graph replay max_shmem >= 164_000 }