perf(cuda): single-dtype decode graph: eliminate per-token AsType conversions #636

Part of #623. Related to #629 (MoE) but scoped to dtype hygiene across all architectures.

Context

The MoE decode-gap investigation (docs/benchmark_results/moe-decode-gap-investigation.md, lines ~38-117) measured decode as ~95% GPU-bound with effective bandwidth at 16-20% of GEMV peak, and counted ~773 diffuse AsType (dtype conversion) ops per generated token with no single fusible lever. Two mitigations already exist:

• src/lib/mlx-cpp/patches-cuda/dtype.cpp changes bf16 type-promotion rules (bf16+fp32 -> bf16) to eliminate copy_v conversion kernels, whose header comment cites ~58% GPU-time overhead on 9B models before the patch.
• convert_quant_scales_bf16_to_f16 promotes quant scales/biases to f16 at load on CUDA so quantized ops see consistent dtypes.

The investigation's standing recommendation is a single-dtype decode graph: every weight, cache, activation, and sampler tensor in one dtype (f16 on CUDA) so no conversion kernels appear at all. Beyond the direct kernel savings, a stable single-dtype graph has fixed topology, which improves MLX's internal CUDA-graph executable reuse (address/topology churn is the known blocker, ml-explore/mlx#2358).

Scope

Reduce per-token AsType count by >= 90% on representative dense and MoE models on CUDA, without quality regressions. Metal behavior unchanged.

Implementation plan

1. Tooling first: extend the existing dtype tracing (MLXCEL_TRACE_DTYPE, src/lib/mlxcel-core/src/generate.rs:1103) or add a graph-dump count (via MLXCEL_EXPORT_DECODE_DOT, generate.rs:1203-1207) that reports AsType-node count per decode step. This is the metric for done.
2. Inventory the top sources on llama-3.1-8b-4bit, qwen3-8b-4bit, and qwen3-30b-a3b-4bit: typical culprits are bf16 checkpoint weights vs f16 activations, fp32 norm accumulations re-cast per call, RoPE frequency tensors in fp32, sampler ops (softmax/categorical) upcasting, and mask tensors materialized in a different dtype each step.
3. Load-time normalization on CUDA: cast weights/norm scales/embeddings to f16 once at load (extending the convert_quant_scales_bf16_to_f16 precedent) behind a backend-conditional policy. Maintain an explicit exception list for known f16-fragile models (e.g. softcap-dependent or wide-dynamic-range models; gemma family norms) which keep bf16 and accept the conversions, driven by a per-model flag in the model implementations.
4. Graph hygiene: hoist per-step constant tensors (masks, rope freqs, scale scalars) to cached buffers with fixed dtype; ensure the sampler chain (src/lib/mlxcel-core/src/sampling.rs fused path) does not round-trip dtypes.
5. Quality gate: run the model parity suite on the touched families plus a perplexity spot check on one bf16-native checkpoint (llama-3.1-8b-bf16) comparing bf16 vs f16-normalized outputs; document the greedy divergence profile.
6. Measure decode tok/s deltas on the three inventory models plus one small model (qwen2.5-0.5b-bf16, most overhead-sensitive).

Acceptance criteria

• AsType-per-token counter tool merged and documented (env-gated, zero cost when off).
• >= 90% AsType reduction on the three inventory models on CUDA; counts before/after committed to a docs/benchmark_results/ note.
• Decode tok/s: measurable improvement on qwen3-30b-a3b-4bit (the documented 16-20%-of-peak case) and no regression anywhere in a representative sweep (scripts/bench_representative.sh).
• Quality: parity suite passes; f16-fragile exception list documented; bf16 models remain available via config/env opt-out.
• Metal numerics untouched (policy is CUDA-conditional); cargo test --features cuda passes.

Validation

cargo build --release --features cuda
MLXCEL_TRACE_DTYPE=1 ./target/release/mlxcel --model ./models/qwen3-30b-a3b-4bit generate "..."   # before/after AsType counts
scripts/bench_representative.sh

References

• Investigation: docs/benchmark_results/moe-decode-gap-investigation.md:38-117.
• Existing dtype patches: src/lib/mlx-cpp/patches-cuda/dtype.cpp; scale conversion convert_quant_scales_bf16_to_f16 (grep in src/lib/mlxcel-core/).
• Upstream CUDA-graph reuse context: [CUDA] Buffer cache and CUDA graph ml-explore/mlx#2358.