Validating MLX Beamforming Against MACH/CUDA

github repo fork

I read a cool blog post about a company building BCIs (https://alephneuro.com/blog/ultrasound-brain) so naturally I decided to check out their github and I came across a repo called "Microbubbles" The microbubbles beamforming path was CUDA-only in practice. The dependency that mattered was MACH, a CUDA/CuPy beamforming kernel. The data format, preprocessing, grid construction, and downstream ULM pipeline were not inherently NVIDIA-specific. The missing piece on a MacBook was the actual delay-and-sum backend.

I added an MLX backend for Apple Silicon and kept the rest of the path fixed. The goal was not to make a new beamformer with different assumptions. The goal was a backend swap: same neutral ultratrace input, same grid math, same transmit-delay handling, and the same output H5 layout: 

acquisitions/<id>/meta/compound_image
acquisitions/<id>/meta/grid/{x,y,z}

Backend Port

The MLX backend implements the same delay-and-sum pieces as the MACH path:

  • neutral IQ preprocessing
  • transmit delay normalization
  • imaging grid construction
  • receive aperture selection from f-number
  • linear interpolation over the sample axis
  • Tukey receive apodization
  • modulation phase correction
  • transmit-angle compounding
  • final layout as (frames, elev, z, x)

The CLI now exposes both backends:

ultratrace-ulm beamform --backend mlx ...
ultratrace-ulm beamform --backend mach ...

The backend dispatch stays small:

def beamform_iq(
    iq_frames: np.ndarray,
    tx_delays: np.ndarray,
    tx_delays_elev: np.ndarray,
    config: NeutralConfig,
    *,
    backend: str = "mach",
    scan_chunk: int = 2048,
    receive_chunk: int = 128,
    frame_chunk: int = 16,
) -> tuple[np.ndarray, Grid]:
    resolved = resolve_beamform_backend(backend)
    if resolved == "mach":
        return beamform_iq_mach(iq_frames, tx_delays, tx_delays_elev, config)
    return beamform_iq_mlx(
        iq_frames,
        tx_delays,
        tx_delays_elev,
        config,
        scan_chunk=scan_chunk,
        receive_chunk=receive_chunk,
        frame_chunk=frame_chunk,
    )

The main risk was not wiring. The risk was a quiet numerical mismatch: one backend looking plausible while drifting from the CUDA reference. So I treated the port as a reference-comparison problem.

The MLX loop is still the same delay calculation, interpolation, apodization, phase correction, and receive sum:

delta = rx[None, :, :] - scan[:, None, :]
horizontal_sq = delta[:, :, 0] * delta[:, :, 0] + delta[:, :, 1] * delta[:, :, 1]
inside = horizontal_sq <= aperture_radius_squared[:, None]
rx_distance = mx.sqrt(horizontal_sq + delta[:, :, 2] * delta[:, :, 2])
physical_tau = scan_tx[:, None] + rx_distance * inv_c
sample_idx = (physical_tau - float(rx_start_s)) * float(sampling_freq_hz)
valid = inside & (sample_idx >= 0.0) & (sample_idx <= float(n_samples - 1))

samples = _linear_interpolate_mlx(ch, sample_idx, valid, mx)
if tukey_alpha > 0.0:
    horizontal = mx.sqrt(mx.maximum(horizontal_sq, 0.0))
    r_norm = horizontal / aperture_radius[:, None]
    samples = samples * _tukey_apodization_mlx(r_norm, tukey_alpha, mx)[:, :, None]

if modulation_freq_hz:
    phase = mx.exp((1j * modulation_freq_rad) * physical_tau)
    samples = samples * phase[:, :, None]

accum = accum + mx.sum(samples, axis=1)
mx.eval(accum)

Synthetic Reference

I started with a deterministic synthetic case. I ran the same small input through MACH on Modal, returned the result as an .npz, and compared it locally against MLX. 

shape: (2, 1, 3, 2)
relative_l2: 1.6909404e-05
magnitude_corr: 1

This covered the basic mechanics: interpolation, apodization, phase correction, and compounding. It was useful, but still too clean. A synthetic test can pass while the real H5 path is wrong.

Real Acquisition

The public sample ultratrace H5 is about 98 GB. I did not want to move the full file just to check one backend. I added an HTTP range reader and extracted only acquisition 0 into a one-acquisition neutral H5.

The extracted input was:

file: sample_acq0.h5
iq_frames: (702, 5, 1, 134, 256) complex64, ~918.6 MB
tx_delays: (5, 134) float32
tx_delays_elev: (5, 8) float32

Both backends used the same deliberately coarse validation grid:

--acq-start 0
--num-acqs 1
--elev-planes 1
--z-coarseness 2.0
--x-coarseness 2.0
--no-large-fov

MLX ran locally on Apple Silicon. MACH ran on a Modal A10G instance. The output files were: 

MLX:  mlx_acq0_real.h5
MACH: mach_acq0_real.h5

Comparison

Both outputs had the same compound image shape:

compound_image shape: (700, 1, 56, 47)

The saved grid arrays matched exactly:

grid x max_abs: 0.0
grid y max_abs: 0.0
grid z max_abs: 0.0

The compound image comparison was:

max_abs: 8.93097
mean_abs: 1.29289
relative_l2: 4.5689e-05
magnitude_corr: 1

That is the result I cared about. The absolute differences are nonzero, as expected across different GPU backends and floating-point execution paths, but the relative error is small and the magnitude correlation is 1 at the printed precision.

Local tests also pass:

.venv/bin/python -m pytest -q
2 passed

Closing Notes

This is not a full production-grid validation. The real-acquisition comparison used: 

elev_planes = 1
z_coarseness = 2.0
x_coarseness = 2.0
large_fov = false

The useful practical outcome is that beamforming now has a MacBook path. MACH can remain the CUDA reference, and MLX can be used for local Apple Silicon work without changing the neutral ultratrace schema or downstream pipeline.