topk_kernels.c File Reference

Top-K selection kernels for MoE router dispatch. More...

#include <stdint.h>
#include <stddef.h>
#include <float.h>
#include <math.h>

Go to the source code of this file.

Functions
int	argmax_f32 (const float *scores, int n)
	Find index of maximum value. More...
void	topk_batched_f32 (const float *scores, int num_tokens, int n_experts, int k, int *indices, float *weights)
	Batched top-K selection for multiple tokens. More...
void	topk_f32 (const float *scores, int n, int k, int *indices, float *values)
	Find top-K indices and values from a score vector. More...
void	topk_softmax_f32 (const float *scores, int n, int k, int *indices, float *weights)
	Find top-K indices with softmax-normalized weights. More...

Detailed Description

Top-K selection kernels for MoE router dispatch.

CK-ENGINE KERNEL RULES:

1. NO malloc/free - memory via bump allocator, pointers passed in
2. NO OpenMP - parallelization at orchestrator/codegen layer
3. API must define: inputs, outputs, workspace, and memory layouts
4. Pure computation - deterministic, no side effects

After changes: make test && make llamacpp-parity-full

Provides efficient top-K selection from a score vector. Used in Mixture-of-Experts models to select which experts process each token.

Operations:

• topk_f32: Find top-K indices and values from N scores
• topk_softmax_f32: Top-K with softmax normalization of selected scores

Definition in file topk_kernels.c.

Function Documentation

argmax_f32()

int argmax_f32	(	const float *	scores,
		int	n
	)

Find index of maximum value.

Parameters

scores	Input scores [n]
n	Number of scores

Returns

Index of maximum value

Definition at line 226 of file topk_kernels.c.

{
    if (!scores || n <= 0) {
        return -1;
    }
 
    int max_idx = 0;
    float max_val = scores[0];
 
#ifdef __AVX512F__
    /* AVX-512 vectorized argmax for large arrays */
    if (n >= 16) {
        __m512 vmax = _mm512_set1_ps(-FLT_MAX);
        __m512i vidx = _mm512_setzero_si512();
        __m512i vcur_max_idx = _mm512_setzero_si512();
 
        int i = 0;
        for (; i + 16 <= n; i += 16) {
            __m512 v = _mm512_loadu_ps(&scores[i]);
            __m512i cur_idx = _mm512_add_epi32(
                _mm512_set1_epi32(i),
                _mm512_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15)
            );
 
            __mmask16 gt_mask = _mm512_cmp_ps_mask(v, vmax, _CMP_GT_OQ);
            vmax = _mm512_mask_blend_ps(gt_mask, vmax, v);
            vcur_max_idx = _mm512_mask_blend_epi32(gt_mask, vcur_max_idx, cur_idx);
        }
 
        /* Horizontal reduction */
        float vals[16];
        int idxs[16];
        _mm512_storeu_ps(vals, vmax);
        _mm512_storeu_si512(idxs, vcur_max_idx);
 
        max_val = vals[0];
        max_idx = idxs[0];
        for (int j = 1; j < 16; j++) {
            if (vals[j] > max_val) {
                max_val = vals[j];
                max_idx = idxs[j];
            }
        }
 
        /* Handle remainder */
        for (; i < n; i++) {
            if (scores[i] > max_val) {
                max_val = scores[i];
                max_idx = i;
            }
        }
 
        return max_idx;
    }
#endif
 
    /* Scalar fallback */
    for (int i = 1; i < n; i++) {
        if (scores[i] > max_val) {
            max_val = scores[i];
            max_idx = i;
        }
    }
 
    return max_idx;
}

topk_batched_f32()

void topk_batched_f32	(	const float *	scores,
		int	num_tokens,
		int	n_experts,
		int	k,
		int *	indices,
		float *	weights
	)

Batched top-K selection for multiple tokens.

Parameters

scores	Input scores [num_tokens, n_experts]
num_tokens	Number of tokens
n_experts	Number of experts
k	Number of experts to select per token
indices	Output: selected expert indices [num_tokens, k]
weights	Output: routing weights [num_tokens, k] (can be NULL for no softmax)

Definition at line 191 of file topk_kernels.c.

{
    if (!scores || !indices || num_tokens <= 0 || n_experts <= 0 || k <= 0) {
        return;
    }
 
    for (int t = 0; t < num_tokens; t++) {
        const float *token_scores = scores + t * n_experts;
        int *token_indices = indices + t * k;
 
        if (weights) {
            float *token_weights = weights + t * k;
            topk_softmax_f32(token_scores, n_experts, k, token_indices, token_weights);
        } else {
            topk_f32(token_scores, n_experts, k, token_indices, NULL);
        }
    }
}

References topk_f32(), and topk_softmax_f32().

topk_f32()

void topk_f32	(	const float *	scores,
		int	n,
		int	k,
		int *	indices,
		float *	values
	)

Find top-K indices and values from a score vector.

Parameters

scores	Input scores [n]
n	Number of scores (e.g., number of experts)
k	Number of top scores to select
indices	Output: indices of top-K scores [k], sorted descending by value
values	Output: top-K score values [k], sorted descending (can be NULL)

Definition at line 49 of file topk_kernels.c.

{
    if (!scores || !indices || n <= 0 || k <= 0) {
        return;
    }
 
    /* Clamp k to n */
    if (k > n) {
        k = n;
    }
 
    /* Initialize with first k elements */
    float local_values[k];
    for (int i = 0; i < k; i++) {
        indices[i] = i;
        local_values[i] = scores[i];
    }
 
    /* Find the minimum in our current top-k */
    int min_idx = 0;
    for (int i = 1; i < k; i++) {
        if (local_values[i] < local_values[min_idx]) {
            min_idx = i;
        }
    }
 
    /* Scan remaining elements */
    for (int i = k; i < n; i++) {
        if (scores[i] > local_values[min_idx]) {
            /* Replace the minimum */
            indices[min_idx] = i;
            local_values[min_idx] = scores[i];
 
            /* Find new minimum */
            min_idx = 0;
            for (int j = 1; j < k; j++) {
                if (local_values[j] < local_values[min_idx]) {
                    min_idx = j;
                }
            }
        }
    }
 
    /* Sort results in descending order (simple insertion sort for small k) */
    for (int i = 1; i < k; i++) {
        float val = local_values[i];
        int idx = indices[i];
        int j = i - 1;
        while (j >= 0 && local_values[j] < val) {
            local_values[j + 1] = local_values[j];
            indices[j + 1] = indices[j];
            j--;
        }
        local_values[j + 1] = val;
        indices[j + 1] = idx;
    }
 
    /* Copy values if output requested */
    if (values) {
        for (int i = 0; i < k; i++) {
            values[i] = local_values[i];
        }
    }
}

Referenced by topk_batched_f32(), and topk_softmax_f32().

topk_softmax_f32()

void topk_softmax_f32	(	const float *	scores,
		int	n,
		int	k,
		int *	indices,
		float *	weights
	)

Find top-K indices with softmax-normalized weights.

Parameters

scores	Input scores [n] (router logits)
n	Number of scores
k	Number of top scores to select
indices	Output: indices of top-K scores [k]
weights	Output: softmax-normalized weights for selected [k], sum to 1.0

Definition at line 134 of file topk_kernels.c.

{
    if (!scores || !indices || !weights || n <= 0 || k <= 0) {
        return;
    }
 
    if (k > n) {
        k = n;
    }
 
    /* First get top-K indices and values */
    float values[k];
    topk_f32(scores, n, k, indices, values);
 
    /* Compute softmax over the selected values */
    /* Find max for numerical stability */
    float max_val = values[0];
    for (int i = 1; i < k; i++) {
        if (values[i] > max_val) {
            max_val = values[i];
        }
    }
 
    /* Compute exp and sum */
    float sum = 0.0f;
    for (int i = 0; i < k; i++) {
        weights[i] = expf(values[i] - max_val);
        sum += weights[i];
    }
 
    /* Normalize */
    float inv_sum = 1.0f / sum;
    for (int i = 0; i < k; i++) {
        weights[i] *= inv_sum;
    }
}

References topk_f32().

Referenced by topk_batched_f32().