From 8ca1a53bdc431cf218deabcd7b6b2671e5ecafb4 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 10 Jul 2026 21:17:38 +0900 Subject: [PATCH 1/4] perf(cuda): single-dtype decode graph to eliminate per-token AsType Add the AsType-per-token counter that is the "done" metric for the CUDA single-dtype decode graph (issue #636), inventory the remaining conversion sources, and remove the one reducible source on the sampling path. The three inventory models already emit 0 AsType per greedy decode step on CUDA: the merged patches-cuda/dtype.cpp bf16 promotion patch (bf16 + fp32 -> bf16) collapses the diffuse scalar/constant promotions that the moe-decode-gap investigation counted at ~773/token on Apple/Metal. The counter makes that state measurable and guards it against regression. Tooling: new C++ traversal count_astype_nodes_pair / astype_breakdown_pair walk the unevaluated (token, logprob) decode graph and count AsType nodes (and a per src->dst dtype breakdown), exposed through the cxx bridge and wired into the decode loop behind MLXCEL_TRACE_ASTYPE. Traversal only, no extra eval, and entirely skipped when the env var is unset. Cross-checked against MLXCEL_EXPORT_DECODE_DOT (grep -c AsType), which agrees at 0. Sampler hygiene: the fused sampler built its temperature, top-k/top-p/min-p sentinel, and threshold scalars as f32, which the CUDA bf16 promotion patch then cast to the logit dtype every decode step. Building each scalar in the logit dtype keeps the chain single-dtype: bit-identical for bf16 logits (the old path cast the scalar to bf16 anyway) and avoids upcasting the whole vocab tensor to f32 for f16 logits. Temperature-sampling AsType drops from 4 to 1 (the residual u32 -> f32 is intrinsic to random::categorical). Greedy uses the untouched argmax path, so the three inventory models stay at 0. Load-time normalization: MLXCEL_CUDA_F16_NORMALIZE adds an opt-in CUDA-only bf16 -> f16 cast of non-quantized weights for fixed-topology / CUDA-graph-reuse experiments, with a conservative f16-fragile exception list (gemma, cohere/command-r, apertus, gpt-oss, and any softcap or logit_scale config) that keeps bf16 regardless. Off by default because the CUDA decode graph is already single-dtype in bf16 and CUDA has native bf16 ALUs, so f16 gives no measured throughput gain (qwen2.5-0.5b-bf16 decode 208 vs 207 tok/s) while narrowing dynamic range. Quantized checkpoints are unaffected, so leaving it unset keeps every bf16 model available. 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. Validated on GB10 (SM 12.1): three inventory models at 0 AsType greedy before and after; 40-token greedy parity byte-identical for llama-3.1-8b-4bit, qwen3-30b-a3b-4bit (env no-op) and llama-3.1-8b-bf16 (bf16 vs f16); env-on/off decode A/B within noise on qwen3-30b-a3b-4bit and qwen2.5-0.5b-bf16; 46 mlxcel-core sampling unit tests pass. Counts and tables in docs/benchmark_results/single-dtype-decode-astype-gb10-2026-07-10.md. Refs #636 --- ...gle-dtype-decode-astype-gb10-2026-07-10.md | 52 ++++++++ docs/environment-variables.md | 2 + src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 125 +++++++++++++++++- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 6 + src/lib/mlxcel-core/src/generate.rs | 18 +++ src/lib/mlxcel-core/src/lib.rs | 9 ++ src/models/sanitize.rs | 86 +++++++++++- 7 files changed, 291 insertions(+), 7 deletions(-) create mode 100644 docs/benchmark_results/single-dtype-decode-astype-gb10-2026-07-10.md 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..1f2b15b0 --- /dev/null +++ b/docs/benchmark_results/single-dtype-decode-astype-gb10-2026-07-10.md @@ -0,0 +1,52 @@ +# 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. + +## 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: scalars now built in the logit dtype so the CUDA promotion patch inserts no per-step conversion. Bit-identical for bf16 logits. | +| `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..613a55d5 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -3,6 +3,13 @@ #include "mlx_cxx_internal.h" +#include "mlx/primitives.h" + +#include +#include +#include +#include + namespace mlx_cxx { using namespace mlx::core; @@ -4346,6 +4353,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); @@ -4367,10 +4459,17 @@ namespace { 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)); + // Build comparison and fill scalars in the working dtype (issue #636): + // a bare `array(f32)` against a bf16/f16 tensor forces the CUDA bf16 + // promotion patch to insert an AsType conversion on the scalar every + // decode step. `probs`/`shifted_cum` follow softmax's output dtype. + auto mask = mlx::core::less_equal( + shifted_cum, mlx::core::array(top_p, shifted_cum.dtype())); 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())); + mask, + sorted_logits, + mlx::core::array(std::numeric_limits::lowest(), x.dtype())); auto unsort_indices = mlx::core::argsort(sorted_indices, -1); return mlx::core::take_along_axis(filtered_sorted, unsort_indices, -1); } @@ -4387,8 +4486,10 @@ namespace { 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); + // Fill sentinel in x's dtype so no per-step AsType is inserted on + // the scalar under the CUDA bf16 promotion patch (issue #636). return {mlx::core::where(mask, x, - mlx::core::array(std::numeric_limits::lowest()))}; + mlx::core::array(std::numeric_limits::lowest(), x.dtype()))}; }; return mlx::core::compile(fn, true); } @@ -4417,9 +4518,19 @@ std::unique_ptr fused_sample( return std::make_unique(mlx::core::argmax(x, -1, false)); } + // Build all sampler scalars in the logit dtype (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: for + // bf16 logits the math is bit-identical (the old path cast the f32 scalar + // to bf16 anyway); for f16 logits it also avoids upcasting the whole + // vocab tensor to f32 through the unpatched f16+f32 rule. + const auto logit_dtype = x.dtype(); + // Temperature scaling if (temperature > 0.0f && temperature != 1.0f) { - x = x / mlx::core::array(temperature); + x = x / mlx::core::array(temperature, logit_dtype); } // Top-k filtering: keep only the k highest-probability tokens @@ -4438,7 +4549,9 @@ 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, + mlx::core::array(std::numeric_limits::lowest(), logit_dtype)); } // Top-p (nucleus) filtering @@ -4449,7 +4562,7 @@ std::unique_ptr fused_sample( // 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, mlx::core::array(min_p, logit_dtype)}); 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..2a3edd71 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,80 @@ 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. The list is intentionally broad: when unsure, keep bf16. The +/// goal is a single-dtype graph for the healthy majority, not full coverage. +fn is_f16_fragile_family(config: &Value) -> bool { + // Softcap / logit-scaling config keys imply wide dynamic range. + for key in [ + "attn_logit_softcapping", + "final_logit_softcapping", + "logit_softcapping", + "logits_soft_cap", + "logit_scale", + ] { + if config + .get(key) + .and_then(Value::as_f64) + .is_some_and(|v| v != 0.0 && v != 1.0) + { + 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 From a8f9922acb49cfecae072f2010590f7f3b5a0c6c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 10 Jul 2026 21:39:58 +0900 Subject: [PATCH 2/4] fix(cuda): gate sampler scalar dtype on non-Metal to keep Metal numerics The issue #636 sampler-chain change built the fused sampler's temperature, top-k/top-p/min-p sentinel, and threshold scalars in the logit dtype unconditionally, which altered Metal temperature-sampling numerics. On Apple Silicon bf16 weights load as f16, logits are f16, and patches-cuda/dtype.cpp deliberately leaves f16+f32 -> f32 unpatched, so main's Metal path upcasts the sampler chain to an f32 softmax via the first bare f32 scalar. Building the scalars in f16 kept the Metal chain in f16, which #636 forbids (Metal numerics must be untouched). Gate the scalar-dtype choice on a runtime `!mlx::core::metal::is_available()` flag (`single_dtype_scalars`), matching the kernel-dispatch pattern already used in this file. On Metal the scalars stay bare `array(f32)` exactly as main; on CUDA (and other non-Metal backends) they are built in the logit dtype so the bf16 promotion patch inserts no per-step AsType. `fused_sample` hoists the flag once and threads it into `top_p_filter` and the compiled min-p kernel (captured at build time, since the backend is process-constant). On CUDA this is bit-identical for bf16 logits (the old path cast the f32 scalar to bf16 anyway); a natively-f16 CUDA checkpoint now runs temperature sampling in f16, which is the intended single-dtype behavior and within fp16 tolerance for the stochastic categorical draw, now noted in the benchmark note. Verified: `cargo check --features cuda --lib --tests` clean; `cargo test --release --features cuda -p mlxcel-core --lib sampling::` 46 passed; CUDA MLXCEL_TRACE_ASTYPE still reports 1 on the qwen3-8b-4bit temp+top-p path (the intrinsic u32->f32) and 0 on greedy. Refs #636 --- ...gle-dtype-decode-astype-gb10-2026-07-10.md | 4 +- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 74 +++++++++++-------- 2 files changed, 48 insertions(+), 30 deletions(-) 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 index 1f2b15b0..ca6c1988 100644 --- 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 @@ -20,13 +20,15 @@ The only remaining reducible AsType source on CUDA is the temperature-sampler ch 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: scalars now built in the logit dtype so the CUDA promotion patch inserts no per-step conversion. Bit-identical for bf16 logits. | +| 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. | diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index 613a55d5..1716fcad 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -4453,23 +4453,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; - // Build comparison and fill scalars in the working dtype (issue #636): - // a bare `array(f32)` against a bf16/f16 tensor forces the CUDA bf16 - // promotion patch to insert an AsType conversion on the scalar every - // decode step. `probs`/`shifted_cum` follow softmax's output dtype. - auto mask = mlx::core::less_equal( - shifted_cum, mlx::core::array(top_p, shifted_cum.dtype())); + 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(), x.dtype())); + 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); } @@ -4479,17 +4482,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); - // Fill sentinel in x's dtype so no per-step AsType is inserted on - // the scalar under the CUDA bf16 promotion patch (issue #636). - return {mlx::core::where(mask, x, - mlx::core::array(std::numeric_limits::lowest(), x.dtype()))}; + 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); } @@ -4518,19 +4527,27 @@ std::unique_ptr fused_sample( return std::make_unique(mlx::core::argmax(x, -1, false)); } - // Build all sampler scalars in the logit dtype (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: for - // bf16 logits the math is bit-identical (the old path cast the f32 scalar - // to bf16 anyway); for f16 logits it also avoids upcasting the whole - // vocab tensor to f32 through the unpatched f16+f32 rule. + // 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` mirrors the runtime + // `!metal::is_available()` gate used by the kernel dispatches 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, logit_dtype); + x = x / sampler_scalar(temperature); } // Top-k filtering: keep only the k highest-probability tokens @@ -4550,19 +4567,18 @@ std::unique_ptr fused_sample( 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(), logit_dtype)); + 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, logit_dtype)}); + auto result = compiled_fn({x, sampler_scalar(min_p)}); x = std::move(result[0]); } From c1d1c3301b79e325458f60d49cc8e5223e629dc7 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 10 Jul 2026 21:54:06 +0900 Subject: [PATCH 3/4] fix(models): check nested text_config for softcap keys in f16-fragile guard Security-review follow-ups: the softcap/logit-scale sniff in is_f16_fragile_family only inspected top-level config keys while the model_type detection already fell back to text_config, so a multimodal checkpoint carrying softcapping under text_config could slip past the guard and get f16-normalized. The key loop now applies the same text_config fallback. Also align the doc comment with actual behavior (unknown families are normalized, which is why the path is opt-in and default-off), fix a comment that claimed the metal::is_available gate already existed in this file, and add the missing cstdint include. --- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 6 ++++-- src/models/sanitize.rs | 21 +++++++++++++-------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index 1716fcad..ef73b2af 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -5,6 +5,7 @@ #include "mlx/primitives.h" +#include #include #include #include @@ -4536,8 +4537,9 @@ std::unique_ptr fused_sample( // 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` mirrors the runtime - // `!metal::is_available()` gate used by the kernel dispatches in this file. + // 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) { diff --git a/src/models/sanitize.rs b/src/models/sanitize.rs index 2a3edd71..c48fd143 100644 --- a/src/models/sanitize.rs +++ b/src/models/sanitize.rs @@ -1730,10 +1730,14 @@ fn env_flag_enabled(name: &str) -> bool { /// /// f16's narrow dynamic range (max ~65504) clips wide-range activations and /// softcapped logits, so these families keep bf16 even when normalization is -/// requested. The list is intentionally broad: when unsure, keep bf16. The -/// goal is a single-dtype graph for the healthy majority, not full coverage. +/// 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. + // 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", @@ -1741,11 +1745,12 @@ fn is_f16_fragile_family(config: &Value) -> bool { "logits_soft_cap", "logit_scale", ] { - if config - .get(key) - .and_then(Value::as_f64) - .is_some_and(|v| v != 0.0 && v != 1.0) - { + 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; } } From 2e3a712a2f9fc66ae42501b8578052a493e78022 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 10 Jul 2026 22:12:36 +0900 Subject: [PATCH 4/4] test(models): cover is_f16_fragile_family text_config softcap fallback Add a table-driven unit test for the CUDA f16-normalization fragile-family guard (issue #636), covering a known-fragile model_type (gemma2), a top-level softcap key, the c1d1c33 text_config-nested softcap fallback that the security review flagged as a multimodal blind spot, and a plain non-fragile llama config. Also runs cargo fmt on the file, which reflowed the pre-existing FRAGILE_SUBSTRINGS const array to match rustfmt's line-width rule. Validation: - cargo check --features cuda --lib --tests -p mlxcel: clean - cargo test --release --features cuda -p mlxcel --lib is_f16_fragile_family -- --test-threads=1: 1 passed - cargo fmt --check on touched Rust files: clean Refs #636 --- src/models/sanitize.rs | 52 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 50 insertions(+), 2 deletions(-) diff --git a/src/models/sanitize.rs b/src/models/sanitize.rs index c48fd143..22bb376c 100644 --- a/src/models/sanitize.rs +++ b/src/models/sanitize.rs @@ -1766,8 +1766,9 @@ fn is_f16_fragile_family(config: &Value) -> bool { .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"]; + const FRAGILE_SUBSTRINGS: &[&str] = &[ + "gemma", "cohere", "command", "apertus", "gpt_oss", "gpt-oss", + ]; FRAGILE_SUBSTRINGS .iter() .any(|needle| model_type.contains(needle)) @@ -2667,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