perf(cuda): single-dtype decode graph to eliminate per-token AsType#732
Merged
Conversation
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
inureyes
added
type:performance
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
… 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.
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
5 tasks
This was referenced Jul 12, 2026
inureyes
added a commit
that referenced
this pull request
Jul 12, 2026
…measurement (#757) ## Summary Closes #755 with the post-reboot re-measurement it gated on, and a verdict: **none of the sweep drops is a code regression**. The gemma-4-26b pair was depressed by the stale ~5.5-day host/driver state and recovers decisively on a fresh boot; glm4-flash and mamba2-130m were never regressed, their sweep readings are single draws from a bimodal host-level throughput distribution whose fast mode reproduces the 0.3.1 records; and the same session's SSM-cluster re-verification shows the 1.4-3.4x "anomaly" is the intended #727 fused-SSM-kernel improvement, misdiagnosed as environmental in the sweep notes. Doc rows for all 13 re-measured models now carry the post-reboot values, and the 21 committed runs live in `benchmarks/cuda_gb10_2026-07-12_postreboot_single_*.csv`. Host: fresh boot 2026-07-12 19:25 (uptime 7 min at first check), mlxcel 0.4.0-rc.1 at `fcb4c20`, MLX pin `57c66cac` (0.32.1), incremental `cargo build --release --features cuda` rebuilt clean before measuring. All script runs used `--cooldown 30`. ## Measurements ### #755 subjects and controls (decode tok/s) | Model | 2026-06-17 (0.3.1) | 2026-07-12 sweep | Post-reboot | Verdict | |---|---|---|---|---| | gemma-4-26b-a4b-it-4bit | 58.59 | 50.19 | **59.88** (n=3: 59.57-60.07) | recovered, environmental | | gemma-4-26b-a4b-it-qat-4bit | 50.33 | 45.29 | **53.68** (n=3: 53.48-53.71) | recovered, environmental | | glm4-flash-4bit | 53.33 | 45.72 | **47.42** (n=3 templated: 46.42-50.07) | within noise floor (below) | | mamba2-130m | 181.05 | 162.23 | **202.98** (n=3: 147.98-204.77) | within noise floor | | qwen3-30b-a3b-4bit (control) | 90.70 | 92.41 | 94.39 | stable | | lfm2-8b-a1b-4bit (control) | 157.73 | 161.87 | 165.53 | stable | ### SSM-cluster re-verification (decode tok/s) | Model | 2026-06-17 | 2026-07-09 single | 2026-07-12 sweep | Post-reboot | |---|---|---|---|---| | granite-4.0-h-350m-4bit | 64.00 | 86.60 | 259.69 | 171.22 | | granite-4.0-h-tiny-4bit | 33.84 | - | 100.28 | 101.37 | | falcon-h1-tiny-90m-instruct-4bit | 102.99 | 110.42 | 413.00 | 354.43 | | nemotron-h-30b-4bit | 40.32 | - | 79.94 | 87.41 | | nemotron-nas-30b-4bit | 37.33 | - | 82.72 | 85.91 | | nemotron-3-nano-omni-30b (reasoning) | 38.45 | - | 80.86 | 82.88 | | plamo-2-1b | 34.36 | 35.14 | 44.54 | 46.84 | The cluster's gains persist on a fresh host, so they are release numbers, not an artifact. The attribution is a timeline fact the sweep notes missed: the low 2026-07-09 singles ran at 15:03, and #727 (fused single-token SSM decode kernel, CUDA port; its own PR measured granite-350m at 4.5x) merged the next day at 16:45 on 2026-07-10. The "no SSM-related code has landed since" premise behind the environmental explanation was simply false, and this PR corrects it on the record. ## Attribution evidence for the non-recovering pair glm4-flash decode over twelve **identical** 100-token greedy runs on the freshly booted host (same binary, same prompt, `--no-chat-template`, identical generated tokens): 39.23, 54.80, 42.33, 52.33, 51.48, 52.86, 40.51, 42.12, 41.28, 49.29, 52.24, 42.54. The distribution is bimodal (a ~41 mode and a ~52.5 mode), and the 0.3.1 record (53.33) sits inside the fast mode. mamba2-130m shows the same flapping (147.98, then 202.98 / 204.77 / 208.25 back-to-back, vs the 181.05 record). There is no stable deficit to attribute; a single sweep run draws one sample from this distribution. The code suspects named in the issue are each ruled out: - **#740 multirow qmv**: `MLXCEL_QMV_MULTIROW=0` reads 49.55, inside the baseline envelope; by design the path keeps `M*B == 1` classic decode on the stock kernel, so B=1 benches never exercise it. - **#727 fused SSM kernel**: `MLXCEL_SSM_CUDA_KERNEL=0` on mamba2-130m reads 208.25 (unchanged); pure mamba2 never uses the fused kernel, exactly as the #727 PR stated. - **#732 single-dtype decode graph**: shipped trace tooling (`MLXCEL_TRACE_ASTYPE`) plus a default-off opt-in normalization (`MLXCEL_CUDA_F16_NORMALIZE`); nothing changes by default. The slow mode is not SM clock capping: sampling `clocks.sm` at 2 Hz during runs caught slow runs (41.28, 42.54 tok/s) at a pinned 2411-2424 MHz, the same clocks as fast runs. The CPU governor is `performance` at max frequency. The host-level mechanism behind the two modes remains unidentified, but the fast mode reproducing the 0.3.1 number rules out a code-level regression, which is what #755 needed to establish. Two further observations recorded in the docs: - The 26B gemma MoE repeats within ±0.5% while the launch-bound small models flap by up to ±25%, so the noise floor is model-class-dependent: single-run deltas on small dense/SSM checkpoints and small MoEs should not be read as regressions without repeats. - glm4-flash's templated greedy output shortened from 100 tokens (0.3.1) to 18 (rc.1), a different greedy continuation rather than a failure, so its sweep-to-sweep decode averages additionally stopped being length-comparable. ## Changes - `benchmarks/cuda_gb10_2026-07-12_postreboot_single_*.csv` (13 files, 21 runs): the post-reboot singles, including the n=3 repeat sets for both gemma variants, glm4-flash, and mamba2-130m. Env-modified A/B runs are excluded from the CSVs and reported here instead. - `docs/benchmark_results/model_tests_gb10.md`: the 13 model rows now carry the post-reboot values (medians where n=3); the SSM-cluster notable-changes bullet is rewritten with the #727 attribution and the correction of the earlier environmental claim; the moderate-drops bullet is rewritten as resolved with the evidence above and the noise-floor guidance. - `docs/benchmark_results/model_tests.md`: Nemotron-H-30B cross-hardware row updated to 87.41 and its footnote now carries the #727/#755 attribution instead of the re-verify caveat. ## Acceptance criteria (from #755) - [x] Post-reboot single-model measurements recorded for the three subjects, mamba2-130m, and both controls (singles CSVs committed; numbers above) - [x] Verdict recorded: environmental for the gemma pair; no-regression (host run-to-run variance) for glm4-flash and mamba2-130m; the SSM-cluster side settled as a real #727 improvement - [x] Doc rows and the unattributed-drop notes updated to the post-reboot values, with this issue referenced - [x] Kill-switch A/B and clock-sampling evidence recorded for the persisting readings (no code fix needed; nothing to fix) - [x] The notes in `docs/benchmark_results/model_tests_gb10.md` reference #755 ## Test plan - [x] Incremental `cargo build --release --features cuda` clean before measuring (no source changes in this PR; docs and benchmark data only) - [x] 13 post-reboot singles + 8 repeat/variance runs + 3 kill-switch A/B runs + 6 clock-sampled runs on GB10 - [x] Doc tables cross-checked against the committed CSV rows
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The CUDA single-dtype decode graph is the objective of issue #636: eliminate the per-token AsType (dtype-conversion) nodes so every weight, cache, activation, and sampler tensor sits in one dtype. This PR ships the AsType-per-token counter that is the "done" metric, inventories the remaining conversion sources on CUDA, and removes the one reducible source (the temperature-sampler chain). The headline result: the three inventory models already emit 0 AsType per greedy decode step on CUDA, because the merged
patches-cuda/dtype.cppbf16 promotion patch 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.What changed
count_astype_nodes_pair/astype_breakdown_pairwalks the unevaluated(token, logprob)decode graph, countsAsTypenodes, and produces a persrc->dstdtype breakdown. Exposed through the cxx bridge (lib.rs,mlx_cxx_bridge.{h,cpp}) and wired into the decode loop behindMLXCEL_TRACE_ASTYPE(generate.rs). Traversal only, no extra eval, entirely skipped when unset. Cross-checked againstMLXCEL_EXPORT_DECODE_DOT(grep -c AsType), which agrees at 0.fused_sample,top_p_filter, and the compiled min-p kernel built their temperature / sentinel / threshold scalars as f32, which the CUDA bf16 promotion patch then cast to the logit dtype every decode step. Each scalar is now built in the logit dtype on non-Metal backends only, gated on a runtime!metal::is_available()flag (the kernel-dispatch pattern already used in this file), keeping the CUDA sampler chain single-dtype. Metal keeps the bare f32 scalars exactly as main, so its temperature-sampling numerics are untouched per the issue AC. On CUDA the change 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.MLXCEL_CUDA_F16_NORMALIZEadds an opt-in, CUDA-only bf16 -> f16 cast of non-quantized weights (sanitize.rs), with a conservative f16-fragile exception list (is_f16_fragile_family) and a global env opt-out. Default off (see rationale below).docs/environment-variables.mddocumentsMLXCEL_TRACE_ASTYPEandMLXCEL_CUDA_F16_NORMALIZE;docs/benchmark_results/single-dtype-decode-astype-gb10-2026-07-10.mdrecords the inventory and before/after counts.Inventory (GB10, CUDA, SM 12.1)
patches-cuda/dtype.cpp(bf16 + fp32 -> bf16) and the consistent-dtype quant path. Nothing to reduce.f32 -> bf16)random::categoricaluniform drawu32 -> f32)f16 <-> f32, 2/layer)Before / after AsType per decode step
The greedy target reduction is vacuously >= 90% (0 before, 0 after): the objective was already met on CUDA by prior merged work. The one path with a real, reducible count (temperature sampling) drops 75%, with the residual
u32 -> f32intrinsic torandom::categorical.Quality: 40-token greedy parity (env off vs on)
Decode tok/s A/B (
mlxcel-bench-decode --prompt x --prompt-tokens 512 --max-tokens 128, 2 runs)Why CUDA f16 normalization is opt-in
The single-dtype objective is already met in bf16 on CUDA, and CUDA has native bf16 compute, so bf16 -> f16 yields no AsType reduction and no throughput gain while narrowing dynamic range. The load-time f16 path is therefore off by default, with the conservative f16-fragile exception list (gemma, cohere/command-r, apertus, gpt-oss, and any softcap/
logit_scaleconfig) keeping bf16 even when enabled. Quantized checkpoints are unaffected, so leaving it unset keeps every bf16 model available. Metal/Apple Silicon numerics are untouched: the always-on Apple policy is unchanged and the new path is gated on the CUDA backend.Test plan
cargo build --release --features cudaMLXCEL_TRACE_ASTYPEbefore/after on the three inventory models (greedy) and the temperature-sampler paths; cross-checked vsMLXCEL_EXPORT_DECODE_DOTcargo test --features cuda -p mlxcel-core --lib sampling::(46 passed)Note: the full representative sweep, perplexity spot check, and full test suite are left to the orchestrator.
Closes #636