fix(tts): cast bf16 tables explicitly where the CUDA overlay demotes

TECHNICAL_REPORTS/2093-cuda-bf16-promotion-tests-20261002.en.md

@@ -0,0 +1,55 @@
+# Technical Report: PR #2093 - Cast bf16 tables explicitly where the CUDA overlay demotes (issue #2087)
+
+**Date**: 2026-10-02
+
+**Status**: Implemented and validated on GB10 (`--features cuda`) and on a CPU-only Linux build; pending merge.
+
+**Languages**: Rust (production code in the Nemotron VoiceChat TTS head, tests)
+
+**Risk Level**: Low (two exact `astype` casts in the TTS path, no-ops on non-CUDA builds; the rest is test code)
+
+## Executive Summary
+
+Gating PR #2077 on GB10 left six lib-test failures that also failed on the base. Five of them were one cause: `src/lib/mlx-cpp/patches-cuda/dtype.cpp`, mlxcel's own CUDA overlay of MLX's promotion table, resolves bf16 with f32 to bf16 instead of upstream's f32. It is not an MLX bug at the pin and not an `mlxcel-core` wrapper: the overlay is deliberate (issue #636, the single-dtype bf16 decode graph) and stays. The tests and two production sites had assumed upstream promotion. The sixth failure, `gelu_approx_matches_mlx_nn_bit_for_bit`, was already fixed on main by #2080 and passes on CUDA unchanged.
+
+## 1. Root Cause
+
+`dtype::BFLOAT16 == 12` and `dtype::FLOAT32 == 10`, so the `left: 12, right: 10` assertions said a mixed op returned bf16. The overlay header states the change: the bf16 row's f32 column and the f32 row's bf16 column both read `bfloat16`. The table is compiled into the MLX library, so the rule holds for every device of a CUDA build, including `MLXCEL_DEVICE=cpu`. ROCm compiles upstream `dtype.cpp` and Metal never sees `patches-cuda/`, so only the `cuda` feature changes the rule.
+
+Two consequences in the Nemotron VoiceChat TTS path were real output-dtype defects on CUDA, not only test artifacts:
+
+- `RvqCodebooks::depthsum_embedding` documents an f32 sum, but added gathered bf16 codebook rows to an f32 `zeros`, so on CUDA the sum came back bf16.
+- `MogHead::infer` gathered `proj_mus` and `low_mat` slabs (stored bf16 by design, see `promotes_to_f32`) and matmul'd them against the f32 activation, so `mu`, and with it the sampled latent, came back bf16.
+
+## 2. Changes
+
+- **Probe test** `audio::f32_weights::tests::mixed_bf16_f32_promotion_follows_the_build`: asserts the dtype of `add` (both operand orders), `matmul` (both orders) and `fast_rms_norm` for an f32 activation against a bf16 operand. It expects f32 without the `cuda` feature and bf16 with it, so a change to either rule fails loudly. The issue asked for f32 on both; that cannot hold while the overlay stays, so the probe pins the contract that does hold and the production code casts explicitly.
+- **`rvq.rs`**: the gathered rows are cast to f32 before the add.
+- **`mog_head.rs`**: the gathered `mus` and `low` slabs are cast to the activation dtype before their matmuls. bf16 to f32 is exact and is the same `astype` upstream promotion inserts, so non-CUDA output is bit-identical, and the cast is a no-op when the dtypes already match.
+- **`f32_weights.rs` tests**: the "promoted" reference is MLX's implicit promotion off CUDA and an explicit in-graph `astype(weight, float32)` on CUDA (`as_promoted`). The module doc now says that on CUDA the load-time precast is required, not only faster.
+- **`fastconformer/tests.rs`**: the reference encoder and projection are built from the bf16 weights cast to f32, which is what upstream promotion computes, instead of relying on implicit promotion.
+- **`tts_tests.rs`**: `mog_infer_shapes_are_finite_with_guidance` loads its weights the way `RvqEarTtsModel` does (`promoted_subset` with `promotes_to_f32`), so the test exercises the production load path rather than an all-bf16 head that production never builds.
+
+## 3. The gelu Test
+
+`gelu_approx_matches_mlx_nn_bit_for_bit` hardcoded values measured on Metal. #2080 (`5486e404`, merged 2026-09-30) replaced them with a same-device transcription of `mlx.nn.gelu_approx` compared bit for bit over 4097 points. The 1-ulp gap was per-backend `tanh`/`power` kernel rounding. Neither side is compiled, there is no matmul, and the graph is element-wise, so fusion, FMA contraction across ops and TF32 are not involved. It passes on CUDA and on the CPU build with no change here.
+
+## 4. Validation
+
+GB10, `--features cuda --release`, each run under `gpu-lock` with `--test-threads=1`:
+
+| Filter | Result |
+|---|---|
+| `audio::f32_weights` (incl. probe) | 4 passed |
+| `audio::fastconformer` | 14 passed |
+| `models::nemotron_voicechat::tts` | 13 passed |
+| `models::gemma3_backbone` (incl. gelu) | 5 passed |
+
+With the two production casts reverted and the tests kept, `depthsum_ignores_mask_index` and `mog_infer_shapes_are_finite_with_guidance` fail on CUDA, so the casts are what makes them pass.
+
+CPU-only Linux build (no GPU feature, separate target directory): the same four filters pass with the same counts (4, 14, 13, 5), and there the probe asserts f32. `MLXCEL_DEVICE=cpu` on the CUDA build was not used as a substitute, because the overlay is compiled into the library and applies on every device of that build.
+
+## 5. Open Items
+
+- **The TTS backbone still demotes on CUDA.** `RvqEarTtsModel` keeps the Gemma norm weights (`1 + w` built in the stored dtype), `bos_emb`, `null_emb` and `audio_prompt_projection_w` as stored bf16, and each meets the f32 stream. Off CUDA they promote it to f32 as the reference does; on CUDA, by the rule the probe pins for `add`, `matmul` and `fast_rms_norm`, they demote it to bf16. The fixes here make the code embeddings and the MoG head f32, but the backbone between them is still bf16 on CUDA. Confirming the impact needs a real Nemotron VoiceChat checkpoint, which this host does not have, so it is left for a follow-up.
+- **The plain Linux CPU lib-test build does not link on main.** `src/lib/mlx-cpp/turbo/kv_inplace_write.cpp` (#1961) references `mlx::core::copy_gpu_inplace`, which a no-GPU MLX build does not define, so `cargo test --release --lib` fails at the link step before running any test. The CPU results above were produced with a linker wrapper that ignores that one unresolved symbol (`-Wl,--unresolved-symbols=ignore-in-object-files`); the symbol is only reached on a GPU stream.

TECHNICAL_REPORTS/2093-cuda-bf16-promotion-tests-20261002.ko.md

@@ -0,0 +1,55 @@
+# 기술 보고서: PR #2093 - CUDA 오버레이가 강등하는 곳에서 bf16 테이블을 명시적으로 캐스트 (issue #2087)
+
+**날짜**: 2026-10-02
+
+**상태**: GB10(`--features cuda`)과 CPU 전용 Linux 빌드에서 구현 및 검증 완료, 머지 대기.
+
+**언어**: Rust (Nemotron VoiceChat TTS 헤드의 프로덕션 코드, 테스트)
+
+**위험도**: 낮음 (TTS 경로에 정확한 `astype` 캐스트 두 개를 추가했고, 비 CUDA 빌드에서는 no-op이다. 나머지는 테스트 코드다)
+
+## 요약
+
+GB10에서 PR #2077을 게이트하던 중 lib 테스트 6건이 실패했고, 베이스에서도 똑같이 실패했다. 이 중 5건은 원인이 하나다. mlxcel이 직접 둔 MLX 승격 테이블의 CUDA 오버레이 `src/lib/mlx-cpp/patches-cuda/dtype.cpp`가 bf16과 f32의 조합을 upstream의 f32가 아니라 bf16으로 결정한다. 고정된 MLX 버전의 버그도 아니고 `mlxcel-core` 래퍼 문제도 아니다. 이 오버레이는 의도된 것(issue #636, 단일 dtype bf16 디코드 그래프)이므로 그대로 둔다. 테스트와 프로덕션 코드 두 곳이 upstream 승격 규칙을 전제하고 있었다. 여섯 번째 실패인 `gelu_approx_matches_mlx_nn_bit_for_bit`는 main에서 이미 #2080으로 수정되었고, 변경 없이 CUDA에서 통과한다.
+
+## 1. 근본 원인
+
+`dtype::BFLOAT16 == 12`, `dtype::FLOAT32 == 10`이므로 `left: 12, right: 10` 단언은 혼합 dtype 연산이 bf16을 돌려줬다는 뜻이다. 오버레이 헤더에 변경 내용이 적혀 있다. bf16 행의 f32 열과 f32 행의 bf16 열이 모두 `bfloat16`이다. 이 테이블은 MLX 라이브러리에 컴파일되므로 CUDA 빌드에서는 `MLXCEL_DEVICE=cpu`를 포함한 모든 디바이스에서 이 규칙이 적용된다. ROCm은 upstream `dtype.cpp`를 컴파일하고 Metal은 `patches-cuda/`를 보지 않으므로, 규칙을 바꾸는 것은 `cuda` 피처뿐이다.
+
+Nemotron VoiceChat TTS 경로의 두 곳은 테스트만의 문제가 아니라 CUDA에서 실제 출력 dtype 결함이었다.
+
+- `RvqCodebooks::depthsum_embedding`은 f32 합을 반환한다고 문서화되어 있지만, gather한 bf16 코드북 행을 f32 `zeros`에 더했으므로 CUDA에서는 합이 bf16이었다.
+- `MogHead::infer`는 `proj_mus`와 `low_mat`(설계상 bf16 그대로 보관, `promotes_to_f32` 참조)에서 gather한 슬랩을 f32 활성값과 matmul했으므로, `mu`와 그에 따라 샘플링된 latent가 bf16이었다.
+
+## 2. 변경 사항
+
+- **프로브 테스트** `audio::f32_weights::tests::mixed_bf16_f32_promotion_follows_the_build`: f32 활성값과 bf16 피연산자에 대해 `add`(양쪽 순서), `matmul`(양쪽 순서), `fast_rms_norm`의 출력 dtype을 단언한다. `cuda` 피처가 없으면 f32, 있으면 bf16을 기대하므로 어느 쪽 규칙이 바뀌어도 바로 실패한다. 이슈는 양쪽 모두 f32를 요구했지만 오버레이를 유지하는 한 성립할 수 없으므로, 프로브는 실제로 성립하는 계약을 고정하고 프로덕션 코드는 명시적으로 캐스트한다.
+- **`rvq.rs`**: gather한 행을 더하기 전에 f32로 캐스트한다.
+- **`mog_head.rs`**: gather한 `mus`와 `low` 슬랩을 matmul 전에 활성값 dtype으로 캐스트한다. bf16에서 f32로의 변환은 정확하고 upstream 승격이 삽입하는 것과 같은 `astype`이므로 비 CUDA 출력은 비트 단위로 같으며, dtype이 이미 같으면 no-op이다.
+- **`f32_weights.rs` 테스트**: "승격된" 기준값은 CUDA가 아니면 MLX의 암묵적 승격, CUDA에서는 그래프 안의 명시적 `astype(weight, float32)`(`as_promoted`)이다. 모듈 문서에 CUDA에서는 로드 시점 사전 캐스트가 단지 빠른 것이 아니라 필수라는 점을 적었다.
+- **`fastconformer/tests.rs`**: 기준 인코더와 프로젝션을 암묵적 승격에 기대지 않고, bf16 가중치를 f32로 캐스트한 값(upstream 승격이 계산하는 값)으로 구성한다.
+- **`tts_tests.rs`**: `mog_infer_shapes_are_finite_with_guidance`가 `RvqEarTtsModel`과 같은 방식(`promoted_subset`과 `promotes_to_f32`)으로 가중치를 로드한다. 프로덕션이 만들지 않는 전부 bf16인 헤드가 아니라 실제 로드 경로를 검증한다.
+
+## 3. gelu 테스트
+
+`gelu_approx_matches_mlx_nn_bit_for_bit`는 Metal에서 측정한 값을 하드코딩하고 있었다. #2080(`5486e404`, 2026-09-30 머지)이 이를 같은 디바이스에서 `mlx.nn.gelu_approx`를 옮겨 쓴 기준과 4097개 점에서 비트 단위로 비교하는 방식으로 바꿨다. 1 ulp 차이는 백엔드별 `tanh`/`power` 커널의 반올림 차이였다. 양쪽 모두 컴파일되지 않고 matmul도 없는 원소별 그래프이므로 fusion, 연산 간 FMA 축약, TF32와는 무관하다. 이 PR의 변경 없이 CUDA와 CPU 빌드에서 통과한다.
+
+## 4. 검증
+
+GB10, `--features cuda --release`, 각 실행은 `gpu-lock` 아래에서 `--test-threads=1`:
+
+| 필터 | 결과 |
+|---|---|
+| `audio::f32_weights` (프로브 포함) | 4 통과 |
+| `audio::fastconformer` | 14 통과 |
+| `models::nemotron_voicechat::tts` | 13 통과 |
+| `models::gemma3_backbone` (gelu 포함) | 5 통과 |
+
+테스트는 그대로 두고 프로덕션 캐스트 두 개만 되돌리면 CUDA에서 `depthsum_ignores_mask_index`와 `mog_infer_shapes_are_finite_with_guidance`가 실패한다. 즉 두 테스트를 통과시키는 것은 이 캐스트다.
+
+CPU 전용 Linux 빌드(GPU 피처 없음, 별도 target 디렉터리): 같은 네 필터가 같은 개수(4, 14, 13, 5)로 통과하며, 여기서는 프로브가 f32를 단언한다. 오버레이가 라이브러리에 컴파일되어 해당 빌드의 모든 디바이스에 적용되므로, CUDA 빌드의 `MLXCEL_DEVICE=cpu`는 대체 검증으로 쓰지 않았다.
+
+## 5. 남은 항목
+
+- **CUDA에서 TTS 백본은 여전히 bf16으로 강등된다.** `RvqEarTtsModel`은 Gemma 노름 가중치(`1 + w`를 저장 dtype으로 계산), `bos_emb`, `null_emb`, `audio_prompt_projection_w`를 bf16 그대로 두며, 각각이 f32 스트림과 만난다. CUDA가 아니면 reference처럼 스트림을 f32로 승격하지만, CUDA에서는 프로브가 `add`, `matmul`, `fast_rms_norm`에 대해 고정한 규칙에 따라 bf16으로 강등한다. 이번 수정으로 코드 임베딩과 MoG 헤드는 f32가 되었지만, 그 사이의 백본은 CUDA에서 여전히 bf16이다. 영향을 확인하려면 실제 Nemotron VoiceChat 체크포인트가 필요한데 이 호스트에는 없으므로 후속 작업으로 남긴다.
+- **main에서 일반 Linux CPU lib 테스트 빌드가 링크되지 않는다.** `src/lib/mlx-cpp/turbo/kv_inplace_write.cpp`(#1961)가 GPU 없는 MLX 빌드에는 정의되지 않은 `mlx::core::copy_gpu_inplace`를 참조하므로, `cargo test --release --lib`는 테스트를 실행하기 전에 링크 단계에서 실패한다. 위 CPU 결과는 이 미해결 심볼 하나를 무시하는 링커 래퍼(`-Wl,--unresolved-symbols=ignore-in-object-files`)로 얻었다. 이 심볼은 GPU 스트림에서만 도달한다.

src/audio/f32_weights.rs

@@ -26,6 +26,14 @@
 //! bit-identical, and each frame skips the cast traffic. The price is twice
 //! the resident size of the promoted weights.
 //!
+//! On CUDA builds the precast is required, not only faster: mlxcel's CUDA
+//! overlay of MLX's promotion table (`src/lib/mlx-cpp/patches-cuda/dtype.cpp`,
+//! kept for the single-dtype bf16 decode graph of issue #636) resolves bf16
+//! with f32 to bf16, so a bf16 weight left as stored would demote the f32
+//! activation stream instead of promoting the weight (issue #2087). Code that
+//! needs f32 from a stored bf16 tensor outside these subsets casts it
+//! explicitly.
+//!
 //! A weight may be promoted only when every op that reads it promotes it
 //! against an f32 activation. Weights combined with other weights in their
 //! own dtype first (for example Gemma's `1 + weight` norms) or read by a
@@ -104,6 +112,55 @@ mod tests {
         mlxcel_core::astype(a, dtype::FLOAT32)
     }
 
+    /// The bf16 weight `w` as upstream MLX promotes it against an f32
+    /// activation. Off CUDA that is MLX's own promotion inside the op, so `w`
+    /// is passed as is and the tests compare the precast with it. The CUDA
+    /// overlay resolves bf16 with f32 to bf16 instead (see
+    /// `mixed_bf16_f32_promotion_follows_the_build`), so there the upstream
+    /// promotion is spelled out as an unevaluated in-graph `astype`.
+    fn as_promoted(w: &MlxArray) -> UniquePtr<MlxArray> {
+        if cfg!(feature = "cuda") {
+            f32_of(w)
+        } else {
+            mlxcel_core::copy(w)
+        }
+    }
+
+    /// Issue #2087 probe: the dtype an f32 activation meets a bf16 operand in.
+    /// Upstream MLX gives f32. mlxcel's CUDA overlay
+    /// (`src/lib/mlx-cpp/patches-cuda/dtype.cpp`, issue #636) gives bf16 on
+    /// every device of a CUDA build, because the promotion table is compiled
+    /// into the library; ROCm and Metal use the upstream table. Production
+    /// code must therefore cast explicitly wherever it needs f32 from a stored
+    /// bf16 tensor, and this test fails loudly if either rule changes.
+    #[test]
+    fn mixed_bf16_f32_promotion_follows_the_build() {
+        let x = rand(1, &[2, 16]);
+        let w = bf16(&rand(2, &[16, 16]));
+        let g = bf16(&rand(3, &[16]));
+        let want = if cfg!(feature = "cuda") {
+            dtype::BFLOAT16
+        } else {
+            dtype::FLOAT32
+        };
+        let outputs = [
+            ("add(f32, bf16)", mlxcel_core::add(&x, &g)),
+            ("add(bf16, f32)", mlxcel_core::add(&g, &x)),
+            ("matmul(f32, bf16)", mlxcel_core::matmul(&x, &w)),
+            ("matmul(bf16, f32)", mlxcel_core::matmul(&w, &f32_of(&w))),
+            (
+                "rms_norm(f32, bf16)",
+                mlxcel_core::fast_rms_norm(&x, &g, 1e-5),
+            ),
+        ];
+        for (name, out) in &outputs {
+            assert_eq!(mlxcel_core::array_dtype(out), want, "{name}");
+        }
+        // An explicit cast is backend independent and exact.
+        let cast = mlxcel_core::add(&x, &f32_of(&g));
+        assert_eq!(mlxcel_core::array_dtype(&cast), dtype::FLOAT32);
+    }
+
     #[test]
     fn precast_matmul_and_addmm_match_the_promoted_graph() {
         let w = bf16(&rand(1, &[48, 64]));
@@ -112,11 +169,17 @@ mod tests {
         let b32 = f32_of(&b);
         for rows in [1, 2, 7] {
             let x = rand(3 + rows as u64, &[1, rows, 64]);
-            let promoted = mlxcel_core::matmul(&x, &mlxcel_core::transpose(&w));
+            let promoted = mlxcel_core::matmul(&x, &mlxcel_core::transpose(&as_promoted(&w)));
             let precast = mlxcel_core::matmul(&x, &mlxcel_core::transpose(&w32));
             assert_eq!(mlxcel_core::array_dtype(&promoted), dtype::FLOAT32);
             assert_eq!(bytes(&promoted), bytes(&precast));
-            let promoted = mlxcel_core::addmm(&b, &x, &mlxcel_core::transpose(&w), 1.0, 1.0);
+            let promoted = mlxcel_core::addmm(
+                &as_promoted(&b),
+                &x,
+                &mlxcel_core::transpose(&as_promoted(&w)),
+                1.0,
+                1.0,
+            );
             let precast = mlxcel_core::addmm(&b32, &x, &mlxcel_core::transpose(&w32), 1.0, 1.0);
             assert_eq!(bytes(&promoted), bytes(&precast));
         }
@@ -127,17 +190,18 @@ mod tests {
         let x = rand(9, &[1, 12, 16]);
         for (shape, groups) in [([32, 1, 16], 1), ([16, 9, 1], 16)] {
             let w = bf16(&rand(10, &shape));
-            let promoted = mlxcel_core::try_conv1d(&x, &w, 1, 0, 1, groups).unwrap();
+            let promoted = mlxcel_core::try_conv1d(&x, &as_promoted(&w), 1, 0, 1, groups).unwrap();
             let precast = mlxcel_core::try_conv1d(&x, &f32_of(&w), 1, 0, 1, groups).unwrap();
             assert_eq!(bytes(&promoted), bytes(&precast));
         }
         let g = bf16(&rand(11, &[16]));
         let beta = bf16(&rand(12, &[16]));
         let (g32, beta32) = (f32_of(&g), f32_of(&beta));
-        // SAFETY: all four pointers refer to arrays that outlive the calls.
+        let (gp, betap) = (as_promoted(&g), as_promoted(&beta));
+        // SAFETY: every pointer refers to an array that outlives the calls.
         let (promoted, precast) = unsafe {
             (
-                mlxcel_core::fast_layer_norm(&x, &*g, &*beta, 1e-5),
+                mlxcel_core::fast_layer_norm(&x, &*gp, &*betap, 1e-5),
                 mlxcel_core::fast_layer_norm(&x, &*g32, &*beta32, 1e-5),
             )
         };