Flash-LLM

Background and Motivation

Flash-LLM addresses the challenge of exploiting unstructured weight sparsity in LLM inference using GPU tensor cores. While SpInfer handles static weight sparsity through bitmap encoding, Flash-LLM tackles a different problem: the fundamental architectural mismatch between unstructured sparsity and tensor cores.

Tensor cores provide 16× higher throughput than SIMT cores on A100 GPUs for FP16 operations, but they cannot skip arbitrary element-level computations. Previous SIMT-based approaches like Sputnik can exploit sparsity through computation skipping but are bottlenecked by SIMT peak performance. The result: SIMT solutions struggle to outperform dense tensor core implementations (cuBLAS) until sparsity exceeds 85-90%.

Flash-LLM's key observation centers on the nature of LLM inference workloads. The MatMuls are "skinny" where the N dimension equals batch size (typically 8-64). This creates a computational intensity problem:

For small N, this severely limits CI regardless of M. Profiling OPT-66B shows tensor cores achieve only 5-10% utilization at batch sizes 8-32 while memory bandwidth is fully saturated. The system is memory-bound, not compute-bound.

With sparse loading at sparsity ratio $\beta$, computational intensity improves to:

At 80% sparsity, you load 5× less weight data while performing the same computation. Since you're memory-bottlenecked, the redundant computation on zeros costs essentially nothing. The tensor cores were idle anyway.

The Core Strategy: Load-as-Sparse, Compute-as-Dense

Flash-LLM reconciles the architectural mismatch through functional separation: SIMT cores handle sparse-to-dense transformation (they're flexible with irregular patterns), while tensor cores handle computation (they require dense inputs but provide massive throughput).

The workflow per thread block:

1. Load sparse weight tile from global memory (compressed, reduced bandwidth)
2. SIMT cores transform sparse to dense in shared memory
3. Load dense activation tile to shared memory
4. Tensor cores compute on dense inputs

This tolerates redundant computation because memory bandwidth is the bottleneck. By reducing global memory traffic through sparse loading, effective tensor core utilization improves from 5-10% to 24-42%.

Tiled-CSL: The Sparse Format

Flash-LLM uses Tiled-CSL, a tile-oriented sparse encoding simpler than SpInfer's hierarchical bitmaps.

Structure:

• TileOffsets array: Starting offset of each tile's data (4 bytes per tile)
• NonZeros array: Each element is 32 bits: [16-bit value | 16-bit location]

The location directly encodes position within the tile, enabling straightforward extraction: extract the location bits, use as shared memory offset, write the value bits.

Storage overhead:

Sparse-to-Dense Transformation Pipeline

The transformation happens in five stages per iteration. Each iteration processes one K_TB-sized chunk of the K dimension.

Stage 1: gmem2reg (Load sparse encoding to registers)

Each thread loads sparse elements from global memory into private registers. Using registers as intermediate buffers provides three benefits:

1. Enables instruction-level parallelism. All LDG instructions execute concurrently rather than serially waiting on each load.
2. Avoids turn-around shared memory accesses (write-then-read same locations).
3. Allows reordering writes to shared memory for bank conflict avoidance.

Stage 2: rst_smem (Reset shared memory)

Zero-initialize the shared memory region for the dense A_Tile using cp.async. The sparse encoding only specifies nonzero locations, so you must first fill with zeros, then overwrite nonzero positions.

Using cp.async allows zero-filling to happen asynchronously while threads continue to other work.

Stage 3: extract (Write registers to shared memory)

Extract sparse elements from registers and write to target locations:

#pragma unroll
for (int i = 0; i < #REG; i++) {
    if (i >= nnz_thread) break;
    // idx extracts location, v extracts value
    A[idx(Reg[i])] = v(Reg[i]);  
}

The #pragma unroll with constant bound #REG is critical. GPU registers aren't addressable by variable indices. Using a compile-time constant allows full loop unrolling, making all register accesses statically determinable.

Stage 4: ld_dense (Load dense activation tile)

Loads dense B_Tile directly from global memory to shared memory using cp.async. This stage runs in parallel with extract since they write to different shared memory regions.

Stage 5: smem2tc (Load to tensor cores and compute)

Uses ldmatrix to cooperatively load 8×8 matrices from shared memory to registers in the layout tensor cores expect, then executes mma.sync instructions for computation. Each mma.sync performs a 16×8×16 matrix multiply-accumulate.

Two-Level Overlapping and Async Operations

Flash-LLM achieves performance through overlapping at two levels:

Inter-iteration overlapping: Double buffering enables preparing iteration N+1 while computing iteration N. While tensor cores compute using buffer A1/B1, CUDA cores load and transform data into buffer A2/B2.

Intra-iteration overlapping: Within a single iteration:

• rst_smem and gmem2reg run in parallel (different execution units)
• extract and ld_dense run in parallel (different shared memory targets)
• Multiple cp.async groups track different operations independently

Synchronization points:

// Barrier 1: Between rst_smem and extract
cp_async_wait<1>();   // Wait for rst_smem to complete
__syncthreads();      // Block-level barrier

// Barrier 2: Between iterations  
cp_async_wait<0>();   // Wait for all cp.async operations
__syncthreads();      // Block-level barrier

The <1> and <0> parameters specify how many cp.async groups to wait for, enabling fine-grained control over which asynchronous operations must complete.

Metadata prefetching:

// Pre-fetch for first iteration
int start_prefetch = offset[1];
int nnz_prefetch = offset[2] - offset[1];

for (int id = 0; id < K_Global/K_TB; id++) {
    int start = start_prefetch;  // Use pre-fetched
    int nnz = nnz_prefetch;
    
    // Pre-fetch for NEXT iteration
    start_prefetch = offset[id + 2];
    nnz_prefetch = offset[id + 3] - offset[id + 2];
    
    CopyGlobal2Reg(A.nz + start, nnz);  // Uses current values
}

By loading the next iteration's tile metadata while processing the current iteration, global memory latency (200-400 cycles) is hidden behind computation.

Bank Conflict Avoidance Through Data Reordering

Shared memory on GPUs is divided into 32 banks that can be accessed simultaneously. When multiple threads in a warp access the same bank, accesses serialize, creating bank conflicts.

The ldmatrix instruction requires specific data layouts for conflict-free loading. An 8×8 matrix maps to banks in a pattern where each thread needs elements from specific banks. Random sparsity means nonzeros could map to arbitrary banks during the extract stage, causing conflicts when multiple threads write to the same bank.

Flash-LLM's solution: ahead-of-time data reordering

Before storing weights in Tiled-CSL format, reorder nonzeros within each tile so elements targeting different banks are grouped for the same warp:

1. Compute target shared memory bank for each nonzero based on its location in the tile
2. Bucket nonzeros by target bank (32 buckets)
3. Form groups of 32 nonzeros by iteratively selecting from the bucket with most remaining elements

This ensures that when a warp executes extract, its 32 writes target different banks, enabling completion in one shared memory wavefront instead of multiple serialized accesses.

The reordering happens offline during weight preprocessing. It changes data placement in global memory but doesn't affect the final shared memory layout or computation results.

Limitations

Flash-LLM is designed for static weight sparsity where patterns are fixed after pruning. It does not support dynamic activation sparsity, which would require runtime pattern detection and adaptive indexing. The sparse-to-dense transformation assumes sparsity patterns are known ahead of time and encoded in Tiled-CSL format during model preparation.