Technical Summary: Meganeura – Portable GPU Training and Inference through Vulkan and Metal

Problem Statement

The current machine learning ecosystem suffers from a fragmentation between training and deployment. Training typically relies on dynamic Python frameworks (e.g., PyTorch) and vendor-specific libraries (CUDA, ROCm), while deployment often requires exporting models to separate runtimes (TensorRT, Core ML, mobile engines). This "split stack" necessitates duplicate work in model conversion, operator coverage, numerical policy validation, and debugging. Furthermore, deployed models are usually frozen; adapting them on-device requires a round-trip through a separate training stack, preventing continuous, in-place learning.

The paper investigates whether a single, compact native compiler can span both training and inference phases across diverse consumer GPUs using general-purpose graphics APIs (Vulkan and Metal) rather than proprietary compute stacks.

Methodology

The author presents Meganeura, a native Rust-based system that unifies inference and training under a single typed static graph.

• Architecture: The system uses a mid-level typed graph where nodes represent operations like matrix multiplication, convolution, and attention. It employs reverse-mode automatic differentiation to generate backward graphs.
• Compilation Pipeline: The graph is rewritten (using greedy rewriting and optional equality saturation), specialized for specific hardware capabilities (e.g., cooperative matrices, tile geometry), and lowered to WGSL (WebGPU Shading Language).
• Backend: WGSL is validated and translated by Naga and executed via Blade, a native graphics abstraction layer that targets Vulkan (Linux, Windows, Android) and Metal (Apple).
• Execution: The system generates a static dispatch sequence with explicit buffer ownership and memory planning. It does not rely on a Python runtime or dynamic graph capture at runtime.
• Evaluation Protocol: The author used Inferena, a cross-framework harness, to compare Meganeura against PyTorch across five devices (NVIDIA RTX 5070, AMD RX 7900 XT, AMD Radeon 780M APU, Intel iGPU, Apple M3) and five workloads (SmolLM2, SmolVLA, Diffusion U-Net, ResNet-50, Whisper-tiny).
• Correctness Gates: Results were validated against strict numerical contracts:
  • Forward: Relative L2 error < 1%.
  • Backward: Total gradient-norm error < 5% and per-parameter gradient-norm error < 5%.
  • Two arithmetic modes were tested: Strict f32 (controlled precision) and Accelerated (allowing TF32 on NVIDIA and f16 cooperative matrices where supported).

Key Contributions

1. Unified Train-to-Deploy Stack: A compact executable system (34.5 KLOC Rust, 6.1 KLOC WGSL) that supports inference, reverse-mode differentiation, SGD/Adam updates, and checkpoint transfer without a Python runtime or CUDA/ROCm dependencies.
2. Controlled Comparative Evaluation: A rigorous comparison across 50 device-workload-mode cells. 48 cells passed both forward and backward correctness gates. The study isolates performance gaps to specific kernel families rather than attributing them to the graphics APIs themselves.
3. Performance Decomposition: The paper attributes performance differences to concrete factors: kernel coverage (specifically convolution derivatives and attention backward), dispatch structure, and arithmetic policy (cooperative matrix selection), rather than inherent API limitations.
4. Physical Deployment Case Study: A demonstration of DinoVision, an Android XR application that shares a graphics queue between rendering and a Meganeura-trained decoder, validating the stack's ability to coexist with graphics workloads.

Results

Performance Portability

• AMD Devices: Meganeura achieved near-parity or better performance compared to compiled ROCm PyTorch. On the RX 7900 XT, four of five inference workloads were within 1.10× of PyTorch, and three training workloads were faster.
• NVIDIA Devices: In strict f32 mode, Meganeura won 12 of 20 minimal-latency cells. However, training gaps were larger (up to 4.6× on ResNet-50), primarily due to missing optimized convolution derivative kernels and the lack of TF32 support in the Meganeura path.
• Apple Silicon: Inference latency was competitive (0.93–1.88×), but training lagged (1.20–2.84×) against eager MPS, largely due to backward attention kernel inefficiencies.
• Intel iGPU: Since the vendor XPU path was non-functional in the test environment, Meganeura outperformed the CPU fallback in 12 of 15 cells.
• Compilation Speed: Meganeura compiled workloads in 0.1–2.4 seconds, significantly faster than torch.compile (6–96 seconds).

Correctness and Stability

• Success Rate: 48 of 50 cells passed both forward and backward gates.
• Anomalies: Two failures occurred on the AMD 780M APU for the Whisper backward pass. Cross-device gradient analysis suggested the reference path (PyTorch/ROCm) was the source of the error, not Meganeura, though a third implementation is needed for resolution.
• Numerical Precision: In strict f32 mode, the worst forward L2 error was 0.0043% (gate: 1%), and the worst gradient-norm error was 0.64% (gate: 5%).

Systems and Deployment

• Artifact Size: The stripped Meganeura binary is 13 MiB, compared to ~4.66 GiB for a full PyTorch CUDA runtime closure.
• Memory: Meganeura's static memory planner reduced physical allocation by 4–48% in training and 6–79% in inference compared to logical totals, though dynamic allocators in PyTorch sometimes achieved lower peaks for specific activation-heavy workloads.
• Co-tenancy: The DinoVision case study demonstrated that Meganeura can share a Vulkan queue with OpenXR rendering, reducing worker latency by 20.2% when using two chunks compared to one.

Significance and Claims

The paper claims that general consumer graphics APIs (Vulkan and Metal) are sufficient substrates for a compact, shared train-to-deploy stack that achieves useful, and sometimes vendor-competitive, performance.

• Not a Replacement for Research: The author explicitly states Meganeura is not a replacement for dynamic Python research workflows, distributed cloud-scale training, or the full operator breadth of mature frameworks like PyTorch.
• Specific Gaps Identified: The remaining performance gaps are attributed to kernel coverage (specifically convolution derivatives on NVIDIA and attention backward on Metal) and scheduling policies (cooperative matrix selection), not to the graphics APIs themselves.
• Portability vs. Performance: The results suggest that while a single implementation may not match peak vendor performance on every hardware configuration, it offers superior performance portability for minimal-batch latency and significantly reduces the engineering overhead of maintaining separate training and deployment stacks.
• Live Adaptation: The work establishes the technical substrate required for "live adaptation" on deployed devices, where models can be fine-tuned or personalized in-place without leaving the application's codebase.

The paper concludes that the "split stack" problem can be solved by a unified native compiler, shifting the engineering burden from managing multiple conversion boundaries to optimizing a single, inspectable compiler and artifact.