Make ONNX Runtime optional and run CUDA natively with cuda-oxide kernels by praveenperera · Pull Request #36 · avencera/speakrs

praveenperera · 2026-10-04T15:20:25Z

Summary

This PR makes ONNX Runtime optional, so the CUDA modes no longer use it. They run on a native NVIDIA backend instead, sped up by three custom cuda-oxide kernels. Overall, the CUDA modes run about 2.96× (cuda) and 2.69× (cuda-fast) faster than master. CPU and CoreML output stays byte-identical to master, and DER does not change.

What changed

Choosing a backend is now required. ONNX Runtime sits behind the cpu, migraphx and load-dynamic features. coreml and cuda don't use it. Building with no backend feature is a compile error, so a backend has to be chosen explicitly.
The PLDA setup no longer needs BLAS. It uses a small Rust solver: Cholesky, pivoted LU and a symmetric eigensolver. It matches LAPACK to about 1e-14 on the shipped models. Inputs at extreme scales are rescaled first, or return a typed error.
Native CUDA backend (breaking). ExecutionMode::Cuda and CudaFast run on cudarc with cuBLAS and cuDNN 9, loaded at run time. Models load from safetensors weights, and CUDA graphs are used.
- Precision is configurable per model through RuntimeConfig. Segmentation runs in FP32 by default, and embedding in TF32.
- It needs a Turing (sm_75) or newer GPU, but no CUDA toolkit to build.
Three cuda-oxide kernels. They ship as committed sm_75 PTX, so end users don't need nightly Rust. Higher tiers (cuda-sm80, cuda-sm90, cuda-sm120) are opt-in features and off by default. cuDNN still handles every other layer, batch and precision combination.

Kernel	Coverage	Speed vs cuDNN
ResNet 3×3 conv	C32 and strided convs in FP32 and TF32; C64 at batch 7+ in both, batch 1 in FP32	1.19–3.1× per conv
BiLSTM recurrence	all batches, FP32	1.10–2.95× per stage vs `PersistStaticSmallH`
SincNet conv+abs+pool	all batches, FP32	2.1× per operator, bit-identical output

A locked qualification harness. It decides when a custom kernel may replace a cuDNN or cuBLAS call (cargo xtask cuda-qualify, scripts/cuda/qualify/). Production selects a kernel only for the layer, batch and precision combinations the harness qualified, because it reads the same coverage declarations.
- The harness checks layer and stage parity, determinism, Compute Sanitizer, the profile and CUDA graph rules, an f64-referenced secret-input check, and A/B/A/B timing with a no-slowdown bar.
- Large test assets stay outside git, identified by hash.
GPU images and benchmarks. The GPU images no longer include ONNX Runtime, and CI fails if it returns to the GPU binary's dependency tree. The GPU benchmark no longer needs the S3 bucket: it downloads public models and datasets and prints results in the log.
Docs. ADR 002 now correctly says the old ORT CUDA path ran with TF32 and NHWC enabled.

Why

Dropping ONNX Runtime from CUDA and CoreML builds removes a large native dependency and its download step. It also lets the native backend run faster than ORT did, while keeping DER unchanged.

Results

VoxConverse-dev, all 216 files, run A/B/A/B:

Step	`cuda`	`cuda-fast`
1. Native backend in FP32 vs `master`'s ORT CUDA, matched per file	1.81×	1.62×
2. Shipped defaults vs native FP32 (TF32 embedding, `PersistStaticSmallH`, CUDA graphs), full-set RTFx	300 → 399 (1.33×)	547 → 748 (1.37×)
3. Adding the kernels on top, full-set RTFx	404 → 496 (1.23×)	761 → 924 (1.21×)
Total vs `master`	1.81 × 1.33 × 1.23 ≈ 2.96×	1.62 × 1.37 × 1.21 ≈ 2.69×
DER change vs `master`	neutral	neutral

The total is the product of the three measured steps, not a direct A/B: master exceeded the GPU box's memory limit in the one-process full-set run, so its full-set RTFx could not be measured.
DER is unchanged at each step. The native FP32 backend is DER-neutral against master, the default precision choices were each gated per file, and the kernels change DER by +0.0003 points in cuda and 0 in cuda-fast.
No file's DER moved by more than 0.5 points.
CPU output is byte-identical to master (280 of 280 compared) and so is CoreML (120 of 120). CoreML also starts in about half the time and uses about 90 MB less memory.

Validation

just fmt, just clippy, cargo test --workspace, and a warnings-as-errors docs build.
cargo check for every supported backend combination with --no-default-features. The intended compile errors fire for invalid combinations.
The GPU suites, harness qualification and benchmarks ran on an RTX 5070 Ti (sm_120) with the sm_75 PTX forced.
Each kernel passed the qualification harness and an independent audit. A final review covered the whole change.

Known limits

Real Turing (sm_75) hardware is untested. Qualification forced the sm_75 PTX on sm_120.
The LSTM kernel uses cooperative grids. On pre-Hopper GPUs, two pipelines sharing one GPU can still contend. That is the same class of risk as cuDNN's persistent RNN algorithms, and it is documented in the code.
The GPU image has not been rebuilt for this change.

PLDA setup needs two inverses and one generalized symmetric eigensolve on 128x128 matrices, once per pipeline. That alone pulled in ndarray-linalg with statically linked MKL or OpenBLAS and four backend features. A small Cholesky, pivoted LU, and tred2/tql2 solver now does this with no dependencies. Results match LAPACK to about 1e-14 on phi and 2e-12 on the transform, and diarization output is unchanged. Breaking: the default-linalg, intel-mkl, openblas-static, and openblas-system features are removed.

macOS builds only need CoreML, yet every build compiled, linked, and downloaded ONNX Runtime, and CoreML modes still built ORT sessions at load time. ONNX Runtime is now an optional dependency enabled only by the cpu, cuda, migraphx, and load-dynamic features. Each build must pick a backend explicitly; a build without one fails with a compile error that lists the choices. Each model holds a single backend chosen at load, so CoreML modes never create ORT sessions or download ONNX files. A crate-owned InferenceError replaces ort::Error in the API. CoreML and CPU diarization output is byte-identical to before on VoxConverse-dev samples. CoreML startup roughly halves and peak memory drops by about 90 MB. Breaking: no default backend, ExecutionMode::Cpu needs the cpu feature, and ort::Error is replaced in public signatures.

The native CUDA backend loads safetensors weights instead of ONNX models, and its kernels need layer-by-layer references to check parity against ONNX Runtime. export_weights.py converts the segmentation and multi-mask embedding ONNX models into the safetensors files the runtime loads, along with a manifest of each graph. make_reference.py records ONNX Runtime CPU inputs, outputs, and every intermediate tensor for the fixture audio. Both are deterministic and download only public models.

ONNX Runtime's CUDA provider ran the filterbank DFT on the CPU, copied data back and forth every batch, needed over 15 GB of GPU memory in one process, and tied users to a matching ORT, CUDA, and cuDNN install. CUDA modes now run natively through cudarc: cuBLAS and cuDNN for the heavy layers and cuda-oxide kernels, shipped as committed PTX, for the filterbank, pooling, SincNet front end, and epilogues. The cuda feature no longer pulls in ONNX Runtime and builds without a CUDA toolkit. Models load from safetensors weights. On all 216 VoxConverse-dev files, DER matches the ONNX Runtime path and diarization runs 1.81x faster in cuda mode and 1.62x faster in cuda-fast mode. Defaults: FP32 segmentation (TF32 worsened one file by 4.5 DER points), TF32 embedding, the persistent SmallH LSTM, and CUDA graphs. Kernels target sm_75; cuda-sm80, cuda-sm90, and cuda-sm120 opt in to newer GPU tiers. Breaking: CUDA modes need the new safetensors assets, a Turing or newer GPU, and cuDNN 9 at run time.

The CUDA benchmark binary no longer uses ONNX Runtime's CUDA provider, so the images stop downloading and copying the ORT GPU shared libraries. The native backend loads cuBLAS, cuDNN 9, and NVRTC from the CUDA runtime base image. These Dockerfile changes have not been built yet.

The gpuq canary staged models and datasets from a Tigris bucket and uploaded results there, so it could not run without that bucket and its AWS credentials. speakrs-bm now downloads the model assets for the selected CUDA modes anonymously from Hugging Face at the revision pinned in src/models.rs. Datasets come from their public sources through the existing xtask acquisition. Results stay on the worker and a compact summary is printed between SPEAKRS_RESULTS markers in the captured log. The workload needs no secrets. The canary's --impls speakrs now selects the native CUDA implementation; it previously matched no implementation.

The ONNX Runtime CUDA EP enabled TF32 and NHWC, not tf32=false. Note that CUDA modes now run on the native backend with per-model math.

Disable default CPU features in both CUDA build stages, since the default xtask feature enables ONNX Runtime. Check the GPU binary dependency tree in CI so ORT cannot return unnoticed.

Use exact power-of-two scaling outside the safe f64 range and restore eigenvalues and B-normalized vectors. Return typed errors when a scale would lose nonzero values. Keep normal-scale arithmetic unchanged. Cover extreme and generalized problems with relative residuals, and pin the shipped scale factors. Refresh only the linalg source hash in the qualification lock.

Show the CUDA PTX environment name as inline code and keep ORT-only items out of the CUDA feature labels. Resolve the VBx variant links so the strict docs build passes. Refresh only the four changed doc-file hashes in the qualification lock. Keep all kernel and qualification harness bytes unchanged.

coderabbitai · 2026-10-04T15:20:34Z

Important

Review skipped

Too many files!

This PR contains 190 files, which is 90 over the limit of 100.

To get a review, reduce the PR to 100 files or fewer by splitting it into smaller PRs or changing its base branch.

Upgrade to a paid plan to raise the limit.

⚙️ Run configuration

Configuration used: defaults
Review profile: CHILL
Plan: Advanced
Run ID: 25d3a6ec-804f-459a-acc0-e34ce2ee5d93

📥 Commits

Reviewing files that changed from the base of the PR and between ab98f35 and 4423732.

⛔ Files ignored due to path filters (6)

Cargo.lock is excluded by !**/*.lock
crates/speakrs-cuda-kernels/Cargo.lock is excluded by !**/*.lock
tests/cuda_qualify/baselines/lstm-sm75-initcheck.log is excluded by !**/*.log
tests/cuda_qualify/baselines/lstm-sm75-memcheck.log is excluded by !**/*.log
tests/cuda_qualify/baselines/lstm-sm75-racecheck.log is excluded by !**/*.log
tests/cuda_qualify/device/Cargo.lock is excluded by !**/*.lock

📒 Files selected for processing (190)

.github/workflows/ci.yml
.github/workflows/package.yml
.gitignore
CHANGELOG.md
Cargo.toml
Dockerfile.gpu
README.md
adr/002-consolidate-cuda-modes.md
crates/speakrs-cuda-kernels/Cargo.toml
crates/speakrs-cuda-kernels/rust-toolchain.toml
crates/speakrs-cuda-kernels/src/embedding.rs
crates/speakrs-cuda-kernels/src/fbank.rs
crates/speakrs-cuda-kernels/src/lib.rs
crates/speakrs-cuda-kernels/src/lstm.rs
crates/speakrs-cuda-kernels/src/probe.rs
crates/speakrs-cuda-kernels/src/resnet.rs
crates/speakrs-cuda-kernels/src/segmentation.rs
crates/speakrs-cuda-kernels/src/sincnet.rs
deny.toml
docker/gpuq-canary.Dockerfile
docker/gpuq-workload.sh
examples/README.md
examples/assign_transcript_speakers.rs
examples/diarize_wav.rs
examples/print_turns.rs
examples/queued.rs
examples/speaker_airtime.rs
gpuq.toml
justfile
scripts/cuda/README.md
scripts/cuda/export_weights.py
scripts/cuda/make_reference.py
scripts/cuda/qualify/LOCK
scripts/cuda/qualify/QUALIFIED.json
scripts/cuda/qualify/README.md
scripts/cuda/qualify/SANITIZER_POLICY.json
scripts/cuda/qualify/SCOPE.json
scripts/cuda/qualify/assets.py
scripts/cuda/qualify/freeze_baselines.py
scripts/cuda/qualify/freeze_control.py
scripts/cuda/qualify/gates.py
scripts/cuda/qualify/lock.py
scripts/cuda/qualify/nvtx.cpp
scripts/cuda/qualify/parse_trace.py
scripts/cuda/qualify/ptx.py
scripts/cuda/qualify/qualified.py
scripts/cuda/qualify/qualify.py
scripts/cuda/qualify/scan.py
scripts/cuda/setup-gpu-box.sh
src/clustering/plda.rs
src/clustering/vbx.rs
src/inference.rs
src/inference/coreml.rs
src/inference/cuda.rs
src/inference/cuda/blas.rs
src/inference/cuda/buffer.rs
src/inference/cuda/candidate.rs
src/inference/cuda/candidate/conv.rs
src/inference/cuda/candidate/lstm.rs
src/inference/cuda/candidate/lstm/layout.rs
src/inference/cuda/candidate/sinc.rs
src/inference/cuda/candidate_test_support.rs
src/inference/cuda/candidate_tests.rs
src/inference/cuda/dispatch.rs
src/inference/cuda/dnn.rs
src/inference/cuda/embedding.rs
src/inference/cuda/embedding/dispatch.rs
src/inference/cuda/embedding/kernels.rs
src/inference/cuda/embedding/trunk.rs
src/inference/cuda/embedding/trunk/test_support.rs
src/inference/cuda/error.rs
src/inference/cuda/fbank.rs
src/inference/cuda/fbank/constants.rs
src/inference/cuda/implementation.rs
src/inference/cuda/implementation/tests.rs
src/inference/cuda/kernels.rs
src/inference/cuda/math.rs
src/inference/cuda/options.rs
src/inference/cuda/probe.rs
src/inference/cuda/ptx/embedding.manifest
src/inference/cuda/ptx/embedding.sm75.ptx
src/inference/cuda/ptx/fbank.manifest
src/inference/cuda/ptx/fbank.sm75.ptx
src/inference/cuda/ptx/lstm.manifest
src/inference/cuda/ptx/lstm.sm75.ptx
src/inference/cuda/ptx/probe.manifest
src/inference/cuda/ptx/probe.sm75.ptx
src/inference/cuda/ptx/probe.sm80.ptx
src/inference/cuda/ptx/resnet.manifest
src/inference/cuda/ptx/resnet.sm75.ptx
src/inference/cuda/ptx/segmentation.manifest
src/inference/cuda/ptx/segmentation.sm75.ptx
src/inference/cuda/ptx/sincnet.manifest
src/inference/cuda/ptx/sincnet.sm75.ptx
src/inference/cuda/runtime.rs
src/inference/cuda/segmentation.rs
src/inference/cuda/segmentation/dispatch.rs
src/inference/cuda/segmentation/graph.rs
src/inference/cuda/segmentation/graph/test_support.rs
src/inference/cuda/segmentation/kernels.rs
src/inference/cuda/segmentation/rnn.rs
src/inference/cuda/segmentation/shape.rs
src/inference/cuda/segmentation/weights.rs
src/inference/cuda/session.rs
src/inference/cuda/test_support.rs
src/inference/cuda/tests.rs
src/inference/cuda/tests/embedding.rs
src/inference/cuda/tests/fbank.rs
src/inference/cuda/tests/resnet.rs
src/inference/cuda/tests/runtime.rs
src/inference/cuda/tests/segmentation.rs
src/inference/cuda/tier.rs
src/inference/cuda/weights.rs
src/inference/cuda/weights/test_support.rs
src/inference/embedding.rs
src/inference/embedding/batch.rs
src/inference/embedding/buffers.rs
src/inference/embedding/cuda.rs
src/inference/embedding/fbank.rs
src/inference/embedding/load.rs
src/inference/embedding/load/sessions.rs
src/inference/embedding/native.rs
src/inference/embedding/native/loaders.rs
src/inference/embedding/native/run.rs
src/inference/embedding/onnx.rs
src/inference/embedding/onnx/fbank_pool.rs
src/inference/embedding/onnx/plan.rs
src/inference/embedding/onnx/run.rs
src/inference/embedding/paths.rs
src/inference/embedding/plan.rs
src/inference/embedding/run.rs
src/inference/embedding/session.rs
src/inference/embedding/tail.rs
src/inference/embedding/tensor.rs
src/inference/error.rs
src/inference/geometry.rs
src/inference/ort_runtime.rs
src/inference/segmentation.rs
src/inference/segmentation/cuda.rs
src/inference/segmentation/native.rs
src/inference/segmentation/onnx.rs
src/inference/segmentation/parallel.rs
src/inference/segmentation/parallel/batch.rs
src/inference/segmentation/parallel/single.rs
src/inference/segmentation/run.rs
src/inference/segmentation/tensor.rs
src/lib.rs
src/linalg.rs
src/models.rs
src/pipeline.rs
src/pipeline/builder.rs
src/pipeline/concurrent.rs
src/pipeline/config.rs
src/pipeline/tests.rs
src/pipeline/types/error.rs
src/powerset.rs
tests/cuda_pipeline.rs
tests/cuda_qualify/ASSETS.json
tests/cuda_qualify/baselines/__init__.py
tests/cuda_qualify/baselines/lstm-sm75.json
tests/cuda_qualify/baselines/test_cache.py
tests/cuda_qualify/device/Cargo.toml
tests/cuda_qualify/device/controls.sm75.ptx
tests/cuda_qualify/device/qualify.sm75.ptx
tests/cuda_qualify/device/rust-toolchain.toml
tests/cuda_qualify/device/src/lib.rs
tests/cuda_qualify/device/uninit_shared.sm75.ptx
tests/cuda_qualify/embedding.rs
tests/cuda_qualify/probe.rs
tests/cuda_qualify/reference.rs
tests/cuda_qualify/scan_fixtures/phase_cheat.rs
tests/cuda_qualify/scan_fixtures/unscanned_call.rs
tests/cuda_qualify/segmentation.rs
tests/cuda_qualify/test_assets.py
tests/cuda_qualify/test_gates.py
tests/cuda_qualify/test_processes.py
tests/cuda_qualify/test_qualified.py
tests/end_to_end.rs
tests/gpuq-workload.sh
tests/queued.rs
tests/support/mod.rs
xtask/Cargo.toml
xtask/README.md
xtask/src/bin/speakrs_bm.rs
xtask/src/cli.rs
xtask/src/commands.rs
xtask/src/commands/cuda_kernels.rs
xtask/src/commands/cuda_qualify.rs
xtask/src/commands/diarize.rs
xtask/src/commands/models.rs

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

Review on demand using usage pricing

Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

greptile-apps · 2026-10-04T15:27:08Z

Confidence Score: 5/5

[Critical risk] Replaces ONNX Runtime with native CUDA backend for inference.

The PR appears safe to merge based on this review; no new blocking issue or outstanding previous finding was identified.

Summary

The PR makes ONNX Runtime optional, adds native CUDA inference and qualified custom kernels, and replaces the PLDA BLAS dependency. Since the previous review, it adds production fallback for devices that cannot host the LSTM candidate and a CI qualification-manifest check.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[CUDA inference] --> B{Qualified coverage?}
  B -- No --> L[Library implementation]
  B -- Yes --> P{Candidate plan succeeds?}
  P -- Yes --> K[Custom kernel]
  P -- Device unsupported --> L
  P -- Other error --> E[Propagate error]

_{Reviews (2) · Last reviewed commit: "Tighten the qualification manifest check"}

The host-only PLDA tests read the PLDA model files, which this job did not fetch.

Keep device limits distinct from real planning errors. Allow Library fallback in production, but fail explicit qualification selections.

Bind production PTX, source files and coverage to accepted records. Encode the Library-noise rule and reject incomplete evidence.

Require the source-evidence note on each manifest entry, correct the Scratch and SideStream signatures in the harness docs, and keep CI push runs on master.

praveenperera · 2026-10-04T18:41:53Z

@greptileai

praveenperera added 14 commits October 2, 2026 15:01

Add qualified cuda-oxide device kernels

a18ceef

Add locked CUDA qualification harness

e7330e9

Enable kernels for qualified CUDA triples

9b8b644

Document qualified CUDA kernel coverage

80c5f11

Correct CUDA precision notes in ADR 002

fc2327f

The ONNX Runtime CUDA EP enabled TF32 and NHWC, not tf32=false. Note that CUDA modes now run on the native backend with per-model math.

Exclude ORT from the GPU image build

d0465b1

Disable default CPU features in both CUDA build stages, since the default xtask feature enables ONNX Runtime. Check the GPU binary dependency tree in CI so ORT cannot return unnoticed.

greptile-apps Bot reviewed Oct 4, 2026

View reviewed changes

Comment thread src/inference/cuda/candidate/lstm.rs

Comment thread .github/workflows/ci.yml

praveenperera added 4 commits October 4, 2026 10:28

Download PLDA fixtures in the CUDA-only CI job

bd07803

The host-only PLDA tests read the PLDA model files, which this job did not fetch.

Fall back for unsupported CUDA candidates

774a6a8

Keep device limits distinct from real planning errors. Allow Library fallback in production, but fail explicit qualification selections.

Enforce CUDA qualification in CI

5311146

Bind production PTX, source files and coverage to accepted records. Encode the Library-noise rule and reject incomplete evidence.

Tighten the qualification manifest check

4423732

Require the source-evidence note on each manifest entry, correct the Scratch and SideStream signatures in the harness docs, and keep CI push runs on master.

praveenperera marked this pull request as ready for review October 5, 2026 11:20

praveenperera merged commit c98d2d6 into master Oct 8, 2026
17 checks passed

praveenperera deleted the cuda-kernels-int branch October 8, 2026 00:15

Navigation Menu

Repository navigation

Make ONNX Runtime optional and run CUDA natively with cuda-oxide kernels - #36

Make ONNX Runtime optional and run CUDA natively with cuda-oxide kernels#36
praveenperera merged 18 commits into
masterfrom
cuda-kernels-int

praveenperera commented Oct 4, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Oct 4, 2026 •

edited

Loading

Review skipped

Uh oh!

greptile-apps Bot commented Oct 4, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

praveenperera commented Oct 4, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

praveenperera commented Oct 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

Why

Results

Validation

Known limits

Uh oh!

coderabbitai Bot commented Oct 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

greptile-apps Bot commented Oct 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Uh oh!

Uh oh!

Uh oh!

praveenperera commented Oct 4, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

praveenperera commented Oct 4, 2026 •

edited

Loading

coderabbitai Bot commented Oct 4, 2026 •

edited

Loading

greptile-apps Bot commented Oct 4, 2026 •

edited

Loading