GPU puzzle challenge - Softmax multistage

dolewhip · July 11, 2025, 12:12am

I was looking at puzzle 16

And did the challenge 1.1, 1.2

I did a multi stage and an online_softmax. Didn’t do the two pass but saw it in the max kernels. maybe later do a two pass etc.

github.com/quinnavila/algo_land

SOFTMAX.md

main


---

> **TL;DR:** This document explains an efficient, multi-stage GPU softmax implementation in Mojo. It uses a generic, multi-pass parallel reduction to handle inputs larger than a single GPU block.
>
> **View the code: [`kernels/softmax.mojo`](kernels/softmax.mojo)**

### How to Run

To compile and run this example, first activate the project environment and then execute the Mojo file:

```bash
pixi shell
mojo kernels/softmax.mojo
```

---

### Implementation Overview

This file has been truncated. show original

Have fun

dolewhip · July 12, 2025, 7:34pm

I think I have challenge I done

I wasn’t entirely sure if 1.2 describes a two pass but I did 1, 2 as described. Also created progression kernels but plan to add some markdown.

Reduction progression (great reference material in pmpp ch10)

github.com/quinnavila/algo_land

REDUCTION.md

main

 
---

## TL;DR: `kernels/reduction.mojo`

> **TL;DR:** This document explains a generic, multi-pass parallel reduction algorithm in Mojo. It uses a shared-memory kernel and a host-side loop to efficiently find the sum or maximum of an array of any size.
>
> **View the code: [`kernels/reduction.mojo`](kernels/reduction.mojo)**

---

## A Scalable Parallel Reduction in Mojo

### Goal

To build a reusable GPU function, `reduce`, that can efficiently find the sum or maximum of an array of any size.

### The Algorithm: A Multi-Stage Tournament

The algorithm finds a single result from a large array by running a "tournament" in multiple stages.

This file has been truncated. show original

github.com/quinnavila/algo_land

kernels/reduction.mojo

main

from gpu.host import DeviceContext
from gpu.id import block_dim, block_idx, thread_idx
from gpu import barrier
from layout import Layout, LayoutTensor, UNKNOWN_VALUE, RuntimeLayout
from layout.tensor_builder import LayoutTensorBuild as tb
from math import ceildiv, inf
from random import random_float64
from time import perf_counter_ns
from testing import assert_almost_equal
from memory import UnsafePointer
from utils.index import Index

alias dtype = DType.float32
alias BLOCK_SIZE = 256
alias ELEMENTS_PER_THREAD = 8
alias ELEMENTS_PER_BLOCK = BLOCK_SIZE * ELEMENTS_PER_THREAD


fn reduction_kernel[
    op: fn (Float32, Float32) -> Float32,

This file has been truncated. show original

Challenge 1 Part 1

github.com/quinnavila/algo_land

kernels/softmax_layout_tensor.mojo

main

from gpu.host import DeviceContext
from gpu.id import block_dim, block_idx, thread_idx
from gpu import barrier
from layout import Layout, LayoutTensor, UNKNOWN_VALUE, RuntimeLayout
from layout.tensor_builder import LayoutTensorBuild as tb
from math import ceildiv, exp, inf
from random import random_float64
from time import perf_counter_ns
from testing import assert_almost_equal
from memory import UnsafePointer
from utils.index import Index

alias dtype = DType.float32
alias BLOCK_SIZE = 256
alias ELEMENTS_PER_THREAD = 8
alias ELEMENTS_PER_BLOCK = BLOCK_SIZE * ELEMENTS_PER_THREAD
alias static_layout = Layout.row_major(UNKNOWN_VALUE, 1)


fn reduction_kernel[

This file has been truncated. show original

Challenge I Part 2

github.com/quinnavila/algo_land

SOFTMAX_BATCHED.md

main


---

## Batched Softmax on the GPU: Row vs. Column

This document compares two strategies for applying softmax to a 2D tensor: **row-wise** and **column-wise**. The choice reveals how memory layout dictates GPU performance.

### The Core Challenge: Memory Access

For a matrix in standard **row-major layout**, accessing elements is not equal.

*   **Row-wise Access:** `[R0,C0], [R0,C1], [R0,C2]...` is **contiguous**. This is fast and ideal for GPUs.
*   **Column-wise Access:** `[R0,C0], [R1,C0], [R2,C0]...` is **strided**. This is slow and inefficient for GPUs.

This difference forces two completely different algorithmic strategies.

```mermaid
graph TD
    subgraph "Row-Major Memory"
        direction LR

This file has been truncated. show original

github.com/quinnavila/algo_land

kernels/softmax_batched_row.mojo

main

from gpu.host import DeviceContext
from gpu.id import block_dim, block_idx, thread_idx
from gpu import barrier
from layout import Layout, LayoutTensor
from layout.tensor_builder import LayoutTensorBuild as tb
from math import ceildiv, exp, inf
from random import random_float64
from time import perf_counter_ns
from testing import assert_almost_equal

alias dtype = DType.float32
alias ROWS = 128
alias COLS = 1024
alias BLOCK_SIZE = 256
alias ELEMENTS_PER_THREAD = ceildiv(COLS, BLOCK_SIZE)

alias static_layout = Layout.row_major(ROWS, COLS)


fn row_softmax_kernel(

This file has been truncated. show original

github.com/quinnavila/algo_land

kernels/softmax_batched_col.mojo

main

# File: batched_col.mojo

from gpu.host import DeviceContext
from gpu.id import block_dim, block_idx, thread_idx
from layout import Layout, LayoutTensor
from math import ceildiv, exp, inf
from random import random_float64
from time import perf_counter_ns
from testing import assert_almost_equal

alias dtype = DType.float32
alias ROWS = 128
alias COLS = 1024
alias BLOCK_SIZE_X = 32
alias BLOCK_SIZE_Y = 8

alias static_layout = Layout.row_major(ROWS, COLS)
alias col_stats_layout = Layout.row_major(1, COLS)

This file has been truncated. show original

Have Fun!

Topic		Replies	Views
Looking for examples of mulit-gpu usage with Mojo GPU Programming gpu	7	412	September 3, 2025
MAX' GPU ReduceOps for non-inner axes MAX discussion	7	97	August 12, 2025
Examples of custom CPU / GPU operations in Mojo MAX discussion , 24_6	29	1247	October 6, 2025
New GPU programming recipes GPU Programming gpu , modular-content	0	215	March 14, 2025
Examples of programming GPU functions using the Mojo MAX Driver API MAX discussion , gpu , 25_1	5	454	April 26, 2025

Related topics