Add work buffer size calculation for the ds4_comp op (#2307)

This commit is contained in:
Kawrakow 2026-08-13 15:23:38 +02:00 committed by GitHub
parent ff141691a1
commit 3c949f3399
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 2 deletions

View File

@ -28972,8 +28972,15 @@ struct ggml_cplan ggml_graph_plan(const struct ggml_cgraph * cgraph, int n_threa
} break;
case GGML_OP_INDEXER_TOPK:
{
size_t size = iqk_idx_topk_work_buffer_size(node, n_tasks);
cur = MAX(cur, size);
cur = iqk_idx_topk_work_buffer_size(node, n_tasks);
} break;
case GGML_OP_DS4_COMP:
{
if (node->op_params[0] == 0) {
struct ggml_tensor * idx = node->src[2];
int ratio = idx->ne[0] / (2*node->ne[1]);
cur = 32*(4*ratio + 3)*sizeof(float)*n_tasks;
}
} break;
case GGML_OP_CROSS_ENTROPY_LOSS:
{