FlashMLA

To maximize GPU utilization, the new FlashMLA kernel uses a technique called "seesaw" scheduling. This is a variant of ping-pong scheduling designed to overcome register constraints.

The Problem: Register Pressure

In MLA, the output matrix must be stored in registers to satisfy WGMMA instruction requirements. A $64 \times 512$ output matrix occupies 32,768 32-bit registers. Since an SM only has 65,536 32-bit registers, it is impossible to maintain two full output matrices for traditional interleaved ping-pong scheduling.

The Solution: Seesaw Scheduling

The kernel splits the output matrix vertically into two halves: $O_L$ (left) and $O_R$ (right), each $64 \times 256$. It then uses two warpgroups to interleave operations on two KV blocks ($K_0, V_0$ and $K_1, V_1$) such that CUDA Core operations (softmax/scaling) and Tensor Core operations (matmul) overlap.

The mathematical flow for one step is:

1. Warpgroup 0 computes $\vec{p}_0 = \vec{q} K_0^\intercal / qk_scale$.
2. Warpgroup 1 computes $\vec{p}_1 = \vec{q} K_1^\intercal / qk_scale$.
3. Warpgroup 0 updates the running max $m$ and performs softmax on $\vec{p}_0$, then updates the left output $\vec{o}_L$.
4. Warpgroup 1 updates the running max $m$, performs softmax on $\vec{p}_1$, and updates the right output $\vec{o}_R$.
5. The warpgroups then perform cross-updates (e.g., Warpgroup 1 updating $\vec{o}_R$ with components from $\vec{p}_0$) to ensure mathematical equivalence to the online softmax algorithm used in FlashAttention.

This approach allows for overlapping memory access (via TMA) and computation, achieving up to 80% Tensor Core utilization.

To support extended context lengths (up to 128K tokens) while reducing GPU memory pressure, DeepSeek-V3.2 uses a fine-grained FP8 quantization for the KVCache.

Format Details:

• Quantized Part: The first 512 elements of each token's KV Cache are quantized using tile-level quantization (tile size $1 \times 128$). This consists of 512 float8_e4m3 values and 4 float32 scale factors.
• Unquantized Part (RoPE): The remaining 64 elements (the RoPE part) are stored in bfloat16 to prevent precision loss.
• Total Size: Each token's KVCache occupies 656 bytes (512 $\times$ float8_e4m3 + 4 $\times$ float32 + 64 $\times$ bfloat16).

Kernel Execution Flow:

1. Dequantize the 512 float8_e4m3 values into bfloat16.
2. Concatenate them with the 64 bfloat16 RoPE values.
3. Perform MQA (Multi-Query Attention) calculations using bfloat16 matrix multiplication-add (MMA) operations (inputs in bfloat16, outputs in float32).

The FP8 sparse decoding kernel for Hopper GPUs is designed to overcome the 'dequantization-bound' bottleneck where CUDA Core dequantization cannot keep up with Tensor Core MMA operations. It utilizes a technique called Crossover enabled by Distributed Shared Memory (DSM).

The Crossover Mechanism: Since DeepSeek-V3.2 uses Multi-Query Attention (MQA), all 128 query heads attend to the same key heads. The kernel exploits this by launching CTAs (CUDA Thread Blocks) in clusters of size 2. Each CTA in the cluster is responsible for 64 query heads.

Implementation Steps via DSM:

1. Load: Each CTA loads half of the quantized K/V from global memory using a wide 128-bit __ldg load.
2. Dequantize: Each CTA dequantizes its assigned half on the CUDA Cores.
3. Store & Exchange:
   • The CTA stores its dequantized K/V into its own shared memory.
   • Simultaneously, it uses st.async to write that dequantized data into the shared memory of the other CTA in the cluster.
4. Synchronize: The kernel uses a cluster transaction barrier to ensure data exchange is complete.
5. Compute: After synchronization, each CTA has the full set of dequantized K/V values in its local shared memory to perform MMA operations.

This approach effectively halves the dequantization workload per CTA.

MLA decoding is performed in two steps:

1. Call get_mla_metadata once before the decoding loop to generate tile scheduler metadata.
2. Call flash_mla_with_kvcache in each decoding step.

Parameters for get_mla_metadata:

• cache_seqlens: Sequence lengths in the cache.
• s_q: Number of query tokens per sequence (use 1 if speculative decoding is disabled).
• h_kv: Number of KV heads.
• h_q: Number of query heads.
• is_fp8: Boolean indicating if FP8 is used.
• topk: Sparsity parameter.

Parameters for flash_mla_with_kvcache:

• q_i: Query tensor.
• kvcache_i: KV cache tensor.
• block_table: Block table for KV cache.
• cache_seqlens: Sequence lengths.
• dv: Value dimension.
• tile_scheduler_metadata: Metadata from get_mla_metadata.
• num_splits: Number of splits.
• is_causal: Boolean for causal masking.
• is_fp8_kvcache: Boolean for FP8 KV cache usage.
• indices: (Optional) 3D tensor of shape (batch_size, seq_len_q, topk) for sparse attention. If provided, block_table is ignored. Invalid entries should be set to -1.

Return Values: Returns a tuple (out, lse) where out is the attention result and lse is the log-sum-exp value.

from flash_mla import get_mla_metadata, flash_mla_with_kvcache

tile_scheduler_metadata, num_splits = get_mla_metadata(
    cache_seqlens,
    s_q * h_q // h_kv,
    h_kv,
    h_q,
    is_fp8,
    topk,
)

for i in range(num_layers):
    ...
    o_i, lse_i = flash_mla_with_kvcache(
        q_i, kvcache_i, block_table, cache_seqlens, dv,
        tile_scheduler_metadata, num_splits,
        is_causal, is_fp8_kvcache, indices,
    )
    ...