FlashMLA

from flash_mla import get_mla_metadata, flash_mla_with_kvcache

# Generate metadata for efficient scheduling
tile_scheduler_metadata, num_splits = get_mla_metadata(
    cache_seqlens,  # Sequence lengths in the KV cache
    s_q * h_q // h_kv,  # Number of query heads per KV head
    h_kv  # Number of KV heads
)

# Apply FlashMLA in each transformer layer
for i in range(num_layers):
    # Process queries and prepare inputs
    q_i = process_queries(...)
    
    # Call the FlashMLA kernel with KV cache
    o_i, lse_i = flash_mla_with_kvcache(
        q_i,                    # Query tensor
        kvcache_i,              # KV cache tensor
        block_table,            # Paged KV cache block table
        cache_seqlens,          # Sequence lengths in KV cache
        dv,                     # Value dimension
        tile_scheduler_metadata, # Pre-computed metadata
        num_splits,             # Number of splits for processing
        causal=True,            # Use causal attention mask
    )
    
    # Continue with the rest of the transformer layer
    output_i = process_output(o_i, ...)

FlashMLA provides a clean API for integrating efficient multi-layer attention into transformer models. The code above demonstrates how to generate metadata for optimized scheduling and apply the FlashMLA kernel within transformer layers. The pre-computed metadata enables efficient workload distribution across GPU resources.

The metadata generation process is a critical optimization in FlashMLA. It analyzes sequence lengths and head configurations to determine the optimal execution strategy. This pre-computation step enables efficient workload distribution, memory access patterns, and tile configurations that maximize GPU utilization during the attention computation.

The FlashMLA kernel is structured to maximize computational efficiency and memory throughput. It consists of three main components: the Tile Scheduler that distributes work across GPU resources, the Block Table Manager that handles paged KV cache access, and the Memory Manager that optimizes data movement. The computation is divided into three stages: Q-K matrix multiplication, softmax with scaling, and output computation, all optimized for the NVIDIA Hopper architecture.