刚刚,DeepSeek开源FlashMLA,推理加速核心技术,Star量飞涨中
- 2025-02-24 10:28:00
- 刘大牛 转自文章
- 503
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_kvcachetile_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 |
| 网址: | ai.tmqcjr.com |