刚刚,DeepSeek开源FlashMLA,推理加速核心技术,Star量飞涨中
浏览252次
点赞0次
收藏0次


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



声明:本文转载自机器之心,转载目的在于传递更多信息,并不代表本社区赞同其观点和对其真实性负责,本文只提供参考并不构成任何建议,若有版权等问题,点击这里。
AI 中文社