刚刚,DeepSeek开源FlashMLA,推理加速核心技术,Star量飞涨中
- 2025-02-24 10:28:00
- 刘大牛 转自文章
- 273


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



发表评论
文章分类
联系我们
联系人: | 透明七彩巨人 |
---|---|
Email: | weok168@gmail.com |