5.2 CUDA Online Softmax implementation and optimization

Traditional Safe Softmax needs three passes over the input. Online Softmax uses a neat recurrence so that max and sum are maintained in a single pass, cutting memory reads from 3 down to 2. This article derives the algorithm, then implements and optimizes the CUDA Kernel step by step — the core prerequisite for understanding FlashAttention.

📑 Contents

• 1. Why Online Softmax
• 2. Online Softmax algorithm derivation
• 3. Version V0: single-thread Online Softmax
• 4. Version V1: Block-level parallelism + Warp Shuffle merge reduce
• 5. Version V2: vectorized float4 loads
• 6. Version V3: register cache to remove the second read
• 7. Version V4: multi-row parallelism + Grid Stride
• 8. Performance comparison and the link to FlashAttention
• Summary
• Self-check list
• References

---

1. Why Online Softmax

1.1 The cost of three passes

In the previous article, Safe Softmax used three passes:

$$\underbrace{m = \max_j(x_j)}_{\text{第1遍}} \quad \rightarrow \quad \underbrace{d = \sum_j e^{x_j - m}}_{\text{第2遍}} \quad \rightarrow \quad \underbrace{y_i = \frac{e^{x_i - m}}{d}}_{\text{第3遍}}$$

Each pass reads the input once, so total memory traffic is $3N$. For a Memory-Bound op, dropping one pass is roughly a 33% gain. The catch: the second pass (sum) needs the first pass’s result (max), so they look hard to merge.

Analogy: you want to normalize everyone’s scores in a class. The usual way is: flip through the roster once for the highest score, flip again for the adjusted total, then flip a third time to write the normalized scores — three full passes. Can you do it in two, or even one?

1.2 The core insight

The key insight of Online Softmax is that max and sum can be maintained in a single pass, at the cost of “correcting” the accumulated sum whenever a new maximum appears.

Imagine you flip through the roster while accumulating “adjusted scores.” At person 50 you find a score higher than the previous max — every prior adjusted score is now off. But you know exactly how much: each previous $e^{x_i - m_{\text{old}}}$ should become $e^{x_i - m_{\text{new}}}$, so the total only needs a correction factor $e^{m_{\text{old}} - m_{\text{new}}}$.

That is the essence of Online Softmax — correct with a multiply instead of recomputing.

---

2. Online Softmax algorithm derivation

2.1 Deriving the recurrence

Suppose the first $k$ elements are already processed, and the maintained state is:

• $m_k = \max(x_0, x_1, \ldots, x_{k-1})$ (max of the first $k$ elements)
• $d_k = \sum_{j=0}^{k-1} e^{x_j - m_k}$ (exp-sum under the current max)

When processing element $k+1$, namely $x_k$:

Step 1: update the max

$$m_{k+1} = \max(m_k, x_k)$$

Step 2: correct the existing sum and add the new element

$$d_{k+1} = d_k \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}}$$

Derivation:

$$\begin{aligned} d_{k+1} &= \sum_{j=0}^{k} e^{x_j - m_{k+1}} \end{aligned}$$ $$\begin{aligned} &= \sum_{j=0}^{k-1} e^{x_j - m_{k+1}} + e^{x_k - m_{k+1}} \end{aligned}$$ $$\begin{aligned} &= \sum_{j=0}^{k-1} e^{x_j - m_k} \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}} \end{aligned}$$ $$\begin{aligned} &= d_k \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}} \end{aligned}$$

💡 Tip: the key is the correction factor $e^{m_k - m_{k+1}}$. If the new element is not a new max ($m_{k+1} = m_k$), then $e^{m_k - m_{k+1}} = e^0 = 1$, and the correction collapses to a simple add — no extra cost.

2.2 Full algorithm pseudocode

// Pass 1: Online compute max and sum together
m = -∞
d = 0
for j = 0 to N-1:
    m_new = max(m, x[j])
    d = d * exp(m - m_new) + exp(x[j] - m_new)
    m = m_new

// Pass 2: normalize output
for j = 0 to N-1:
    y[j] = exp(x[j] - m) / d

Three passes become two: the first pass gets both $m$ and $d$, the second does normalization. Memory reads drop from $3N$ to $2N$, a theoretical 33% gain.

2.3 Numerical stability

Online Softmax has the same numerical stability as Safe Softmax:

• $m$ is always the max of the elements seen so far, so every exp argument $x_j - m \leq 0$ and cannot overflow
• The correction factor $e^{m_k - m_{k+1}} \leq 1$ (because $m_{k+1} \geq m_k$), so it cannot overflow either
• The final $d$ is mathematically equivalent to the three-pass Safe Softmax result

2.4 Test environment

Same as the previous article: A100 80GB SXM4, test matrix $(M, N) = (4096, 4096)$.

---

3. Version V0: single-thread Online Softmax

3.1 Approach

The most direct implementation: each thread handles one row, serially maintains $(m, d)$, then does one normalization pass.

3.2 Kernel implementation

// V0: single-thread Online Softmax (one thread per row)
__global__ void online_softmax_v0(float* input, float* output, int M, int N) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M) return;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Pass 1: Online compute max and sum together
    float m = -INFINITY;
    float d = 0.0f;

    for (int i = 0; i < N; i++) {
        float xi = x[i];
        float m_new = fmaxf(m, xi);
        d = d * expf(m - m_new) + expf(xi - m_new);
        m = m_new;
    }

    // Pass 2: normalize
    float inv_d = 1.0f / d;
    for (int i = 0; i < N; i++) {
        y[i] = expf(x[i] - m) * inv_d;
    }
}

3.3 Performance analysis

Like Safe Softmax V0, one thread per full row has severe under-parallelism and uncoalesced access. Compared with three-pass V0, though, the input is only read twice.

⚠️ Note: each Online recurrence step needs an expf for the correction factor, which adds compute. Softmax is Memory-Bound, so in high-performance versions that extra compute is fully hidden behind memory latency.

V0 measured bandwidth utilization: ~7% (~143 GB/s) — only a correctness baseline.

---

4. Version V1: Block-level parallelism + Warp Shuffle merge reduce

4.1 Parallelization challenge

The Online Softmax recurrence is serial — each step’s $(m, d)$ depends on the previous one. How to parallelize?

Key observation: let each thread process a subset of elements and maintain its own local $(m_{\text{local}}, d_{\text{local}})$, then merge those local results.

Merge of two local results $(m_1, d_1)$ and $(m_2, d_2)$:

$$\begin{aligned} m_{\text{merged}} &= \max(m_1, m_2) \end{aligned}$$ $$\begin{aligned} d_{\text{merged}} &= d_1 \cdot e^{m_1 - m_{\text{merged}}} + d_2 \cdot e^{m_2 - m_{\text{merged}}} \end{aligned}$$

This merge is associative (tree reduce works), so it parallelizes well.

4.2 Merge reduce with Warp Shuffle

Following the Safe Softmax optimizations from the previous article, reduction is best done in registers: Warp Shuffle has low latency, no Bank Conflict, and natural warp sync, so it beats Shared Memory reduce.

Unlike a plain Reduce, Online Softmax’s reduce is not a simple max or sum — it merges $(m, d)$ together. __shfl_down_sync can shuffle the two values separately, then each round applies the merge formula above. The overall structure still uses the two-level “Warp Shuffle + Shared Memory across warps” pattern.

4.3 Kernel implementation

// Online merge reduce within a Warp
__device__ void warpReduceOnline(float& m, float& d, const int Warp_Threads = 32) {
    for (int offset = Warp_Threads / 2; offset > 0; offset >>= 1) {
        float m2 = __shfl_down_sync(0xffffffff, m, offset);
        float d2 = __shfl_down_sync(0xffffffff, d, offset);

        float m_new = fmaxf(m, m2);
        d = d * expf(m - m_new) + d2 * expf(m2 - m_new);
        m = m_new;
    }
}

// V1: Block-level parallel + Warp Shuffle two-level merge reduce
__global__ void online_softmax_v1(float* input, float* output, int M, int N, const int Warp_Size) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % Warp_Size;
    int wid  = tid / Warp_Size;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Pass 1: each thread Online-processes a segment
    float local_m = -INFINITY;
    float local_d = 0.0f;

    for (int i = tid; i < N; i += blockDim.x) {
        float xi = x[i];
        float m_new = fmaxf(local_m, xi);
        local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
        local_m = m_new;
    }

    // Level 1: merge within Warp
    warpReduceOnline(local_m, local_d, Warp_Size);

    // Exchange across Warps via Shared Memory
    __shared__ float warp_m[32];
    __shared__ float warp_d[32];

    if (lane == 0) {
        warp_m[wid] = local_m;
        warp_d[wid] = local_d;
    }
    __syncthreads();

    // Level 2: Warp 0 does the final merge
    int num_warps = blockDim.x / Warp_Size;
    if (wid == 0) {
        local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
        local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
        warpReduceOnline(local_m, local_d, num_warps);
    }

    // Broadcast final result
    __shared__ float final_m, final_d;
    if (tid == 0) {
        final_m = local_m;
        final_d = local_d;
    }
    __syncthreads();

    float row_max = final_m;
    float inv_sum = 1.0f / final_d;

    // Pass 2: normalize output
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - row_max) * inv_sum;
    }
}

In the second-level Warp reduce, watch the offset bounds: when lane > num_warps, local_m = -INFINITY. Suppose num_warps = 8 and offset = 16: thread 9 gets thread 25’s m2 = -INFINITY, while thread 9 itself also has m = -INFINITY, so m_new = -INFINITY, m - m_new = -INFINITY - (-INFINITY) = NaN, expf(NaN) = NaN, and the final sum becomes Nan, so the result is NaN.

4.4 Comparison with three-pass Safe Softmax

⚠️ Note: merge order affects floating-point precision (addition is not strictly associative). Softmax is insensitive to tiny errors, so tree-reduce precision is fine.

💡 Tip: Online reduce is heavier than a plain sum-reduce (2× expf per step), but the saved memory bandwidth dwarfs the extra compute. For Memory-Bound ops, removing one global read is decisive.

V1 measured bandwidth utilization: ~58% (~1183 GB/s), about 8.3× over single-thread V0, and about 1.4× over the matching three-pass Safe Softmax version.

---

5. Version V2: vectorized float4 loads

5.1 Improvement idea

Same optimization path as Safe Softmax: use float4 vectorized loads to cut instruction count. In the Online recurrence phase, load 4 elements at a time and update $(m, d)$ one by one.

5.2 Kernel implementation

// V2: float4 vectorized load + Online Softmax
__global__ void online_softmax_v2(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    float* x = input  + row * N;
    float* y = output + row * N;

    float4* x4 = reinterpret_cast<float4*>(x);
    float4* y4 = reinterpret_cast<float4*>(y);
    int N4 = N / 4;

    // Pass 1: vectorized Online recurrence
    float local_m = -INFINITY;
    float local_d = 0.0f;

    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];

        // Update (m, d) for 4 elements in order
        float vals[4] = {data.x, data.y, data.z, data.w};
        for (int k = 0; k < 4; k++) {
            float m_new = fmaxf(local_m, vals[k]);
            local_d = local_d * expf(local_m - m_new) + expf(vals[k] - m_new);
            local_m = m_new;
        }
    }
    // Tail handling
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        float xi = x[i];
        float m_new = fmaxf(local_m, xi);
        local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
        local_m = m_new;
    }

    // Two-level Warp Shuffle merge (same as V1)
    warpReduceOnline(local_m, local_d);

    __shared__ float warp_m[32];
    __shared__ float warp_d[32];
    if (lane == 0) {
        warp_m[wid] = local_m;
        warp_d[wid] = local_d;
    }
    __syncthreads();

    int num_warps = blockDim.x / 32;
    if (wid == 0) {
        local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
        local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
        warpReduceOnline(local_m, local_d);
    }

    __shared__ float final_m, final_d;
    if (tid == 0) {
        final_m = local_m;
        final_d = local_d;
    }
    __syncthreads();

    float row_max = final_m;
    float inv_sum = 1.0f / final_d;

    // Pass 2: vectorized normalize
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        float4 result;
        result.x = expf(data.x - row_max) * inv_sum;
        result.y = expf(data.y - row_max) * inv_sum;
        result.z = expf(data.z - row_max) * inv_sum;
        result.w = expf(data.w - row_max) * inv_sum;
        y4[i] = result;
    }
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - row_max) * inv_sum;
    }
}

5.3 How float4 works with the Online recurrence

float4 loads 4 elements at once, but the Online recurrence still updates $(m, d)$ element by element. That looks like no win, but in practice:

1. Memory: float4 moves the same data with fewer instructions, easing instruction-issue pressure
2. Compute: four expf calls can pipeline; the compiler can schedule them better
3. Overall: loop iterations drop 4×, so loop overhead (branches, address math) falls a lot

V2 measured bandwidth utilization: ~70% (~1427 GB/s), about 1.21× over V1.

---

6. Version V3: register cache to remove the second read

6.1 Ultimate optimization idea

Two passes still mean the input is read twice. Can we read only once?

Yes, if each row is not too long ($N / \text{blockDim} \leq$ a few dozen): each thread handles a bounded number of elements, so data from the first pass can be stashed in registers, and the second pass normalizes from registers — no second global memory read.

That is a true one-pass Softmax — one global read and one write.

6.2 Kernel implementation

// V3: register cache for One-Pass Softmax
// For small N / blockDim.x (elements per thread <= MAX_ELEMS)
#define MAX_ELEMS_PER_THREAD 32

__global__ void online_softmax_v3(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Register array caches input data
    float reg_cache[MAX_ELEMS_PER_THREAD];

    // Pass 1 (only global read): load + Online recurrence
    float local_m = -INFINITY;
    float local_d = 0.0f;
    int count = 0;

    for (int i = tid; i < N; i += blockDim.x) {
        float xi = x[i];
        reg_cache[count] = xi;  // cache into registers
        count++;

        float m_new = fmaxf(local_m, xi);
        local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
        local_m = m_new;
    }

    // Two-level Warp Shuffle merge
    warpReduceOnline(local_m, local_d);

    __shared__ float warp_m[32];
    __shared__ float warp_d[32];
    if (lane == 0) {
        warp_m[wid] = local_m;
        warp_d[wid] = local_d;
    }
    __syncthreads();

    int num_warps = blockDim.x / 32;
    if (wid == 0) {
        local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
        local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
        warpReduceOnline(local_m, local_d);
    }

    __shared__ float final_m, final_d;
    if (tid == 0) {
        final_m = local_m;
        final_d = local_d;
    }
    __syncthreads();

    float row_max = final_m;
    float inv_sum = 1.0f / final_d;

    // Read from register cache and write out (no second global read)
    int idx = 0;
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(reg_cache[idx] - row_max) * inv_sum;
        idx++;
    }
}

6.3 Register pressure analysis

⚠️ Note: the hard limit is register capacity. Each thread uses MAX_ELEMS_PER_THREAD float registers as cache:

A100 has 65536 32-bit registers per SM, and a per-thread cap of 255. If MAX_ELEMS_PER_THREAD is too large, registers spill to Local Memory (which is really global memory) and things get slower.

📌 Key point: V3 fits $N \leq 8192$ with blockDim=256. For larger $N$, fall back to V2 (two-pass).

V3 measured bandwidth utilization: ~82% (~1672 GB/s) ($N=4096$, blockDim=256), about 1.17× over V2.

---

7. Version V4: multi-row parallelism + Grid Stride

7.1 When it applies

Earlier versions use “one Block per row” and launch $M$ Blocks. When $M$ is huge (tens of thousands of rows), that causes two problems: Grid scheduling overhead grows linearly, and each Block exits after one row so Shared Memory buffers and register state cannot be reused across rows.

The fix is Grid Stride: a fixed number of Blocks (usually a multiple of the SM count) loop over all rows, still stacking the earlier optimizations (vectorized loads + Warp Shuffle merge reduce). Algorithmically V4 is identical to V2 (two passes); the only difference is Block-to-row mapping — V2 is 1:1, V4 is 1:many.

⚠️ Note: Grid Stride does not fix the case where $M$ is tiny (e.g. batch=1) and $N$ is huge, so Block count is too low — that needs multiple Blocks cooperating on one row (segmented reduce within a row + two-phase Kernel or atomic merge).

7.2 Kernel implementation

// V4: Grid Stride multi-row + all optimizations
__global__ void online_softmax_v4(float* input, float* output, int M, int N) {
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    // Grid Stride over all rows
    for (int row = blockIdx.x; row < M; row += gridDim.x) {
        float* x = input  + row * N;
        float* y = output + row * N;

        float4* x4 = reinterpret_cast<float4*>(x);
        float4* y4 = reinterpret_cast<float4*>(y);
        int N4 = N / 4;

        // Pass 1: vectorized Online recurrence
        float local_m = -INFINITY;
        float local_d = 0.0f;

        for (int i = tid; i < N4; i += blockDim.x) {
            float4 data = x4[i];
            float vals[4] = {data.x, data.y, data.z, data.w};
            for (int k = 0; k < 4; k++) {
                float m_new = fmaxf(local_m, vals[k]);
                local_d = local_d * expf(local_m - m_new) + expf(vals[k] - m_new);
                local_m = m_new;
            }
        }
        for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
            float xi = x[i];
            float m_new = fmaxf(local_m, xi);
            local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
            local_m = m_new;
        }

        // Two-level Warp Shuffle merge
        warpReduceOnline(local_m, local_d);

        __shared__ float warp_m[32];
        __shared__ float warp_d[32];
        if (lane == 0) {
            warp_m[wid] = local_m;
            warp_d[wid] = local_d;
        }
        __syncthreads();

        int num_warps = blockDim.x / 32;
        if (wid == 0) {
            local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
            local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
            warpReduceOnline(local_m, local_d);
        }

        __shared__ float final_m, final_d;
        if (tid == 0) {
            final_m = local_m;
            final_d = local_d;
        }
        __syncthreads();

        float row_max = final_m;
        float inv_sum = 1.0f / final_d;

        // Pass 2: vectorized normalize
        for (int i = tid; i < N4; i += blockDim.x) {
            float4 data = x4[i];
            float4 result;
            result.x = expf(data.x - row_max) * inv_sum;
            result.y = expf(data.y - row_max) * inv_sum;
            result.z = expf(data.z - row_max) * inv_sum;
            result.w = expf(data.w - row_max) * inv_sum;
            y4[i] = result;
        }
        for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
            y[i] = expf(x[i] - row_max) * inv_sum;
        }

        __syncthreads();  // Ensure Shared Memory is ready before the next row
    }
}

Launch pattern:

int num_sms;
cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, 0);
int grid_size = min(M, num_sms * 4);  // maximize GPU utilization
int block_size = 256;
online_softmax_v4<<<grid_size, block_size>>>(d_input, d_output, M, N);

7.3 Benefits of Grid Stride

When $M$ is large, Grid Stride lets Blocks loop over many rows:

• Fixed Grid size avoids oversized Grid scheduling overhead
• Shared Memory and register state inside a Block can be reused across rows
• The GPU stays fully occupied

When $M$ is small, grid_size = M, and behavior collapses to V2 — one Block per row.

V4 measured bandwidth utilization: ~72% (~1468 GB/s) ($M=4096, N=4096$), roughly on par with V2 (the theoretical ceiling of a two-pass scheme at large $N$). Grid Stride’s real win is scheduling overhead when $M$ is huge; at this article’s test scale ($M=4096$) the gain is limited.

---

8. Performance comparison and the link to FlashAttention

8.1 Version summary

📌 Key point: V3’s One-Pass wins when $N \leq 8192$, but is limited by register capacity. V2/V4 Two-Pass is more general and still hits 70%+ bandwidth across $N$ sizes.

8.2 Link to FlashAttention

FlashAttention’s core idea is a generalization of Online Softmax. Full Self-Attention is:

$$\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$$

The traditional path computes the full $QK^T$ matrix ($N \times N$), Softmax, then multiplies by $V$. That requires storing the whole $N \times N$ Attention matrix in HBM — 64GB when $N=128K$.

FlashAttention uses the Online Softmax recurrence: load $K$ and $V$ in tiles into SRAM (Shared Memory), compute partial $QK^T$, incrementally update Softmax state $(m, d)$, then correct the output. The process never stores the full Attention matrix.

graph LR
    A["Online Softmax<br>recurrence"] --> B["Tile QK^T<br>incrementally update (m, d)"]
    B --> C["Correct output O<br>O = O * correction + new_block"]
    C --> D["FlashAttention<br>IO-aware Attention"]

The Online Softmax correction factor $e^{m_{\text{old}} - m_{\text{new}}}$ is used in FlashAttention to correct prior tiles’ contribution to the output:

$$O_{\text{new}} = O_{\text{old}} \cdot \frac{d_{\text{old}}}{d_{\text{new}}} \cdot e^{m_{\text{old}} - m_{\text{new}}} + \frac{e^{m_{\text{block}} - m_{\text{new}}}}{d_{\text{new}}} \cdot P_{\text{block}} \cdot V_{\text{block}}$$

💡 Tip: once you understand this article’s Online Softmax recurrence and merge formula, FlashAttention’s tiling is just that idea plus an incremental matrix-multiply update — the core math is the same.

8.3 Engineering recommendations

---

📝 Summary

Online Softmax’s contribution is a compact recurrence that merges Safe Softmax’s three passes into two:

1. Math: $d_{k+1} = d_k \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}}$ maintains max and sum in one pass
2. Parallelism: $(m, d)$ merge is associative, so Warp Shuffle two-level reduce parallelizes efficiently (V1)
3. Approaching bandwidth limits: vectorized loads (V2) cut instructions; register cache (V3) reaches One-Pass; Grid Stride (V4) handles any batch
4. Path to FlashAttention: the Online Softmax recurrence is the mathematical foundation of FlashAttention’s tiled compute

From an engineering view, Online Softmax brings Memory-Bound memory traffic down to the theoretical minimum; with Warp Shuffle and vectorized loads, A100 reaches 70–82% bandwidth utilization.

---

🎯 Self-check list

• Explain which pass Online Softmax saves versus three-pass Safe Softmax, and why
• Manually derive $d_{k+1} = d_k \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}}$
• Write code that merges two local $(m_1, d_1)$ and $(m_2, d_2)$ into global $(m, d)$
• Explain why $(m, d)$ merge is associative and thus parallelizable
• Implement warpReduceOnline with __shfl_down_sync
• State V3 register-cache applicability and register-pressure limits
• Explain the math link between Online Softmax and FlashAttention tiling
• Choose a Softmax version given $M$ and $N$
• Use Nsight Compute to compare memory throughput of Two-Pass vs Three-Pass
• Extend the Online Softmax $(m, d)$ recurrence to other incremental-compute settings

---

📚 References

• Online normalizer calculation for softmax - Milakov & Gimelshein, 2018
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• From Online Softmax to FlashAttention - UW CSE 599M
• NVIDIA CUDA C++ Programming Guide - Warp Shuffle Functions
• GPU optimization series in depth: Softmax optimization
• FlashAttention core principles explained
• NVIDIA CUB Library