attention_kernels_sliding.c File Reference

Sliding-window flash attention kernels split from attention_kernels.c. More...

#include "ckernel_engine.h"
#include "ck_threadpool.h"
#include <math.h>
#include <stdlib.h>

Go to the source code of this file.

Macros
#define	CK_SLIDING_FLASH_IMPL   attention_flash_query_sliding
#define	SLIDING_DECODE_IMPL   attention_flash_query_sliding
#define	SLIDING_DECODE_IMPL_GEMMA4   attention_flash_query_sliding
#define	SLIDING_FLASH_IMPL   attention_flash_query_sliding
#define	SLIDING_FLASH_IMPL_GEMMA4   attention_flash_query_sliding

Functions
static void	attention_flash_query_sliding (const float *q_vec, const float *k_head, const float *v_head, int query_pos, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float *out_vec, int sliding_window)
void	attention_forward_causal_head_major_gqa_flash_strided_sliding (const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)
void	attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4 (const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)
static void	attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl (const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window, int output_token_major)
void	attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_token_output (const float *q, const float *k, const float *v, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)
void	attention_forward_causal_head_major_shared_kv_sliding_gemma4 (const float *q, float *output, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)
void	attention_forward_decode_head_major_gqa_flash_sliding (const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window)
void	attention_forward_decode_head_major_gqa_flash_sliding_gemma4 (const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window)
void	attention_forward_decode_head_major_shared_kv_sliding_gemma4 (const float *q_token, const float *k_cache, const float *v_cache, float *out_token, int num_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window)
static size_t	attention_output_index (int h, int t, int num_heads, int num_tokens, int aligned_head_dim, int output_token_major)
static int	ck_env_int_default (const char *name, int fallback)
static void	ck_sliding_attention_compute_one (const ck_sliding_attention_args_t *a, int job)
static int	ck_sliding_attention_parallel_disabled (void)
static int	ck_sliding_attention_pick_threads (ck_threadpool_t *pool, int total_jobs, int num_tokens, int head_dim)
static void	ck_sliding_attention_work_fn (int ith, int nth, void *args)
static size_t	qkv_index (int h, int t, int d, int num_tokens, int aligned_head_dim)

Detailed Description

Sliding-window flash attention kernels split from attention_kernels.c.

Definition in file attention_kernels_sliding.c.

Macro Definition Documentation

CK_SLIDING_FLASH_IMPL

#define CK_SLIDING_FLASH_IMPL   attention_flash_query_sliding

SLIDING_DECODE_IMPL

#define SLIDING_DECODE_IMPL   attention_flash_query_sliding

SLIDING_DECODE_IMPL_GEMMA4

#define SLIDING_DECODE_IMPL_GEMMA4   attention_flash_query_sliding

SLIDING_FLASH_IMPL

#define SLIDING_FLASH_IMPL   attention_flash_query_sliding

SLIDING_FLASH_IMPL_GEMMA4

#define SLIDING_FLASH_IMPL_GEMMA4   attention_flash_query_sliding

Function Documentation

attention_flash_query_sliding()

static void attention_flash_query_sliding ( const float *  q_vec, const float *  k_head, const float *  v_head, int  query_pos, int  kv_tokens, int  head_dim, int  aligned_head_dim, float  scale, float *  out_vec, int  sliding_window  )

static

Definition at line 392 of file attention_kernels_sliding.c.

{
    float m = -INFINITY;
    float s = 0.0f;
 
    int window_start = 0;
    if (sliding_window > 0) {
        window_start = query_pos - sliding_window + 1;
        if (window_start < 0) window_start = 0;
    }
 
    for (int d = 0; d < head_dim; ++d) {
        out_vec[d] = 0.0f;
    }
 
    int effective_kv_end = query_pos < kv_tokens ? query_pos : kv_tokens - 1;
    for (int j = window_start; j <= effective_kv_end; ++j) {
        const float *k_vec = k_head + (size_t)j * (size_t)aligned_head_dim;
        const float *v_vec = v_head + (size_t)j * (size_t)aligned_head_dim;
 
        float dot = 0.0f;
        for (int d = 0; d < head_dim; ++d) {
            dot += q_vec[d] * k_vec[d];
        }
        float score = dot * scale;
 
        if (score > m) {
            float exp_m = (m == -INFINITY) ? 0.0f : expf(m - score);
            s *= exp_m;
            for (int d = 0; d < head_dim; ++d) {
                out_vec[d] *= exp_m;
            }
            s += 1.0f;
            for (int d = 0; d < head_dim; ++d) {
                out_vec[d] += v_vec[d];
            }
            m = score;
        } else {
            float e = expf(score - m);
            s += e;
            for (int d = 0; d < head_dim; ++d) {
                out_vec[d] += e * v_vec[d];
            }
        }
    }
 
    float inv_s = 1.0f / s;
    for (int d = 0; d < head_dim; ++d) {
        out_vec[d] *= inv_s;
    }
    for (int d = head_dim; d < aligned_head_dim; ++d) {
        out_vec[d] = 0.0f;
    }
}

References score.

attention_forward_causal_head_major_gqa_flash_strided_sliding()

void attention_forward_causal_head_major_gqa_flash_strided_sliding	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	kv_stride_tokens,
		int	sliding_window
	)

Flash attention forward with sliding window (prefill)

Test:

test_attention.py::TestAttentionForward::test_sliding_window_prefill

Sliding-window attention for prefill: each token attends to the last W tokens. When sliding_window <= 0, behaves like regular causal attention.

After changes: make test

Definition at line 566 of file attention_kernels_sliding.c.

{
    if (!q || !k || !v || !output) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || num_tokens <= 0) {
        return;
    }
    if (kv_stride_tokens < num_tokens) {
        return;
    }
 
    /* Debug escape hatch:
     * For Gemma bring-up, force the proven non-sliding flash kernel to isolate
     * whether divergence is caused by sliding-window implementation details.
     */
    if (getenv("CK_FORCE_NONSLIDING_ATTN")) {
        attention_forward_causal_head_major_gqa_flash_strided(
            q, k, v, output, num_heads, num_kv_heads, num_tokens,
            head_dim, aligned_head_dim, kv_stride_tokens
        );
        return;
    }
 
    const float scale = 1.0f / sqrtf((float)head_dim);
    const int T = num_tokens;
    const size_t kv_head_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
 
    const int total_jobs = num_heads * T;
    ck_threadpool_t *pool = ck_threadpool_global();
    const int active = ck_sliding_attention_pick_threads(pool, total_jobs, T, head_dim);
    if (pool && active > 1) {
        ck_sliding_attention_args_t args = {
            .q = q,
            .k = k,
            .v = v,
            .output = output,
            .num_heads = num_heads,
            .num_kv_heads = num_kv_heads,
            .num_tokens = T,
            .head_dim = head_dim,
            .aligned_head_dim = aligned_head_dim,
            .kv_stride_tokens = kv_stride_tokens,
            .sliding_window = sliding_window,
            .output_token_major = 0,
            .scale = scale,
        };
        ck_threadpool_dispatch_n(pool, active, ck_sliding_attention_work_fn, &args);
        return;
    }
 
#if defined(__AVX512F__)
    #define SLIDING_FLASH_IMPL attention_flash_query_sliding_avx512
#elif defined(__AVX2__)
    #define SLIDING_FLASH_IMPL attention_flash_query_sliding_avx2
#elif defined(__AVX__)
    #define SLIDING_FLASH_IMPL attention_flash_query_sliding_avx
#else
    #define SLIDING_FLASH_IMPL attention_flash_query_sliding
#endif
 
    for (int h = 0; h < num_heads; ++h) {
        int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        const float *k_head = k + (size_t)kv_head * kv_head_stride;
        const float *v_head = v + (size_t)kv_head * kv_head_stride;
 
        for (int i = 0; i < T; ++i) {
            const float *q_vec = q + qkv_index(h, i, 0, T, aligned_head_dim);
            float *out_vec = output + attention_output_index(
                h, i, num_heads, T, aligned_head_dim, 0);
            SLIDING_FLASH_IMPL(q_vec, k_head, v_head,
                               /*query_pos=*/i,
                               /*kv_tokens=*/T,
                               head_dim, aligned_head_dim,
                               scale, out_vec,
                               sliding_window);
        }
    }
 
#undef SLIDING_FLASH_IMPL
}

References attention_forward_causal_head_major_gqa_flash_strided(), attention_output_index(), ck_sliding_attention_pick_threads(), ck_sliding_attention_work_fn(), ck_threadpool_dispatch_n(), ck_threadpool_global(), qkv_index(), and SLIDING_FLASH_IMPL.

attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4()

void attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	kv_stride_tokens,
		int	sliding_window
	)

Definition at line 763 of file attention_kernels_sliding.c.

{
    attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl(
        q, k, v, output, num_heads, num_kv_heads, num_tokens, head_dim,
        aligned_head_dim, kv_stride_tokens, sliding_window,
        /*output_token_major=*/0);
}

References attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl().

Referenced by attention_forward_causal_head_major_shared_kv_sliding_gemma4().

attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl()

static void attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl ( const float *  q, const float *  k, const float *  v, float *  output, int  num_heads, int  num_kv_heads, int  num_tokens, int  head_dim, int  aligned_head_dim, int  kv_stride_tokens, int  sliding_window, int  output_token_major  )

static

Flash attention decode with sliding window

Test:

test_attention.py::TestAttentionForward::test_sliding_window_decode

Single query token attends to the last W tokens in the KV cache. For decode: effective_kv_tokens = min(kv_tokens, sliding_window)

After changes: make test

Definition at line 668 of file attention_kernels_sliding.c.

{
    if (!q || !k || !v || !output) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || num_tokens <= 0) {
        return;
    }
    if (kv_stride_tokens < num_tokens) {
        return;
    }
 
    if (getenv("CK_FORCE_NONSLIDING_ATTN")) {
        if (output_token_major) {
            attention_forward_causal_head_major_gqa_flash_strided_gemma4_token_output(
                q, k, v, output, num_heads, num_kv_heads, num_tokens,
                head_dim, aligned_head_dim, kv_stride_tokens);
        } else {
            attention_forward_causal_head_major_gqa_flash_strided_gemma4(
                q, k, v, output, num_heads, num_kv_heads, num_tokens,
                head_dim, aligned_head_dim, kv_stride_tokens);
        }
        return;
    }
 
    const float scale = 1.0f;
    const int T = num_tokens;
    const size_t kv_head_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
 
    const int total_jobs = num_heads * T;
    ck_threadpool_t *pool = ck_threadpool_global();
    const int active = ck_sliding_attention_pick_threads(pool, total_jobs, T, head_dim);
    if (pool && active > 1) {
        ck_sliding_attention_args_t args = {
            .q = q,
            .k = k,
            .v = v,
            .output = output,
            .num_heads = num_heads,
            .num_kv_heads = num_kv_heads,
            .num_tokens = T,
            .head_dim = head_dim,
            .aligned_head_dim = aligned_head_dim,
            .kv_stride_tokens = kv_stride_tokens,
            .sliding_window = sliding_window,
            .output_token_major = output_token_major,
            .scale = scale,
        };
        ck_threadpool_dispatch_n(pool, active, ck_sliding_attention_work_fn, &args);
        return;
    }
 
#if defined(__AVX512F__)
    #define SLIDING_FLASH_IMPL_GEMMA4 attention_flash_query_sliding_avx512
#elif defined(__AVX2__)
    #define SLIDING_FLASH_IMPL_GEMMA4 attention_flash_query_sliding_avx2
#elif defined(__AVX__)
    #define SLIDING_FLASH_IMPL_GEMMA4 attention_flash_query_sliding_avx
#else
    #define SLIDING_FLASH_IMPL_GEMMA4 attention_flash_query_sliding
#endif
 
    for (int h = 0; h < num_heads; ++h) {
        int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        const float *k_head = k + (size_t)kv_head * kv_head_stride;
        const float *v_head = v + (size_t)kv_head * kv_head_stride;
 
        for (int i = 0; i < T; ++i) {
            const float *q_vec = q + qkv_index(h, i, 0, T, aligned_head_dim);
            float *out_vec = output + attention_output_index(
                h, i, num_heads, T, aligned_head_dim, output_token_major);
            SLIDING_FLASH_IMPL_GEMMA4(q_vec, k_head, v_head,
                                      /*query_pos=*/i,
                                      /*kv_tokens=*/T,
                                      head_dim, aligned_head_dim,
                                      scale, out_vec,
                                      sliding_window);
        }
    }
 
#undef SLIDING_FLASH_IMPL_GEMMA4
}

References attention_forward_causal_head_major_gqa_flash_strided_gemma4(), attention_forward_causal_head_major_gqa_flash_strided_gemma4_token_output(), attention_output_index(), ck_sliding_attention_pick_threads(), ck_sliding_attention_work_fn(), ck_threadpool_dispatch_n(), ck_threadpool_global(), qkv_index(), and SLIDING_FLASH_IMPL_GEMMA4.

Referenced by attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4(), and attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_token_output().

attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_token_output()

void attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_token_output	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	kv_stride_tokens,
		int	sliding_window
	)

Definition at line 782 of file attention_kernels_sliding.c.

{
    attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl(
        q, k, v, output, num_heads, num_kv_heads, num_tokens, head_dim,
        aligned_head_dim, kv_stride_tokens, sliding_window,
        /*output_token_major=*/1);
}

References attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl().

attention_forward_causal_head_major_shared_kv_sliding_gemma4()

void attention_forward_causal_head_major_shared_kv_sliding_gemma4	(	const float *	q,
		float *	output,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	kv_stride_tokens,
		int	sliding_window
	)

Definition at line 801 of file attention_kernels_sliding.c.

{
    attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4(
        q, q, q, output, num_heads, num_heads, num_tokens,
        head_dim, aligned_head_dim, kv_stride_tokens, sliding_window
    );
}

References attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4().

attention_forward_decode_head_major_gqa_flash_sliding()

void attention_forward_decode_head_major_gqa_flash_sliding	(	const float *	q_token,
		const float *	k_cache,
		const float *	v_cache,
		float *	out_token,
		int	num_heads,
		int	num_kv_heads,
		int	kv_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		int	sliding_window
	)

Definition at line 817 of file attention_kernels_sliding.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || cache_capacity <= 0) {
        return;
    }
    if (kv_tokens <= 0 || kv_tokens > cache_capacity || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    /* Debug escape hatch:
     * Route decode through non-sliding flash attention when requested to
     * quickly A/B check sliding-window kernel correctness.
     */
    if (getenv("CK_FORCE_NONSLIDING_ATTN")) {
        attention_forward_decode_head_major_gqa_flash(
            q_token, k_cache, v_cache, out_token,
            num_heads, num_kv_heads, kv_tokens, cache_capacity,
            head_dim, aligned_head_dim
        );
        return;
    }
 
    const float scale = 1.0f / sqrtf((float)head_dim);
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
 
    // Compute effective KV tokens based on sliding window
    int effective_kv_tokens = kv_tokens;
    if (sliding_window > 0 && sliding_window < kv_tokens) {
        effective_kv_tokens = sliding_window;
    }
 
    // Guard against empty window (shouldn't happen with kv_tokens >= 1)
    if (effective_kv_tokens <= 0) {
        return;
    }
 
    // Offset to start reading from the last effective_kv_tokens entries
    int kv_start_offset = kv_tokens - effective_kv_tokens;
 
#if defined(__AVX512F__)
    #define SLIDING_DECODE_IMPL attention_flash_query_sliding_avx512
#elif defined(__AVX2__)
    #define SLIDING_DECODE_IMPL attention_flash_query_sliding_avx2
#elif defined(__AVX__)
    #define SLIDING_DECODE_IMPL attention_flash_query_sliding_avx
#else
    #define SLIDING_DECODE_IMPL attention_flash_query_sliding
#endif
 
    for (int h = 0; h < num_heads; ++h) {
        int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        const float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
        // Offset K/V pointer to start from the first token in the sliding window
        const float *k_head = k_cache + (size_t)kv_head * head_stride
                            + (size_t)kv_start_offset * (size_t)aligned_head_dim;
        const float *v_head = v_cache + (size_t)kv_head * head_stride
                            + (size_t)kv_start_offset * (size_t)aligned_head_dim;
        float *out_head = out_token + (size_t)h * (size_t)aligned_head_dim;
 
        // Use query_pos relative to the windowed KV (last token = effective_kv_tokens - 1)
        // sliding_window = 0 since we've already windowed via K/V pointer offset
        SLIDING_DECODE_IMPL(q_head, k_head, v_head,
                            /*query_pos=*/effective_kv_tokens - 1,
                            /*kv_tokens=*/effective_kv_tokens,
                            head_dim, aligned_head_dim,
                            scale, out_head,
                            /*sliding_window=*/0);
    }
 
#undef SLIDING_DECODE_IMPL
}

References attention_forward_decode_head_major_gqa_flash(), and SLIDING_DECODE_IMPL.

attention_forward_decode_head_major_gqa_flash_sliding_gemma4()

void attention_forward_decode_head_major_gqa_flash_sliding_gemma4	(	const float *	q_token,
		const float *	k_cache,
		const float *	v_cache,
		float *	out_token,
		int	num_heads,
		int	num_kv_heads,
		int	kv_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		int	sliding_window
	)

Definition at line 903 of file attention_kernels_sliding.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || cache_capacity <= 0) {
        return;
    }
    if (kv_tokens <= 0 || kv_tokens > cache_capacity || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    if (getenv("CK_FORCE_NONSLIDING_ATTN")) {
        attention_forward_decode_head_major_gqa_flash_gemma4(
            q_token, k_cache, v_cache, out_token,
            num_heads, num_kv_heads, kv_tokens, cache_capacity,
            head_dim, aligned_head_dim
        );
        return;
    }
 
    const float scale = 1.0f;
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
    int effective_kv_tokens = kv_tokens;
    if (sliding_window > 0 && sliding_window < kv_tokens) {
        effective_kv_tokens = sliding_window;
    }
    if (effective_kv_tokens <= 0) {
        return;
    }
    int kv_start_offset = kv_tokens - effective_kv_tokens;
 
#if defined(__AVX512F__)
    #define SLIDING_DECODE_IMPL_GEMMA4 attention_flash_query_sliding_avx512
#elif defined(__AVX2__)
    #define SLIDING_DECODE_IMPL_GEMMA4 attention_flash_query_sliding_avx2
#elif defined(__AVX__)
    #define SLIDING_DECODE_IMPL_GEMMA4 attention_flash_query_sliding_avx
#else
    #define SLIDING_DECODE_IMPL_GEMMA4 attention_flash_query_sliding
#endif
 
    for (int h = 0; h < num_heads; ++h) {
        int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        const float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
        const float *k_head = k_cache + (size_t)kv_head * head_stride
                            + (size_t)kv_start_offset * (size_t)aligned_head_dim;
        const float *v_head = v_cache + (size_t)kv_head * head_stride
                            + (size_t)kv_start_offset * (size_t)aligned_head_dim;
        float *out_head = out_token + (size_t)h * (size_t)aligned_head_dim;
 
        SLIDING_DECODE_IMPL_GEMMA4(q_head, k_head, v_head,
                                   /*query_pos=*/effective_kv_tokens - 1,
                                   /*kv_tokens=*/effective_kv_tokens,
                                   head_dim, aligned_head_dim,
                                   scale, out_head,
                                   /*sliding_window=*/0);
    }
 
#undef SLIDING_DECODE_IMPL_GEMMA4
}

References attention_forward_decode_head_major_gqa_flash_gemma4(), and SLIDING_DECODE_IMPL_GEMMA4.

Referenced by attention_forward_decode_head_major_shared_kv_sliding_gemma4().

attention_forward_decode_head_major_shared_kv_sliding_gemma4()

void attention_forward_decode_head_major_shared_kv_sliding_gemma4	(	const float *	q_token,
		const float *	k_cache,
		const float *	v_cache,
		float *	out_token,
		int	num_heads,
		int	kv_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		int	sliding_window
	)

Definition at line 976 of file attention_kernels_sliding.c.

{
    attention_forward_decode_head_major_gqa_flash_sliding_gemma4(
        q_token, k_cache, v_cache, out_token, num_heads, num_heads,
        kv_tokens, cache_capacity, head_dim, aligned_head_dim, sliding_window
    );
}

References attention_forward_decode_head_major_gqa_flash_sliding_gemma4().

attention_output_index()

static size_t attention_output_index ( int  h, int  t, int  num_heads, int  num_tokens, int  aligned_head_dim, int  output_token_major  )

inlinestatic

Definition at line 26 of file attention_kernels_sliding.c.

{
    if (output_token_major) {
        return ((size_t)t * (size_t)num_heads + (size_t)h)
             * (size_t)aligned_head_dim;
    }
    return qkv_index(h, t, 0, num_tokens, aligned_head_dim);
}

References qkv_index().

Referenced by attention_forward_causal_head_major_gqa_flash_strided_sliding(), attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl(), and ck_sliding_attention_compute_one().

ck_env_int_default()

static int ck_env_int_default ( const char *  name, int  fallback  )

static

Definition at line 472 of file attention_kernels_sliding.c.

{
    const char *v = getenv(name);
    if (!v || !v[0]) return fallback;
    char *end = NULL;
    long parsed = strtol(v, &end, 10);
    if (end == v || (end && *end != '\0')) return fallback;
    if (parsed < 0) parsed = 0;
    if (parsed > 1 << 20) parsed = 1 << 20;
    return (int)parsed;
}

References end.

Referenced by ck_sliding_attention_pick_threads().

ck_sliding_attention_compute_one()

static void ck_sliding_attention_compute_one ( const ck_sliding_attention_args_t *  a, int  job  )

static

Definition at line 509 of file attention_kernels_sliding.c.

{
    if (!a || job < 0) return;
 
    const int T = a->num_tokens;
    const size_t kv_head_stride = (size_t)a->kv_stride_tokens * (size_t)a->aligned_head_dim;
 
#if defined(__AVX512F__)
    #define CK_SLIDING_FLASH_IMPL attention_flash_query_sliding_avx512
#elif defined(__AVX2__)
    #define CK_SLIDING_FLASH_IMPL attention_flash_query_sliding_avx2
#elif defined(__AVX__)
    #define CK_SLIDING_FLASH_IMPL attention_flash_query_sliding_avx
#else
    #define CK_SLIDING_FLASH_IMPL attention_flash_query_sliding
#endif
 
    const int h = job / T;
    const int i = job - h * T;
    const int kv_head = (int)((long long)h * (long long)a->num_kv_heads /
                              (long long)a->num_heads);
    const float *k_head = a->k + (size_t)kv_head * kv_head_stride;
    const float *v_head = a->v + (size_t)kv_head * kv_head_stride;
    const float *q_vec = a->q + qkv_index(h, i, 0, T, a->aligned_head_dim);
    float *out_vec = a->output + attention_output_index(
        h, i, a->num_heads, T, a->aligned_head_dim, a->output_token_major);
 
    CK_SLIDING_FLASH_IMPL(q_vec, k_head, v_head,
                          /*query_pos=*/i,
                          /*kv_tokens=*/T,
                          a->head_dim, a->aligned_head_dim,
                          a->scale, out_vec,
                          a->sliding_window);
 
#undef CK_SLIDING_FLASH_IMPL
}

References attention_output_index(), CK_SLIDING_FLASH_IMPL, and qkv_index().

Referenced by ck_sliding_attention_work_fn().

ck_sliding_attention_parallel_disabled()

static int ck_sliding_attention_parallel_disabled ( void  )

static

Definition at line 484 of file attention_kernels_sliding.c.

{
    const char *v = getenv("CK_DISABLE_SLIDING_ATTN_PARALLEL");
    return v && v[0] && v[0] != '0';
}

Referenced by ck_sliding_attention_pick_threads().

ck_sliding_attention_pick_threads()

static int ck_sliding_attention_pick_threads ( ck_threadpool_t *  pool, int  total_jobs, int  num_tokens, int  head_dim  )

static

Definition at line 490 of file attention_kernels_sliding.c.

{
    if (!pool || total_jobs <= 0) return 1;
    if (ck_sliding_attention_parallel_disabled()) return 1;
    if (ck_threadpool_thread_id(pool) > 0) return 1;
 
    const int min_tokens = ck_env_int_default("CK_SLIDING_ATTN_PARALLEL_MIN_TOKENS", 128);
    if (num_tokens < min_tokens || head_dim < 8) return 1;
 
    int active = ck_threadpool_n_threads(pool);
    const int cap = ck_env_int_default("CK_SLIDING_ATTN_THREAD_CAP", active);
    if (cap > 0 && active > cap) active = cap;
    if (active > total_jobs) active = total_jobs;
    return active > 1 ? active : 1;
}

References ck_env_int_default(), ck_sliding_attention_parallel_disabled(), ck_threadpool_n_threads(), and ck_threadpool_thread_id().

Referenced by attention_forward_causal_head_major_gqa_flash_strided_sliding(), and attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl().

ck_sliding_attention_work_fn()

static void ck_sliding_attention_work_fn ( int  ith, int  nth, void *  args  )

static

Definition at line 547 of file attention_kernels_sliding.c.

{
    const ck_sliding_attention_args_t *a = (const ck_sliding_attention_args_t *)args;
    if (!a || ith < 0 || nth <= 0 || ith >= nth) return;
    const int total_jobs = a->num_heads * a->num_tokens;
    for (int job = ith; job < total_jobs; job += nth) {
        ck_sliding_attention_compute_one(a, job);
    }
}

References ck_sliding_attention_compute_one().

Referenced by attention_forward_causal_head_major_gqa_flash_strided_sliding(), and attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl().

qkv_index()

static size_t qkv_index ( int  h, int  t, int  d, int  num_tokens, int  aligned_head_dim  )

inlinestatic

Definition at line 16 of file attention_kernels_sliding.c.

{
    return ((size_t)h * (size_t)num_tokens + (size_t)t) * (size_t)aligned_head_dim
         + (size_t)d;
}

Referenced by attention_forward_causal_head_major_gqa_flash_strided_sliding(), attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl(), attention_output_index(), and ck_sliding_attention_compute_one().