| 400 | |
| 401 | template <typename T> |
| 402 | void MHAForwardProxyBase::exec_internal(MHA_PROXY_FORWARD_EXEC_PARAM) { |
| 403 | auto wksp_bundle = get_workspace_bundle( |
| 404 | MHA_PROXY_FORWARD_TENSOR_TO_LAYOUT_CALL, workspace.raw_ptr); |
| 405 | auto mask_bundle = get_mask_reservespace_bundle( |
| 406 | MHA_PROXY_FORWARD_TENSOR_TO_LAYOUT_CALL, mask_reservespace.raw_ptr()); |
| 407 | auto othr_bundle = get_othr_reservespace_bundle( |
| 408 | MHA_PROXY_FORWARD_TENSOR_TO_LAYOUT_CALL, othr_reservespace.raw_ptr()); |
| 409 | |
| 410 | m_matmul_opr->param().transposeA = false; |
| 411 | m_matmul_opr->param().transposeB = false; |
| 412 | TensorND q, k, v; |
| 413 | if (param.qproj_size) { |
| 414 | if (param.num_heads == 1) { |
| 415 | q = TensorND{othr_bundle.get_workspace(0).raw_ptr, m_q_layout}; |
| 416 | } else { |
| 417 | q = TensorND{wksp_bundle.get_workspace(0).raw_ptr, m_q_layout}; |
| 418 | } |
| 419 | TensorND qweight{qkvo_weight_bias.ptr<T>() + m_wq_off, m_wq_layout}; |
| 420 | matmul_exec(m_matmul_opr, queries, qweight, q, wksp_bundle.get_workspace(1)); |
| 421 | if (param.qbias) { |
| 422 | m_add_opr->exec(q, {qkvo_weight_bias.ptr<T>() + m_bq_off, m_bq_layout}); |
| 423 | } |
| 424 | } else { |
| 425 | q = TensorND{queries.raw_ptr(), queries.layout}; |
| 426 | } |
| 427 | if (param.kproj_size) { |
| 428 | if (param.num_heads == 1) { |
| 429 | k = TensorND{othr_bundle.get_workspace(1).raw_ptr, m_k_layout}; |
| 430 | } else { |
| 431 | k = TensorND{wksp_bundle.get_workspace(2).raw_ptr, m_k_layout}; |
| 432 | } |
| 433 | TensorND kweight{qkvo_weight_bias.ptr<T>() + m_wk_off, m_wk_layout}; |
| 434 | matmul_exec(m_matmul_opr, keys, kweight, k, wksp_bundle.get_workspace(3)); |
| 435 | if (param.kbias) { |
| 436 | m_add_opr->exec(k, {qkvo_weight_bias.ptr<T>() + m_bk_off, m_bk_layout}); |
| 437 | } |
| 438 | } else { |
| 439 | k = TensorND{keys.raw_ptr(), keys.layout}; |
| 440 | } |
| 441 | if (param.vproj_size) { |
| 442 | if (param.num_heads == 1) { |
| 443 | v = TensorND{othr_bundle.get_workspace(2).raw_ptr, m_v_layout}; |
| 444 | } else { |
| 445 | v = TensorND{wksp_bundle.get_workspace(4).raw_ptr, m_v_layout}; |
| 446 | } |
| 447 | TensorND vweight{qkvo_weight_bias.ptr<T>() + m_wv_off, m_wv_layout}; |
| 448 | matmul_exec(m_matmul_opr, values, vweight, v, wksp_bundle.get_workspace(5)); |
| 449 | if (param.vbias) { |
| 450 | m_add_opr->exec(v, {qkvo_weight_bias.ptr<T>() + m_bv_off, m_bv_layout}); |
| 451 | } |
| 452 | } else { |
| 453 | v = TensorND{values.raw_ptr(), values.layout}; |
| 454 | } |
| 455 | |
| 456 | // nq/nk/nv: norm to multihead |
| 457 | auto relayout_to_multihead = [&](TensorND& q, TensorND& nq) { |
| 458 | size_t batch = q.layout[0]; |
| 459 | size_t seq = q.layout[1]; |
nothing calls this directly
no test coverage detected