Skip to content

Mojo GPU Programming — Knowledge Base Article

Source: YouTube — "Mojo GPU Programming" series Tags: mojo, gpu, cuda, kernel, parallel-computing, ai-infra Status: Stable

Overview

Mojo is a programming language purpose-built for AI workloads. It combines Python's simplicity and ecosystem with C/C++-level performance, and—critically—runs on both CPUs and GPUs without requiring separate toolchains or hardware-specific rewrites.

The core philosophy: write once, deploy anywhere (CPU, NVIDIA GPU, AMD GPU) using one language.

Why Mojo for GPU Programming?

Problem Mojo Solution
Python + CUDA requires multiple toolchains Single language, single compiler
Production optimisation requires TensorRT, ONNX, hardware-specific code Native GPU kernel support
Multiple hardware targets → multiple codebases Write once, target any accelerator

GPU Architecture Fundamentals (Context)

A Mojo GPU program uses the same conceptual model as CUDA:

  • Thread — single smallest unit of compute
  • Block — group of threads sharing memory and executing concurrently
  • Grid — group of blocks covering the entire workload
Grid
├── Block 0
│   ├── Thread 0, Thread 1, Thread 2, Thread 3
├── Block 1
│   ├── Thread 0, Thread 1, Thread 2, Thread 3
├── Block 2
│   ├── Thread 0, Thread 1, Thread 2, Thread 3
└── Block 3
    ├── Thread 0, Thread 1, Thread 2, Thread 3

Blocks and threads can be organised in 1D, 2D, or 3D layouts for natural data mapping (e.g., 2D for matrix operations).

CPU vs GPU

CPU GPU
Few cores, sophisticated per-core logic Thousands/millions of simple cores
Handles DB, networking, UI Handles parallel math (matrix multiply, division)
Fewer threads, complex work per thread Massive parallelism, simple work per thread

Host / Device Model

  • Host = CPU (issues commands, consumes results)
  • Device = GPU (executes kernels in parallel)

Pipeline: CPU launches kernel → GPU executes on N threads → results returned to CPU

Writing a GPU Kernel in Mojo

1. Import and detect accelerator

from module import necessary_libs

# Check for GPU presence
if has_accelerator():
    print("GPU detected")

2. Define the kernel

A kernel is a regular Mojo function that runs on the GPU:

fn print_threads():
    # Access thread ID and block ID
    thread_id = ...  # built-in
    block_id  = ...  # built-in
    print(f"Block {block_id}, Thread {thread_id}")

3. Launch the kernel

# Get reference to the GPU (device context)
device = get_device()

# Queue the kernel with thread/block configuration
# 4 blocks × 4 threads = 16 total threads
launch_kernel(print_threads, blocks=4, threads_per_block=4)

# Wait for all threads to complete
synchronize()

# Print results
print("All threads completed")

4. Thread/Block layout

The configuration blocks=4, threads_per_block=4 yields:

Block 0: Threads 0,1,2,3
Block 1: Threads 0,1,2,3
Block 2: Threads 0,1,2,3
Block 3: Threads 0,1,2,3

2D layout extends this so each block is addressed as (block_x, block_y) and each thread within it as (thread_x, thread_y). 3D adds a z dimension.

Key Takeaways

  1. Mojo is the one language for AI — runs on all hardware, Python syntax, C speed.
  2. Kernel = function that executes on the GPU.
  3. Thread = smallest GPU execution unit.
  4. Block = group of threads (share memory, execute together).
  5. Grid = collection of blocks covering the entire compute workload.
  6. Thread/block count matters — must match data size to maximise efficiency.
  7. Always call synchronize() after launching a kernel to collect results before downstream processing.

Deployment Workflow

Research / Training (Mojo / Python)
        │
        ▼
One Mojo codebase
        │
        ├── CPU target
        ├── NVIDIA GPU target
        └── AMD GPU target
        │
        ▼
Production inference — no rewrites needed

References