速看！DeepSeek开源周第一弹：FlashMLA加速GPU计算！

git clone https://github.com/deepseek-ai/FlashMLA.git

python setup.py install

python tests/test_flash_mla.py

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)

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

Hopper GPUs
CUDA 12.3 and above
PyTorch 2.0 and above