| 4597 | } |
| 4598 | |
| 4599 | static struct ggml_tensor * llm_build_sparse_axpy( |
| 4600 | struct ggml_context * ctx, |
| 4601 | struct ggml_tensor * w_t, |
| 4602 | struct ggml_tensor * x, |
| 4603 | struct ggml_tensor * sparse_idx, |
| 4604 | struct ggml_tensor * wt_gpu, |
| 4605 | struct ggml_tensor * gpu_index, |
| 4606 | struct ggml_tensor * gpu_bucket, |
| 4607 | const llm_build_cb_short & cb, |
| 4608 | const char * name, |
| 4609 | bool full_gpu) { |
| 4610 | std::string full_name = "ffn_" + std::string(name) + "_sparse"; |
| 4611 | ggml_tensor * out = nullptr; |
| 4612 | |
| 4613 | #ifdef GGML_USE_HIPBLAS |
| 4614 | // WARNING: THIS IS A HACK! |
| 4615 | // if wt_gpu->data is null |
| 4616 | // inference fails when model exceeds 40B on rocm device |
| 4617 | // so we just let wt_gpu->data point to itself |
| 4618 | |
| 4619 | wt_gpu->data = wt_gpu; |
| 4620 | |
| 4621 | #endif |
| 4622 | |
| 4623 | #ifdef GGML_USE_CUBLAS |
| 4624 | // Full offloading fast path |
| 4625 | if (full_gpu) { |
| 4626 | GGML_ASSERT(wt_gpu && "full_gpu but no wt_gpu"); |
| 4627 | out = ggml_axpy(ctx, wt_gpu, x, sparse_idx, NULL); |
| 4628 | ggml_cuda_assign_buffers_no_alloc(out); |
| 4629 | cb(out, (full_name).c_str()); |
| 4630 | return out; |
| 4631 | } |
| 4632 | #endif |
| 4633 | |
| 4634 | // TODO: should pass NULL as hybrid_aux when hybrid_split is false |
| 4635 | out = ggml_axpy(ctx, w_t, x, sparse_idx, gpu_index); |
| 4636 | cb(out, full_name.c_str()); |
| 4637 | |
| 4638 | #ifdef GGML_USE_CUBLAS |
| 4639 | if (wt_gpu) { |
| 4640 | ggml_tensor * out_gpu = ggml_axpy(ctx, wt_gpu, x, sparse_idx, gpu_bucket); |
| 4641 | cb(out_gpu, (full_name + "_gpu").c_str()); |
| 4642 | ggml_cuda_assign_buffers_no_alloc(out_gpu); |
| 4643 | out = ggml_add(ctx, out, out_gpu); |
| 4644 | ggml_cuda_assign_buffers_no_alloc(out); |
| 4645 | cb(out, (full_name + "_merged").c_str()); |
| 4646 | } |
| 4647 | #endif |
| 4648 | |
| 4649 | return out; |
| 4650 | } |
| 4651 | |
| 4652 | static struct ggml_tensor * llm_build_ffn_sparse( |
| 4653 | struct ggml_context * ctx, |