Improving Online Softmax for GPU Beam Search

Overview

Implementing Online Softmax Normalization with some enhancements:

Improvement	Technique	Expected Speedup
1. Warp Shuffle	Replace shared memory reductions with __shfl_down_sync	10-20%
2. Vectorized Memory	Use float4 loads for 4× fewer memory transactions	10-30%
3. Mixed Precision	FP16/BF16 input with FP32 accumulation for 2× bandwidth	30-50%
4. Improved TopK	Bitonic sort network instead of insertion sort	2-3× for K>10

Project Structure

online-softmax-improved/
├── include/
│   ├── common.cuh           # Utilities, timing, verification
│   ├── baseline.cuh         # Original paper's algorithms
│   ├── warp_primitives.cuh  # Improvement #1: Warp shuffle
│   ├── vectorized.cuh       # Improvement #2: float4 access
│   ├── mixed_precision.cuh  # Improvement #3: FP16/BF16
│   └── topk_improved.cuh    # Improvement #4: Bitonic TopK
├── src/
│   └── benchmark.cu         # Performance benchmarks
├── tests/
│   └── correctness.cu       # Correctness verification
├── results/                  # Benchmark outputs
├── Makefile
└── README.md

Algorithms

Algorithm 3 - Online Softmax:

for each element x_i:
    m_new = max(m, x_i)
    d = d * exp(m - m_new) + exp(x_i - m_new)
    m = m_new

This computes the softmax normalizer in a single pass, enabling:

• 2 passes instead of 3 (vs safe softmax)
• 3 memory accesses per element instead of 4

Algorithm 4 - Fused Softmax + TopK:

• Compute online normalizer AND track top-K simultaneously
• Only 1 memory access per element (major win for beam search)

Improvements

1. Warp Shuffle Reductions

// Original: shared memory reduction
__shared__ float s_max[BLOCK_SIZE];
s_max[tid] = local_max;
__syncthreads();
for (int s = blockDim.x/2; s > 0; s >>= 1) { ... }

// Improved: warp shuffle (no shared memory)
float warp_max = warp_reduce_max(local_max);  // __shfl_down_sync

2. Vectorized Memory Access

// Original: scalar loads
float xi = x[i];

// Improved: vector loads (4× fewer transactions)
float4 vals = reinterpret_cast<const float4*>(x)[i];

3. Mixed Precision

// Load FP16, compute in FP32, store FP16
half2 vals = x2[i];
float v0 = __half2float(vals.x);
// ... FP32 accumulation ...
y2[i] = make_half2(__float2half(out0), __float2half(out1));

4. Bitonic Sort TopK

// Original: O(K) insertion sort per element
// Improved: O(log²K) parallel bitonic network for merging
bitonic_sort_desc<2*K>(combined_vals, combined_idxs);

Building

Requirements

• CUDA Toolkit 11.0+ (tested with CUDA 12.x)
• RTX 3060 or compatible GPU (SM 86)
• GCC 9+ or compatible compiler

Compile

# Build all
make all

# Build with debug symbols
make DEBUG=1 all

# Build specific target
make benchmark
make test

Run

# Run correctness tests
make run_test

# Run full benchmark suite
make run_benchmark

# Run specific benchmarks
make run_benchmark_softmax
make run_benchmark_topk

# Profile with Nsight Compute
make profile

Benchmark Configurations

Softmax-Only

Batch Size	Vocabulary	Use Case
4000	100-1000	Small models
4000	32000-50000	LLMs (GPT, LLaMA)
10	50000	Online inference

Softmax + TopK (Beam Search)

Batch Size	Vocabulary	K	Use Case
4000	32000	5	Standard beam search
4000	50000	10	Large beam

Expected Results

RTX 3060 Specifications

• SMs: 28
• Memory Bandwidth: 360 GB/s
• Shared Memory: 100 KB/SM
• L2 Cache: 3 MB

Theoretical Analysis

Softmax arithmetic intensity: ~0.625 FLOPS/byte (memory-bound)

Kernel	Memory Accesses/Element	Expected BW Utilization
Safe Softmax	4	~50%
Online Softmax	3	~60%
+ Vectorized	3 (coalesced)	~75%
+ FP16	1.5 (half precision)	~85%

Key Files Explained

common.cuh

• CUDA_CHECK macro for error handling
• GpuTimer class using CUDA events
• softmax_cpu reference implementation
• Data generation and verification utilities

baseline.cuh

• naive_softmax_kernel - Unsafe, can overflow
• safe_softmax_kernel - Standard 3-pass implementation
• online_softmax_kernel - Paper's Algorithm 3
• online_softmax_topk_kernel - Paper's Algorithm 4

warp_primitives.cuh

• warp_reduce_max/sum - Single-warp reductions via shuffle
• OnlineSoftmaxPair - (max, sum) pair for online algorithm
• block_reduce_online_pair - Hierarchical warp+shared reduction

vectorized.cuh

• float4 load/store for memory bandwidth optimization
• process_float4 - Online update for 4 values at once
• Handles non-aligned remainder elements

mixed_precision.cuh

• half2 vectorized loads (2× elements per transaction)
• FP32 accumulation for numerical stability
• BF16 variant for better dynamic range

topk_improved.cuh

• bitonic_sort_desc - Parallel sorting network
• merge_topk_bitonic - Efficient K-way merge
• warp_reduce_topk - Warp-level TopK reduction
• online_softmax_topk_optimized_kernel - All improvements combined

References

1. Milakov, M., & Gimelshein, N. (2018). Online normalizer calculation for softmax. arXiv:1805.02867
2. Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention. NeurIPS.
3. NVIDIA. (2024). CUDA C++ Programming Guide.
4. NVIDIA. (2024). Ampere Tuning Guide.

License

Educational use for GPU computing coursework.