diff --git a/docs/benchmark_results/single-dtype-decode-astype-gb10-2026-07-10.md b/docs/benchmark_results/single-dtype-decode-astype-gb10-2026-07-10.md new file mode 100644 index 00000000..ca6c1988 --- /dev/null +++ b/docs/benchmark_results/single-dtype-decode-astype-gb10-2026-07-10.md @@ -0,0 +1,54 @@ +# Single-dtype decode graph: per-token AsType inventory on CUDA (GB10) + +Issue #636. Host: GB10 (NVIDIA GB10, SM 12.1), CUDA backend, MLX pin per `main` at 70b2b29. All counts from `MLXCEL_TRACE_ASTYPE`, which walks the first decode step's unevaluated `(token, logprob)` graph and counts `AsType` nodes (traversal only, no extra eval). Cross-checked against `MLXCEL_EXPORT_DECODE_DOT` (`grep -c AsType` on the DOT) which agrees at 0. + +## Headline finding + +On CUDA the greedy quantized decode graph is already single-dtype: the three inventory models emit **0 AsType nodes per token**. The ~773 AsType/token figure in `moe-decode-gap-investigation.md` was measured on Apple/Metal (M1 Ultra); the CUDA-only `src/lib/mlx-cpp/patches-cuda/dtype.cpp` promotion patch (`bf16 + fp32 -> bf16`) already collapses the diffuse scalar/constant promotions that produced it. The counter tool merged here makes that state measurable and guards it against regression. + +The only remaining reducible AsType source on CUDA is the temperature-sampler chain, addressed by building the fused sampler's scalar constants in the logit dtype. + +## Before / after AsType per decode step + +| Model | Path | Before | After | Reduction | +|-------|------|-------:|------:|----------:| +| llama-3.1-8b-4bit | greedy (temp 0) | 0 | 0 | already single-dtype | +| qwen3-8b-4bit | greedy (temp 0) | 0 | 0 | already single-dtype | +| qwen3-30b-a3b-4bit | greedy (temp 0) | 0 | 0 | already single-dtype | +| qwen3-8b-4bit | temp 0.8 + top-k 40 + top-p 0.9 | 4 | 1 | 75% | +| qwen3-8b-4bit | temp 0.8 + min-p 0.05 | 4 | 1 | 75% | + +The residual `1` on the sampling path is the intrinsic `u32 -> f32` inside `mlx::core::random::categorical` (uniform bit-draw to float for the gumbel key), not a logit dtype round-trip. + +On CUDA a natively-f16 checkpoint now runs temperature sampling in f16 rather than the previous f32 upcast. This is exactly the single-dtype behavior the issue asks for on CUDA, and it is within fp16 tolerance: an f16 softmax over well-separated logits on an inherently stochastic categorical draw. The scalar-dtype change is gated on `!metal::is_available()`, so Metal keeps its bare-f32-scalar behavior (the unpatched f16+f32 rule upcasts the chain to an f32 softmax) and its temperature-sampling numerics are unchanged. + +## Inventory: where the AsType come from + +| Source | Model class | Before | Disposition | +|--------|-------------|-------:|-------------| +| Model body (weights/norms/rope/attention) on CUDA | quantized dense + MoE | 0 | Already single-dtype via the `patches-cuda/dtype.cpp` bf16 promotion patch and the consistent-dtype quant path. Nothing to reduce. | +| Model body, bf16-native checkpoints on CUDA | dense bf16 | 0 | Single-dtype bf16; CUDA keeps bf16 (native ALUs). Optional f16 normalization available via env, no AsType change. | +| Fused sampler scalars (temperature, top-k/top-p/min-p sentinels) | any, temperature sampling | 3 (`f32 -> bf16`) | Removed on non-Metal backends: scalars built in the logit dtype so the CUDA promotion patch inserts no per-step conversion. Bit-identical for bf16 logits. Metal keeps the bare f32 scalars unchanged (gated on `!metal::is_available()`), so its numerics are untouched. | +| `random::categorical` uniform draw | any, temperature sampling | 1 (`u32 -> f32`) | Intrinsic to random generation; left as-is. | +| RMSNorm fp32 accumulation | gemma family (f16-fragile) | 52 (`f16 <-> f32`, 2 per layer) | Intentionally kept: gemma norms + softcap need the wider accumulation. On the f16-fragile exception list. | + +## Quality: greedy parity + +40-token greedy (temp 0), `MLXCEL_CUDA_F16_NORMALIZE` off vs on, generated text compared byte-for-byte: + +- llama-3.1-8b-4bit: identical (quantized, env is a no-op). +- qwen3-30b-a3b-4bit: identical (quantized, env is a no-op). +- llama-3.1-8b-bf16: identical (bf16 vs f16-normalized both produce the same 40 tokens on this healthy family). + +## Decode tok/s A/B (env on vs off) + +`mlxcel-bench-decode --prompt x --prompt-tokens 512 --max-tokens 128`, two runs each: + +| Model | off (bf16) | on | Note | +|-------|-----------:|---:|------| +| qwen3-30b-a3b-4bit | 92.44 / 90.59 | 92.84 / 91.63 | quantized: env is a no-op, within run-to-run noise | +| qwen2.5-0.5b-bf16 | 208.05 / 207.79 | 206.82 / 207.33 | bf16 vs f16 within noise, confirming CUDA native bf16 (no f16 throughput win) | + +## Why CUDA f16 normalization is opt-in, not default + +The single-dtype objective is already met in bf16 on CUDA, and CUDA has native bf16 compute, so casting bf16 -> f16 yields no AsType reduction and no throughput gain while narrowing dynamic range. The load-time f16 path (`MLXCEL_CUDA_F16_NORMALIZE`) is therefore off by default, with a conservative f16-fragile exception list (gemma, cohere/command-r, apertus, gpt-oss, and any softcap/`logit_scale` config) that keeps bf16 even when enabled. Metal/Apple Silicon numerics are untouched: the always-on Apple bf16 -> f16 policy is unchanged and the new path is gated on the CUDA backend. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 0a1fcac3..cc6fa477 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -231,6 +231,7 @@ recommended as normal deployment settings. | `MLXCEL_COMPILED_QGELU_MLP` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | Compiles the affine-quantized GeGLU MLP (gate/up/gelu/down) into one `mx::compile` graph so the tanh-approx GELU's ~14 element-wise ops per layer collapse into a single fused `Compiled` primitive. Covers the Gemma family (gemma/gemma2/gemma3/gemma4). The group_size=64/bits=4 case was already compiled on every shape and is unchanged; other affine quantizations (notably the group_size=64/bits=8 MLP weights in Gemma 4 mixed-precision checkpoints, issue #680) are compiled only on the single-token decode call, because compiling the 8-bit prefill GEMM measured 8-9% slower and +0.2-0.7 GB peak on GB10 (the shapeless fused graph forces a decode-oriented qmm kernel onto the large prefill matmul). GB10 gemma-4-12b decode is weight-bandwidth-bound (~94% GPU-busy), so this primitive-count cut is throughput-neutral there (measured <1% from a ~530-primitive/step reduction); it still removes real CPU dispatch and helps op-count-bound backends. Set to `0` to force the op-at-a-time fallback (A/B + rollback). It also gates the NVFP4 scaled fused MLP variant (`MLXCEL_DISABLE_FUSED_GLOBAL_SCALE`), forcing its eager fold when set to `0`. | | `MLXCEL_DISABLE_FUSED_GLOBAL_SCALE` | `1`/`true`/`on`/`yes` (case-insensitive) disable; unset or any other value keeps the fold | off (fold on) | Rolls back the NVFP4 global-scale fold for Gemma 4 (issues #698/#705). By default the fused MLP and per-layer-input-gate C++ paths fold each per-projection `weight_scale_2` sidecar (from the direct ModelOpt transcode, issue #693/#697) into the fused kernel at the mathematically correct points: the gate scale before the GeGLU activation, the up scale on the up product, and the down scale on the fused output, each reproducing `apply_global_scale` byte-for-byte. Native NVFP4 prefill now uses a shape-specific scaled MLP graph, and standalone `UnifiedLinear` sidecar projections can apply qmm + global scale + dense bias through one C++ helper. When set, sidecar-carrying paths fall back to the op-at-a-time `UnifiedLinear::forward` scalar application (the pre-#698 bypass). Greedy temp-0 decode is token-identical across the two paths on `gemma-4-31b-it-nvfp4`; the fold removes element-wise dispatches on the gemma-4 path where CUDA graphs are disabled (#688). See `docs/benchmark_results/nvfp4-direct-transcode-gb10-2026-07-08.md` and `docs/benchmark_results/nvfp4-native-prefill-m1ultra-2026-07-09.md`. | | `MLXCEL_FUSED_XIELU` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | Fused single-launch Metal xIELU kernel for the Apertus MLP activation (#409), on by default since the M5 Max validation. `MLP::forward` routes through one Metal dispatch covering the ~11 elementwise ops in `apertus_xielu` (square, minimum, expm1, where, and neighbors) instead of the per-op graph. Greedy temp-0 decode is byte-identical to the elementwise path on Apple Silicon: every intermediate stays in the input dtype (bf16) and the kernel reproduces MLX's `expm1f` exactly. Measured decode speedup on M1 Ultra (+2.7%, Apertus-8B 83.4 to 85.7 tok/s) and M5 Max (+1.9%, 112.0 to 114.2 tok/s), with no regression. Set to `0` to force the elementwise path. On non-Metal back-ends the FFI falls back to an equivalent elementwise graph, so the flag is safe to set everywhere. Apertus only; no other model family is affected. | +| `MLXCEL_CUDA_F16_NORMALIZE` | `1`/`true`/`on`/`yes` enable; unset or `0`/`false`/`off`/`no` keep bf16 | off (opt-in) | **CUDA only.** Opt-in load-time bf16 -> f16 normalization of non-quantized weights for the single-dtype decode graph (issue #636). Off by default: the merged `patches-cuda/dtype.cpp` bf16 promotion patch already yields a 0-AsType single-dtype bf16 decode graph on CUDA, and CUDA GPUs have native bf16 ALUs (unlike Apple Silicon), so f16 offers no measured throughput gain (qwen2.5-0.5b-bf16 decode 208 vs 207 tok/s bf16 vs f16 on GB10) and narrows dynamic range. When enabled, healthy dense families cast to f16 for fixed-topology / CUDA-graph-reuse experiments; f16-fragile families (gemma, cohere/command-r, apertus, gpt-oss, and any config with a nonzero softcap or `logit_scale`) stay bf16 regardless. Quantized checkpoints are unaffected (their conversion path is skipped), so leaving this unset keeps all bf16 models available as before. Metal/Apple Silicon is governed by the separate always-on bf16->f16 policy and is untouched. | ## Block-diffusion diagnostic variables @@ -246,6 +247,7 @@ throughput measurements. Use them for diagnosis, not capacity planning. | Variable | Values | Default | Purpose | |----------|--------|---------|---------| | `MLXCEL_TRACE_DTYPE` | presence enables | off | Prints selected tensor dtypes/shapes during generation. | +| `MLXCEL_TRACE_ASTYPE` | presence enables; `2`/`break` adds breakdown | off | Prints the AsType (dtype-conversion) node count in the first decode step's graph (the single-dtype decode-graph metric, issue #636). Set to `2` or a value containing `break` to also dump the per src->dst dtype breakdown. Graph traversal only, no extra eval; zero cost when unset. Do not combine with `MLXCEL_TRACE_DTYPE`, which pre-evaluates the logits and collapses the graph before the count. | | `MLXCEL_FORCE_SYNC` | presence enables | off | Forces synchronous decode evaluation. Also disables the server `BatchScheduler`'s lookahead decode pipeline (issue #632), falling back to the pre-pipeline synchronous tick. | | `MLXCEL_PROFILE_PIPELINE` | presence enables | off | Emits high-level generation pipeline timing. | | `MLXCEL_PROFILE_PIPELINE_DETAIL` | presence enables | off | Adds per-step pipeline timing detail. | diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index cb7f238f..ef73b2af 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -3,6 +3,14 @@ #include "mlx_cxx_internal.h" +#include "mlx/primitives.h" + +#include +#include +#include +#include +#include + namespace mlx_cxx { using namespace mlx::core; @@ -4346,6 +4354,91 @@ void export_to_dot_pair(rust::Str path, const MlxArray& a, const MlxArray& b) { mlx::core::export_to_dot(os, arrs); } +namespace { + inline const char* astype_dtype_short_name(int32_t d) { + switch (d) { + case 0: return "bool"; + case 1: return "u8"; + case 2: return "u16"; + case 3: return "u32"; + case 4: return "u64"; + case 5: return "i8"; + case 6: return "i16"; + case 7: return "i32"; + case 8: return "i64"; + case 9: return "f16"; + case 10: return "f32"; + case 11: return "f64"; + case 12: return "bf16"; + case 13: return "c64"; + default: return "?"; + } + } + + // Depth-first walk over the unevaluated graph rooted at `outputs`, + // deduplicated by array id(). Counts AsType primitive nodes, the total + // node count, and a per (src_dtype -> dst_dtype) AsType breakdown. + void collect_astype_stats( + const std::vector& outputs, + uint64_t& astype_count, + uint64_t& total_nodes, + std::map, uint64_t>& breakdown) { + std::unordered_set visited; + std::vector stack(outputs.begin(), outputs.end()); + while (!stack.empty()) { + array a = stack.back(); + stack.pop_back(); + if (!visited.insert(a.id()).second) { + continue; + } + total_nodes++; + if (a.has_primitive()) { + if (std::strcmp(a.primitive().name(), "AsType") == 0) { + astype_count++; + int32_t dst = from_dtype(a.dtype()); + int32_t src = a.inputs().empty() + ? -1 + : from_dtype(a.inputs()[0].dtype()); + breakdown[std::make_pair(src, dst)]++; + } + for (const auto& in : a.inputs()) { + stack.push_back(in); + } + } + } + } +} // namespace + +// Count AsType (dtype-conversion) nodes in the unevaluated graph that produces +// the given pair of arrays. This is the "done" metric for the single-dtype +// decode graph work (issue #636): a decode step traced with a consistent dtype +// produces zero (or near-zero) AsType nodes. Traversal only; no eval. +uint64_t count_astype_nodes_pair(const MlxArray& a, const MlxArray& b) { + std::vector outs = {a.inner, b.inner}; + uint64_t astype_count = 0; + uint64_t total_nodes = 0; + std::map, uint64_t> breakdown; + collect_astype_stats(outs, astype_count, total_nodes, breakdown); + return astype_count; +} + +// Human-readable AsType breakdown for the same graph: a header line with the +// AsType and total node counts, followed by one line per src->dst dtype pair. +rust::String astype_breakdown_pair(const MlxArray& a, const MlxArray& b) { + std::vector outs = {a.inner, b.inner}; + uint64_t astype_count = 0; + uint64_t total_nodes = 0; + std::map, uint64_t> breakdown; + collect_astype_stats(outs, astype_count, total_nodes, breakdown); + std::ostringstream os; + os << "astype_nodes=" << astype_count << " total_nodes=" << total_nodes; + for (const auto& kv : breakdown) { + os << "\n " << astype_dtype_short_name(kv.first.first) << "->" + << astype_dtype_short_name(kv.first.second) << " : " << kv.second; + } + return rust::String(os.str()); +} + // Set default stream for subsequent operations void set_default_stream(const MlxStream& stream) { mlx::core::set_default_stream(stream.inner); @@ -4361,16 +4454,26 @@ bool is_gpu_available() { // which causes "CumSum cannot infer output shapes" when used inside // mlx::core::compile with shapeless=true. namespace { - array top_p_filter(const array& x, float top_p) { + // `single_dtype_scalars` is true on non-Metal backends (issue #636): build + // the comparison/fill scalars in the working dtype so the CUDA bf16 + // promotion patch inserts no per-step AsType on the scalar. On Metal it is + // false, keeping the bare `array(f32)` scalars exactly as before so the + // unpatched f16+f32 rule upcasts the chain to an f32 softmax unchanged. + array top_p_filter(const array& x, float top_p, bool single_dtype_scalars) { auto probs = mlx::core::softmax(x, -1); auto sorted_indices = mlx::core::argsort(mlx::core::negative(probs), -1); auto sorted_probs = mlx::core::take_along_axis(probs, sorted_indices, -1); auto cum_probs = mlx::core::cumsum(sorted_probs, -1, false, true); auto shifted_cum = cum_probs - sorted_probs; - auto mask = mlx::core::less_equal(shifted_cum, mlx::core::array(top_p)); + auto top_p_scalar = single_dtype_scalars + ? mlx::core::array(top_p, shifted_cum.dtype()) + : mlx::core::array(top_p); + auto fill_scalar = single_dtype_scalars + ? mlx::core::array(std::numeric_limits::lowest(), x.dtype()) + : mlx::core::array(std::numeric_limits::lowest()); + auto mask = mlx::core::less_equal(shifted_cum, top_p_scalar); auto sorted_logits = mlx::core::take_along_axis(x, sorted_indices, -1); - auto filtered_sorted = mlx::core::where( - mask, sorted_logits, mlx::core::array(std::numeric_limits::lowest())); + auto filtered_sorted = mlx::core::where(mask, sorted_logits, fill_scalar); auto unsort_indices = mlx::core::argsort(sorted_indices, -1); return mlx::core::take_along_axis(filtered_sorted, unsort_indices, -1); } @@ -4380,15 +4483,23 @@ namespace { namespace { static std::function(const std::vector&)> get_compiled_min_p_filter() { - auto fn = [](const std::vector& inputs) -> std::vector { + // Backend is process-constant, so capture the gate once at build time. + // Non-Metal fills the sentinel in x's dtype so no per-step AsType is + // inserted on the scalar under the CUDA bf16 promotion patch (issue + // #636); Metal keeps the bare f32 sentinel exactly as before. + const bool single_dtype_scalars = !mlx::core::metal::is_available(); + auto fn = [single_dtype_scalars]( + const std::vector& inputs) -> std::vector { const auto& x = inputs[0]; const auto& min_p_arr = inputs[1]; auto probs = mlx::core::softmax(x, -1); auto max_prob = mlx::core::max(probs, -1, true); auto threshold = mlx::core::multiply(max_prob, min_p_arr); auto mask = mlx::core::greater_equal(probs, threshold); - return {mlx::core::where(mask, x, - mlx::core::array(std::numeric_limits::lowest()))}; + auto fill_scalar = single_dtype_scalars + ? mlx::core::array(std::numeric_limits::lowest(), x.dtype()) + : mlx::core::array(std::numeric_limits::lowest()); + return {mlx::core::where(mask, x, fill_scalar)}; }; return mlx::core::compile(fn, true); } @@ -4417,9 +4528,28 @@ std::unique_ptr fused_sample( return std::make_unique(mlx::core::argmax(x, -1, false)); } + // Build sampler scalars in the logit dtype on non-Metal backends only + // (issue #636). On CUDA the bf16 promotion patch keeps bf16 op boundaries + // in bf16, but a bare `array(f32)` scalar against a bf16/f16 logit tensor + // still forces an AsType conversion on the scalar every decode step; + // constructing each scalar in `x.dtype()` keeps the fused sampler chain + // single-dtype (bit-identical for bf16 logits since the old path cast the + // f32 scalar to bf16 anyway). Metal is left exactly as before: dtype.cpp + // deliberately leaves f16+f32 -> f32 unpatched there, so a bare f32 scalar + // upcasts the chain to an f32 softmax, and issue #636 requires Metal + // numerics untouched. `single_dtype_scalars` uses the same runtime + // `!metal::is_available()` selection as the fused-kernel ports in + // mlx_cxx_kernels.cpp (first use of that gate in this file). + const bool single_dtype_scalars = !mlx::core::metal::is_available(); + const auto logit_dtype = x.dtype(); + auto sampler_scalar = [&](float value) { + return single_dtype_scalars ? mlx::core::array(value, logit_dtype) + : mlx::core::array(value); + }; + // Temperature scaling if (temperature > 0.0f && temperature != 1.0f) { - x = x / mlx::core::array(temperature); + x = x / sampler_scalar(temperature); } // Top-k filtering: keep only the k highest-probability tokens @@ -4438,18 +4568,19 @@ std::unique_ptr fused_sample( auto kth_idx = mlx::core::slice(indices, start, stop); auto threshold = mlx::core::take_along_axis(x, kth_idx, -1); auto mask = mlx::core::greater_equal(x, threshold); - x = mlx::core::where(mask, x, mlx::core::array(std::numeric_limits::lowest())); + x = mlx::core::where( + mask, x, sampler_scalar(std::numeric_limits::lowest())); } // Top-p (nucleus) filtering if (top_p > 0.0f && top_p < 1.0f) { - x = top_p_filter(x, top_p); + x = top_p_filter(x, top_p, single_dtype_scalars); } // Min-p filtering — compiled kernel if (min_p > 0.0f && min_p < 1.0f) { static auto compiled_fn = get_compiled_min_p_filter(); - auto result = compiled_fn({x, mlx::core::array(min_p)}); + auto result = compiled_fn({x, sampler_scalar(min_p)}); x = std::move(result[0]); } diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 5bb88ced..a3558be4 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1255,6 +1255,12 @@ void async_eval_pair(const MlxArray& a, const MlxArray& b); // Export a pair of unevaluated arrays as a DOT graph for profiling. void export_to_dot_pair(rust::Str path, const MlxArray& a, const MlxArray& b); +// Count AsType (dtype-conversion) nodes in the unevaluated graph for a pair. +uint64_t count_astype_nodes_pair(const MlxArray& a, const MlxArray& b); + +// Human-readable AsType breakdown (counts + per src->dst dtype pair) for a pair. +rust::String astype_breakdown_pair(const MlxArray& a, const MlxArray& b); + // Set default stream for subsequent operations void set_default_stream(const MlxStream& stream); diff --git a/src/lib/mlxcel-core/src/generate.rs b/src/lib/mlxcel-core/src/generate.rs index c8cf3cdf..d5f4d1e8 100644 --- a/src/lib/mlxcel-core/src/generate.rs +++ b/src/lib/mlxcel-core/src/generate.rs @@ -1186,6 +1186,12 @@ impl CxxGenerator { // Hoist env var checks out of the hot loop to avoid per-token syscalls. let trace_dtype = std::env::var("MLXCEL_TRACE_DTYPE").is_ok(); + // AsType-per-token counter (issue #636). When set, the first decode + // step reports the number of AsType (dtype-conversion) nodes in the + // decode+sampler graph; `=2` (or any value containing "break") also + // dumps the per src->dst dtype breakdown. Traversal only, no extra + // eval, and entirely skipped when the env var is unset. + let trace_astype = std::env::var("MLXCEL_TRACE_ASTYPE").ok(); let force_sync = std::env::var("MLXCEL_FORCE_SYNC").is_ok(); let profile_pipeline = std::env::var("MLXCEL_PROFILE_PIPELINE").is_ok(); let profile_pipeline_detail = std::env::var("MLXCEL_PROFILE_PIPELINE_DETAIL").is_ok(); @@ -1335,6 +1341,18 @@ impl CxxGenerator { { ffi::export_to_dot_pair(&path, &next_tok, &next_log); } + if n == 0 + && let Some(mode) = trace_astype.as_deref() + { + // Count on the unevaluated decode+sampler graph so no + // conversion is hidden by a prior eval. + let count = ffi::count_astype_nodes_pair(&next_tok, &next_log); + eprintln!("[ASTYPE] decode astype_nodes={count}"); + if mode.contains("break") || mode == "2" { + let breakdown = ffi::astype_breakdown_pair(&next_tok, &next_log); + eprintln!("[ASTYPE] {breakdown}"); + } + } // Optional Metal GPU capture of one warm decode token for // per-kernel profiling vs mlx-lm. Fires at n==2 so // all decode kernels are JIT-cached. Requires the process to be diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index c83d6b93..8da3eff0 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -1858,6 +1858,15 @@ mod ffi { /// Export two unevaluated arrays as a DOT graph for profiling. fn export_to_dot_pair(path: &str, a: &MlxArray, b: &MlxArray); + /// Count AsType (dtype-conversion) nodes in the unevaluated graph that + /// produces the given pair of arrays. Metric for the single-dtype + /// decode graph work (issue #636); traversal only, no eval. + fn count_astype_nodes_pair(a: &MlxArray, b: &MlxArray) -> u64; + + /// Human-readable AsType breakdown (counts plus per src->dst dtype pair) + /// for the graph that produces the given pair of arrays. + fn astype_breakdown_pair(a: &MlxArray, b: &MlxArray) -> String; + /// Set default stream for subsequent operations fn set_default_stream(stream: &MlxStream); diff --git a/src/models/sanitize.rs b/src/models/sanitize.rs index a056bb63..22bb376c 100644 --- a/src/models/sanitize.rs +++ b/src/models/sanitize.rs @@ -1670,7 +1670,17 @@ pub fn load_text_weights>( // attributed to bf16 scales (Apertus-2509) was actually a separate xIELU // read_scalar bf16 bug, fixed in the apertus loader; Apertus, Seed-OSS, and // every other bf16-scale quant decode correctly with no scale promotion. - if should_convert_bf16_to_f16() && !is_bitnet && !is_quantized { + // + // On CUDA the base policy keeps bf16 (native bf16 ALUs, and the merged + // patches-cuda/dtype.cpp promotion patch already yields a 0-AsType, + // single-dtype bf16 decode graph). Issue #636 adds an opt-in load-time + // bf16 -> f16 normalization for the fixed-topology / CUDA-graph-reuse + // experiment, gated by MLXCEL_CUDA_F16_NORMALIZE and skipped for + // f16-fragile families (see `cuda_f16_normalize_for_config`). The default + // stays bf16 so bf16 models remain available unchanged. + let convert_bf16 = + should_convert_bf16_to_f16() || cuda_f16_normalize_for_config(parsed_config.as_ref()); + if convert_bf16 && !is_bitnet && !is_quantized { let had_bf16 = if keep_gemma3n_mlp_bf16 { convert_bf16_weights_with_keep(&mut weights, gemma3n_language_mlp_bf16_key) } else { @@ -1684,6 +1694,86 @@ pub fn load_text_weights>( Ok(weights) } +/// True when opt-in CUDA load-time bf16 -> f16 normalization applies to this +/// model (issue #636). +/// +/// CUDA-only and default-off. Metal / Apple Silicon is driven by +/// [`should_convert_bf16_to_f16`] and is unaffected. Returns true only when: +/// - the MLX CUDA backend is active at runtime, and +/// - `MLXCEL_CUDA_F16_NORMALIZE` is set to a truthy value (opt-in; unset or a +/// falsy value keeps bf16 so bf16 models remain available), and +/// - the model is not on the conservative f16-fragile exception list. +fn cuda_f16_normalize_for_config(config: Option<&Value>) -> bool { + if !mlxcel_core::cuda_is_available() { + return false; + } + if !env_flag_enabled("MLXCEL_CUDA_F16_NORMALIZE") { + return false; + } + !config.is_some_and(is_f16_fragile_family) +} + +/// Parse a boolean-ish environment flag. Unset, empty, `0`, `false`, `off`, and +/// `no` (any case) are false; anything else is true. +fn env_flag_enabled(name: &str) -> bool { + std::env::var(name).ok().is_some_and(|raw| { + let v = raw.trim(); + !(v.is_empty() + || v == "0" + || v.eq_ignore_ascii_case("false") + || v.eq_ignore_ascii_case("off") + || v.eq_ignore_ascii_case("no")) + }) +} + +/// Conservative f16-fragile family check for CUDA f16 normalization. +/// +/// f16's narrow dynamic range (max ~65504) clips wide-range activations and +/// softcapped logits, so these families keep bf16 even when normalization is +/// requested. Detection is heuristic (softcap/logit-scale config keys plus a +/// known-family substring list): families it does not recognize are treated +/// as healthy and normalized, which is why the whole path stays opt-in and +/// default-off. Extend the list when a new fragile family surfaces. +fn is_f16_fragile_family(config: &Value) -> bool { + // Softcap / logit-scaling config keys imply wide dynamic range. Checked + // at the top level and under `text_config`, mirroring the `model_type` + // fallback below, so nested multimodal configs are not a blind spot. + for key in [ + "attn_logit_softcapping", + "final_logit_softcapping", + "logit_softcapping", + "logits_soft_cap", + "logit_scale", + ] { + let capped = |c: &Value| { + c.get(key) + .and_then(Value::as_f64) + .is_some_and(|v| v != 0.0 && v != 1.0) + }; + if capped(config) || config.get("text_config").is_some_and(capped) { + return true; + } + } + let model_type = config + .get("model_type") + .and_then(Value::as_str) + .or_else(|| { + config + .get("text_config") + .and_then(|c| c.get("model_type")) + .and_then(Value::as_str) + }) + .unwrap_or(""); + // Gemma (norms + softcap), Cohere/Command-R (logit scale), Apertus (xIELU + // x^2), and gpt-oss (wide dynamic range) are the known-fragile families. + const FRAGILE_SUBSTRINGS: &[&str] = &[ + "gemma", "cohere", "command", "apertus", "gpt_oss", "gpt-oss", + ]; + FRAGILE_SUBSTRINGS + .iter() + .any(|needle| model_type.contains(needle)) +} + /// Returns true when bf16 tensors should be cast to f16 at load time. /// /// All Apple Silicon GPUs (M1–M5) lack native bf16 ALU hardware. Metal's @@ -2578,6 +2668,53 @@ mod tests { ); } + /// Issue #636 / security-review follow-up (c1d1c33): the f16-fragile guard + /// must catch known-fragile `model_type`s, a top-level softcap/logit_scale + /// key, and the same softcap key nested under `text_config` (the + /// multimodal-checkpoint blind spot the review flagged), while leaving a + /// plain healthy family unaffected. + #[test] + fn is_f16_fragile_family_covers_model_type_and_softcap_cases() { + let cases: &[(Value, bool, &str)] = &[ + ( + serde_json::json!({ "model_type": "gemma2" }), + true, + "gemma model_type is on the known-fragile substring list", + ), + ( + serde_json::json!({ + "model_type": "llama", + "final_logit_softcapping": 30.0 + }), + true, + "top-level softcap key implies wide dynamic range", + ), + ( + serde_json::json!({ + "model_type": "llama", + "text_config": { + "final_logit_softcapping": 30.0 + } + }), + true, + "softcap nested under text_config must not slip past the guard", + ), + ( + serde_json::json!({ "model_type": "llama" }), + false, + "plain llama with no softcap/logit_scale is not fragile", + ), + ]; + + for (config, expected, label) in cases { + assert_eq!( + is_f16_fragile_family(config), + *expected, + "{label}: {config:?}" + ); + } + } + #[test] fn strip_gemma4_kv_shared_works_on_quantized_suffixes() { // Quantized checkpoints store k_proj/v_proj/k_norm as three separate