Programming Languages # Mojo as a Python Superset: Performance on Production ML Workloads 2026-02-119 min read Lucio Durán Engineering Manager & AI Solutions Architect Also available in: [Español]() ## Mojo Language Fundamentals [Mojo]() is a programming language built on MLIR (Multi-Level Intermediate Representation), the compiler infrastructure that also powers TensorFlow, PyTorch compiler graphs, and LLVM. It's designed as a Python superset, you can import Python modules directly, with systems programming features: ownership, SIMD types, compile-time metaprogramming, and manual memory management. Representative Mojo code: ```mojo from math import sqrt from sys.info import simdwidthof from algorithm import vectorize alias float_type = DType.float32 alias simd_width = simdwidthof[float_type]() fn euclidean_distance( a: DTypePointer[float_type], b: DTypePointer[float_type], n: Int ) -> Float32: var sum = SIMD[float_type, simd_width](0) @parameter fn compute_chunk[simd_width: Int](i: Int): let diff = a.load[width=simd_width](i) - b.load[width=simd_width](i) sum += diff * diff vectorize[compute_chunk, simd_width](n) return sqrt(sum.reduce_add()) ``` A few things to unpack: 1. **` alias `** defines compile-time constants. ` simd_width ` is resolved at compile time based on your CPU's SIMD register width (4 for SSE, 8 for AVX2, 16 for AVX-512). 2. **` SIMD[float_type, simd_width] `** is an explicit SIMD vector type. When you write ` a.load[width=simd_width](i) `, you're loading 8 floats at once into a 256-bit register (on AVX2). 3. **` @parameter `** makes the inner function a compile-time parameter, enabling loop unrolling and specialization. 4. **` vectorize `** is a standard library function that handles the tail loop (when ` n ` isn't a multiple of ` simd_width `). This is not Python with type hints. This is a systems language that looks like Python. ## The Ownership Model: Not Rust, Not C++ Mojo has ownership semantics, but they're deliberately simpler than Rust's borrow checker. Three modes: ```mojo # `owned` - the function takes ownership, caller can't use the value after fn consume(owned data: Tensor[float_type]): # data is ours, will be destroyed when function ends process(data) # `borrowed` (default) - immutable reference, no copy fn inspect(data: Tensor[float_type]): # data is read-only, still owned by caller print(data.shape()) # `inout` - mutable reference fn modify(inout data: Tensor[float_type]): # data can be modified in place data *= 2.0 fn main(): var tensor = Tensor[float_type](1000, 1000) inspect(tensor) # borrowed, tensor still valid modify(tensor) # inout, tensor modified in place consume(tensor^) # owned, tensor moved (^ is the transfer operator) # tensor is no longer valid here ``` The ` ^ ` transfer operator is Mojo's version of ` std::move `. It's explicit, you can see every ownership transfer in the code. This is simpler than Rust's implicit moves and explicit borrows, at the cost of less safety. Mojo won't catch use-after-move at compile time in all cases the way Rust does. ## Real Benchmarks: ML Inference Pipeline The following benchmark uses a text embedding pipeline, the kind of thing you'd use for RAG (Retrieval Augmented Generation) in production. The pipeline: tokenize text, run it through a transformer model, normalize the output embeddings, compute cosine similarities against a vector database. **Test setup:** - Hardware: c6i.4xlarge (16 vCPU, 32GB, Intel Ice Lake with AVX-512) - Model: sentence-transformers/all-MiniLM-L6-v2 (384-dim embeddings) - Dataset: 10,000 text passages, average 128 tokens - Task: Embed all passages and compute top-10 nearest neighbors for 100 queries ### Embedding Computation ```mojo # Mojo: Batched embedding normalization + cosine similarity from tensor import Tensor from algorithm import vectorize, parallelize from math import sqrt fn normalize_embeddings( inout embeddings: Tensor[DType.float32], num_vectors: Int, dim: Int ): @parameter fn normalize_row(row: Int): var norm = SIMD[DType.float32, 8](0) let offset = row * dim @parameter fn accumulate[simd_width: Int](i: Int): let v = embeddings.load[width=simd_width](offset + i) norm += (v * v).reduce_add() vectorize[accumulate, 8](dim) let inv_norm = 1.0 / sqrt(norm.reduce_add()) @parameter fn scale[simd_width: Int](i: Int): embeddings.store[width=simd_width]( offset + i, embeddings.load[width=simd_width](offset + i) * inv_norm ) vectorize[scale, 8](dim) parallelize[normalize_row](num_vectors) fn cosine_similarity_batch( queries: Tensor[DType.float32], database: Tensor[DType.float32], num_queries: Int, num_db: Int, dim: Int, inout results: Tensor[DType.float32] ): @parameter fn process_query(q: Int): let q_offset = q * dim for d in range(num_db): let d_offset = d * dim var dot = SIMD[DType.float32, 8](0) @parameter fn dot_product[simd_width: Int](i: Int): dot += ( queries.load[width=simd_width](q_offset + i) * database.load[width=simd_width](d_offset + i) ).reduce_add() vectorize[dot_product, 8](dim) results.store(q * num_db + d, dot.reduce_add()) parallelize[process_query](num_queries) ``` ### Benchmark Results | Operation | Python (NumPy) | Python (PyTorch) | Mojo | Rust (ndarray) | C++ (Eigen) | | --- | --- | --- | --- | --- | --- | | Tokenization (10K texts) | 2.8s | 2.8s | 2.9s\* | N/A | N/A | | Embedding (10K, batched) | 4.2s | 3.8s | 3.9s\*\* | N/A | N/A | | Normalize (10K x 384) | 12ms | 8ms | 1.4ms | 1.6ms | 1.3ms | | Cosine sim (100 x 10K) | 89ms | 42ms | 6.1ms | 6.8ms | 5.9ms | | Top-10 selection | 34ms | 28ms | 4.2ms | 4.8ms | 4.1ms | | **End-to-end** | **7.13s** | **6.70s** | **6.55s** | - | - | \*Mojo tokenization calls the Python tokenizers library via interop \*\*Mojo embedding uses PyTorch model via interop The headline: Mojo is 6-8x faster than NumPy and 2-4x faster than PyTorch on the pure numerical operations (normalize, cosine similarity, top-k selection). But it's within 10% of Rust and C++ on those same operations. And the end-to-end pipeline time is dominated by the model inference step, which everyone calls PyTorch for. This is the reality nobody talks about. For the operations where Mojo shines, vectorized math on dense tensors, you're competing with highly optimized BLAS/LAPACK implementations that NumPy and PyTorch already use via MKL or OpenBLAS. Mojo's advantage is when you have custom compute kernels that can't be expressed as standard BLAS calls. ## Where Mojo Actually Wins **Custom preprocessing pipelines.** A text cleaning pipeline that applies regex, Unicode normalization, token filtering, and feature hashing. In Python, this was 340ms for 10K documents. In Mojo with parallelization: 28ms. A 12x speedup, because the Python version couldn't easily parallelize the regex + string manipulation. ```mojo from python import Python from algorithm import parallelize fn preprocess_batch( texts: PythonObject, num_texts: Int ) -> Tensor[DType.float32]: var features = Tensor[DType.float32](num_texts, 768) @parameter fn process_text(i: Int): let text = texts[i] # Clean text (Mojo string operations) var cleaned = text.lower() cleaned = remove_special_chars(cleaned) cleaned = normalize_unicode(cleaned) # Feature hashing (pure Mojo, vectorized) let tokens = tokenize_whitespace(cleaned) hash_features(tokens, features, i) parallelize[process_text](num_texts) return features ``` **Custom loss functions and metrics.** If you're writing research code with non-standard loss functions, Mojo lets you write them with explicit SIMD without dropping to C++: ```mojo fn focal_loss( predictions: Tensor[DType.float32], targets: Tensor[DType.float32], gamma: Float32, alpha: Float32, n: Int ) -> Float32: var total_loss = SIMD[DType.float32, 8](0) @parameter fn compute[simd_width: Int](i: Int): let p = predictions.load[width=simd_width](i) let t = targets.load[width=simd_width](i) # Focal loss: -alpha * (1-pt)^gamma * log(pt) let pt = t * p + (1 - t) * (1 - p) let focal_weight = (1 - pt) ** gamma let ce = -(t * log(p) + (1 - t) * log(1 - p)) total_loss += alpha * focal_weight * ce vectorize[compute, 8](n) return total_loss.reduce_add() / n ``` ## The Python Interop: Seamless but Not Free The interop story is Mojo's biggest selling point and its most misunderstood feature. You can import any Python module: ```mojo from python import Python fn main() raises: let np = Python.import_module("numpy") let torch = Python.import_module("torch") let transformers = Python.import_module("transformers") # This works, you're calling Python objects from Mojo let model = transformers.AutoModel.from_pretrained("bert-base-uncased") let tokenizer = transformers.AutoTokenizer.from_pretrained("bert-base-uncased") let inputs = tokenizer("Hello world", return_tensors="pt") let outputs = model(**inputs) # But this crosses the Python/Mojo boundary # Every call has marshalling overhead let embeddings = outputs.last_hidden_state.detach().numpy() # Convert to Mojo tensor for fast operations var mojo_tensor = numpy_to_mojo_tensor(embeddings) normalize_embeddings(mojo_tensor, 1, 768) ``` The overhead of crossing the Python/Mojo boundary is ~2-5 microseconds per call. That's negligible for coarse operations (model inference) but devastating if you do it in a tight loop. Measurement results: ``` Boundary crossing overhead (per call): Simple scalar: 2.1 µs Small array (100 elements): 3.8 µs Large array (1M elements): 4.2 µs + copy time ``` The array copy is the real cost for large tensors. Mojo can share memory with NumPy arrays (zero-copy), but you need to use ` DTypePointer ` and manage the lifetime carefully. ## Mojo vs Rust: The Honest Comparison Here is a direct comparison: **Mojo wins on:** Python ecosystem access, ML-specific ergonomics (SIMD types, tensor operations, ` @parameter ` metaprogramming), ease of adoption for Python developers. **Rust wins on:** Memory safety guarantees (Mojo's ownership is simpler but less safe), ecosystem maturity (crates.io has 140K+ crates), tooling (cargo, clippy, rust-analyzer are best-in-class), cross-platform deployment. **They tie on:** Raw numerical performance. Both compile to LLVM IR (Mojo via MLIR). For the same algorithm with explicit SIMD, you'll get within 5-10% of each other. The real question is: who writes the code? For teams of Python developers building ML infrastructure, Mojo is the path of least resistance. For teams of systems programmers building general-purpose infrastructure, Rust is the better choice. The language capabilities overlap, but the developer experience diverges sharply. ## What's Missing Mojo is pre-1.0 and it shows: **Standard library is sparse.** No async runtime, no HTTP client, no database drivers. You call Python for all of that, which works but negates the performance advantage. **Error messages are cryptic.** The MLIR-based compiler produces error messages that reference intermediate representations instead of source code. This is improving but still rough. **Debugger support is limited.** No stepping debugger of production quality exists yet. Print debugging is the norm. **Package management is nascent.** Modular has a package system, but the community package ecosystem is tiny compared to pip or cargo. **Closed-source compiler.** The Mojo standard library is open-source, but the compiler is proprietary. This is a significant risk factor for long-term adoption. If Modular goes under, you're stuck. ## Assessment Mojo is the most interesting new language for ML infrastructure in years. But the "Python superset" framing oversells it. In practice, the pattern is: Mojo for the hot paths and Python for everything else. That's a valid and useful pattern, it's what Cython and C extensions have done for decades, but Mojo does it with much better ergonomics. If you're building ML inference servers, custom data pipelines, or numerical computing libraries, Mojo is worth serious evaluation. If you're building web servers, CLI tools, or general-purpose applications, look elsewhere. The ecosystem isn't there yet, and the Python interop overhead means you're not getting Mojo-speed for anything that touches the Python world. The 68,000x number is misleading without context. The real speedup on production ML workloads is 2-10x over optimized Python, and 0.9-1.1x compared to C++/Rust. That 2-10x matters when it saves GPU hours, but it is not the revolution the marketing suggests. It is an evolution, and a good one. ## Frequently asked questions ### Is Mojo a replacement for Python? No. Mojo is a superset, valid Python is (mostly) valid Mojo. The value proposition is writing performance-critical code in Mojo while keeping your existing Python ecosystem. Think of it as 'Python for the 95% of your code that doesn't need to be fast, Mojo for the 5% that does.' ### How does Mojo compare to Cython or Numba? Cython and Numba are Python accelerators that add type annotations for JIT/AOT compilation. Mojo is a new language with its own compiler (MLIR-based) that happens to be Python-compatible. Mojo gives you ownership semantics, SIMD intrinsics, and compile-time metaprogramming that Cython/Numba can't match. ### Should I learn Mojo for ML engineering? If you write custom kernels, data preprocessing pipelines, or inference servers, yes, it's worth exploring. If you primarily use PyTorch/JAX at a high level, Mojo won't change your workflow much yet. The ecosystem is still young. ## Tips 1. **Tip:** Start Mojo files with ` from python import Python ` to call any Python library directly. NumPy, PyTorch, and scikit-learn all work, but the interop has marshalling overhead, so don't call Python in hot loops. 2. **Tip:** Use ` SIMD[DType.float32, 8] ` for explicit vectorization instead of relying on auto-vectorization. Mojo's SIMD types map directly to hardware SIMD registers, you control the width, the compiler generates the instructions. 3. **Tip:** The ` @parameter ` decorator in Mojo evaluates at compile time, similar to Zig's comptime. Use it for loop unrolling and specialization: ` @parameter for i in range(4): ` unrolls the loop at compile time. - mojo - python - mlir - simd - performance - machine-learning - systems-programming - rust - cpp - benchmarks [Back to Blog]()