deepseek_kernels.c File Reference

Scalar reference kernels for DeepSeek-style research ops. More...

#include <math.h>
#include <float.h>
#include <limits.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include "bf16_utils.h"
#include "ck_threadpool.h"

Go to the source code of this file.

Functions
void	deepseek_csa_attention_backward_f32 (const float *d_out, const float *q, const float *k, const float *v, const int *indices, const float *attn, float *d_q, float *d_k, float *d_v, int query_tokens, int key_tokens, int heads, int dim, int top_k, float scale)
void	deepseek_csa_attention_f32 (const float *q, const float *k, const float *v, const int *indices, float *out, float *attn, int query_tokens, int key_tokens, int heads, int dim, int top_k, float scale)
void	deepseek_dsa_topk_softmax_backward_f32 (const int *indices, const float *weights, const float *d_weights, float *d_scores, int tokens, int heads, int key_count, int top_k)
void	deepseek_dsa_topk_softmax_f32 (const float *scores, int *indices, float *weights, int tokens, int heads, int key_count, int top_k)
void	deepseek_hybrid_attention_f32 (const float *q, const float *k, const float *v, const int *indices, float *out, float *attn, int query_tokens, int key_tokens, int heads, int dim, int top_k, float scale, int mode)
void	deepseek_mhc_mix_backward_f32 (const float *d_out, const float *streams, const float *mix, float *d_streams, float *d_mix, int tokens, int n_streams, int dim)
void	deepseek_mhc_mix_f32 (const float *streams, const float *mix, float *out, int tokens, int n_streams, int dim)
void	deepseek_mla_attention_decode_f32 (const float *q, const float *k_cache, const float *v_cache, float *output, int num_heads, int num_kv_heads, int cache_len, int qk_head_dim, int v_head_dim, int max_seq_len, int cache_stride)
void	deepseek_mla_attention_decode_f32_workspace (const float *q, const float *k_cache, const float *v_cache, float *output, int num_heads, int num_kv_heads, int cache_len, int qk_head_dim, int v_head_dim, int max_seq_len, int cache_stride, float scale, float *scores, size_t scores_bytes)
void	deepseek_mla_attention_f32 (const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int qk_head_dim, int v_head_dim)
void	deepseek_mla_attention_f32_parallel_dispatch (const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int qk_head_dim, int v_head_dim, float scale, float *scores, size_t scores_bytes)
void	deepseek_mla_attention_f32_workspace (const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int qk_head_dim, int v_head_dim, float scale, float *scores, size_t scores_bytes)
void	deepseek_mla_kv_cache_batch_store_f32 (float *k_cache, float *v_cache, const float *k, const float *v, int num_tokens, int num_kv_heads, int qk_head_dim, int v_head_dim, int max_seq_len, int cache_stride)
void	deepseek_mla_kv_cache_store_f32 (float *k_cache, float *v_cache, const float *k, const float *v, int pos, int num_kv_heads, int qk_head_dim, int v_head_dim, int max_seq_len, int cache_stride)
void	deepseek_mla_kv_decompress_bf16 (const float *compressed_kv, const uint16_t *kv_b_proj, float *k_nope, float *value, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int v_dim)
void	deepseek_mla_kv_decompress_bf16_parallel_dispatch (const float *compressed_kv, const uint16_t *kv_b_proj, float *k_nope, float *value, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int v_dim)
void	deepseek_mla_kv_decompress_bf16_token_range (const float *compressed_kv, const uint16_t *kv_b_proj, float *k_nope, float *value, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int v_dim, int token_begin, int token_end)
void	deepseek_mla_kv_decompress_f32 (const float *compressed_kv, const float *kv_b_proj, float *k_nope, float *value, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int v_dim)
void	deepseek_mla_partial_rope_concat_f32 (const float *q_nope, const float *q_pe, const float *k_nope, const float *k_pe, const float *cos, const float *sin, float *query, float *key, int tokens, int heads, int qk_nope_dim, int qk_rope_dim)
void	deepseek_mla_partial_rope_concat_packed_bf16_storage (const float *q_packed, const float *k_nope, const float *kv_a_packed, const float *cos, const float *sin, float *query, float *key, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int qk_rope_dim)
void	deepseek_mla_partial_rope_concat_packed_f32 (const float *q_packed, const float *k_nope, const float *kv_a_packed, const float *cos, const float *sin, float *query, float *key, int tokens, int heads, int kv_lora_rank, int qk_nope_dim, int qk_rope_dim)
static size_t	ds_mhc_idx (int t, int s, int d, int n_streams, int dim)
static size_t	ds_mix_idx (int t, int out_s, int in_s, int n_streams)
static void	ds_mla_apply_kimi_rope (const float *src, float *dst, const float *cos_row, const float *sin_row, int dim)
static void	ds_mla_attention_f32_query_range (const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int qk_head_dim, int v_head_dim, float scale, float *scores, int query_begin, int query_end, int query_step)
static void	ds_mla_attention_f32_work (int ith, int nth, void *opaque)
static float	ds_mla_bf16_round (float value)
static void	ds_mla_kv_decompress_bf16_rows (int begin, int end, void *opaque)
static size_t	ds_mla_thd_idx (int t, int h, int d, int heads, int dim)
static size_t	ds_mla_tok_idx (int t, int d, int dim)
static size_t	ds_qkv_idx (int token, int head, int d, int heads, int dim)
static void	ds_softmax (float *x, int n)
void	topk_softmax_backward_f32 (const int *indices, const float *weights, const float *d_weights, float *d_scores, int num_tokens, int n_experts_or_keys, int k)
	Backward for hard top-k followed by softmax over selected values.

Detailed Description

Scalar reference kernels for DeepSeek-style research ops.

These kernels intentionally prioritize explicit math contracts over speed. They are used to pin PyTorch parity before adding SIMD/threaded variants.

Definition in file deepseek_kernels.c.

Function Documentation

deepseek_csa_attention_backward_f32()

void deepseek_csa_attention_backward_f32	(	const float *	d_out,
		const float *	q,
		const float *	k,
		const float *	v,
		const int *	indices,
		const float *	attn,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		int	query_tokens,
		int	key_tokens,
		int	heads,
		int	dim,
		int	top_k,
		float	scale
	)

Definition at line 621 of file deepseek_kernels.c.

{
    if (!d_out || !q || !k || !v || !indices || !attn || !d_q || !d_k || !d_v ||
        query_tokens <= 0 || key_tokens <= 0 || heads <= 0 || dim <= 0 || top_k <= 0) return;
 
    const size_t q_count = (size_t)query_tokens * (size_t)heads * (size_t)dim;
    const size_t kv_count = (size_t)key_tokens * (size_t)heads * (size_t)dim;
    for (size_t i = 0; i < q_count; ++i) d_q[i] = 0.0f;
    for (size_t i = 0; i < kv_count; ++i) {
        d_k[i] = 0.0f;
        d_v[i] = 0.0f;
    }
 
    for (int tq = 0; tq < query_tokens; ++tq) {
        for (int h = 0; h < heads; ++h) {
            float d_attn[top_k];
            float attn_dot = 0.0f;
            for (int j = 0; j < top_k; ++j) {
                const int tk = indices[((size_t)tq * (size_t)heads + (size_t)h) * (size_t)top_k + (size_t)j];
                float da = 0.0f;
                if (tk >= 0 && tk < key_tokens) {
                    const float a = attn[((size_t)tq * (size_t)heads + (size_t)h) * (size_t)top_k + (size_t)j];
                    for (int d = 0; d < dim; ++d) {
                        const float go = d_out[ds_qkv_idx(tq, h, d, heads, dim)];
                        da += go * v[ds_qkv_idx(tk, h, d, heads, dim)];
                        d_v[ds_qkv_idx(tk, h, d, heads, dim)] += a * go;
                    }
                    attn_dot += a * da;
                }
                d_attn[j] = da;
            }
 
            for (int j = 0; j < top_k; ++j) {
                const int tk = indices[((size_t)tq * (size_t)heads + (size_t)h) * (size_t)top_k + (size_t)j];
                if (tk < 0 || tk >= key_tokens) continue;
                const float a = attn[((size_t)tq * (size_t)heads + (size_t)h) * (size_t)top_k + (size_t)j];
                const float d_score = a * (d_attn[j] - attn_dot);
                for (int d = 0; d < dim; ++d) {
                    const float qv = q[ds_qkv_idx(tq, h, d, heads, dim)];
                    const float kv = k[ds_qkv_idx(tk, h, d, heads, dim)];
                    d_q[ds_qkv_idx(tq, h, d, heads, dim)] += scale * d_score * kv;
                    d_k[ds_qkv_idx(tk, h, d, heads, dim)] += scale * d_score * qv;
                }
            }
        }
    }
}

References ds_qkv_idx().

deepseek_csa_attention_f32()

void deepseek_csa_attention_f32	(	const float *	q,
		const float *	k,
		const float *	v,
		const int *	indices,
		float *	out,
		float *	attn,
		int	query_tokens,
		int	key_tokens,
		int	heads,
		int	dim,
		int	top_k,
		float	scale
	)

Definition at line 562 of file deepseek_kernels.c.

{
    if (!q || !k || !v || !indices || !out || query_tokens <= 0 || key_tokens <= 0 ||
        heads <= 0 || dim <= 0 || top_k <= 0) return;
 
    for (int tq = 0; tq < query_tokens; ++tq) {
        for (int h = 0; h < heads; ++h) {
            float local_scores[top_k];
            int valid = 0;
            for (int j = 0; j < top_k; ++j) {
                const int tk = indices[((size_t)tq * (size_t)heads + (size_t)h) * (size_t)top_k + (size_t)j];
                if (tk < 0 || tk >= key_tokens) {
                    local_scores[j] = -FLT_MAX;
                    continue;
                }
                float dot = 0.0f;
                for (int d = 0; d < dim; ++d) {
                    dot += q[ds_qkv_idx(tq, h, d, heads, dim)] * k[ds_qkv_idx(tk, h, d, heads, dim)];
                }
                local_scores[j] = dot * scale;
                valid++;
            }
 
            float *out_row = out + ds_qkv_idx(tq, h, 0, heads, dim);
            for (int d = 0; d < dim; ++d) out_row[d] = 0.0f;
            if (valid == 0) {
                if (attn) {
                    for (int j = 0; j < top_k; ++j) {
                        attn[((size_t)tq * (size_t)heads + (size_t)h) * (size_t)top_k + (size_t)j] = 0.0f;
                    }
                }
                continue;
            }
 
            ds_softmax(local_scores, top_k);
            for (int j = 0; j < top_k; ++j) {
                const int tk = indices[((size_t)tq * (size_t)heads + (size_t)h) * (size_t)top_k + (size_t)j];
                const float a = (tk >= 0 && tk < key_tokens) ? local_scores[j] : 0.0f;
                if (attn) attn[((size_t)tq * (size_t)heads + (size_t)h) * (size_t)top_k + (size_t)j] = a;
                if (a == 0.0f) continue;
                for (int d = 0; d < dim; ++d) {
                    out_row[d] += a * v[ds_qkv_idx(tk, h, d, heads, dim)];
                }
            }
        }
    }
}

References ds_qkv_idx(), and ds_softmax().

Referenced by deepseek_hybrid_attention_f32().

deepseek_dsa_topk_softmax_backward_f32()

void deepseek_dsa_topk_softmax_backward_f32	(	const int *	indices,
		const float *	weights,
		const float *	d_weights,
		float *	d_scores,
		int	tokens,
		int	heads,
		int	key_count,
		int	top_k
	)

Definition at line 153 of file deepseek_kernels.c.

{
    if (!indices || !weights || !d_weights || !d_scores ||
        tokens <= 0 || heads <= 0 || key_count <= 0 || top_k <= 0) return;
 
    topk_softmax_backward_f32(indices,
                              weights,
                              d_weights,
                              d_scores,
                              tokens * heads,
                              key_count,
                              top_k);
}

References topk_softmax_backward_f32().

deepseek_dsa_topk_softmax_f32()

void deepseek_dsa_topk_softmax_f32	(	const float *	scores,
		int *	indices,
		float *	weights,
		int	tokens,
		int	heads,
		int	key_count,
		int	top_k
	)

Definition at line 86 of file deepseek_kernels.c.

{
    if (!scores || !indices || !weights || tokens <= 0 || heads <= 0 ||
        key_count <= 0 || top_k <= 0) return;
 
    if (top_k > key_count) top_k = key_count;
 
    for (int t = 0; t < tokens; ++t) {
        for (int h = 0; h < heads; ++h) {
            const float *row = scores + ((size_t)t * (size_t)heads + (size_t)h) * (size_t)key_count;
            int *idx = indices + ((size_t)t * (size_t)heads + (size_t)h) * (size_t)top_k;
            float *w = weights + ((size_t)t * (size_t)heads + (size_t)h) * (size_t)top_k;
 
            for (int k = 0; k < top_k; ++k) {
                idx[k] = -1;
                w[k] = -FLT_MAX;
            }
 
            for (int j = 0; j < key_count; ++j) {
                const float v = row[j];
                int pos = -1;
                for (int k = 0; k < top_k; ++k) {
                    if (idx[k] < 0 || v > w[k] || (v == w[k] && j < idx[k])) {
                        pos = k;
                        break;
                    }
                }
                if (pos >= 0) {
                    for (int k = top_k - 1; k > pos; --k) {
                        idx[k] = idx[k - 1];
                        w[k] = w[k - 1];
                    }
                    idx[pos] = j;
                    w[pos] = v;
                }
            }
 
            float max_v = w[0];
            for (int k = 1; k < top_k; ++k) if (w[k] > max_v) max_v = w[k];
            float sum = 0.0f;
            for (int k = 0; k < top_k; ++k) {
                w[k] = expf(w[k] - max_v);
                sum += w[k];
            }
            if (sum > 0.0f) {
                const float inv = 1.0f / sum;
                for (int k = 0; k < top_k; ++k) w[k] *= inv;
            }
        }
    }
}

deepseek_hybrid_attention_f32()

void deepseek_hybrid_attention_f32	(	const float *	q,
		const float *	k,
		const float *	v,
		const int *	indices,
		float *	out,
		float *	attn,
		int	query_tokens,
		int	key_tokens,
		int	heads,
		int	dim,
		int	top_k,
		float	scale,
		int	mode
	)

Definition at line 683 of file deepseek_kernels.c.

{
    if (mode != 0) {
        deepseek_csa_attention_f32(q, k, v, indices, out, attn,
                                   query_tokens, key_tokens, heads, dim, top_k, scale);
        return;
    }
 
    int dense_indices[query_tokens * heads * key_tokens];
    for (int tq = 0; tq < query_tokens; ++tq) {
        for (int h = 0; h < heads; ++h) {
            for (int tk = 0; tk < key_tokens; ++tk) {
                dense_indices[((size_t)tq * (size_t)heads + (size_t)h) * (size_t)key_tokens + (size_t)tk] = tk;
            }
        }
    }
    deepseek_csa_attention_f32(q, k, v, dense_indices, out, attn,
                               query_tokens, key_tokens, heads, dim, key_tokens, scale);
}

References deepseek_csa_attention_f32().

deepseek_mhc_mix_backward_f32()

void deepseek_mhc_mix_backward_f32	(	const float *	d_out,
		const float *	streams,
		const float *	mix,
		float *	d_streams,
		float *	d_mix,
		int	tokens,
		int	n_streams,
		int	dim
	)

Definition at line 53 of file deepseek_kernels.c.

{
    if (!d_out || !streams || !mix || !d_streams || !d_mix ||
        tokens <= 0 || n_streams <= 0 || dim <= 0) return;
 
    const size_t stream_count = (size_t)tokens * (size_t)n_streams * (size_t)dim;
    const size_t mix_count = (size_t)tokens * (size_t)n_streams * (size_t)n_streams;
    for (size_t i = 0; i < stream_count; ++i) d_streams[i] = 0.0f;
    for (size_t i = 0; i < mix_count; ++i) d_mix[i] = 0.0f;
 
    for (int t = 0; t < tokens; ++t) {
        for (int os = 0; os < n_streams; ++os) {
            for (int is = 0; is < n_streams; ++is) {
                float d_mix_acc = 0.0f;
                const float m = mix[ds_mix_idx(t, os, is, n_streams)];
                for (int d = 0; d < dim; ++d) {
                    const float go = d_out[ds_mhc_idx(t, os, d, n_streams, dim)];
                    d_streams[ds_mhc_idx(t, is, d, n_streams, dim)] += m * go;
                    d_mix_acc += go * streams[ds_mhc_idx(t, is, d, n_streams, dim)];
                }
                d_mix[ds_mix_idx(t, os, is, n_streams)] = d_mix_acc;
            }
        }
    }
}

References ds_mhc_idx(), and ds_mix_idx().

deepseek_mhc_mix_f32()

void deepseek_mhc_mix_f32	(	const float *	streams,
		const float *	mix,
		float *	out,
		int	tokens,
		int	n_streams,
		int	dim
	)

Definition at line 30 of file deepseek_kernels.c.

{
    if (!streams || !mix || !out || tokens <= 0 || n_streams <= 0 || dim <= 0) return;
 
    for (int t = 0; t < tokens; ++t) {
        for (int os = 0; os < n_streams; ++os) {
            for (int d = 0; d < dim; ++d) {
                float acc = 0.0f;
                for (int is = 0; is < n_streams; ++is) {
                    acc += mix[ds_mix_idx(t, os, is, n_streams)] *
                           streams[ds_mhc_idx(t, is, d, n_streams, dim)];
                }
                out[ds_mhc_idx(t, os, d, n_streams, dim)] = acc;
            }
        }
    }
}

References ds_mhc_idx(), and ds_mix_idx().

deepseek_mla_attention_decode_f32()

void deepseek_mla_attention_decode_f32	(	const float *	q,
		const float *	k_cache,
		const float *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	cache_len,
		int	qk_head_dim,
		int	v_head_dim,
		int	max_seq_len,
		int	cache_stride
	)

Definition at line 1044 of file deepseek_kernels.c.

{
    if (cache_len <= 0 || (size_t)cache_len > SIZE_MAX / sizeof(float)) return;
    const size_t scores_bytes = (size_t)cache_len * sizeof(float);
    float *scores = (float *)malloc(scores_bytes);
    if (!scores) return;
    deepseek_mla_attention_decode_f32_workspace(
        q, k_cache, v_cache, output, num_heads, num_kv_heads, cache_len,
        qk_head_dim, v_head_dim, max_seq_len, cache_stride,
        1.0f / sqrtf((float)qk_head_dim), scores, scores_bytes);
    free(scores);
}

References deepseek_mla_attention_decode_f32_workspace().

deepseek_mla_attention_decode_f32_workspace()

void deepseek_mla_attention_decode_f32_workspace	(	const float *	q,
		const float *	k_cache,
		const float *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	cache_len,
		int	qk_head_dim,
		int	v_head_dim,
		int	max_seq_len,
		int	cache_stride,
		float	scale,
		float *	scores,
		size_t	scores_bytes
	)

Definition at line 978 of file deepseek_kernels.c.

{
    if (!q || !k_cache || !v_cache || !output || num_heads <= 0 ||
        num_kv_heads <= 0 || cache_len <= 0 || qk_head_dim <= 0 ||
        v_head_dim <= 0 || max_seq_len <= 0 || cache_stride <= 0 ||
        !isfinite(scale) || scale <= 0.0f) {
        return;
    }
    if (qk_head_dim > cache_stride || v_head_dim > cache_stride) {
        return;
    }
    if ((size_t)cache_len > SIZE_MAX / sizeof(float) || !scores ||
        scores_bytes < (size_t)cache_len * sizeof(float)) {
        return;
    }
 
    for (int h = 0; h < num_heads; ++h) {
        const int kv_h = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        const float *q_vec = q + (size_t)h * (size_t)qk_head_dim;
 
        float max_score = -FLT_MAX;
        for (int j = 0; j < cache_len; ++j) {
            const float *k_vec = k_cache + ((size_t)kv_h * (size_t)max_seq_len + (size_t)j) * (size_t)cache_stride;
            float dot = 0.0f;
            for (int d = 0; d < qk_head_dim; ++d) {
                dot += q_vec[d] * k_vec[d];
            }
            const float score = dot * scale;
            scores[j] = score;
            if (score > max_score) max_score = score;
        }
 
        float sum = 0.0f;
        for (int j = 0; j < cache_len; ++j) {
            const float e = expf(scores[j] - max_score);
            scores[j] = e;
            sum += e;
        }
 
        const float inv_sum = sum > 0.0f ? (1.0f / sum) : 0.0f;
        float *out = output + (size_t)h * (size_t)v_head_dim;
        for (int d = 0; d < v_head_dim; ++d) out[d] = 0.0f;
        for (int j = 0; j < cache_len; ++j) {
            const float w = scores[j] * inv_sum;
            const float *v_vec = v_cache + ((size_t)kv_h * (size_t)max_seq_len + (size_t)j) * (size_t)cache_stride;
            for (int d = 0; d < v_head_dim; ++d) {
                out[d] += w * v_vec[d];
            }
        }
    }
 
}

References score.

Referenced by deepseek_mla_attention_decode_f32().

deepseek_mla_attention_f32()

void deepseek_mla_attention_f32	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	qk_head_dim,
		int	v_head_dim
	)

Definition at line 879 of file deepseek_kernels.c.

{
    if (num_tokens <= 0 || (size_t)num_tokens > SIZE_MAX / sizeof(float)) return;
    const size_t scores_bytes = (size_t)num_tokens * sizeof(float);
    float *scores = (float *)malloc(scores_bytes);
    if (!scores) return;
    deepseek_mla_attention_f32_workspace(
        q, k, v, output, num_heads, num_kv_heads, num_tokens,
        qk_head_dim, v_head_dim, 1.0f / sqrtf((float)qk_head_dim),
        scores, scores_bytes);
    free(scores);
}

References deepseek_mla_attention_f32_workspace().

deepseek_mla_attention_f32_parallel_dispatch()

void deepseek_mla_attention_f32_parallel_dispatch	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	qk_head_dim,
		int	v_head_dim,
		float	scale,
		float *	scores,
		size_t	scores_bytes
	)

Definition at line 839 of file deepseek_kernels.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    const char *disabled = getenv("CK_DISABLE_MLA_PARALLEL_PREFILL");
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > num_heads) active = num_heads;
    const size_t score_row_bytes =
        num_tokens > 0 ? (size_t)num_tokens * sizeof(float) : 0;
    if ((disabled && disabled[0] && strcmp(disabled, "0") != 0) ||
        !pool || active <= 1 || num_tokens <= 0 || num_heads <= 0 ||
        num_tokens > INT_MAX / num_heads ||
        score_row_bytes == 0 || (size_t)active > SIZE_MAX / score_row_bytes ||
        scores_bytes < (size_t)active * score_row_bytes) {
        deepseek_mla_attention_f32_workspace(
            q, k, v, output, num_heads, num_kv_heads, num_tokens,
            qk_head_dim, v_head_dim, scale, scores, scores_bytes);
        return;
    }
    ds_mla_attention_f32_args_t args = {
        .q = q, .k = k, .v = v, .output = output,
        .num_heads = num_heads, .num_kv_heads = num_kv_heads,
        .num_tokens = num_tokens, .qk_head_dim = qk_head_dim,
        .v_head_dim = v_head_dim, .scale = scale, .scores = scores,
    };
    atomic_init(&args.next_token, 0);
    ck_threadpool_dispatch_n(pool, active, ds_mla_attention_f32_work, &args);
}

References ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), deepseek_mla_attention_f32_workspace(), and ds_mla_attention_f32_work().

deepseek_mla_attention_f32_workspace()

void deepseek_mla_attention_f32_workspace	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	qk_head_dim,
		int	v_head_dim,
		float	scale,
		float *	scores,
		size_t	scores_bytes
	)

Definition at line 768 of file deepseek_kernels.c.

{
    if (!q || !k || !v || !output || num_heads <= 0 || num_kv_heads <= 0 ||
        num_tokens <= 0 || qk_head_dim <= 0 || v_head_dim <= 0 ||
        num_tokens > INT_MAX / num_heads ||
        !isfinite(scale) || scale <= 0.0f) {
        return;
    }
    if ((size_t)num_tokens > SIZE_MAX / sizeof(float) || !scores ||
        scores_bytes < (size_t)num_tokens * sizeof(float)) {
        return;
    }
 
    ds_mla_attention_f32_query_range(
        q, k, v, output, num_heads, num_kv_heads, num_tokens,
        qk_head_dim, v_head_dim, scale, scores,
        0, num_tokens * num_heads, 1);
}

References ds_mla_attention_f32_query_range().

Referenced by deepseek_mla_attention_f32(), and deepseek_mla_attention_f32_parallel_dispatch().

deepseek_mla_kv_cache_batch_store_f32()

void deepseek_mla_kv_cache_batch_store_f32	(	float *	k_cache,
		float *	v_cache,
		const float *	k,
		const float *	v,
		int	num_tokens,
		int	num_kv_heads,
		int	qk_head_dim,
		int	v_head_dim,
		int	max_seq_len,
		int	cache_stride
	)

Definition at line 900 of file deepseek_kernels.c.

{
    if (!k_cache || !v_cache || !k || !v || num_tokens <= 0 ||
        num_kv_heads <= 0 || qk_head_dim <= 0 || v_head_dim <= 0 ||
        max_seq_len <= 0 || cache_stride <= 0) {
        return;
    }
    if (qk_head_dim > cache_stride || v_head_dim > cache_stride) {
        return;
    }
    if (num_tokens > max_seq_len) {
        num_tokens = max_seq_len;
    }
 
    for (int t = 0; t < num_tokens; ++t) {
        for (int h = 0; h < num_kv_heads; ++h) {
            const float *k_src = k + ((size_t)t * (size_t)num_kv_heads + (size_t)h) * (size_t)qk_head_dim;
            const float *v_src = v + ((size_t)t * (size_t)num_kv_heads + (size_t)h) * (size_t)v_head_dim;
            float *k_dst = k_cache + ((size_t)h * (size_t)max_seq_len + (size_t)t) * (size_t)cache_stride;
            float *v_dst = v_cache + ((size_t)h * (size_t)max_seq_len + (size_t)t) * (size_t)cache_stride;
            for (int d = 0; d < qk_head_dim; ++d) k_dst[d] = k_src[d];
            for (int d = qk_head_dim; d < cache_stride; ++d) k_dst[d] = 0.0f;
            for (int d = 0; d < v_head_dim; ++d) v_dst[d] = v_src[d];
            for (int d = v_head_dim; d < cache_stride; ++d) v_dst[d] = 0.0f;
        }
    }
}

deepseek_mla_kv_cache_store_f32()

void deepseek_mla_kv_cache_store_f32	(	float *	k_cache,
		float *	v_cache,
		const float *	k,
		const float *	v,
		int	pos,
		int	num_kv_heads,
		int	qk_head_dim,
		int	v_head_dim,
		int	max_seq_len,
		int	cache_stride
	)

Definition at line 937 of file deepseek_kernels.c.

{
    if (!k_cache || !v_cache || !k || !v || pos < 0 ||
        num_kv_heads <= 0 || qk_head_dim <= 0 || v_head_dim <= 0 ||
        max_seq_len <= 0 || cache_stride <= 0) {
        return;
    }
    if (pos >= max_seq_len || qk_head_dim > cache_stride || v_head_dim > cache_stride) {
        return;
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_src = k + ((size_t)h * (size_t)qk_head_dim);
        const float *v_src = v + ((size_t)h * (size_t)v_head_dim);
        float *k_dst = k_cache + ((size_t)h * (size_t)max_seq_len + (size_t)pos) * (size_t)cache_stride;
        float *v_dst = v_cache + ((size_t)h * (size_t)max_seq_len + (size_t)pos) * (size_t)cache_stride;
 
        for (int d = 0; d < qk_head_dim; ++d) {
            k_dst[d] = k_src[d];
        }
        for (int d = qk_head_dim; d < cache_stride; ++d) {
            k_dst[d] = 0.0f;
        }
        for (int d = 0; d < v_head_dim; ++d) {
            v_dst[d] = v_src[d];
        }
        for (int d = v_head_dim; d < cache_stride; ++d) {
            v_dst[d] = 0.0f;
        }
    }
}

deepseek_mla_kv_decompress_bf16()

void deepseek_mla_kv_decompress_bf16	(	const float *	compressed_kv,
		const uint16_t *	kv_b_proj,
		float *	k_nope,
		float *	value,
		int	tokens,
		int	heads,
		int	kv_lora_rank,
		int	qk_nope_dim,
		int	v_dim
	)

Definition at line 270 of file deepseek_kernels.c.

{
    if (tokens <= 0) return;
    deepseek_mla_kv_decompress_bf16_token_range(
        compressed_kv, kv_b_proj, k_nope, value, tokens, heads,
        kv_lora_rank, qk_nope_dim, v_dim, 0, tokens);
}

References deepseek_mla_kv_decompress_bf16_token_range().

Referenced by deepseek_mla_kv_decompress_bf16_parallel_dispatch().

deepseek_mla_kv_decompress_bf16_parallel_dispatch()

void deepseek_mla_kv_decompress_bf16_parallel_dispatch	(	const float *	compressed_kv,
		const uint16_t *	kv_b_proj,
		float *	k_nope,
		float *	value,
		int	tokens,
		int	heads,
		int	kv_lora_rank,
		int	qk_nope_dim,
		int	v_dim
	)

Definition at line 308 of file deepseek_kernels.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    const char *disabled = getenv("CK_DISABLE_MLA_PARALLEL_PREFILL");
    if ((disabled && disabled[0] && strcmp(disabled, "0") != 0) ||
        !pool || ck_threadpool_n_threads(pool) <= 1 || tokens < 2) {
        deepseek_mla_kv_decompress_bf16(
            compressed_kv, kv_b_proj, k_nope, value, tokens, heads,
            kv_lora_rank, qk_nope_dim, v_dim);
        return;
    }
    ds_mla_kv_decompress_bf16_args_t args = {
        .compressed_kv = compressed_kv,
        .kv_b_proj = kv_b_proj,
        .k_nope = k_nope,
        .value = value,
        .tokens = tokens,
        .heads = heads,
        .kv_lora_rank = kv_lora_rank,
        .qk_nope_dim = qk_nope_dim,
        .v_dim = v_dim,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > tokens) active = tokens;
    int grain = tokens / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, tokens, grain, ds_mla_kv_decompress_bf16_rows, &args);
}

References ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), deepseek_mla_kv_decompress_bf16(), and ds_mla_kv_decompress_bf16_rows().

deepseek_mla_kv_decompress_bf16_token_range()

void deepseek_mla_kv_decompress_bf16_token_range	(	const float *	compressed_kv,
		const uint16_t *	kv_b_proj,
		float *	k_nope,
		float *	value,
		int	tokens,
		int	heads,
		int	kv_lora_rank,
		int	qk_nope_dim,
		int	v_dim,
		int	token_begin,
		int	token_end
	)

Definition at line 225 of file deepseek_kernels.c.

{
    if (!compressed_kv || !kv_b_proj || !k_nope || !value ||
        tokens <= 0 || heads <= 0 || kv_lora_rank <= 0 || qk_nope_dim <= 0 || v_dim <= 0 ||
        token_begin < 0 || token_begin >= token_end || token_end > tokens) {
        return;
    }
 
    const int out_per_head = qk_nope_dim + v_dim;
    for (int t = token_begin; t < token_end; ++t) {
        for (int h = 0; h < heads; ++h) {
            for (int d = 0; d < qk_nope_dim; ++d) {
                const int out_col = h * out_per_head + d;
                float acc = 0.0f;
                for (int r = 0; r < kv_lora_rank; ++r) {
                    acc += bf16_to_float(kv_b_proj[(size_t)out_col * (size_t)kv_lora_rank + (size_t)r]) *
                           compressed_kv[ds_mla_tok_idx(t, r, kv_lora_rank)];
                }
                k_nope[ds_mla_thd_idx(t, h, d, heads, qk_nope_dim)] =
                    bf16_to_float(float_to_bf16(acc));
            }
            for (int d = 0; d < v_dim; ++d) {
                const int out_col = h * out_per_head + qk_nope_dim + d;
                float acc = 0.0f;
                for (int r = 0; r < kv_lora_rank; ++r) {
                    acc += bf16_to_float(kv_b_proj[(size_t)out_col * (size_t)kv_lora_rank + (size_t)r]) *
                           compressed_kv[ds_mla_tok_idx(t, r, kv_lora_rank)];
                }
                value[ds_mla_thd_idx(t, h, d, heads, v_dim)] =
                    bf16_to_float(float_to_bf16(acc));
            }
        }
    }
}

References bf16_to_float(), ds_mla_thd_idx(), ds_mla_tok_idx(), and float_to_bf16().

Referenced by deepseek_mla_kv_decompress_bf16(), and ds_mla_kv_decompress_bf16_rows().

deepseek_mla_kv_decompress_f32()

void deepseek_mla_kv_decompress_f32	(	const float *	compressed_kv,
		const float *	kv_b_proj,
		float *	k_nope,
		float *	value,
		int	tokens,
		int	heads,
		int	kv_lora_rank,
		int	qk_nope_dim,
		int	v_dim
	)

Definition at line 185 of file deepseek_kernels.c.

{
    if (!compressed_kv || !kv_b_proj || !k_nope || !value ||
        tokens <= 0 || heads <= 0 || kv_lora_rank <= 0 || qk_nope_dim <= 0 || v_dim <= 0) {
        return;
    }
 
    const int out_per_head = qk_nope_dim + v_dim;
    for (int t = 0; t < tokens; ++t) {
        for (int h = 0; h < heads; ++h) {
            for (int d = 0; d < qk_nope_dim; ++d) {
                const int out_col = h * out_per_head + d;
                float acc = 0.0f;
                for (int r = 0; r < kv_lora_rank; ++r) {
                    acc += kv_b_proj[(size_t)out_col * (size_t)kv_lora_rank + (size_t)r] *
                           compressed_kv[ds_mla_tok_idx(t, r, kv_lora_rank)];
                }
                k_nope[ds_mla_thd_idx(t, h, d, heads, qk_nope_dim)] = acc;
            }
            for (int d = 0; d < v_dim; ++d) {
                const int out_col = h * out_per_head + qk_nope_dim + d;
                float acc = 0.0f;
                for (int r = 0; r < kv_lora_rank; ++r) {
                    acc += kv_b_proj[(size_t)out_col * (size_t)kv_lora_rank + (size_t)r] *
                           compressed_kv[ds_mla_tok_idx(t, r, kv_lora_rank)];
                }
                value[ds_mla_thd_idx(t, h, d, heads, v_dim)] = acc;
            }
        }
    }
}

References ds_mla_thd_idx(), and ds_mla_tok_idx().

deepseek_mla_partial_rope_concat_f32()

void deepseek_mla_partial_rope_concat_f32	(	const float *	q_nope,
		const float *	q_pe,
		const float *	k_nope,
		const float *	k_pe,
		const float *	cos,
		const float *	sin,
		float *	query,
		float *	key,
		int	tokens,
		int	heads,
		int	qk_nope_dim,
		int	qk_rope_dim
	)

Definition at line 367 of file deepseek_kernels.c.

{
    if (!q_nope || !q_pe || !k_nope || !k_pe || !cos || !sin || !query || !key ||
        tokens <= 0 || heads <= 0 || qk_nope_dim <= 0 || qk_rope_dim <= 0 || (qk_rope_dim % 2) != 0) {
        return;
    }
 
    const int q_head_dim = qk_nope_dim + qk_rope_dim;
    for (int t = 0; t < tokens; ++t) {
        const float *cos_row = cos + (size_t)t * (size_t)(qk_rope_dim / 2);
        const float *sin_row = sin + (size_t)t * (size_t)(qk_rope_dim / 2);
        for (int h = 0; h < heads; ++h) {
            float *q_out = query + ds_mla_thd_idx(t, h, 0, heads, q_head_dim);
            float *k_out = key + ds_mla_thd_idx(t, h, 0, heads, q_head_dim);
            const float *qn = q_nope + ds_mla_thd_idx(t, h, 0, heads, qk_nope_dim);
            const float *qp = q_pe + ds_mla_thd_idx(t, h, 0, heads, qk_rope_dim);
            const float *kn = k_nope + ds_mla_thd_idx(t, h, 0, heads, qk_nope_dim);
            const float *kp = k_pe + ds_mla_tok_idx(t, 0, qk_rope_dim);
            for (int d = 0; d < qk_nope_dim; ++d) {
                q_out[d] = qn[d];
                k_out[d] = kn[d];
            }
            ds_mla_apply_kimi_rope(qp, q_out + qk_nope_dim, cos_row, sin_row, qk_rope_dim);
            ds_mla_apply_kimi_rope(kp, k_out + qk_nope_dim, cos_row, sin_row, qk_rope_dim);
        }
    }
}

References ds_mla_apply_kimi_rope(), ds_mla_thd_idx(), and ds_mla_tok_idx().

deepseek_mla_partial_rope_concat_packed_bf16_storage()

void deepseek_mla_partial_rope_concat_packed_bf16_storage	(	const float *	q_packed,
		const float *	k_nope,
		const float *	kv_a_packed,
		const float *	cos,
		const float *	sin,
		float *	query,
		float *	key,
		int	tokens,
		int	heads,
		int	kv_lora_rank,
		int	qk_nope_dim,
		int	qk_rope_dim
	)

Definition at line 473 of file deepseek_kernels.c.

{
    if (!q_packed || !k_nope || !kv_a_packed || !cos || !sin || !query || !key ||
        tokens <= 0 || heads <= 0 || kv_lora_rank <= 0 || qk_nope_dim <= 0 ||
        qk_rope_dim <= 0 || (qk_rope_dim % 2) != 0) {
        return;
    }
 
    const int q_head_dim = qk_nope_dim + qk_rope_dim;
    const int kv_a_dim = kv_lora_rank + qk_rope_dim;
    const int half = qk_rope_dim / 2;
    for (int t = 0; t < tokens; ++t) {
        const float *cos_row = cos + (size_t)t * (size_t)half;
        const float *sin_row = sin + (size_t)t * (size_t)half;
        const float *kp = kv_a_packed + (size_t)t * (size_t)kv_a_dim + (size_t)kv_lora_rank;
        for (int h = 0; h < heads; ++h) {
            const float *q_in = q_packed + ds_mla_thd_idx(t, h, 0, heads, q_head_dim);
            const float *kn = k_nope + ds_mla_thd_idx(t, h, 0, heads, qk_nope_dim);
            float *q_out = query + ds_mla_thd_idx(t, h, 0, heads, q_head_dim);
            float *k_out = key + ds_mla_thd_idx(t, h, 0, heads, q_head_dim);
 
            for (int d = 0; d < qk_nope_dim; ++d) {
                q_out[d] = ds_mla_bf16_round(q_in[d]);
                k_out[d] = ds_mla_bf16_round(kn[d]);
            }
 
            float q_pe_tmp[256];
            for (int i = 0; i < qk_rope_dim; ++i) {
                q_pe_tmp[i] = ds_mla_bf16_round(q_in[qk_nope_dim + i]);
            }
            for (int i = 0; i < half; ++i) {
                const float q_first = q_pe_tmp[2 * i];
                const float q_second = q_pe_tmp[2 * i + 1];
                const float k_first = ds_mla_bf16_round(kp[2 * i]);
                const float k_second = ds_mla_bf16_round(kp[2 * i + 1]);
                const float c = ds_mla_bf16_round(cos_row[i]);
                const float s = ds_mla_bf16_round(sin_row[i]);
 
                const float q_first_cos = ds_mla_bf16_round(q_first * c);
                const float q_second_sin = ds_mla_bf16_round(q_second * s);
                const float q_second_cos = ds_mla_bf16_round(q_second * c);
                const float q_first_sin = ds_mla_bf16_round(q_first * s);
                const float k_first_cos = ds_mla_bf16_round(k_first * c);
                const float k_second_sin = ds_mla_bf16_round(k_second * s);
                const float k_second_cos = ds_mla_bf16_round(k_second * c);
                const float k_first_sin = ds_mla_bf16_round(k_first * s);
 
                q_out[qk_nope_dim + i] = ds_mla_bf16_round(q_first_cos - q_second_sin);
                q_out[qk_nope_dim + half + i] = ds_mla_bf16_round(q_second_cos + q_first_sin);
                k_out[qk_nope_dim + i] = ds_mla_bf16_round(k_first_cos - k_second_sin);
                k_out[qk_nope_dim + half + i] = ds_mla_bf16_round(k_second_cos + k_first_sin);
            }
        }
    }
}

References ds_mla_bf16_round(), and ds_mla_thd_idx().

deepseek_mla_partial_rope_concat_packed_f32()

void deepseek_mla_partial_rope_concat_packed_f32	(	const float *	q_packed,
		const float *	k_nope,
		const float *	kv_a_packed,
		const float *	cos,
		const float *	sin,
		float *	query,
		float *	key,
		int	tokens,
		int	heads,
		int	kv_lora_rank,
		int	qk_nope_dim,
		int	qk_rope_dim
	)

Definition at line 406 of file deepseek_kernels.c.

{
    if (!q_packed || !k_nope || !kv_a_packed || !cos || !sin || !query || !key ||
        tokens <= 0 || heads <= 0 || kv_lora_rank <= 0 || qk_nope_dim <= 0 ||
        qk_rope_dim <= 0 || qk_rope_dim > 256 || (qk_rope_dim % 2) != 0) {
        return;
    }
 
    const int q_head_dim = qk_nope_dim + qk_rope_dim;
    const int kv_a_dim = kv_lora_rank + qk_rope_dim;
    const int half = qk_rope_dim / 2;
    if (qk_rope_dim > 256) {
        return;
    }
    for (int t = 0; t < tokens; ++t) {
        const float *cos_row = cos + (size_t)t * (size_t)half;
        const float *sin_row = sin + (size_t)t * (size_t)half;
        const float *kp = kv_a_packed + (size_t)t * (size_t)kv_a_dim + (size_t)kv_lora_rank;
        for (int h = 0; h < heads; ++h) {
            const float *q_in = q_packed + ds_mla_thd_idx(t, h, 0, heads, q_head_dim);
            const float *kn = k_nope + ds_mla_thd_idx(t, h, 0, heads, qk_nope_dim);
            float *q_out = query + ds_mla_thd_idx(t, h, 0, heads, q_head_dim);
            float *k_out = key + ds_mla_thd_idx(t, h, 0, heads, q_head_dim);
 
            for (int d = 0; d < qk_nope_dim; ++d) {
                q_out[d] = q_in[d];
                k_out[d] = kn[d];
            }
 
            const float *qp = q_in + qk_nope_dim;
            float q_pe_tmp[256];
            if (qk_rope_dim > (int)(sizeof(q_pe_tmp) / sizeof(q_pe_tmp[0]))) {
                return;
            }
            for (int i = 0; i < qk_rope_dim; ++i) q_pe_tmp[i] = qp[i];
            for (int i = 0; i < half; ++i) {
                const float q_first = q_pe_tmp[2 * i];
                const float q_second = q_pe_tmp[2 * i + 1];
                const float k_first = kp[2 * i];
                const float k_second = kp[2 * i + 1];
                const float c = cos_row[i];
                const float ss = sin_row[i];
                q_out[qk_nope_dim + i] = q_first * c - q_second * ss;
                q_out[qk_nope_dim + half + i] = q_second * c + q_first * ss;
                k_out[qk_nope_dim + i] = k_first * c - k_second * ss;
                k_out[qk_nope_dim + half + i] = k_second * c + k_first * ss;
            }
        }
    }
}

References ds_mla_thd_idx().

ds_mhc_idx()

static size_t ds_mhc_idx ( int  t, int  s, int  d, int  n_streams, int  dim  )

inlinestatic

Definition at line 20 of file deepseek_kernels.c.

{
    return ((size_t)t * (size_t)n_streams + (size_t)s) * (size_t)dim + (size_t)d;
}

Referenced by deepseek_mhc_mix_backward_f32(), and deepseek_mhc_mix_f32().

ds_mix_idx()

static size_t ds_mix_idx ( int  t, int  out_s, int  in_s, int  n_streams  )

inlinestatic

Definition at line 25 of file deepseek_kernels.c.

{
    return ((size_t)t * (size_t)n_streams + (size_t)out_s) * (size_t)n_streams + (size_t)in_s;
}

Referenced by deepseek_mhc_mix_backward_f32(), and deepseek_mhc_mix_f32().

ds_mla_apply_kimi_rope()

static void ds_mla_apply_kimi_rope ( const float *  src, float *  dst, const float *  cos_row, const float *  sin_row, int  dim  )

static

Definition at line 350 of file deepseek_kernels.c.

{
    const int half = dim / 2;
    for (int i = 0; i < half; ++i) {
        const float x_first = src[2 * i];
        const float x_second = src[2 * i + 1];
        const float c = cos_row[i];
        const float s = sin_row[i];
        dst[i] = x_first * c - x_second * s;
        dst[half + i] = x_second * c + x_first * s;
    }
}

Referenced by deepseek_mla_partial_rope_concat_f32().

ds_mla_attention_f32_query_range()

static void ds_mla_attention_f32_query_range ( const float *  q, const float *  k, const float *  v, float *  output, int  num_heads, int  num_kv_heads, int  num_tokens, int  qk_head_dim, int  v_head_dim, float  scale, float *  scores, int  query_begin, int  query_end, int  query_step  )

static

Definition at line 716 of file deepseek_kernels.c.

{
    for (int query = query_begin; query < query_end; query += query_step) {
        const int t = query / num_heads;
        const int h = query - t * num_heads;
        const int kv_h = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        const float *q_vec = q + ds_mla_thd_idx(t, h, 0, num_heads, qk_head_dim);
 
        float max_score = -FLT_MAX;
        for (int j = 0; j <= t; ++j) {
            const float *k_vec = k + ds_mla_thd_idx(j, kv_h, 0, num_kv_heads, qk_head_dim);
            float dot = 0.0f;
            for (int d = 0; d < qk_head_dim; ++d) {
                dot += q_vec[d] * k_vec[d];
            }
            const float score = dot * scale;
            scores[j] = score;
            if (score > max_score) max_score = score;
        }
 
        float sum = 0.0f;
        for (int j = 0; j <= t; ++j) {
            const float e = expf(scores[j] - max_score);
            scores[j] = e;
            sum += e;
        }
        const float inv_sum = sum > 0.0f ? (1.0f / sum) : 0.0f;
        float *out = output + ((size_t)t * (size_t)num_heads + (size_t)h) * (size_t)v_head_dim;
        for (int d = 0; d < v_head_dim; ++d) out[d] = 0.0f;
        for (int j = 0; j <= t; ++j) {
            const float weight = scores[j] * inv_sum;
            const float *v_vec = v + ds_mla_thd_idx(j, kv_h, 0, num_kv_heads, v_head_dim);
            for (int d = 0; d < v_head_dim; ++d) {
                out[d] += weight * v_vec[d];
            }
        }
    }
}

References ds_mla_thd_idx(), and score.

Referenced by deepseek_mla_attention_f32_workspace(), and ds_mla_attention_f32_work().

ds_mla_attention_f32_work()

static void ds_mla_attention_f32_work ( int  ith, int  nth, void *  opaque  )

static

Definition at line 813 of file deepseek_kernels.c.

{
    ds_mla_attention_f32_args_t *args =
        (ds_mla_attention_f32_args_t *)opaque;
    float *thread_scores = args->scores + (size_t)ith * (size_t)args->num_tokens;
    (void)nth;
    /* One token row contains every query head and costs O(t). Dynamic claims
     * keep asymmetric cores busy without sharing score scratch. Query rows are
     * independent, and every row retains its original reduction order. */
    for (;;) {
        const int token = atomic_fetch_add_explicit(
            &args->next_token,
            1,
            memory_order_relaxed);
        if (token >= args->num_tokens) break;
        ds_mla_attention_f32_query_range(
            args->q, args->k, args->v, args->output,
            args->num_heads, args->num_kv_heads, args->num_tokens,
            args->qk_head_dim, args->v_head_dim, args->scale,
            thread_scores,
            token * args->num_heads,
            (token + 1) * args->num_heads,
            1);
    }
}

References ds_mla_attention_f32_query_range(), and token.

Referenced by deepseek_mla_attention_f32_parallel_dispatch().

ds_mla_bf16_round()

static float ds_mla_bf16_round ( float  value )

inlinestatic

Definition at line 468 of file deepseek_kernels.c.

{
    return bf16_to_float(float_to_bf16(value));
}

References bf16_to_float(), and float_to_bf16().

Referenced by deepseek_mla_partial_rope_concat_packed_bf16_storage().

ds_mla_kv_decompress_bf16_rows()

static void ds_mla_kv_decompress_bf16_rows ( int  begin, int  end, void *  opaque  )

static

Definition at line 298 of file deepseek_kernels.c.

{
    const ds_mla_kv_decompress_bf16_args_t *args =
        (const ds_mla_kv_decompress_bf16_args_t *)opaque;
    deepseek_mla_kv_decompress_bf16_token_range(
        args->compressed_kv, args->kv_b_proj, args->k_nope, args->value,
        args->tokens, args->heads, args->kv_lora_rank, args->qk_nope_dim,
        args->v_dim, begin, end);
}

References deepseek_mla_kv_decompress_bf16_token_range(), and end.

Referenced by deepseek_mla_kv_decompress_bf16_parallel_dispatch().

ds_mla_thd_idx()

static size_t ds_mla_thd_idx ( int  t, int  h, int  d, int  heads, int  dim  )

inlinestatic

Definition at line 180 of file deepseek_kernels.c.

{
    return ((size_t)t * (size_t)heads + (size_t)h) * (size_t)dim + (size_t)d;
}

Referenced by deepseek_mla_kv_decompress_bf16_token_range(), deepseek_mla_kv_decompress_f32(), deepseek_mla_partial_rope_concat_f32(), deepseek_mla_partial_rope_concat_packed_bf16_storage(), deepseek_mla_partial_rope_concat_packed_f32(), and ds_mla_attention_f32_query_range().

ds_mla_tok_idx()

static size_t ds_mla_tok_idx ( int  t, int  d, int  dim  )

inlinestatic

Definition at line 175 of file deepseek_kernels.c.

{
    return (size_t)t * (size_t)dim + (size_t)d;
}

Referenced by deepseek_mla_kv_decompress_bf16_token_range(), deepseek_mla_kv_decompress_f32(), and deepseek_mla_partial_rope_concat_f32().

ds_qkv_idx()

static size_t ds_qkv_idx ( int  token, int  head, int  d, int  heads, int  dim  )

inlinestatic

Definition at line 541 of file deepseek_kernels.c.

{
    return ((size_t)token * (size_t)heads + (size_t)head) * (size_t)dim + (size_t)d;
}

References token.

Referenced by deepseek_csa_attention_backward_f32(), and deepseek_csa_attention_f32().

ds_softmax()

static void ds_softmax ( float *  x, int  n  )

static

Definition at line 546 of file deepseek_kernels.c.

{
    if (n <= 0) return;
    float max_v = x[0];
    for (int i = 1; i < n; ++i) if (x[i] > max_v) max_v = x[i];
    float sum = 0.0f;
    for (int i = 0; i < n; ++i) {
        x[i] = expf(x[i] - max_v);
        sum += x[i];
    }
    if (sum > 0.0f) {
        const float inv = 1.0f / sum;
        for (int i = 0; i < n; ++i) x[i] *= inv;
    }
}

Referenced by deepseek_csa_attention_f32().

topk_softmax_backward_f32()

void topk_softmax_backward_f32 ( const int *  indices, const float *  weights, const float *  d_weights, float *  d_scores, int  num_tokens, int  n_experts_or_keys, int  k  )

extern

Backward for hard top-k followed by softmax over selected values.

Matches PyTorch behavior for: values, indices = torch.topk(scores, k, dim=-1) weights = torch.softmax(values, dim=-1)

The hard selected indices are treated as fixed for this backward pass. Gradients are scattered only to selected scores; unselected scores are zero.

Definition at line 408 of file topk_kernels.c.

{
    if (!indices || !weights || !d_weights || !d_scores ||
        num_tokens <= 0 || n_experts_or_keys <= 0 || k <= 0) {
        return;
    }
 
    const size_t total = (size_t)num_tokens * (size_t)n_experts_or_keys;
    for (size_t i = 0; i < total; ++i) {
        d_scores[i] = 0.0f;
    }
 
    for (int t = 0; t < num_tokens; ++t) {
        const int *row_indices = indices + (size_t)t * (size_t)k;
        const float *row_weights = weights + (size_t)t * (size_t)k;
        const float *row_d_weights = d_weights + (size_t)t * (size_t)k;
        float *row_d_scores = d_scores + (size_t)t * (size_t)n_experts_or_keys;
 
        float dot = 0.0f;
        for (int i = 0; i < k; ++i) {
            const int idx = row_indices[i];
            if (idx >= 0 && idx < n_experts_or_keys) {
                dot += row_weights[i] * row_d_weights[i];
            }
        }
 
        for (int i = 0; i < k; ++i) {
            const int idx = row_indices[i];
            if (idx >= 0 && idx < n_experts_or_keys) {
                row_d_scores[idx] += row_weights[i] * (row_d_weights[i] - dot);
            }
        }
    }
}

Referenced by deepseek_dsa_topk_softmax_backward_f32().