Skip to content

fix: route nemotron omni audio-encoder convs through fallible FFI#439

Merged
inureyes merged 2 commits into
mainfrom
fix/issue-435-nemotron-omni-conv-fallible
Jun 25, 2026
Merged

fix: route nemotron omni audio-encoder convs through fallible FFI#439
inureyes merged 2 commits into
mainfrom
fix/issue-435-nemotron-omni-conv-fallible

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

Extends the #427/#434 defense-in-depth conv hardening to the Nemotron-H Nano Omni Parakeet/Conformer audio encoder. Its data-dependent conv2d/conv1d ops are now built through the fallible try_conv2d/try_conv1d FFI wrappers so a crafted audio request whose runtime-derived shape mismatches the conv weights returns a per-request Err instead of an uncaught C++ throw that std::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: subsampling conv2d (SubsamplingConv2DLayer::forward) and the convolution module's pointwise1 / depthwise / pointwise2 conv1d (ParakeetConvModule::forward), each with a descriptive .map_err(...).
  • src/audio/nemotron_h_nano_omni/encoder.rs: thread Result<_, String> up the forward chain that transitively calls a conv: SubsamplingConv2DLayer::forward, ParakeetSubsamplingConv2D::forward, ParakeetConvModule::forward, ParakeetEncoderBlock::forward, and the public NemotronOmniSoundEncoder::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 caller compute_nemotron_h_nano_omni_audio_embeddings already maps it to an anyhow request 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-level abort_safety_tests module asserting SubsamplingConv2DLayer::forward with a channel-mismatched weight returns Err (complements the try_conv2d/try_conv1d FFI 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 existing encoder.forward(...) call sites for the new Result return.

Notes

Test plan

  • cargo check --bin mlxcel --features metal,accelerate
  • cargo 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. new abort_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

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
@inureyes inureyes added type:enhancement New features, capabilities, or significant additions priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers status:review Under review labels Jun 25, 2026
…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.
@inureyes

Copy link
Copy Markdown
Member Author

PR Finalization Complete

Summary

Tests: added conv_module_pointwise1_returns_err_on_channel_mismatch to the abort_safety_tests module in src/audio/nemotron_h_nano_omni/encoder.rs. It constructs a ParakeetConvModule directly (same module, access to private fields) with a 4-channel pointwise1 weight and a 2-channel input, confirming the first try_conv1d call returns Err on shape mismatch rather than terminating. Pairs with the existing conv2d test in the same block. Both pass.

Documentation: extended the Fault isolation paragraph in docs/audio-api.md and the residual-abort-vector entry plus the References list 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.

Lint/format: cargo fmt --all --check clean (exit 0). cargo clippy --lib --features metal,accelerate -- -D warnings clean (exit 0).

Verification commands run

cargo fmt --all --check           # exit 0
cargo clippy --lib --features metal,accelerate -- -D warnings  # exit 0
cargo test --lib --features metal,accelerate \
  'nemotron_h_nano_omni::encoder::abort_safety_tests'
# 2 tests: subsampling_conv2d_layer_returns_err_on_channel_mismatch ... ok
#          conv_module_pointwise1_returns_err_on_channel_mismatch ... ok

All checks passing. Ready for merge.

@inureyes inureyes added status:done Completed and removed status:review Under review labels Jun 25, 2026
@inureyes
inureyes merged commit 263a759 into main Jun 25, 2026
5 checks passed
@inureyes
inureyes deleted the fix/issue-435-nemotron-omni-conv-fallible branch June 25, 2026 15:12
inureyes added a commit that referenced this pull request Jun 25, 2026
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers priority:medium Medium priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix: route nemotron_h_nano_omni audio-encoder conv2d/conv1d through the fallible FFI variants so conv shape faults don't abort the server

1 participant