| 377 | } |
| 378 | |
| 379 | void MHAForwardProxyBase::deduce_layout(MHA_PROXY_FORWARD_LAYOUT_PARAM) { |
| 380 | if (!layout_ismatch(MHA_PROXY_FORWARD_CALL)) { |
| 381 | layout_refill(MHA_PROXY_FORWARD_CALL); |
| 382 | } |
| 383 | attn_weight = m_nx_layout; |
| 384 | out = m_out_layout; |
| 385 | size_t mask_size = get_mask_reservespace_in_bytes(MHA_PROXY_FORWARD_CALL); |
| 386 | size_t othr_size = get_othr_reservespace_in_bytes(MHA_PROXY_FORWARD_CALL); |
| 387 | mask_reservespace = TensorLayout{{mask_size}, dtype::Uint8()}; |
| 388 | othr_reservespace = TensorLayout{{othr_size / queries.dtype.size()}, queries.dtype}; |
| 389 | } |
| 390 | |
| 391 | void MHAForwardProxyBase::exec(MHA_PROXY_FORWARD_EXEC_PARAM) { |
| 392 | #define cb(DType) \ |
no test coverage detected