zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Split cuDNN helpers into a separate header (#2491)

Browse Source 
* Add RAII managed CudaGraph class

* Implement forward rms_norm with cuDNN

* Revert back to old rms norm kernel
This commit is contained in:
Cheng
2025-08-20 09:29:28 +09:00
committed by GitHub
parent cea9369610
commit 65d0d40232
8 changed files with 527 additions and 302 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/cuda/device.h
View File

@@ -21,7 +21,7 @@ class CommandEncoder {
struct CaptureContext {
CaptureContext(CommandEncoder& enc);
~CaptureContext();
cudaGraph_t graph;
CudaGraph graph;
CommandEncoder& enc;
bool discard{false};
};
@@ -115,7 +115,7 @@ class CommandEncoder {
Device& device_;
CudaStream stream_;
cudaGraph_t graph_;
CudaGraph graph_;
Worker worker_;
char node_count_{0};
char graph_node_count_{0};