Oneflow-Inc / one-codegeex

Apache License 2.0
7 stars 1 forks source link

try to add fused_multi_head_attention_inference_v2 #16

Closed BBuf closed 1 year ago

BBuf commented 1 year ago

这个pr将fmha升级为使用fused_multi_head_attention_inference_v2 op,好处是我们不需要在fmha之前做一系列的形状变换工作,只需要在fused_multi_head_attention_inference_v2配置一下q,k,v的layerout就可以了,可以降低eager view机制设置到的调度开销,并且fused_codegeex_qkv_reshape也不再需要了。