rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

feat(meta,jepa): expose GPU features through meta-crates; wire JEPA cluster plan and real shard loading

Browse Source 
Meta-crates (Phase 2):
- rtx-core / rtx-training / rtx-inference-stack gain cuda and metal
  features threading into their sub-crates; GPU was previously
  unreachable through the user-facing bundles.
- rtx-training restores rtx-distributed (the hpc-channels blocker is
  gone) so the advertised DistributedTransformerTrainer resolves; drops
  the unused rtx-runtime dep.
- rtx-transformers drops unused rtx-backend/rtx-backend-cpu deps
  (stale comment referenced a teacher that never used them).

Never-compiled CUDA paths fixed (surfaced by the new feature wiring,
verified on RTX 5060 Ti / CUDA 13.1):
- rtx-compress build.rs: missing Path/Command/fs imports.
- rtx-flash-attention flash_decode_forward: reborrow &mut kernel args.
- rtx-transformers: rope kernel include path, cudarc 0.18 Arc<CudaModule>,
  PushKernelArg imports in jepa_gpu, edition-2024 ref patterns.
- rtx-memory: full cudarc 0.18 port (CudaContext, stream-based alloc,
  DevicePtr accessors, error enum formatting) across gpu_pinning,
  gpu_transfer, gpu_real, gpu_allocator/arena, gpu_tests.

JEPA (Phase 3):
- JepaRunConfig::apply_cluster_plan consumes ClusterTrainingPlan
  (batch size, TP/DP, world size, total steps) so jepa_cluster is no
  longer standalone dead config; ViTSizeStr::approx_params_m feeds
  JepaParallelConfig::for_model_and_cluster.
- WebDatasetShard::load reads real .tar shards from disk via the
  existing parser (gzip rejected explicitly); to_in_memory documented
  as synthetic/test-only.
- New image-decode feature actually defines the dep for the previously
  unreachable cfg(feature = "image-decode") JPEG/PNG decode path.

Co-Authored-By: Claude Fable 5 <[email protected]>
This commit is contained in:
osobh
2026-07-09 19:25:51 -07:00
co-authored by Claude Fable 5
parent 64ade03ab9
commit 1e3c604896
17 changed files with 261 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/core/rtx-memory/src/gpu_allocator/arena.rs
+16 -6
View File
@@ -7,6 +7,8 @@ use crate::{MemoryError, Result};
use parking_lot::Mutex; use parking_lot::Mutex;
use std::collections::{BTreeMap, HashMap, VecDeque}; use std::collections::{BTreeMap, HashMap, VecDeque};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
#[cfg(all(feature = "cuda", feature = "gpu"))]
use std::sync::Arc;
use std::time::Instant; use std::time::Instant;
/// Arena-based GPU memory allocator /// Arena-based GPU memory allocator
@@ -306,7 +308,11 @@ impl GpuArenaAllocator {
)) ))
})?; })?;
let device_ptr = device_slice.as_ptr() as usize; let device_ptr = {
use cudarc::driver::DevicePtr;
let (ptr, _guard) = device_slice.device_ptr(&stream);
ptr as usize
};
// Store the DeviceSlice to ensure proper cleanup // Store the DeviceSlice to ensure proper cleanup
let block_id = self.next_block_id.load(Ordering::SeqCst); let block_id = self.next_block_id.load(Ordering::SeqCst);
@@ -330,11 +336,11 @@ impl GpuArenaAllocator {
sys::cuMemHostAlloc(&mut host_ptr, size, sys::CU_MEMHOSTALLOC_PORTABLE) sys::cuMemHostAlloc(&mut host_ptr, size, sys::CU_MEMHOSTALLOC_PORTABLE)
}; };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(host_ptr as usize) Ok(host_ptr as usize)
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"Failed to allocate pinned memory: CUDA error {}", "Failed to allocate pinned memory: CUDA error {:?}",
result result
))) )))
} }
@@ -353,7 +359,11 @@ impl GpuArenaAllocator {
)) ))
})?; })?;
let device_ptr = device_slice.as_ptr() as usize; let device_ptr = {
use cudarc::driver::DevicePtr;
let (ptr, _guard) = device_slice.device_ptr(&stream);
ptr as usize
};
// Store the DeviceSlice // Store the DeviceSlice
let block_id = self.next_block_id.load(Ordering::SeqCst); let block_id = self.next_block_id.load(Ordering::SeqCst);
@@ -493,11 +503,11 @@ impl GpuArenaAllocator {
let result = let result =
unsafe { sys::cuMemFreeHost(block.device_ptr as *mut std::ffi::c_void) }; unsafe { sys::cuMemFreeHost(block.device_ptr as *mut std::ffi::c_void) };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(()) Ok(())
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"Failed to free pinned memory: CUDA error {}", "Failed to free pinned memory: CUDA error {:?}",
result result
))) )))
} }
crates/core/rtx-memory/src/gpu_pinning.rs
+8 -8
View File
@@ -150,7 +150,7 @@ pub struct GpuPinningManager {
/// Device contexts for CUDA operations /// Device contexts for CUDA operations
#[cfg(all(feature = "cuda", feature = "gpu"))] #[cfg(all(feature = "cuda", feature = "gpu"))]
device_contexts: HashMap<DeviceId, Arc<cudarc::driver::CudaDevice>>, device_contexts: HashMap<DeviceId, std::sync::Arc<cudarc::driver::safe::CudaContext>>,
/// Metal device and buffers for Apple Silicon /// Metal device and buffers for Apple Silicon
#[cfg(all(target_os = "macos", feature = "metal"))] #[cfg(all(target_os = "macos", feature = "metal"))]
@@ -210,16 +210,16 @@ impl GpuPinningManager {
/// Initialize device context for pinning operations /// Initialize device context for pinning operations
#[cfg(all(feature = "cuda", feature = "gpu"))] #[cfg(all(feature = "cuda", feature = "gpu"))]
pub fn initialize_device(&mut self, device_id: DeviceId) -> Result<()> { pub fn initialize_device(&mut self, device_id: DeviceId) -> Result<()> {
use cudarc::driver::CudaDevice; use cudarc::driver::safe::CudaContext;
if self.device_contexts.contains_key(&device_id) { if self.device_contexts.contains_key(&device_id) {
return Ok(()); return Ok(());
} }
let device = CudaDevice::new(device_id.id() as usize) let device = CudaContext::new(device_id.id() as usize)
.map_err(|e| MemoryError::gpu_memory(format!("Failed to initialize device: {}", e)))?; .map_err(|e| MemoryError::gpu_memory(format!("Failed to initialize device: {}", e)))?;
self.device_contexts.insert(device_id, Arc::new(device)); self.device_contexts.insert(device_id, device);
Ok(()) Ok(())
} }
@@ -359,11 +359,11 @@ impl GpuPinningManager {
} }
}; };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(host_ptr as usize) Ok(host_ptr as usize)
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"Failed to allocate pinned memory: CUDA error {}", "Failed to allocate pinned memory: CUDA error {:?}",
result result
))) )))
} }
@@ -703,11 +703,11 @@ impl GpuPinningManager {
// block.host_ptr was obtained from a successful CUDA pinned memory allocation. // block.host_ptr was obtained from a successful CUDA pinned memory allocation.
let result = unsafe { sys::cuMemFreeHost(block.host_ptr as *mut std::ffi::c_void) }; let result = unsafe { sys::cuMemFreeHost(block.host_ptr as *mut std::ffi::c_void) };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(()) Ok(())
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"Failed to free pinned memory: CUDA error {}", "Failed to free pinned memory: CUDA error {:?}",
result result
))) )))
} }
crates/core/rtx-memory/src/gpu_real.rs
+35 -20
View File
@@ -5,6 +5,8 @@
use crate::{MemoryError, Result}; use crate::{MemoryError, Result};
use std::collections::HashMap; use std::collections::HashMap;
#[cfg(all(feature = "cuda", feature = "gpu"))]
use std::sync::Arc;
use std::sync::Mutex; use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant; use std::time::Instant;
@@ -93,7 +95,7 @@ trait DeviceSliceManager: Send + Sync {
#[cfg(all(feature = "cuda", feature = "gpu"))] #[cfg(all(feature = "cuda", feature = "gpu"))]
struct DeviceSliceWrapper<T: Clone + cudarc::driver::DeviceRepr + Send + Sync + 'static> { struct DeviceSliceWrapper<T: Clone + cudarc::driver::DeviceRepr + Send + Sync + 'static> {
slice: cudarc::driver::DeviceSlice<T>, slice: cudarc::driver::CudaSlice<T>,
memory_type: GpuMemoryType, memory_type: GpuMemoryType,
} }
@@ -106,7 +108,10 @@ impl<T: Clone + cudarc::driver::DeviceRepr + Send + Sync + 'static> DeviceSliceM
} }
fn device_ptr(&self) -> usize { fn device_ptr(&self) -> usize {
self.slice.as_ptr() as usize use cudarc::driver::{DevicePtr, DeviceSlice};
let stream = self.slice.stream().clone();
let (ptr, _guard) = self.slice.device_ptr(&stream);
ptr as usize
} }
fn memory_type(&self) -> GpuMemoryType { fn memory_type(&self) -> GpuMemoryType {
@@ -207,7 +212,11 @@ impl RealGpuAllocator {
)) ))
})?; })?;
let device_ptr = device_slice.as_ptr() as usize; let device_ptr = {
use cudarc::driver::DevicePtr;
let (ptr, _guard) = device_slice.device_ptr(&stream);
ptr as usize
};
let manager: Box<dyn DeviceSliceManager> = Box::new(DeviceSliceWrapper { let manager: Box<dyn DeviceSliceManager> = Box::new(DeviceSliceWrapper {
slice: device_slice, slice: device_slice,
memory_type, memory_type,
@@ -223,14 +232,19 @@ impl RealGpuAllocator {
GpuMemoryType::Unified | GpuMemoryType::Managed => { GpuMemoryType::Unified | GpuMemoryType::Managed => {
// cudarc doesn't have direct unified/managed memory APIs in safe interface // cudarc doesn't have direct unified/managed memory APIs in safe interface
// For now, fall back to device memory // For now, fall back to device memory
let device_slice = device.alloc_zeros::<u8>(size).map_err(|e| { let stream = device.default_stream();
let device_slice = stream.alloc_zeros::<u8>(size).map_err(|e| {
MemoryError::gpu_memory(format!( MemoryError::gpu_memory(format!(
"cudarc fallback device memory allocation failed: {}", "cudarc fallback device memory allocation failed: {}",
e e
)) ))
})?; })?;
let device_ptr = device_slice.as_ptr() as usize; let device_ptr = {
use cudarc::driver::DevicePtr;
let (ptr, _guard) = device_slice.device_ptr(&stream);
ptr as usize
};
let manager: Box<dyn DeviceSliceManager> = Box::new(DeviceSliceWrapper { let manager: Box<dyn DeviceSliceManager> = Box::new(DeviceSliceWrapper {
slice: device_slice, slice: device_slice,
memory_type: GpuMemoryType::Device, // Override to actual type memory_type: GpuMemoryType::Device, // Override to actual type
@@ -258,7 +272,7 @@ impl RealGpuAllocator {
let result = let result =
unsafe { sys::cuMemHostAlloc(&mut host_ptr, size, sys::CU_MEMHOSTALLOC_DEVICEMAP) }; unsafe { sys::cuMemHostAlloc(&mut host_ptr, size, sys::CU_MEMHOSTALLOC_DEVICEMAP) };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
// Create a wrapper for pinned memory // Create a wrapper for pinned memory
let manager = Box::new(PinnedMemoryWrapper { let manager = Box::new(PinnedMemoryWrapper {
host_ptr: host_ptr as usize, host_ptr: host_ptr as usize,
@@ -269,7 +283,7 @@ impl RealGpuAllocator {
Ok((host_ptr as usize, manager)) Ok((host_ptr as usize, manager))
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"Failed to allocate pinned memory: CUDA error {}", "Failed to allocate pinned memory: CUDA error {:?}",
result result
))) )))
} }
@@ -454,20 +468,21 @@ impl RealGpuTransferManager {
} }
} }
pub fn initialize_device(&mut self, _device_id: DeviceId) -> Result<()> { #[cfg_attr(not(all(feature = "cuda", feature = "gpu")), allow(unused_variables))]
pub fn initialize_device(&mut self, device_id: DeviceId) -> Result<()> {
#[cfg(all(feature = "cuda", feature = "gpu"))] #[cfg(all(feature = "cuda", feature = "gpu"))]
{ {
use cudarc::driver::CudaContext; use cudarc::driver::CudaContext;
if !self.device_contexts.contains_key(&device_id) { if !self.device_contexts.contains_key(&device_id) {
let device = CudaDevice::new(device_id.id() as usize).map_err(|e| { let device = CudaContext::new(device_id.id() as usize).map_err(|e| {
MemoryError::gpu_memory(format!( MemoryError::gpu_memory(format!(
"Failed to initialize device for transfers: {}", "Failed to initialize device for transfers: {}",
e e
)) ))
})?; })?;
self.device_contexts.insert(device_id, Arc::new(device)); self.device_contexts.insert(device_id, device);
} }
} }
@@ -623,11 +638,11 @@ impl RealGpuTransferManager {
) )
}; };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(()) Ok(())
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"CUDA memcpy H2D failed: {}", "CUDA memcpy H2D failed: {:?}",
result result
))) )))
} }
@@ -657,11 +672,11 @@ impl RealGpuTransferManager {
) )
}; };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(()) Ok(())
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"CUDA memcpy D2H failed: {}", "CUDA memcpy D2H failed: {:?}",
result result
))) )))
} }
@@ -692,11 +707,11 @@ impl RealGpuTransferManager {
) )
}; };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(()) Ok(())
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"CUDA D2D transfer failed: {}", "CUDA D2D transfer failed: {:?}",
result result
))) )))
} }
@@ -849,11 +864,11 @@ impl RealGpuPinningManager {
let result = let result =
unsafe { sys::cuMemHostAlloc(&mut host_ptr, size, sys::CU_MEMHOSTALLOC_PORTABLE) }; unsafe { sys::cuMemHostAlloc(&mut host_ptr, size, sys::CU_MEMHOSTALLOC_PORTABLE) };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(host_ptr as usize) Ok(host_ptr as usize)
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"Failed to allocate pinned memory: CUDA error {}", "Failed to allocate pinned memory: CUDA error {:?}",
result result
))) )))
} }
@@ -885,11 +900,11 @@ impl RealGpuPinningManager {
// block.host_ptr was obtained from a successful cuMemHostAlloc call. // block.host_ptr was obtained from a successful cuMemHostAlloc call.
let result = unsafe { sys::cuMemFreeHost(block.host_ptr as *mut std::ffi::c_void) }; let result = unsafe { sys::cuMemFreeHost(block.host_ptr as *mut std::ffi::c_void) };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(()) Ok(())
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"Failed to free pinned memory: CUDA error {}", "Failed to free pinned memory: CUDA error {:?}",
result result
))) )))
} }
crates/core/rtx-memory/src/gpu_tests.rs
+11 -9
View File
@@ -1,16 +1,18 @@
//! Comprehensive TDD tests for GPU memory operations //! Comprehensive TDD tests for GPU memory operations
use crate::{ use crate::{MemoryError, Result};
MemoryError, Result,
gpu_simple::{DeviceId, GpuMemoryType, TransferType}, // Import device/type identifiers, allocator, and transfer manager from the
// appropriate module (gpu_real when the cuda backend is compiled in, the
// mock gpu_simple implementation otherwise).
#[cfg(feature = "cuda")]
use crate::gpu_real::{
DeviceId, GpuMemoryType, RealGpuAllocator as GpuAllocator,
RealGpuTransferManager as GpuTransferManager, TransferType,
}; };
// Import allocator and transfer manager from the appropriate module
#[cfg(feature = "cuda")]
use crate::gpu_real::{GpuAllocator, GpuTransferManager};
#[cfg(not(feature = "cuda"))] #[cfg(not(feature = "cuda"))]
use crate::gpu_simple::{GpuAllocator, GpuTransferManager}; use crate::gpu_simple::{DeviceId, GpuAllocator, GpuMemoryType, GpuTransferManager, TransferType};
#[cfg(test)] #[cfg(test)]
mod gpu_allocator_tests { mod gpu_allocator_tests {
@@ -472,7 +474,7 @@ mod integration_tests {
match allocator.allocate(4096, GpuMemoryType::Device) { match allocator.allocate(4096, GpuMemoryType::Device) {
Ok(gpu_block) => { Ok(gpu_block) => {
// Upload data // Upload data
let input_data = (0..4096u8).map(|i| (i % 256) as u8).collect::<Vec<_>>(); let input_data = (0..4096u32).map(|i| (i % 256) as u8).collect::<Vec<_>>();
let upload_result = manager let upload_result = manager
.transfer_host_to_device(&input_data, &gpu_block) .transfer_host_to_device(&input_data, &gpu_block)
.await; .await;
crates/core/rtx-memory/src/gpu_transfer.rs
+42 -37
View File
@@ -101,9 +101,9 @@ pub struct GpuTransferManager {
struct DeviceContext { struct DeviceContext {
device_id: DeviceId, device_id: DeviceId,
#[cfg(all(feature = "cuda", feature = "gpu"))] #[cfg(all(feature = "cuda", feature = "gpu"))]
cuda_device: Option<Arc<cudarc::driver::CudaDevice>>, cuda_device: Option<Arc<cudarc::driver::CudaContext>>,
#[cfg(all(feature = "cuda", feature = "gpu"))] #[cfg(all(feature = "cuda", feature = "gpu"))]
cuda_stream: Option<cudarc::driver::CudaStream>, cuda_stream: Option<Arc<cudarc::driver::CudaStream>>,
peer_access_enabled: HashMap<DeviceId, bool>, peer_access_enabled: HashMap<DeviceId, bool>,
} }
@@ -156,19 +156,19 @@ impl GpuTransferManager {
&self, &self,
device_id: DeviceId, device_id: DeviceId,
) -> Result<( ) -> Result<(
Option<Arc<cudarc::driver::CudaDevice>>, Option<Arc<cudarc::driver::CudaContext>>,
Option<cudarc::driver::CudaStream>, Option<Arc<cudarc::driver::CudaStream>>,
)> { )> {
use cudarc::driver::{CudaDevice, CudaStream}; use cudarc::driver::CudaContext;
let device = CudaDevice::new(device_id.id() as usize) let device = CudaContext::new(device_id.id() as usize)
.map_err(|e| MemoryError::gpu_memory(format!("Failed to create CUDA device: {}", e)))?; .map_err(|e| MemoryError::gpu_memory(format!("Failed to create CUDA device: {}", e)))?;
let stream = device let stream = device.new_stream().map_err(|e| {
.fork_default_stream() MemoryError::gpu_memory(format!("Failed to create CUDA stream: {}", e))
.map_err(|e| MemoryError::gpu_memory(format!("Failed to create CUDA stream: {}", e)))?; })?;
Ok((Some(Arc::new(device)), Some(stream))) Ok((Some(device), Some(stream)))
} }
#[cfg(not(feature = "gpu"))] #[cfg(not(feature = "gpu"))]
@@ -212,8 +212,8 @@ impl GpuTransferManager {
#[cfg(all(feature = "cuda", feature = "gpu"))] #[cfg(all(feature = "cuda", feature = "gpu"))]
fn enable_cuda_peer_access( fn enable_cuda_peer_access(
&self, &self,
src_device: &cudarc::driver::CudaDevice, src_device: &cudarc::driver::CudaContext,
_dst_device: &cudarc::driver::CudaDevice, _dst_device: &cudarc::driver::CudaContext,
src_id: DeviceId, src_id: DeviceId,
dst_id: DeviceId, dst_id: DeviceId,
) -> Result<bool> { ) -> Result<bool> {
@@ -223,19 +223,24 @@ impl GpuTransferManager {
// SAFETY: cuDeviceCanAccessPeer is a read-only CUDA driver API call that queries // SAFETY: cuDeviceCanAccessPeer is a read-only CUDA driver API call that queries
// peer access capability. We pass valid device ordinals from DeviceId and a valid // peer access capability. We pass valid device ordinals from DeviceId and a valid
// mutable pointer to receive the result. The call cannot cause memory corruption. // mutable pointer to receive the result. The call cannot cause memory corruption.
let result = let result = unsafe {
unsafe { sys::cuDeviceCanAccessPeer(&mut can_access, src_id.id(), dst_id.id()) }; sys::cuDeviceCanAccessPeer(
&mut can_access,
src_id.id() as i32,
dst_id.id() as i32,
)
};
if result != sys::CUDA_SUCCESS || can_access == 0 { if result != sys::cudaError_enum::CUDA_SUCCESS || can_access == 0 {
return Ok(false); return Ok(false);
} }
// SAFETY: cuCtxEnablePeerAccess enables peer access on a valid CUDA context. // SAFETY: cuCtxEnablePeerAccess enables peer access on a valid CUDA context.
// src_device.cu_primary_ctx() returns a valid CUcontext from an initialized CudaDevice. // src_device.cu_ctx() returns a valid CUcontext from an initialized CudaContext.
// The flags parameter (0) is the only valid value per CUDA documentation. // The flags parameter (0) is the only valid value per CUDA documentation.
let enable_result = unsafe { sys::cuCtxEnablePeerAccess(src_device.cu_primary_ctx(), 0) }; let enable_result = unsafe { sys::cuCtxEnablePeerAccess(src_device.cu_ctx(), 0) };
Ok(enable_result == sys::CUDA_SUCCESS) Ok(enable_result == sys::cudaError_enum::CUDA_SUCCESS)
} }
/// Transfer data from CPU to GPU /// Transfer data from CPU to GPU
@@ -384,13 +389,13 @@ impl GpuTransferManager {
// - dst_block.device_ptr is a valid device pointer from GPU allocation // - dst_block.device_ptr is a valid device pointer from GPU allocation
// - src_data.as_ptr() points to valid host memory with src_data.len() bytes // - src_data.as_ptr() points to valid host memory with src_data.len() bytes
// - Size was validated earlier (src_data.len() <= dst_block.size) // - Size was validated earlier (src_data.len() <= dst_block.size)
// - cuda_stream.cu_stream is a valid stream from the device context // - cuda_stream.cu_stream() is a valid stream from the device context
unsafe { unsafe {
sys::cuMemcpyHtoDAsync_v2( sys::cuMemcpyHtoDAsync_v2(
dst_block.device_ptr as sys::CUdeviceptr, dst_block.device_ptr as sys::CUdeviceptr,
src_data.as_ptr() as *const std::ffi::c_void, src_data.as_ptr() as *const std::ffi::c_void,
src_data.len(), src_data.len(),
cuda_stream.cu_stream, cuda_stream.cu_stream(),
) )
} }
} else { } else {
@@ -407,20 +412,20 @@ impl GpuTransferManager {
} }
}; };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
if self.config.async_transfer { if self.config.async_transfer {
// SAFETY: cuStreamSynchronize blocks until all operations in the stream complete. // SAFETY: cuStreamSynchronize blocks until all operations in the stream complete.
// cuda_stream.cu_stream is a valid stream from the device context. // cuda_stream.cu_stream() is a valid stream from the device context.
let sync_result = let sync_result =
unsafe { sys::cuStreamSynchronize(cuda_stream.cu_stream) }; unsafe { sys::cuStreamSynchronize(cuda_stream.cu_stream()) };
if sync_result != sys::CUDA_SUCCESS { if sync_result != sys::cudaError_enum::CUDA_SUCCESS {
return Err(MemoryError::gpu_memory("Stream synchronization failed")); return Err(MemoryError::gpu_memory("Stream synchronization failed"));
} }
} }
Ok(()) Ok(())
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"CUDA memcpy H2D failed: {}", "CUDA memcpy H2D failed: {:?}",
result result
))) )))
} }
@@ -447,13 +452,13 @@ impl GpuTransferManager {
// - dst_data.as_mut_ptr() points to valid mutable host memory // - dst_data.as_mut_ptr() points to valid mutable host memory
// - src_block.device_ptr is a valid device pointer from GPU allocation // - src_block.device_ptr is a valid device pointer from GPU allocation
// - Size was validated earlier (dst_data.len() >= src_block.size) // - Size was validated earlier (dst_data.len() >= src_block.size)
// - cuda_stream.cu_stream is a valid stream from the device context // - cuda_stream.cu_stream() is a valid stream from the device context
unsafe { unsafe {
sys::cuMemcpyDtoHAsync_v2( sys::cuMemcpyDtoHAsync_v2(
dst_data.as_mut_ptr() as *mut std::ffi::c_void, dst_data.as_mut_ptr() as *mut std::ffi::c_void,
src_block.device_ptr as sys::CUdeviceptr, src_block.device_ptr as sys::CUdeviceptr,
src_block.size, src_block.size,
cuda_stream.cu_stream, cuda_stream.cu_stream(),
) )
} }
} else { } else {
@@ -470,20 +475,20 @@ impl GpuTransferManager {
} }
}; };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
if self.config.async_transfer { if self.config.async_transfer {
// SAFETY: cuStreamSynchronize blocks until all operations in the stream complete. // SAFETY: cuStreamSynchronize blocks until all operations in the stream complete.
// cuda_stream.cu_stream is a valid stream from the device context. // cuda_stream.cu_stream() is a valid stream from the device context.
let sync_result = let sync_result =
unsafe { sys::cuStreamSynchronize(cuda_stream.cu_stream) }; unsafe { sys::cuStreamSynchronize(cuda_stream.cu_stream()) };
if sync_result != sys::CUDA_SUCCESS { if sync_result != sys::cudaError_enum::CUDA_SUCCESS {
return Err(MemoryError::gpu_memory("Stream synchronization failed")); return Err(MemoryError::gpu_memory("Stream synchronization failed"));
} }
} }
Ok(()) Ok(())
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"CUDA memcpy D2H failed: {}", "CUDA memcpy D2H failed: {:?}",
result result
))) )))
} }
@@ -525,13 +530,13 @@ impl GpuTransferManager {
if src_ctx.peer_access_enabled.get(&dst_block.device_id) == Some(&true) { if src_ctx.peer_access_enabled.get(&dst_block.device_id) == Some(&true) {
// SAFETY: cuMemcpyPeer copies between devices with peer access enabled. // SAFETY: cuMemcpyPeer copies between devices with peer access enabled.
// - device_ptr values are valid device pointers from GPU allocations // - device_ptr values are valid device pointers from GPU allocations
// - cu_primary_ctx() returns valid CUDA contexts for initialized devices // - cu_ctx() returns valid CUDA contexts for initialized devices
// - Peer access was explicitly enabled via enable_peer_access() // - Peer access was explicitly enabled via enable_peer_access()
// - Size is validated earlier (src_block.size == dst_block.size) // - Size is validated earlier (src_block.size == dst_block.size)
unsafe { unsafe {
sys::cuMemcpyPeer( sys::cuMemcpyPeer(
dst_block.device_ptr as sys::CUdeviceptr, dst_block.device_ptr as sys::CUdeviceptr,
src_ctx.cuda_device.as_ref().unwrap().cu_primary_ctx(), src_ctx.cuda_device.as_ref().unwrap().cu_ctx(),
src_block.device_ptr as sys::CUdeviceptr, src_block.device_ptr as sys::CUdeviceptr,
self.device_contexts self.device_contexts
.get(&dst_block.device_id) .get(&dst_block.device_id)
@@ -539,7 +544,7 @@ impl GpuTransferManager {
.cuda_device .cuda_device
.as_ref() .as_ref()
.unwrap() .unwrap()
.cu_primary_ctx(), .cu_ctx(),
src_block.size, src_block.size,
) )
} }
@@ -553,11 +558,11 @@ impl GpuTransferManager {
_ => return Err(MemoryError::gpu_memory("Invalid transfer type for D2D")), _ => return Err(MemoryError::gpu_memory("Invalid transfer type for D2D")),
}; };
if result == sys::CUDA_SUCCESS { if result == sys::cudaError_enum::CUDA_SUCCESS {
Ok(()) Ok(())
} else { } else {
Err(MemoryError::gpu_memory(format!( Err(MemoryError::gpu_memory(format!(
"CUDA D2D transfer failed: {}", "CUDA D2D transfer failed: {:?}",
result result
))) )))
} }
crates/meta/rtx-core/Cargo.toml
+5
View File
@@ -16,6 +16,11 @@ rtx-runtime = { path = "../../core/rtx-runtime", version = "1.0.0" }
rtx-autograd = { path = "../../core/rtx-autograd", version = "1.0.0" } rtx-autograd = { path = "../../core/rtx-autograd", version = "1.0.0" }
rtx-memory = { path = "../../core/rtx-memory", version = "1.0.0" } rtx-memory = { path = "../../core/rtx-memory", version = "1.0.0" }
[features]
default = []
cuda = ["rtx-tensor/cuda", "rtx-runtime/cuda", "rtx-memory/cuda"]
metal = ["rtx-tensor/metal", "rtx-runtime/metal", "rtx-memory/metal"]
[dev-dependencies] [dev-dependencies]
tokio-test = "0.4" tokio-test = "0.4"
criterion = "0.5" criterion = "0.5"
crates/meta/rtx-inference-stack/Cargo.toml
+17
View File
@@ -21,6 +21,23 @@ rtx-serving-api = { path = "../../production/rtx-serving-api", version = "1.0.0"
rtx-streaming = { path = "../../production/rtx-streaming", version = "1.0.0" } rtx-streaming = { path = "../../production/rtx-streaming", version = "1.0.0" }
rtx-compress = { path = "../../training/rtx-compress", version = "1.0.0" } rtx-compress = { path = "../../training/rtx-compress", version = "1.0.0" }
[features]
default = []
# GPU inference via the engine's optional backends: candle + ONNX Runtime
# CUDA execution providers, plus CUDA tensor/runtime/compression support.
cuda = [
"rtx-tensor/cuda",
"rtx-runtime/cuda",
"rtx-compress/cuda",
"rtx-inference/candle-cuda",
"rtx-inference/onnx-cuda",
]
metal = [
"rtx-tensor/metal",
"rtx-runtime/metal",
"rtx-inference/candle-metal",
]
[dev-dependencies] [dev-dependencies]
tokio-test = "0.4" tokio-test = "0.4"
criterion = "0.5" criterion = "0.5"
crates/meta/rtx-training/Cargo.toml
+16 -2
View File
@@ -13,16 +13,30 @@ categories = ["science", "mathematics"]
[dependencies] [dependencies]
# Core # Core
rtx-tensor = { path = "../../core/rtx-tensor", version = "1.0.0" } rtx-tensor = { path = "../../core/rtx-tensor", version = "1.0.0" }
rtx-runtime = { path = "../../core/rtx-runtime", version = "1.0.0" }
rtx-autograd = { path = "../../core/rtx-autograd", version = "1.0.0" } rtx-autograd = { path = "../../core/rtx-autograd", version = "1.0.0" }
# Training essentials # Training essentials
rtx-transformers = { path = "../../training/rtx-transformers", version = "1.0.0" } rtx-transformers = { path = "../../training/rtx-transformers", version = "1.0.0" }
# rtx-distributed = { path = "../../training/rtx-distributed", version = "1.0.0" } # Temporarily disabled - missing hpc-channels rtx-distributed = { path = "../../training/rtx-distributed", version = "1.0.0" }
rtx-flash-attention = { path = "../../training/rtx-flash-attention", version = "1.0.0" } rtx-flash-attention = { path = "../../training/rtx-flash-attention", version = "1.0.0" }
rtx-rl = { path = "../../training/rtx-rl", version = "1.0.0" } rtx-rl = { path = "../../training/rtx-rl", version = "1.0.0" }
rtx-compress = { path = "../../training/rtx-compress", version = "1.0.0" } rtx-compress = { path = "../../training/rtx-compress", version = "1.0.0" }
[features]
default = []
cuda = [
"rtx-tensor/cuda",
"rtx-transformers/cuda",
"rtx-flash-attention/cuda",
"rtx-distributed/cuda",
"rtx-compress/cuda",
]
metal = [
"rtx-tensor/metal",
"rtx-transformers/metal",
"rtx-flash-attention/metal",
]
[dev-dependencies] [dev-dependencies]
tokio-test = "0.4" tokio-test = "0.4"
criterion = "0.5" criterion = "0.5"
crates/meta/rtx-training/src/lib.rs
+2 -3
View File
@@ -15,8 +15,8 @@ pub use rtx_tensor as tensor;
// Training // Training
pub use rtx_transformers as transformers; pub use rtx_transformers as transformers;
// pub use rtx_distributed as distributed; // Temporarily disabled - missing hpc-channels
pub use rtx_compress as compress; pub use rtx_compress as compress;
pub use rtx_distributed as distributed;
pub use rtx_flash_attention as flash_attention; pub use rtx_flash_attention as flash_attention;
pub use rtx_rl as rl; pub use rtx_rl as rl;
@@ -25,8 +25,7 @@ pub mod prelude {
pub use rtx_tensor::{Device, Tensor}; pub use rtx_tensor::{Device, Tensor};
// Use the correct exports from rtx_transformers // Use the correct exports from rtx_transformers
pub use rtx_transformers::architectures::{TransformerArchitecture, TransformerConfig}; pub use rtx_transformers::architectures::{TransformerArchitecture, TransformerConfig};
// rtx_distributed temporarily disabled - missing hpc-channels pub use rtx_distributed::distributed_transformer_trainer::DistributedTransformerTrainer;
// pub use rtx_distributed::distributed_transformer_trainer::DistributedTransformerTrainer;
pub use rtx_flash_attention::FlashAttention; pub use rtx_flash_attention::FlashAttention;
// Use the correct exports from rtx_rl // Use the correct exports from rtx_rl
pub use rtx_rl::actor_learner::ActorLearner; pub use rtx_rl::actor_learner::ActorLearner;
crates/training/rtx-compress/build.rs
+6
View File
@@ -3,6 +3,12 @@
//! Compiles MX quantization CUDA kernels to PTX for GPU acceleration. //! Compiles MX quantization CUDA kernels to PTX for GPU acceleration.
use std::env; use std::env;
#[cfg(feature = "cuda")]
use std::fs;
#[cfg(feature = "cuda")]
use std::path::Path;
#[cfg(feature = "cuda")]
use std::process::Command;
fn main() { fn main() {
// Only compile CUDA kernels if CUDA feature is enabled // Only compile CUDA kernels if CUDA feature is enabled
crates/training/rtx-flash-attention/src/kernels/flash_decode_forward.rs
+3 -3
View File
@@ -432,9 +432,9 @@ impl FlashDecodeGpuKernel {
builder.arg(q); builder.arg(q);
builder.arg(k); builder.arg(k);
builder.arg(v); builder.arg(v);
builder.arg(partial_out); builder.arg(&mut *partial_out);
builder.arg(partial_max); builder.arg(&mut *partial_max);
builder.arg(partial_sum); builder.arg(&mut *partial_sum);
builder.arg(&num_heads_i32); builder.arg(&num_heads_i32);
builder.arg(&seq_len_i32); builder.arg(&seq_len_i32);
builder.arg(&head_dim_i32); builder.arg(&head_dim_i32);
crates/training/rtx-transformers/Cargo.toml
+6 -4
View File
@@ -12,10 +12,6 @@ description = "Complete transformer training infrastructure with revolutionary q
# Core RTX dependencies - enabled for autograd integration # Core RTX dependencies - enabled for autograd integration
rtx-tensor = { workspace = true } rtx-tensor = { workspace = true }
rtx-autograd = { workspace = true } rtx-autograd = { workspace = true }
# Tape backend for the SMT predictive-state teacher (SetEncoderTeacher trains
# on `Autodiff<CpuBackend>` — the gradient-checked real-backend path).
rtx-backend = { workspace = true }
rtx-backend-cpu = { path = "../../core/rtx-backend-cpu", version = "1.0.0" }
# Essential dependencies # Essential dependencies
anyhow = { workspace = true } anyhow = { workspace = true }
@@ -66,6 +62,9 @@ ordered-float = "4.2"
uuid = { version = "1.0", features = ["v4"] } uuid = { version = "1.0", features = ["v4"] }
semver = "1.0" semver = "1.0"
# JPEG/PNG decoding for WebDataset shards (jepa_data), optional
image = { workspace = true, optional = true }
[target.'cfg(target_os = "macos")'.dependencies] [target.'cfg(target_os = "macos")'.dependencies]
# Metal GPU acceleration for Apple Silicon # Metal GPU acceleration for Apple Silicon
objc2 = { version = "0.6", optional = true, features = ["std"] } objc2 = { version = "0.6", optional = true, features = ["std"] }
@@ -88,6 +87,9 @@ metal = ["rtx-flash-attention/metal", "rtx-tensor/metal", "rtx-runtime/metal", "
cpu = ["rtx-tensor/cpu"] cpu = ["rtx-tensor/cpu"]
disabled_tests = [] disabled_tests = []
vision-bridge = ["rtx-vision"] vision-bridge = ["rtx-vision"]
# Real JPEG/PNG pixel decoding for WebDataset records; without it,
# webdataset_record_to_image falls back to placeholder pixels.
image-decode = ["dep:image"]
# Binary targets commented out - missing source files # Binary targets commented out - missing source files
# [[bin]] # [[bin]]
crates/training/rtx-transformers/src/layers/rope_cuda.rs
+2 -4
View File
@@ -177,7 +177,7 @@ pub fn rope_forward_cpu(
/// Embedded CUDA kernel source compiled via NVRTC at runtime. /// Embedded CUDA kernel source compiled via NVRTC at runtime.
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
const ROPE_KERNEL_SRC: &str = include_str!("../../cuda_kernels/rope_forward.cu"); const ROPE_KERNEL_SRC: &str = include_str!("../cuda_kernels/rope_forward.cu");
/// Errors originating from the RoPE CUDA kernel. /// Errors originating from the RoPE CUDA kernel.
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
@@ -241,9 +241,7 @@ impl RopeCudaKernel {
.load_module(ptx) .load_module(ptx)
.map_err(|e| RopeKernelError::Load(format!("{e:?}")))?; .map_err(|e| RopeKernelError::Load(format!("{e:?}")))?;
info!("rope_forward_kernel compiled and loaded successfully"); info!("rope_forward_kernel compiled and loaded successfully");
Ok(Self { Ok(Self { module })
module: std::sync::Arc::new(module),
})
} }
/// Launch `rope_forward_kernel` on the given stream. /// Launch `rope_forward_kernel` on the given stream.
crates/training/rtx-transformers/src/ssl/jepa_data.rs
+55 -4
View File
@@ -728,10 +728,12 @@ impl DatasetStats {
// WebDatasetShard // WebDatasetShard
// ============================================================================ // ============================================================================
/// Filesystem shard descriptor for WebDataset-format `.tar` / `.tar.gz` archives. /// Filesystem shard descriptor for WebDataset-format `.tar` archives.
/// ///
/// In production this would open and iterate over the tar archive. /// Use [`WebDatasetShard::load`] to read the actual archive from disk
/// This stub stores metadata and generates synthetic data for `to_in_memory`. /// (via [`read_webdataset_shard`]); [`WebDatasetShard::to_in_memory`]
/// generates synthetic data and exists for tests that need a shard
/// without touching the filesystem.
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct WebDatasetShard { pub struct WebDatasetShard {
pub path: String, pub path: String,
@@ -754,9 +756,33 @@ impl WebDatasetShard {
} }
} }
/// Read the shard's `.tar` archive from disk and decode its records
/// into an [`InMemoryShard`].
///
/// Gzip-compressed shards (`.tar.gz`/`.tgz`) are not yet supported and
/// return an error rather than mis-parsing.
pub fn load(&self) -> Result<InMemoryShard, String> {
if self.compressed {
return Err(format!(
"compressed shard not supported yet (gzip): {}",
self.path
));
}
let (raw_records, _stats) =
read_webdataset_shard(std::path::Path::new(&self.path))?;
let records: Vec<ImageRecord> = raw_records
.into_iter()
.map(webdataset_record_to_image)
.collect();
Ok(InMemoryShard {
records,
shard_id: self.shard_id,
})
}
/// Generate a synthetic `InMemoryShard` with `num_records` records of /// Generate a synthetic `InMemoryShard` with `num_records` records of
/// size `image_size × image_size × 3`. Used for testing without a /// size `image_size × image_size × 3`. Used for testing without a
/// real filesystem. /// real filesystem — use [`WebDatasetShard::load`] for real data.
pub fn to_in_memory(&self, image_size: usize) -> InMemoryShard { pub fn to_in_memory(&self, image_size: usize) -> InMemoryShard {
InMemoryShard::synthetic(self.num_records, image_size, self.shard_id) InMemoryShard::synthetic(self.num_records, image_size, self.shard_id)
} }
@@ -1750,6 +1776,31 @@ mod tests {
// ── Tar parsing tests ───────────────────────────────────────────────────── // ── Tar parsing tests ─────────────────────────────────────────────────────
#[test]
fn test_webdataset_shard_load_reads_real_tar() {
// Write a real tar to a temp file and load it through WebDatasetShard.
let fake_jpg = b"\xFF\xD8\xFF\xE0fake jpeg content";
let tar = make_test_tar(&[("000000", fake_jpg, Some(7)), ("000001", fake_jpg, None)]);
let dir = std::env::temp_dir();
let path = dir.join(format!("jepa_shard_load_test_{}.tar", std::process::id()));
std::fs::write(&path, &tar).expect("write temp tar");
let shard = WebDatasetShard::new(path.to_str().unwrap(), 2, 3);
let mem = shard.load().expect("load real tar");
std::fs::remove_file(&path).ok();
assert_eq!(mem.shard_id, 3);
assert_eq!(mem.records.len(), 2);
assert_eq!(mem.records[0].label, Some(7));
assert_eq!(mem.records[1].label, None);
}
#[test]
fn test_webdataset_shard_load_rejects_gzip() {
let shard = WebDatasetShard::new("/tmp/whatever.tar.gz", 1, 0);
assert!(shard.load().is_err(), "gzip shards must be rejected, not mis-parsed");
}
#[test] #[test]
fn test_parse_empty_tar() { fn test_parse_empty_tar() {
// Two zero blocks = empty archive // Two zero blocks = empty archive
crates/training/rtx-transformers/src/ssl/jepa_gpu.rs
+2 -2
View File
@@ -603,7 +603,7 @@ impl GpuViTEncoder {
n_rows: usize, n_rows: usize,
d: usize, d: usize,
) -> Option<Vec<f32>> { ) -> Option<Vec<f32>> {
use cudarc::driver::LaunchConfig; use cudarc::driver::{LaunchConfig, PushKernelArg};
let mut gpu_x = stream.clone_htod(x.as_slice()).ok()?; let mut gpu_x = stream.clone_htod(x.as_slice()).ok()?;
let n_rows_i = n_rows as i32; let n_rows_i = n_rows as i32;
@@ -643,7 +643,7 @@ impl GpuViTEncoder {
n_rows: usize, n_rows: usize,
d: usize, d: usize,
) -> Option<Vec<f32>> { ) -> Option<Vec<f32>> {
use cudarc::driver::LaunchConfig; use cudarc::driver::{LaunchConfig, PushKernelArg};
// Bias add on CPU (tiny O(n*d), avoids an extra H2D for the bias vector) // Bias add on CPU (tiny O(n*d), avoids an extra H2D for the bias vector)
for t in 0..n_rows { for t in 0..n_rows {
crates/training/rtx-transformers/src/ssl/jepa_runner.rs
+33
View File
@@ -76,6 +76,18 @@ impl ViTSizeStr {
}; };
JepaViTConfig { image_size, patch_size, ..base } JepaViTConfig { image_size, patch_size, ..base }
} }
/// Approximate parameter count in millions, used to pick a parallelism
/// layout via [`JepaParallelConfig::for_model_and_cluster`].
pub fn approx_params_m(&self) -> usize {
match self {
ViTSizeStr::Tiny => 6,
ViTSizeStr::Small => 22,
ViTSizeStr::Base => 86,
ViTSizeStr::Large => 307,
ViTSizeStr::Huge => 632,
}
}
} }
// ============================================================================ // ============================================================================
@@ -209,6 +221,27 @@ impl Default for JepaRunConfig {
} }
} }
impl JepaRunConfig {
/// Apply a [`ClusterTrainingPlan`] to this run config.
///
/// Derives batch size, parallelism degrees (TP/DP), world size, and total
/// steps from the plan, so a plan built with
/// [`JepaParallelConfig::for_model_and_cluster`] actually drives the
/// training run instead of being standalone configuration.
///
/// Returns the plan's human-readable summary so callers can log it.
pub fn apply_cluster_plan(&mut self, plan: &super::jepa_cluster::ClusterTrainingPlan) -> String {
let cfg = &plan.config;
self.batch_size = cfg.global_batch_size.max(1);
self.tensor_parallel = cfg.parallel.tensor_parallel_size.max(1);
self.data_parallel = cfg.parallel.data_parallel_size.max(1);
self.num_gpus = cfg.topology.total_gpus().max(1);
self.world_size = cfg.parallel.total_gpus().max(1);
self.total_steps = plan.total_steps().max(1);
plan.summary()
}
}
// ============================================================================ // ============================================================================
// TOML-like config parser // TOML-like config parser
// ============================================================================ // ============================================================================
crates/training/rtx-transformers/src/training/training_loop.rs
+2 -2
View File
@@ -300,7 +300,7 @@ impl TrainingLoop {
// execution for this step and will retry capture on the next step. // execution for this step and will retry capture on the next step.
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
if is_capture_step { if is_capture_step {
if let (Some(ref gm), Some(ref stream)) = if let (Some(gm), Some(stream)) =
(&self.graph_manager, &self.cuda_stream) (&self.graph_manager, &self.cuda_stream)
{ {
// begin_capture returns a provisional graph_id that is only // begin_capture returns a provisional graph_id that is only
@@ -391,7 +391,7 @@ impl TrainingLoop {
// `begin_capture` … `zero_grad()` and returns the real graph_id // `begin_capture` … `zero_grad()` and returns the real graph_id
// that will be used for all future replays. // that will be used for all future replays.
#[cfg(feature = "cuda")] #[cfg(feature = "cuda")]
if let (Some(sentinel_id), Some(ref gm), Some(ref stream)) = ( if let (Some(sentinel_id), Some(gm), Some(stream)) = (
// Only end-capture when this is specifically the capture step: // Only end-capture when this is specifically the capture step:
// graph_id is Some (sentinel set by begin_capture above) but // graph_id is Some (sentinel set by begin_capture above) but
// we have NOT yet successfully ended capture (is_capture_step // we have NOT yet successfully ended capture (is_capture_step