zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-28 00:36:32 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

increase cache size

Browse Source
This commit is contained in:
Awni Hannun 2025-08-26 07:49:09 -07:00
parent 4ba4544549
commit c093fa72c8
1 changed files with 13 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
mlx/backend/cuda/device.cpp
View File

@ -29,7 +29,7 @@ void check_cudnn_error(const char* name, cudnnStatus_t err) {
int cuda_graph_cache_size() {
static int cache_size = []() {
return env::get_var("MLX_CUDA_GRAPH_CACHE_SIZE", 100);
return env::get_var("MLX_CUDA_GRAPH_CACHE_SIZE", 400);
}();
return cache_size;
}
@ -41,7 +41,6 @@ bool use_cuda_graphs() {
return use_graphs;
}
} // namespace
Device::Device(int device) : device_(device) {
@ -242,13 +241,7 @@ void CommandEncoder::add_kernel_node(
void** params) {
if (!use_cuda_graphs()) {
CHECK_CUDA_ERROR(cudaLaunchKernel(
func,
grid_dim,
block_dim,
params,
smem_bytes,
stream()
));
func, grid_dim, block_dim, params, smem_bytes, stream()));
return;
}
cudaKernelNodeParams kernel_params = {0};
@ -268,18 +261,17 @@ void CommandEncoder::add_kernel_node(
void** params) {
if (!use_cuda_graphs()) {
CHECK_CUDA_ERROR(cuLaunchKernel(
func,
grid_dim.x,
grid_dim.y,
grid_dim.z,
block_dim.x,
block_dim.y,
block_dim.z,
smem_bytes,
stream(),
params,
nullptr
));
func,
grid_dim.x,
grid_dim.y,
grid_dim.z,
block_dim.x,
block_dim.y,
block_dim.z,
smem_bytes,
stream(),
params,
nullptr));
return;
}