diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ade729df..5e463697 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -29,7 +29,7 @@ permissions: {} # documented in CLAUDE.md). Hoisted to workflow scope so a future bump # touches one literal here instead of one per job. env: - MLX_EXPECTED_COMMIT: "57c66cac7cb3e5b1eb350488a61f1506b40d39f8" + MLX_EXPECTED_COMMIT: "b7c3dd6d27f45b5365b08a840310187dc503f1db" jobs: # ============================================================================ diff --git a/src/distributed/pipeline/stage_executor/mod.rs b/src/distributed/pipeline/stage_executor/mod.rs index f78dc7bb..f4f36c27 100644 --- a/src/distributed/pipeline/stage_executor/mod.rs +++ b/src/distributed/pipeline/stage_executor/mod.rs @@ -168,15 +168,16 @@ impl LoadedStageExecutor { stage.stage_index ); - // Issue #688 (M2 hardening): defense-in-depth CUDA-graph disable for a - // Gemma 4 pipeline stage, mirroring the non-PP backend-seam load sites. + // Issue #688 (M2 hardening), extended to DeepSeek-V2 by issue #824: + // defense-in-depth CUDA-graph disable for a hazard-family pipeline stage, + // mirroring the non-PP backend-seam load sites. // Both the in-process and the remote-process pipeline loaders reach this // common stage-load chokepoint before any stage weights are realised. The // authoritative env write already happens on the main startup thread in // `start_server`; this idempotent call (guarded by `var_os`) only takes // effect if a future pipeline entry ever bypasses that hoist, so the env is // set before the stage's first GPU eval regardless. - crate::loading::maybe_disable_cuda_graphs_for_gemma4_for_path(model_dir); + crate::loading::maybe_disable_cuda_graphs_for_model_for_path(model_dir); let filter = LayerFilter::from_stage(stage); let family = resolve_stage_family(model_dir)?; diff --git a/src/lib/mlx-cpp/CMakeLists.txt b/src/lib/mlx-cpp/CMakeLists.txt index 38ba43d5..2f908071 100644 --- a/src/lib/mlx-cpp/CMakeLists.txt +++ b/src/lib/mlx-cpp/CMakeLists.txt @@ -92,7 +92,7 @@ else() FetchContent_Declare( mlx GIT_REPOSITORY "https://github.com/ml-explore/mlx.git" - GIT_TAG 57c66cac7cb3e5b1eb350488a61f1506b40d39f8) + GIT_TAG b7c3dd6d27f45b5365b08a840310187dc503f1db) # Use FetchContent_Populate + add_subdirectory so we can apply source # overlays before the MLX build system processes the files. diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu new file mode 100644 index 00000000..07463d87 --- /dev/null +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu @@ -0,0 +1,551 @@ +// mlxcel overlay patch (issue #824): CUDA RMSNorm kernel reverted to the +// upstream pre-#3792 implementation (ml-explore/mlx commit 1700b39a). +// +// Upstream MLX #3792 ("Fix CUDA RMSNorm small-row dispatch") and the later +// #3850 rewrite ("RMSNorm forward speed up") both regressed the small-axis +// CUDA RMSNorm for the DeepSeek-V2 MLA `kv_a_layernorm` (axis_size == the +// kv_lora_rank of 512, which with N_READS == 8 for half-precision activations +// lands on the multi-warp block_dim == 64 dispatch config). The result was +// deterministically-wrong RMSNorm output that degenerated DeepSeek-V2 text +// generation (bisected to a5a684d / #3792; verified coherent at 1700b39). The +// upstream fast kernel is not restored here because its small-axis path is the +// broken one; this overlay keeps the MLX pin at the latest commit and only +// swaps this single kernel file back to the last-known-good version. Re-sync or +// drop this overlay once the upstream small-axis RMSNorm bug is fixed (report +// filed as a follow-up). The helpers this file relies on (load_vector, +// BlockBroadcastReduce, dispatch_float_types, contiguous_copy_gpu) are +// unchanged between 1700b39 and the pinned commit, so it compiles as-is. +// +// ---- begin upstream mlx/backend/cuda/rms_norm.cu @ 1700b39a ---- +// Copyright © 2025 Apple Inc. + +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/backend/cuda/reduce/reduce.cuh" +#include "mlx/backend/gpu/copy.h" +#include "mlx/dtype_utils.h" +#include "mlx/fast_primitives.h" + +#include +#include +#include + +namespace mlx::core { + +namespace cu { + +namespace cg = cooperative_groups; + +inline __device__ float2 plus_f2(const float2& a, const float2& b) { + return {a.x + b.x, a.y + b.y}; +} + +// Similar to cub::BlockReduce, but result is broadcasted to every thread. +template +struct BlockBroadcastReduce { + using TempStorage = T[std::max(BLOCK_DIM / WARP_SIZE, 1)]; + + cg::thread_block& block; + TempStorage& temp; + + template + __device__ T Reduce(const T& input, const Op& op, const T& init_value) { + auto warp = cg::tiled_partition(block); + T x = cg::reduce(warp, input, op); + if constexpr (BLOCK_DIM > GROUP_DIM) { + if (warp.thread_rank() == 0) { + temp[warp.meta_group_rank()] = x; + } + block.sync(); + x = warp.thread_rank() < warp.meta_group_size() ? temp[warp.thread_rank()] + : init_value; + return cg::reduce(warp, x, op); + } else { + return x; + } + } + + __device__ T Sum(const T& input) { + return Reduce(input, cg::plus{}, T{}); + } +}; + +template +__global__ void rms_norm_small( + const T* x, + const T* w, + T* out, + float eps, + uint32_t axis_size, + uint32_t n_rows, + int64_t w_stride) { + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + + using BlockReduceT = BlockBroadcastReduce; + __shared__ typename BlockReduceT::TempStorage temp; + + auto row = + (grid.block_rank() * block.dim_threads().y) + block.thread_index().y; + if (row >= n_rows) { + return; + } + x += row * axis_size; + out += row * axis_size; + + // Normalizer. + float normalizer = 0; + auto index = block.thread_index().x; + auto xn = load_vector(x, index, axis_size, T(0)); +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + float t = static_cast(xn[i]); + normalizer += t * t; + } + + normalizer = BlockReduceT{block, temp}.Sum(normalizer); + normalizer = rsqrt(normalizer / axis_size + eps); + + // Outputs. + auto wn = load_vector(w, index, axis_size, w_stride, T(0)); +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + float y = static_cast(xn[i]) * normalizer; + xn[i] = wn[i] * static_cast(y); + } + store_vector(out, index, xn, axis_size); +} + +template +__global__ void rms_norm( + const T* x, + const T* w, + T* out, + float eps, + uint32_t axis_size, + int64_t w_stride) { + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + + using BlockReduceT = BlockBroadcastReduce; + __shared__ typename BlockReduceT::TempStorage temp; + + x += grid.block_rank() * axis_size; + out += grid.block_rank() * axis_size; + + // Normalizer. + float normalizer = 0; + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) { + auto index = r * BLOCK_DIM + block.thread_rank(); + auto xn = load_vector(x, index, axis_size, T(0)); +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + float t = static_cast(xn[i]); + normalizer += t * t; + } + } + normalizer = BlockReduceT{block, temp}.Sum(normalizer); + normalizer = rsqrt(normalizer / axis_size + eps); + + // Outputs. + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) { + auto index = r * BLOCK_DIM + block.thread_rank(); + auto xn = load_vector(x, index, axis_size, T(0)); + auto wn = load_vector(w, index, axis_size, w_stride, T(0)); +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + float y = static_cast(xn[i]) * normalizer; + xn[i] = wn[i] * static_cast(y); + } + store_vector(out, index, xn, axis_size); + } +} + +template < + typename T, + bool HAS_W, + int BLOCK_DIM, + int REDUCE_DIM, + int N_READS = 4> +__global__ void rms_norm_vjp_small( + const T* x, + const T* w, + const T* g, + T* gx, + T* gw, + float eps, + int32_t axis_size, + int32_t n_rows, + int64_t w_stride) { + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + + using BlockReduceF2 = BlockBroadcastReduce; + __shared__ typename BlockReduceF2::TempStorage temp; + + auto row = + (grid.block_rank() * block.dim_threads().y) + block.thread_index().y; + if (row >= n_rows) { + return; + } + + x += row * axis_size; + g += row * axis_size; + gx += row * axis_size; + gw += row * axis_size; + + // Normalizer. + float2 factors = {}; + auto index = block.thread_index().x; + auto xn = load_vector(x, index, axis_size, T(0)); + auto gn = load_vector(g, index, axis_size, T(0)); + auto wn = load_vector(w, index, axis_size, w_stride, T(0)); + for (int i = 0; i < N_READS; i++) { + float t = static_cast(xn[i]); + float wi = wn[i]; + float gi = gn[i]; + float wg = wi * gi; + factors = plus_f2(factors, {wg * t, t * t}); + } + + factors = BlockReduceF2{block, temp}.Reduce(factors, plus_f2, {}); + float meangwx = factors.x / axis_size; + float normalizer = rsqrt(factors.y / axis_size + eps); + float normalizer3 = normalizer * normalizer * normalizer; + + // Outputs. + for (int i = 0; i < N_READS; i++) { + float xi = xn[i]; + float wi = wn[i]; + float gi = gn[i]; + xn[i] = static_cast(normalizer * wi * gi - xi * meangwx * normalizer3); + if constexpr (HAS_W) { + wn[i] = static_cast(gi * xi * normalizer); + } + } + store_vector(gx, index, xn, axis_size); + if constexpr (HAS_W) { + store_vector(gw, index, wn, axis_size); + } +} + +template +__global__ void rms_norm_vjp( + const T* x, + const T* w, + const T* g, + T* gx, + T* gw, + float eps, + int32_t axis_size, + int64_t w_stride) { + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + + using BlockReduceF2 = BlockBroadcastReduce; + __shared__ typename BlockReduceF2::TempStorage temp; + + x += grid.block_rank() * axis_size; + g += grid.block_rank() * axis_size; + gx += grid.block_rank() * axis_size; + gw += grid.block_rank() * axis_size; + + // Normalizer. + float2 factors = {}; + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) { + auto index = r * BLOCK_DIM + block.thread_rank(); + auto xn = load_vector(x, index, axis_size, T(0)); + auto gn = load_vector(g, index, axis_size, T(0)); + auto wn = load_vector(w, index, axis_size, w_stride, T(0)); + for (int i = 0; i < N_READS; i++) { + float t = static_cast(xn[i]); + float wi = wn[i]; + float gi = gn[i]; + float wg = wi * gi; + factors = plus_f2(factors, {wg * t, t * t}); + } + } + factors = BlockReduceF2{block, temp}.Reduce(factors, plus_f2, {}); + float meangwx = factors.x / axis_size; + float normalizer = rsqrt(factors.y / axis_size + eps); + float normalizer3 = normalizer * normalizer * normalizer; + + // Outputs. + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) { + auto index = r * BLOCK_DIM + block.thread_rank(); + auto xn = load_vector(x, index, axis_size, T(0)); + auto gn = load_vector(g, index, axis_size, T(0)); + auto wn = load_vector(w, index, axis_size, w_stride, T(0)); + for (int i = 0; i < N_READS; i++) { + float xi = xn[i]; + float wi = wn[i]; + float gi = gn[i]; + xn[i] = static_cast(normalizer * wi * gi - xi * meangwx * normalizer3); + if constexpr (HAS_W) { + wn[i] = static_cast(gi * xi * normalizer); + } + } + store_vector(gx, index, xn, axis_size); + if constexpr (HAS_W) { + store_vector(gw, index, wn, axis_size); + } + } +} + +} // namespace cu + +namespace fast { + +bool RMSNorm::use_fallback(Stream s) { + return s.device == Device::cpu; +} + +template +void dispatch_group_dim(int axis_size, F&& f) { + if (axis_size <= n_per_thread * 8) { + f(std::integral_constant{}, + std::integral_constant(), + std::integral_constant()); + } else if (axis_size <= n_per_thread * 16) { + f(std::integral_constant{}, + std::integral_constant(), + std::integral_constant()); + } else if (axis_size <= n_per_thread * 32) { + f(std::integral_constant{}, + std::integral_constant(), + std::integral_constant()); + } else if (axis_size <= n_per_thread * 32 * 2) { + f(std::integral_constant{}, + std::integral_constant(), + std::integral_constant()); + } else if (axis_size <= n_per_thread * 32 * 4) { + f(std::integral_constant{}, + std::integral_constant(), + std::integral_constant()); + } else if (axis_size <= n_per_thread * 32 * 8) { + f(std::integral_constant{}, + std::integral_constant(), + std::integral_constant()); + } else if (axis_size <= n_per_thread * 32 * 16) { + f(std::integral_constant{}, + std::integral_constant(), + std::integral_constant()); + } else { + f(std::integral_constant{}, + std::integral_constant(), + std::integral_constant()); + } +} + +// TODO: There are duplicate code with backend/metal/normalization.cpp +void RMSNorm::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + nvtx3::scoped_range r("RMSNorm::eval_gpu"); + auto& s = stream(); + auto& out = outputs[0]; + auto& encoder = cu::get_command_encoder(s); + + // Make sure that the last dimension is contiguous. + auto set_output = [&s, &out, &encoder](const array& x) { + bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1; + if (no_copy && x.ndim() > 1) { + auto s = x.strides()[x.ndim() - 2]; + no_copy &= (s == 0 || s == x.shape().back()); + } + if (no_copy) { + if (x.is_donatable()) { + out.copy_shared_buffer(x); + } else { + out.set_data( + cu::malloc_async(x.data_size() * x.itemsize(), encoder), + x.data_size(), + x.strides(), + x.flags()); + } + return x; + } else { + array x_copy = contiguous_copy_gpu(x, s); + out.copy_shared_buffer(x_copy); + return x_copy; + } + }; + + const array x = set_output(inputs[0]); + const array& w = inputs[1]; + + int32_t axis_size = x.shape().back(); + int32_t n_rows = x.data_size() / axis_size; + int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0; + + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_output_array(out); + dispatch_float_types(out.dtype(), "rms_norm", [&](auto type_tag) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + if (axis_size <= N_READS * 1024) { + dispatch_group_dim( + axis_size, [&](auto group_dim, auto n_groups, auto groups_per_block) { + constexpr int block_dim = n_groups() * group_dim(); + auto kernel = + cu::rms_norm_small; + auto n_blocks = + (n_rows + groups_per_block() - 1) / groups_per_block(); + encoder.add_kernel_node( + kernel, + n_blocks, + {block_dim, groups_per_block()}, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(out), + eps_, + axis_size, + n_rows, + w_stride); + }); + } else { + auto kernel = cu::rms_norm; + encoder.add_kernel_node( + kernel, + n_rows, + 1024, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(out), + eps_, + axis_size, + w_stride); + } + }); +} + +void RMSNormVJP::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + nvtx3::scoped_range r("RMSNormVJP::eval_gpu"); + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + // Ensure row contiguity. We could relax this step by checking that the array + // is contiguous (no broadcasts or holes) and that the input strides are the + // same as the cotangent strides but for now this is simpler. + auto check_input = [&s](const array& x, bool& copied) { + if (x.flags().row_contiguous) { + copied = false; + return x; + } + copied = true; + return contiguous_copy_gpu(x, s); + }; + bool donate_x = inputs[0].is_donatable(); + bool donate_g = inputs[2].is_donatable(); + bool copied; + auto x = check_input(inputs[0], copied); + donate_x |= copied; + const array& w = inputs[1]; + bool g_copied; + auto g = check_input(inputs[2], g_copied); + donate_g |= g_copied; + array& gx = outputs[0]; + array& gw = outputs[1]; + + // Check whether we had a weight. + bool has_w = w.ndim() != 0; + + // Allocate space for the outputs. + bool g_in_gx = false; + if (donate_x) { + gx.copy_shared_buffer(x); + } else if (donate_g) { + gx.copy_shared_buffer(g); + g_in_gx = true; + } else { + gx.set_data(cu::malloc_async(gx.nbytes(), encoder)); + } + if (g_copied && !g_in_gx) { + encoder.add_temporary(g); + } + + int32_t axis_size = x.shape().back(); + int32_t n_rows = x.data_size() / axis_size; + int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0; + + // Allocate a temporary to store the gradients for w and allocate the output + // gradient accumulators. + array gw_temp = + (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w; + if (has_w) { + if (!g_in_gx && donate_g) { + gw_temp.copy_shared_buffer(g); + } else { + gw_temp.set_data(cu::malloc_async(gw_temp.nbytes(), encoder)); + encoder.add_temporary(gw_temp); + } + } + + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(g); + encoder.set_output_array(gx); + encoder.set_output_array(gw_temp); + dispatch_float_types(gx.dtype(), "rms_norm_vjp", [&](auto type_tag) { + dispatch_bool(has_w, [&](auto has_w_constant) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + if (axis_size <= N_READS * 1024) { + dispatch_group_dim( + axis_size, + [&](auto group_dim, auto n_groups, auto groups_per_block) { + constexpr int block_dim = group_dim() * n_groups(); + auto kernel = cu::rms_norm_vjp_small< + DataType, + has_w_constant.value, + block_dim, + group_dim(), + N_READS>; + auto n_blocks = + (n_rows + groups_per_block() - 1) / groups_per_block(); + encoder.add_kernel_node( + kernel, + n_blocks, + {block_dim, groups_per_block()}, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(g), + gpu_ptr(gx), + gpu_ptr(gw_temp), + eps_, + axis_size, + n_rows, + w_stride); + }); + } else { + auto kernel = + cu::rms_norm_vjp; + encoder.add_kernel_node( + kernel, + n_rows, + 1024, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(g), + gpu_ptr(gx), + gpu_ptr(gw_temp), + eps_, + axis_size, + w_stride); + } + }); + }); + + if (has_w) { + ReductionPlan plan( + ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size}); + col_reduce(encoder, gw_temp, gw, Reduce::ReduceType::Sum, {0}, plan); + } +} + +} // namespace fast + +} // namespace mlx::core diff --git a/src/lib/mlxcel-core/build.rs b/src/lib/mlxcel-core/build.rs index 676cb820..7596983d 100644 --- a/src/lib/mlxcel-core/build.rs +++ b/src/lib/mlxcel-core/build.rs @@ -176,7 +176,7 @@ fn main() { } /// Expected MLX git commit — must match GIT_TAG in mlx-cpp/CMakeLists.txt. -const MLX_EXPECTED_COMMIT: &str = "57c66cac7cb3e5b1eb350488a61f1506b40d39f8"; +const MLX_EXPECTED_COMMIT: &str = "b7c3dd6d27f45b5365b08a840310187dc503f1db"; /// Purge stale cached MLX build artifacts before CMake runs. /// diff --git a/src/loading/mod.rs b/src/loading/mod.rs index 8d78a94c..dd668a59 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -358,7 +358,13 @@ pub fn context_window_from_config(config: &serde_json::Value) -> Option { None } -/// Force CUDA graph capture off for Gemma 4 checkpoints (issue #688). +/// Force CUDA graph capture off for model families whose captured decode +/// collapses on this hardware: Gemma 4 (issue #688) and DeepSeek-V2 (issue +/// #824). The DeepSeek-V2 case is the same class of hazard: with graph capture +/// on, the naive-MLA incremental decode degenerates into repeated tokens even +/// though the identical forward is deterministically coherent with capture off +/// (verified via `MLX_USE_CUDA_GRAPHS=0`). The Gemma 4 analysis below applies to +/// that family; DeepSeek-V2 is added to the same lever for the same reason. /// /// mlxcel's Gemma 4 incremental (single-query, KV-cache) decode has a CUDA-graph /// read-before-write hazard: greedy (temperature 0) decode is nondeterministic @@ -376,29 +382,39 @@ pub fn context_window_from_config(config: &serde_json::Value) -> Option { /// This writes `MLX_USE_CUDA_GRAPHS=0` before the first GPU eval so MLX's /// process-wide `use_cuda_graphs()` static (read once, on the first graph-eligible /// eval) picks it up. An explicit operator value (either direction) is respected. -/// Non-Gemma-4 loads are untouched. +/// Loads of unaffected families are untouched. /// /// The sound home for this write is the main startup thread, upstream of any /// worker spawn: the server path calls -/// [`maybe_disable_cuda_graphs_for_gemma4_for_path`] from `start_server` before the +/// [`maybe_disable_cuda_graphs_for_model_for_path`] from `start_server` before the /// generation or pipeline worker is created (issue #688 M1/M2 hardening), and the /// CLI `generate`/`run` path loads on the main thread. The per-load-site calls that /// invoke this helper remain as idempotent defense-in-depth. /// /// Invariant: this env-based lever assumes one model per process. The server is /// single-model-per-process today (no in-process hot-swap in `start_server`), so a -/// non-Gemma eval never latches `use_cuda_graphs = true` before a later Gemma 4 -/// load. If in-process model hot-swap is ever added, a non-Gemma model loaded first -/// would make a subsequent Gemma 4 `set_var` a silent no-op (the static is already -/// latched) and this approach would need revisiting. -fn maybe_disable_cuda_graphs_for_gemma4(model_type: ModelType) { - let is_gemma4 = matches!( - model_type, - ModelType::Gemma4 | ModelType::Gemma4VLM | ModelType::Gemma4Unified - ); - if !is_gemma4 { +/// non-hazard-family eval never latches `use_cuda_graphs = true` before a later +/// hazard-family (Gemma 4 or DeepSeek-V2) load. If in-process model hot-swap is +/// ever added, a non-hazard-family model loaded first would make a subsequent +/// hazard-family `set_var` a silent no-op (the static is already latched) and this +/// approach would need revisiting. +fn maybe_disable_cuda_graphs_for_model(model_type: ModelType) { + // Families whose CUDA-graph-captured decode collapses on this hardware and + // must run with graph capture off. Gemma 4 (issue #688) and DeepSeek-V2 + // (issue #824: the naive-MLA decode path degenerates into repeated tokens + // under graph capture even though the same forward is coherent with capture + // off). Other families (Gemma 3, Qwen 3.5, etc.) are unaffected and keep + // capture on. + let hazard = match model_type { + ModelType::Gemma4 | ModelType::Gemma4VLM | ModelType::Gemma4Unified => { + Some(("Gemma 4", "#688")) + } + ModelType::DeepSeekV2 | ModelType::DeepSeekVL2 => Some(("DeepSeek-V2", "#824")), + _ => None, + }; + let Some((family, issue)) = hazard else { return; - } + }; // Respect an explicit operator override (either direction). if std::env::var_os("MLX_USE_CUDA_GRAPHS").is_some() { return; @@ -407,7 +423,7 @@ fn maybe_disable_cuda_graphs_for_gemma4(model_type: ModelType) { // Rust 2024 only when another thread runs getenv/setenv concurrently. The // authoritative write is hoisted to the main startup thread, upstream of any // worker spawn and of the first GPU eval that latches MLX's `use_cuda_graphs` - // static: `maybe_disable_cuda_graphs_for_gemma4_for_path` runs in `start_server` + // static: `maybe_disable_cuda_graphs_for_model_for_path` runs in `start_server` // for every serve path (batched, legacy, tensor-parallel, XLA, in-process and // remote pipeline-parallel), and the CLI `generate`/`run` path loads on the main // thread. On the server path this per-load-site call runs inside the spawned @@ -421,13 +437,13 @@ fn maybe_disable_cuda_graphs_for_gemma4(model_type: ModelType) { // bring-up. unsafe { std::env::set_var("MLX_USE_CUDA_GRAPHS", "0") }; tracing::info!( - "Gemma 4 detected: disabling CUDA graph capture (MLX_USE_CUDA_GRAPHS=0) to \ - avoid the decode-path graph hazard from issue #688. Set MLX_USE_CUDA_GRAPHS \ + "{family} detected: disabling CUDA graph capture (MLX_USE_CUDA_GRAPHS=0) to \ + avoid the decode-path graph hazard from issue {issue}. Set MLX_USE_CUDA_GRAPHS \ explicitly to override." ); } -/// Path-based entry to [`maybe_disable_cuda_graphs_for_gemma4`] for the main-thread +/// Path-based entry to [`maybe_disable_cuda_graphs_for_model`] for the main-thread /// hoist (issue #688 M1/M2 hardening). /// /// Peeks the on-disk model type and, for a Gemma 4 checkpoint, performs the @@ -436,10 +452,10 @@ fn maybe_disable_cuda_graphs_for_gemma4(model_type: ModelType) { /// detection error (unreadable or absent `config.json`) is a no-op, which preserves /// the baseline for non-model paths and defers to the per-load-site calls. Kept as a /// thin wrapper so the Gemma-4 detection, operator-override guard, idempotency, and -/// logging all stay in the single [`maybe_disable_cuda_graphs_for_gemma4`] helper. -pub(crate) fn maybe_disable_cuda_graphs_for_gemma4_for_path(model_path: &Path) { +/// logging all stay in the single [`maybe_disable_cuda_graphs_for_model`] helper. +pub(crate) fn maybe_disable_cuda_graphs_for_model_for_path(model_path: &Path) { if let Ok(model_type) = get_model_type(model_path) { - maybe_disable_cuda_graphs_for_gemma4(model_type); + maybe_disable_cuda_graphs_for_model(model_type); } } @@ -448,7 +464,7 @@ pub fn load_model(model_path: &Path) -> Result<(LoadedModel, MlxcelTokenizer)> { let model_path = resolve_model_dir(model_path); let model_path = model_path.as_path(); let model_type = get_model_type(model_path)?; - maybe_disable_cuda_graphs_for_gemma4(model_type); + maybe_disable_cuda_graphs_for_model(model_type); // Whisper is an encoder-decoder ASR model, not a text generator. It is // wired into the speech-to-text audio endpoints at server startup rather @@ -555,7 +571,7 @@ pub fn load_model_with_tensor_parallel( let model_path = resolve_model_dir(model_path); let model_path = model_path.as_path(); let support = validate_supported_runtime(model_path, shard_config.clone(), adapter_path)?; - maybe_disable_cuda_graphs_for_gemma4(support.summary.model_type); + maybe_disable_cuda_graphs_for_model(support.summary.model_type); let model = match support.summary.model_type { ModelType::Llama | ModelType::Qwen2 => LoadedModel::TensorParallelLlama( TensorParallelLlamaModel::from_model_dir(model_path, shard_config.clone())?, @@ -599,7 +615,7 @@ pub fn load_model_with_adapter( let model_path = resolve_model_dir(model_path); let model_path = model_path.as_path(); // Disable CUDA graphs for Gemma 4 before any weight realisation (issue #688). - maybe_disable_cuda_graphs_for_gemma4(get_model_type(model_path)?); + maybe_disable_cuda_graphs_for_model(get_model_type(model_path)?); // Load base weights let base_weights = mlxcel_core::weights::load_weights_from_dir(model_path) .map_err(|e| anyhow::anyhow!("{}", e))?; @@ -617,7 +633,7 @@ pub fn load_model_with_adapter( /// Build a model from pre-loaded weights (used by adapter loading) fn load_model_from_weights(model_path: &Path, weights: &mut WeightMap) -> Result { let model_type = get_model_type(model_path)?; - maybe_disable_cuda_graphs_for_gemma4(model_type); + maybe_disable_cuda_graphs_for_model(model_type); let config_path = model_path.join("config.json"); let config_str = std::fs::read_to_string(&config_path)?; let config_str = sanitize_config_json(&config_str); diff --git a/src/server/startup.rs b/src/server/startup.rs index 35488bb5..21c87f41 100644 --- a/src/server/startup.rs +++ b/src/server/startup.rs @@ -1563,18 +1563,19 @@ fn install_surgery_pipeline_for_server(startup: &ServerStartupConfig) -> Result< pub async fn start_server(mut startup: ServerStartupConfig) -> Result<()> { initialize_server_logging(&startup)?; - // Issue #688 (M1/M2 hardening): disable CUDA graph capture for Gemma 4 here, - // on the main startup thread, before any generation or pipeline worker is - // spawned and before the first GPU eval latches MLX's process-wide - // `use_cuda_graphs` static. Performing the env write on this thread (rather than - // only at the per-load-site calls, which run inside the spawned worker) keeps it - // sound under Rust 2024's concurrent-getenv rule. This single chokepoint covers - // every serve path that flows through `start_server`: the batched, legacy, + // Issue #688 (M1/M2 hardening), extended to DeepSeek-V2 by issue #824: disable + // CUDA graph capture for hazard-family models (Gemma 4, DeepSeek-V2) here, on + // the main startup thread, before any generation or pipeline worker is spawned + // and before the first GPU eval latches MLX's process-wide `use_cuda_graphs` + // static. Performing the env write on this thread (rather than only at the + // per-load-site calls, which run inside the spawned worker) keeps it sound + // under Rust 2024's concurrent-getenv rule. This single chokepoint covers every + // serve path that flows through `start_server`: the batched, legacy, // tensor-parallel and XLA workers, the in-process pipeline-parallel worker, and // the remote pipeline stage worker (the `serve_remote_pipeline_stage` branch - // below) — all load `startup.model_path`. Non-Gemma-4 models and an explicit - // `MLX_USE_CUDA_GRAPHS` operator override are left untouched. - crate::loading::maybe_disable_cuda_graphs_for_gemma4_for_path(&startup.model_path); + // below) — all load `startup.model_path`. Unaffected model families and an + // explicit `MLX_USE_CUDA_GRAPHS` operator override are left untouched. + crate::loading::maybe_disable_cuda_graphs_for_model_for_path(&startup.model_path); super::media::configure_image_input_limits(super::media::ImageInputLimits { max_payload_bytes: startup.max_image_payload_size, diff --git a/src/vision/processors/minimax_m3.rs b/src/vision/processors/minimax_m3.rs index bd78e656..4e4f1ee5 100644 --- a/src/vision/processors/minimax_m3.rs +++ b/src/vision/processors/minimax_m3.rs @@ -267,7 +267,7 @@ mod tests { mlxcel_core::eval(&pixel_values); let shape = mlxcel_core::array_shape(&pixel_values); let (t, gh, gw) = grid[0]; - let expected_rows = (t * gh * gw) as i32; + let expected_rows = t * gh * gw; assert_eq!(shape, vec![expected_rows, 1176]); }