fix: route nemotron omni audio-encoder convs through fallible FFI#439
Conversation
The Nemotron-H Nano Omni Parakeet/Conformer audio encoder built its data-dependent conv2d/conv1d ops through the non-fallible `mlxcel_core::conv2d`/`conv1d`. MLX validates conv shapes eagerly at graph-build time, so a crafted audio request whose runtime-derived post-subsampling shape mismatches the conv weights throws an uncaught C++ exception across the cxx boundary and `std::terminate`-aborts the whole server for all clients. This is the same remote-DoS class addressed for the Gemma 4 Conformer in #427/#434 and for the audio-synthesis matmul in #384. Route the four data-dependent conv calls in `src/audio/nemotron_h_nano_omni/encoder.rs` through the already-existing `try_conv2d`/`try_conv1d` wrappers (added in #434), which are declared `-> Result` so cxx catches MLX's eager shape throw and returns it as a Rust `Err`: the subsampling conv2d stack and the convolution module's pointwise1, depthwise, and pointwise2 conv1d calls. Thread `Result<_, String>` up the encoder forward chain so the error reaches a request boundary instead of an unwrap: `SubsamplingConv2DLayer::forward`, `ParakeetSubsamplingConv2D::forward`, `ParakeetConvModule::forward`, `ParakeetEncoderBlock::forward`, and the public `NemotronOmniSoundEncoder::forward` now return `Result`. The forward methods that never transitively call a conv (LayerNorm, BatchNorm, FeedForward, Attention, GLU/norm helpers) stay infallible. The VLM embedding builder `NemotronHNanoOmniVlModel::extract_audio_features` already returns `Result`, so it propagates the encoder error with `?`; its request-path call site in `compute_nemotron_h_nano_omni_audio_embeddings` already maps it to an `anyhow` request error, so a conv shape fault now surfaces as a per-request `Err` rather than a process abort. Add an encoder-level abort-safety test next to the encoder (`abort_safety_tests`) asserting that `SubsamplingConv2DLayer::forward` with a channel-mismatched weight returns `Err`, complementing the existing `try_conv2d`/`try_conv1d` FFI Err tests from #434. Update the three existing encoder-test call sites for the new `Result` return. Closes #435
…docs Add `conv_module_pointwise1_returns_err_on_channel_mismatch` to the abort_safety_tests module in the Nemotron-H Nano Omni encoder. It directly constructs a `ParakeetConvModule` with a 4-channel pointwise1 weight and feeds a 2-channel input, confirming the `try_conv1d` path returns `Err` on a shape fault rather than reaching `std::terminate`. Complements the existing conv2d test in the same module. Extend the Fault isolation section of `docs/audio-api.md` and the residual-abort-vector entry in `docs/adr/0003-release-panic-unwind-with-core-thread-abort.md` to record that PR #439 (issue #435) applied `try_conv2d`/`try_conv1d` coverage to the Nemotron-H Nano Omni Conformer/Parakeet encoder, covering all four data-dependent conv calls on that path.
PR Finalization CompleteSummaryTests: added Documentation: extended the Fault isolation paragraph in Lint/format: Verification commands runAll checks passing. Ready for merge. |
The initial v0.3.3 cut (4d5f4fb) was never tagged, and 16 more commits landed on main afterward. Fold those fixes into the v0.3.3 CHANGELOG and debian/changelog entries and set the release date to 2026-06-25: - N-gram loop detection (#433) and Nemotron-H Nano Omni audio input (#443) - Prefill masks sized from the live window under --max-kv-size trim (#418, #420, #422, #431) - Double-transpose crash on mlx-community conv checkpoints (#429) - conv1d/conv2d and nemotron audio-encoder convs fallible at the FFI boundary (#434, #439) - Gemma 4 audio placed in the user turn (#438, #440) - mistral4 MLA backbone routing and 2D MoE token flattening (#423, #425) - Bump actions/checkout from 6 to 7 (#395)
Summary
Extends the #427/#434 defense-in-depth conv hardening to the Nemotron-H Nano Omni Parakeet/Conformer audio encoder. Its data-dependent
conv2d/conv1dops are now built through the fallibletry_conv2d/try_conv1dFFI wrappers so a crafted audio request whose runtime-derived shape mismatches the conv weights returns a per-requestErrinstead of an uncaught C++ throw thatstd::terminate-aborts the server for all clients (the remote-DoS class documented in ADR 0003, first addressed for audio synthesis in #384).What changed
src/audio/nemotron_h_nano_omni/encoder.rs: route the four data-dependent conv calls through the already-existing (#434) fallible variants: subsamplingconv2d(SubsamplingConv2DLayer::forward) and the convolution module's pointwise1 / depthwise / pointwise2conv1d(ParakeetConvModule::forward), each with a descriptive.map_err(...).src/audio/nemotron_h_nano_omni/encoder.rs: threadResult<_, String>up the forward chain that transitively calls a conv:SubsamplingConv2DLayer::forward,ParakeetSubsamplingConv2D::forward,ParakeetConvModule::forward,ParakeetEncoderBlock::forward, and the publicNemotronOmniSoundEncoder::forward. Conv-free forward methods (LayerNorm, BatchNorm, FeedForward, Attention, GLU/norm helpers) stay infallible, so no per-op overhead is added to statically-shaped paths.src/vision/nemotron_h_nano_omni_vl.rs:NemotronHNanoOmniVlModel::extract_audio_features(already-> Result) propagates the encoder error with?; its request-path callercompute_nemotron_h_nano_omni_audio_embeddingsalready maps it to ananyhowrequest error, so the fault surfaces as a per-request error at the request boundary.src/audio/nemotron_h_nano_omni/encoder.rs: add an encoder-levelabort_safety_testsmodule assertingSubsamplingConv2DLayer::forwardwith a channel-mismatched weight returnsErr(complements thetry_conv2d/try_conv1dFFI Err tests added in fix: make conv2d/conv1d fallible at the FFI boundary so conv shape faults don't abort the server #434).src/audio/nemotron_h_nano_omni/encoder_tests.rs: update the three existingencoder.forward(...)call sites for the newResultreturn.Notes
try_conv2d/try_conv1dalready exist (merged in fix: make conv2d/conv1d fallible at the FFI boundary so conv shape faults don't abort the server #434), so this PR is Rust-only.Test plan
cargo check --bin mlxcel --features metal,acceleratecargo test -p mlxcel-core --lib --features metal,accelerate ffi_tests::try_conv(4 passed: the conv2d/conv1d Err + Ok-parity tests)cargo test --lib --features metal,accelerate nemotron_h_nano_omni::encoder(10 passed, incl. newabort_safety_tests::subsampling_conv2d_layer_returns_err_on_channel_mismatch)cargo clippy --lib --tests --features metal,accelerate -- -D warnings(clean)cargo fmt --check -p mlxcel(clean)Closes #435