FlashMLA Documentation

Technical Documentation: deepseek-ai/FlashMLA

ℹ️ Provenance: Multi-Tier Ingestion · CodeWiki Reference

Overview (`README.md`) Copy Section

FlashMLA

Breaking change notice (2026.09.30): In the 2026.09.30 release, we removed support for the Hopper architecture and for earlier models (including DeepSeek V3 / V3.2 / V4.0), and we changed the FP8 / FP4 KV cache format. This version is therefore not compatible with previous ones. If you need to run those models or use the old KV cache format, please switch to this commit.

Introduction Copy Chapter

FlashMLA is DeepSeek's library of optimized attention kernels, powering inference of the DeepSeek-V4.1 model on NVIDIA GPUs and Huawei Ascend NPUs. This repo contains token-level sparse attention kernels for prefill and for decoding with an FP8 / FP4 KV cache, a fused kernel that combines Q-norm, Q-RoPE, attention, O-RoPE (conjugate) and the cast to FP8, and dense attention kernels for prefill and backward. The sparse kernels power DeepSeek Sparse Attention (DSA), as introduced in this paper.

FlashMLA 是 DeepSeek 的高性能注意力算子库，支撑着 DeepSeek-V4.1 模型在 NVIDIA GPU 与华为昇腾 NPU 上的推理。该仓库支持 prefill 与使用了 FP8 或 FP4 格式的 KV 缓存的 decoding 计算，以及一个高性能的融合了 Attention 前后的小操作的融合 kernel。此外，该仓库还提供了稠密 Attention 的 prefill 与反向传播算子。

News Copy Chapter

• 2026.09.30 Release of Ascend Attention Kernels: We've released sparse attention prefill and decoding kernels for the Huawei Ascend 950 NPU, which achieve up to 410 TFlops (95% hardware peak) and 360 TFlops (83% of hardware peak) during prefill and decoding, respectively. We've also released a deep-dive technical report on the algorithms and optimization techniques behind these kernels.

  我们发布了华为昇腾 NPU 平台上的稀疏注意力 Prefill 和 Decoding 算子，其在 Prefill 时能够达到 410 TFlops（95% 硬件极限），在 Decoding 时能达到 360 TFlops（83% 硬件极限）。此外，我们还发布了 一篇简短的技术报告，介绍了这个算子背后所使用到的算法与优化技术。

  We also optimize the performance of the fused norm-RoPE-attn-RoPE-cast kernel under decoding settings by around 10% - 15%.

  我们还将 fused norm-RoPE-attn-RoPE-cast kernel 的 decoding 的性能优化了 10% - 15%。

  Note that this release removes support for the Hopper architecture and for earlier models (including DeepSeek V3 / V3.2 / V4.0), and changes the FP8 / FP4 KV cache format, so it is not compatible with previous versions. If you need to run those models or use the old KV cache format, please switch to this commit.

  请注意，这个版本移除了对 NVIDIA Hopper 架构显卡与前代模型（包括 DeepSeek V3 / V3.2 / V4.0）的支持，并且改变了 KV cache 的格式。如果您希望使用 Hopper 显卡、推理前代模型或者使用旧的 KV cache 格式，请切换至 这个 commit。

• 2026.09.10 Release of DeepSeek v4.1's Attention Kernels: We've released attention kernels for DeepSeek-V4.1, including both prefill and decoding (with FP8 or FP4 KV cache). We've also released a fused-norm-rope-attn-rope-cast kernel which fuses Q-norm (only used in V4, not V4.1), Q-RoPE, core attention, O-RoPE (conjugate), and cast-to-fp8, while retaining the same performance.

• 2025.09.29 Release of Sparse Attention Kernels: With the launch of DeepSeek-V3.2, we are releasing the corresponding token-level sparse attention kernels. These kernels power the model's DeepSeek Sparse Attention (DSA) and achieve up to 640 TFlops during prefilling and 410 TFlops during decoding. We also release a deep-dive blog for our new FP8 sparse decoding kernel. Check it out here.

• 2025.08.01 Kernels for MHA on SM100: Thanks to NVIDIA's PR for MHA forward / backward kernels on SM100!

• 2025.04.22 Deep-Dive Blog: We'd love to share the technical details behind the new FlashMLA kernel! Check out our deep-dive write-up here.

• 2025.04.22 Performance Update: We're excited to announce the new release of Flash MLA, which delivers 5% ~ 15% performance improvement for compute-bound workloads, achieving up to 660 TFlops on NVIDIA H800 SXM5 GPUs. The interface of the new version is fully compatible with the old one. Simply upgrade to the new version for an immediate performance boost! 🚀🚀🚀

Performance Copy Chapter

Test & benchmark the fused norm RoPE attn RoPE cast kernel (Sparse): Copy Section

python tests/test-fused-norm-rope-attn-rope-cast.py

TileLang, Tile-Kernels, and DeepGEMM are required to run this test script.

This kernel fuses Q-norm (only used in V4, not V4.1), Q-RoPE, core attention, O-RoPE (conjugate) and the cast to FP8 into a single kernel, which saves the overhead of those small kernels. Although it fuses many small operations, it still achieves the same or even slightly higher TFlops, at the cost of having to permute the Q_b and Wv weights in advance. It achieves up to 1460 TFlops during prefill and 950 TFlops during decoding on B200 with CUDA 13.3.

This kernel currently supports CUDA only; Ascend is not supported.

Test & benchmark MLA prefill (Sparse): Copy Section

python tests/test-sparse-prefill.py

It achieves up to 1350 TFlops on B200 with CUDA 13.3. In practice, we highly recommend using the fused-norm-rope-attn-rope-cast kernel for better performance.

On the Huawei Ascend 950 NPU, it achieves up to 410 TFlops, which is 95% of the theoretical hardware peak.

Test & benchmark MLA decoding (Sparse): Copy Section

python tests/test-sparse-decode.py

It achieves up to 1024 TFlops on B200 with CUDA 13.3. In practice, we highly recommend using the fused-norm-rope-attn-rope-cast kernel for better performance.

On the Huawei Ascend 950 NPU, it achieves up to 360 TFlops, which is 83% of the theoretical hardware peak.

Test & benchmark MHA prefill (Dense): Copy Section

python tests/test_fmha_sm100.py

It achieves up to 1460 TFlops in forward and 1000 TFlops in backward computation on B200, as reported by NVIDIA.

This kernel currently supports CUDA only; Ascend is not supported.

Requirements Copy Chapter

For the NVIDIA platform:

• NVIDIA SM100 / SM103 GPU
• CUDA 13.1 and above (CUDA 13.2 and above is recommended, since it provides the native cvt.rn.bf16x2 conversions that the quantized decoding kernels use; on CUDA 13.1 they fall back to a slower widening through FP32)
• PyTorch 2.10 and above

For the Huawei platform:

• Huawei Ascend 950 NPU
• CANN 9.2.0 and above (provides the bisheng compiler and the Ascend C headers used at build time)
• torch_npu
• PyTorch 2.10 and above

FlashMLA CUDA wheels use PyTorch's Stable C++ ABI and CPython's Limited API (abi3, minimum CPython 3.10). A CUDA wheel built once can therefore be used with PyTorch 2.10+ and CPython 3.10+ without recompilation, subject to the GPU architectures compiled into it. On the Ascend platform, however, the wheel obtains the native device stream through an Ascend-only torch_npu bridge, and must therefore be built for a torch_npu version compatible with the installed PyTorch.

Installation Copy Chapter

git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla
cd flash-mla
git submodule update --init --recursive
pip install -v . --no-build-isolation

--no-build-isolation is required: setup.py imports torch (and torch_npu on the Ascend platform) while building, and this repository does not declare them as PEP 518 build requirements.

The build target platform is detected automatically (/dev/davinci_manager means Ascend, anything else means CUDA) and can be overridden with FLASH_MLA_BUILD_TARGET_PLATFORM=CUDA or FLASH_MLA_BUILD_TARGET_PLATFORM=ASCEND.

On the Ascend platform, ASCEND_HOME_PATH must point to the root of the CANN installation (default: /usr/local/Ascend/ascend-toolkit/latest).

After the CUDA extension is built, a register-spill check runs over the produced .so. The CUTLASS FMHA kernels are exempt, but if any kernel defined by this repository spills registers the build fails with "Register spilling detected. Build failed!"; export FLASH_MLA_SKIP_REG_SPILL_CHECK=1 to skip the check.

Usage Copy Chapter

DeepSeek Sparse Attention (DSA) Decoding Copy Section

To use the DSA decoding kernels, call get_mla_metadata once before the decoding loop to get the tile scheduler metadata. Then, call flash_mla_with_kvcache in each decoding step. For example:

from flash_mla import get_mla_metadata, flash_mla_with_kvcache

tile_scheduler_metadata, num_splits = get_mla_metadata()    # A placeholder only. The actual scheduling metadata is generated on the first call to flash_mla_with_kvcache

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,
        indices=indices,
    )
    ...

Where

• s_q is the number of q tokens per q sequence. If MTP (speculative decoding) is disabled, it should be 1.
• h_kv is the number of key-value heads.
• h_q is the number of query heads.

FP8 / FP4 KV Cache Format

For decoding, this kernel currently supports only the FP8 and FP4 KV cache formats. Unquantized (bfloat16) KV cache format is not supported.

For DeepSeek V4.1 (head_dim = 512), the format is detected from the last dimension of k_cache (i.e. the bytes per token): 528 (V4.1) or 288 (V4.1 fp4). In both formats, each token stores its quantized raw data first, followed immediately by its scales:

• V4.1 (528 bytes per token): the first 512 bytes are float8_e4m3 values covering all 512 dimensions. The 64 RoPE dimensions are quantized as well, so there is no bfloat16 part. The last 16 bytes are float8_e8m0 scales, one for every 32 consecutive float8_e4m3 values.
• V4.1 fp4 (288 bytes per token): the first 256 bytes contain 512 e2m1 values, 2 values per byte, the even-indexed one in the low nibble. The last 32 bytes are float8_e4m3 scales, one for every 16 consecutive e2m1 values. This format is only valid for extra_k_cache when k_cache is in the V4.1 format.

See tests/quant.py for quantization and dequantization details.

indices Tensor

The indices tensor enables token-level sparse attention by instructing the kernel to compute attention only for specified tokens.

• Shape: indices should be a 3D tensor of shape (batch_size, seq_len_q, topk).
• Format: indices_in_kvcache[i][j][k] = (the index of the page block where token t resides) * page_block_size + (the offset of token t within the page block), where t is the k-th token for the j-th query sequence in the i-th batch. Since the index of the page block has already been encoded into indices_in_kvcache, the kernel does not use the block_table parameter; you still have to pass it, since it is a required positional argument of flash_mla_with_kvcache (None is fine).
• Invalid entries: Set invalid indices to -1.

Return Values

The kernel returns (out, lse), where:

• out is the attention result, [batch_size, seq_len_q, h_q, head_dim_v], bfloat16.
• lse is the log-sum-exp value of the attention scores for each query head, [batch_size, h_q, seq_len_q], float32. Its shape is transposed with respect to out.

See tests/test-sparse-decode.py for complete examples.

DeepSeek Sparse Attention (DSA) Prefill Copy Section

For the DSA prefill kernel, call flash_mla_sparse_fwd directly with the following parameters:

• q: Query tensor of shape [s_q, h_q, d_qk]
• kv: Key-Value tensor of shape [s_kv, h_kv, d_qk]
• indices: Indices tensor of shape [s_q, h_kv, topk]
• sm_scale: A scalar value
• d_v: Optional, the dimension of the value vectors. Can only be 512, which is also the default.
• attn_sink: Optional, [h_q], float32. If provided, the output is additionally scaled by exp(lse) / (exp(lse) + exp(attn_sink)). It has no effect on the returned lse and max_logits.
• topk_length: Optional, [s_q], int32. If provided, the i-th query token only attends to the first topk_length[i] indices.

Note on batching: This kernel does not support a batch dimension. For multi-batch inference, reshape the input tensors and adjust the indices parameter to simulate batch processing.

Invalid indices: Set invalid entries in indices to -1 or any number >= s_kv.

Return Values and Equivalent PyTorch Code:
The kernel returns (out, max_logits, lse), where max_logits and lse are in the natural logarithm. This is equivalent to the following PyTorch operations:

Q: [s_q, h_q, d_qk], bfloat16
kv: [s_kv, h_kv, d_qk], bfloat16
indices: [s_q, h_kv, topk], int32

kv = kv.squeeze(1)  # [s_kv, d_qk], h_kv must be 1
indices = indices.squeeze(1)    # [s_q, topk]
invalid = (indices < 0) | (indices >= s_kv)     # [s_q, topk], the kernel ignores these entries
indices = indices.masked_fill(invalid, 0)       # So that the gather below stays in range
focused_kv = kv[indices]    # For the i-th sequence (s_q), the corresponding KV tokens are selected from the KV cache based on indices[i, :]. This operation results in a tensor of shape [s_q, topk, d_qk].

P = (Q @ focused_kv.transpose(-1, -2)) * sm_scale    # [s_q, h_q, topk]
P = P.masked_fill(invalid.unsqueeze(1), float('-inf'))
max_logits = P.max(dim=-1).values   # [s_q, h_q]
lse = torch.logsumexp(P, dim=-1)    # [s_q, h_q]
S = torch.softmax(P, dim=-1)        # [s_q, h_q, topk]
out = S @ focused_kv  # [s_q, h_q, d_qk]

return (out, max_logits, lse)

A query token that has no valid index is an edge case on top of the code above: the kernel returns max_logits = -inf, lse = +inf and an all-zero output for it, whereas the pseudo-code would produce a NaN output.

See tests/test-sparse-prefill.py for a complete example.

Dense MHA Prefill Copy Section

This kernel implements the standard dense Multi-Head Attention (MHA) forward and backward operations. It can be called using:

• flash_attn_varlen_func
• flash_attn_varlen_qkvpacked_func
• flash_attn_varlen_kvpacked_func

The usage is similar to the flash_attn package, with two differences: the two packed variants take an additional required head_dim_qk argument, and all three return (out, lse) instead of only out. See tests/test_fmha_sm100.py for a complete example of flash_attn_varlen_func.

Fused norm + RoPE + attn + RoPE + cast kernel Copy Section

In the DeepSeek-V4.1 release, we also provide a fused kernel that combines Q-norm (only used in V4, not in V4.1), Q-RoPE, core attention, O-RoPE (conjugate) and the cast to FP8 into a single kernel. It removes the extra time spent on these small kernels while keeping the same or even slightly higher TFlops, at the cost of having to permute the Q_b and Wv weights in advance.

In DeepSeek-V4.1 attention, Q ([hidden_size]) is first projected to [q_lora_rank] (the Q_a projection) and then to [num_attention_heads, head_dim] (the Q_b projection). After core attention, the output ([num_attention_heads, head_dim]) is reshaped to [o_groups, num_attention_heads // o_groups * head_dim], and each of its rows is projected to [o_lora_rank] (the Wv projection), giving an [o_groups, o_lora_rank] matrix. That matrix is reshaped to [o_groups * o_lora_rank] and finally projected to [hidden_size] (the Wo projection). This kernel requires the Q_b and Wv weights to be permuted.

To permute the Q_b weight:

import torch
import tile_kernels
from flash_mla import fused_norm_rope_attn_rope_cast

h_q, d_q = 64, 512          # Q heads and Q head dimension
q_lora_rank = 1536
scale_gran = 128

# q_b_proj: [h_q * d_q, q_lora_rank], bfloat16
q_b_proj = torch.randn((h_q * d_q, q_lora_rank), dtype=torch.bfloat16, device='cuda')

# Quantize the weight to FP8 with per-token scale factors, in DeepGEMM's layout
q_b_proj_fp8, q_b_sf = tile_kernels.quant.per_token_cast(
    q_b_proj, 'e4m3', scale_gran,
    use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True,
)

# Permute the weight and its scale factors into the layout required by the fused kernel
q_b_proj_fp8, q_b_sf = fused_norm_rope_attn_rope_cast.permute_q_b_proj(
    (q_b_proj_fp8, q_b_sf), h_q, d_q,
)

To permute the Wv weight:

import deep_gemm
import torch
import tile_kernels
from flash_mla import fused_norm_rope_attn_rope_cast

n_wv_group, wv_group_size, d_o = 8, 8, 512      # n_wv_group * wv_group_size == h_q
wv_proj_out_dim = 512                           # o_lora_rank
scale_gran = 32

# wv_proj: [n_wv_group * wv_proj_out_dim, wv_group_size * d_o], bfloat16
# (it is viewed as [n_wv_group, wv_proj_out_dim, wv_group_size * d_o] further below)
wv_proj = torch.randn((n_wv_group * wv_proj_out_dim, wv_group_size * d_o),
                      dtype=torch.bfloat16, device='cuda')

# Quantize the weight to FP8, and put its scale factors into the layout that DeepGEMM's einsum expects
wv_proj_fp8, wv_sf = tile_kernels.quant.per_token_cast(
    wv_proj, 'e4m3', scale_gran,
    use_tma_aligned_col_major_sf=False, round_sf=True, use_packed_ue8m0=False,
)
wv_sf = deep_gemm.transform_sf_into_required_layout(
    wv_sf.view(n_wv_group, wv_proj_out_dim, wv_group_size * d_o // scale_gran),
    wv_proj_out_dim, wv_group_size * d_o,
    num_groups=n_wv_group, recipe=(1, 1, scale_gran), is_sfa=False,
)
wv_proj_fp8 = wv_proj_fp8.view(n_wv_group, wv_proj_out_dim, wv_group_size * d_o)

# Permute the weight and its scale factors into the layout required by the fused kernel
wv_proj_fp8, wv_sf = fused_norm_rope_attn_rope_cast.permute_wv_proj(
    (wv_proj_fp8, wv_sf), wv_group_size, d_o,
)

And finally, to use the fused kernel:

/* Detailed source-code truncated for AI context efficiency. */

You may refer to the fused kernel's test script (tests/test-fused-norm-rope-attn-rope-cast.py) for a complete example.

Acknowledgement Copy Chapter

FlashMLA is inspired by FlashAttention 2&3 and cutlass projects.

Community Support Copy Chapter

MetaX Copy Section

For MetaX GPUs, visit the official website: MetaX.

The corresponding FlashMLA version can be found at: MetaX-MACA/FlashMLA

Moore Threads Copy Section

For the Moore Threads GPU, visit the official website: Moore Threads.

The corresponding FlashMLA version is available on GitHub: MooreThreads/MT-flashMLA.

Hygon DCU Copy Section

For the Hygon DCU, visit the official website: Hygon Developer.

The corresponding FlashMLA version is available here: OpenDAS/MLAttention.

Intellifusion Copy Section

For the Intellifusion NNP, visit the official website: Intellifusion.

The corresponding FlashMLA version is available on Gitee: Intellifusion/tyllm.

Iluvatar Corex Copy Section

For Iluvatar Corex GPUs, visit the official website: Iluvatar Corex.

The corresponding FlashMLA version is available on GitHub: Deep-Spark/FlashMLA

AMD Instinct Copy Section

For AMD Instinct GPUs, visit the official website: AMD Instinct.

The corresponding FlashMLA version can be found at: AITER/MLA

Citation Copy Chapter

@misc{flashmla2025,
      title={FlashMLA: Efficient Multi-head Latent Attention Kernels},
      author={Jiashi Li and Shengyu Liu and Yuanhang Sun},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}},
}

Overview (`csrc/cuda_kernels/sm100/decode/sparse/head128/README.md`) Copy Section

Located under csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations since it shares the same .cu file with sm100/prefill/sparse/fwd_for_small_topk/head128

20250422 New Kernel Deep Dive (`docs/20250422-new-kernel-deep-dive.md`) Copy Section

A Deep-Dive Into the New Flash MLA Kernel

In the previous version of the Flash MLA kernel, we have achieved impressive performance: 3000 GB/s in memory-intensive settings and 580 TFlops in compute-bound settings. Now, we're pushing these numbers even further, reaching up to 660 TFlops.

In this blog, we present a deep dive into the new kernel, explaining the optimizations and techniques behind this performance boost. We'll first explain why the MLA kernel is compute-bound despite being a decoding-stage attention kernel, then discuss our high-level kernel schedule design, and finally cover the technical details of the new kernel.

A Theoretical Analysis of the MLA Algorithm Copy Chapter

GPU kernels can be classified as either compute-bound (limited by floating-point operations per second, FLOPs) or memory-bound (limited by memory bandwidth). To identify the kernel's bottleneck, we calculate the ratio of FLOPs to memory bandwidth (FLOPs/byte) and compare it with the GPU's capacity.

Assume the number of q heads is $h_q$, the number of q tokens per request is $s_q$ (should be 1 if MTP / speculative decoding is disabled), the number of kv tokens per request is $s_k\ (s_k \gg h_q s_q)$, and the head dimensions of K and V are $d_k$ and $d_v$ respectively. The number of FLOPs is roughly $2 (h_q s_q \cdot d_k \cdot s_k + h_q s_q \cdot s_k \cdot d_v) = 2 h_q s_q s_k (d_k+d_v)$, and the memory access volume (in bytes) is $\mathop{\text{sizeof}}(\text{bfloat16}) \times (h_q s_q d_k + s_k d_k + h_q s_q d_v) \approx 2s_k d_k$. Thus, the compute-memory ratio is $h_q s_q \cdot \frac{d_k+d_v}{d_k} \approx 2 h_q s_q$.

An NVIDIA H800 SXM5 GPU has a peak memory bandwidth of 3.35 TB/s and peak FLOPs of 990 TFlops. However, due to throttling (reducing to ~1600 MHz in our case), the practical peak FLOPs drops to ~865 TFlops. Therefore, when $h_qs_q \ge \frac{1}{2} \cdot \frac{865}{3.35} = 128$, the kernel is compute-bound; otherwise, it's memory-bound.

According to the overview of DeepSeek's Online Inference System, we don't use Tensor Parallel for decoding instances, meaning $h_q$ is 128 and the kernel is compute-bound. Thus, we need to optimize the kernel for compute-bound settings.

High-Level Design of the New Kernel Copy Chapter

To fully utilize GPU compute resources, we need to overlap CUDA Core operations with Tensor Core operations and memory access with computation, keeping the Tensor Core constantly busy. This requires redesigning the kernel's "schedule."

FlashAttention-3's paper introduces ping-pong scheduling and intra-warpgroup GEMM-softmax pipelining to overlap block-wise matmul and CUDA Core operations. However, these techniques can't be directly applied here due to resource constraints. The output matrix (scaled and accumulated during each mainloop round, similar to FlashAttention's algorithm) must be stored in registers due to WGMMA instruction requirements. Each $64 \times 512$ output matrix occupies 32,768 32-bit registers. With only 65,536 32-bit registers per SM, we can store only one output matrix per SM. This eliminates the possibility of having two output matrices and letting them use CUDA Core and Tensor Core in a interleaved manner. We need to find another clever way to overlap CUDA Core and Tensor Core computation.

(You might pause here to ponder - perhaps you can find a better solution than ours!)

Our solution involves an additional mathematical transformation beyond FlashAttention's online softmax and accumulation approach. In each step, we take two KV blocks (called $K_0$, $K_1$, $V_0$, and $V_1$). Since the output matrix occupies 32,768 registers (too many for one warpgroup), we split it vertically into $O_L$ and $O_R$ (each $64 \times 256$). We similarly split $V_0$ and $V_1$ into $V_{0L}$, $V_{0R}$, $V_{1L}$, and $V_{1R}$ (each $64 \times 256$). The output matrix is then computed as follows:

0. Maintain a running max $m$ (initialized to $-\infty$, shared between the two warpgroups) and output matrices $\vec o_L, \vec o_R$ (initialized to 0).
1. [0] Compute $\vec p_0 = \vec q K_0^\intercal / qk\_scale$.
2. [1] Compute $\vec p_1 = \vec q K_1^\intercal / qk\_scale$.
3. [0] Compute $mp_0 = \max(\vec p_0)$, $m\_new_0 = \max(m, mp_0)$, and $scale_0 = \exp(m\_new_0 - m)$. Update $m \gets m\_new_0$.
4. [0] Perform softmax on $\vec p_0$: $\vec p_0 \gets \exp(\vec p_0 - m\_new_0)$.
5. [0] Update $\vec o_L \gets \vec o_L \cdot scale_0 + \vec p_0 V_{0L}$.
6. [1] Compute $mp_1 = \max(\vec p_1)$, $m\_new_1 = \max(m, mp_1)$, and $scale_1 = \exp(m\_new_1 - m)$. Update $m \gets m\_new_1$.
7. [1] Perform softmax on $\vec p_1$: $\vec p_1 \gets \exp(\vec p_1 - m\_new_1)$.
8. [1] Update $\vec o_R \gets \vec o_R \cdot (scale_0 \cdot scale_1) + \vec p_1 V_{1R}$.
9. [0] Update $\vec p_0 \gets \vec p_0 \cdot scale_1$.
10. [1] Update $\vec o_R \gets \vec o_R + \vec p_0 V_{0R}$.
11. [0] Update $\vec o_L \gets \vec o_L \cdot scale_1 + \vec p_1 V_{1L}$.

Note: We assume one q head for simplicity, so $\vec q$ and $\vec o$ are vectors. Bracketed numbers indicate the warpgroup performing the operation. Assume $\vec o_L$ resides in warpgroup 0's register and $\vec o_R$ resides in warpgroup 1's register.

This schedule can be viewed as a "ping-pong" variant using one output matrix—we call it "seesaw" scheduling. It's mathematically equivalent to FlashAttention's online softmax algorithm. This schedule allows us to overlap CUDA Core and Tensor Core operations by interleaving the two warpgroups, and also allows us to overlap memory access with computation since we can launch the corresponding Tensor Memory Accelerator (TMA) instructions right after data is no longer needed.

The complete schedule is shown below (remember that in MLA, $K$ and $V$ are the same with different names):

Discussion of Technical Details Copy Chapter

This section covers technical details of the new kernel.

First, although the kernel targets compute-bound scenarios (where memory bandwidth isn't the bottleneck), we can't ignore memory latency. If the data is not ready when we want to use it, we have to wait. To solve this problem, we employ the following techniques:

• Fine-grained TMA copy - GEMM pipelining: For a $64 \times 576$ K block, we launch 9 TMA copies (each moving a $64 \times 64$ block). GEMM operations begin as soon as each TMA copy completes (When the first TMA copy is done, we can start the first GEMM operation, and so on), improving memory latency tolerance.
• Cache hints: Using cute::TMA::CacheHintSm90::EVICT_FIRST for TMA copies improves L2 cache hit rates, as shown by experiments.

These optimizations achieve up to 80% Tensor Core utilization (of the throttled theoretical peak) and 3 TB/s memory bandwidth on an H800 SXM5 GPU. While slightly slower (~2%) than the old ping-pong buffer version in memory-bound settings, this is acceptable.

Other performance improvements include:

• Programmatic Dependent Launch. We use programmatic dependent launch to overlap splitkv_mla and combine kernels.
• Tile Scheduler. We implement a tile scheduler to allocate jobs (requests and blocks) to SMs. This ensures a balanced load across SMs.

Acknowledgements Copy Chapter

FlashMLA's algorithm and scheduling are inspired by FlashAttention, Flash-Decoding, and CUTLASS, as well as many projects behind them. We thank the authors for their great work.

Citation Copy Chapter

@misc{flashmla2025,
      title={FlashMLA: Efficient MLA decoding kernels},
      author={Jiashi Li and Shengyu Liu and Yuanhang Sun},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}},
}

20250929 Hopper Fp8 Sparse Deep Dive (`docs/20250929-hopper-fp8-sparse-deep-dive.md`) Copy Section

A Deep Dive Into The Flash MLA FP8 Decoding Kernel on Hopper

With the release of DeepSeek-V3.2, we have doubled the context length of our models from 64K tokens to 128K tokens. This puts significant pressure on GPU memory (a single request with 128K tokens requires a KVCache of size $576 \times 2 \times 62 \times 128 \times 1024 = 8.72\ \mathrm{GiB}$), which can lead to out-of-memory (OOM) errors or under-utilized GPUs due to small batch sizes. To address this, we introduced FP8 KVCache for DeepSeek-V3.2.

However, writing a high-performance decoding kernel is challenging due to the need for dequantization and its sparse memory access patterns. In this blog, we share the story behind our new FP8 sparse decoding kernel for Hopper GPUs. We will first explain our FP8 KVCache format, then provide a theoretical analysis of clock cycles, and finally detail the techniques used in our new kernel.

The FP8 KVCache Format Copy Chapter

Recall that the decoding phase of the Multi-head Latent Attention (MLA) algorithm operates similarly to Multi-Query Attention (MQA), with 128 query heads and 1 key head, where head_dim_k = 576 and head_dim_v = 512 respectively. To reduce the size of the KVCache while maintaining accuracy, we use a fine-grained quantization method. Specifically, we apply tile-level quantization (with a tile size of $1 \times 128$) to the first 512 elements in each token's KV Cache. This results in 512 float8_e4m3 values and 4 float32 scale factors. For the remaining 64 elements (the RoPE part), we do not apply quantization as they are sensitive to precision loss. Therefore, in GPU memory, each token's KVCache occupies 656 bytes, consisting of 512 float8_e4m3s, 4 float32s, and 64 bfloat16s.

Inside the kernel, we first dequantize the 512 float8_e4m3 values into 512 bfloat16s. We then concatenate them with the 64 original bfloat16 values from the RoPE part. Finally, we perform the MQA calculation using matrix multiplication-add (MMA) operations in bfloat16 precision (i.e., the inputs to the MMAs are in bfloat16 and the outputs are in float32. This applies to both the QK gemm and the attention-score-V gemm).

Theoretical Analysis of Clock Cycles Copy Chapter

The main challenge is that Tensor Cores (which handle MMA calculations) are extremely fast, while the dequantization process, performed on CUDA Cores, struggles to keep up.

The basic unit on an NVIDIA GPU is the Stream Multiprocessor (SM). You can think of each SM as an independent core on the GPU. For simplicity, let's focus on a single SM. Each SM can process 4096 MMA Flops per clock cycle (calculated as 989 TFlops / 1830 MHz / 132 SMs on H800). In our kernel, each CTA runs on one SM, and each SM is only mapped to one CTA. If we assign each CTA (CUDA Thread Block) to process 64 query heads, it only requires $64 \times (576+512) \times 2 / 4096 \approx 34$ cycles for MMA operations per K/V token.

However, because the H800 cannot directly cast float8_e4m3 to bfloat16, dequantizing the KVCache for one token requires the following steps:

1. Convert float8_e4m3 to half
2. Convert half to float32
3. Convert float32 to bfloat16
4. Multiply the converted bfloat16 value by the float32 scale factor

According to NVIDIA's documentation, we need at least $(\frac{1}{64} + \frac{1}{64} + \frac{1}{16} + \frac{1}{256}) \times 512 \approx 50$ cycles for dequantizing each token! This is significantly more than the 34 cycles required for the MMA operations, meaning the kernel is dequantization-bound. If left unaddressed, dequantization would become the performance bottleneck, leaving the powerful Tensor Cores underutilized.

Crossover Copy Chapter

Before we continue, it's important to note a key fact: every query head within the same query token attends to the same key heads, because this is Multi-Query Attention (MQA).

Recall that each CTA processes 64 query heads, while DeepSeek-V3.2 has a total of 128 query heads. If we can find a way to "share" the dequantized K/V values between two CTAs that are processing different sets of query heads, then each CTA would only need to dequantize half of the KV cache – which is fantastic! We call this method "crossover", since the idea was actually inspired by Chromosomal crossover during Meiosis.

The next question is, how do we implement this in CUDA? Before NVIDIA's Hopper architecture, the only options for data exchange between CTAs were global memory or the L2 cache, which are slow. However, the powerful Distributed Shared Memory gave us a new solution.

Distributed Shared Memory to the Rescue Copy Chapter

Distributed Shared Memory (DSM) is a new feature introduced with the Hopper architecture, alongside the CTA Cluster (thread block cluster). CTAs within the same cluster can directly access each other's shared memory. For more details, you can refer to NVIDIA Hopper Architecture In-Depth.

Here is how we use it: We launch CTAs in clusters of size 2. Each CTA within a cluster is responsible for 64 query heads from the same query token. Each CTA performs the following steps:

1. Loads half of the quantized K/V from global memory. We use a wide __ldg load with a width of 128 bits to improve performance.
2. Dequantizes its assigned half on the CUDA Cores.
3. Stores the dequantized K/V into its own shared memory.
4. Simultaneously uses st.async to write the dequantized K/V into the shared memory of the other CTA in the cluster.

For synchronization between these operations, we rely on the cluster transaction barrier, another powerful programming primitive available in CTA Clusters. After the data exchange is complete, each CTA has the full set of dequantized K and V values available in its own shared memory, which it can then use to perform the MMA operations.

Performance Copy Chapter

Using these techniques, we achieved 410 TFLOPS in a compute-bound configuration (batch_size=128, num_heads=128, s_q=2, topk=2048) on H800 SXM5 GPUs. This is a significant improvement over the 250 TFLOPS achieved by our previous FP8 sparse decoding kernel without the crossover technique.

Although this number is still below the 640 TFLOPS peak of our previous bfloat16 dense decoding kernel, one reason is that it's a sparse kernel, and its topk is only 2048. With a smaller topk, the relative overhead of the kernel's prologue and epilogue becomes larger compared with dense decoding with long context length. If we set topk to a larger value, such as 32768, this kernel can achieve up to 460 TFLOPS.

From another perspective, the execution time of this kernel in the configuration mentioned above is comparable to that of the dense decoding kernel when the sequence length is around 3000. When the sequence length exceeds 3000, the performance advantage of our new kernel becomes even more significant. This also highlights the effectiveness of our DeepSeek Sparse Attention algorithm.

20260930 Ascend Prefill Deep Dive Zh (`docs/20260930-ascend-prefill-deep-dive-zh.md`) Copy Section

Ascend Sparse Attention Forward 算法与优化技术简析

前言 Copy Chapter

DeepSeek 于 2026 年 9 月 30 日开源了华为昇腾平台的推理基础设施与核心组件，其中就包括本仓库（deepseek-ai/FlashMLA）中的昇腾平台的稀疏 Attention 的 prefill 与 decoding 算子。在 DeepSeek v4.1 模型的典型工况下，该算子可以在 prefill 时达到 410 TFlops，decoding 时达到 360 TFlops，分别达到了硬件理论性能上限的 95% 与 83%。

在本文中，我将系统地介绍该算子的设计思路、流水排布与优化技巧。

算法回顾 Copy Chapter

本算子用于计算 DeepSeek V4.1 模型中使用的 DeepSeek Sparse Attention (DSA) 架构中的 Attention 的前向传播，包括 Prefill 和 Decoding。其接受若干 q token 与 kv token，按照 indices 表的指示，让每个 q token 只关注部分“最重要的” kv token，从而在保证模型效果的同时大幅节约算力。

数学建模 Copy Section

在本算子中，不同的 q token 总是分开处理的。为了简便，我们后面只考虑单个 Q token。

给定：

• Q: [h_q, d_qk]

• KV: [s_kv, d_qk]

• indices: [topk]

• sm_scale: 一个标量

前向传播需要计算的是：

• gathered_kv = KV.index_select(indices, dim=0), [topk, d_qk]

• P = Q @ gathered_kv^T, [h_q, topk]

• S = softmax(P*sm_scale, dim=-1), [h_q, topk]

• out = S @ gathered_kv[:, :d_v], [h_q, d_v]

计算流程 Copy Section

由于 softmax 具有全局依赖，我们使用 Flash Attention 同款的 online softmax 算法进行计算，动态更新当前的每行的 P 的最大值、每行的 lse (log sum exp)、以及迄今为止的输出累加器。具体步骤如下：

1. 将 topk 个 indices 切成若干块，每块大小为 B_TOPK（B_TOPK 一般取 64 或 96）

2. 维护三个东西：running_max, running_max_for_softmax 和 running_sumexp，shape 均为 [h_q]，初始值为 $-\infty$, $-\infty$ 和 0。另维护一个 output accumulator out_accum, shape 为 [h_q, d_v]。这几个东西都是 float32 的。

3. 对于每个块：

   1. 根据 indices[B_TOPK * k, B_TOPK * k + B_TOPK] 稀疏地从 global memory 中读取 KV，得到这个块的 gathered_kv, [B_TOPK, d_qk]

   2. 计算 P = Q @ gathered_kv^T, [h_q, B_TOPK]

   3. 计算 P 的每行的最大值 cur_max, [h_q]

   4. 更新 running_max：running_max = max(running_max, cur_max)（也就是说 running_max 总是反映当前每行的最大值）

   5. 检查是否有某一行 i 满足 running_max[i] - running_max_for_softmax[i] > RESCALE_THRES，其中 RESCALE_THRES 为一超参数，在我们的 kernel 中一般取 $6$

   6. 如果有某一行 i 触发了这个条件，进入 online softmax 更新流程

      1. out_accum <- out_accum * exp(running_max_for_softmax - running_max).unsqueeze(-1)

      2. running_sumexp <- running_sumexp * exp(running_max_for_softmax - running_max)

      3. running_max_for_softmax <- running_max

      4. 这相当于更新了 online softmax 所使用的 max，因此 exp(running_max - running_max_for_softmax) 总是 $\le \exp(6)$ 的

      5. 如果没有任何一行 i 触发了这个条件，就跳过这一步。该优化技术被称为“skip scale”

   7. 计算 S = exp(P - running_max_for_softmax)，并将 S cast 成 bfloat16，然后与 gathered_kv 相乘，加到 out_accum 上（out_accum += S.to(torch.bfloat16) @ gathered_kv）

   8. 更新 running_sumexp += S.sum(dim=-1)

讨论 Copy Chapter

在正式开始介绍算子的流水设计、优化技巧之前，我们先进行几个简单的讨论，确定算子的总体设计方向：

每个 AI Core 负责多少个 q head Copy Section

第一个问题是：每个 AI Core 负责多少个 q head 合适呢？

先来算一下算存比：按照 CUBE Core 打满（4096 FMA per cycle per AICore）、L2 总带宽 5 TB/s 计算：

d = 512
l2_total_bw = 5e12
num_cores = 32
freq = 1650e6
l2_bw = l2_total_bw / num_cores / freq  # Per AI Core per cycle

for h_q in [16, 32, 64, 128]:
    for topk in [512, 1024]:
        cube_cycles = h_q * topk * (d+d) // 4096
        qo_mem_vol = h_q * d * 2 * 2
        kv_mem_vol = topk * d * 2
        mem_cycles = (qo_mem_vol + kv_mem_vol) / l2_bw
        print(f"{h_q:3d}, {topk=:4d}, {cube_cycles=:6d}, {mem_cycles=:6.0f}")

16, topk= 512, cube_cycles=  2048, mem_cycles=  5883
 16, topk=1024, cube_cycles=  4096, mem_cycles= 11419
 32, topk= 512, cube_cycles=  4096, mem_cycles=  6229
 32, topk=1024, cube_cycles=  8192, mem_cycles= 11765
 64, topk= 512, cube_cycles=  8192, mem_cycles=  6921
 64, topk=1024, cube_cycles= 16384, mem_cycles= 12457
128, topk= 512, cube_cycles= 16384, mem_cycles=  8305
128, topk=1024, cube_cycles= 32768, mem_cycles= 13841

可以看出 head >= 64 才是 Compute bound。对于 head=64 的情况，compute 和 memory 大致均衡。

同时，如果 per AI Core 的 head 取 128 的话，SRAM 容量方面会有诸多问题。具体来说，L0C 的容量会被 O（128 x 512, float32, 256KB）完全占满，因此 accumulator 必须放在 VECTOR Core 上面，因此必须用两个 AIV，以及必须要做细粒度 Pipeline（Pipeline 拷贝 - 累加），很难写。以及，根据 micro-benchmark 的结果，MMAD 的 M 取 64 就足以打满 CUBE Core 的吞吐，因此 head 取 128 是不必要的。

FA 的 output accumulator 在哪里做 Copy Section

Flash Attention 需要在 running max 和当前缩放使用的 scaling max 差异过大的时候，需要把 output 乘上 $\exp(\text{new_max} - \text{old_scaling_max})$，称为“rescale”。但是，L0C 上的数据只能被 FixPipe 搬出或被 CUBE Core 用作累加器，无法原地乘以标量。可行的做法有两种：

1. 在 VECTOR Core 上维护一个 accumulator，每次需要 rescale-O 的时候，把 CUBE 上的 accumulation 结果拿过来，并加在当前的 accumulator 上面

2. accumulator 始终在 CUBE Core 上面，当需要 rescale 的时候，把 accumulator 拉到 vector 上面，rescale，并放回到 CUBE 上面（通过乘一个单位矩阵）

   1. 这种做法需要消耗额外的 CUBE 算力，有损性能。把一个 64x512 的矩阵从 UB 搬到 L0C 上，至少需要做一个 64x64x512 的 float 矩阵乘法

   2. 并且，如果希望不损失精度，需要关闭 HF32 mode，此时 CUBE 算力只有 256 FMA / cycle / AI Core，64x64x512 的矩阵乘法需要 8192 周期，根本无法接受

因此，最终选择了方案一。

能不能用 1C1V Copy Section

下一个问题是，每个 CUBE Core 在计算时配合上多少个 VECTOR Core 比较合理呢？华为昇腾 950 的每个 AI Core 具有一个 CUBE Core 和两个 VECTOR Core，倘若能把每个 CUBE Core 所需的 VECTOR Core 数量控制在一个（1C1V），那么另一个 VECTOR Core 便可以跑一些其他的轻量操作（比如模型中的其它模块或跨卡通信），提高整体的 MFU。

经过一些分析，答案是不能。原因如下：

1. MTE2 Issue Queue Entry bound。在华为昇腾 950 NPU 上，单个 VECTOR Core 在同一时刻最多只能有 16 个“已发出但还未完成”的 MTE2 拷贝请求。哪怕用上了 gather2 技巧（见下文），一次也只能拷贝 32 个 token，太少了。

2. Dequantize bound。在 decoding 的时候，算子需要读取 FP8 或 FP4 格式的 KV Cache，将其反量化为 bfloat16 格式，再送入 CUBE Core 计算。根据微基准测试（micro-benchmark），使用 1C1V 时，反量化的速度赶不上 CUBE Core 计算的速度，导致反量化成为瓶颈。

3. FixPipe bound。当只使用一个 CUBE Core 和一个 VECTOR Core 的时候，FixPipe 向外搬运的速度只有 128 Byte/cycle/AI Core。此时每个块的 MMA 计算周期为 MMA_M * B_TOPK * (512+512) / 4096 = B_TOPK * 16，FixPipe 搬运所需的周期为 (MMA_M * 512 + MMA_M * B_TOPK) * sizeof(float) / 128 = 1024 + B_TOPK * 2。其中 MMA_M 代表每个 AI Core 处理的 q head 数量，根据上文的讨论，应该是 64；B_TOPK 为 Flash Attention 中，沿着 KV 方向循环时的 Tile size。如果希望 FixPipe 不要成为带宽瓶颈的话，需要 B_TOPK * 16 > 1024 + B_TOPK * 2，得到 B_TOPK > 73.14。为了留出足够的余量，B_TOPK 需要选取 96 或 128，这一方面使得 UB 空间更加紧张，另一方面也加重了上面提到的 MTE2 Issue Queue Entry Bound 的问题。

因此，最终我们还是决定为每个 CUBE Core 配备两个 VECTOR Core，以分摊 MTE2 发射、反量化、FixPipe 搬运、softmax 等操作的开销。

KV 的 ND2NZ Layout 转换何时进行 Copy Section

下一个问题：KV 的 ND2NZ Layout 转换何时进行呢？

KV 在显存中是以 ND 格式存储的，shape 为 [s_kv, d_qk]，stride 为 [d_qk, 1]。但是，CUBE Core 要求输入的矩阵使用 NZ layout。正常来说，我们会在把数据从显存搬运到 L1 buffer 中时，随路完成 ND2NZ 转换，但我们的 KV 是使用 MTE2 一个一个 token 地“稀疏”的从显存上搬到片上存储上的，无法启用随路的 ND2NZ 转换。

有一个思路是，在发起 GM -> L1 拷贝时，将一个 token（512 个 bfloat16）视为 32 个 32B 的数据块，这 32 个数据块在显存上是连续的，在 L1 buffer 中则以 B_TOPK * 32B 为 stride 排布。但后续实验显示，AI Core 内部的 BIU（负责发送访存指令的单元）只会在源和目标都连续的时候才会触发访存合并，而像这里这种“只有源连续，而目的不连续”的情况是无法合并访存的，导致 L2 带宽利用率很低，整体瓶颈位于 L2 带宽。

所以，唯一的做法就是将 KV 先从显存拷贝到 UB，使用 SIMD VF 将其从 ND Layout 转为 NZ Layout，然后再拷贝到 L1。

如何进行 skip-scale 的同步 Copy Section

上文“算法回顾”一章中提到过“skip-scale”这一技术，其思想为，只要现在用来做 online softmax 的 max 没有比真正的 max 小太多，就可以跳过本轮的 output 的更新与累加。

这个操作在华为 NPU 上的意义更大 —— 因为 NPU 上 scale-O 很麻烦，需要使用 FixPipe 将 O 从 L0C 拷贝到 UB，然后通过 VF 累加到 UB 的 out_accum 上，浪费 FixPipe、UB 与 SIMD VF 的带宽。

但是，想在 NPU 上实现这个操作有一个难点：“这轮要不要 skip-scale”的信号是 VECTOR Core 发起的，消费者却是 CUBE Core 与 VECTOR Core 二者（因为 CUBE Core 需要决定是否将 L0C 上的 O 清空）。最终，我们决定使用 SS buffer 与 cross core flag 完成这里的 skip-scale 信号传递。

流水线 Copy Chapter

至此，Kernel 的总体设计方向已经定稿：64 个 head / AI Core，1C2V，output accumulator 放在 UB 上，KV 经由 UB 和 SIMD VF 完成 ND2NZ 的中转。

下一步便是流水线的设计了。鉴于这个 kernel 应该是 CUBE Core 的算力 bound，我们的目标就是：让 CUBE Core 上的矩阵乘法计算单元时时刻刻都能打满，并且不要让其他操作（比如 softmax）成为瓶颈。

最基本的流水思路大致是这样（下面使用中括号中的数字来指代这是第几个 KV block 对应的操作）：

1. 计算 P[i] = Q @ K[i]^T（@ 代表矩阵乘法）

2. 计算 O += S[i-1] @ V[i-1]

3. 计算 P[i+1] = Q @ K[i+1]^T

4. 计算 O += S[i] @ V[i]

5. ...

也即，相当于将前一个块的 O += S @ V 延后计算，以给 softmax 操作留出时间：P[i] 在 (1) 结束后就已就绪，直到 (4) 开始时才会被用到，期间有充足的时间（大约 1024 周期）供 FixPipe 将 P[i] 从 L0C 拷贝至 UB、softmax、以及将 S[i] 从 UB 拷贝至 L1 使用。

index 的读取（从显存拷贝到 UB）、index 的预处理（预处理下文所说的 gather2 的 MTE 拷贝参数）、MTE2 拷贝的发起、以及 skip-scale 技术中的 SS buffer 信号同步则会让这个流水线变得复杂一些。CUBE Core 一侧的流水如上文所示，VECTOR Core 一侧的流水如下排布：

1. 发起第 i+5 个 KV 块的 index 读取

2. 发起第 i+3 个 KV 块的 ND2NZ 转换与 UB -> L1 的 MTE3 拷贝

3. 等待 P[i+1] = Q @ k[i+1]^T 完成，发起第 i+1 个 KV 块的 softmax 计算，并将结果

4. 发起第 i+5 个 KV 块的 index 预处理

5. 发起第 i 个 KV 块的 O scale，如果这一块没有 skip-scale O 的话

6. 和 SIMD VF 同步，获取第 i+1 个块是否需要 skip-scale O，将其写入 SS buffer

7. 发起第 i+5 个 KV 块的 MTE2 拷贝，将 KV 从显存拷贝至 UB

为何这个流水如此复杂？

• 一是因为，算子依靠 SS buffer 传递 skip-scale 的信号，但能读写 SS buffer 的只有 scalar core，而 skip-scale 的信号由 SIMD VF 生成，因此这会导致 SIMD VF -> Scalar Core 的阻塞式同步。倘若此处没有处理好，便会导致 scalar core 等待 SIMD VF 完成、其他单元出现空闲的情况。

• 二是因为，华为的各个硬件管线（PIPE）都是顺序发射、顺序执行的，比如对于 PIPE_V 来说，发射 SIMD VF 的顺序与执行的顺序必须完全一致，如果前面的 SIMD VF 由于某些条件不满足（比如 asc_sync_wait 条件尚不满足）而不能发射的话，就会阻塞整条硬件管线。

因此，我们需要谨慎设计各个操作的发起顺序，以让各种硬件资源得到充分的利用。

优化技巧 Copy Chapter

gather2 Copy Section

本算子需要从显存中以 token 为粒度读取 KV。倘若我们为每个 KV 都发起一次 MTE2 拷贝请求，那么总拷贝请求数就会过多，超过 MTE2 的 Issue Queue 队列深度（16）。为此，我们使用了一种技巧：设置 MTE2 请求的 burst_count=2，src_stride 为两个 token 开始地址之差，这样就可以一次性拷贝两个 KV token 了，节约 MTE2 拷贝请求数量并降低发射 overhead。MTE2 拷贝的各个参数使用 SIMD VF 计算得到，避免阻塞 scalar 核心。

Decoding KV Cache Layout 优化 Copy Section

在 先前的 Flash MLA 中，KV Cache 中的每个块（KV Block）的前面是该块中所有 token 的 data（也即量化得到的 FP8 或 FP4），然后才是所有 token 的 scale。但是这种设计对于华为昇腾来说并不友好：每个 token 需要两次 MTE2 拷贝才能拿到 data 和 scale。所以，在本算子中，我们将一个 token 的 data 和 scale 相邻放置，这样就不必分别拷贝 data 和 scale 了。

高速 ND2NZ Copy Section

前文提到，我们需要一个 SIMD VF 来完成 KV 的 ND2NZ layout 转换。这个 SIMD VF 的效率直接关系到整个算子的效率，特别是启动第一个 QK 矩阵乘法之前的“空窗期”的长度。为此，我们对该 SIMD VF 做出了如下优化：

• 在 SIMD VF 中，使用 asc_loadalign 读取 ND 格式的 KV，并使用 vsstb 将其中的每个 32B 数据块以一定间隔写回 UB，如此便完成了 ND2NZ 的转换。

• 在输出的 NZ buffer 中，额外 pad 一行，以保证写回时没有 UB 的 bank conflict。

• 使用 asc_loadalign_postupdate 与 vsstb 自带的“自增 pointer”的功能，免除更新指针的开销，防止 SIMD VF 的小 Scalar Core 成为新的瓶颈。

• 一次性处理一对相邻的 256B 元素，以利用满 UB 的所有 bank 的带宽。

最终该 VF 可以打满 UB 的读、写吞吐。

icache miss 问题 Copy Section

华为昇腾 NPU 上的 Instruction Cache (icache) 大小相对有限，只有 32K (CUBE Core)、16K (VECTOR Core) 或 8K (SIMD VF)，因此需要注意控制代码的总体积，防止因为代码过大而导致指令缓存出现频繁的换入换出。值得一提的技巧包括在 SIMD VF 中使用硬件循环指令（VLOOP）而不是 #pragma unroll 展开循环、以及需要小心函数的内联会导致代码体积成倍数地增长（因此在 VECTOR Core 的结尾把 Systolic array 的 main loop 和 drain loop 给 fuse 在了一起）。

L1 bank 配平 Copy Section

Ascend 950 NPU 的 L1 分为两个 bank，分别掌控着 0 - 256K 以及 256K - 512K 两段地址空间。为了充分利用这两个 bank 的带宽，我把每个 Tensor（比如 Q、量化好的 K 等）都切成两份，分别放在两个 bank 中，以随时随地地打满两个 bank 的带宽。

skip-scale Copy Section

Skip scale 技术具有许多好处，包括：

• 减少 FixPipe 的拷贝量与 VF 计算量

• 降低对 UB 读、写带宽的压力

• 通过减少计算与数据搬运量，降低算子运行时所需的功率，进而获得更高的频率并提高性能

20260930 Ascend Prefill Deep Dive (`docs/20260930-ascend-prefill-deep-dive.md`) Copy Section

A Deep Dive Into the Ascend Sparse Attention Forward Kernel

Introduction Copy Chapter

On September 30, 2026, DeepSeek open-sourced its inference infrastructure and core components for the Huawei Ascend platform, including the Ascend sparse attention prefill and decoding kernels in this repository (deepseek-ai/FlashMLA). Under typical workloads of the DeepSeek V4.1 model, this kernel reaches 410 TFlops during prefill and 360 TFlops during decoding, i.e. 95% and 83% of the hardware's theoretical peak performance, respectively.

In this article, I will systematically describe the design rationale, pipeline schedule, and optimization techniques of this kernel.

Algorithm Recap Copy Chapter

This kernel computes the forward pass of the attention in the DeepSeek Sparse Attention (DSA) architecture used by the DeepSeek V4.1 model, covering both prefill and decoding. It takes a number of q tokens and kv tokens and, following the indices table, lets each q token attend to only the "most important" kv tokens, thereby saving a large amount of compute while preserving model quality.

Mathematical Formulation Copy Section

In this kernel, different q tokens are always processed separately. For simplicity, we consider only a single Q token in the remainder of this section.

Given:

• Q: [h_q, d_qk]

• KV: [s_kv, d_qk]

• indices: [topk]

• sm_scale: a scalar

the forward pass must compute:

• gathered_kv = KV.index_select(indices, dim=0), [topk, d_qk]

• P = Q @ gathered_kv^T, [h_q, topk]

• S = softmax(P*sm_scale, dim=-1), [h_q, topk]

• out = S @ gathered_kv[:, :d_v], [h_q, d_v]

Computation Flow Copy Section

Because softmax has a global dependency, we compute it with the same online softmax algorithm as Flash Attention, dynamically updating the current per-row maximum of P, the per-row lse (log sum exp), and the output accumulator computed so far. The steps are as follows:

1. Split the topk indices into blocks of size B_TOPK (B_TOPK is usually 64 or 96)

2. Maintain three quantities: running_max, running_max_for_softmax, and running_sumexp, all of shape [h_q] and initialized to $-\infty$, $-\infty$, and 0. Also maintain an output accumulator out_accum of shape [h_q, d_v]. All of these are float32.

3. For each block:

   1. Read KV sparsely from global memory according to indices[B_TOPK * k, B_TOPK * k + B_TOPK], obtaining this block's gathered_kv, [B_TOPK, d_qk]

   2. Compute P = Q @ gathered_kv^T, [h_q, B_TOPK]

   3. Compute the per-row maximum of P, cur_max, [h_q]

   4. Update running_max: running_max = max(running_max, cur_max) (in other words, running_max always reflects the current per-row maximum)

   5. Check whether some row i satisfies running_max[i] - running_max_for_softmax[i] > RESCALE_THRES, where RESCALE_THRES is a hyperparameter, usually set to $6$ in our kernel

   6. If some row i triggers this condition, enter the online softmax update

      1. out_accum <- out_accum * exp(running_max_for_softmax - running_max).unsqueeze(-1)

      2. running_sumexp <- running_sumexp * exp(running_max_for_softmax - running_max)

      3. running_max_for_softmax <- running_max

      4. This is equivalent to updating the max used by the online softmax, so that exp(running_max - running_max_for_softmax) is always $\le \exp(6)$

      5. If no row i triggers this condition, skip this step. This optimization is called "skip scale"

   7. Compute S = exp(P - running_max_for_softmax), cast S to bfloat16, then multiply it with gathered_kv and add the result to out_accum (out_accum += S.to(torch.bfloat16) @ gathered_kv)

   8. Update running_sumexp += S.sum(dim=-1)

Discussion Copy Chapter

Before we formally begin the pipeline design and optimization techniques of the kernel, let us go through a few simple discussions to settle the overall design direction.

How Many q Heads Should Each AI Core Handle? Copy Section

The first question is: how many q heads per AI Core is appropriate?

Let us first compute the compute-to-memory ratio, assuming the CUBE Core runs at full throughput (4096 FMA per cycle per AI Core) and the total L2 bandwidth is 5 TB/s:

d = 512
l2_total_bw = 5e12
num_cores = 32
freq = 1650e6
l2_bw = l2_total_bw / num_cores / freq  # Per AI Core per cycle

for h_q in [16, 32, 64, 128]:
    for topk in [512, 1024]:
        cube_cycles = h_q * topk * (d+d) // 4096
        qo_mem_vol = h_q * d * 2 * 2
        kv_mem_vol = topk * d * 2
        mem_cycles = (qo_mem_vol + kv_mem_vol) / l2_bw
        print(f"{h_q:3d}, {topk=:4d}, {cube_cycles=:6d}, {mem_cycles=:6.0f}")

16, topk= 512, cube_cycles=  2048, mem_cycles=  5883
 16, topk=1024, cube_cycles=  4096, mem_cycles= 11419
 32, topk= 512, cube_cycles=  4096, mem_cycles=  6229
 32, topk=1024, cube_cycles=  8192, mem_cycles= 11765
 64, topk= 512, cube_cycles=  8192, mem_cycles=  6921
 64, topk=1024, cube_cycles= 16384, mem_cycles= 12457
128, topk= 512, cube_cycles= 16384, mem_cycles=  8305
128, topk=1024, cube_cycles= 32768, mem_cycles= 13841

As we can see, the kernel is compute-bound only when head >= 64. At head=64, compute and memory are roughly balanced.

Meanwhile, taking 128 heads per AI Core would cause many problems in terms of SRAM capacity. Specifically, the L0C would be completely filled by O ($128 \times 512$, float32, 256KB), so the accumulator would have to be kept on the VECTOR Core; that in turn requires two AIVs and a fine-grained pipeline (pipelined copy - accumulate), which is very hard to write. In addition, according to micro-benchmark results, an MMAD with M = 64 is already enough to saturate CUBE Core throughput, so 128 heads is unnecessary.

Where Should the FA Output Accumulator Live? Copy Section

Whenever the running max differs too much from the scaling max currently in use, Flash Attention has to multiply the output by $\exp(\text{new_max} - \text{old_scaling_max})$; this is called a "rescale". However, data in L0C can only be moved out by the FixPipe or used by the CUBE Core as an accumulator, and cannot be multiplied by a scalar in place. There are two feasible approaches:

1. Maintain an accumulator on the VECTOR Core, and whenever O needs to be rescaled, bring over the accumulation result from the CUBE Core and add it to the current accumulator

2. Keep the accumulator on the CUBE Core at all times, and when a rescale is needed, pull the accumulator up to the VECTOR Core, rescale it, and put it back onto the CUBE Core (by multiplying by an identity matrix)

   1. This consumes extra CUBE throughput and hurts performance. Moving a 64x512 matrix from UB to L0C requires at least one 64x64x512 float matmul

   2. Moreover, if no precision is to be lost, HF32 mode must be turned off; the CUBE Core then delivers only 256 FMA / cycle / AI Core, and the 64x64x512 matmul takes 8192 cycles, which is entirely unacceptable

We therefore chose option 1 in the end.

Can We Use 1C1V? Copy Section

The next question is how many VECTOR Cores each CUBE Core should be paired with during computation. Each AI Core of the Huawei Ascend 950 has one CUBE Core and two VECTOR Cores. If the number of VECTOR Cores required per CUBE Core could be kept at one (1C1V), the other VECTOR Core could run some other lightweight work (such as other modules in the model or cross-device communication), improving the overall MFU.

After some analysis, the answer is no, for the following reasons:

1. MTE2 Issue Queue Entry bound. On the Huawei Ascend 950 NPU, a single VECTOR Core can have at most 16 "issued but not yet completed" MTE2 copy requests at any moment. Even with the gather2 trick (see below), only 32 tokens can be in flight at a time, which is far too few.

2. Dequantize bound. During decoding, the kernel needs to read the KV Cache in FP8 or FP4 format, dequantize it to bfloat16, and then feed it to the CUBE Core. According to micro-benchmarks, with 1C1V, dequantization cannot keep up with the CUBE Core, making dequantization the bottleneck.

3. FixPipe bound. When only one CUBE Core and one VECTOR Core are used, the FixPipe moves data out at only 128 Byte/cycle/AI Core. In this case the MMA compute cycles per block are MMA_M * B_TOPK * (512+512) / 4096 = B_TOPK * 16, while the cycles needed for the FixPipe to move the data out are (MMA_M * 512 + MMA_M * B_TOPK) * sizeof(float) / 128 = 1024 + B_TOPK * 2. Here MMA_M is the number of q heads handled by each AI Core, which per the discussion above should be 64, and B_TOPK is the tile size when looping along the KV dimension in Flash Attention. For the FixPipe not to become the bandwidth bottleneck, we need B_TOPK * 16 > 1024 + B_TOPK * 2, which gives B_TOPK > 73.14. To leave enough margin, B_TOPK has to be 96 or 128, which on one hand makes UB space even tighter and on the other hand worsens the MTE2 Issue Queue Entry bound mentioned above.

Therefore, we decided in the end to give each CUBE Core two VECTOR Cores, so as to share the overhead of MTE2 issue, dequantization, FixPipe transfers, softmax, and so on.

When Should the ND2NZ Layout Conversion of KV Happen? Copy Section

The next question: when should the ND2NZ layout conversion of KV happen?

KV is stored in device memory in ND format, with shape [s_kv, d_qk] and stride [d_qk, 1]. The CUBE Core, however, requires its input matrices in NZ layout. Normally, we would perform the ND2NZ conversion on the fly while moving data from device memory into the L1 buffer, but our KV tokens are moved "sparsely" from device memory onto the chip by MTE2 one token at a time, so on-the-fly ND2NZ conversion cannot be enabled.

One idea is to treat a single token (512 bfloat16 values) as 32 blocks of 32B when issuing the GM -> L1 copy: these 32 blocks are contiguous in device memory, and are laid out in the L1 buffer with a stride of B_TOPK * 32B. Later experiments, however, showed that the BIU inside the AI Core (the unit that issues memory access instructions) triggers access coalescing only when both the source and the destination are contiguous. In a case like this, where only the source is contiguous and the destination is not, accesses cannot be coalesced, so L2 bandwidth utilization is very low and the overall bottleneck lies in L2 bandwidth.

Therefore, the only option is to first copy KV from device memory to UB, convert it from ND layout to NZ layout with a SIMD VF, and then copy it to L1.

How Should the skip-scale Be Synchronized? Copy Section

The "Algorithm Recap" section above mentioned the "skip-scale" technique, whose idea is that as long as the max currently used for online softmax is not much smaller than the true max, this round's update and accumulation of the output can be skipped.

This operation is even more valuable on Huawei NPUs, because scale-O is awkward there: it requires using the FixPipe to copy O from L0C to UB and then accumulating it into out_accum in UB via a VF, wasting the bandwidth of the FixPipe, UB, and the SIMD VF.

However, implementing this operation on the NPU has one difficulty: the signal "should we skip-scale this round" is produced by the VECTOR Core, but its consumers are both the CUBE Core and the VECTOR Core (because the CUBE Core needs it to decide whether to clear O in L0C). In the end, we decided to use an SS buffer and cross-core flags to pass the skip-scale signal.

Pipeline Schedule Copy Chapter

At this point, the kernel's overall design direction is finalized: 64 heads / AI Core, 1C2V, the output accumulator in UB, and KV routed through UB and a SIMD VF for the ND2NZ conversion.

The next step is the pipeline design. Since this kernel should be bound by CUBE Core throughput, our goal is to keep the matmul units on the CUBE Core saturated at all times, and not to let other operations (such as softmax) become the bottleneck.

The most basic pipeline idea is roughly as follows (the bracketed numbers below indicate which KV block an operation belongs to):

1. Compute P[i] = Q @ K[i]^T (@ denotes matrix multiplication)

2. Compute O += S[i-1] @ V[i-1]

3. Compute P[i+1] = Q @ K[i+1]^T

4. Compute O += S[i] @ V[i]

5. ...

In other words, the O += S @ V of the previous block is deferred so as to leave time for the softmax operations: P[i] is ready as soon as (1) finishes, but it is not used until (4) starts, so there is ample time in between (about 1024 cycles) for the FixPipe to copy P[i] from L0C to UB, for softmax, and for copying S[i] from UB to L1.

Reading the indices (copying them from device memory to UB), preprocessing the indices (preparing the MTE copy parameters for the gather2 described below), issuing the MTE2 copies, and synchronizing the SS buffer signal in the skip-scale technique make this pipeline somewhat more complex. The CUBE Core side follows the schedule above, while the VECTOR Core side is scheduled as follows:

1. Issue the index read for KV block i+5

2. Issue the ND2NZ conversion and the UB -> L1 MTE3 copy for KV block i+3

3. Wait for P[i+1] = Q @ k[i+1]^T to complete, issue the softmax computation for KV block i+1, and copy the result to L1

4. Issue the index preprocessing for KV block i+5

5. Issue the O scale for KV block i, if this block does not skip-scale O

6. Synchronize with the SIMD VF, get whether KV block i+1 needs to skip-scale O, and write it into the SS buffer

7. Issue the MTE2 copy for KV block i+5, copying KV from device memory to UB

Why is this pipeline so complex?

• First, the kernel relies on the SS buffer to pass the skip-scale signal, but only the scalar core can read and write the SS buffer, while the skip-scale signal is generated by the SIMD VF; this leads to a blocking synchronization between the SIMD VF and the scalar core. If this is not handled well, the scalar core will wait for the SIMD VF to finish and other units will sit idle.

• Second, Huawei's hardware pipes (PIPEs) issue and execute in order. For PIPE_V, for example, the order in which SIMD VFs are issued must exactly match the order in which they execute, so if an earlier SIMD VF cannot be issued because some condition is not satisfied (for instance, the asc_sync_wait condition is not yet met), the whole hardware pipe is blocked.

We therefore need to design the issue order of the various operations carefully, so that all hardware resources are fully utilized.

Optimization Techniques Copy Chapter

gather2 Copy Section

This kernel needs to read KV from device memory at token granularity. If we issued one MTE2 copy request for every KV token, the total number of copy requests would be too large and would exceed the MTE2 issue queue depth (16). To avoid this, we use a trick: set the MTE2 request's burst_count=2 and src_stride to the difference between the start addresses of the two tokens, so that two KV tokens can be copied at once, saving MTE2 copy requests and reducing issue overhead. The parameters of the MTE2 copy are computed with a SIMD VF, so the scalar core is not blocked.

Decoding KV Cache Layout Optimization Copy Section

In the previous Flash MLA, each block in the KV Cache (KV Block) stores the data of all tokens in the block first (i.e. the quantized FP8 or FP4 values), followed by the scales of all tokens. This design is not friendly to the Huawei Ascend: each token needs two MTE2 copies to fetch its data and scale. So in this kernel, we place a token's data and scale next to each other, so that data and scale no longer have to be copied separately.

Fast ND2NZ Copy Section

As mentioned above, we need a SIMD VF to perform the ND2NZ layout conversion of KV. The efficiency of this SIMD VF directly affects the efficiency of the whole kernel, especially the length of the "idle window" before the first QK matmul is launched. We therefore made the following optimizations to this SIMD VF:

• In the SIMD VF, use asc_loadalign to read the ND-format KV and vsstb to write each 32B data block back to UB at a certain interval, which completes the ND2NZ conversion.

• Pad one extra row in the output NZ buffer, to ensure that the write-back has no UB bank conflict.

• Use asc_loadalign_postupdate and the "auto-increment pointer" feature of vsstb to eliminate the pointer-update overhead, so that the SIMD VF's small scalar core does not become a new bottleneck.

• Process a pair of adjacent 256B elements at a time, so as to use the bandwidth of all UB banks to the full.

In the end, this VF achieves peak UB read and write throughput.

The icache Miss Problem Copy Section

Huawei Ascend NPUs have a relatively small instruction cache (icache): only 32K (CUBE Core), 16K (VECTOR Core), or 8K (SIMD VF). We therefore need to keep the total code size under control, to prevent code that is too large from causing frequent instruction cache thrashing. Useful tricks include using the hardware loop instruction (VLOOP) in the SIMD VF instead of unrolling loops with #pragma unroll, and being careful that function inlining can multiply code size (which is why the systolic array's main loop and drain loop were fused together at the end of the VECTOR Core).

L1 Bank Balancing Copy Section

The Ascend 950 NPU's L1 is split into two banks, which respectively control the 0 - 256K and 256K - 512K address ranges. To make full use of the bandwidth of these two banks, I split every tensor (such as Q and the quantized K) into two halves and place them in the two banks respectively, so that the bandwidth of both banks is saturated at all times.

skip-scale Copy Section

The skip-scale technique has many benefits, including:

• Reducing the volume of FixPipe copies and the amount of VF computation

• Reducing the pressure on UB read and write bandwidth

• Reducing the computation and data movement, and thus the power drawn by the kernel, which in turn allows a higher frequency and improves performance

--- METRICS ---

• Files Extracted: 6
• Estimated Token Budget: ~16096 tokens
• Canonical Reference: https://codewiki.google/github.com/deepseek-ai/FlashMLA