华为计算微信公众号
昇腾AI开发者公众号
华为计算微博
华为计算今日头条
flash_attention.py 测试用例，基于GPU（CUDA）使用flash_attn_with_kvcache接口进行测试，调用代码如下：
# q: [batch_size, q_seq_len, q_head_num, head_dim] # k: [batch_size, max_kv_len, kv_head_num, head_dim] # v: [batch_size, max_kv_len, kv_head_num, head_dim] out = flash_attn_with_kvcache( q=q, k_cache=k_cache, v_cache=v_cache, cache_seqlens=cache_lens, cache_batch_idx=cache_slot_ids, softmax_scale=self.softmax_scale, causal=True )
flash_attn_with_kvcache接口描述
基于NPU可以使用什么接口进行替代？如何传参？
我要发帖子
点赞
收藏
回复
分享
举报
已解决 已解决
flash_attention.py 测试用例，基于GPU（CUDA）使用flash_attn_with_kvcache接口进行测试，调用代码如下：
# q: [batch_size, q_seq_len, q_head_num, head_dim] # k: [batch_size, max_kv_len, kv_head_num, head_dim] # v: [batch_size, max_kv_len, kv_head_num, head_dim] out = flash_attn_with_kvcache( q=q, k_cache=k_cache, v_cache=v_cache, cache_seqlens=cache_lens, cache_batch_idx=cache_slot_ids, softmax_scale=self.softmax_scale, causal=True )flash_attn_with_kvcache接口描述
基于NPU可以使用什么接口进行替代？如何传参？