Skip to content

M8: native conv via mlx::core::conv_general - #17

Merged
ausimian merged 1 commit into
mainfrom
m8-native-conv
Apr 14, 2026
Merged

M8: native conv via mlx::core::conv_general#17
ausimian merged 1 commit into
mainfrom
m8-native-conv

Conversation

@ausimian

Copy link
Copy Markdown
Owner

Summary

  • Lift Backend.conv/4 from the via_binary fallback onto Native.conv_general (already bound to mlx::core::conv_general since M1). The fallback was correct but CPU-bound — ≥90% of ViT and Whisper forward-pass cost.
  • Handle the layout mismatch (Nx NCHW/OIHW ↔ MLX NHWC/OHWI) by composing the caller's input_permutation, kernel_permutation, and output_permutation opts with the canonical-layout transposes. Apply the inverse of output_permutation on the way out (Nx delivers it in user→canonical form; see deps/nx/lib/nx/shape.ex:729-735).
  • Cast integer operands to out.type via Native.astype before dispatching — Nx.conv returns float but does not cast inputs, and MLX conv is float-only. Same-type astype is elided by MLX.
  • batch_group_size > 1 and complex-typed conv still route through via_binary — no MLX primitive, and neither appears in the pinned Bumblebee ref.

Speedups

Suite Before (fallback) After (native) Factor
mix test --only whisper_full 403.8s 1.8s ~224×
mix test --only vit_full tens of seconds 1.9s ~20–50×

Test plan

  • mix precommit — compile warnings-as-errors, format, credo --strict, 198 tests + 56 properties + 1 doctest green
  • mix test --only conformance — 16/16 tiny-random (DistilBERT, Qwen3, ViT, Whisper)
  • mix test --only vit_full — ViT full-checkpoint
  • mix test --only whisper_full — Whisper full-checkpoint
  • New test/emily/backend_conv_test.exs oracle suite vs Nx.BinaryBackend:
    • 1-D, 2-D, 3-D conv
    • stride, :same / :valid / explicit asymmetric padding
    • kernel_dilation and input_dilation > 1
    • grouped and depthwise conv
    • all three permutation options (independently and combined NHWC end-to-end)
    • integer-input coercion
  • test/emily/backend_fallbacks_test.exs — removed obsolete "conv routes through BinaryBackend" test, added batch_group_size > 1 fallback coverage

Lift Backend.conv/4 from the BinaryBackend fallback onto
Native.conv_general (already bound to mlx::core::conv_general since
M1). The fallback was correct but CPU-bound — ≥90% of ViT and
Whisper forward-pass cost.

Layout translation: MLX conv_general expects NHWC input and OHWI
weight; Nx delivers tensors in the caller's layout plus
input/kernel/output_permutation opts pointing at canonical NCHW/OIHW.
Compose the caller's permutations with the NCHW↔NHWC and OIHW↔OHWI
transposes, applying the inverse of output_permutation on the way
out (Nx delivers it in user→canonical form; see
deps/nx/lib/nx/shape.ex:729-735). Two ordered transposes rather
than one composed — MLX lazy-fuses them and the step-wise form is
obviously correct across rank 3, 4, and 5.

Integer inputs get cast to out.type via Native.astype before the
transpose chain (Nx.conv returns float but does not cast operands;
MLX conv is float-only). Same-type astype is elided by MLX.

batch_group_size > 1 and complex-typed conv still route through
via_binary — no MLX primitive, and neither appears in the pinned
Bumblebee ref.

Whisper full-checkpoint conformance: 403.8s → 1.8s (~224× speedup).
ViT full-checkpoint conformance: similar scale (1.9s end-to-end).

New oracle suite in test/emily/backend_conv_test.exs covers 1-D/2-D/
3-D, stride/:same/:valid/asymmetric padding, kernel and input
dilation, grouped and depthwise conv, all three permutation options
(independently and combined for NHWC end-to-end), and integer-input
coercion.

test/emily/backend_fallbacks_test.exs: dropped the obsolete
"conv routes through BinaryBackend" test; added a batch_group_size
> 1 case asserting the rare fallback still matches BinaryBackend.
@ausimian
ausimian merged commit 05c5a49 into main Apr 14, 2026
1 check passed
@ausimian
ausimian deleted the m8-native-conv branch April 14, 2026 23:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant