

BF16 块大小为 64 的分页 kvcache
Hopper GPU CUDA 12.3 及以上版本 PyTorch 2.0 及以上版本
安装
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,
)
...


