File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change 1919 - uses : dtolnay/rust-toolchain@stable
2020 with :
2121 components : clippy, rustfmt
22+ - uses : Swatinem/rust-cache@v2
23+ with :
24+ workspaces : crates
2225 - name : Check formatting
2326 run : cargo fmt --workspace --check
2427 working-directory : crates
Original file line number Diff line number Diff line change 3535 target : x86_64-unknown-linux-gnu
3636 - os : ubuntu-24.04-arm
3737 target : aarch64-unknown-linux-gnu
38- - os : macos-13
38+ - os : macos-latest
3939 target : x86_64-apple-darwin
4040 - os : macos-latest
4141 target : aarch64-apple-darwin
4747 with :
4848 target : ${{ matrix.target }}
4949 args : --release --out dist
50+ sccache : true
5051 manylinux : auto
5152 before-script-linux : |
5253 yum install -y libcurl-devel openssl-devel || (apt-get update && apt-get install -y libcurl4-openssl-dev libssl-dev pkg-config)
Original file line number Diff line number Diff line change @@ -15,12 +15,13 @@ use crate::GpuError;
1515/// (e.g. to hand to PyTorch via `from_blob`).
1616pub struct DeviceBuffer {
1717 inner : CudaSlice < u8 > ,
18+ stream : Arc < CudaStream > ,
1819}
1920
2021impl DeviceBuffer {
2122 /// Raw device pointer as `u64`, suitable for FFI / PyTorch interop.
2223 pub fn device_ptr ( & self ) -> u64 {
23- let ( ptr, _sync) = self . inner . device_ptr ( ) ;
24+ let ( ptr, _sync) = self . inner . device_ptr ( & self . stream ) ;
2425 ptr as u64
2526 }
2627
@@ -93,7 +94,7 @@ impl GpuCopier {
9394 reason : e. to_string ( ) ,
9495 }
9596 } ) ?;
96- return Ok ( DeviceBuffer { inner : empty } ) ;
97+ return Ok ( DeviceBuffer { inner : empty, stream : Arc :: clone ( & self . stream ) } ) ;
9798 }
9899
99100 // Build a host slice from the raw pointer (no ownership taken).
@@ -111,7 +112,7 @@ impl GpuCopier {
111112 } ) ?;
112113
113114 tracing:: trace!( bytes = len, "queued H2D copy" ) ;
114- Ok ( DeviceBuffer { inner : device_buf } )
115+ Ok ( DeviceBuffer { inner : device_buf, stream : Arc :: clone ( & self . stream ) } )
115116 }
116117
117118 /// Block until all work on this copier's stream has completed.
Original file line number Diff line number Diff line change @@ -15,7 +15,6 @@ use crate::pinned::PinnedBuffer;
1515use crate :: GpuError ;
1616
1717/// A tensor whose data now lives in GPU device memory.
18- #[ derive( Debug ) ]
1918pub struct GpuTensor {
2019 /// Tensor name from the safetensors file.
2120 pub name : String ,
@@ -31,6 +30,18 @@ pub struct GpuTensor {
3130 _device_buf : crate :: copy:: DeviceBuffer ,
3231}
3332
33+ impl std:: fmt:: Debug for GpuTensor {
34+ fn fmt ( & self , f : & mut std:: fmt:: Formatter < ' _ > ) -> std:: fmt:: Result {
35+ f. debug_struct ( "GpuTensor" )
36+ . field ( "name" , & self . name )
37+ . field ( "device_ptr" , & self . device_ptr )
38+ . field ( "len" , & self . len )
39+ . field ( "dtype" , & self . dtype )
40+ . field ( "shape" , & self . shape )
41+ . finish ( )
42+ }
43+ }
44+
3445/// Orchestrates the full streaming pipeline:
3546///
3647/// `Consumer` -> stage in pinned host memory -> async H2D copy -> [`GpuTensor`]
You can’t perform that action at this time.
0 commit comments