perf(cuda): single-dtype decode graph to eliminate per-token AsType

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.

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. |