From df686a7351cfb6005cca734112025843e54eb9c4 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 2 Oct 2026 17:04:44 +0900 Subject: [PATCH 1/2] fix(tts): cast bf16 tables explicitly where the CUDA overlay demotes 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 #636) resolves bf16 with f32 to bf16 on every device of a CUDA build. Five lib tests and two Nemotron VoiceChat TTS sites assumed upstream's f32: the RVQ depth-sum and the MoG head's gathered proj_mus/low_mat slabs returned bf16 on CUDA. Both now cast the gathered rows explicitly (exact, bit-identical off CUDA). A probe test pins the per-build rule for add, matmul and fast_rms_norm (f32 off CUDA, bf16 with it). The f32_weights and perception tests build their reference with an explicit in-graph cast on CUDA, and the MoG test loads weights through promoted_subset as the model does. The gelu failure was already fixed by #2080. Validated on GB10 (--features cuda) and a CPU-only build: f32_weights, fastconformer, nemotron_voicechat::tts and gemma3_backbone all pass; clippy clean. Closes #2087 --- ...7-cuda-bf16-promotion-tests-20261002.en.md | 55 ++++++++++++++ ...7-cuda-bf16-promotion-tests-20261002.ko.md | 55 ++++++++++++++ src/audio/f32_weights.rs | 74 +++++++++++++++++-- src/audio/fastconformer/tests.rs | 25 ++++--- src/models/nemotron_voicechat/tts/mog_head.rs | 9 ++- src/models/nemotron_voicechat/tts/rvq.rs | 4 + .../nemotron_voicechat/tts/tts_tests.rs | 38 +++++++--- 7 files changed, 229 insertions(+), 31 deletions(-) create mode 100644 TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.en.md create mode 100644 TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.ko.md diff --git a/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.en.md b/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.en.md new file mode 100644 index 000000000..e560ba19c --- /dev/null +++ b/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.en.md @@ -0,0 +1,55 @@ +# Technical Report: Issue #2087 - The six lib tests that failed under `--features cuda` + +**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. diff --git a/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.ko.md b/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.ko.md new file mode 100644 index 000000000..7e6a7521a --- /dev/null +++ b/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.ko.md @@ -0,0 +1,55 @@ +# 기술 보고서: Issue #2087 - `--features cuda`에서 실패하던 lib 테스트 6건 + +**날짜**: 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 스트림에서만 도달한다. diff --git a/src/audio/f32_weights.rs b/src/audio/f32_weights.rs index 87b46b4ee..c08ace707 100644 --- a/src/audio/f32_weights.rs +++ b/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 { + 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), ) }; diff --git a/src/audio/fastconformer/tests.rs b/src/audio/fastconformer/tests.rs index 2dfe42f57..99fffd4f9 100644 --- a/src/audio/fastconformer/tests.rs +++ b/src/audio/fastconformer/tests.rs @@ -457,26 +457,27 @@ fn attention_loads_non_default_quantization() { } /// Issue #2045 (P1): perception holds its bf16 weights as f32, which must be -/// bit-identical to MLX promoting the bf16 weights inside every op. +/// bit-identical to MLX promoting the bf16 weights inside every op. The +/// reference spells that promotion out as an in-graph cast, because CUDA +/// builds resolve bf16 + f32 to bf16 instead (issue #2087). #[test] fn perception_f32_weights_match_the_promoted_bf16_path() { let mut w = tiny_weights("p.encoder"); let mut s = 5u64; put(&mut w, "p.proj.weight", &[12, 8], &mut s, 0.3); put(&mut w, "p.proj.bias", &[12], &mut s, 0.1); - let w: WeightMap = w - .iter() - .map(|(k, v)| { - ( - k.clone(), - mlxcel_core::astype(v, mlxcel_core::dtype::BFLOAT16), - ) - }) - .collect(); + let cast = |w: &WeightMap, dt| -> WeightMap { + w.iter() + .map(|(k, v)| (k.clone(), mlxcel_core::astype(v, dt))) + .collect() + }; + let w = cast(&w, mlxcel_core::dtype::BFLOAT16); let perception = VoiceChatPerception::from_weights(&w, "p", &tiny_args(), (64, 4)).unwrap(); + let promoted = cast(&w, mlxcel_core::dtype::FLOAT32); let encoder = - FastConformerEncoder::from_weights(&w, "p.encoder", &tiny_args(), (64, 4)).unwrap(); - let proj = mlxcel_core::layers::UnifiedLinear::from_weights(&w, "p.proj", 64, 4).unwrap(); + FastConformerEncoder::from_weights(&promoted, "p.encoder", &tiny_args(), (64, 4)).unwrap(); + let proj = + mlxcel_core::layers::UnifiedLinear::from_weights(&promoted, "p.proj", 64, 4).unwrap(); let input = mel(24, 16); let out = perception.forward(&input, 24).unwrap(); let (encoded, _) = encoder.forward(&input, 24).unwrap(); diff --git a/src/models/nemotron_voicechat/tts/mog_head.rs b/src/models/nemotron_voicechat/tts/mog_head.rs index 2d3eb2180..2e3c4e402 100644 --- a/src/models/nemotron_voicechat/tts/mog_head.rs +++ b/src/models/nemotron_voicechat/tts/mog_head.rs @@ -192,10 +192,15 @@ impl MogHead { mlxcel_core::astype(&mlxcel_core::reshape(&component, &[-1]), dtype::INT32); let flat_x = mlxcel_core::reshape(&x, &[-1, h]); - let mus = mlxcel_core::take(&self.mus, &flat_component, 0); + // The gathered tables stay as stored (bf16); cast the gathered slabs + // to the activation dtype so `mu` stays f32 on CUDA builds, whose + // promotion table resolves bf16 + f32 to bf16 (issue #2087). Exact, + // and a no-op when the dtypes already match. + let act = mlxcel_core::array_dtype(&flat_x); + let mus = mlxcel_core::astype(&mlxcel_core::take(&self.mus, &flat_component, 0), act); let mu = mlxcel_core::matmul(&mus, &mlxcel_core::expand_dims(&flat_x, -1)); let mu = mlxcel_core::squeeze_axis(&mu, -1); - let low = mlxcel_core::take(&self.low_mat, &flat_component, 0); + let low = mlxcel_core::astype(&mlxcel_core::take(&self.low_mat, &flat_component, 0), act); let mu = mlxcel_core::matmul(&low, &mlxcel_core::expand_dims(&mu, -1)); let mu = mlxcel_core::reshape(&mlxcel_core::squeeze_axis(&mu, -1), &[b, t, self.out_size]); diff --git a/src/models/nemotron_voicechat/tts/rvq.rs b/src/models/nemotron_voicechat/tts/rvq.rs index f2b573922..1afb1c04d 100644 --- a/src/models/nemotron_voicechat/tts/rvq.rs +++ b/src/models/nemotron_voicechat/tts/rvq.rs @@ -113,6 +113,10 @@ impl RvqCodebooks { let mut result = mlxcel_core::zeros(&[s[0], s[1], self.latent_size], dtype::FLOAT32); for (idx, table) in self.padded.iter().enumerate() { let rows = mlxcel_core::take(table, &Self::column(code, idx as i32), 0); + // Cast the gathered rows explicitly: CUDA builds resolve + // bf16 + f32 to bf16 (issue #2087). The cast is exact and is the + // same `astype` upstream promotion inserts elsewhere. + let rows = mlxcel_core::astype(&rows, dtype::FLOAT32); result = mlxcel_core::add(&result, &rows); } Ok(result) diff --git a/src/models/nemotron_voicechat/tts/tts_tests.rs b/src/models/nemotron_voicechat/tts/tts_tests.rs index 0fecb225b..30e3fd3c4 100644 --- a/src/models/nemotron_voicechat/tts/tts_tests.rs +++ b/src/models/nemotron_voicechat/tts/tts_tests.rs @@ -321,27 +321,41 @@ fn mog_infer_shapes_are_finite_with_guidance() { w.insert(k, v); }; for l in 0..2 { - put(format!("m.mlp_stack.{l}.pre_norm.weight"), &[h], 0.1); - put(format!("m.mlp_stack.{l}.post_norm.weight"), &[h], 0.1); + put(format!("mog_head.mlp_stack.{l}.pre_norm.weight"), &[h], 0.1); put( - format!("m.mlp_stack.{l}.mlp.gate_proj.weight"), + format!("mog_head.mlp_stack.{l}.post_norm.weight"), + &[h], + 0.1, + ); + put( + format!("mog_head.mlp_stack.{l}.mlp.gate_proj.weight"), + &[16, h], + 0.3, + ); + put( + format!("mog_head.mlp_stack.{l}.mlp.up_proj.weight"), &[16, h], 0.3, ); - put(format!("m.mlp_stack.{l}.mlp.up_proj.weight"), &[16, h], 0.3); put( - format!("m.mlp_stack.{l}.mlp.down_proj.weight"), + format!("mog_head.mlp_stack.{l}.mlp.down_proj.weight"), &[h, 16], 0.3, ); } - put("m.mlp_stack.2.weight".into(), &[h], 0.1); - put("m.proj_logits.weight".into(), &[5, h], 0.3); - put("m.proj_mus.weight".into(), &[15, h], 0.3); - put("m.proj_logs.weight".into(), &[1, h], 0.3); - put("m.proj_else.weight".into(), &[out, h], 0.3); - put("m.low_mat".into(), &[5, out, 3], 0.3); - let head = MogHead::from_weights(&w, "m", h as usize, out as usize, &cfg, 64, 4).unwrap(); + put("mog_head.mlp_stack.2.weight".into(), &[h], 0.1); + put("mog_head.proj_logits.weight".into(), &[5, h], 0.3); + put("mog_head.proj_mus.weight".into(), &[15, h], 0.3); + put("mog_head.proj_logs.weight".into(), &[1, h], 0.3); + put("mog_head.proj_else.weight".into(), &[out, h], 0.3); + put("mog_head.low_mat".into(), &[5, out, 3], 0.3); + // Load the weights as `RvqEarTtsModel` does: the projections are held + // as f32 and the gathered tables (`proj_mus`, `low_mat`) as stored, so + // the f32 result also depends on the head casting the gathered slabs + // (issue #2087: CUDA builds resolve bf16 + f32 to bf16). + let w = crate::audio::f32_weights::promoted_subset(&w, "", super::model::promotes_to_f32); + let head = + MogHead::from_weights(&w, "mog_head", h as usize, out as usize, &cfg, 64, 4).unwrap(); let x = rand(&mut key, &[2, 1, h], 1.0); let (mu, logs) = head.infer(&x, 0.2, 0.95).unwrap(); assert_eq!(mlxcel_core::array_shape(&mu), vec![1, 1, out]); From b77d00611b9198c21eaa9503e100e7d261b4698d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 2 Oct 2026 17:06:43 +0900 Subject: [PATCH 2/2] docs: key the #2087 technical report by PR #2093 --- ...1002.en.md => 2093-cuda-bf16-promotion-tests-20261002.en.md} | 2 +- ...1002.ko.md => 2093-cuda-bf16-promotion-tests-20261002.ko.md} | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename TECHNICAL_REPORTS/{2087-cuda-bf16-promotion-tests-20261002.en.md => 2093-cuda-bf16-promotion-tests-20261002.en.md} (98%) rename TECHNICAL_REPORTS/{2087-cuda-bf16-promotion-tests-20261002.ko.md => 2093-cuda-bf16-promotion-tests-20261002.ko.md} (98%) diff --git a/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.en.md b/TECHNICAL_REPORTS/2093-cuda-bf16-promotion-tests-20261002.en.md similarity index 98% rename from TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.en.md rename to TECHNICAL_REPORTS/2093-cuda-bf16-promotion-tests-20261002.en.md index e560ba19c..75a606678 100644 --- a/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.en.md +++ b/TECHNICAL_REPORTS/2093-cuda-bf16-promotion-tests-20261002.en.md @@ -1,4 +1,4 @@ -# Technical Report: Issue #2087 - The six lib tests that failed under `--features cuda` +# Technical Report: PR #2093 - Cast bf16 tables explicitly where the CUDA overlay demotes (issue #2087) **Date**: 2026-10-02 diff --git a/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.ko.md b/TECHNICAL_REPORTS/2093-cuda-bf16-promotion-tests-20261002.ko.md similarity index 98% rename from TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.ko.md rename to TECHNICAL_REPORTS/2093-cuda-bf16-promotion-tests-20261002.ko.md index 7e6a7521a..428ae58ee 100644 --- a/TECHNICAL_REPORTS/2087-cuda-bf16-promotion-tests-20261002.ko.md +++ b/TECHNICAL_REPORTS/2093-cuda-bf16-promotion-tests-20261002.ko.md @@ -1,4 +1,4 @@ -# 기술 보고서: Issue #2087 - `--features cuda`에서 실패하던 lib 테스트 6건 +# 기술 보고서: PR #2093 - CUDA 오버레이가 강등하는 곳에서 bf16 테이블을 명시적으로 캐스트 (issue #2087) **날짜**: 2026-10-02