Attention score/softmax/output kernels with SIMD (SSE/AVX/AVX512) More...

#include "bf16_utils.h"
#include "attention_oracle_ggml.h"
#include "ckernel_engine.h"
#include "ck_threadpool.h"
#include <dlfcn.h>
#include <limits.h>
#include <math.h>
#include <float.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include "ck_speed_profiles.h"
#include <string.h>

Macros
#define	CK_ENABLE_LLAMA_CPP_PARITY 0

#define	CK_GGML_FA_TILE_KV 64

#define	CK_GGML_FA_TILE_Q 64

#define	CK_GGML_FA_TILE_Q_LARGE 336

#define	CK_GGML_FA_TILE_Q_LARGE_MIN_TOKENS 1536

#define	CK_MUSE_PREFILL_QUERY_TILE 1024

#define	CK_NOINLINE

#define	CK_OPTNONE

#define	FLASH_QUERY_IMPL attention_flash_query_causal

#define	FLASH_QUERY_IMPL attention_flash_query_causal

#define	FLASH_QUERY_IMPL_DECODE attention_flash_query_causal

#define	RTLD_DEFAULT ((void *)0)

Typedefs
typedef float(*	ck_attention_math_f32_fn) (float)

Functions
void	attention_backward_causal_head_major (const float d_output, const float q, const float k, const float v, const float attn_weights, float d_q, float d_k, float d_v, float *d_scores, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)

void	attention_backward_causal_head_major_gqa (const float d_output, const float q, const float k, const float v, const float attn_weights, float d_q, float d_k, float d_v, float *d_scores, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)

void	attention_backward_causal_head_major_gqa_bf16 (const uint16_t d_output, float d_x, const uint16_t q, const uint16_t k, const uint16_t v, const float attn_weights, float d_q, float d_k, float d_v, float d_scores, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window, float scratch_d_output, float scratch_q, float scratch_k, float scratch_v)

static void	attention_flash_query_causal (const float q_vec, const float k_head, const float v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float out_vec)

static void	attention_flash_query_causal_exact (const float q_vec, const float k_head, const float v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float out_vec)

static void	attention_flash_query_causal_exact_f16kv (const float q_vec, const float k_head, const float v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float out_vec)

static void	attention_flash_query_causal_exact_prerounded_f16kv (const float q_vec, const float k_head, const float v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float out_vec)

void	attention_forward_causal_head_major (const float q, const float k, const float v, float scores, float *output, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)

void	attention_forward_causal_head_major_exact (const float q, const float k, const float v, float scores, float *output, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)

void	attention_forward_causal_head_major_gqa (const float q, const float k, const float v, float scores, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)

void	attention_forward_causal_head_major_gqa_bf16 (const uint16_t q, const uint16_t k, const uint16_t v, float scores, float output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window, float scratch_q, float scratch_k, float scratch_v)

void	attention_forward_causal_head_major_gqa_exact (const float q, const float k, const float v, float scores, float *output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int aligned_context_window)

void	attention_forward_causal_head_major_gqa_flash (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)

void	attention_forward_causal_head_major_gqa_flash_strided (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)

void	attention_forward_causal_head_major_gqa_flash_strided_f16kv (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)

void	attention_forward_causal_head_major_gqa_flash_strided_f16kv_serial (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)

void	attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace (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, float *rounded_kv, size_t rounded_kv_bytes)

void	attention_forward_causal_head_major_gqa_flash_strided_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)

void	attention_forward_causal_head_major_gqa_flash_strided_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)

void	attention_forward_causal_head_major_gqa_flash_strided_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)

void	attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace (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, float scores, size_t scores_bytes, float value_columns, size_t value_columns_bytes, float *scaled_scores, size_t scaled_scores_bytes)

void	attention_forward_causal_head_major_gqa_muse_eager_bf16_storage (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, float scores, size_t scores_bytes, uint16_t key_rows, size_t key_rows_bytes, uint16_t value_columns, size_t value_columns_bytes, uint16_t gemm_rows, size_t gemm_rows_bytes, uint16_t *gemm_scores, size_t gemm_scores_bytes)

void	attention_forward_causal_head_major_gqa_muse_eager_bf16_storage_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, float scores, size_t scores_bytes, uint16_t key_rows, size_t key_rows_bytes, uint16_t value_columns, size_t value_columns_bytes, uint16_t gemm_rows, size_t gemm_rows_bytes, uint16_t *gemm_scores, size_t gemm_scores_bytes)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_append_f16cache_auto_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float token_workspace, size_t token_workspace_bytes, void gqa_workspace, size_t gqa_workspace_bytes, int route_num_heads, int route_num_kv_heads, int route_head_dim, int route_query_tokens, int route_min_kv_tokens, int route_workers, int route_query_tile_size, int route_concurrent_query_tiles)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_append_f16cache_gemma4_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int query_tile_size, int concurrent_query_tiles, void *workspace, size_t workspace_bytes)

size_t	attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes (int num_heads, int num_kv_heads, int head_dim, int workers, int query_tile_size, int concurrent_query_tiles)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_append_f16cache_qtile64_schedule (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_prefill_schedule_t schedule)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_append_f16cache_sliding_gemma4_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window, ck_attention_reduction_t reduction, float *token_workspace, size_t token_workspace_bytes)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_full_bf16cache_pytorch_contract (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_contract_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float token_workspace, size_t token_workspace_bytes, const int segment_lengths, int num_segments)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_gemma4_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float token_workspace, size_t token_workspace_bytes, const int segment_lengths, int num_segments)

ck_attention_status_t	attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_sliding_gemma4_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window, ck_attention_reduction_t reduction, float token_workspace, size_t token_workspace_bytes, const int segment_lengths, int num_segments)

void	attention_forward_causal_head_major_shared_kv_gemma4 (const float q, float output, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens)

void	attention_forward_chunk_head_major_gqa_flash_gemma4 (const float q_chunk, const float k_cache, const float v_cache, float out_chunk, int num_heads, int num_kv_heads, int q_tokens, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim)

ck_attention_status_t	attention_forward_decode_head_major_gqa_bf16cache_pytorch_contract (const float q_token, const uint16_t k_cache, const uint16_t 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, ck_attention_reduction_t reduction)

ck_attention_status_t	attention_forward_decode_head_major_gqa_f16cache_gemma4_contract (const float q_token, const uint16_t k_cache, const uint16_t 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, ck_attention_reduction_t reduction)

ck_attention_status_t	attention_forward_decode_head_major_gqa_f16cache_regular_noscale (const float q_token, const uint16_t k_cache, const uint16_t 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)

ck_attention_status_t	attention_forward_decode_head_major_gqa_f16cache_sliding_gemma4_contract (const float q_token, const uint16_t k_cache, const uint16_t 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, ck_attention_reduction_t reduction)

void	attention_forward_decode_head_major_gqa_flash (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)

void	attention_forward_decode_head_major_gqa_flash_f16cache (const float q_token, const uint16_t k_cache, const uint16_t 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)

ck_attention_status_t	attention_forward_decode_head_major_gqa_flash_f16cache_contract (const float q_token, const uint16_t k_cache, const uint16_t 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, ck_attention_reduction_t reduction)

void	attention_forward_decode_head_major_gqa_flash_f16cache_split (const float q_token, const uint16_t k_cache, const uint16_t 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 split_chunks)

static void	attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned (const float q_token, const uint16_t k_cache, const uint16_t 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 split_chunks, int partition_tokens, int kv_start, float attention_scale)

void	attention_forward_decode_head_major_gqa_flash_f16kv (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)

void	attention_forward_decode_head_major_gqa_flash_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)

void	attention_forward_decode_head_major_gqa_llama_regular_sliding_workspace (const float q, const float k, const float v, float output, int num_heads, int num_kv_heads, int live_tokens, int kv_stride_tokens, int head_dim, int aligned_head_dim, int sliding_window, float scores, size_t scores_bytes, float value_columns, size_t value_columns_bytes, float *scaled_scores, size_t scaled_scores_bytes)

void	attention_forward_decode_head_major_gqa_muse_eager_bf16_storage (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, float scores, size_t scores_bytes, uint16_t key_rows, size_t key_rows_bytes, uint16_t value_columns, size_t value_columns_bytes, uint16_t gemm_rows, size_t gemm_rows_bytes, uint16_t *gemm_scores, size_t gemm_scores_bytes)

void	attention_forward_decode_head_major_gqa_muse_eager_bf16_storage_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, float scores, size_t scores_bytes, uint16_t key_rows, size_t key_rows_bytes, uint16_t value_columns, size_t value_columns_bytes, uint16_t gemm_rows, size_t gemm_rows_bytes, uint16_t *gemm_scores, size_t gemm_scores_bytes)

void	attention_forward_decode_head_major_gqa_regular (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)
	WARNING: This is NOT true flash attention!

void	attention_forward_decode_head_major_shared_kv_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)

void	attention_forward_full_head_major_gqa_exact_strided (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)

void	attention_forward_full_head_major_gqa_flash (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)

void	attention_forward_full_head_major_gqa_flash_strided (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)

void	attention_forward_full_head_major_gqa_flash_strided_bf16_storage (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)

void	attention_forward_full_head_major_gqa_flash_strided_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)

void	attention_forward_full_head_major_gqa_ggml_strided (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)

void	attention_forward_full_head_major_gqa_ggml_strided_workspace (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, float score_rows, size_t score_rows_bytes, float v_columns, size_t v_columns_bytes, float *probability_row, size_t probability_row_bytes)

void	attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage (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)

void	attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage_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)

void	attention_forward_full_head_major_gqa_sdpa_bf16_storage (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)

void	attention_forward_full_head_major_gqa_tiled336_f16kv_fp32_strided (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)

void	attention_forward_full_head_major_gqa_tiled64_f16kv_fp32_strided (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)

void	attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided (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)

static void	attention_forward_head_major_gqa_flash_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 causal, int round_full_kv_fp16, int output_token_major, float scale)

static int	attention_forward_head_major_gqa_unfused_f16_strict (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 causal, int output_token_major, float scale, int debug_layer_id)

void	attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_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 visual_start, int visual_tokens)

static void	attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_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 visual_start, int visual_tokens, int output_token_major)

void	attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_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 visual_start, int visual_tokens)

int	attention_forward_query_key_head_major_f32 (const float query, const float key, const float value, float output, float *score_scratch, int num_heads, int query_tokens, int key_tokens, int head_dim, float scale)

int	attention_forward_query_key_head_major_f32_decode_heads (const float query, const float key, const float value, float output, float *score_scratch, int num_heads, int query_tokens, int key_tokens, int head_dim, float scale)

int	attention_forward_query_key_head_major_f32_packed_k (const float query, const float key, const float value, float output, float score_scratch, float key_transpose_scratch, int num_heads, int query_tokens, int key_tokens, int head_dim, float scale)

int	attention_forward_query_key_head_major_tiled_f16kv_fp32 (const float query, const float key, const float value, float output, int num_heads, int query_tokens, int key_tokens, int head_dim, float scale)

void	attention_forward_sparse_token_major_gqa_bf16cache_pytorch_cpu_flash_contract (const float query, const uint16_t key_cache, const uint16_t value_cache, const float selected_indices, float output, float score_scratch, int rows, int query_heads, int kv_heads, int head_dim, int selection_width, int context_length, int position)

static size_t	attention_output_index (int h, int token, int num_heads, int num_tokens, int aligned_head_dim, int token_major)

static void	attention_query_full_exact_regular (const float q_vec, const float k_head, const float v_cols, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float score_row, float *out_vec, int layer_id, int head_id, int query_id)

static void	attention_query_full_ggml_regular (const float q_vec, const float k_head, const float v_cols, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float score_row, float *out_vec, int layer_id, int head_id, int query_id)

static size_t	ck_attention_align64_size (size_t value)

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

int	ck_attention_bf16_pytorch_gqa_available (void)

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

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

static float	ck_attention_dot_f16_llama (const uint16_t x, const uint16_t y, int n)

static float	ck_attention_dot_f16_unfused_llama (const uint16_t x, const uint16_t y, int n)

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

static size_t	ck_attention_f16_prefill_gqa_reuse_worker_bytes (int num_heads, int num_kv_heads, int head_dim, int workers, int query_tile_size, int concurrent_query_tiles)

static ck_attention_status_t	ck_attention_f16_prefill_qtile64_dispatch (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int cache_is_bf16, float attention_scale, size_t q_head_stride, size_t output_head_stride, ck_attention_prefill_schedule_t schedule)

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

static float	ck_attention_f16_reduce_expf (float value)

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

static void	ck_attention_flash_query_auto (const float q_vec, const float k_head, const float v_head, int kv_tokens, int head_dim, int aligned_head_dim, float scale, float out_vec)

static ck_attention_status_t	ck_attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_schedule_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, ck_attention_reduction_t reduction, float token_workspace, size_t token_workspace_bytes, const int segment_lengths, int num_segments, ck_attention_prefill_schedule_t schedule)

static ck_attention_status_t	ck_attention_forward_decode_f16cache_gemma4 (const float q_token, const uint16_t k_cache, const uint16_t 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, ck_attention_reduction_t reduction)

static void	ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided (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 query_tile_size)

static int	ck_attention_forward_query_key_head_major_f32_run (const float query, const float key, const float value, float output, float score_scratch, float key_transpose_scratch, int num_heads, int query_tokens, int key_tokens, int head_dim, float scale)

static int	ck_attention_full_bf16_pytorch_flash (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 output_token_major)

static int	ck_attention_full_bf16_sdpa_amx_range (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 head_begin, int head_step, int output_token_major)

static int	ck_attention_full_bf16_sdpa_tiled (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)

static int	ck_attention_full_bf16_sdpa_tiled_range (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 head_begin, int head_step)

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

static void	ck_attention_full_tiled_f16kv_fp32_range (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 query_tile_size, int ith, int nth)

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

static int	ck_attention_ggml_out_graph_enabled (void)

static void	ck_attention_gqa_team_barrier_wait (ck_attention_gqa_team_barrier_t *barrier)

static void	ck_attention_llama_regular_impl (const float q, const float k, const float v, float output, int num_heads, int num_kv_heads, int query_tokens, int live_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window, float scores, size_t scores_bytes, float value_columns, size_t value_columns_bytes, float *scaled_scores, size_t scaled_scores_bytes, int batched_prefill)

static void	ck_attention_llama_regular_query (const float query, const float key_head, const float value_columns, float output, float scores, float scaled_scores, int live_tokens, int padded_tokens, int query_position, int head_dim, int aligned_head_dim, int sliding_window, int batched_prefill)

static void	ck_attention_mad_f16_llama (uint16_t y, const uint16_t x, float scale, int n)

static void	ck_attention_matmul_f32_accum (float c, const float a, const float *b, int m, int k, int n)

static void	ck_attention_muse_eager_bf16_storage_impl (const float q, const float k, const float v, float output, int num_heads, int num_kv_heads, int query_tokens, int kv_tokens, int kv_stride_tokens, int head_dim, int aligned_head_dim, int sliding_window, float scores, size_t scores_bytes, uint16_t key_rows, size_t key_rows_bytes, uint16_t value_columns, size_t value_columns_bytes, uint16_t gemm_rows, size_t gemm_rows_bytes, uint16_t *gemm_scores, size_t gemm_scores_bytes)

static int	ck_attention_parallel_enabled (int total_queries, int num_tokens, int head_dim)

static int	ck_attention_pick_active_threads (const ck_threadpool_t *pool, int total_queries, int num_tokens)

static ck_attention_status_t	ck_attention_prefill_append_f16cache_gemma4_workspace (const float q, const uint16_t k_cache, const uint16_t v_cache, float output, int num_heads, int num_kv_heads, int q_tokens, int past_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window, ck_attention_reduction_t reduction, float token_workspace, size_t token_workspace_bytes, const int segment_lengths, int num_segments)

float	ck_attention_pytorch_sdpa_scale_f32 (int head_dim)

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

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

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

static float	ck_attention_reference_expf (float value)

static int	ck_attention_reverse_out_dot_enabled (void)

static void	ck_attention_scale_f16_llama (uint16_t *y, float scale, int n)

int	ck_attention_sparse_bf16_pytorch_gqa_available (void)

static float	ck_attention_strict_scale_f32 (int head_dim)

static int	ck_attention_strict_unfused_f16_enabled (void)

static void	ck_attention_trace (const char *branch, int layer_id, int head_id)

static void	ck_attention_trace_float (const char *tag, int layer_id, int head_id, float value)

static void	ck_attention_trace_query (const char *tag, int layer_id, int head_id, int query_id, int value)

static float	ck_attention_u16_cache_to_f32 (uint16_t value, int cache_is_bf16)

static int	ck_attention_vec_dump_enabled (void)

static void	ck_attention_vec_dump_exact_query (const float q_vec, const float k_head, const float *out_vec, int kv_tokens, int head_dim, int aligned_head_dim, float scale, int layer_id, int head_id, int query_id)

static int	ck_attention_vec_dump_next_layer_id (void)

static int	ck_attention_vec_dump_parse_env_int (const char name, int out)

static void	ck_attention_vec_dump_selected_query (const float raw_scores, const float probs, const float out_vec, const float v_cols, int kv_tokens, int head_dim, int layer_id, int head_id, int query_id)

static int	ck_attention_vec_dump_should_emit (int layer_id, int head_id, int query_id)

static void	ck_attention_vec_dump_tensor (const char name, int layer_id, int query_id, const float data, size_t elem_count)

static int	ck_attention_vec_dump_vcols_enabled (void)

static float	ck_bf16_dot_contract (const float a, const float b, int count)

static float	ck_ggml_vec_dot_f32_contig (const float x, const float y, int n)

static double	ck_ggml_vec_soft_max_row (int n, float y, const float x, float max)

static int	ck_llama_kv_pad_256 (int live_tokens, int capacity)

static float	ck_llama_regular_dot_f16 (const float a, const float b, int count)

static float	ck_llama_regular_gemm_f16 (const float probability, const float value_column, int count)

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

static void	ck_local_fp16_to_fp32_2d (const uint16_t src, float dst, int rows, int cols, int src_stride, int dst_stride)

static void	ck_local_fp16_to_fp32_row (const uint16_t src, float dst, int n)

static size_t	ck_muse_attention_checked_product (size_t left, size_t right, const char *label)

static void	ck_round_fp16_buffer (const float src, float dst, size_t count)

static float	ck_round_fp16_scalar (float x)

static float	ck_vec_dot_f32_reverse_strict (const float x, const float y, int n)

static float	ck_vec_dot_f32_strict (const float x, const float y, int n)

static float	ck_vec_dot_f32x_f32_to_f32_via_f64 (const float x, const float y, int n)

static float	ck_vec_max_f32_contig (const float *x, int n)

static void	ck_vec_scale_f32_inplace (float *x, int n, float scale)

static void	convert_bf16_tensor_to_buf (const uint16_t src, float dst, size_t count)

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

static size_t	score_index (int h, int i, int j, int aligned_context_window)

Variables
static int	ck_attention_vec_dump_layer_seq = 0

static const char	ck_attention_vec_dump_magic [8] = {'C', 'K', 'D', 'M', 'P', '\0', '\0', '\0'}

static const uint32_t	ck_attention_vec_dump_version = 1u

Detailed Description

Attention score/softmax/output kernels with SIMD (SSE/AVX/AVX512)

CK-ENGINE KERNEL RULES:

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

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

Attention: softmax(Q @ K^T / sqrt(d)) @ V Supports GQA (grouped-query attention) with head broadcasting.

Definition in file attention_kernels.c.

Macro Definition Documentation

◆ CK_ENABLE_LLAMA_CPP_PARITY

#define CK_ENABLE_LLAMA_CPP_PARITY 0

Definition at line 19 of file attention_kernels.c.

◆ CK_GGML_FA_TILE_KV

#define CK_GGML_FA_TILE_KV 64

Definition at line 2890 of file attention_kernels.c.

◆ CK_GGML_FA_TILE_Q

#define CK_GGML_FA_TILE_Q 64

Definition at line 2888 of file attention_kernels.c.

◆ CK_GGML_FA_TILE_Q_LARGE

#define CK_GGML_FA_TILE_Q_LARGE 336

Definition at line 2889 of file attention_kernels.c.

◆ CK_GGML_FA_TILE_Q_LARGE_MIN_TOKENS

#define CK_GGML_FA_TILE_Q_LARGE_MIN_TOKENS 1536

Definition at line 2891 of file attention_kernels.c.

◆ CK_MUSE_PREFILL_QUERY_TILE

#define CK_MUSE_PREFILL_QUERY_TILE 1024

Definition at line 9051 of file attention_kernels.c.

◆ CK_NOINLINE

#define CK_NOINLINE

Definition at line 402 of file attention_kernels.c.

◆ CK_OPTNONE

#define CK_OPTNONE

Definition at line 410 of file attention_kernels.c.

◆ FLASH_QUERY_IMPL [1/2]

#define FLASH_QUERY_IMPL attention_flash_query_causal

◆ FLASH_QUERY_IMPL [2/2]

#define FLASH_QUERY_IMPL attention_flash_query_causal

◆ FLASH_QUERY_IMPL_DECODE

#define FLASH_QUERY_IMPL_DECODE attention_flash_query_causal

◆ RTLD_DEFAULT

#define RTLD_DEFAULT ((void *)0)

Definition at line 32 of file attention_kernels.c.

Typedef Documentation

◆ ck_attention_math_f32_fn

typedef float(* ck_attention_math_f32_fn) (float)

Definition at line 1241 of file attention_kernels.c.

Function Documentation

◆ attention_backward_causal_head_major()

void attention_backward_causal_head_major	(	const float *	d_output,
		const float *	q,
		const float *	k,
		const float *	v,
		const float *	attn_weights,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		float *	d_scores,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	aligned_context_window
	)

Causal attention backward (non-GQA version)

Test:

test_attention_backward.py::TestAttentionBackward::test_backward

test_attention_backward.py::TestAttentionBackward::test_backward_vs_separate

test_parity.py::test_attention_backward_parity

Non-GQA version where num_heads == num_kv_heads. Simpler than GQA, no head broadcasting needed.

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

Definition at line 9931 of file attention_kernels.c.

{
    attention_backward_causal_head_major_gqa(
        d_output, q, k, v, attn_weights,
        d_q, d_k, d_v, d_scores,
        num_heads, num_heads,  // num_kv_heads == num_heads
        num_tokens, head_dim, aligned_head_dim, aligned_context_window);
}

References attention_backward_causal_head_major_gqa().

◆ attention_backward_causal_head_major_gqa()

void attention_backward_causal_head_major_gqa	(	const float *	d_output,
		const float *	q,
		const float *	k,
		const float *	v,
		const float *	attn_weights,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		float *	d_scores,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	aligned_context_window
	)

GQA causal attention backward (score-matrix version)

Test:

test_attention_backward.py::TestAttentionBackwardGQA::test_gqa_backward

test_attention_backward.py::TestAttentionBackwardGQA::test_gqa_vs_separate

test_parity.py::test_attention_backward_parity

Computes dQ, dK, dV given dOutput and attention weights. Supports grouped-query attention with head broadcasting.

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

Definition at line 9792 of file attention_kernels.c.

{
    const float scale = 1.0f / sqrtf((float)head_dim);
    int T = num_tokens;
    int H = num_heads;
    int H_kv = num_kv_heads;
    int hd = head_dim;
    int ad = aligned_head_dim;
    int aw = aligned_context_window;
 
    const size_t d_q_elems = (size_t)H * (size_t)T * (size_t)ad;
    const size_t kv_elems = (size_t)H_kv * (size_t)T * (size_t)ad;
    /* Zero the aligned outputs so padded lanes never leak garbage to downstream GEMMs. */
    for (size_t idx = 0; idx < d_q_elems; ++idx) {
        d_q[idx] = 0.0f;
    }
    for (size_t idx = 0; idx < kv_elems; ++idx) {
        d_k[idx] = 0.0f;
        d_v[idx] = 0.0f;
    }
 
    // Process each query head
    for (int h = 0; h < H; ++h) {
        // Which KV head does this query head use?
        int kv_h = (int)((long long)h * (long long)H_kv / (long long)H);
 
        // ----------------------------------------------------------------
        // Step 1: d_weights = d_output @ V^T  and  d_v += weights^T @ d_output
        // ----------------------------------------------------------------
        // For each query position i, compute d_weights[i, j] for j <= i
        // and accumulate d_v[j] contributions
 
        for (int i = 0; i < T; ++i) {
            size_t d_out_base = qkv_index(h, i, 0, T, ad);
 
            for (int j = 0; j <= i; ++j) {
                size_t v_base = qkv_index(kv_h, j, 0, T, ad);
                size_t w_idx = score_index(h, i, j, aw);
                float w = attn_weights[w_idx];
 
                // d_weights[h, i, j] = d_output[h, i, :] @ v[kv_h, j, :]^T
                float dot = 0.0f;
                for (int dd = 0; dd < hd; ++dd) {
                    dot += d_output[d_out_base + dd] * v[v_base + dd];
                }
                d_scores[w_idx] = dot;
 
                // d_v[kv_h, j, :] += weights[h, i, j] * d_output[h, i, :]
                for (int dd = 0; dd < hd; ++dd) {
                    d_v[v_base + dd] += w * d_output[d_out_base + dd];
                }
            }
 
            // Zero out upper triangle of d_scores
            for (int j = i + 1; j < T; ++j) {
                d_scores[score_index(h, i, j, aw)] = 0.0f;
            }
            /* Scores scratch uses aligned_context_window, zero the padded columns. */
            for (int j = T; j < aw; ++j) {
                d_scores[score_index(h, i, j, aw)] = 0.0f;
            }
        }
 
        // ----------------------------------------------------------------
        // Step 2: Backward through softmax (in-place on d_scores for this head)
        // ----------------------------------------------------------------
        // d_scores = softmax_backward(d_scores, attn_weights)
        // Formula: d_score[i,j] = w[i,j] * (d_w[i,j] - sum_k(w[i,k] * d_w[i,k]))
 
        for (int i = 0; i < T; ++i) {
            int base = h * aw * aw + i * aw;
 
            // Compute dot product: sum_j w[i,j] * d_w[i,j]
            float dot_product = 0.0f;
            for (int j = 0; j <= i; ++j) {
                float wt = attn_weights[base + j];
                float dw = d_scores[base + j];
                dot_product += wt * dw;
            }
 
            // Apply softmax backward formula
            for (int j = 0; j <= i; ++j) {
                float wt = attn_weights[base + j];
                float dw = d_scores[base + j];
                d_scores[base + j] = wt * (dw - dot_product);
            }
        }
 
        // ----------------------------------------------------------------
        // Step 3: d_q = d_scores @ K * scale
        //         d_k += d_scores^T @ Q * scale
        // ----------------------------------------------------------------
 
        for (int i = 0; i < T; ++i) {
            size_t d_q_base = qkv_index(h, i, 0, T, ad);
            size_t q_base = qkv_index(h, i, 0, T, ad);
 
            // d_q[h, i, :] = sum_j d_scores[h, i, j] * k[kv_h, j, :] * scale
            // d_k[kv_h, j, :] += d_scores[h, i, j] * q[h, i, :] * scale
            for (int j = 0; j <= i; ++j) {
                size_t k_base = qkv_index(kv_h, j, 0, T, ad);
                size_t d_k_base = qkv_index(kv_h, j, 0, T, ad);
                float ds = d_scores[score_index(h, i, j, aw)] * scale;
 
                for (int dd = 0; dd < hd; ++dd) {
                    d_q[d_q_base + dd] += ds * k[k_base + dd];
                    d_k[d_k_base + dd] += ds * q[q_base + dd];
                }
            }
        }
    }
}

References qkv_index(), and score_index().

Referenced by attention_backward_causal_head_major(), attention_backward_causal_head_major_gqa_bf16(), and ck_layer_backward_rmsnorm_swiglu().

◆ attention_backward_causal_head_major_gqa_bf16()

void attention_backward_causal_head_major_gqa_bf16	(	const uint16_t *	d_output,
		float *	d_x,
		const uint16_t *	q,
		const uint16_t *	k,
		const uint16_t *	v,
		const float *	attn_weights,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		float *	d_scores,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	aligned_context_window,
		float *	scratch_d_output,
		float *	scratch_q,
		float *	scratch_k,
		float *	scratch_v
	)

BF16 attention backward with caller-provided scratch buffers

Test:: bf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_backward

Accepts BF16 inputs, converts to FP32, runs FP32 backward. Caller provides scratch buffers (no per-call malloc).

After changes: make test

Definition at line 9739 of file attention_kernels.c.

{
    (void)d_x;
    const size_t head_elems = (size_t)num_heads * (size_t)num_tokens * (size_t)aligned_head_dim;
    const size_t kv_elems = (size_t)num_kv_heads * (size_t)num_tokens * (size_t)aligned_head_dim;
 
    if (!scratch_d_output || !scratch_q || !scratch_k || !scratch_v) return;
 
    convert_bf16_tensor_to_buf(d_output, scratch_d_output, head_elems);
    convert_bf16_tensor_to_buf(q, scratch_q, head_elems);
    convert_bf16_tensor_to_buf(k, scratch_k, kv_elems);
    convert_bf16_tensor_to_buf(v, scratch_v, kv_elems);
 
    attention_backward_causal_head_major_gqa(scratch_d_output, scratch_q, scratch_k, scratch_v,
                                             attn_weights,
                                             d_q, d_k, d_v, d_scores,
                                             num_heads, num_kv_heads,
                                             num_tokens, head_dim,
                                             aligned_head_dim, aligned_context_window);
    /* No free - caller owns scratch buffers */
}

References attention_backward_causal_head_major_gqa(), and convert_bf16_tensor_to_buf().

◆ attention_flash_query_causal()

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

static

Definition at line 2015 of file attention_kernels.c.

{
    // Online softmax:
    //   m = running max, s = running sum(exp(score - m))
    //   out = sum(exp(score - m) * v)
    float m = -INFINITY;
    float s = 0.0f;
 
    for (int d = 0; d < head_dim; ++d) {
        out_vec[d] = 0.0f;
    }
 
    for (int j = 0; j < kv_tokens; ++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.

Referenced by ck_attention_flash_query_auto().

◆ attention_flash_query_causal_exact()

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

static

Definition at line 2077 of file attention_kernels.c.

{
    if (kv_tokens <= 0) {
        for (int d = 0; d < aligned_head_dim; ++d) {
            out_vec[d] = 0.0f;
        }
        return;
    }
 
    for (int d = 0; d < aligned_head_dim; ++d) {
        out_vec[d] = 0.0f;
    }
 
    float max_score = -INFINITY;
    for (int j = 0; j < kv_tokens; ++j) {
        const float *k_vec = k_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 > max_score) {
            max_score = score;
        }
    }
 
    float sum = 0.0f;
    for (int j = 0; j < kv_tokens; ++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;
        float w = expf(score - max_score);
        sum += w;
        for (int d = 0; d < head_dim; ++d) {
            out_vec[d] += w * v_vec[d];
        }
    }
 
    if (sum > 0.0f) {
        float inv_sum = 1.0f / sum;
        for (int d = 0; d < head_dim; ++d) {
            out_vec[d] *= inv_sum;
        }
    } else {
        for (int d = 0; d < head_dim; ++d) {
            out_vec[d] = 0.0f;
        }
    }
    for (int d = head_dim; d < aligned_head_dim; ++d) {
        out_vec[d] = 0.0f;
    }
}

References score.

Referenced by attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl().

◆ attention_flash_query_causal_exact_f16kv()

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

static

Definition at line 2193 of file attention_kernels.c.

{
    if (kv_tokens <= 0) {
        for (int d = 0; d < aligned_head_dim; ++d) {
            out_vec[d] = 0.0f;
        }
        return;
    }
 
    for (int d = 0; d < aligned_head_dim; ++d) {
        out_vec[d] = 0.0f;
    }
 
    // Mirror llama.cpp GGML flash-attention more closely:
    // - Q is converted through FP16 before the KQ dot
    // - V accumulation is rounded through FP16 at each update
    // - the softmax accumulator uses the online max/sum form
    float sum = 0.0f;
    float max_score = -INFINITY;
    for (int j = 0; j < kv_tokens; ++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 += ck_round_fp16_scalar(q_vec[d]) * ck_round_fp16_scalar(k_vec[d]);
        }
        float score = dot * scale;
 
        const float prev_max = max_score;
        float max_scale = 1.0f;
        float value_scale = 1.0f;
 
        if (score > max_score) {
            max_score = score;
            max_scale = isfinite(prev_max) ? expf(prev_max - max_score) : 0.0f;
            for (int d = 0; d < head_dim; ++d) {
                out_vec[d] = ck_round_fp16_scalar(out_vec[d] * max_scale);
            }
        } else {
            value_scale = expf(score - max_score);
        }
 
        for (int d = 0; d < head_dim; ++d) {
            const float v_rounded = ck_round_fp16_scalar(v_vec[d]);
            const float updated = out_vec[d] + value_scale * v_rounded;
            out_vec[d] = ck_round_fp16_scalar(updated);
        }
 
        sum = sum * max_scale + value_scale;
    }
 
    if (sum > 0.0f) {
        float inv_sum = 1.0f / sum;
        for (int d = 0; d < head_dim; ++d) {
            out_vec[d] *= inv_sum;
        }
    } else {
        for (int d = 0; d < head_dim; ++d) {
            out_vec[d] = 0.0f;
        }
    }
    for (int d = head_dim; d < aligned_head_dim; ++d) {
        out_vec[d] = 0.0f;
    }
}

References ck_round_fp16_scalar(), and score.

Referenced by attention_forward_decode_head_major_gqa_flash_f16kv(), attention_forward_decode_head_major_gqa_regular(), and ck_attention_causal_f16kv_work().

◆ attention_flash_query_causal_exact_prerounded_f16kv()

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

static

Definition at line 2268 of file attention_kernels.c.

{
    if (kv_tokens <= 0) {
        memset(out_vec, 0, (size_t)aligned_head_dim * sizeof(float));
        return;
    }
 
    float *q_rounded = (float *)alloca((size_t)head_dim * sizeof(float));
    ck_round_fp16_buffer(q_vec, q_rounded, (size_t)head_dim);
    memset(out_vec, 0, (size_t)aligned_head_dim * sizeof(float));
 
    float sum = 0.0f;
    float max_score = -INFINITY;
    for (int j = 0; j < kv_tokens; ++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_rounded[d] * k_vec[d];
        }
        const float score = dot * scale;
        const float prev_max = max_score;
        float max_scale = 1.0f;
        float value_scale = 1.0f;
        if (score > max_score) {
            max_score = score;
            max_scale = isfinite(prev_max) ? expf(prev_max - max_score) : 0.0f;
            for (int d = 0; d < head_dim; ++d) {
                out_vec[d] = ck_round_fp16_scalar(out_vec[d] * max_scale);
            }
        } else {
            value_scale = expf(score - max_score);
        }
        for (int d = 0; d < head_dim; ++d) {
            const float updated = out_vec[d] + value_scale * v_vec[d];
            out_vec[d] = ck_round_fp16_scalar(updated);
        }
        sum = sum * max_scale + value_scale;
    }
 
    if (sum > 0.0f) {
        const float inv_sum = 1.0f / sum;
        for (int d = 0; d < head_dim; ++d) {
            out_vec[d] *= inv_sum;
        }
    } else {
        memset(out_vec, 0, (size_t)head_dim * sizeof(float));
    }
}

References ck_round_fp16_buffer(), ck_round_fp16_scalar(), and score.

Referenced by ck_attention_causal_f16kv_work().

◆ attention_forward_causal_head_major()

void attention_forward_causal_head_major	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	scores,
		float *	output,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	aligned_context_window
	)

Causal attention forward (score-matrix version)

Test:

test_attention.py::TestAttentionForward::test_causal_forward

test_attention.py::TestAttentionForward::test_gqa_broadcast

test_attention.py::TestAttentionForward::test_exact_vs_fast

test_parity.py::test_attention_parity

Computes softmax(Q @ K^T / sqrt(d)) @ V with causal masking. Uses O(N^2) memory for scores matrix.

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

Definition at line 1355 of file attention_kernels.c.

{
    const float scale = 1.0f / sqrtf((float)head_dim);
 
    // Phase 1: compute scaled dot-product scores Q·K^T / sqrt(d_k),
    // lower triangle only (j <= i).
    for (int h = 0; h < num_heads; ++h) {
        for (int i = 0; i < num_tokens; ++i) {
            for (int j = 0; j <= i; ++j) {
                float dot = 0.0f;
                size_t base_q = qkv_index(h, i, 0, num_tokens, aligned_head_dim);
                size_t base_k = qkv_index(h, j, 0, num_tokens, aligned_head_dim);
 
                for (int d = 0; d < head_dim; ++d) {
                    dot += q[base_q + d] * k[base_k + d];
                }
 
                scores[score_index(h, i, j, aligned_context_window)] = dot * scale;
            }
 
            // Ensure upper triangle is zeroed so there are no stale values
            // before the softmax kernel runs.
            for (int j = i + 1; j < num_tokens; ++j) {
                scores[score_index(h, i, j, aligned_context_window)] = 0.0f;
            }
        }
    }
 
    // Phase 2: apply causal row-wise softmax in-place over j <= i.
    causal_softmax_head_major(scores,
                              num_heads,
                              num_tokens,
                              aligned_context_window);
 
    // Phase 3: attention weights · V.
    for (int h = 0; h < num_heads; ++h) {
        for (int i = 0; i < num_tokens; ++i) {
            size_t out_base = qkv_index(h, i, 0, num_tokens, aligned_head_dim);
 
            // Zero the full aligned head slice so padded dims stay clean.
            for (int d = 0; d < aligned_head_dim; ++d) {
                output[out_base + d] = 0.0f;
            }
 
            // Weighted sum over causal positions.
            for (int j = 0; j <= i; ++j) {
                float w = scores[score_index(h, i, j, aligned_context_window)];
                size_t v_base = qkv_index(h, j, 0, num_tokens, aligned_head_dim);
 
                for (int d = 0; d < head_dim; ++d) {
                    output[out_base + d] += w * v[v_base + d];
                }
            }
        }
    }
}

References causal_softmax_head_major(), qkv_index(), and score_index().

◆ attention_forward_causal_head_major_exact()

void attention_forward_causal_head_major_exact	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	scores,
		float *	output,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	aligned_context_window
	)

Causal attention forward (exact version using stdlib expf)

Test:

test_attention.py::TestAttentionForward::test_exact_single

test_attention.py::TestAttentionForward::test_exact_vs_fast

Uses standard library expf for numerical accuracy reference. Slower but provides maximum accuracy.

After changes: make test

Definition at line 1431 of file attention_kernels.c.

{
    const float scale = 1.0f / sqrtf((float)head_dim);
 
    // Phase 1: compute scaled dot-product scores Q·K^T / sqrt(d_k),
    // lower triangle only (j <= i).
    for (int h = 0; h < num_heads; ++h) {
        for (int i = 0; i < num_tokens; ++i) {
            for (int j = 0; j <= i; ++j) {
                float dot = 0.0f;
                size_t base_q = qkv_index(h, i, 0, num_tokens, aligned_head_dim);
                size_t base_k = qkv_index(h, j, 0, num_tokens, aligned_head_dim);
 
                for (int d = 0; d < head_dim; ++d) {
                    dot += q[base_q + d] * k[base_k + d];
                }
 
                scores[score_index(h, i, j, aligned_context_window)] = dot * scale;
            }
 
            // Ensure upper triangle is zeroed so there are no stale values
            // before the softmax kernel runs.
            for (int j = i + 1; j < num_tokens; ++j) {
                scores[score_index(h, i, j, aligned_context_window)] = 0.0f;
            }
        }
    }
 
    // Phase 2: apply causal row-wise softmax using exact expf.
    causal_softmax_head_major_exact(scores,
                                     num_heads,
                                     num_tokens,
                                     aligned_context_window);
 
    // Phase 3: attention weights · V.
    for (int h = 0; h < num_heads; ++h) {
        for (int i = 0; i < num_tokens; ++i) {
            size_t out_base = qkv_index(h, i, 0, num_tokens, aligned_head_dim);
 
            // Zero the full aligned head slice so padded dims stay clean.
            for (int d = 0; d < aligned_head_dim; ++d) {
                output[out_base + d] = 0.0f;
            }
 
            // Weighted sum over causal positions.
            for (int j = 0; j <= i; ++j) {
                float w = scores[score_index(h, i, j, aligned_context_window)];
                size_t v_base = qkv_index(h, j, 0, num_tokens, aligned_head_dim);
 
                for (int d = 0; d < head_dim; ++d) {
                    output[out_base + d] += w * v[v_base + d];
                }
            }
        }
    }
}

References causal_softmax_head_major_exact(), qkv_index(), and score_index().

◆ attention_forward_causal_head_major_gqa()

void attention_forward_causal_head_major_gqa	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	scores,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	aligned_context_window
	)

GQA causal attention forward (score-matrix version)

Test:

test_attention.py::TestAttentionForward::test_gqa_forward

test_attention.py::TestAttentionForward::test_gqa_broadcast

test_attention_backward.py::TestAttentionBackwardGQA::test_gqa_backward

test_parity.py::test_attention_gqa_parity

Grouped-query attention: Q has num_heads, K/V have num_kv_heads. Each query head maps to a KV head via ratio.

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

Definition at line 1509 of file attention_kernels.c.

{
    const float scale = 1.0f / sqrtf((float)head_dim);
 
    for (int h = 0; h < num_heads; ++h) {
        int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        for (int i = 0; i < num_tokens; ++i) {
            for (int j = 0; j <= i; ++j) {
                float dot = 0.0f;
                size_t base_q = qkv_index(h, i, 0, num_tokens, aligned_head_dim);
                size_t base_k = qkv_index(kv_head, j, 0, num_tokens, aligned_head_dim);
 
                for (int d = 0; d < head_dim; ++d) {
                    dot += q[base_q + d] * k[base_k + d];
                }
 
                scores[score_index(h, i, j, aligned_context_window)] = dot * scale;
            }
 
            for (int j = i + 1; j < num_tokens; ++j) {
                scores[score_index(h, i, j, aligned_context_window)] = 0.0f;
            }
        }
    }
 
    causal_softmax_head_major(scores,
                              num_heads,
                              num_tokens,
                              aligned_context_window);
 
    for (int h = 0; h < num_heads; ++h) {
        int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        for (int i = 0; i < num_tokens; ++i) {
            size_t out_base = qkv_index(h, i, 0, num_tokens, aligned_head_dim);
            for (int d = 0; d < aligned_head_dim; ++d) {
                output[out_base + d] = 0.0f;
            }
 
            for (int j = 0; j <= i; ++j) {
                float w = scores[score_index(h, i, j, aligned_context_window)];
                size_t v_base = qkv_index(kv_head, j, 0, num_tokens, aligned_head_dim);
 
                for (int d = 0; d < head_dim; ++d) {
                    output[out_base + d] += w * v[v_base + d];
                }
            }
        }
    }
}

References causal_softmax_head_major(), qkv_index(), and score_index().

Referenced by ck_layer_forward_rmsnorm_swiglu(), ck_layer_forward_rmsnorm_swiglu_q4_k(), ck_layer_forward_rmsnorm_swiglu_quant(), and ck_layer_forward_rmsnorm_swiglu_ref().

◆ attention_forward_causal_head_major_gqa_bf16()

void attention_forward_causal_head_major_gqa_bf16	(	const uint16_t *	q,
		const uint16_t *	k,
		const uint16_t *	v,
		float *	scores,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	aligned_context_window,
		float *	scratch_q,
		float *	scratch_k,
		float *	scratch_v
	)

BF16 GQA causal attention forward

Test:

bf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_forward

bf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_gqa

bf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_flash

Accepts BF16 inputs, converts to FP32, uses exact softmax. Caller provides scratch buffers (no per-call malloc).

After changes: make test

Definition at line 1651 of file attention_kernels.c.

{
    const size_t q_elems = (size_t)num_heads * (size_t)num_tokens * (size_t)aligned_head_dim;
    const size_t kv_elems = (size_t)num_kv_heads * (size_t)num_tokens * (size_t)aligned_head_dim;
 
    if (!scratch_q || !scratch_k || !scratch_v) return;
 
    convert_bf16_tensor_to_buf(q, scratch_q, q_elems);
    convert_bf16_tensor_to_buf(k, scratch_k, kv_elems);
    convert_bf16_tensor_to_buf(v, scratch_v, kv_elems);
 
    // Use exact version to avoid fast exp approximation error accumulating
    // with BF16 precision loss.
    attention_forward_causal_head_major_gqa_exact(scratch_q, scratch_k, scratch_v,
                                                   scores, output,
                                                   num_heads, num_kv_heads,
                                                   num_tokens, head_dim,
                                                   aligned_head_dim, aligned_context_window);
    /* No free - caller owns scratch buffers */
}

References attention_forward_causal_head_major_gqa_exact(), and convert_bf16_tensor_to_buf().

◆ attention_forward_causal_head_major_gqa_exact()

void attention_forward_causal_head_major_gqa_exact	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	scores,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	aligned_context_window
	)

GQA causal attention forward (exact version using stdlib expf)

Test:

test_attention.py::TestAttentionForward::test_gqa_exact

bf16/test_attention_bf16.py::TestAttentionBF16::test_bf16_gqa

Uses standard library expf for numerical accuracy reference. Used by BF16 wrapper to avoid approximation error accumulation.

After changes: make test

Definition at line 1579 of file attention_kernels.c.

{
    const float scale = 1.0f / sqrtf((float)head_dim);
 
    for (int h = 0; h < num_heads; ++h) {
        int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        for (int i = 0; i < num_tokens; ++i) {
            for (int j = 0; j <= i; ++j) {
                float dot = 0.0f;
                size_t base_q = qkv_index(h, i, 0, num_tokens, aligned_head_dim);
                size_t base_k = qkv_index(kv_head, j, 0, num_tokens, aligned_head_dim);
 
                for (int d = 0; d < head_dim; ++d) {
                    dot += q[base_q + d] * k[base_k + d];
                }
 
                scores[score_index(h, i, j, aligned_context_window)] = dot * scale;
            }
 
            for (int j = i + 1; j < num_tokens; ++j) {
                scores[score_index(h, i, j, aligned_context_window)] = 0.0f;
            }
        }
    }
 
    // Use exact softmax with standard library expf
    causal_softmax_head_major_exact(scores,
                                     num_heads,
                                     num_tokens,
                                     aligned_context_window);
 
    for (int h = 0; h < num_heads; ++h) {
        int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        for (int i = 0; i < num_tokens; ++i) {
            size_t out_base = qkv_index(h, i, 0, num_tokens, aligned_head_dim);
            for (int d = 0; d < aligned_head_dim; ++d) {
                output[out_base + d] = 0.0f;
            }
 
            for (int j = 0; j <= i; ++j) {
                float w = scores[score_index(h, i, j, aligned_context_window)];
                size_t v_base = qkv_index(kv_head, j, 0, num_tokens, aligned_head_dim);
 
                for (int d = 0; d < head_dim; ++d) {
                    output[out_base + d] += w * v[v_base + d];
                }
            }
        }
    }
}

References causal_softmax_head_major_exact(), qkv_index(), and score_index().

Referenced by attention_forward_causal_head_major_gqa_bf16(), ck_layer_forward_rmsnorm_swiglu(), and ck_layer_forward_rmsnorm_swiglu_ref().

◆ attention_forward_causal_head_major_gqa_flash()

void attention_forward_causal_head_major_gqa_flash	(	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
	)

Definition at line 4444 of file attention_kernels.c.

{
    attention_forward_head_major_gqa_flash_impl(q, k, v, output,
                                                num_heads, num_kv_heads,
                                                num_tokens, head_dim,
                                                aligned_head_dim,
                                                /*kv_stride_tokens=*/num_tokens,
                                                /*causal=*/1,
                                                /*round_full_kv_fp16=*/0,
                                                /*output_token_major=*/0,
                                                1.0f / sqrtf((float)head_dim));
}

References attention_forward_head_major_gqa_flash_impl().

◆ attention_forward_causal_head_major_gqa_flash_strided()

void attention_forward_causal_head_major_gqa_flash_strided	(	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
	)

Flash attention forward with custom KV stride (for KV cache)

Test:

test_flash_attention.py::TestFlashAttention::test_flash_strided

test_kv_cache_attention.py::TestKVCacheAttention::test_flash_attention

Variant with configurable kv_stride_tokens for KV cache layouts where K/V may not be contiguous in memory.

After changes: make test

Definition at line 4496 of file attention_kernels.c.

{
    attention_forward_head_major_gqa_flash_impl(q, k, v, output,
                                                num_heads, num_kv_heads,
                                                num_tokens, head_dim,
                                                aligned_head_dim,
                                                kv_stride_tokens,
                                                /*causal=*/1,
                                                /*round_full_kv_fp16=*/0,
                                                /*output_token_major=*/0,
                                                1.0f / sqrtf((float)head_dim));
}

References attention_forward_head_major_gqa_flash_impl().

◆ attention_forward_causal_head_major_gqa_flash_strided_f16kv()

void attention_forward_causal_head_major_gqa_flash_strided_f16kv	(	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
	)

Definition at line 5417 of file attention_kernels.c.

{
    size_t elements = 0;
    if (num_kv_heads > 0 && kv_stride_tokens > 0 && aligned_head_dim > 0 &&
        (size_t)num_kv_heads <= SIZE_MAX / (size_t)kv_stride_tokens) {
        const size_t rows = (size_t)num_kv_heads * (size_t)kv_stride_tokens;
        if (rows <= SIZE_MAX / (size_t)aligned_head_dim) {
            elements = rows * (size_t)aligned_head_dim;
        }
    }
    float *workspace = elements > 0 && elements <= SIZE_MAX / (2 * sizeof(float))
        ? (float *)malloc(2 * elements * sizeof(float))
        : NULL;
    if (!workspace) {
        attention_forward_causal_head_major_gqa_flash_strided_f16kv_serial(
            q, k, v, output, num_heads, num_kv_heads, num_tokens,
            head_dim, aligned_head_dim, kv_stride_tokens);
        return;
    }
    attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace(
        q, k, v, output, num_heads, num_kv_heads, num_tokens, head_dim,
        aligned_head_dim, kv_stride_tokens, workspace,
        2 * elements * sizeof(float));
    free(workspace);
}

References attention_forward_causal_head_major_gqa_flash_strided_f16kv_serial(), and attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace().

◆ attention_forward_causal_head_major_gqa_flash_strided_f16kv_serial()

void attention_forward_causal_head_major_gqa_flash_strided_f16kv_serial	(	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
	)

Definition at line 5361 of file attention_kernels.c.

{
    if (!q || !k || !v || !output || num_heads <= 0 || num_kv_heads <= 0 ||
        num_tokens <= 0 || kv_stride_tokens < num_tokens) {
        return;
    }
    ck_attention_causal_f16kv_args_t args = {
        q, k, v, output, num_heads, num_kv_heads, num_tokens,
        head_dim, aligned_head_dim, kv_stride_tokens, 0,
    };
    ck_attention_causal_f16kv_work(0, 1, &args);
}

References ck_attention_causal_f16kv_work().

Referenced by attention_forward_causal_head_major_gqa_flash_strided_f16kv().

◆ attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace()

void attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace	(	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,
		float *	rounded_kv,
		size_t	rounded_kv_bytes
	)

Definition at line 5377 of file attention_kernels.c.

{
    if (!q || !k || !v || !output || num_heads <= 0 || num_kv_heads <= 0 ||
        num_tokens <= 0 || kv_stride_tokens < num_tokens) {
        return;
    }
    size_t elements = 0;
    if ((size_t)num_kv_heads <= SIZE_MAX / (size_t)kv_stride_tokens) {
        const size_t rows = (size_t)num_kv_heads * (size_t)kv_stride_tokens;
        if (rows <= SIZE_MAX / (size_t)aligned_head_dim) {
            elements = rows * (size_t)aligned_head_dim;
        }
    }
    if (elements == 0 || elements > SIZE_MAX / (2 * sizeof(float)) ||
        !rounded_kv || rounded_kv_bytes < 2 * elements * sizeof(float)) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: FP16-KV attention workspace is too small\n");
        abort();
    }
    ck_round_fp16_buffer(k, rounded_kv, elements);
    ck_round_fp16_buffer(v, rounded_kv + elements, elements);
    ck_attention_causal_f16kv_args_t args = {
        q, rounded_kv, rounded_kv + elements, output,
        num_heads, num_kv_heads, num_tokens,
        head_dim, aligned_head_dim, kv_stride_tokens, 1,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    const int workers = pool ? ck_threadpool_n_threads(pool) : 1;
    const int active = workers < num_heads ? workers : num_heads;
    if (active > 1) {
        ck_threadpool_dispatch_n(pool, active, ck_attention_causal_f16kv_work, &args);
    } else {
        ck_attention_causal_f16kv_work(0, 1, &args);
    }
}

References ck_attention_causal_f16kv_work(), ck_round_fp16_buffer(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by attention_forward_causal_head_major_gqa_flash_strided_f16kv().

◆ attention_forward_causal_head_major_gqa_flash_strided_gemma4()

void attention_forward_causal_head_major_gqa_flash_strided_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
	)

Definition at line 5098 of file attention_kernels.c.

{
    (void)head_dim;
    attention_forward_head_major_gqa_flash_impl(q, k, v, output,
                                                num_heads, num_kv_heads,
                                                num_tokens, head_dim,
                                                aligned_head_dim,
                                                kv_stride_tokens,
                                                /*causal=*/1,
                                                /*round_full_kv_fp16=*/0,
                                                /*output_token_major=*/0,
                                                1.0f);
}

References attention_forward_head_major_gqa_flash_impl().

Referenced by attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl(), and attention_forward_causal_head_major_shared_kv_gemma4().

◆ attention_forward_causal_head_major_gqa_flash_strided_gemma4_token_output()

void attention_forward_causal_head_major_gqa_flash_strided_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
	)

Definition at line 5121 of file attention_kernels.c.

{
    attention_forward_head_major_gqa_flash_impl(q, k, v, output,
                                                num_heads, num_kv_heads,
                                                num_tokens, head_dim,
                                                aligned_head_dim,
                                                kv_stride_tokens,
                                                /*causal=*/1,
                                                /*round_full_kv_fp16=*/0,
                                                /*output_token_major=*/1,
                                                1.0f);
}

References attention_forward_head_major_gqa_flash_impl().

Referenced by attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl().

◆ attention_forward_causal_head_major_gqa_flash_strided_token_output()

void attention_forward_causal_head_major_gqa_flash_strided_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
	)

Definition at line 4518 of file attention_kernels.c.

{
    attention_forward_head_major_gqa_flash_impl(q, k, v, output,
                                                num_heads, num_kv_heads,
                                                num_tokens, head_dim,
                                                aligned_head_dim,
                                                kv_stride_tokens,
                                                /*causal=*/1,
                                                /*round_full_kv_fp16=*/0,
                                                /*output_token_major=*/1,
                                                1.0f / sqrtf((float)head_dim));
}

References attention_forward_head_major_gqa_flash_impl().

◆ attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace()

void attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace	(	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,
		float *	scores,
		size_t	scores_bytes,
		float *	value_columns,
		size_t	value_columns_bytes,
		float *	scaled_scores,
		size_t	scaled_scores_bytes
	)

Definition at line 6086 of file attention_kernels.c.

{
    if (!q || !k || !v || !output || !scores || !value_columns || !scaled_scores ||
        num_heads <= 0 || num_kv_heads <= 0 || num_tokens <= 0 || head_dim <= 0 ||
        aligned_head_dim < head_dim || kv_stride_tokens < num_tokens ||
        num_heads > INT_MAX / num_tokens) return;
 
    const size_t score_elements = scores_bytes / sizeof(float);
    const size_t scaled_elements = scaled_scores_bytes / sizeof(float);
    const size_t value_elements = value_columns_bytes / sizeof(float);
    if ((size_t)num_kv_heads > SIZE_MAX / (size_t)head_dim) return;
    const size_t value_rows = (size_t)num_kv_heads * (size_t)head_dim;
    const size_t value_capacity = value_rows > 0 ? value_elements / value_rows : 0;
    const size_t padded_size = num_tokens < 256
        ? 256 : (((size_t)num_tokens + 255) / 256) * 256;
    const int padded = padded_size > (size_t)INT_MAX ? 0 : (int)padded_size;
 
    /* Older direct callers provide one score row and one value bank. Keep that
     * ABI usable while generated v8 runtimes provide the parallel workspace. */
    if (padded <= 0 || value_capacity < (size_t)padded ||
        score_elements < (size_t)padded * 2 ||
        scaled_elements < (size_t)padded * 2 ||
        value_elements < value_rows * (size_t)padded) {
        ck_attention_llama_regular_impl(
            q, k, v, output, num_heads, num_kv_heads, num_tokens, num_tokens,
            head_dim, aligned_head_dim, kv_stride_tokens, sliding_window,
            scores, scores_bytes, value_columns, value_columns_bytes,
            scaled_scores, scaled_scores_bytes, 1);
        return;
    }
 
    if (ck_attention_vec_dump_enabled()) {
        ck_attention_llama_regular_impl(
            q, k, v, output, num_heads, num_kv_heads, num_tokens, num_tokens,
            head_dim, aligned_head_dim, kv_stride_tokens, sliding_window,
            scores, (size_t)padded * sizeof(float),
            value_columns, (size_t)head_dim * (size_t)padded * sizeof(float),
            scaled_scores, (size_t)padded * sizeof(float), 1);
        return;
    }
 
    const size_t kv_head_stride =
        (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
    for (int kv_head = 0; kv_head < num_kv_heads; ++kv_head) {
        const float *value_head = v + (size_t)kv_head * kv_head_stride;
        float *packed_head = value_columns +
            (size_t)kv_head * (size_t)head_dim * (size_t)padded;
        for (int dim = 0; dim < head_dim; ++dim) {
            float *column = packed_head + (size_t)dim * (size_t)padded;
            int token = 0;
            for (; token < num_tokens; ++token) {
                column[token] =
                    value_head[(size_t)token * (size_t)aligned_head_dim + (size_t)dim];
            }
            for (; token < padded; ++token) column[token] = 0.0f;
        }
    }
 
    ck_llama_regular_prefill_args_t args = {
        .q = q,
        .k = k,
        .value_columns = value_columns,
        .output = output,
        .scores = scores,
        .scaled_scores = scaled_scores,
        .num_heads = num_heads,
        .num_kv_heads = num_kv_heads,
        .num_tokens = num_tokens,
        .head_dim = head_dim,
        .aligned_head_dim = aligned_head_dim,
        .kv_stride_tokens = kv_stride_tokens,
        .sliding_window = sliding_window,
        .padded = padded,
    };
    atomic_init(&args.next_job, 0);
 
    size_t scratch_lanes = score_elements / (size_t)padded;
    const size_t scaled_lanes = scaled_elements / (size_t)padded;
    if (scaled_lanes < scratch_lanes) scratch_lanes = scaled_lanes;
    int active = scratch_lanes > (size_t)INT_MAX ? INT_MAX : (int)scratch_lanes;
    const int jobs = num_heads * num_tokens;
    if (active > jobs) active = jobs;
    ck_threadpool_t *pool = ck_threadpool_global();
    const int pool_threads = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > pool_threads) active = pool_threads;
    if (active < 1) active = 1;
    if (pool && active > 1) {
        ck_threadpool_dispatch_n(pool, active, ck_llama_regular_prefill_work, &args);
    } else {
        ck_llama_regular_prefill_work(0, 1, &args);
    }
}

References ck_attention_llama_regular_impl(), ck_attention_vec_dump_enabled(), ck_llama_regular_prefill_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), and token.

◆ attention_forward_causal_head_major_gqa_muse_eager_bf16_storage()

void attention_forward_causal_head_major_gqa_muse_eager_bf16_storage	(	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,
		float *	scores,
		size_t	scores_bytes,
		uint16_t *	key_rows,
		size_t	key_rows_bytes,
		uint16_t *	value_columns,
		size_t	value_columns_bytes,
		uint16_t *	gemm_rows,
		size_t	gemm_rows_bytes,
		uint16_t *	gemm_scores,
		size_t	gemm_scores_bytes
	)

Definition at line 9206 of file attention_kernels.c.

{
    ck_attention_muse_eager_bf16_storage_impl(
        q, k, v, output, num_heads, num_kv_heads, num_tokens, num_tokens,
        kv_stride_tokens, head_dim, aligned_head_dim, 0,
        scores, scores_bytes, key_rows, key_rows_bytes,
        value_columns, value_columns_bytes, gemm_rows, gemm_rows_bytes,
        gemm_scores, gemm_scores_bytes);
}

References ck_attention_muse_eager_bf16_storage_impl().

◆ attention_forward_causal_head_major_gqa_muse_eager_bf16_storage_sliding()

void attention_forward_causal_head_major_gqa_muse_eager_bf16_storage_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,
		float *	scores,
		size_t	scores_bytes,
		uint16_t *	key_rows,
		size_t	key_rows_bytes,
		uint16_t *	value_columns,
		size_t	value_columns_bytes,
		uint16_t *	gemm_rows,
		size_t	gemm_rows_bytes,
		uint16_t *	gemm_scores,
		size_t	gemm_scores_bytes
	)

Definition at line 9224 of file attention_kernels.c.

{
    ck_attention_muse_eager_bf16_storage_impl(
        q, k, v, output, num_heads, num_kv_heads, num_tokens, num_tokens,
        kv_stride_tokens, head_dim, aligned_head_dim, sliding_window,
        scores, scores_bytes, key_rows, key_rows_bytes,
        value_columns, value_columns_bytes, gemm_rows, gemm_rows_bytes,
        gemm_scores, gemm_scores_bytes);
}

References ck_attention_muse_eager_bf16_storage_impl().

◆ attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction
	)

Definition at line 9429 of file attention_kernels.c.

{
    if (q_tokens >= CK_GGML_FA_TILE_Q) {
        return attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace(
            q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
            past_tokens, cache_capacity, head_dim, aligned_head_dim, reduction,
            NULL, 0);
    }
    if (num_heads <= 0 || aligned_head_dim <= 0 ||
        (size_t) num_heads > SIZE_MAX / (size_t) aligned_head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    const size_t token_elements = (size_t) num_heads * (size_t) aligned_head_dim;
    if (token_elements > SIZE_MAX / (2 * sizeof(float))) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    const size_t workspace_bytes = 2 * token_elements * sizeof(float);
    float *workspace = (float *) malloc(workspace_bytes);
    if (!workspace) return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    const ck_attention_status_t status =
        attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace(
            q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
            past_tokens, cache_capacity, head_dim, aligned_head_dim, reduction,
            workspace, workspace_bytes);
    free(workspace);
    return status;
}

References attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, and CK_GGML_FA_TILE_Q.

◆ attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes
	)

Definition at line 9352 of file attention_kernels.c.

{
    if (!q || !k_cache || !v_cache || !output ||
        num_heads <= 0 || num_kv_heads <= 0 || q_tokens <= 0 ||
        past_tokens < 0 || past_tokens + q_tokens > cache_capacity ||
        head_dim <= 0 || aligned_head_dim < head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    if (reduction != CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA) {
        return CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT;
    }
 
    if (q_tokens >= CK_GGML_FA_TILE_Q) {
        ck_attention_f16_prefill_qtile64_args_t args = {
            .q = q, .k_cache = k_cache, .v_cache = v_cache, .output = output,
            .num_heads = num_heads, .num_kv_heads = num_kv_heads,
            .q_tokens = q_tokens, .past_tokens = past_tokens,
            .cache_capacity = cache_capacity, .head_dim = head_dim,
            .aligned_head_dim = aligned_head_dim, .cache_is_bf16 = 1,
        };
        ck_threadpool_t *pool = ck_threadpool_global();
        int active = pool ? ck_threadpool_n_threads(pool) : 1;
        if (active > num_kv_heads) active = num_kv_heads;
        if (pool && active > 1 && ck_threadpool_thread_id(pool) <= 0) {
            ck_threadpool_dispatch_n(pool, active, ck_attention_f16_prefill_qtile64_work, &args);
        } else {
            ck_attention_f16_prefill_qtile64_work(0, 1, &args);
        }
        const size_t count = (size_t)num_heads * (size_t)q_tokens * (size_t)aligned_head_dim;
        for (size_t i = 0; i < count; ++i) {
            output[i] = bf16_to_float(float_to_bf16(output[i]));
        }
        return CK_ATTENTION_STATUS_OK;
    }
 
    const size_t token_elems = (size_t)num_heads * (size_t)aligned_head_dim;
    if (token_elems > SIZE_MAX / (2 * sizeof(float)) ||
        !token_workspace || token_workspace_bytes < 2 * token_elems * sizeof(float)) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    float *q_token = token_workspace;
    float *out_token = token_workspace + token_elems;
    ck_attention_status_t status = CK_ATTENTION_STATUS_OK;
    for (int t = 0; t < q_tokens; ++t) {
        for (int h = 0; h < num_heads; ++h) {
            memcpy(q_token + (size_t)h * (size_t)aligned_head_dim,
                   q + ((size_t)h * (size_t)q_tokens + (size_t)t) * (size_t)aligned_head_dim,
                   (size_t)aligned_head_dim * sizeof(float));
        }
        status = attention_forward_decode_head_major_gqa_bf16cache_pytorch_contract(
            q_token, k_cache, v_cache, out_token,
            num_heads, num_kv_heads, past_tokens + t + 1, cache_capacity,
            head_dim, aligned_head_dim, reduction);
        if (status != CK_ATTENTION_STATUS_OK) break;
        for (int h = 0; h < num_heads; ++h) {
            memcpy(output + ((size_t)h * (size_t)q_tokens + (size_t)t) * (size_t)aligned_head_dim,
                   out_token + (size_t)h * (size_t)aligned_head_dim,
                   (size_t)aligned_head_dim * sizeof(float));
        }
    }
    return status;
}

References attention_forward_decode_head_major_gqa_bf16cache_pytorch_contract(), bf16_to_float(), ck_attention_f16_prefill_qtile64_work(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT, CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA, CK_GGML_FA_TILE_Q, ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_thread_id(), count, and float_to_bf16().

Referenced by attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract().

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_auto_workspace()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_auto_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes,
		void *	gqa_workspace,
		size_t	gqa_workspace_bytes,
		int	route_num_heads,
		int	route_num_kv_heads,
		int	route_head_dim,
		int	route_query_tokens,
		int	route_min_kv_tokens,
		int	route_workers,
		int	route_query_tile_size,
		int	route_concurrent_query_tiles
	)

Definition at line 7686 of file attention_kernels.c.

{
    if (route_num_heads <= 0 || route_num_kv_heads <= 0 ||
        route_num_heads % route_num_kv_heads != 0 || route_head_dim <= 0 ||
        route_query_tokens <= 0 || route_min_kv_tokens < route_query_tokens ||
        route_workers < route_num_kv_heads ||
        route_workers % route_num_kv_heads != 0 ||
        route_query_tile_size < 16 || route_query_tile_size > 128 ||
        route_query_tile_size % 16 != 0 ||
        route_concurrent_query_tiles <= 0 ||
        route_concurrent_query_tiles > 4) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    ck_threadpool_t *pool = ck_threadpool_global();
    const int active = pool ? ck_threadpool_n_threads(pool) : 1;
    const int kv_tokens = past_tokens + q_tokens;
    if (reduction == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64 &&
        num_heads == route_num_heads && num_kv_heads == route_num_kv_heads &&
        head_dim == route_head_dim && aligned_head_dim == route_head_dim &&
        q_tokens == route_query_tokens && kv_tokens >= route_min_kv_tokens &&
        active == route_workers) {
        return attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config(
            q, k_cache, v_cache, output, num_heads, num_kv_heads,
            q_tokens, past_tokens, cache_capacity, head_dim, aligned_head_dim,
            route_query_tile_size, route_concurrent_query_tiles,
            gqa_workspace, gqa_workspace_bytes);
    }
    return attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace(
        q, k_cache, v_cache, output, num_heads, num_kv_heads,
        q_tokens, past_tokens, cache_capacity, head_dim, aligned_head_dim,
        reduction, token_workspace, token_workspace_bytes);
}

References attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace(), attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64, ck_threadpool_global(), and ck_threadpool_n_threads().

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction
	)

Definition at line 8155 of file attention_kernels.c.

{
    if (reduction == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64 &&
        q_tokens >= CK_GGML_FA_TILE_Q) {
        return attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace(
            q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
            past_tokens, cache_capacity, head_dim, aligned_head_dim, reduction,
            NULL, 0);
    }
 
    if (num_heads <= 0 || aligned_head_dim <= 0 ||
        (size_t) num_heads > SIZE_MAX / (size_t) aligned_head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    const size_t token_elements = (size_t) num_heads * (size_t) aligned_head_dim;
    if (token_elements > SIZE_MAX / (2 * sizeof(float))) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    const size_t workspace_bytes = 2 * token_elements * sizeof(float);
    float *workspace = (float *) malloc(workspace_bytes);
    if (!workspace) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    const ck_attention_status_t status =
        attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace(
            q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
            past_tokens, cache_capacity, head_dim, aligned_head_dim, reduction,
            workspace, workspace_bytes);
    free(workspace);
    return status;
}

References attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64, and CK_GGML_FA_TILE_Q.

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes
	)

Definition at line 7770 of file attention_kernels.c.

{
    if (!q || !k_cache || !v_cache || !output ||
        num_heads <= 0 || num_kv_heads <= 0 || q_tokens <= 0 ||
        past_tokens < 0 || past_tokens + q_tokens > cache_capacity ||
        head_dim <= 0 || aligned_head_dim < head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
 
    if (reduction == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64 &&
        q_tokens < CK_GGML_FA_TILE_Q) {
        reduction = CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE;
    }
 
    if (reduction == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64) {
        return ck_attention_f16_prefill_qtile64_dispatch(
            q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
            past_tokens, cache_capacity, head_dim, aligned_head_dim, 0, 0.0f,
            (size_t) q_tokens * (size_t) aligned_head_dim,
            (size_t) q_tokens * (size_t) aligned_head_dim,
            CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES);
    }
 
    if (reduction != CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE &&
        reduction != CK_ATTN_REDUCTION_F16_ONLINE_FP32_MERGE &&
        reduction != CK_ATTN_REDUCTION_FP32_ONLINE) {
        return CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT;
    }
 
    const size_t token_elems = (size_t) num_heads * (size_t) aligned_head_dim;
    if (token_elems > SIZE_MAX / (2 * sizeof(float))) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    const size_t required_bytes = 2 * token_elems * sizeof(float);
    if (!token_workspace || token_workspace_bytes < required_bytes) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    float *q_token = token_workspace;
    float *out_token = token_workspace + token_elems;
 
    ck_attention_status_t status = CK_ATTENTION_STATUS_OK;
    for (int t = 0; t < q_tokens; ++t) {
        for (int h = 0; h < num_heads; ++h) {
            const float *src = q +
                ((size_t) h * (size_t) q_tokens + (size_t) t) * (size_t) aligned_head_dim;
            memcpy(q_token + (size_t) h * (size_t) aligned_head_dim,
                   src,
                   (size_t) aligned_head_dim * sizeof(float));
        }
 
        status = attention_forward_decode_head_major_gqa_flash_f16cache_contract(
            q_token, k_cache, v_cache, out_token,
            num_heads, num_kv_heads, past_tokens + t + 1, cache_capacity,
            head_dim, aligned_head_dim, reduction);
        if (status != CK_ATTENTION_STATUS_OK) {
            break;
        }
 
        for (int h = 0; h < num_heads; ++h) {
            float *dst = output +
                ((size_t) h * (size_t) q_tokens + (size_t) t) * (size_t) aligned_head_dim;
            memcpy(dst,
                   out_token + (size_t) h * (size_t) aligned_head_dim,
                   (size_t) aligned_head_dim * sizeof(float));
        }
    }
 
    return status;
}

References attention_forward_decode_head_major_gqa_flash_f16cache_contract(), ck_attention_f16_prefill_qtile64_dispatch(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT, CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES, CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64, CK_ATTN_REDUCTION_F16_ONLINE_FP32_MERGE, CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE, CK_ATTN_REDUCTION_FP32_ONLINE, and CK_GGML_FA_TILE_Q.

Referenced by attention_forward_causal_head_major_gqa_prefill_append_f16cache_auto_workspace(), and attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract().

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_gemma4_workspace()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_gemma4_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes
	)

Definition at line 7962 of file attention_kernels.c.

{
    return ck_attention_prefill_append_f16cache_gemma4_workspace(
        q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
        past_tokens, cache_capacity, head_dim, aligned_head_dim, 0, reduction,
        token_workspace, token_workspace_bytes, NULL, 0);
}

References ck_attention_prefill_append_f16cache_gemma4_workspace().

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		int	query_tile_size,
		int	concurrent_query_tiles,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 7585 of file attention_kernels.c.

{
    if (!q || !k_cache || !v_cache || !output || !workspace || num_heads <= 0 ||
        num_kv_heads <= 0 || num_heads % num_kv_heads != 0 ||
        query_tile_size <= 0 || q_tokens < query_tile_size || past_tokens < 0 ||
        past_tokens + q_tokens > cache_capacity || head_dim <= 0 ||
        aligned_head_dim < head_dim || query_tile_size < 16 ||
        query_tile_size > 128 || query_tile_size % 16 != 0 ||
        concurrent_query_tiles <= 0 || concurrent_query_tiles > 4) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    const int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (!pool || active < num_kv_heads || active % num_kv_heads != 0 ||
        ck_threadpool_thread_id(pool) > 0) {
        return CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT;
    }
    const size_t required_workspace =
        attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes(
            num_heads, num_kv_heads, head_dim, active,
            query_tile_size, concurrent_query_tiles);
    if (required_workspace == 0 || workspace_bytes < required_workspace) {
        return CK_ATTENTION_STATUS_INSUFFICIENT_WORKSPACE;
    }
    const int lanes = active / num_kv_heads;
    const size_t tile_elems =
        (size_t) num_kv_heads * (size_t) head_dim * CK_GGML_FA_TILE_KV;
    uintptr_t workspace_address = (uintptr_t) workspace;
    workspace_address = (workspace_address + 63u) & ~(uintptr_t) 63u;
    unsigned char *workspace_base = (unsigned char *) workspace_address;
    size_t cursor = 0;
    float *shared_k_tiles = (float *) (workspace_base + cursor);
    cursor += tile_elems * sizeof(float);
    cursor = ck_attention_align64_size(cursor);
    float *shared_v_tiles = (float *) (workspace_base + cursor);
    cursor += tile_elems * sizeof(float);
    cursor = ck_attention_align64_size(cursor);
    ck_attention_gqa_team_barrier_t *barriers =
        (ck_attention_gqa_team_barrier_t *) (workspace_base + cursor);
    cursor += (size_t) num_kv_heads * sizeof(*barriers);
    cursor = ck_attention_align64_size(cursor);
    unsigned char *worker_workspace = workspace_base + cursor;
    const size_t worker_workspace_stride =
        ck_attention_f16_prefill_gqa_reuse_worker_bytes(
            num_heads, num_kv_heads, head_dim, active,
            query_tile_size, concurrent_query_tiles);
    for (int group = 0; group < num_kv_heads; ++group) {
        __atomic_store_n(&barriers[group].arrived, 0, __ATOMIC_RELAXED);
        __atomic_store_n(&barriers[group].phase, 0, __ATOMIC_RELAXED);
        barriers[group].lanes = lanes;
    }
 
    ck_attention_f16_prefill_gqa_reuse_args_t args = {
        .base = {
            .q = q,
            .k_cache = k_cache,
            .v_cache = v_cache,
            .output = output,
            .num_heads = num_heads,
            .num_kv_heads = num_kv_heads,
            .q_tokens = q_tokens,
            .past_tokens = past_tokens,
            .cache_capacity = cache_capacity,
            .head_dim = head_dim,
            .aligned_head_dim = aligned_head_dim,
            .cache_is_bf16 = 0,
            .schedule = CK_ATTN_PREFILL_SCHEDULE_GQA_SHARED_KV_TILES,
            .q_head_stride =
                (size_t) q_tokens * (size_t) aligned_head_dim,
            .output_head_stride =
                (size_t) q_tokens * (size_t) aligned_head_dim,
        },
        .shared_k_tiles = shared_k_tiles,
        .shared_v_tiles = shared_v_tiles,
        .barriers = barriers,
        .worker_workspace = worker_workspace,
        .worker_workspace_stride = worker_workspace_stride,
        .query_tile_size = query_tile_size,
        .concurrent_query_tiles = concurrent_query_tiles,
    };
    ck_threadpool_dispatch_n(
        pool, active, ck_attention_f16_prefill_gqa_reuse_work, &args);
    return CK_ATTENTION_STATUS_OK;
}

References attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes(), ck_attention_align64_size(), ck_attention_f16_prefill_gqa_reuse_work(), ck_attention_f16_prefill_gqa_reuse_worker_bytes(), CK_ATTENTION_STATUS_INSUFFICIENT_WORKSPACE, CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT, CK_ATTN_PREFILL_SCHEDULE_GQA_SHARED_KV_TILES, CK_GGML_FA_TILE_KV, ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), and ck_threadpool_thread_id().

Referenced by attention_forward_causal_head_major_gqa_prefill_append_f16cache_auto_workspace().

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes()

size_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes	(	int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	workers,
		int	query_tile_size,
		int	concurrent_query_tiles
	)

Definition at line 7064 of file attention_kernels.c.

{
    const size_t worker_bytes =
        ck_attention_f16_prefill_gqa_reuse_worker_bytes(
            num_heads, num_kv_heads, head_dim, workers,
            query_tile_size, concurrent_query_tiles);
    if (worker_bytes == 0) return 0;
    size_t bytes = 63u;
    bytes = ck_attention_align64_size(bytes) +
        (size_t) num_kv_heads * (size_t) head_dim * CK_GGML_FA_TILE_KV *
        sizeof(float);
    bytes = ck_attention_align64_size(bytes) +
        (size_t) num_kv_heads * CK_GGML_FA_TILE_KV * (size_t) head_dim *
        sizeof(float);
    bytes = ck_attention_align64_size(bytes) +
        (size_t) num_kv_heads * sizeof(ck_attention_gqa_team_barrier_t);
    bytes = ck_attention_align64_size(bytes) +
        (size_t) workers * worker_bytes;
    return bytes;
}

References ck_attention_align64_size(), ck_attention_f16_prefill_gqa_reuse_worker_bytes(), and CK_GGML_FA_TILE_KV.

Referenced by attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config().

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_qtile64_schedule()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_qtile64_schedule	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_prefill_schedule_t	schedule
	)

Definition at line 7743 of file attention_kernels.c.

{
    if (!q || !k_cache || !v_cache || !output || num_heads <= 0 ||
        num_kv_heads <= 0 || q_tokens < CK_GGML_FA_TILE_Q || past_tokens < 0 ||
        past_tokens + q_tokens > cache_capacity || head_dim <= 0 ||
        aligned_head_dim < head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    return ck_attention_f16_prefill_qtile64_dispatch(
        q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
        past_tokens, cache_capacity, head_dim, aligned_head_dim, 0, 0.0f,
        (size_t) q_tokens * (size_t) aligned_head_dim,
        (size_t) q_tokens * (size_t) aligned_head_dim, schedule);
}

References ck_attention_f16_prefill_qtile64_dispatch(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, and CK_GGML_FA_TILE_Q.

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_sliding_gemma4_workspace()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_append_f16cache_sliding_gemma4_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		int	sliding_window,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes
	)

Definition at line 7975 of file attention_kernels.c.

{
    return ck_attention_prefill_append_f16cache_gemma4_workspace(
        q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
        past_tokens, cache_capacity, head_dim, aligned_head_dim, sliding_window,
        reduction, token_workspace, token_workspace_bytes, NULL, 0);
}

References ck_attention_prefill_append_f16cache_gemma4_workspace().

◆ attention_forward_causal_head_major_gqa_prefill_full_bf16cache_pytorch_contract()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_full_bf16cache_pytorch_contract	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction
	)

Definition at line 9469 of file attention_kernels.c.

{
    if (!q || !k_cache || !v_cache || !output ||
        num_heads <= 0 || num_kv_heads <= 0 ||
        num_heads % num_kv_heads != 0 ||
        q_tokens <= 0 || past_tokens < 0 ||
        past_tokens + q_tokens > cache_capacity ||
        head_dim <= 0 || aligned_head_dim < head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    if (reduction != CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA) {
        return CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT;
    }
 
#if defined(__AVX512F__)
    pthread_once(&ck_pytorch_attention_once, ck_bind_pytorch_attention_primitives);
    return ck_attention_prefill_bf16_pytorch_math_gqa_full(
        q, k_cache, v_cache, output,
        num_heads, num_kv_heads, q_tokens, past_tokens, cache_capacity,
        head_dim, aligned_head_dim);
#else
    return CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT;
#endif
}

References CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT, and CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA.

◆ attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_contract_workspace()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_contract_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes,
		const int *	segment_lengths,
		int	num_segments
	)

Definition at line 8130 of file attention_kernels.c.

{
    return ck_attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_schedule_workspace(
        q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
        past_tokens, cache_capacity, head_dim, aligned_head_dim, reduction,
        token_workspace, token_workspace_bytes, segment_lengths, num_segments,
        CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES);
}

References ck_attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_schedule_workspace(), and CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES.

◆ attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_gemma4_workspace()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_gemma4_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes,
		const int *	segment_lengths,
		int	num_segments
	)

Definition at line 8003 of file attention_kernels.c.

{
    return ck_attention_prefill_append_f16cache_gemma4_workspace(
        q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
        past_tokens, cache_capacity, head_dim, aligned_head_dim, 0,
        reduction, token_workspace, token_workspace_bytes,
        segment_lengths, num_segments);
}

References ck_attention_prefill_append_f16cache_gemma4_workspace().

◆ attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_sliding_gemma4_workspace()

ck_attention_status_t attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_sliding_gemma4_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		int	sliding_window,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes,
		const int *	segment_lengths,
		int	num_segments
	)

Definition at line 7988 of file attention_kernels.c.

{
    return ck_attention_prefill_append_f16cache_gemma4_workspace(
        q, k_cache, v_cache, output, num_heads, num_kv_heads, q_tokens,
        past_tokens, cache_capacity, head_dim, aligned_head_dim, sliding_window,
        reduction, token_workspace, token_workspace_bytes,
        segment_lengths, num_segments);
}

References ck_attention_prefill_append_f16cache_gemma4_workspace().

◆ attention_forward_causal_head_major_shared_kv_gemma4()

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

Definition at line 5144 of file attention_kernels.c.

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

References attention_forward_causal_head_major_gqa_flash_strided_gemma4().

◆ attention_forward_chunk_head_major_gqa_flash_gemma4()

void attention_forward_chunk_head_major_gqa_flash_gemma4	(	const float *	q_chunk,
		const float *	k_cache,
		const float *	v_cache,
		float *	out_chunk,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	kv_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim
	)

Definition at line 6299 of file attention_kernels.c.

{
    if (!q_chunk || !k_cache || !v_cache || !out_chunk) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || q_tokens <= 0 || kv_tokens <= 0 || cache_capacity <= 0) {
        return;
    }
    if (kv_tokens > cache_capacity || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    const float scale = 1.0f;
    const size_t cache_head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
    const size_t q_head_stride = (size_t)q_tokens * (size_t)aligned_head_dim;
 
    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_cache + (size_t)kv_head * cache_head_stride;
        const float *v_head = v_cache + (size_t)kv_head * cache_head_stride;
        for (int t = 0; t < q_tokens; ++t) {
            const float *q_head = q_chunk + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
            float *out_head = out_chunk + (size_t)h * q_head_stride + (size_t)t * (size_t)aligned_head_dim;
            attention_flash_decode(out_head,
                                   q_head,
                                   k_head,
                                   v_head,
                                   1,
                                   kv_tokens,
                                   1,
                                   aligned_head_dim,
                                   scale);
        }
    }
}

References attention_flash_decode().

◆ attention_forward_decode_head_major_gqa_bf16cache_pytorch_contract()

ck_attention_status_t attention_forward_decode_head_major_gqa_bf16cache_pytorch_contract	(	const float *	q_token,
		const uint16_t *	k_cache,
		const uint16_t *	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,
		ck_attention_reduction_t	reduction
	)

Definition at line 9278 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token ||
        num_heads <= 0 || num_kv_heads <= 0 || kv_tokens <= 0 ||
        cache_capacity <= 0 || kv_tokens > cache_capacity ||
        head_dim <= 0 || aligned_head_dim < head_dim ||
        num_heads % num_kv_heads != 0) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    if (reduction != CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA) {
        return CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT;
    }
 
#if defined(__AVX512F__)
    pthread_once(&ck_pytorch_attention_once, ck_bind_pytorch_attention_primitives);
    return ck_attention_decode_bf16_pytorch_math_gqa(
        q_token, k_cache, v_cache, out_token, num_heads, num_kv_heads,
        kv_tokens, cache_capacity, head_dim, aligned_head_dim);
#else
    const float scale = 1.0f / sqrtf((float)head_dim);
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
    for (int h = 0; h < num_heads; ++h) {
        const 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 uint16_t *k_head = k_cache + (size_t)kv_head * head_stride;
        const uint16_t *v_head = v_cache + (size_t)kv_head * head_stride;
        float *out_head = out_token + (size_t)h * (size_t)aligned_head_dim;
 
        float maximum = -INFINITY;
        for (int token = 0; token < kv_tokens; ++token) {
            const uint16_t *k_row = k_head + (size_t)token * (size_t)aligned_head_dim;
            float dot = 0.0f;
            for (int d = 0; d < head_dim; ++d) {
                dot = fmaf(q_head[d], bf16_to_float(k_row[d]), dot);
            }
            const float score = dot * scale;
            if (score > maximum) maximum = score;
        }
 
        for (int d = 0; d < aligned_head_dim; ++d) out_head[d] = 0.0f;
        float denominator = 0.0f;
        for (int token = 0; token < kv_tokens; ++token) {
            const uint16_t *k_row = k_head + (size_t)token * (size_t)aligned_head_dim;
            const uint16_t *v_row = v_head + (size_t)token * (size_t)aligned_head_dim;
            float dot = 0.0f;
            for (int d = 0; d < head_dim; ++d) {
                dot = fmaf(q_head[d], bf16_to_float(k_row[d]), dot);
            }
            const float probability = expf(dot * scale - maximum);
            denominator += probability;
            for (int d = 0; d < head_dim; ++d) {
                out_head[d] = fmaf(probability, bf16_to_float(v_row[d]), out_head[d]);
            }
        }
        const float inverse = denominator > 0.0f ? 1.0f / denominator : 0.0f;
        for (int d = 0; d < head_dim; ++d) {
            out_head[d] = bf16_to_float(float_to_bf16(out_head[d] * inverse));
        }
    }
    return CK_ATTENTION_STATUS_OK;
#endif
}

References bf16_to_float(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT, CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA, float_to_bf16(), score, and token.

Referenced by attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace().

◆ attention_forward_decode_head_major_gqa_f16cache_gemma4_contract()

ck_attention_status_t attention_forward_decode_head_major_gqa_f16cache_gemma4_contract	(	const float *	q_token,
		const uint16_t *	k_cache,
		const uint16_t *	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,
		ck_attention_reduction_t	reduction
	)

Definition at line 6910 of file attention_kernels.c.

{
    return ck_attention_forward_decode_f16cache_gemma4(
        q_token, k_cache, v_cache, out_token, num_heads, num_kv_heads,
        kv_tokens, cache_capacity, head_dim, aligned_head_dim, 0, reduction);
}

References ck_attention_forward_decode_f16cache_gemma4().

◆ attention_forward_decode_head_major_gqa_f16cache_regular_noscale()

ck_attention_status_t attention_forward_decode_head_major_gqa_f16cache_regular_noscale	(	const float *	q_token,
		const uint16_t *	k_cache,
		const uint16_t *	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 6844 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token ||
        num_heads <= 0 || num_kv_heads <= 0 || num_heads % num_kv_heads != 0 ||
        kv_tokens <= 0 || cache_capacity <= 0 || kv_tokens > cache_capacity ||
        head_dim <= 0 || head_dim > 4096 || aligned_head_dim < head_dim ||
        sliding_window < 0 ||
        (size_t) cache_capacity > SIZE_MAX / (size_t) aligned_head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
 
    const size_t cache_head_stride = (size_t) cache_capacity * (size_t) aligned_head_dim;
    const int first_token = sliding_window > 0 && kv_tokens > sliding_window
        ? kv_tokens - sliding_window : 0;
    uint16_t *q_half = (uint16_t *) alloca((size_t) head_dim * sizeof(uint16_t));
    for (int h = 0; h < num_heads; ++h) {
        const int kv_head = h / (num_heads / num_kv_heads);
        const float *q_head = q_token + (size_t) h * (size_t) aligned_head_dim;
        const uint16_t *k_head = k_cache + (size_t) kv_head * cache_head_stride;
        const uint16_t *v_head = v_cache + (size_t) kv_head * cache_head_stride;
        float *out_head = out_token + (size_t) h * (size_t) aligned_head_dim;
        for (int d = 0; d < head_dim; ++d) q_half[d] = CK_FP32_TO_FP16(q_head[d]);
 
        float maximum = -INFINITY;
        for (int token = first_token; token < kv_tokens; ++token) {
            const float score = ck_attention_dot_f16_llama(
                q_half, k_head + (size_t) token * (size_t) aligned_head_dim, head_dim);
            if (score > maximum) maximum = score;
        }
        if (!isfinite(maximum)) return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
 
        double total = 0.0;
        for (int token = first_token; token < kv_tokens; ++token) {
            const float score = ck_attention_dot_f16_llama(
                q_half, k_head + (size_t) token * (size_t) aligned_head_dim, head_dim);
            total += (double) ck_attention_reference_expf(score - maximum);
        }
        if (!(total > 0.0) || !isfinite(total)) return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
 
        for (int d = 0; d < head_dim; ++d) out_head[d] = 0.0f;
        for (int token = first_token; token < kv_tokens; ++token) {
            const float score = ck_attention_dot_f16_llama(
                q_half, k_head + (size_t) token * (size_t) aligned_head_dim, head_dim);
            const float probability = ck_attention_reference_expf(score - maximum) / (float) total;
            const float rounded_probability = CK_FP16_TO_FP32(CK_FP32_TO_FP16(probability));
            const uint16_t *v_row = v_head + (size_t) token * (size_t) aligned_head_dim;
            for (int d = 0; d < head_dim; ++d) {
                out_head[d] = fmaf(rounded_probability, CK_FP16_TO_FP32(v_row[d]), out_head[d]);
            }
        }
        for (int d = head_dim; d < aligned_head_dim; ++d) out_head[d] = 0.0f;
    }
    return CK_ATTENTION_STATUS_OK;
}

References ck_attention_dot_f16_llama(), ck_attention_reference_expf(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_FP16_TO_FP32, CK_FP32_TO_FP16, score, and token.

Referenced by ck_attention_forward_decode_f16cache_gemma4().

◆ attention_forward_decode_head_major_gqa_f16cache_sliding_gemma4_contract()

ck_attention_status_t attention_forward_decode_head_major_gqa_f16cache_sliding_gemma4_contract	(	const float *	q_token,
		const uint16_t *	k_cache,
		const uint16_t *	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,
		ck_attention_reduction_t	reduction
	)

Definition at line 6921 of file attention_kernels.c.

{
    return ck_attention_forward_decode_f16cache_gemma4(
        q_token, k_cache, v_cache, out_token, num_heads, num_kv_heads,
        kv_tokens, cache_capacity, head_dim, aligned_head_dim,
        sliding_window, reduction);
}

References ck_attention_forward_decode_f16cache_gemma4().

◆ attention_forward_decode_head_major_gqa_flash()

void attention_forward_decode_head_major_gqa_flash	(	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
	)

Flash attention decode (single token attends to KV cache)

Test:

test_flash_attention.py::TestFlashAttention::test_flash_decode

test_kv_cache_attention.py::TestKVCacheAttention::test_flash_decode

test_fused_attention_decode.py::TestFusedAttentionDecode::test_flash_decode

test_attention.py::TestAttentionForward::test_flash_decode

Single query token attends to kv_tokens in KV cache. Uses true flash attention from attention_flash_true.c.

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

Definition at line 6212 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || kv_tokens <= 0 || cache_capacity <= 0) {
        return;
    }
    if (kv_tokens > cache_capacity || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    const float scale = 1.0f / sqrtf((float)head_dim);
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
 
    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;
        const float *v_head = v_cache + (size_t)kv_head * head_stride;
        float *out_head = out_token + (size_t)h * (size_t)aligned_head_dim;
 
        attention_flash_decode(out_head,
                               q_head,
                               k_head,
                               v_head,
                               1,
                               kv_tokens,
                               1,
                               aligned_head_dim,
                               scale);
    }
}

References attention_flash_decode().

◆ attention_forward_decode_head_major_gqa_flash_f16cache()

void attention_forward_decode_head_major_gqa_flash_f16cache	(	const float *	q_token,
		const uint16_t *	k_cache,
		const uint16_t *	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
	)

Definition at line 9506 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || kv_tokens <= 0 || cache_capacity <= 0) {
        return;
    }
    if (kv_tokens > cache_capacity || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    const float scale = 1.0f / sqrtf((float)head_dim);
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
    const size_t scratch_elems = (size_t)kv_tokens * (size_t)aligned_head_dim;
    const size_t scratch_bytes = scratch_elems * sizeof(float);
    const size_t max_stack_bytes = 1024u * 1024u;
 
    if (scratch_bytes * 2u > max_stack_bytes) {
        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 uint16_t *k_head = k_cache + (size_t)kv_head * head_stride;
            const uint16_t *v_head = v_cache + (size_t)kv_head * head_stride;
            float *out_head = out_token + (size_t)h * (size_t)aligned_head_dim;
 
            for (int d = 0; d < aligned_head_dim; ++d) {
                out_head[d] = 0.0f;
            }
 
            float max_score = -INFINITY;
            for (int j = 0; j < kv_tokens; ++j) {
                const uint16_t *k_vec = k_head + (size_t)j * (size_t)aligned_head_dim;
                float dot = 0.0f;
                for (int d = 0; d < head_dim; ++d) {
                    dot += q_head[d] * CK_FP16_TO_FP32(k_vec[d]);
                }
                const float score = dot * scale;
                if (score > max_score) {
                    max_score = score;
                }
            }
 
            float sum = 0.0f;
            for (int j = 0; j < kv_tokens; ++j) {
                const uint16_t *k_vec = k_head + (size_t)j * (size_t)aligned_head_dim;
                const uint16_t *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_head[d] * CK_FP16_TO_FP32(k_vec[d]);
                }
                const float w = expf(dot * scale - max_score);
                sum += w;
                for (int d = 0; d < head_dim; ++d) {
                    out_head[d] += w * CK_FP16_TO_FP32(v_vec[d]);
                }
            }
 
            if (sum > 0.0f) {
                const float inv_sum = 1.0f / sum;
                for (int d = 0; d < head_dim; ++d) {
                    out_head[d] *= inv_sum;
                }
            }
            for (int d = head_dim; d < aligned_head_dim; ++d) {
                out_head[d] = 0.0f;
            }
        }
        return;
    }
 
    float *k_head_fp32 = (float *)alloca(scratch_bytes);
    float *v_head_fp32 = (float *)alloca(scratch_bytes);
 
    for (int kv_head = 0; kv_head < num_kv_heads; ++kv_head) {
        const uint16_t *k_head = k_cache + (size_t)kv_head * head_stride;
        const uint16_t *v_head = v_cache + (size_t)kv_head * head_stride;
        const int q_begin = (int)((long long)kv_head * (long long)num_heads / (long long)num_kv_heads);
        const int q_end = (int)((long long)(kv_head + 1) * (long long)num_heads / (long long)num_kv_heads);
 
        ck_local_fp16_to_fp32_2d(k_head, k_head_fp32, kv_tokens, aligned_head_dim, aligned_head_dim, aligned_head_dim);
        ck_local_fp16_to_fp32_2d(v_head, v_head_fp32, kv_tokens, aligned_head_dim, aligned_head_dim, aligned_head_dim);
 
        for (int h = q_begin; h < q_end; ++h) {
            const float *q_head = q_token + (size_t)h * (size_t)aligned_head_dim;
            float *out_head = out_token + (size_t)h * (size_t)aligned_head_dim;
 
            attention_flash_decode(out_head,
                                   q_head,
                                   k_head_fp32,
                                   v_head_fp32,
                                   1,
                                   kv_tokens,
                                   1,
                                   aligned_head_dim,
                                   scale);
        }
    }
}

References attention_flash_decode(), CK_FP16_TO_FP32, ck_local_fp16_to_fp32_2d(), and score.

Referenced by attention_forward_decode_head_major_gqa_flash_f16cache_contract().

◆ attention_forward_decode_head_major_gqa_flash_f16cache_contract()

ck_attention_status_t attention_forward_decode_head_major_gqa_flash_f16cache_contract	(	const float *	q_token,
		const uint16_t *	k_cache,
		const uint16_t *	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,
		ck_attention_reduction_t	reduction
	)

Definition at line 6933 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token ||
        num_heads <= 0 || num_kv_heads <= 0 || kv_tokens <= 0 ||
        cache_capacity <= 0 || kv_tokens > cache_capacity ||
        head_dim <= 0 || aligned_head_dim < head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
 
    switch (reduction) {
    case CK_ATTN_REDUCTION_FP32_ONLINE:
        attention_forward_decode_head_major_gqa_flash_f16cache(
            q_token, k_cache, v_cache, out_token,
            num_heads, num_kv_heads, kv_tokens, cache_capacity,
            head_dim, aligned_head_dim);
        return CK_ATTENTION_STATUS_OK;
 
    case CK_ATTN_REDUCTION_F16_ONLINE_FP32_MERGE: {
        const int partition_alignment = 256;
        const int partition_tokens =
            ((kv_tokens + partition_alignment - 1) / partition_alignment) * partition_alignment;
        const int split_chunks = partition_tokens >= 512 ? ck_get_num_threads() : 1;
        attention_forward_decode_head_major_gqa_flash_f16cache_split(
            q_token, k_cache, v_cache, out_token,
            num_heads, num_kv_heads, kv_tokens, cache_capacity,
            head_dim, aligned_head_dim, split_chunks);
        return CK_ATTENTION_STATUS_OK;
    }
 
    case CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE:
        attention_forward_decode_head_major_gqa_flash_f16cache_split(
            q_token, k_cache, v_cache, out_token,
            num_heads, num_kv_heads, kv_tokens, cache_capacity,
            head_dim, aligned_head_dim, 1);
        return CK_ATTENTION_STATUS_OK;
 
    default:
        return CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT;
    }
}

References attention_forward_decode_head_major_gqa_flash_f16cache(), attention_forward_decode_head_major_gqa_flash_f16cache_split(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT, CK_ATTN_REDUCTION_F16_ONLINE_FP32_MERGE, CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE, CK_ATTN_REDUCTION_FP32_ONLINE, and ck_get_num_threads().

Referenced by attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace(), and ck_attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_schedule_workspace().

◆ attention_forward_decode_head_major_gqa_flash_f16cache_split()

void attention_forward_decode_head_major_gqa_flash_f16cache_split	(	const float *	q_token,
		const uint16_t *	k_cache,
		const uint16_t *	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	split_chunks
	)

Definition at line 6784 of file attention_kernels.c.

{
    const int partition_alignment = 256;
    const int partition_tokens =
        ((kv_tokens + partition_alignment - 1) / partition_alignment) * partition_alignment;
    attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned(
        q_token, k_cache, v_cache, out_token,
        num_heads, num_kv_heads, kv_tokens, cache_capacity,
        head_dim, aligned_head_dim, split_chunks, partition_tokens, 0,
        ck_attention_strict_scale_f32(head_dim));
}

References attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned(), and ck_attention_strict_scale_f32().

Referenced by attention_forward_decode_head_major_gqa_flash_f16cache_contract().

◆ attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned()

static void attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned	(	const float *	q_token,
		const uint16_t *	k_cache,
		const uint16_t *	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	split_chunks,
		int	partition_tokens,
		int	kv_start,
		float	attention_scale
	)

static

Definition at line 6662 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token ||
        num_heads <= 0 || num_kv_heads <= 0 || kv_tokens <= 0 ||
        cache_capacity <= 0 || kv_tokens > cache_capacity ||
        kv_start < 0 || kv_start >= kv_tokens ||
        head_dim <= 0 || aligned_head_dim < head_dim ||
        partition_tokens < kv_tokens - kv_start || !isfinite(attention_scale)) {
        return;
    }
 
    const int live_tokens = kv_tokens - kv_start;
 
    if (split_chunks < 1) {
        split_chunks = 1;
    }
    if (split_chunks > live_tokens) {
        split_chunks = live_tokens;
    }
 
    const int partial_stride = aligned_head_dim + 2;
    const size_t max_partial_bytes = 1024u * 1024u;
    while (split_chunks > 1 &&
           (size_t) num_heads * (size_t) split_chunks * (size_t) partial_stride * sizeof(float) >
               max_partial_bytes) {
        split_chunks = (split_chunks + 1) / 2;
    }
 
    const int chunk_size = (partition_tokens + split_chunks - 1) / split_chunks;
    const int scheduled_chunks = (live_tokens + chunk_size - 1) / chunk_size;
    const size_t resolved_count =
        (size_t) num_heads * (size_t) scheduled_chunks * (size_t) partial_stride;
    float *partials = (float *) alloca(resolved_count * sizeof(float));
    ck_attention_f16_split_args_t args = {
        .q_token = q_token,
        .k_cache = k_cache,
        .v_cache = v_cache,
        .partials = partials,
        .num_heads = num_heads,
        .num_kv_heads = num_kv_heads,
        .kv_start = kv_start,
        .kv_tokens = kv_tokens,
        .cache_capacity = cache_capacity,
        .head_dim = head_dim,
        .aligned_head_dim = aligned_head_dim,
        .split_chunks = split_chunks,
        .scheduled_chunks = scheduled_chunks,
        .partition_tokens = partition_tokens,
        .attention_scale = attention_scale,
    };
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int active_threads = pool ? ck_threadpool_n_threads(pool) : 1;
    const int total_jobs = num_heads * scheduled_chunks;
    if (active_threads > total_jobs) {
        active_threads = total_jobs;
    }
    if (pool && active_threads > 1) {
        ck_threadpool_dispatch_n(pool, active_threads, ck_attention_f16_split_work, &args);
    } else {
        ck_attention_f16_split_work(0, 1, &args);
    }
 
    for (int h = 0; h < num_heads; ++h) {
        float *out_head = out_token + (size_t) h * (size_t) aligned_head_dim;
        float final_max = -INFINITY;
        float final_sum = 0.0f;
        for (int d = 0; d < aligned_head_dim; ++d) {
            out_head[d] = 0.0f;
        }
 
        for (int chunk = 0; chunk < scheduled_chunks; ++chunk) {
            const float *partial = partials +
                ((size_t) h * (size_t) scheduled_chunks + (size_t) chunk) * (size_t) partial_stride;
            const float chunk_max = partial[0];
            const float chunk_sum = partial[1];
            if (chunk_sum == 0.0f) {
                continue;
            }
            const float new_max = fmaxf(final_max, chunk_max);
            const float old_scale = isfinite(final_max)
                ? ck_attention_f16_reduce_expf(final_max - new_max)
                : 0.0f;
            const float chunk_scale =
                ck_attention_f16_reduce_expf(chunk_max - new_max);
            for (int d = 0; d < head_dim; ++d) {
                const float chunk_term = partial[2 + d] * chunk_scale;
                out_head[d] = fmaf(out_head[d], old_scale, chunk_term);
            }
#if defined(__INTEL_LLVM_COMPILER)
            const float scaled_chunk_sum = chunk_sum * chunk_scale;
            final_sum = fmaf(final_sum, old_scale, scaled_chunk_sum);
#else
            final_sum = fmaf(
                final_sum, old_scale, chunk_sum * chunk_scale);
#endif
            final_max = new_max;
        }
 
        if (final_sum > 0.0f) {
            const float inv_sum = 1.0f / final_sum;
            for (int d = 0; d < head_dim; ++d) {
                out_head[d] *= inv_sum;
            }
        }
    }
}

References ck_attention_f16_reduce_expf(), ck_attention_f16_split_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by attention_forward_decode_head_major_gqa_flash_f16cache_split(), and ck_attention_forward_decode_f16cache_gemma4().

◆ attention_forward_decode_head_major_gqa_flash_f16kv()

void attention_forward_decode_head_major_gqa_flash_f16kv	(	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
	)

Definition at line 6361 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || kv_tokens <= 0 || cache_capacity <= 0) {
        return;
    }
    if (kv_tokens > cache_capacity || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    const float scale = 1.0f / sqrtf((float)head_dim);
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
 
    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;
        const float *v_head = v_cache + (size_t)kv_head * head_stride;
        float *out_head = out_token + (size_t)h * (size_t)aligned_head_dim;
 
        attention_flash_query_causal_exact_f16kv(q_head,
                                                 k_head,
                                                 v_head,
                                                 kv_tokens,
                                                 head_dim,
                                                 aligned_head_dim,
                                                 scale,
                                                 out_head);
    }
}

References attention_flash_query_causal_exact_f16kv().

◆ attention_forward_decode_head_major_gqa_flash_gemma4()

void attention_forward_decode_head_major_gqa_flash_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
	)

Definition at line 6256 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || kv_tokens <= 0 || cache_capacity <= 0) {
        return;
    }
    if (kv_tokens > cache_capacity || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    const float scale = 1.0f;
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
 
    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;
        const float *v_head = v_cache + (size_t)kv_head * head_stride;
        float *out_head = out_token + (size_t)h * (size_t)aligned_head_dim;
 
        attention_flash_decode(out_head,
                               q_head,
                               k_head,
                               v_head,
                               1,
                               kv_tokens,
                               1,
                               aligned_head_dim,
                               scale);
    }
}

References attention_flash_decode().

Referenced by attention_forward_decode_head_major_gqa_flash_sliding_gemma4(), and attention_forward_decode_head_major_shared_kv_gemma4().

◆ attention_forward_decode_head_major_gqa_llama_regular_sliding_workspace()

void attention_forward_decode_head_major_gqa_llama_regular_sliding_workspace	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	live_tokens,
		int	kv_stride_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	sliding_window,
		float *	scores,
		size_t	scores_bytes,
		float *	value_columns,
		size_t	value_columns_bytes,
		float *	scaled_scores,
		size_t	scaled_scores_bytes
	)

Definition at line 6185 of file attention_kernels.c.

{
    ck_attention_llama_regular_impl(
        q, k, v, output, num_heads, num_kv_heads, 1, live_tokens,
        head_dim, aligned_head_dim, kv_stride_tokens, sliding_window,
        scores, scores_bytes, value_columns, value_columns_bytes,
        scaled_scores, scaled_scores_bytes, 0);
}

References ck_attention_llama_regular_impl().

◆ attention_forward_decode_head_major_gqa_muse_eager_bf16_storage()

void attention_forward_decode_head_major_gqa_muse_eager_bf16_storage	(	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,
		float *	scores,
		size_t	scores_bytes,
		uint16_t *	key_rows,
		size_t	key_rows_bytes,
		uint16_t *	value_columns,
		size_t	value_columns_bytes,
		uint16_t *	gemm_rows,
		size_t	gemm_rows_bytes,
		uint16_t *	gemm_scores,
		size_t	gemm_scores_bytes
	)

Definition at line 9242 of file attention_kernels.c.

{
    ck_attention_muse_eager_bf16_storage_impl(
        q_token, k_cache, v_cache, out_token, num_heads, num_kv_heads,
        1, kv_tokens, cache_capacity, head_dim, aligned_head_dim, 0,
        scores, scores_bytes, key_rows, key_rows_bytes,
        value_columns, value_columns_bytes, gemm_rows, gemm_rows_bytes,
        gemm_scores, gemm_scores_bytes);
}

References ck_attention_muse_eager_bf16_storage_impl().

◆ attention_forward_decode_head_major_gqa_muse_eager_bf16_storage_sliding()

void attention_forward_decode_head_major_gqa_muse_eager_bf16_storage_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,
		float *	scores,
		size_t	scores_bytes,
		uint16_t *	key_rows,
		size_t	key_rows_bytes,
		uint16_t *	value_columns,
		size_t	value_columns_bytes,
		uint16_t *	gemm_rows,
		size_t	gemm_rows_bytes,
		uint16_t *	gemm_scores,
		size_t	gemm_scores_bytes
	)

Definition at line 9260 of file attention_kernels.c.

{
    ck_attention_muse_eager_bf16_storage_impl(
        q_token, k_cache, v_cache, out_token, num_heads, num_kv_heads,
        1, kv_tokens, cache_capacity, head_dim, aligned_head_dim,
        sliding_window, scores, scores_bytes, key_rows, key_rows_bytes,
        value_columns, value_columns_bytes, gemm_rows, gemm_rows_bytes,
        gemm_scores, gemm_scores_bytes);
}

References ck_attention_muse_eager_bf16_storage_impl().

◆ attention_forward_decode_head_major_gqa_regular()

void attention_forward_decode_head_major_gqa_regular	(	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
	)

WARNING: This is NOT true flash attention!

This function is named "flash" but implements regular attention with O(n) complexity. It's kept for reference and as a fallback.

TRUE flash attention is implemented in attention_flash_true.c

Test:

test_kv_cache_attention.py::TestKVCacheAttention::test_regular_decode

test_attention.py::TestAttentionForward::test_regular_decode

Regular attention decode (score-matrix version) for fallback.

After changes: make test

Definition at line 9629 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || kv_tokens <= 0 || cache_capacity <= 0) {
        return;
    }
    if (kv_tokens > cache_capacity) {
        return;
    }
 
    const int strict = ck_strict_parity_enabled();
    const float scale = strict
        ? ck_attention_strict_scale_f32(head_dim)
        : 1.0f / sqrtf((float) head_dim);
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
 
    // Select SIMD implementation based on compile-time CPU features
#if defined(__AVX512F__)
    #define FLASH_QUERY_IMPL_DECODE attention_flash_query_causal_avx512
#elif defined(__AVX2__)
    #define FLASH_QUERY_IMPL_DECODE attention_flash_query_causal_avx2
#elif defined(__AVX__)
    #define FLASH_QUERY_IMPL_DECODE attention_flash_query_causal_avx
#else
    #define FLASH_QUERY_IMPL_DECODE attention_flash_query_causal
#endif
 
#pragma omp parallel for schedule(static) if(num_heads > 1)
    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_vec = q_token + (size_t)h * (size_t)aligned_head_dim;
        const float *k_head = k_cache + (size_t)kv_head * head_stride;
        const float *v_head = v_cache + (size_t)kv_head * head_stride;
        float *out_vec = out_token + (size_t)h * (size_t)aligned_head_dim;
 
        if (strict) {
            attention_flash_query_causal_exact_f16kv(q_vec,
                                                     k_head,
                                                     v_head,
                                                     kv_tokens,
                                                     head_dim,
                                                     aligned_head_dim,
                                                     scale,
                                                     out_vec);
            continue;
        }
 
        FLASH_QUERY_IMPL_DECODE(q_vec, k_head, v_head,
                                 kv_tokens, head_dim, aligned_head_dim,
                                 scale, out_vec);
    }
 
#undef FLASH_QUERY_IMPL_DECODE
}

References attention_flash_query_causal_exact_f16kv(), ck_attention_strict_scale_f32(), ck_strict_parity_enabled(), and FLASH_QUERY_IMPL_DECODE.

◆ attention_forward_decode_head_major_shared_kv_gemma4()

void attention_forward_decode_head_major_shared_kv_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
	)

Definition at line 6345 of file attention_kernels.c.

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

References attention_forward_decode_head_major_gqa_flash_gemma4().

◆ attention_forward_full_head_major_gqa_exact_strided()

void attention_forward_full_head_major_gqa_exact_strided	(	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
	)

Definition at line 5446 of file attention_kernels.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;
    }
 
    const float scale = ck_strict_parity_enabled()
        ? ck_attention_strict_scale_f32(head_dim)
        : 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 debug_layer_id = ck_strict_parity_enabled() ? ck_attention_vec_dump_next_layer_id() : -1;
    float *score_row = (float *) alloca((size_t) T * sizeof(float));
    float *v_cols = (float *) alloca((size_t) head_dim * (size_t) T * sizeof(float));
    for (int h = 0; h < num_heads; ++h) {
        const 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;
 
#if CK_ENABLE_LLAMA_CPP_PARITY
        float *out_head = output + (size_t) h * (size_t) T * (size_t) aligned_head_dim;
        if (ck_strict_parity_enabled() &&
            ck_attention_head_full_ggml_graph_oracle_regular(
                q + (size_t) h * (size_t) T * (size_t) aligned_head_dim,
                k_head,
                v_head,
                out_head,
                T,
                head_dim,
                aligned_head_dim,
                scale)) {
            ck_attention_trace("regular_graph_oracle", debug_layer_id, h);
            continue;
        }
#endif
 
        for (int d = 0; d < head_dim; ++d) {
            float *dst_col = v_cols + (size_t) d * (size_t) T;
            for (int j = 0; j < T; ++j) {
                dst_col[j] = v_head[(size_t) j * (size_t) aligned_head_dim + (size_t) d];
            }
        }
 
        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 + qkv_index(h, i, 0, T, aligned_head_dim);
            attention_query_full_exact_regular(q_vec,
                                               k_head,
                                               v_cols,
                                               T,
                                               head_dim,
                                               aligned_head_dim,
                                               scale,
                                               score_row,
                                               out_vec,
                                               debug_layer_id,
                                               h,
                                               i);
        }
    }
}

References attention_query_full_exact_regular(), ck_attention_head_full_ggml_graph_oracle_regular(), ck_attention_strict_scale_f32(), ck_attention_trace(), ck_attention_vec_dump_next_layer_id(), ck_strict_parity_enabled(), and qkv_index().

◆ attention_forward_full_head_major_gqa_flash()

void attention_forward_full_head_major_gqa_flash	(	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
	)

Definition at line 4465 of file attention_kernels.c.

{
    attention_forward_head_major_gqa_flash_impl(q, k, v, output,
                                                num_heads, num_kv_heads,
                                                num_tokens, head_dim,
                                                aligned_head_dim,
                                                /*kv_stride_tokens=*/num_tokens,
                                                /*causal=*/0,
                                                /*round_full_kv_fp16=*/1,
                                                /*output_token_major=*/0,
                                                1.0f / sqrtf((float)head_dim));
}

References attention_forward_head_major_gqa_flash_impl().

◆ attention_forward_full_head_major_gqa_flash_strided()

void attention_forward_full_head_major_gqa_flash_strided	(	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
	)

Definition at line 4541 of file attention_kernels.c.

{
    attention_forward_head_major_gqa_flash_impl(q, k, v, output,
                                                num_heads, num_kv_heads,
                                                num_tokens, head_dim,
                                                aligned_head_dim,
                                                kv_stride_tokens,
                                                /*causal=*/0,
                                                /*round_full_kv_fp16=*/1,
                                                /*output_token_major=*/0,
                                                1.0f / sqrtf((float)head_dim));
}

References attention_forward_head_major_gqa_flash_impl().

◆ attention_forward_full_head_major_gqa_flash_strided_bf16_storage()

void attention_forward_full_head_major_gqa_flash_strided_bf16_storage	(	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
	)

Definition at line 4994 of file attention_kernels.c.

{
    attention_forward_head_major_gqa_flash_impl(
        q, k, v, output,
        num_heads, num_kv_heads, num_tokens, head_dim,
        aligned_head_dim, kv_stride_tokens,
        /*causal=*/0,
        /*round_full_kv_fp16=*/0,
        /*output_token_major=*/0,
        1.0f / sqrtf((float)head_dim)
    );
    const size_t count = (size_t)num_heads * (size_t)num_tokens
                       * (size_t)aligned_head_dim;
    for (size_t i = 0; i < count; ++i) {
        output[i] = bf16_to_float(float_to_bf16(output[i]));
    }
}

References attention_forward_head_major_gqa_flash_impl(), bf16_to_float(), count, and float_to_bf16().

◆ attention_forward_full_head_major_gqa_flash_strided_gemma4()

void attention_forward_full_head_major_gqa_flash_strided_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
	)

Definition at line 5158 of file attention_kernels.c.

{
    (void)head_dim;
    attention_forward_head_major_gqa_flash_impl(q, k, v, output,
                                                num_heads, num_kv_heads,
                                                num_tokens, head_dim,
                                                aligned_head_dim,
                                                kv_stride_tokens,
                                                /*causal=*/0,
                                                /*round_full_kv_fp16=*/1,
                                                /*output_token_major=*/0,
                                                1.0f);
}

References attention_forward_head_major_gqa_flash_impl().

◆ attention_forward_full_head_major_gqa_ggml_strided()

void attention_forward_full_head_major_gqa_ggml_strided	(	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
	)

Definition at line 5764 of file attention_kernels.c.

{
    if (num_heads <= 0 || num_tokens <= 0 || head_dim <= 0 ||
        (size_t) num_heads > SIZE_MAX / (size_t) num_tokens ||
        (size_t) head_dim > SIZE_MAX / (size_t) num_tokens) {
        return;
    }
    const size_t score_elements = (size_t) num_heads * (size_t) num_tokens;
    const size_t columns_per_head = (size_t) head_dim * (size_t) num_tokens;
    if ((size_t) num_heads > SIZE_MAX / columns_per_head) return;
    const size_t column_elements = (size_t) num_heads * columns_per_head;
    if (score_elements > SIZE_MAX - column_elements ||
        score_elements + column_elements > SIZE_MAX - (size_t) num_tokens) return;
    const size_t total_elements = score_elements + column_elements + (size_t) num_tokens;
    if (total_elements > SIZE_MAX / sizeof(float)) return;
    float *workspace = (float *) malloc(total_elements * sizeof(float));
    if (!workspace) return;
    float *score_rows = workspace;
    float *v_columns = score_rows + score_elements;
    float *probability_row = v_columns + column_elements;
    attention_forward_full_head_major_gqa_ggml_strided_workspace(
        q, k, v, output, num_heads, num_kv_heads, num_tokens, head_dim,
        aligned_head_dim, kv_stride_tokens,
        score_rows, score_elements * sizeof(float),
        v_columns, column_elements * sizeof(float),
        probability_row, (size_t) num_tokens * sizeof(float));
    free(workspace);
}

References attention_forward_full_head_major_gqa_ggml_strided_workspace().

◆ attention_forward_full_head_major_gqa_ggml_strided_workspace()

void attention_forward_full_head_major_gqa_ggml_strided_workspace	(	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,
		float *	score_rows,
		size_t	score_rows_bytes,
		float *	v_columns,
		size_t	v_columns_bytes,
		float *	probability_row,
		size_t	probability_row_bytes
	)

Definition at line 5523 of file attention_kernels.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;
    }
 
    const size_t T = (size_t) num_tokens;
    if ((size_t) num_heads > SIZE_MAX / T ||
        (size_t) head_dim > SIZE_MAX / T) {
        return;
    }
    const size_t score_elements = (size_t) num_heads * T;
    const size_t columns_per_head = (size_t) head_dim * T;
    if ((size_t) num_heads > SIZE_MAX / columns_per_head ||
        score_elements > SIZE_MAX / sizeof(float) ||
        columns_per_head * (size_t) num_heads > SIZE_MAX / sizeof(float) ||
        T > SIZE_MAX / sizeof(float)) {
        return;
    }
    const size_t column_elements = columns_per_head * (size_t) num_heads;
    if (!score_rows || score_rows_bytes < score_elements * sizeof(float) ||
        !v_columns || v_columns_bytes < column_elements * sizeof(float) ||
        !probability_row || probability_row_bytes < T * sizeof(float)) {
        return;
    }
 
    const int strict = ck_strict_parity_enabled();
    const float scale = strict
        ? ck_attention_strict_scale_f32(head_dim)
        : 1.0f / sqrtf((float) head_dim);
    const int token_count = num_tokens;
    const size_t kv_head_stride = (size_t) kv_stride_tokens * (size_t) aligned_head_dim;
    const int debug_layer_id = strict ? ck_attention_vec_dump_next_layer_id() : -1;
#if CK_ENABLE_LLAMA_CPP_PARITY
    ck_ggml_vec_dot_f32_fn dot_fn = NULL;
    ck_ggml_vec_soft_max_f32_fn softmax_fn = NULL;
    ck_ggml_compute_forward_mul_mat_fn mul_mat_fn = NULL;
    ck_ggml_compute_forward_soft_max_fn softmax_compute_fn = NULL;
    if (strict) {
        if (ck_attention_full_ggml_graph_oracle_multihead(q,
                                                          k,
                                                          v,
                                                          output,
                                                          num_heads,
                                                          num_kv_heads,
                                                          num_tokens,
                                                          head_dim,
                                                          aligned_head_dim,
                                                          kv_stride_tokens,
                                                          scale)) {
            return;
        }
        dot_fn = ck_resolve_ggml_vec_dot_f32();
        softmax_fn = ck_resolve_ggml_vec_soft_max_f32();
        mul_mat_fn = ck_resolve_ggml_compute_forward_mul_mat();
        softmax_compute_fn = ck_resolve_ggml_compute_forward_soft_max();
    }
#endif
    if (strict) {
        float *score_row = score_rows;
        float *v_cols = v_columns;
#if CK_ENABLE_LLAMA_CPP_PARITY
        float *prob_row = probability_row;
#endif
 
        for (int h = 0; h < num_heads; ++h) {
            const 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;
 
#if CK_ENABLE_LLAMA_CPP_PARITY
            float *out_head = output + (size_t) h * (size_t) T * (size_t) aligned_head_dim;
            if (ck_attention_head_full_ggml_graph_oracle_regular(
                    q + (size_t) h * (size_t) T * (size_t) aligned_head_dim,
                    k_head,
                    v_head,
                    out_head,
                    token_count,
                    head_dim,
                    aligned_head_dim,
                    scale)) {
                continue;
            }
#endif
 
            for (int d = 0; d < head_dim; ++d) {
                float *dst_col = v_cols + (size_t) d * (size_t) token_count;
                for (int j = 0; j < token_count; ++j) {
                    dst_col[j] = v_head[(size_t) j * (size_t) aligned_head_dim + (size_t) d];
                }
            }
 
#if CK_ENABLE_LLAMA_CPP_PARITY
            if (dot_fn && softmax_fn && ck_attention_ggml_out_graph_enabled()) {
                if (attention_head_full_dyn_ggml_regular_graph_out(
                        q + (size_t) h * (size_t) T * (size_t) aligned_head_dim,
                        k_head,
                        v_cols,
                        token_count,
                        head_dim,
                        aligned_head_dim,
                        scale,
                        score_row,
                        prob_row,
                        out_head,
                        dot_fn,
                        softmax_fn,
                        debug_layer_id,
                        h)) {
                    if (token_count > 0) {
                        ck_attention_trace("dyn_ggml_regular_graph_out", debug_layer_id, h);
                        ck_attention_trace_float("scale", debug_layer_id, h, scale);
                    }
                    continue;
                }
            }
#endif
 
            for (int i = 0; i < token_count; ++i) {
                const float *q_vec = q + qkv_index(h, i, 0, token_count, aligned_head_dim);
                float *out_vec = output + qkv_index(h, i, 0, token_count, aligned_head_dim);
#if CK_ENABLE_LLAMA_CPP_PARITY
                if (dot_fn && softmax_fn) {
                    if (i == 0) {
                        ck_attention_trace("dyn_ggml_regular", debug_layer_id, h);
                        ck_attention_trace_float("scale", debug_layer_id, h, scale);
                    }
                    attention_query_full_dyn_ggml_regular(q_vec,
                                                          k_head,
                                                          v_cols,
                                                          token_count,
                                                          head_dim,
                                                          aligned_head_dim,
                                                          scale,
                                                          score_row,
                                                          prob_row,
                                                          out_vec,
                                                          dot_fn,
                                                          softmax_fn,
                                                          debug_layer_id,
                                                          h,
                                                          i);
                } else if (dot_fn && mul_mat_fn && softmax_compute_fn) {
                    if (i == 0) {
                        ck_attention_trace("ggml_compute_regular", debug_layer_id, h);
                    }
                    attention_query_full_ggml_compute_regular(q_vec,
                                                              k_head,
                                                              v_cols,
                                                              token_count,
                                                              head_dim,
                                                              aligned_head_dim,
                                                              scale,
                                                              score_row,
                                                              prob_row,
                                                              out_vec,
                                                              dot_fn,
                                                              mul_mat_fn,
                                                              softmax_compute_fn);
                } else
#endif
                {
                    if (i == 0) {
                        ck_attention_trace("ggml_regular", debug_layer_id, h);
                    }
                    attention_query_full_ggml_regular(q_vec,
                                                      k_head,
                                                      v_cols,
                                                      token_count,
                                                      head_dim,
                                                      aligned_head_dim,
                                                      scale,
                                                      score_row,
                                                      out_vec,
                                                      debug_layer_id,
                                                      h,
                                                      i);
                }
            }
        }
        return;
    }
 
#pragma omp parallel for schedule(static) if(num_heads > 1)
    for (int h = 0; h < num_heads; ++h) {
        float *score_row = score_rows + (size_t) h * (size_t) token_count;
        float *v_cols = v_columns + (size_t) h * (size_t) head_dim * (size_t) token_count;
        const 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;
 
        if (v_cols) {
            for (int d = 0; d < head_dim; ++d) {
                float *dst_col = v_cols + (size_t) d * (size_t) token_count;
                for (int j = 0; j < token_count; ++j) {
                    dst_col[j] = v_head[(size_t) j * (size_t) aligned_head_dim + (size_t) d];
                }
            }
        }
 
        for (int i = 0; i < token_count; ++i) {
            const float *q_vec = q + qkv_index(h, i, 0, token_count, aligned_head_dim);
            float *out_vec = output + qkv_index(h, i, 0, token_count, aligned_head_dim);
            attention_query_full_ggml_regular(q_vec,
                                              k_head,
                                              v_cols,
                                              token_count,
                                              head_dim,
                                              aligned_head_dim,
                                              scale,
                                              score_row,
                                              out_vec,
                                              -1,
                                              h,
                                              i);
        }
    }
}

References attention_query_full_ggml_regular(), ck_attention_full_ggml_graph_oracle_multihead(), ck_attention_ggml_out_graph_enabled(), ck_attention_head_full_ggml_graph_oracle_regular(), ck_attention_strict_scale_f32(), ck_attention_trace(), ck_attention_trace_float(), ck_attention_vec_dump_next_layer_id(), ck_strict_parity_enabled(), and qkv_index().

Referenced by attention_forward_full_head_major_gqa_ggml_strided().

◆ attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage()

void attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage	(	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
	)

Definition at line 5047 of file attention_kernels.c.

{
    if (ck_attention_full_bf16_pytorch_flash(
            q, k, v, output, num_heads, num_kv_heads, num_tokens,
            head_dim, aligned_head_dim, kv_stride_tokens,
            /*output_token_major=*/0)) {
        return;
    }
    fprintf(stderr,
            "HARD KERNEL CONTRACT FAULT: PyTorch CPU-flash BF16 attention "
            "requires AMX-BF16, AVX-512, D=72/A=72 and a token multiple of 16; "
            "no numerically different fallback is permitted\n");
    abort();
}

References ck_attention_full_bf16_pytorch_flash().

◆ attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage_token_output()

void attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage_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
	)

Definition at line 5072 of file attention_kernels.c.

{
    if (ck_attention_full_bf16_pytorch_flash(
            q, k, v, output, num_heads, num_kv_heads, num_tokens,
            head_dim, aligned_head_dim, kv_stride_tokens,
            /*output_token_major=*/1)) {
        return;
    }
    fprintf(stderr,
            "HARD KERNEL CONTRACT FAULT: PyTorch CPU-flash BF16 token-output "
            "attention requires AMX-BF16, AVX-512, D=72/A=72 and a token "
            "multiple of 16; no numerically different fallback is permitted\n");
    abort();
}

References ck_attention_full_bf16_pytorch_flash().

◆ attention_forward_full_head_major_gqa_sdpa_bf16_storage()

void attention_forward_full_head_major_gqa_sdpa_bf16_storage	(	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
	)

Definition at line 5022 of file attention_kernels.c.

{
    if (ck_attention_full_bf16_sdpa_tiled(
            q, k, v, output, num_heads, num_kv_heads, num_tokens,
            head_dim, aligned_head_dim, kv_stride_tokens)) {
        const size_t count = (size_t)num_heads * (size_t)num_tokens
                           * (size_t)aligned_head_dim;
        for (size_t i = 0; i < count; ++i) {
            output[i] = bf16_to_float(float_to_bf16(output[i]));
        }
        return;
    }
    fprintf(stderr, "CK numerical contract failure: BF16 tiled SDPA received invalid dimensions or could not allocate scratch\n");
}

References bf16_to_float(), ck_attention_full_bf16_sdpa_tiled(), count, and float_to_bf16().

◆ attention_forward_full_head_major_gqa_tiled336_f16kv_fp32_strided()

void attention_forward_full_head_major_gqa_tiled336_f16kv_fp32_strided	(	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
	)

Definition at line 3638 of file attention_kernels.c.

{
    ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(
        q, k, v, output, num_heads, num_kv_heads, num_tokens,
        head_dim, aligned_head_dim, kv_stride_tokens, CK_GGML_FA_TILE_Q_LARGE);
}

References ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(), and CK_GGML_FA_TILE_Q_LARGE.

◆ attention_forward_full_head_major_gqa_tiled64_f16kv_fp32_strided()

void attention_forward_full_head_major_gqa_tiled64_f16kv_fp32_strided	(	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
	)

Definition at line 3621 of file attention_kernels.c.

{
    ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(
        q, k, v, output, num_heads, num_kv_heads, num_tokens,
        head_dim, aligned_head_dim, kv_stride_tokens, CK_GGML_FA_TILE_Q);
}

References ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(), and CK_GGML_FA_TILE_Q.

◆ attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided()

void attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided	(	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
	)

Definition at line 3601 of file attention_kernels.c.

{
    const int query_tile_size = num_tokens >= CK_GGML_FA_TILE_Q_LARGE_MIN_TOKENS
        ? CK_GGML_FA_TILE_Q_LARGE
        : CK_GGML_FA_TILE_Q;
    ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(
        q, k, v, output, num_heads, num_kv_heads, num_tokens,
        head_dim, aligned_head_dim, kv_stride_tokens, query_tile_size);
}

References ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(), CK_GGML_FA_TILE_Q, CK_GGML_FA_TILE_Q_LARGE, and CK_GGML_FA_TILE_Q_LARGE_MIN_TOKENS.

◆ attention_forward_head_major_gqa_flash_impl()

static void attention_forward_head_major_gqa_flash_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	causal,
		int	round_full_kv_fp16,
		int	output_token_major,
		float	scale
	)

static

Definition at line 4183 of file attention_kernels.c.

{
    if (!q || !k || !v || !output) {
        return;
    }
    if (num_heads <= 0 || num_kv_heads <= 0 || num_tokens <= 0 ||
        head_dim <= 0 || aligned_head_dim < head_dim) {
        return;
    }
    if (kv_stride_tokens < num_tokens) {
        return;
    }
 
    const int T = num_tokens;
    const size_t kv_head_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
 
    if (ck_strict_parity_enabled()) {
        const float strict_scale = ck_attention_strict_scale_f32(head_dim);
        const int debug_layer_id = ck_attention_vec_dump_enabled()
            ? ck_attention_vec_dump_next_layer_id()
            : -1;
        if (causal && ck_attention_strict_unfused_f16_enabled()) {
            if (attention_forward_head_major_gqa_unfused_f16_strict(
                    q, k, v, output,
                    num_heads, num_kv_heads, num_tokens,
                    head_dim, aligned_head_dim, kv_stride_tokens,
                    causal, output_token_major, strict_scale, debug_layer_id)) {
                return;
            }
        }
#if CK_ENABLE_LLAMA_CPP_PARITY
        if (!causal && !output_token_major &&
            ck_attention_full_ggml_graph_oracle_multihead(q,
                                                          k,
                                                          v,
                                                          output,
                                                          num_heads,
                                                          num_kv_heads,
                                                          num_tokens,
                                                          head_dim,
                                                          aligned_head_dim,
                                                          kv_stride_tokens,
                                                          strict_scale)) {
            return;
        }
#endif
        float *score_row = (float *) alloca((size_t) T * sizeof(float));
        float *logit_row = (float *) alloca((size_t) T * sizeof(float));
        float *v_cols = (float *) alloca((size_t) head_dim * (size_t) T * sizeof(float));
        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 d = 0; d < head_dim; ++d) {
                float *dst_col = v_cols + (size_t) d * (size_t) T;
                for (int j = 0; j < T; ++j) {
                    dst_col[j] = v_head[(size_t) j * (size_t) aligned_head_dim + (size_t) d];
                }
            }
 
            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 + qkv_index(h, i, 0, T, aligned_head_dim);
                const int kv_tokens = causal ? (i + 1) : T;
                if (kv_tokens <= 0) {
                    for (int d = 0; d < aligned_head_dim; ++d) {
                        out_vec[d] = 0.0f;
                    }
                    continue;
                }
                for (int j = 0; j < kv_tokens; ++j) {
                    const float *k_vec = k_head + (size_t) j * (size_t) aligned_head_dim;
                    score_row[j] = ck_ggml_vec_dot_f32_contig(q_vec, k_vec, head_dim);
                }
                float *raw_dump = NULL;
                if (ck_attention_vec_dump_should_emit(debug_layer_id, h, i)) {
                    raw_dump = (float *) alloca((size_t) kv_tokens * sizeof(float));
                    memcpy(raw_dump, score_row, (size_t) kv_tokens * sizeof(float));
                }
                memcpy(logit_row, score_row, (size_t) kv_tokens * sizeof(float));
                ck_vec_scale_f32_inplace(logit_row, kv_tokens, strict_scale);
                const float max_score = ck_vec_max_f32_contig(logit_row, kv_tokens);
                const double sum = ck_ggml_vec_soft_max_row(kv_tokens, score_row, logit_row, max_score);
                if (sum > 0.0) {
                    const float inv_sum = (float) (1.0 / sum);
                    ck_vec_scale_f32_inplace(score_row, kv_tokens, inv_sum);
                    for (int d = 0; d < head_dim; ++d) {
                        const float *v_col = v_cols + (size_t) d * (size_t) T;
                        out_vec[d] = ck_ggml_vec_dot_f32_contig(score_row, v_col, kv_tokens);
                    }
                } else {
                    for (int d = 0; d < head_dim; ++d) {
                        out_vec[d] = 0.0f;
                    }
                }
                for (int d = head_dim; d < aligned_head_dim; ++d) {
                    out_vec[d] = 0.0f;
                }
                if (raw_dump) {
                    ck_attention_vec_dump_selected_query(raw_dump, score_row, out_vec, v_cols,
                                                         kv_tokens, head_dim,
                                                         debug_layer_id, h, i);
                }
            }
        }
        return;
    }
 
    /*
     * Full GGML-style attention consumes FP32 Q with K/V rounded through
     * FP16.  Materialize that semantic input once so serial, threaded, and
     * ISA-specific implementations all see the same values.  Causal callers
     * retain their existing FP32 K/V contract.
     */
    float *rounded_kv = NULL;
    const float *compute_k = k;
    const float *compute_v = v;
    if (!causal && round_full_kv_fp16) {
        if ((size_t) kv_stride_tokens > SIZE_MAX / (size_t) num_kv_heads) {
            return;
        }
        const size_t rows = (size_t) num_kv_heads * (size_t) kv_stride_tokens;
        if ((size_t) aligned_head_dim > SIZE_MAX / rows) {
            return;
        }
        const size_t elements = rows * (size_t) aligned_head_dim;
        if (elements > SIZE_MAX / (2 * sizeof(float))) {
            return;
        }
        rounded_kv = (float *) malloc(2 * elements * sizeof(float));
        if (!rounded_kv) {
            return;
        }
        compute_k = rounded_kv;
        compute_v = rounded_kv + elements;
        ck_round_fp16_buffer(k, rounded_kv, elements);
        ck_round_fp16_buffer(v, rounded_kv + elements, elements);
    }
 
    // Select SIMD implementation based on compile-time CPU features
#if defined(__AVX512F__)
    #define FLASH_QUERY_IMPL attention_flash_query_causal_avx512
#elif defined(__AVX2__)
    #define FLASH_QUERY_IMPL attention_flash_query_causal_avx2
#elif defined(__AVX__)
    #define FLASH_QUERY_IMPL attention_flash_query_causal_avx
#else
    #define FLASH_QUERY_IMPL attention_flash_query_causal
#endif
 
    const int total_queries = num_heads * T;
#if defined(__AVX512F__)
    if (!causal && !output_token_major && head_dim == 72 && aligned_head_dim >= 72 && ck_attention_qblock8_enabled()) {
        ck_threadpool_t *pool = ck_threadpool_global();
        const int q_blocks = (T + 7) / 8;
        const int total_blocks = num_heads * q_blocks;
        int active = ck_attention_pick_active_threads(pool, total_blocks, T);
        if (pool && active > 1) {
            ck_attention_qblock4_args_t args = {
                .q = q,
                .k = compute_k,
                .v = compute_v,
                .output = output,
                .num_heads = num_heads,
                .num_kv_heads = num_kv_heads,
                .num_tokens = num_tokens,
                .aligned_head_dim = aligned_head_dim,
                .kv_stride_tokens = kv_stride_tokens,
                .scale = scale,
            };
            ck_threadpool_dispatch_n(pool, active, ck_attention_full_qblock8_work, &args);
            free(rounded_kv);
            return;
        }
    }
 
    if (!causal && !output_token_major && head_dim == 72 && aligned_head_dim >= 72 && ck_attention_qblock4_enabled()) {
        ck_threadpool_t *pool = ck_threadpool_global();
        const int q_blocks = (T + 3) / 4;
        const int total_blocks = num_heads * q_blocks;
        int active = ck_attention_pick_active_threads(pool, total_blocks, T);
        if (pool && active > 1) {
            ck_attention_qblock4_args_t args = {
                .q = q,
                .k = compute_k,
                .v = compute_v,
                .output = output,
                .num_heads = num_heads,
                .num_kv_heads = num_kv_heads,
                .num_tokens = num_tokens,
                .aligned_head_dim = aligned_head_dim,
                .kv_stride_tokens = kv_stride_tokens,
                .scale = scale,
            };
            ck_threadpool_dispatch_n(pool, active, ck_attention_full_qblock4_work, &args);
            free(rounded_kv);
            return;
        }
    }
#endif
    if (ck_attention_parallel_enabled(total_queries, T, head_dim)) {
        ck_threadpool_t *pool = ck_threadpool_global();
        const int active = ck_attention_pick_active_threads(pool, total_queries, T);
        if (pool && active > 1) {
            ck_attention_parallel_args_t args = {
                .q = q,
                .k = compute_k,
                .v = compute_v,
                .output = output,
                .num_heads = num_heads,
                .num_kv_heads = num_kv_heads,
                .num_tokens = num_tokens,
                .head_dim = head_dim,
                .aligned_head_dim = aligned_head_dim,
                .kv_stride_tokens = kv_stride_tokens,
                .causal = causal,
                .output_token_major = output_token_major,
                .scale = scale,
            };
            ck_threadpool_dispatch_n(pool, active, ck_attention_full_grid_work, &args);
            free(rounded_kv);
            return;
        }
    }
 
    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 = compute_k + (size_t)kv_head * kv_head_stride;
        const float *v_head = compute_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);
            const int kv_tokens = causal ? (i + 1) : T;
            FLASH_QUERY_IMPL(q_vec, k_head, v_head,
                             kv_tokens,
                             head_dim, aligned_head_dim,
                             scale, out_vec);
        }
    }
 
    free(rounded_kv);
 
#undef FLASH_QUERY_IMPL
}

References attention_forward_head_major_gqa_unfused_f16_strict(), attention_output_index(), ck_attention_full_ggml_graph_oracle_multihead(), ck_attention_full_grid_work(), ck_attention_parallel_enabled(), ck_attention_pick_active_threads(), ck_attention_strict_scale_f32(), ck_attention_strict_unfused_f16_enabled(), ck_attention_vec_dump_enabled(), ck_attention_vec_dump_next_layer_id(), ck_attention_vec_dump_selected_query(), ck_attention_vec_dump_should_emit(), ck_ggml_vec_dot_f32_contig(), ck_ggml_vec_soft_max_row(), ck_round_fp16_buffer(), ck_strict_parity_enabled(), ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_vec_max_f32_contig(), ck_vec_scale_f32_inplace(), FLASH_QUERY_IMPL, and qkv_index().

◆ attention_forward_head_major_gqa_unfused_f16_strict()

static int attention_forward_head_major_gqa_unfused_f16_strict	(	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	causal,
		int	output_token_major,
		float	scale,
		int	debug_layer_id
	)

static

Definition at line 4095 of file attention_kernels.c.

{
    const int T = num_tokens;
    const size_t kv_head_stride = (size_t) kv_stride_tokens * (size_t) aligned_head_dim;
    const size_t kv_half_count = (size_t) T * (size_t) aligned_head_dim;
    uint16_t *k_half = (uint16_t *) malloc(kv_half_count * sizeof(uint16_t));
    uint16_t *v_half = (uint16_t *) malloc(kv_half_count * sizeof(uint16_t));
    if (!k_half || !v_half) {
        free(k_half);
        free(v_half);
        return 0;
    }
    uint16_t *q_half = (uint16_t *) alloca((size_t) aligned_head_dim * sizeof(uint16_t));
    uint16_t *prob_half = (uint16_t *) alloca((size_t) T * sizeof(uint16_t));
    uint16_t *v_col_half = (uint16_t *) alloca((size_t) T * sizeof(uint16_t));
    float *raw_scores = (float *) alloca((size_t) T * sizeof(float));
    float *logits = (float *) alloca((size_t) T * sizeof(float));
    int cached_kv_head = -1;
 
    for (int h = 0; h < num_heads; ++h) {
        const 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;
        if (kv_head != cached_kv_head) {
            for (size_t idx = 0; idx < kv_half_count; ++idx) {
                k_half[idx] = CK_FP32_TO_FP16(k_head[idx]);
                v_half[idx] = CK_FP32_TO_FP16(v_head[idx]);
            }
            cached_kv_head = kv_head;
        }
 
        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);
            const int kv_tokens = causal ? i + 1 : T;
            for (int d = 0; d < aligned_head_dim; ++d) {
                q_half[d] = CK_FP32_TO_FP16(q_vec[d]);
            }
            for (int j = 0; j < kv_tokens; ++j) {
                raw_scores[j] = ck_attention_dot_f16_unfused_llama(
                    q_half,
                    k_half + (size_t) j * (size_t) aligned_head_dim,
                    head_dim);
                logits[j] = raw_scores[j] * scale;
            }
 
            const float max_score = ck_vec_max_f32_contig(logits, kv_tokens);
            const double sum = ck_ggml_vec_soft_max_row(kv_tokens, logits, logits, max_score);
            const float inv_sum = sum > 0.0 ? (float) (1.0 / sum) : 0.0f;
            for (int j = 0; j < kv_tokens; ++j) {
                logits[j] *= inv_sum;
                prob_half[j] = CK_FP32_TO_FP16(logits[j]);
            }
            for (int d = 0; d < head_dim; ++d) {
                for (int j = 0; j < kv_tokens; ++j) {
                    v_col_half[j] = v_half[(size_t) j * (size_t) aligned_head_dim + (size_t) d];
                }
                out_vec[d] = ck_attention_dot_f16_unfused_llama(prob_half, v_col_half, kv_tokens);
            }
            for (int d = head_dim; d < aligned_head_dim; ++d) {
                out_vec[d] = 0.0f;
            }
            ck_attention_vec_dump_selected_query(raw_scores, logits, out_vec, NULL,
                                                 kv_tokens, head_dim,
                                                 debug_layer_id, h, i);
        }
    }
    free(k_half);
    free(v_half);
    return 1;
}

References attention_output_index(), ck_attention_dot_f16_unfused_llama(), ck_attention_vec_dump_selected_query(), CK_FP32_TO_FP16, ck_ggml_vec_soft_max_row(), ck_vec_max_f32_contig(), and qkv_index().

Referenced by attention_forward_head_major_gqa_flash_impl().

◆ attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4()

void attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_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	visual_start,
		int	visual_tokens
	)

Definition at line 5290 of file attention_kernels.c.

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

References attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl().

◆ attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl()

static void attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_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	visual_start,
		int	visual_tokens,
		int	output_token_major
	)

static

Definition at line 5182 of file attention_kernels.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 || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
    if (visual_start < 0 || visual_tokens <= 0 || visual_start >= num_tokens) {
        attention_forward_head_major_gqa_flash_impl(
            q, k, v, output, num_heads, num_kv_heads, num_tokens, head_dim,
            aligned_head_dim, kv_stride_tokens, /*causal=*/1,
            /*round_full_kv_fp16=*/0, output_token_major, 1.0f);
        return;
    }
 
    int visual_end = visual_start + visual_tokens;
    if (visual_end > num_tokens) {
        visual_end = num_tokens;
    }
    if (visual_end <= visual_start) {
        attention_forward_head_major_gqa_flash_impl(
            q, k, v, output, num_heads, num_kv_heads, num_tokens, head_dim,
            aligned_head_dim, kv_stride_tokens, /*causal=*/1,
            /*round_full_kv_fp16=*/0, output_token_major, 1.0f);
        return;
    }
 
    const int T = num_tokens;
    const size_t kv_head_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
    const float scale = 1.0f;
 
    if (ck_strict_parity_enabled()) {
        const int debug_layer_id = ck_attention_vec_dump_enabled()
            ? ck_attention_vec_dump_next_layer_id()
            : -1;
        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);
                const int in_visual = (i >= visual_start && i < visual_end);
                const int kv_tokens = in_visual ? visual_end : (i + 1);
                attention_flash_query_causal_exact(q_vec, k_head, v_head,
                                                   kv_tokens,
                                                   head_dim, aligned_head_dim,
                                                   scale, out_vec);
                ck_attention_vec_dump_exact_query(q_vec, k_head, out_vec,
                                                  kv_tokens,
                                                  head_dim, aligned_head_dim,
                                                  scale,
                                                  debug_layer_id, h, i);
            }
        }
        return;
    }
 
#if defined(__AVX512F__)
    #define FLASH_QUERY_IMPL attention_flash_query_causal_avx512
#elif defined(__AVX2__)
    #define FLASH_QUERY_IMPL attention_flash_query_causal_avx2
#elif defined(__AVX__)
    #define FLASH_QUERY_IMPL attention_flash_query_causal_avx
#else
    #define FLASH_QUERY_IMPL attention_flash_query_causal
#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);
            const int in_visual = (i >= visual_start && i < visual_end);
            const int kv_tokens = in_visual ? visual_end : (i + 1);
            FLASH_QUERY_IMPL(q_vec, k_head, v_head,
                             kv_tokens,
                             head_dim, aligned_head_dim,
                             scale, out_vec);
        }
    }
 
#undef FLASH_QUERY_IMPL
}

References attention_flash_query_causal_exact(), attention_forward_head_major_gqa_flash_impl(), attention_output_index(), ck_attention_vec_dump_enabled(), ck_attention_vec_dump_exact_query(), ck_attention_vec_dump_next_layer_id(), ck_strict_parity_enabled(), FLASH_QUERY_IMPL, and qkv_index().

Referenced by attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4(), and attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_token_output().

◆ attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_token_output()

void attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_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	visual_start,
		int	visual_tokens
	)

Definition at line 5302 of file attention_kernels.c.

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

References attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl().

◆ attention_forward_query_key_head_major_f32()

int attention_forward_query_key_head_major_f32	(	const float *	query,
		const float *	key,
		const float *	value,
		float *	output,
		float *	score_scratch,
		int	num_heads,
		int	query_tokens,
		int	key_tokens,
		int	head_dim,
		float	scale
	)

Definition at line 210 of file attention_kernels.c.

{
    return ck_attention_forward_query_key_head_major_f32_run(
        query, key, value, output, score_scratch, NULL, num_heads,
        query_tokens, key_tokens, head_dim, scale);
}

References ck_attention_forward_query_key_head_major_f32_run().

◆ attention_forward_query_key_head_major_f32_decode_heads()

int attention_forward_query_key_head_major_f32_decode_heads	(	const float *	query,
		const float *	key,
		const float *	value,
		float *	output,
		float *	score_scratch,
		int	num_heads,
		int	query_tokens,
		int	key_tokens,
		int	head_dim,
		float	scale
	)

Definition at line 260 of file attention_kernels.c.

{
    if (query == NULL || key == NULL || value == NULL || output == NULL ||
        score_scratch == NULL) return -1;
    if (num_heads <= 0 || query_tokens <= 0 || key_tokens <= 0 ||
        head_dim <= 0 || !isfinite(scale)) return -2;
 
    ck_attention_query_key_f32_head_args_t args = {
        .query = query,
        .key = key,
        .value = value,
        .output = output,
        .score_scratch = score_scratch,
        .num_heads = num_heads,
        .query_tokens = query_tokens,
        .key_tokens = key_tokens,
        .head_dim = head_dim,
        .scale = scale,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active <= 1) {
        ck_attention_query_key_f32_head_work(0, 1, &args);
        return 0;
    }
    if (active > num_heads) active = num_heads;
    ck_threadpool_dispatch_n(pool, active,
        ck_attention_query_key_f32_head_work, &args);
    return 0;
}

References ck_attention_query_key_f32_head_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

◆ attention_forward_query_key_head_major_f32_packed_k()

int attention_forward_query_key_head_major_f32_packed_k	(	const float *	query,
		const float *	key,
		const float *	value,
		float *	output,
		float *	score_scratch,
		float *	key_transpose_scratch,
		int	num_heads,
		int	query_tokens,
		int	key_tokens,
		int	head_dim,
		float	scale
	)

Definition at line 301 of file attention_kernels.c.

{
    if (key_transpose_scratch == NULL) return -1;
    return ck_attention_forward_query_key_head_major_f32_run(
        query, key, value, output, score_scratch, key_transpose_scratch,
        num_heads, query_tokens, key_tokens, head_dim, scale);
}

References ck_attention_forward_query_key_head_major_f32_run().

◆ attention_forward_query_key_head_major_tiled_f16kv_fp32()

int attention_forward_query_key_head_major_tiled_f16kv_fp32	(	const float *	query,
		const float *	key,
		const float *	value,
		float *	output,
		int	num_heads,
		int	query_tokens,
		int	key_tokens,
		int	head_dim,
		float	scale
	)

Definition at line 3655 of file attention_kernels.c.

{
    if (!query || !key || !value || !output || num_heads <= 0 ||
        query_tokens <= 0 || key_tokens != query_tokens || head_dim <= 0) {
        return -1;
    }
 
    const float contract_scale = ck_attention_strict_scale_f32(head_dim);
    if (scale != contract_scale) {
        return -1;
    }
 
    const int query_tile_size = query_tokens >= CK_GGML_FA_TILE_Q_LARGE_MIN_TOKENS
        ? CK_GGML_FA_TILE_Q_LARGE
        : CK_GGML_FA_TILE_Q;
    ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(
        query, key, value, output,
        num_heads, num_heads, query_tokens,
        head_dim, head_dim, key_tokens, query_tile_size);
    return 0;
}

References ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(), ck_attention_strict_scale_f32(), CK_GGML_FA_TILE_Q, CK_GGML_FA_TILE_Q_LARGE, and CK_GGML_FA_TILE_Q_LARGE_MIN_TOKENS.

◆ attention_forward_sparse_token_major_gqa_bf16cache_pytorch_cpu_flash_contract()

void attention_forward_sparse_token_major_gqa_bf16cache_pytorch_cpu_flash_contract	(	const float *	query,
		const uint16_t *	key_cache,
		const uint16_t *	value_cache,
		const float *	selected_indices,
		float *	output,
		float *	score_scratch,
		int	rows,
		int	query_heads,
		int	kv_heads,
		int	head_dim,
		int	selection_width,
		int	context_length,
		int	position
	)

Definition at line 8966 of file attention_kernels.c.

{
    (void)score_scratch;
    if (!query || !key_cache || !value_cache || !selected_indices || !output ||
        rows <= 0 || query_heads <= 0 || kv_heads <= 0 || head_dim <= 0 ||
        selection_width <= 0 || context_length <= 0 || position < 0 ||
        position + rows > context_length || query_heads % kv_heads != 0) {
        fprintf(stderr, "CK sparse BF16 CPU-flash attention: invalid contract\n");
        abort();
    }
 
#if defined(__AVX512F__)
    for (int row = 0; row < rows; ++row) {
        const int visible_tokens = position + row + 1;
        const ck_attention_status_t status =
            ck_attention_decode_bf16_pytorch_cpu_flash_masked(
                query + (size_t)row * (size_t)query_heads * (size_t)head_dim,
                key_cache, value_cache,
                output + (size_t)row * (size_t)query_heads * (size_t)head_dim,
                selected_indices + (size_t)row * (size_t)selection_width,
                selection_width, query_heads, kv_heads, visible_tokens,
                context_length, head_dim, head_dim);
        if (status != CK_ATTENTION_STATUS_OK) {
            fprintf(stderr,
                    "CK sparse BF16 CPU-flash attention unavailable (status=%d)\n",
                    (int)status);
            abort();
        }
    }
#else
    fprintf(stderr,
            "CK sparse BF16 CPU-flash attention requires AVX-512F; "
            "no numerically different fallback is permitted\n");
    abort();
#endif
}

References CK_ATTENTION_STATUS_OK.

◆ attention_output_index()

static size_t attention_output_index	(	int	h,
		int	token,
		int	num_heads,
		int	num_tokens,
		int	aligned_head_dim,
		int	token_major
	)

inlinestatic

Definition at line 339 of file attention_kernels.c.

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

References qkv_index(), and token.

Referenced by attention_forward_head_major_gqa_flash_impl(), attention_forward_head_major_gqa_unfused_f16_strict(), attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl(), and ck_attention_full_grid_work().

◆ attention_query_full_exact_regular()

static void attention_query_full_exact_regular	(	const float *	q_vec,
		const float *	k_head,
		const float *	v_cols,
		int	kv_tokens,
		int	head_dim,
		int	aligned_head_dim,
		float	scale,
		float *	score_row,
		float *	out_vec,
		int	layer_id,
		int	head_id,
		int	query_id
	)

static

Definition at line 2321 of file attention_kernels.c.

{
    if (kv_tokens <= 0) {
        for (int d = 0; d < aligned_head_dim; ++d) {
            out_vec[d] = 0.0f;
        }
        return;
    }
 
    for (int j = 0; j < kv_tokens; ++j) {
        const float *k_vec = k_head + (size_t) j * (size_t) aligned_head_dim;
        score_row[j] = ck_vec_dot_f32_strict(q_vec, k_vec, head_dim);
    }
    float *raw_dump = NULL;
    if (ck_attention_vec_dump_should_emit(layer_id, head_id, query_id)) {
        raw_dump = (float *) alloca((size_t) kv_tokens * sizeof(float));
        memcpy(raw_dump, score_row, (size_t) kv_tokens * sizeof(float));
    }
 
    float max_score = -INFINITY;
    for (int j = 0; j < kv_tokens; ++j) {
        const float score = score_row[j] * scale;
        score_row[j] = score;
        if (score > max_score) {
            max_score = score;
        }
    }
 
    double sum = 0.0;
    for (int j = 0; j < kv_tokens; ++j) {
        const float w = expf(score_row[j] - max_score);
        score_row[j] = w;
        volatile double next = sum + (double) w;
        sum = next;
    }
 
    if (sum > 0.0) {
        const float inv_sum = (float) (1.0 / sum);
        for (int j = 0; j < kv_tokens; ++j) {
            score_row[j] *= inv_sum;
        }
        for (int d = 0; d < head_dim; ++d) {
            const float *v_col = v_cols + (size_t) d * (size_t) kv_tokens;
            const float dot = ck_vec_dot_f32x_f32_to_f32_via_f64(score_row, v_col, kv_tokens);
            out_vec[d] = dot;
        }
    } else {
        for (int d = 0; d < head_dim; ++d) {
            out_vec[d] = 0.0f;
        }
    }
 
    for (int d = head_dim; d < aligned_head_dim; ++d) {
        out_vec[d] = 0.0f;
    }
    if (raw_dump) {
        ck_attention_vec_dump_selected_query(raw_dump, score_row, out_vec, v_cols, kv_tokens, head_dim,
                                             layer_id, head_id, query_id);
    }
}

References ck_attention_vec_dump_selected_query(), ck_attention_vec_dump_should_emit(), ck_vec_dot_f32_strict(), ck_vec_dot_f32x_f32_to_f32_via_f64(), and score.

Referenced by attention_forward_full_head_major_gqa_exact_strided().

◆ attention_query_full_ggml_regular()

static void attention_query_full_ggml_regular	(	const float *	q_vec,
		const float *	k_head,
		const float *	v_cols,
		int	kv_tokens,
		int	head_dim,
		int	aligned_head_dim,
		float	scale,
		float *	score_row,
		float *	out_vec,
		int	layer_id,
		int	head_id,
		int	query_id
	)

static

Definition at line 2393 of file attention_kernels.c.

{
    if (kv_tokens <= 0) {
        for (int d = 0; d < aligned_head_dim; ++d) {
            out_vec[d] = 0.0f;
        }
        return;
    }
 
    for (int j = 0; j < kv_tokens; ++j) {
        const float *k_vec = k_head + (size_t) j * (size_t) aligned_head_dim;
        score_row[j] = ck_ggml_vec_dot_f32_contig(q_vec, k_vec, head_dim);
    }
    float *raw_dump = NULL;
    if (ck_attention_vec_dump_should_emit(layer_id, head_id, query_id)) {
        raw_dump = (float *) alloca((size_t) kv_tokens * sizeof(float));
        memcpy(raw_dump, score_row, (size_t) kv_tokens * sizeof(float));
    }
 
    float *logit_row = (float *) alloca((size_t) kv_tokens * sizeof(float));
    memcpy(logit_row, score_row, (size_t) kv_tokens * sizeof(float));
    ck_vec_scale_f32_inplace(logit_row, kv_tokens, scale);
    const float max_score = ck_vec_max_f32_contig(logit_row, kv_tokens);
    const double sum = ck_ggml_vec_soft_max_row(kv_tokens, score_row, logit_row, max_score);
    if (sum > 0.0) {
        const float inv_sum = (float) (1.0 / sum);
        ck_vec_scale_f32_inplace(score_row, kv_tokens, inv_sum);
        for (int d = 0; d < head_dim; ++d) {
            const float *v_col = v_cols + (size_t) d * (size_t) kv_tokens;
            out_vec[d] = ck_ggml_vec_dot_f32_contig(score_row, v_col, kv_tokens);
        }
    } else {
        for (int d = 0; d < head_dim; ++d) {
            out_vec[d] = 0.0f;
        }
    }
 
    for (int d = head_dim; d < aligned_head_dim; ++d) {
        out_vec[d] = 0.0f;
    }
    if (raw_dump) {
        ck_attention_vec_dump_selected_query(raw_dump, score_row, out_vec, v_cols, kv_tokens, head_dim,
                                             layer_id, head_id, query_id);
    }
}

References ck_attention_vec_dump_selected_query(), ck_attention_vec_dump_should_emit(), ck_ggml_vec_dot_f32_contig(), ck_ggml_vec_soft_max_row(), ck_vec_max_f32_contig(), and ck_vec_scale_f32_inplace().

Referenced by attention_forward_full_head_major_gqa_ggml_strided_workspace().

◆ ck_attention_align64_size()

static size_t ck_attention_align64_size ( size_t value )

inlinestatic

Definition at line 7021 of file attention_kernels.c.

{
    return (value + 63u) & ~(size_t) 63u;
}

Referenced by attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config(), attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes(), ck_attention_f16_prefill_gqa_reuse_work(), and ck_attention_f16_prefill_gqa_reuse_worker_bytes().

◆ ck_attention_bf16_pytorch_flash_work()

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

static

Definition at line 4932 of file attention_kernels.c.

{
    ck_attention_bf16_sdpa_args_t *args = (ck_attention_bf16_sdpa_args_t *)opaque;
    if (!ck_attention_full_bf16_sdpa_amx_range(
            args->q, args->k, args->v, args->output,
            args->num_heads, args->num_kv_heads, args->num_tokens,
            args->head_dim, args->aligned_head_dim, args->kv_stride_tokens,
            ith, nth, args->output_token_major)) {
        __atomic_store_n(&args->failed, 1, __ATOMIC_RELAXED);
    }
}

References ck_attention_full_bf16_sdpa_amx_range().

Referenced by ck_attention_full_bf16_pytorch_flash().

◆ ck_attention_bf16_pytorch_gqa_available()

int ck_attention_bf16_pytorch_gqa_available ( void )

Definition at line 9025 of file attention_kernels.c.

{
#if defined(__AVX512F__)
    pthread_once(&ck_pytorch_attention_once, ck_bind_pytorch_attention_primitives);
    return ck_pytorch_sgemm_batch != NULL &&
           ck_pytorch_attention_expf16 != NULL;
#else
    return 1;
#endif
}

Referenced by ck_attention_sparse_bf16_pytorch_gqa_available().

◆ ck_attention_bf16_sdpa_work()

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

static

Definition at line 4920 of file attention_kernels.c.

{
    ck_attention_bf16_sdpa_args_t *args = (ck_attention_bf16_sdpa_args_t *)opaque;
    if (!ck_attention_full_bf16_sdpa_tiled_range(
            args->q, args->k, args->v, args->output,
            args->num_heads, args->num_kv_heads, args->num_tokens,
            args->head_dim, args->aligned_head_dim, args->kv_stride_tokens,
            ith, nth)) {
        __atomic_store_n(&args->failed, 1, __ATOMIC_RELAXED);
    }
}

References ck_attention_full_bf16_sdpa_tiled_range().

Referenced by ck_attention_full_bf16_sdpa_tiled().

◆ ck_attention_causal_f16kv_work()

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

static

Definition at line 5328 of file attention_kernels.c.

{
    const ck_attention_causal_f16kv_args_t *args =
        (const ck_attention_causal_f16kv_args_t *)opaque;
    const float scale = 1.0f / sqrtf((float)args->head_dim);
    const int T = args->num_tokens;
    const size_t kv_head_stride =
        (size_t)args->kv_stride_tokens * (size_t)args->aligned_head_dim;
 
    for (int h = ith; h < args->num_heads; h += nth) {
        const int kv_head = (int)((long long)h * (long long)args->num_kv_heads /
                                  (long long)args->num_heads);
        const float *k_head = args->k + (size_t)kv_head * kv_head_stride;
        const float *v_head = args->v + (size_t)kv_head * kv_head_stride;
 
        for (int i = 0; i < T; ++i) {
            const float *q_vec = args->q +
                qkv_index(h, i, 0, T, args->aligned_head_dim);
            float *out_vec = args->output +
                qkv_index(h, i, 0, T, args->aligned_head_dim);
            if (args->inputs_prerounded) {
                attention_flash_query_causal_exact_prerounded_f16kv(
                    q_vec, k_head, v_head, /*kv_tokens=*/i + 1,
                    args->head_dim, args->aligned_head_dim, scale, out_vec);
            } else {
                attention_flash_query_causal_exact_f16kv(
                    q_vec, k_head, v_head, /*kv_tokens=*/i + 1,
                    args->head_dim, args->aligned_head_dim, scale, out_vec);
            }
        }
    }
}

References attention_flash_query_causal_exact_f16kv(), attention_flash_query_causal_exact_prerounded_f16kv(), and qkv_index().

Referenced by attention_forward_causal_head_major_gqa_flash_strided_f16kv_serial(), and attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace().

◆ ck_attention_dot_f16_llama()

static float ck_attention_dot_f16_llama	(	const uint16_t *	x,
		const uint16_t *	y,
		int	n
	)

inlinestatic

Definition at line 6426 of file attention_kernels.c.

{
    int i = 0;
#if defined(__AVX512F__)
    // Match ggml's AVX-512 FP16-to-FP32 path: four 16-lane
    // accumulators per 64 values, then its fixed pairwise tree.
    __m512 sum0 = _mm512_setzero_ps();
    __m512 sum1 = _mm512_setzero_ps();
    __m512 sum2 = _mm512_setzero_ps();
    __m512 sum3 = _mm512_setzero_ps();
    const int n64 = n & ~63;
    for (; i < n64; i += 64) {
        const __m512 x0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *) (x + i)));
        const __m512 y0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *) (y + i)));
        const __m512 x1 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *) (x + i + 16)));
        const __m512 y1 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *) (y + i + 16)));
        const __m512 x2 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *) (x + i + 32)));
        const __m512 y2 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *) (y + i + 32)));
        const __m512 x3 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *) (x + i + 48)));
        const __m512 y3 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *) (y + i + 48)));
        sum0 = _mm512_fmadd_ps(x0, y0, sum0);
        sum1 = _mm512_fmadd_ps(x1, y1, sum1);
        sum2 = _mm512_fmadd_ps(x2, y2, sum2);
        sum3 = _mm512_fmadd_ps(x3, y3, sum3);
    }
    sum0 = _mm512_add_ps(sum0, sum2);
    sum1 = _mm512_add_ps(sum1, sum3);
    sum0 = _mm512_add_ps(sum0, sum1);
    const float vector_result = _mm512_reduce_add_ps(sum0);
#elif defined(__AVX2__) && defined(__F16C__)
    // Match ggml_vec_dot_f16's AVX reduction contract: four independent
    // accumulators per 32 values, followed by its fixed pairwise tree.
    __m256 sum0 = _mm256_setzero_ps();
    __m256 sum1 = _mm256_setzero_ps();
    __m256 sum2 = _mm256_setzero_ps();
    __m256 sum3 = _mm256_setzero_ps();
    const int n32 = n & ~31;
    for (; i < n32; i += 32) {
        const __m256 x0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (x + i)));
        const __m256 y0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (y + i)));
        const __m256 x1 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (x + i + 8)));
        const __m256 y1 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (y + i + 8)));
        const __m256 x2 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (x + i + 16)));
        const __m256 y2 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (y + i + 16)));
        const __m256 x3 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (x + i + 24)));
        const __m256 y3 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (y + i + 24)));
#if defined(__FMA__)
        sum0 = _mm256_fmadd_ps(x0, y0, sum0);
        sum1 = _mm256_fmadd_ps(x1, y1, sum1);
        sum2 = _mm256_fmadd_ps(x2, y2, sum2);
        sum3 = _mm256_fmadd_ps(x3, y3, sum3);
#else
        sum0 = _mm256_add_ps(sum0, _mm256_mul_ps(x0, y0));
        sum1 = _mm256_add_ps(sum1, _mm256_mul_ps(x1, y1));
        sum2 = _mm256_add_ps(sum2, _mm256_mul_ps(x2, y2));
        sum3 = _mm256_add_ps(sum3, _mm256_mul_ps(x3, y3));
#endif
    }
    sum0 = _mm256_add_ps(sum0, sum2);
    sum1 = _mm256_add_ps(sum1, sum3);
    sum0 = _mm256_add_ps(sum0, sum1);
    const __m128 pair = _mm_add_ps(
        _mm256_castps256_ps128(sum0),
        _mm256_extractf128_ps(sum0, 1));
    const __m128 half = _mm_hadd_ps(pair, pair);
    const float vector_result = _mm_cvtss_f32(_mm_hadd_ps(half, half));
#else
    const float vector_result = 0.0f;
#endif
    /* ggml_vec_dot_f16 stores the SIMD reduction in ggml_float (double) and
     * accumulates any scalar tail there. This matters for head dimensions
     * below or not divisible by the active ISA step (for example D=32 on
     * AVX-512 and padded D=80). */
    double result = (double) vector_result;
    for (; i < n; ++i) {
        const float product = CK_FP16_TO_FP32(x[i]) * CK_FP16_TO_FP32(y[i]);
        result += (double) product;
    }
    return (float) result;
}

References CK_FP16_TO_FP32, and product().

Referenced by attention_forward_decode_head_major_gqa_f16cache_regular_noscale(), and ck_attention_f16_split_work().

◆ ck_attention_dot_f16_unfused_llama()

static float ck_attention_dot_f16_unfused_llama	(	const uint16_t *	x,
		const uint16_t *	y,
		int	n
	)

inlinestatic

Definition at line 1179 of file attention_kernels.c.

{
    int i = 0;
#if defined(__AVX2__) && defined(__F16C__)
    __m256 sum0 = _mm256_setzero_ps();
    __m256 sum1 = _mm256_setzero_ps();
    __m256 sum2 = _mm256_setzero_ps();
    __m256 sum3 = _mm256_setzero_ps();
    const int n32 = n & ~31;
    for (; i < n32; i += 32) {
        const __m256 x0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (x + i)));
        const __m256 y0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (y + i)));
        const __m256 x1 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (x + i + 8)));
        const __m256 y1 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (y + i + 8)));
        const __m256 x2 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (x + i + 16)));
        const __m256 y2 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (y + i + 16)));
        const __m256 x3 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (x + i + 24)));
        const __m256 y3 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) (y + i + 24)));
#if defined(__FMA__)
        sum0 = _mm256_fmadd_ps(x0, y0, sum0);
        sum1 = _mm256_fmadd_ps(x1, y1, sum1);
        sum2 = _mm256_fmadd_ps(x2, y2, sum2);
        sum3 = _mm256_fmadd_ps(x3, y3, sum3);
#else
        sum0 = _mm256_add_ps(sum0, _mm256_mul_ps(x0, y0));
        sum1 = _mm256_add_ps(sum1, _mm256_mul_ps(x1, y1));
        sum2 = _mm256_add_ps(sum2, _mm256_mul_ps(x2, y2));
        sum3 = _mm256_add_ps(sum3, _mm256_mul_ps(x3, y3));
#endif
    }
    sum0 = _mm256_add_ps(sum0, sum2);
    sum1 = _mm256_add_ps(sum1, sum3);
    sum0 = _mm256_add_ps(sum0, sum1);
    const __m128 pair = _mm_add_ps(
        _mm256_castps256_ps128(sum0),
        _mm256_extractf128_ps(sum0, 1));
    const __m128 half = _mm_hadd_ps(pair, pair);
    float result = _mm_cvtss_f32(_mm_hadd_ps(half, half));
#else
    float result = 0.0f;
#endif
    for (; i < n; ++i) {
        result += CK_FP16_TO_FP32(x[i]) * CK_FP16_TO_FP32(y[i]);
    }
    return result;
}

References CK_FP16_TO_FP32.

Referenced by attention_forward_head_major_gqa_unfused_f16_strict().

◆ ck_attention_f16_prefill_gqa_reuse_work()

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

static

Definition at line 7292 of file attention_kernels.c.

{
    const ck_attention_f16_prefill_gqa_reuse_args_t *args =
        (const ck_attention_f16_prefill_gqa_reuse_args_t *) opaque;
    const ck_attention_f16_prefill_qtile64_args_t *base = &args->base;
    const int lanes = nth / base->num_kv_heads;
    const int kv_group = ith / lanes;
    const int lane = ith % lanes;
    const int heads_per_group = base->num_heads / base->num_kv_heads;
    const int head_begin = kv_group * heads_per_group;
    const int query_tiles =
        (base->q_tokens + args->query_tile_size - 1) / args->query_tile_size;
    const int kv_tokens = base->past_tokens + base->q_tokens;
    const float scale = ck_attention_strict_scale_f32(base->head_dim);
    const size_t kv_head_stride =
        (size_t) base->cache_capacity * (size_t) base->aligned_head_dim;
    const uint16_t *k_head =
        base->k_cache + (size_t) kv_group * kv_head_stride;
    const uint16_t *v_head =
        base->v_cache + (size_t) kv_group * kv_head_stride;
    float *shared_k = args->shared_k_tiles +
        (size_t) kv_group * (size_t) base->head_dim * CK_GGML_FA_TILE_KV;
    float *shared_v = args->shared_v_tiles +
        (size_t) kv_group * CK_GGML_FA_TILE_KV * (size_t) base->head_dim;
    ck_attention_gqa_team_barrier_t *barrier = &args->barriers[kv_group];
    const int max_group_jobs =
        heads_per_group * args->concurrent_query_tiles;
    const int local_capacity = (max_group_jobs + lanes - 1) / lanes;
    unsigned char *worker = args->worker_workspace +
        (size_t) ith * args->worker_workspace_stride;
    size_t cursor = 0;
    cursor = ck_attention_align64_size(cursor);
    int *local_heads = (int *) (worker + cursor);
    cursor += (size_t) local_capacity * sizeof(int);
    int *local_iq = (int *) (worker + cursor);
    cursor += (size_t) local_capacity * sizeof(int);
    int *local_rows = (int *) (worker + cursor);
    cursor += (size_t) local_capacity * sizeof(int);
    cursor = ck_attention_align64_size(cursor);
    float *q_tiles = (float *) (worker + cursor);
    cursor += (size_t) local_capacity * (size_t) args->query_tile_size *
        (size_t) base->head_dim * sizeof(float);
    cursor = ck_attention_align64_size(cursor);
    float *kq_tiles = (float *) (worker + cursor);
    cursor += (size_t) local_capacity * (size_t) args->query_tile_size *
        CK_GGML_FA_TILE_KV * sizeof(float);
    cursor = ck_attention_align64_size(cursor);
    float *vkq_tiles = (float *) (worker + cursor);
    cursor += (size_t) local_capacity * (size_t) args->query_tile_size *
        (size_t) base->head_dim * sizeof(float);
    cursor = ck_attention_align64_size(cursor);
    float *sum_rows = (float *) (worker + cursor);
    cursor += (size_t) local_capacity * (size_t) args->query_tile_size *
        sizeof(float);
    cursor = ck_attention_align64_size(cursor);
    float *max_rows = (float *) (worker + cursor);
 
    for (int query_tile_begin = 0; query_tile_begin < query_tiles;
         query_tile_begin += args->concurrent_query_tiles) {
        int batch_tiles = query_tiles - query_tile_begin;
        if (batch_tiles > args->concurrent_query_tiles) {
            batch_tiles = args->concurrent_query_tiles;
        }
        const int group_jobs = heads_per_group * batch_tiles;
        int local_jobs = 0;
 
        for (int group_job = lane; group_job < group_jobs; group_job += lanes) {
            const int local = local_jobs++;
            const int group_head = group_job / batch_tiles;
            const int batch_tile = group_job % batch_tiles;
            const int h = head_begin + group_head;
            const int iq =
                (query_tile_begin + batch_tile) * args->query_tile_size;
            int tile_rows = base->q_tokens - iq;
            if (tile_rows > args->query_tile_size) {
                tile_rows = args->query_tile_size;
            }
            local_heads[local] = h;
            local_iq[local] = iq;
            local_rows[local] = tile_rows;
 
            float *q_tile = q_tiles +
                (size_t) local * (size_t) args->query_tile_size *
                (size_t) base->head_dim;
            float *vkq = vkq_tiles +
                (size_t) local * (size_t) args->query_tile_size *
                (size_t) base->head_dim;
            float *sum_row = sum_rows +
                (size_t) local * (size_t) args->query_tile_size;
            float *max_row = max_rows +
                (size_t) local * (size_t) args->query_tile_size;
            memset(q_tile, 0,
                   (size_t) args->query_tile_size * (size_t) base->head_dim *
                       sizeof(float));
            memset(vkq, 0,
                   (size_t) args->query_tile_size * (size_t) base->head_dim *
                       sizeof(float));
            for (int tq = 0; tq < args->query_tile_size; ++tq) {
                sum_row[tq] = 0.0f;
                max_row[tq] = -INFINITY;
            }
            for (int tq = 0; tq < tile_rows; ++tq) {
                const float *q_vec = base->q +
                    (size_t) h * base->q_head_stride +
                    (size_t) (iq + tq) * (size_t) base->aligned_head_dim;
                memcpy(q_tile + (size_t) tq * (size_t) base->head_dim,
                       q_vec, (size_t) base->head_dim * sizeof(float));
            }
        }
 
        int kv_limit = base->past_tokens +
            (query_tile_begin + batch_tiles) * args->query_tile_size;
        if (kv_limit > kv_tokens) kv_limit = kv_tokens;
 
        for (int ik = 0; ik < kv_limit; ik += CK_GGML_FA_TILE_KV) {
            int kv_tile = kv_limit - ik;
            if (kv_tile > CK_GGML_FA_TILE_KV) kv_tile = CK_GGML_FA_TILE_KV;
            const int packed_elems = CK_GGML_FA_TILE_KV * base->head_dim;
            for (int flat = lane; flat < packed_elems; flat += lanes) {
                const int tk = flat / base->head_dim;
                const int d = flat % base->head_dim;
                float k_value = 0.0f;
                float v_value = 0.0f;
                if (tk < kv_tile) {
                    const uint16_t *k_vec = k_head +
                        (size_t) (ik + tk) * (size_t) base->aligned_head_dim;
                    const uint16_t *v_vec = v_head +
                        (size_t) (ik + tk) * (size_t) base->aligned_head_dim;
                    k_value = ck_attention_u16_cache_to_f32(
                        k_vec[d], base->cache_is_bf16);
                    v_value = ck_attention_u16_cache_to_f32(
                        v_vec[d], base->cache_is_bf16);
                }
                shared_k[(size_t) d * CK_GGML_FA_TILE_KV + (size_t) tk] =
                    k_value;
                shared_v[(size_t) tk * (size_t) base->head_dim + (size_t) d] =
                    v_value;
            }
            ck_attention_gqa_team_barrier_wait(barrier);
 
            for (int local = 0; local < local_jobs; ++local) {
                const int iq = local_iq[local];
                const int tile_rows = local_rows[local];
                float *q_tile = q_tiles +
                    (size_t) local * (size_t) args->query_tile_size *
                    (size_t) base->head_dim;
                float *kq = kq_tiles +
                    (size_t) local * (size_t) args->query_tile_size *
                    CK_GGML_FA_TILE_KV;
                float *vkq = vkq_tiles +
                    (size_t) local * (size_t) args->query_tile_size *
                    (size_t) base->head_dim;
                float *sum_row = sum_rows +
                    (size_t) local * (size_t) args->query_tile_size;
                float *max_row = max_rows +
                    (size_t) local * (size_t) args->query_tile_size;
                memset(kq, 0,
                       (size_t) args->query_tile_size * CK_GGML_FA_TILE_KV *
                           sizeof(float));
                ck_attention_matmul_f32_accum(
                    kq, q_tile, shared_k, args->query_tile_size,
                    base->head_dim, CK_GGML_FA_TILE_KV);
                ck_vec_scale_f32_inplace(
                    kq, args->query_tile_size * CK_GGML_FA_TILE_KV, scale);
 
                for (int tq = 0; tq < args->query_tile_size; ++tq) {
                    float *kq_row = kq + (size_t) tq * CK_GGML_FA_TILE_KV;
                    const int last_valid_key = base->past_tokens + iq + tq;
                    for (int tk = 0; tk < CK_GGML_FA_TILE_KV; ++tk) {
                        const int key = ik + tk;
                        if (tk >= kv_tile || tq >= tile_rows || key > last_valid_key) {
                            kq_row[tk] = -INFINITY;
                        }
                    }
                    const float tile_max =
                        ck_vec_max_f32_contig(kq_row, CK_GGML_FA_TILE_KV);
                    if (tile_max == -INFINITY) {
                        memset(kq_row, 0,
                               CK_GGML_FA_TILE_KV * sizeof(float));
                        continue;
                    }
                    const float old_max = max_row[tq];
                    const float new_max = old_max > tile_max ? old_max : tile_max;
                    if (new_max > old_max) {
                        const float ms =
                            ck_attention_reference_expf(old_max - new_max);
                        ck_vec_scale_f32_inplace(
                            vkq + (size_t) tq * (size_t) base->head_dim,
                            base->head_dim, ms);
                        sum_row[tq] *= ms;
                    }
                    max_row[tq] = new_max;
                    sum_row[tq] = (float) (
                        (double) sum_row[tq] +
                        ck_ggml_vec_soft_max_row(
                            CK_GGML_FA_TILE_KV, kq_row, kq_row, new_max));
                }
                ck_attention_matmul_f32_accum(
                    vkq, kq, shared_v, args->query_tile_size,
                    CK_GGML_FA_TILE_KV, base->head_dim);
            }
            ck_attention_gqa_team_barrier_wait(barrier);
        }
 
        for (int local = 0; local < local_jobs; ++local) {
            const int h = local_heads[local];
            const int iq = local_iq[local];
            const int tile_rows = local_rows[local];
            float *vkq = vkq_tiles +
                (size_t) local * (size_t) args->query_tile_size *
                (size_t) base->head_dim;
            float *sum_row = sum_rows +
                (size_t) local * (size_t) args->query_tile_size;
            for (int tq = 0; tq < tile_rows; ++tq) {
                float *out_vec = base->output +
                    (size_t) h * base->output_head_stride +
                    (size_t) (iq + tq) * (size_t) base->aligned_head_dim;
                const float inv_sum =
                    sum_row[tq] == 0.0f ? 0.0f : 1.0f / sum_row[tq];
                for (int d = 0; d < base->head_dim; ++d) {
                    out_vec[d] = vkq[
                        (size_t) tq * (size_t) base->head_dim + (size_t) d] *
                        inv_sum;
                }
                for (int d = base->head_dim; d < base->aligned_head_dim; ++d) {
                    out_vec[d] = 0.0f;
                }
            }
        }
    }
}

References ck_attention_align64_size(), ck_attention_gqa_team_barrier_wait(), ck_attention_matmul_f32_accum(), ck_attention_reference_expf(), ck_attention_strict_scale_f32(), ck_attention_u16_cache_to_f32(), CK_GGML_FA_TILE_KV, ck_ggml_vec_soft_max_row(), ck_vec_max_f32_contig(), and ck_vec_scale_f32_inplace().

Referenced by attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config().

◆ ck_attention_f16_prefill_gqa_reuse_worker_bytes()

static size_t ck_attention_f16_prefill_gqa_reuse_worker_bytes	(	int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	workers,
		int	query_tile_size,
		int	concurrent_query_tiles
	)

static

Definition at line 7026 of file attention_kernels.c.

{
    if (num_heads <= 0 || num_kv_heads <= 0 ||
        num_heads % num_kv_heads != 0 || workers < num_kv_heads ||
        workers % num_kv_heads != 0 || head_dim <= 0 ||
        query_tile_size <= 0 || concurrent_query_tiles <= 0) {
        return 0;
    }
    const int lanes = workers / num_kv_heads;
    const int group_jobs =
        (num_heads / num_kv_heads) * concurrent_query_tiles;
    const size_t local_capacity =
        (size_t) (group_jobs + lanes - 1) / (size_t) lanes;
    size_t bytes = 0;
    bytes = ck_attention_align64_size(bytes) +
        3u * local_capacity * sizeof(int);
    bytes = ck_attention_align64_size(bytes) +
        local_capacity * (size_t) query_tile_size * (size_t) head_dim *
        sizeof(float);
    bytes = ck_attention_align64_size(bytes) +
        local_capacity * (size_t) query_tile_size * CK_GGML_FA_TILE_KV *
        sizeof(float);
    bytes = ck_attention_align64_size(bytes) +
        local_capacity * (size_t) query_tile_size * (size_t) head_dim *
        sizeof(float);
    bytes = ck_attention_align64_size(bytes) +
        local_capacity * (size_t) query_tile_size * sizeof(float);
    bytes = ck_attention_align64_size(bytes) +
        local_capacity * (size_t) query_tile_size * sizeof(float);
    return ck_attention_align64_size(bytes);
}

References ck_attention_align64_size(), and CK_GGML_FA_TILE_KV.

Referenced by attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config(), and attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes().

◆ ck_attention_f16_prefill_qtile64_dispatch()

static ck_attention_status_t ck_attention_f16_prefill_qtile64_dispatch	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		int	cache_is_bf16,
		float	attention_scale,
		size_t	q_head_stride,
		size_t	output_head_stride,
		ck_attention_prefill_schedule_t	schedule
	)

static

Definition at line 7524 of file attention_kernels.c.

{
    if (schedule < CK_ATTN_PREFILL_SCHEDULE_KV_HEADS ||
        schedule > CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    ck_attention_f16_prefill_qtile64_args_t args = {
        .q = q,
        .k_cache = k_cache,
        .v_cache = v_cache,
        .output = output,
        .num_heads = num_heads,
        .num_kv_heads = num_kv_heads,
        .q_tokens = q_tokens,
        .past_tokens = past_tokens,
        .cache_capacity = cache_capacity,
        .head_dim = head_dim,
        .aligned_head_dim = aligned_head_dim,
        .cache_is_bf16 = cache_is_bf16,
        .attention_scale = attention_scale,
        .schedule = schedule,
        .q_head_stride = q_head_stride,
        .output_head_stride = output_head_stride,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    const int query_tiles =
        (q_tokens + CK_GGML_FA_TILE_Q - 1) / CK_GGML_FA_TILE_Q;
    int available_jobs = num_heads;
    if (schedule == CK_ATTN_PREFILL_SCHEDULE_KV_HEADS) {
        available_jobs = num_kv_heads;
    } else if (schedule == CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES ||
               schedule == CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES) {
        available_jobs = num_heads * query_tiles;
    }
    if (active > available_jobs) active = available_jobs;
    if (pool && active > 1 && ck_threadpool_thread_id(pool) <= 0) {
        ck_threadpool_dispatch_n(
            pool, active, ck_attention_f16_prefill_qtile64_work, &args);
    } else {
        ck_attention_f16_prefill_qtile64_work(0, 1, &args);
    }
    return CK_ATTENTION_STATUS_OK;
}

References ck_attention_f16_prefill_qtile64_work(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES, CK_ATTN_PREFILL_SCHEDULE_KV_HEADS, CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES, CK_GGML_FA_TILE_Q, ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), and ck_threadpool_thread_id().

Referenced by attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace(), attention_forward_causal_head_major_gqa_prefill_append_f16cache_qtile64_schedule(), and ck_attention_prefill_append_f16cache_gemma4_workspace().

◆ ck_attention_f16_prefill_qtile64_work()

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

static

Definition at line 7113 of file attention_kernels.c.

{
    const ck_attention_f16_prefill_qtile64_args_t *args =
        (const ck_attention_f16_prefill_qtile64_args_t *) opaque;
    const int query_tiles =
        (args->q_tokens + CK_GGML_FA_TILE_Q - 1) / CK_GGML_FA_TILE_Q;
    const int total_jobs = args->num_heads * query_tiles;
    int job_begin = 0;
    int job_end = 0;
 
    if (args->schedule == CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES &&
        nth >= args->num_kv_heads) {
        const int kv_group = (ith * args->num_kv_heads) / nth;
        const int worker_begin =
            (kv_group * nth + args->num_kv_heads - 1) / args->num_kv_heads;
        const int worker_end =
            ((kv_group + 1) * nth + args->num_kv_heads - 1) /
            args->num_kv_heads;
        const int lane = ith - worker_begin;
        const int lanes = worker_end - worker_begin;
        const int head_begin =
            (kv_group * args->num_heads + args->num_kv_heads - 1) /
            args->num_kv_heads;
        const int head_end =
            ((kv_group + 1) * args->num_heads + args->num_kv_heads - 1) /
            args->num_kv_heads;
        const int group_jobs = (head_end - head_begin) * query_tiles;
        job_begin = head_begin * query_tiles + (group_jobs * lane) / lanes;
        job_end = head_begin * query_tiles + (group_jobs * (lane + 1)) / lanes;
    } else if (args->schedule == CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES) {
        job_begin = (total_jobs * ith) / nth;
        job_end = (total_jobs * (ith + 1)) / nth;
    } else {
        const int head_begin = (args->num_heads * ith) / nth;
        const int head_end = (args->num_heads * (ith + 1)) / nth;
        job_begin = head_begin * query_tiles;
        job_end = head_end * query_tiles;
    }
    if (job_begin >= job_end) return;
 
    const int kv_tokens = args->past_tokens + args->q_tokens;
    const float scale = args->attention_scale != 0.0f
        ? args->attention_scale : ck_attention_strict_scale_f32(args->head_dim);
    const size_t kv_head_stride =
        (size_t) args->cache_capacity * (size_t) args->aligned_head_dim;
 
    float *q_tile = (float *) alloca(
        (size_t) CK_GGML_FA_TILE_Q * (size_t) args->head_dim * sizeof(float));
    float *k_tile = (float *) alloca(
        (size_t) args->head_dim * (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
    float *v_tile = (float *) alloca(
        (size_t) CK_GGML_FA_TILE_KV * (size_t) args->head_dim * sizeof(float));
    float *kq = (float *) alloca(
        (size_t) CK_GGML_FA_TILE_Q * (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
    float *vkq = (float *) alloca(
        (size_t) CK_GGML_FA_TILE_Q * (size_t) args->head_dim * sizeof(float));
 
    for (int job = job_begin; job < job_end; ++job) {
        const int h = job / query_tiles;
        const int iq = (job % query_tiles) * CK_GGML_FA_TILE_Q;
        const int kv_head =
            (int) ((long long) h * (long long) args->num_kv_heads /
                   (long long) args->num_heads);
        const uint16_t *k_head = args->k_cache + (size_t) kv_head * kv_head_stride;
        const uint16_t *v_head = args->v_cache + (size_t) kv_head * kv_head_stride;
 
        const int tile_rows =
            (args->q_tokens - iq) < CK_GGML_FA_TILE_Q
                ? (args->q_tokens - iq)
                : CK_GGML_FA_TILE_Q;
            float sum_row[CK_GGML_FA_TILE_Q];
            float max_row[CK_GGML_FA_TILE_Q];
 
            for (int tq = 0; tq < CK_GGML_FA_TILE_Q; ++tq) {
                sum_row[tq] = 0.0f;
                max_row[tq] = -INFINITY;
            }
            memset(q_tile, 0,
                   (size_t) CK_GGML_FA_TILE_Q * (size_t) args->head_dim * sizeof(float));
            memset(vkq, 0,
                   (size_t) CK_GGML_FA_TILE_Q * (size_t) args->head_dim * sizeof(float));
 
            for (int tq = 0; tq < tile_rows; ++tq) {
                const float *q_vec = args->q +
                    (size_t) h * args->q_head_stride +
                    (size_t) (iq + tq) * (size_t) args->aligned_head_dim;
                memcpy(q_tile + (size_t) tq * (size_t) args->head_dim,
                       q_vec,
                       (size_t) args->head_dim * sizeof(float));
            }
 
            for (int ik = 0; ik < kv_tokens; ik += CK_GGML_FA_TILE_KV) {
                const int kv_tile =
                    (kv_tokens - ik) < CK_GGML_FA_TILE_KV
                        ? (kv_tokens - ik)
                        : CK_GGML_FA_TILE_KV;
                memset(k_tile, 0,
                       (size_t) args->head_dim * (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
                memset(v_tile, 0,
                       (size_t) CK_GGML_FA_TILE_KV * (size_t) args->head_dim * sizeof(float));
                memset(kq, 0,
                       (size_t) CK_GGML_FA_TILE_Q * (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
 
                for (int tk = 0; tk < kv_tile; ++tk) {
                    const uint16_t *k_vec = k_head +
                        (size_t) (ik + tk) * (size_t) args->aligned_head_dim;
                    const uint16_t *v_vec = v_head +
                        (size_t) (ik + tk) * (size_t) args->aligned_head_dim;
                    for (int d = 0; d < args->head_dim; ++d) {
                        k_tile[(size_t) d * (size_t) CK_GGML_FA_TILE_KV + (size_t) tk] =
                            ck_attention_u16_cache_to_f32(k_vec[d], args->cache_is_bf16);
                        v_tile[(size_t) tk * (size_t) args->head_dim + (size_t) d] =
                            ck_attention_u16_cache_to_f32(v_vec[d], args->cache_is_bf16);
                    }
                }
 
                ck_attention_matmul_f32_accum(
                    kq, q_tile, k_tile,
                    CK_GGML_FA_TILE_Q, args->head_dim, CK_GGML_FA_TILE_KV);
                ck_vec_scale_f32_inplace(
                    kq, CK_GGML_FA_TILE_Q * CK_GGML_FA_TILE_KV, scale);
 
                for (int tq = 0; tq < CK_GGML_FA_TILE_Q; ++tq) {
                    float *kq_row = kq +
                        (size_t) tq * (size_t) CK_GGML_FA_TILE_KV;
                    const int last_valid_key = args->past_tokens + iq + tq;
                    for (int tk = 0; tk < CK_GGML_FA_TILE_KV; ++tk) {
                        const int key = ik + tk;
                        if (tk >= kv_tile || tq >= tile_rows || key > last_valid_key) {
                            kq_row[tk] = -INFINITY;
                        }
                    }
 
                    const float tile_max =
                        ck_vec_max_f32_contig(kq_row, CK_GGML_FA_TILE_KV);
                    if (tile_max == -INFINITY) {
                        memset(kq_row, 0,
                               (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
                        continue;
                    }
 
                    const float old_max = max_row[tq];
                    const float new_max = old_max > tile_max ? old_max : tile_max;
                    if (new_max > old_max) {
                        const float ms = ck_attention_reference_expf(old_max - new_max);
                        ck_vec_scale_f32_inplace(
                            vkq + (size_t) tq * (size_t) args->head_dim,
                            args->head_dim, ms);
                        sum_row[tq] *= ms;
                    }
                    max_row[tq] = new_max;
                    sum_row[tq] = (float) (
                        (double) sum_row[tq] +
                        ck_ggml_vec_soft_max_row(
                            CK_GGML_FA_TILE_KV, kq_row, kq_row, new_max));
                }
 
                ck_attention_matmul_f32_accum(
                    vkq, kq, v_tile,
                    CK_GGML_FA_TILE_Q, CK_GGML_FA_TILE_KV, args->head_dim);
            }
 
            for (int tq = 0; tq < tile_rows; ++tq) {
                float *out_vec = args->output +
                    (size_t) h * args->output_head_stride +
                    (size_t) (iq + tq) * (size_t) args->aligned_head_dim;
                const float inv_sum =
                    sum_row[tq] == 0.0f ? 0.0f : 1.0f / sum_row[tq];
                for (int d = 0; d < args->head_dim; ++d) {
                    out_vec[d] =
                        vkq[(size_t) tq * (size_t) args->head_dim + (size_t) d] * inv_sum;
                }
                for (int d = args->head_dim; d < args->aligned_head_dim; ++d) {
                    out_vec[d] = 0.0f;
                }
        }
    }
}

References ck_attention_matmul_f32_accum(), ck_attention_reference_expf(), ck_attention_strict_scale_f32(), ck_attention_u16_cache_to_f32(), CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES, CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES, CK_GGML_FA_TILE_KV, CK_GGML_FA_TILE_Q, ck_ggml_vec_soft_max_row(), ck_vec_max_f32_contig(), and ck_vec_scale_f32_inplace().

Referenced by attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace(), ck_attention_f16_prefill_qtile64_dispatch(), and ck_attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_schedule_workspace().

◆ ck_attention_f16_reduce_expf()

static float ck_attention_f16_reduce_expf ( float value )

inlinestatic

Definition at line 6421 of file attention_kernels.c.

{
    return ck_attention_reference_expf(value);
}

References ck_attention_reference_expf().

Referenced by attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned().

◆ ck_attention_f16_split_work()

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

static

Definition at line 6589 of file attention_kernels.c.

{
    ck_attention_f16_split_args_t *args = (ck_attention_f16_split_args_t *) opaque;
    const int partial_stride = args->aligned_head_dim + 2;
    const int total_jobs = args->num_heads * args->scheduled_chunks;
    const int chunk_size =
        (args->partition_tokens + args->split_chunks - 1) / args->split_chunks;
    const size_t head_stride = (size_t) args->cache_capacity * (size_t) args->aligned_head_dim;
    const float scale = args->attention_scale;
    uint16_t *q_half = (uint16_t *) alloca((size_t) args->aligned_head_dim * sizeof(uint16_t));
    uint16_t *acc_half = (uint16_t *) alloca((size_t) args->aligned_head_dim * sizeof(uint16_t));
 
    for (int job = ith; job < total_jobs; job += nth) {
        const int h = job / args->scheduled_chunks;
        const int chunk = job % args->scheduled_chunks;
        const int kv_head = (int) ((long long) h * (long long) args->num_kv_heads /
                                   (long long) args->num_heads);
        const int begin = args->kv_start + chunk * chunk_size;
        const int end = begin < args->kv_tokens
            ? (begin + chunk_size < args->kv_tokens ? begin + chunk_size : args->kv_tokens)
            : begin;
        const float *q_head = args->q_token + (size_t) h * (size_t) args->aligned_head_dim;
        const uint16_t *k_head = args->k_cache + (size_t) kv_head * head_stride;
        const uint16_t *v_head = args->v_cache + (size_t) kv_head * head_stride;
        float *partial = args->partials + (size_t) job * (size_t) partial_stride;
 
        for (int d = 0; d < args->aligned_head_dim; ++d) {
            q_half[d] = CK_FP32_TO_FP16(q_head[d]);
            acc_half[d] = CK_FP32_TO_FP16(0.0f);
        }
 
        float sum = 0.0f;
        float max_score = -INFINITY;
        for (int j = begin; j < end; ++j) {
            const uint16_t *k_vec = k_head + (size_t) j * (size_t) args->aligned_head_dim;
            const uint16_t *v_vec = v_head + (size_t) j * (size_t) args->aligned_head_dim;
            const float dot = ck_attention_dot_f16_llama(q_half, k_vec, args->head_dim);
            const float score = dot * scale;
            const float old_max = max_score;
            float max_scale = 1.0f;
            float value_scale = 1.0f;
 
            if (score > max_score) {
                max_score = score;
                max_scale = isfinite(old_max)
                    ? ck_attention_reference_expf(old_max - max_score)
                    : 0.0f;
                ck_attention_scale_f16_llama(acc_half, max_scale, args->head_dim);
            } else {
                value_scale = ck_attention_reference_expf(score - max_score);
            }
 
            ck_attention_mad_f16_llama(
                acc_half, v_vec, value_scale, args->head_dim);
#if defined(__INTEL_LLVM_COMPILER)
            sum = ck_attention_mul_add_rounded_f32(
                sum, max_scale, value_scale);
#else
            sum = fmaf(sum, max_scale, value_scale);
#endif
        }
 
        partial[0] = max_score;
        partial[1] = sum;
        for (int d = 0; d < args->head_dim; ++d) {
            partial[2 + d] = CK_FP16_TO_FP32(acc_half[d]);
        }
        for (int d = args->head_dim; d < args->aligned_head_dim; ++d) {
            partial[2 + d] = 0.0f;
        }
    }
}

References ck_attention_dot_f16_llama(), ck_attention_mad_f16_llama(), ck_attention_reference_expf(), ck_attention_scale_f16_llama(), CK_FP16_TO_FP32, CK_FP32_TO_FP16, end, and score.

Referenced by attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned().

◆ ck_attention_flash_query_auto()

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

inlinestatic

Definition at line 3715 of file attention_kernels.c.

{
#if defined(__AVX512F__)
    attention_flash_query_causal_avx512(q_vec, k_head, v_head,
                                        kv_tokens, head_dim, aligned_head_dim,
                                        scale, out_vec);
#elif defined(__AVX2__)
    attention_flash_query_causal_avx2(q_vec, k_head, v_head,
                                      kv_tokens, head_dim, aligned_head_dim,
                                      scale, out_vec);
#elif defined(__AVX__)
    attention_flash_query_causal_avx(q_vec, k_head, v_head,
                                     kv_tokens, head_dim, aligned_head_dim,
                                     scale, out_vec);
#else
    attention_flash_query_causal(q_vec, k_head, v_head,
                                 kv_tokens, head_dim, aligned_head_dim,
                                 scale, out_vec);
#endif
}

References attention_flash_query_causal().

Referenced by ck_attention_full_grid_work().

◆ ck_attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_schedule_workspace()

static ck_attention_status_t ck_attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_schedule_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes,
		const int *	segment_lengths,
		int	num_segments,
		ck_attention_prefill_schedule_t	schedule
	)

static

Definition at line 8017 of file attention_kernels.c.

{
    if (!q || !k_cache || !v_cache || !output || !segment_lengths ||
        num_segments <= 0 || num_heads <= 0 || num_kv_heads <= 0 ||
        q_tokens <= 0 || past_tokens < 0 || past_tokens + q_tokens > cache_capacity ||
        head_dim <= 0 || aligned_head_dim < head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    int planned_rows = 0;
    for (int s = 0; s < num_segments; ++s) {
        if (segment_lengths[s] < 0 || segment_lengths[s] > q_tokens - planned_rows) {
            return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
        }
        planned_rows += segment_lengths[s];
    }
    if (planned_rows != q_tokens) return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
 
    const size_t head_stride = (size_t) q_tokens * (size_t) aligned_head_dim;
    const size_t token_elems = (size_t) num_heads * (size_t) aligned_head_dim;
    if (!token_workspace || token_workspace_bytes < 2 * token_elems * sizeof(float)) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
 
    int row_offset = 0;
    for (int s = 0; s < num_segments; ++s) {
        const int rows = segment_lengths[s];
        if (rows == 0) continue;
        if (reduction == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64 &&
            rows >= CK_GGML_FA_TILE_Q) {
            ck_attention_f16_prefill_qtile64_args_t args = {
                .q = q + (size_t) row_offset * (size_t) aligned_head_dim,
                .k_cache = k_cache,
                .v_cache = v_cache,
                .output = output + (size_t) row_offset * (size_t) aligned_head_dim,
                .num_heads = num_heads,
                .num_kv_heads = num_kv_heads,
                .q_tokens = rows,
                .past_tokens = past_tokens + row_offset,
                .cache_capacity = cache_capacity,
                .head_dim = head_dim,
                .aligned_head_dim = aligned_head_dim,
                .cache_is_bf16 = 0,
                .schedule = schedule,
                .q_head_stride = head_stride,
                .output_head_stride = head_stride,
            };
            ck_threadpool_t *pool = ck_threadpool_global();
            int active = pool ? ck_threadpool_n_threads(pool) : 1;
            const int query_tiles =
                (rows + CK_GGML_FA_TILE_Q - 1) / CK_GGML_FA_TILE_Q;
            const int available_jobs =
                schedule == CK_ATTN_PREFILL_SCHEDULE_KV_HEADS
                    ? num_kv_heads
                    : num_heads * query_tiles;
            if (active > available_jobs) active = available_jobs;
            if (pool && active > 1 && ck_threadpool_thread_id(pool) <= 0) {
                ck_threadpool_dispatch_n(
                    pool, active, ck_attention_f16_prefill_qtile64_work, &args);
            } else {
                ck_attention_f16_prefill_qtile64_work(0, 1, &args);
            }
        } else {
            float *q_token = token_workspace;
            float *out_token = token_workspace + token_elems;
            ck_attention_reduction_t selected = reduction;
            if (selected == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64) {
                selected = CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE;
            }
            for (int t = 0; t < rows; ++t) {
                for (int h = 0; h < num_heads; ++h) {
                    memcpy(
                        q_token + (size_t) h * (size_t) aligned_head_dim,
                        q + (size_t) h * head_stride +
                            (size_t) (row_offset + t) * (size_t) aligned_head_dim,
                        (size_t) aligned_head_dim * sizeof(float));
                }
                ck_attention_status_t status =
                    attention_forward_decode_head_major_gqa_flash_f16cache_contract(
                        q_token, k_cache, v_cache, out_token,
                        num_heads, num_kv_heads, past_tokens + row_offset + t + 1,
                        cache_capacity, head_dim, aligned_head_dim, selected);
                if (status != CK_ATTENTION_STATUS_OK) return status;
                for (int h = 0; h < num_heads; ++h) {
                    memcpy(
                        output + (size_t) h * head_stride +
                            (size_t) (row_offset + t) * (size_t) aligned_head_dim,
                        out_token + (size_t) h * (size_t) aligned_head_dim,
                        (size_t) aligned_head_dim * sizeof(float));
                }
            }
        }
        row_offset += rows;
    }
    return CK_ATTENTION_STATUS_OK;
}

References attention_forward_decode_head_major_gqa_flash_f16cache_contract(), ck_attention_f16_prefill_qtile64_work(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_ATTN_PREFILL_SCHEDULE_KV_HEADS, CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64, CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE, CK_GGML_FA_TILE_Q, ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), and ck_threadpool_thread_id().

Referenced by attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_contract_workspace().

◆ ck_attention_forward_decode_f16cache_gemma4()

static ck_attention_status_t ck_attention_forward_decode_f16cache_gemma4	(	const float *	q_token,
		const uint16_t *	k_cache,
		const uint16_t *	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,
		ck_attention_reduction_t	reduction
	)

static

Definition at line 6806 of file attention_kernels.c.

{
    if (!q_token || !k_cache || !v_cache || !out_token ||
        num_heads <= 0 || num_kv_heads <= 0 || kv_tokens <= 0 ||
        cache_capacity <= 0 || kv_tokens > cache_capacity ||
        head_dim <= 0 || aligned_head_dim < head_dim || sliding_window < 0) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    if (reduction == CK_ATTN_REDUCTION_F16_REGULAR_NOSCALE) {
        return attention_forward_decode_head_major_gqa_f16cache_regular_noscale(
            q_token, k_cache, v_cache, out_token, num_heads, num_kv_heads,
            kv_tokens, cache_capacity, head_dim, aligned_head_dim, sliding_window);
    }
    if (reduction != CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE) {
        return CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT;
    }
    const int kv_start = sliding_window > 0 && kv_tokens > sliding_window
        ? kv_tokens - sliding_window : 0;
    const int live_tokens = kv_tokens - kv_start;
    attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned(
        q_token, k_cache, v_cache, out_token,
        num_heads, num_kv_heads, kv_tokens, cache_capacity,
        head_dim, aligned_head_dim, 1, live_tokens, kv_start, 1.0f);
    return CK_ATTENTION_STATUS_OK;
}

References attention_forward_decode_head_major_gqa_f16cache_regular_noscale(), attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT, CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE, and CK_ATTN_REDUCTION_F16_REGULAR_NOSCALE.

Referenced by attention_forward_decode_head_major_gqa_f16cache_gemma4_contract(), attention_forward_decode_head_major_gqa_f16cache_sliding_gemma4_contract(), and ck_attention_prefill_append_f16cache_gemma4_workspace().

◆ ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided()

static void ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided	(	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	query_tile_size
	)

static

Definition at line 3560 of file attention_kernels.c.

{
    if (!q || !k || !v || !output || num_heads <= 0 || num_kv_heads <= 0 ||
        num_tokens <= 0 || head_dim <= 0 || aligned_head_dim < head_dim ||
        kv_stride_tokens < num_tokens || num_heads % num_kv_heads != 0) {
        return;
    }
 
    ck_attention_full_tiled_f16kv_fp32_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,
        .head_dim = head_dim,
        .aligned_head_dim = aligned_head_dim,
        .kv_stride_tokens = kv_stride_tokens,
        .query_tile_size = query_tile_size,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    const int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (pool && active > 1) {
        ck_threadpool_dispatch_n(pool, active, ck_attention_full_tiled_f16kv_fp32_work, &args);
    } else {
        ck_attention_full_tiled_f16kv_fp32_work(0, 1, &args);
    }
}

References ck_attention_full_tiled_f16kv_fp32_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by attention_forward_full_head_major_gqa_tiled336_f16kv_fp32_strided(), attention_forward_full_head_major_gqa_tiled64_f16kv_fp32_strided(), attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided(), and attention_forward_query_key_head_major_tiled_f16kv_fp32().

◆ ck_attention_forward_query_key_head_major_f32_run()

static int ck_attention_forward_query_key_head_major_f32_run	(	const float *	query,
		const float *	key,
		const float *	value,
		float *	output,
		float *	score_scratch,
		float *	key_transpose_scratch,
		int	num_heads,
		int	query_tokens,
		int	key_tokens,
		int	head_dim,
		float	scale
	)

static

Definition at line 153 of file attention_kernels.c.

{
    if (query == NULL || key == NULL || value == NULL || output == NULL ||
        score_scratch == NULL) {
        return -1;
    }
    if (num_heads <= 0 || query_tokens <= 0 || key_tokens <= 0 ||
        head_dim <= 0 || !isfinite(scale)) {
        return -2;
    }
    ck_attention_query_key_f32_args_t args = {
        .query = query,
        .key = key,
        .value = value,
        .output = output,
        .score_scratch = score_scratch,
        .key_transpose_scratch = key_transpose_scratch,
        .num_heads = num_heads,
        .query_tokens = query_tokens,
        .key_tokens = key_tokens,
        .head_dim = head_dim,
        .scale = scale,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (key_transpose_scratch != NULL) {
        int transpose_active = active;
        const int transpose_rows = num_heads * head_dim;
        if (transpose_active > transpose_rows) transpose_active = transpose_rows;
        if (pool && transpose_active > 1) {
            ck_threadpool_dispatch_n(pool, transpose_active,
                ck_attention_query_key_f32_transpose_work, &args);
        } else {
            ck_attention_query_key_f32_transpose_work(0, 1, &args);
        }
    }
    if (active > query_tokens) active = query_tokens;
    if (pool && active > 1) {
        ck_threadpool_dispatch_n(
            pool, active, ck_attention_query_key_f32_work, &args);
    } else {
        ck_attention_query_key_f32_work(0, 1, &args);
    }
    return 0;
}

References ck_attention_query_key_f32_transpose_work(), ck_attention_query_key_f32_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by attention_forward_query_key_head_major_f32(), attention_forward_query_key_head_major_f32_packed_k(), and ck_attention_query_key_f32_head_work().

◆ ck_attention_full_bf16_pytorch_flash()

static int ck_attention_full_bf16_pytorch_flash	(	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	output_token_major
	)

static

Definition at line 4944 of file attention_kernels.c.

{
    ck_attention_bf16_sdpa_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, .head_dim=head_dim,
        .aligned_head_dim=aligned_head_dim, .kv_stride_tokens=kv_stride_tokens,
        .output_token_major=output_token_major,
        .failed=0
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > num_heads) active = num_heads;
    if (pool && active > 1) {
        ck_threadpool_dispatch_n(
            pool, active, ck_attention_bf16_pytorch_flash_work, &args);
    } else {
        ck_attention_bf16_pytorch_flash_work(0, 1, &args);
    }
    return __atomic_load_n(&args.failed, __ATOMIC_RELAXED) == 0;
}

References ck_attention_bf16_pytorch_flash_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage(), and attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage_token_output().

◆ ck_attention_full_bf16_sdpa_amx_range()

static int ck_attention_full_bf16_sdpa_amx_range	(	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	head_begin,
		int	head_step,
		int	output_token_major
	)

static

Definition at line 4690 of file attention_kernels.c.

{
#if defined(__AVX512F__)
    /*
     * This is a schedule contract, not merely an AMX acceleration of the
     * portable kernel. PyTorch CPU flash attention composes FP64 scale
     * evaluation, oneDNN-style BF16 BRGEMMs, AVX-512 exp/reductions, BF16
     * probability storage, and a second BRGEMM. When parity drifts, compare
     * this complete sequence with the PyTorch and oneDNN sources before
     * changing an isolated primitive or tolerance.
     */
    enum { Q_BLOCK = 256, KV_BLOCK = 512, K_PAD = 72, D_PAD = 80 };
    if (!ck_gemm_bf16_amx_available() || head_dim != 72 || aligned_head_dim != 72 ||
        (num_tokens % 16) != 0 || (kv_stride_tokens < num_tokens) ||
        num_heads % num_kv_heads != 0) return 0;
    uint16_t *q_block = malloc((size_t)Q_BLOCK * K_PAD * sizeof(uint16_t));
    uint16_t *k_packed = malloc((size_t)num_tokens * K_PAD * sizeof(uint16_t));
    uint16_t *probabilities = malloc((size_t)Q_BLOCK * KV_BLOCK * sizeof(uint16_t));
    uint16_t *v_packed = malloc((size_t)num_tokens * D_PAD * sizeof(uint16_t));
    float *scores = malloc((size_t)Q_BLOCK * KV_BLOCK * sizeof(float));
    float *destination = malloc((size_t)Q_BLOCK * D_PAD * sizeof(float));
    float *row_max = malloc((size_t)Q_BLOCK * sizeof(float));
    float *row_sum = malloc((size_t)Q_BLOCK * sizeof(float));
    if (!q_block || !k_packed || !probabilities || !v_packed || !scores ||
        !destination || !row_max || !row_sum) {
        free(q_block); free(k_packed); free(probabilities); free(v_packed);
        free(scores); free(destination); free(row_max); free(row_sum);
        return 0;
    }
    const float scale = ck_attention_pytorch_sdpa_scale_f32(head_dim);
    const size_t q_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
    const size_t kv_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
    for (int h = head_begin; h < num_heads; h += head_step) {
        const int kv_head = (int)((long long)h * num_kv_heads / num_heads);
        const float *qh = q + (size_t)h * q_stride;
        const float *kh = k + (size_t)kv_head * kv_stride;
        const float *vh = v + (size_t)kv_head * kv_stride;
        /*
         * K and V are invariant across all query blocks for this head. Pack
         * each 512-key tile once, retaining the exact block-local layouts
         * consumed by the two certified BRGEMMs below.
         */
        for (int n = 0; n < num_tokens; n += KV_BLOCK) {
            const int key_count = num_tokens - n < KV_BLOCK ? num_tokens - n : KV_BLOCK;
            uint16_t *k_block = k_packed + (size_t)n * K_PAD;
            uint16_t *v_block = v_packed + (size_t)n * D_PAD;
            memset(k_block, 0, (size_t)key_count * K_PAD * sizeof(uint16_t));
            memset(v_block, 0, (size_t)D_PAD * key_count * sizeof(uint16_t));
            for (int key = 0; key < key_count; ++key) {
                for (int d = 0; d < head_dim; ++d) {
                    k_block[(size_t)key * K_PAD + d] =
                        float_to_bf16(kh[(size_t)(n + key) * aligned_head_dim + d]);
                    v_block[(size_t)d * key_count + key] =
                        float_to_bf16(vh[(size_t)(n + key) * aligned_head_dim + d]);
                }
            }
        }
        for (int m = 0; m < num_tokens; m += Q_BLOCK) {
            const int query_count = num_tokens - m < Q_BLOCK ? num_tokens - m : Q_BLOCK;
            memset(q_block, 0, (size_t)query_count * K_PAD * sizeof(uint16_t));
            memset(destination, 0, (size_t)query_count * D_PAD * sizeof(float));
            for (int row = 0; row < query_count; ++row) {
                for (int d = 0; d < head_dim; ++d) {
                    q_block[(size_t)row * K_PAD + d] =
                        float_to_bf16(qh[(size_t)(m + row) * aligned_head_dim + d]);
                }
                row_max[row] = -INFINITY;
                row_sum[row] = 0.0f;
            }
            for (int n = 0; n < num_tokens; n += KV_BLOCK) {
                const int key_count = num_tokens - n < KV_BLOCK ? num_tokens - n : KV_BLOCK;
                const uint16_t *k_block = k_packed + (size_t)n * K_PAD;
                const uint16_t *v_block = v_packed + (size_t)n * D_PAD;
                if (!ck_gemm_bf16_fp32out_amx_raw(
                        q_block, k_block, scores,
                        query_count, key_count, K_PAD, 0)) goto fail;
                for (int row = 0; row < query_count; ++row) {
                    float *score_row = scores + (size_t)row * key_count;
                    uint16_t *prob_row = probabilities + (size_t)row * key_count;
                    const float block_max = ck_bf16_sdpa_scale_max_avx512(
                        score_row, key_count, scale);
                    const float merged_max = row_max[row] > block_max
                        ? row_max[row] : block_max;
                    const float old_scale = isfinite(row_max[row])
                        ? expf(row_max[row] - merged_max) : 0.0f;
                    const float block_sum = ck_bf16_sdpa_exp_sum_avx512(
                        score_row, prob_row, key_count, merged_max);
                    row_sum[row] = block_sum + old_scale * row_sum[row];
                    row_max[row] = merged_max;
                    if (n > 0) {
                        __m512 scale_v = _mm512_set1_ps(old_scale);
                        int d = 0;
                        for (; d + 16 <= D_PAD; d += 16) {
                            float *dst = destination + (size_t)row * D_PAD + d;
                            _mm512_storeu_ps(dst, _mm512_mul_ps(_mm512_loadu_ps(dst), scale_v));
                        }
                    }
                }
                if (!ck_gemm_bf16_fp32out_amx_raw(
                        probabilities, v_block, destination,
                        query_count, D_PAD, key_count, n > 0)) goto fail;
            }
            for (int row = 0; row < query_count; ++row) {
                const float reciprocal = row_sum[row] == 0.0f ? 1.0f : 1.0f / row_sum[row];
                float *out_row = output + (
                    output_token_major
                        ? ((size_t)(m + row) * (size_t)num_heads + (size_t)h)
                            * (size_t)aligned_head_dim
                        : ((size_t)h * (size_t)num_tokens + (size_t)(m + row))
                            * (size_t)aligned_head_dim);
                for (int d = 0; d < head_dim; ++d) {
                    out_row[d] = bf16_to_float(
                        float_to_bf16(destination[(size_t)row * D_PAD + d] * reciprocal));
                }
            }
        }
    }
    free(q_block); free(k_packed); free(probabilities); free(v_packed);
    free(scores); free(destination); free(row_max); free(row_sum);
    return 1;
fail:
    free(q_block); free(k_packed); free(probabilities); free(v_packed);
    free(scores); free(destination); free(row_max); free(row_sum);
    return 0;
#else
    (void)q; (void)k; (void)v; (void)output; (void)num_heads;
    (void)num_kv_heads; (void)num_tokens; (void)head_dim;
    (void)aligned_head_dim; (void)kv_stride_tokens; (void)head_begin; (void)head_step;
    (void)output_token_major;
    return 0;
#endif
}

References bf16_to_float(), ck_attention_pytorch_sdpa_scale_f32(), ck_gemm_bf16_amx_available(), ck_gemm_bf16_fp32out_amx_raw(), and float_to_bf16().

Referenced by ck_attention_bf16_pytorch_flash_work().

◆ ck_attention_full_bf16_sdpa_tiled()

static int ck_attention_full_bf16_sdpa_tiled	(	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
	)

static

Definition at line 4970 of file attention_kernels.c.

{
    ck_attention_bf16_sdpa_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, .head_dim=head_dim,
        .aligned_head_dim=aligned_head_dim, .kv_stride_tokens=kv_stride_tokens,
        .output_token_major=0,
        .failed=0
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > num_heads) active = num_heads;
    if (pool && active > 1) {
        ck_threadpool_dispatch_n(pool, active, ck_attention_bf16_sdpa_work, &args);
    } else {
        ck_attention_bf16_sdpa_work(0, 1, &args);
    }
    return __atomic_load_n(&args.failed, __ATOMIC_RELAXED) == 0;
}

References ck_attention_bf16_sdpa_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by attention_forward_full_head_major_gqa_sdpa_bf16_storage().

◆ ck_attention_full_bf16_sdpa_tiled_range()

static int ck_attention_full_bf16_sdpa_tiled_range	(	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	head_begin,
		int	head_step
	)

static

Definition at line 4827 of file attention_kernels.c.

{
    if (!q || !k || !v || !output || num_heads <= 0 || num_kv_heads <= 0 ||
        head_dim <= 0 || num_tokens <= 0 || aligned_head_dim < head_dim ||
        kv_stride_tokens < num_tokens || num_heads % num_kv_heads != 0) return 0;
    if ((size_t)num_tokens > SIZE_MAX / (size_t)head_dim) return 0;
    const size_t v_col_count = (size_t)head_dim * (size_t)num_tokens;
    if (v_col_count > SIZE_MAX / sizeof(float)) return 0;
    if ((size_t)num_tokens > SIZE_MAX / (size_t)aligned_head_dim) return 0;
    if ((size_t)kv_stride_tokens > SIZE_MAX / (size_t)aligned_head_dim) return 0;
    const float scale = ck_attention_pytorch_sdpa_scale_f32(head_dim);
    const size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
    const size_t kv_head_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
    if ((size_t)num_heads > SIZE_MAX / q_head_stride) return 0;
    if ((size_t)num_kv_heads > SIZE_MAX / kv_head_stride) return 0;
    float scores[512];
    float probs[512];
    float *v_cols = (float *)malloc(v_col_count * sizeof(float));
    if (!v_cols) return 0;
    for (int h = head_begin; h < num_heads; h += head_step) {
        const int kv_head = (int)((long long)h * (long long)num_kv_heads / (long long)num_heads);
        const float *qh = q + (size_t)h * q_head_stride;
        const float *kh = k + (size_t)kv_head * kv_head_stride;
        const float *vh = v + (size_t)kv_head * kv_head_stride;
        float *oh = output + (size_t)h * q_head_stride;
        for (int d = 0; d < head_dim; ++d) {
            float *col = v_cols + (size_t)d * (size_t)num_tokens;
            for (int t = 0; t < num_tokens; ++t) {
                col[t] = vh[(size_t)t * (size_t)aligned_head_dim + (size_t)d];
            }
        }
        for (int row = 0; row < num_tokens; ++row) {
            const float *qrow = qh + (size_t)row * (size_t)aligned_head_dim;
            float *dst = oh + (size_t)row * (size_t)aligned_head_dim;
            for (int d = 0; d < head_dim; ++d) dst[d] = 0.0f;
            float running_max = -INFINITY;
            float running_sum = 0.0f;
            for (int n = 0; n < num_tokens; n += 512) {
                const int block = num_tokens - n < 512 ? num_tokens - n : 512;
                float block_max = -INFINITY;
                for (int j = 0; j < block; ++j) {
                    scores[j] = ck_bf16_dot_contract(
                        qrow,
                        kh + (size_t)(n + j) * (size_t)aligned_head_dim,
                        head_dim) * scale;
                    if (scores[j] > block_max) block_max = scores[j];
                }
                const float merged_max = running_max > block_max ? running_max : block_max;
                const float old_scale = isfinite(running_max) ? expf(running_max - merged_max) : 0.0f;
                float block_sum = 0.0f;
                for (int j = 0; j < block; ++j) {
                    const float p = expf(scores[j] - merged_max);
                    block_sum += p;
                    probs[j] = bf16_to_float(float_to_bf16(p));
                }
                for (int d = 0; d < head_dim; ++d) {
                    const float partial = ck_bf16_dot_contract(
                        probs,
                        v_cols + (size_t)d * (size_t)num_tokens + (size_t)n,
                        block);
                    dst[d] = dst[d] * old_scale + partial;
                }
                running_sum = running_sum * old_scale + block_sum;
                running_max = merged_max;
            }
            const float inv = running_sum > 0.0f ? 1.0f / running_sum : 0.0f;
            for (int d = 0; d < head_dim; ++d) dst[d] *= inv;
            for (int d = head_dim; d < aligned_head_dim; ++d) dst[d] = 0.0f;
        }
    }
    free(v_cols);
    return 1;
}

References bf16_to_float(), ck_attention_pytorch_sdpa_scale_f32(), ck_bf16_dot_contract(), and float_to_bf16().

Referenced by ck_attention_bf16_sdpa_work().

◆ ck_attention_full_grid_work()

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

static

Definition at line 4033 of file attention_kernels.c.

{
    ck_attention_parallel_args_t *args = (ck_attention_parallel_args_t *) opaque;
    const int T = args->num_tokens;
    const int total = args->num_heads * T;
    const size_t kv_head_stride = (size_t) args->kv_stride_tokens * (size_t) args->aligned_head_dim;
 
    for (int idx = ith; idx < total; idx += nth) {
        const int h = idx / T;
        const int i = idx - h * T;
        const int kv_head = (int) ((long long) h * (long long) args->num_kv_heads / (long long) args->num_heads);
        const float *k_head = args->k + (size_t) kv_head * kv_head_stride;
        const float *v_head = args->v + (size_t) kv_head * kv_head_stride;
        const float *q_vec = args->q + qkv_index(h, i, 0, T, args->aligned_head_dim);
        float *out_vec = args->output + attention_output_index(
            h, i, args->num_heads, T, args->aligned_head_dim,
            args->output_token_major);
        const int kv_tokens = args->causal ? (i + 1) : T;
        ck_attention_flash_query_auto(q_vec, k_head, v_head,
                                      kv_tokens,
                                      args->head_dim,
                                      args->aligned_head_dim,
                                      args->scale,
                                      out_vec);
    }
}

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

Referenced by attention_forward_head_major_gqa_flash_impl().

◆ ck_attention_full_tiled_f16kv_fp32_range()

static void ck_attention_full_tiled_f16kv_fp32_range	(	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	query_tile_size,
		int	ith,
		int	nth
	)

static

Definition at line 3388 of file attention_kernels.c.

{
    if (!q || !k || !v || !output || num_heads <= 0 || num_kv_heads <= 0 ||
        num_tokens <= 0 || head_dim <= 0 || aligned_head_dim < head_dim ||
        kv_stride_tokens < num_tokens || num_heads % num_kv_heads != 0 ||
        query_tile_size <= 0) {
        return;
    }
 
    const float scale = ck_attention_strict_scale_f32(head_dim);
    const int T = num_tokens;
    const size_t kv_head_stride = (size_t) kv_stride_tokens * (size_t) aligned_head_dim;
 
    float *q_tile = (float *) alloca((size_t) query_tile_size * (size_t) head_dim * sizeof(float));
    float *k_tile = (float *) alloca((size_t) head_dim * (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
    float *v_tile = (float *) alloca((size_t) CK_GGML_FA_TILE_KV * (size_t) head_dim * sizeof(float));
    float *kq = (float *) alloca((size_t) query_tile_size * (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
    float *vkq = (float *) alloca((size_t) query_tile_size * (size_t) head_dim * sizeof(float));
    float *sum_row = (float *) alloca((size_t) query_tile_size * sizeof(float));
    float *max_row = (float *) alloca((size_t) query_tile_size * sizeof(float));
 
    const int total_rows = num_heads * T;
    const int rows_per_worker = (total_rows + nth - 1) / nth;
    int ir = rows_per_worker * ith;
    const int ir1 = (ir + rows_per_worker) < total_rows
        ? (ir + rows_per_worker)
        : total_rows;
 
    while (ir < ir1) {
        const int h = ir / T;
        const int iq = ir - h * T;
        int tile_rows = ir1 - ir;
        if (tile_rows > query_tile_size) tile_rows = query_tile_size;
        if (tile_rows > T - iq) tile_rows = T - iq;
        const 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 tq = 0; tq < query_tile_size; ++tq) {
            sum_row[tq] = 0.0f;
            max_row[tq] = -INFINITY;
        }
 
        memset(vkq, 0, (size_t) query_tile_size * (size_t) head_dim * sizeof(float));
        memset(q_tile, 0, (size_t) query_tile_size * (size_t) head_dim * sizeof(float));
        memset(k_tile, 0, (size_t) head_dim * (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
        memset(v_tile, 0, (size_t) CK_GGML_FA_TILE_KV * (size_t) head_dim * sizeof(float));
 
        for (int tq = 0; tq < tile_rows; ++tq) {
            const float *q_vec = q + qkv_index(h, iq + tq, 0, T, aligned_head_dim);
            memcpy(q_tile + (size_t) tq * (size_t) head_dim, q_vec, (size_t) head_dim * sizeof(float));
        }
 
        for (int ik = 0; ik < T; ik += CK_GGML_FA_TILE_KV) {
                const int kv_tile = (T - ik) < CK_GGML_FA_TILE_KV ? (T - ik) : CK_GGML_FA_TILE_KV;
                memset(kq, 0, (size_t) query_tile_size * (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
 
                for (int tk = 0; tk < kv_tile; ++tk) {
                    const float *k_vec = k_head + (size_t) (ik + tk) * (size_t) aligned_head_dim;
                    const float *v_vec = v_head + (size_t) (ik + tk) * (size_t) aligned_head_dim;
                    for (int d = 0; d < head_dim; ++d) {
                        k_tile[(size_t) d * (size_t) CK_GGML_FA_TILE_KV + (size_t) tk] =
                            ck_round_fp16_scalar(k_vec[d]);
                        v_tile[(size_t) tk * (size_t) head_dim + (size_t) d] =
                            ck_round_fp16_scalar(v_vec[d]);
                    }
                }
 
                ck_attention_matmul_f32_accum(kq,
                                              q_tile,
                                              k_tile,
                                              query_tile_size,
                                              head_dim,
                                              CK_GGML_FA_TILE_KV);
                ck_vec_scale_f32_inplace(kq,
                                         query_tile_size * CK_GGML_FA_TILE_KV,
                                         scale);
 
                if (kv_tile < CK_GGML_FA_TILE_KV) {
                    for (int tq = 0; tq < query_tile_size; ++tq) {
                        float *kq_row = kq + (size_t) tq * (size_t) CK_GGML_FA_TILE_KV;
                        for (int tk = kv_tile; tk < CK_GGML_FA_TILE_KV; ++tk) {
                            kq_row[tk] = -INFINITY;
                        }
                    }
                }
 
                for (int tq = 0; tq < tile_rows; ++tq) {
                    float *kq_row = kq + (size_t) tq * (size_t) CK_GGML_FA_TILE_KV;
                    const float tile_max = ck_vec_max_f32_contig(kq_row, CK_GGML_FA_TILE_KV);
                    if (tile_max == -INFINITY) {
                        memset(kq_row, 0, (size_t) CK_GGML_FA_TILE_KV * sizeof(float));
                        continue;
                    }
 
                    const float old_max = max_row[tq];
                    const float new_max = old_max > tile_max ? old_max : tile_max;
                    if (new_max > old_max) {
                        const float ms = ck_attention_reference_expf(old_max - new_max);
                        ck_vec_scale_f32_inplace(vkq + (size_t) tq * (size_t) head_dim,
                                                 head_dim,
                                                 ms);
                        sum_row[tq] *= ms;
                    }
                    max_row[tq] = new_max;
                    sum_row[tq] = (float) (
                        (double) sum_row[tq] +
                        ck_ggml_vec_soft_max_row(
                            CK_GGML_FA_TILE_KV, kq_row, kq_row, new_max));
                }
 
                ck_attention_matmul_f32_accum(vkq,
                                              kq,
                                              v_tile,
                                              query_tile_size,
                                              CK_GGML_FA_TILE_KV,
                                              head_dim);
            }
 
        for (int tq = 0; tq < tile_rows; ++tq) {
            float *out_vec = output + qkv_index(h, iq + tq, 0, T, aligned_head_dim);
            const float inv_sum = sum_row[tq] == 0.0f ? 0.0f : (1.0f / sum_row[tq]);
            for (int d = 0; d < head_dim; ++d) {
                out_vec[d] = vkq[(size_t) tq * (size_t) head_dim + (size_t) d] * inv_sum;
            }
            for (int d = head_dim; d < aligned_head_dim; ++d) {
                out_vec[d] = 0.0f;
            }
        }
        ir += tile_rows;
    }
}

References ck_attention_matmul_f32_accum(), ck_attention_reference_expf(), ck_attention_strict_scale_f32(), CK_GGML_FA_TILE_KV, ck_ggml_vec_soft_max_row(), ck_round_fp16_scalar(), ck_vec_max_f32_contig(), ck_vec_scale_f32_inplace(), and qkv_index().

Referenced by ck_attention_full_tiled_f16kv_fp32_work().

◆ ck_attention_full_tiled_f16kv_fp32_work()

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

static

Definition at line 3548 of file attention_kernels.c.

{
    ck_attention_full_tiled_f16kv_fp32_args_t *args =
        (ck_attention_full_tiled_f16kv_fp32_args_t *) opaque;
    ck_attention_full_tiled_f16kv_fp32_range(
        args->q, args->k, args->v, args->output,
        args->num_heads, args->num_kv_heads, args->num_tokens,
        args->head_dim, args->aligned_head_dim, args->kv_stride_tokens,
        args->query_tile_size,
        ith, nth);
}

References ck_attention_full_tiled_f16kv_fp32_range().

Referenced by ck_attention_forward_full_head_major_gqa_tiled_f16kv_fp32_strided().

◆ ck_attention_ggml_out_graph_enabled()

static int ck_attention_ggml_out_graph_enabled ( void )

static

Definition at line 500 of file attention_kernels.c.

{
    return 0;
}

Referenced by attention_forward_full_head_major_gqa_ggml_strided_workspace().

◆ ck_attention_gqa_team_barrier_wait()

static void ck_attention_gqa_team_barrier_wait ( ck_attention_gqa_team_barrier_t * barrier )

inlinestatic

Definition at line 7091 of file attention_kernels.c.

{
    const int phase = __atomic_load_n(&barrier->phase, __ATOMIC_RELAXED);
    if (__atomic_fetch_add(&barrier->arrived, 1, __ATOMIC_ACQ_REL) ==
        barrier->lanes - 1) {
        __atomic_store_n(&barrier->arrived, 0, __ATOMIC_RELAXED);
        __atomic_store_n(&barrier->phase, phase + 1, __ATOMIC_RELEASE);
        return;
    }
    while (__atomic_load_n(&barrier->phase, __ATOMIC_ACQUIRE) == phase) {
#if defined(__i386__) || defined(__x86_64__)
        _mm_pause();
#endif
    }
}

Referenced by ck_attention_f16_prefill_gqa_reuse_work().

◆ ck_attention_llama_regular_impl()

static void ck_attention_llama_regular_impl	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	query_tokens,
		int	live_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	kv_stride_tokens,
		int	sliding_window,
		float *	scores,
		size_t	scores_bytes,
		float *	value_columns,
		size_t	value_columns_bytes,
		float *	scaled_scores,
		size_t	scaled_scores_bytes,
		int	batched_prefill
	)

static

Definition at line 5968 of file attention_kernels.c.

{
    if (!q || !k || !v || !output || !scores || !value_columns || !scaled_scores ||
        num_heads <= 0 || num_kv_heads <= 0 || query_tokens <= 0 || live_tokens <= 0 ||
        head_dim <= 0 || aligned_head_dim < head_dim || kv_stride_tokens < live_tokens) return;
    size_t scratch_capacity = scores_bytes / sizeof(float);
    const size_t scaled_capacity = scaled_scores_bytes / sizeof(float);
    const size_t value_capacity = value_columns_bytes /
        ((size_t)head_dim * sizeof(float));
    if (scaled_capacity < scratch_capacity) scratch_capacity = scaled_capacity;
    if (value_capacity < scratch_capacity) scratch_capacity = value_capacity;
    if (!batched_prefill && (size_t)kv_stride_tokens < scratch_capacity) {
        scratch_capacity = (size_t)kv_stride_tokens;
    }
    const int padded = scratch_capacity > (size_t)INT_MAX ? 0 :
        ck_llama_kv_pad_256(live_tokens, (int)scratch_capacity);
    if (padded <= 0 || scores_bytes < (size_t)padded * sizeof(float) ||
        scaled_scores_bytes < (size_t)padded * sizeof(float) ||
        value_columns_bytes < (size_t)head_dim * (size_t)padded * sizeof(float)) return;
 
    const size_t kv_head_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
    const int debug_layer_id = ck_attention_vec_dump_enabled()
        ? ck_attention_vec_dump_next_layer_id() : -1;
    int packed_kv_head = -1;
    for (int head = 0; head < num_heads; ++head) {
        const int kv_head = (int)((long long)head * num_kv_heads / num_heads);
        const float *key_head = k + (size_t)kv_head * kv_head_stride;
        const float *value_head = v + (size_t)kv_head * kv_head_stride;
        if (kv_head != packed_kv_head) {
            for (int dim = 0; dim < head_dim; ++dim) {
                float *column = value_columns + (size_t)dim * (size_t)padded;
                int token = 0;
                for (; token < live_tokens; ++token) {
                    column[token] = value_head[(size_t)token * aligned_head_dim + dim];
                }
                for (; token < padded; ++token) column[token] = 0.0f;
            }
            packed_kv_head = kv_head;
        }
        for (int query = 0; query < query_tokens; ++query) {
            const int position = batched_prefill ? query : live_tokens - 1;
            ck_attention_llama_regular_query(
                q + ((size_t)head * query_tokens + query) * aligned_head_dim,
                key_head, value_columns,
                output + ((size_t)head * query_tokens + query) * aligned_head_dim,
                scores, scaled_scores, live_tokens, padded, position,
                head_dim, aligned_head_dim, sliding_window, batched_prefill);
            if (debug_layer_id >= 0 &&
                ck_attention_vec_dump_should_emit(debug_layer_id, head, query)) {
                ck_attention_vec_dump_selected_query(
                    scaled_scores, scores,
                    output + ((size_t)head * query_tokens + query) * aligned_head_dim,
                    value_columns, padded, head_dim,
                    debug_layer_id, head, query);
            }
        }
    }
}

References ck_attention_llama_regular_query(), ck_attention_vec_dump_enabled(), ck_attention_vec_dump_next_layer_id(), ck_attention_vec_dump_selected_query(), ck_attention_vec_dump_should_emit(), ck_llama_kv_pad_256(), and token.

Referenced by attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace(), and attention_forward_decode_head_major_gqa_llama_regular_sliding_workspace().

◆ ck_attention_llama_regular_query()

static void ck_attention_llama_regular_query	(	const float *	query,
		const float *	key_head,
		const float *	value_columns,
		float *	output,
		float *	scores,
		float *	scaled_scores,
		int	live_tokens,
		int	padded_tokens,
		int	query_position,
		int	head_dim,
		int	aligned_head_dim,
		int	sliding_window,
		int	batched_prefill
	)

static

Definition at line 5930 of file attention_kernels.c.

{
    const int first = sliding_window > 0 && query_position >= sliding_window
        ? query_position - sliding_window + 1 : 0;
    const int last = query_position < live_tokens ? query_position : live_tokens - 1;
    for (int token = 0; token < padded_tokens; ++token) scores[token] = -INFINITY;
    for (int token = first; token <= last; ++token) {
        const float *key = key_head + (size_t)token * (size_t)aligned_head_dim;
        scores[token] = batched_prefill
            ? ck_llama_regular_gemm_f16(query, key, head_dim)
            : ck_llama_regular_dot_f16(query, key, head_dim);
    }
 
    memcpy(scaled_scores, scores, (size_t)padded_tokens * sizeof(float));
    const float scale = 1.0f / sqrtf((float)head_dim);
    ck_vec_scale_f32_inplace(scaled_scores, padded_tokens, scale);
    const float maximum = ck_vec_max_f32_contig(scaled_scores, padded_tokens);
    const double total = ck_ggml_vec_soft_max_row(
        padded_tokens, scores, scaled_scores, maximum);
    if (total > 0.0) {
        ck_vec_scale_f32_inplace(scores, padded_tokens, (float)(1.0 / total));
        for (int dim = 0; dim < head_dim; ++dim) {
            const float *column = value_columns + (size_t)dim * (size_t)padded_tokens;
            output[dim] = batched_prefill
                ? ck_llama_regular_gemm_f16(scores, column, padded_tokens)
                : ck_llama_regular_dot_f16(scores, column, padded_tokens);
        }
    } else {
        memset(output, 0, (size_t)head_dim * sizeof(float));
    }
    for (int dim = head_dim; dim < aligned_head_dim; ++dim) output[dim] = 0.0f;
}

References ck_ggml_vec_soft_max_row(), ck_llama_regular_dot_f16(), ck_llama_regular_gemm_f16(), ck_vec_max_f32_contig(), ck_vec_scale_f32_inplace(), and token.

Referenced by ck_attention_llama_regular_impl(), and ck_llama_regular_prefill_work().

◆ ck_attention_mad_f16_llama()

static void ck_attention_mad_f16_llama	(	uint16_t *	y,
		const uint16_t *	x,
		float	scale,
		int	n
	)

inlinestatic

Definition at line 6543 of file attention_kernels.c.

{
    int i = 0;
#if defined(__AVX512F__)
    const __m512 factor = _mm512_set1_ps(scale);
    for (; i + 63 < n; i += 64) {
        for (int lane = 0; lane < 4; ++lane) {
            const int offset = i + lane * 16;
            const __m512 xv = _mm512_cvtph_ps(
                _mm256_loadu_si256((const __m256i *) (x + offset)));
            __m512 yv = _mm512_cvtph_ps(
                _mm256_loadu_si256((const __m256i *) (y + offset)));
            yv = _mm512_fmadd_ps(xv, factor, yv);
            _mm256_storeu_si256(
                (__m256i *) (y + offset), _mm512_cvtps_ph(yv, 0));
        }
    }
#elif defined(__AVX2__) && defined(__F16C__)
    const __m256 factor = _mm256_set1_ps(scale);
    for (; i + 31 < n; i += 32) {
        for (int lane = 0; lane < 4; ++lane) {
            const int offset = i + lane * 8;
            const __m256 xv = _mm256_cvtph_ps(
                _mm_loadu_si128((const __m128i *) (x + offset)));
            __m256 yv = _mm256_cvtph_ps(
                _mm_loadu_si128((const __m128i *) (y + offset)));
#if defined(__FMA__)
            yv = _mm256_fmadd_ps(xv, factor, yv);
#else
            yv = _mm256_add_ps(yv, _mm256_mul_ps(xv, factor));
#endif
            _mm_storeu_si128(
                (__m128i *) (y + offset), _mm256_cvtps_ph(yv, 0));
        }
    }
#endif
    for (; i < n; ++i) {
        const float product = CK_FP16_TO_FP32(x[i]) * scale;
        const float updated = CK_FP16_TO_FP32(y[i]) + product;
        y[i] = CK_FP32_TO_FP16(updated);
    }
}

References CK_FP16_TO_FP32, CK_FP32_TO_FP16, and product().

Referenced by ck_attention_f16_split_work().

◆ ck_attention_matmul_f32_accum()

static void ck_attention_matmul_f32_accum	(	float *	c,
		const float *	a,
		const float *	b,
		int	m,
		int	k,
		int	n
	)

static

Definition at line 3367 of file attention_kernels.c.

{
    for (int i = 0; i < m; ++i) {
        float *c_row = c + (size_t) i * (size_t) n;
        const float *a_row = a + (size_t) i * (size_t) k;
        for (int kk = 0; kk < k; ++kk) {
            const float a_ik = a_row[kk];
            const float *b_row = b + (size_t) kk * (size_t) n;
            for (int j = 0; j < n; ++j) {
                c_row[j] += a_ik * b_row[j];
            }
        }
    }
}

Referenced by ck_attention_f16_prefill_gqa_reuse_work(), ck_attention_f16_prefill_qtile64_work(), and ck_attention_full_tiled_f16kv_fp32_range().

◆ ck_attention_muse_eager_bf16_storage_impl()

static void ck_attention_muse_eager_bf16_storage_impl	(	const float *	q,
		const float *	k,
		const float *	v,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	query_tokens,
		int	kv_tokens,
		int	kv_stride_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	sliding_window,
		float *	scores,
		size_t	scores_bytes,
		uint16_t *	key_rows,
		size_t	key_rows_bytes,
		uint16_t *	value_columns,
		size_t	value_columns_bytes,
		uint16_t *	gemm_rows,
		size_t	gemm_rows_bytes,
		uint16_t *	gemm_scores,
		size_t	gemm_scores_bytes
	)

static

Definition at line 9053 of file attention_kernels.c.

{
#if defined(__AVX512F__) && defined(USE_ONEDNN)
    if (!q || !k || !v || !output || num_heads <= 0 || num_kv_heads <= 0 ||
        query_tokens <= 0 || kv_tokens < query_tokens ||
        kv_stride_tokens < kv_tokens || head_dim <= 0 ||
        aligned_head_dim != head_dim || num_heads % num_kv_heads != 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid Muse eager attention dimensions\n");
        abort();
    }
 
    const int workspace_rows = query_tokens < CK_MUSE_PREFILL_QUERY_TILE
        ? query_tokens : CK_MUSE_PREFILL_QUERY_TILE;
    const size_t score_count = ck_muse_attention_checked_product(
        (size_t)workspace_rows, (size_t)kv_tokens, "score");
    const size_t kv_matrix_count = ck_muse_attention_checked_product(
        (size_t)kv_tokens, (size_t)head_dim, "KV matrix");
    const size_t row_count = ck_muse_attention_checked_product(
        (size_t)workspace_rows, (size_t)head_dim, "row matrix");
    if (score_count > SIZE_MAX / sizeof(*scores) ||
        kv_matrix_count > SIZE_MAX / sizeof(*key_rows) ||
        row_count > SIZE_MAX / sizeof(*gemm_rows)) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: Muse eager attention workspace "
                "byte size overflow\n");
        abort();
    }
    if (!scores || scores_bytes < score_count * sizeof(*scores) ||
        !key_rows || key_rows_bytes < kv_matrix_count * sizeof(*key_rows) ||
        !value_columns ||
            value_columns_bytes < kv_matrix_count * sizeof(*value_columns) ||
        !gemm_rows || gemm_rows_bytes < row_count * sizeof(*gemm_rows) ||
        !gemm_scores ||
            gemm_scores_bytes < score_count * sizeof(*gemm_scores)) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: Muse eager attention received "
                "an undersized planner-owned workspace\n");
        abort();
    }
    pthread_once(&ck_pytorch_attention_once, ck_bind_pytorch_attention_primitives);
    if (!ck_pytorch_attention_expf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: Muse eager attention requires "
                "Sleef_expf16_u10 from CK_SLEEF_LIBRARY\n");
        abort();
    }
 
    const float scale = 1.0f / sqrtf((float)head_dim);
    const int heads_per_kv = num_heads / num_kv_heads;
    const int query_position_offset = kv_tokens - query_tokens;
    const size_t query_head_stride = ck_muse_attention_checked_product(
        (size_t)query_tokens, (size_t)aligned_head_dim, "query stride");
    const size_t kv_head_stride = ck_muse_attention_checked_product(
        (size_t)kv_stride_tokens, (size_t)aligned_head_dim, "KV stride");
    (void)ck_muse_attention_checked_product(
        (size_t)num_heads, query_head_stride, "query tensor");
    (void)ck_muse_attention_checked_product(
        (size_t)num_kv_heads, kv_head_stride, "KV tensor");
 
    for (int head = 0; head < num_heads; ++head) {
        const int kv_head = head / heads_per_kv;
        const float *key_head = k + (size_t)kv_head * kv_head_stride;
        const float *value_head = v + (size_t)kv_head * kv_head_stride;
        for (int token = 0; token < kv_tokens; ++token) {
            const float *key_row = key_head +
                (size_t)token * (size_t)aligned_head_dim;
            const float *value_row = value_head +
                (size_t)token * (size_t)aligned_head_dim;
            for (int dim = 0; dim < head_dim; ++dim) {
                key_rows[(size_t)token * (size_t)head_dim + (size_t)dim] =
                    float_to_bf16(key_row[dim]);
                value_columns[(size_t)dim * (size_t)kv_tokens + (size_t)token] =
                    float_to_bf16(value_row[dim]);
            }
        }
 
        const float *query_head = q + (size_t)head * query_head_stride;
        float *output_head = output + (size_t)head * query_head_stride;
        for (int query_start = 0; query_start < query_tokens;
             query_start += CK_MUSE_PREFILL_QUERY_TILE) {
            const int tile_rows = query_tokens - query_start < CK_MUSE_PREFILL_QUERY_TILE
                ? query_tokens - query_start : CK_MUSE_PREFILL_QUERY_TILE;
            gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace(
                query_head + (size_t)query_start * (size_t)head_dim,
                key_rows, NULL, scores,
                tile_rows, kv_tokens, head_dim,
                gemm_rows, gemm_rows_bytes, gemm_scores, gemm_scores_bytes,
                NULL, 0);
 
            for (int query_token = 0; query_token < tile_rows; ++query_token) {
                const int query_position = query_position_offset + query_start + query_token;
                const int visible_end = query_position + 1;
                int visible_begin = 0;
                if (sliding_window > 0 && visible_end > sliding_window) {
                    visible_begin = visible_end - sliding_window;
                }
                float *score_row = scores +
                    (size_t)query_token * (size_t)kv_tokens;
                for (int token = 0; token < kv_tokens; ++token) {
                    score_row[token] = bf16_to_float(
                        float_to_bf16(score_row[token] * scale));
                    if (token < visible_begin || token >= visible_end) {
                        score_row[token] = -INFINITY;
                    }
                }
                (void)ck_pytorch_softmax_f32(score_row, kv_tokens);
            }
            gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace(
                scores, value_columns, NULL,
                output_head + (size_t)query_start * (size_t)head_dim,
                tile_rows, head_dim, kv_tokens,
                gemm_scores, gemm_scores_bytes, gemm_rows, gemm_rows_bytes,
                NULL, 0);
        }
    }
#else
    (void)q; (void)k; (void)v; (void)output; (void)num_heads;
    (void)num_kv_heads; (void)query_tokens; (void)kv_tokens;
    (void)kv_stride_tokens; (void)head_dim; (void)aligned_head_dim;
    (void)sliding_window;
    (void)scores; (void)scores_bytes; (void)key_rows; (void)key_rows_bytes;
    (void)value_columns; (void)value_columns_bytes;
    (void)gemm_rows; (void)gemm_rows_bytes;
    (void)gemm_scores; (void)gemm_scores_bytes;
    fprintf(stderr,
            "HARD KERNEL CONTRACT FAULT: Muse eager attention requires "
            "AVX-512 and USE_ONEDNN=1\n");
    abort();
#endif
}

References bf16_to_float(), ck_muse_attention_checked_product(), CK_MUSE_PREFILL_QUERY_TILE, float_to_bf16(), gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage_workspace(), and token.

Referenced by attention_forward_causal_head_major_gqa_muse_eager_bf16_storage(), attention_forward_causal_head_major_gqa_muse_eager_bf16_storage_sliding(), attention_forward_decode_head_major_gqa_muse_eager_bf16_storage(), and attention_forward_decode_head_major_gqa_muse_eager_bf16_storage_sliding().

◆ ck_attention_parallel_enabled()

static int ck_attention_parallel_enabled	(	int	total_queries,
		int	num_tokens,
		int	head_dim
	)

static

Definition at line 4060 of file attention_kernels.c.

{
    const char *disable = getenv("CK_DISABLE_ATTENTION_THREADPOOL");
    if (disable && disable[0] && strcmp(disable, "0") != 0) return 0;
    /*
     * A 128-token GQA prefill can fall below the old 2048-query threshold
     * even though its causal dot products contain ample independent work.
     * Keep smaller prompts serial, but allow the shared pool once at least
     * eight 128-query worker grains are available.
     */
    if (total_queries < 1024 || num_tokens < 128 || head_dim <= 0) return 0;
    return 1;
}

Referenced by attention_forward_head_major_gqa_flash_impl().

◆ ck_attention_pick_active_threads()

static int ck_attention_pick_active_threads	(	const ck_threadpool_t *	pool,
		int	total_queries,
		int	num_tokens
	)

static

Definition at line 4074 of file attention_kernels.c.

{
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
    if (nth <= 1) return 1;
    const char *cap_env = getenv("CK_ATTENTION_THREAD_CAP");
    int cap = cap_env && cap_env[0] ? atoi(cap_env) : 24;
    if (cap < 1) cap = 1;
    if (cap > nth) cap = nth;
    int active = (total_queries + 127) / 128;
    if (num_tokens >= 1024 && active < 8) active = 8;
    if (active > cap) active = cap;
    if (active > nth) active = nth;
    return active < 1 ? 1 : active;
}

References ck_threadpool_n_threads().

Referenced by attention_forward_head_major_gqa_flash_impl().

◆ ck_attention_prefill_append_f16cache_gemma4_workspace()

static ck_attention_status_t ck_attention_prefill_append_f16cache_gemma4_workspace	(	const float *	q,
		const uint16_t *	k_cache,
		const uint16_t *	v_cache,
		float *	output,
		int	num_heads,
		int	num_kv_heads,
		int	q_tokens,
		int	past_tokens,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim,
		int	sliding_window,
		ck_attention_reduction_t	reduction,
		float *	token_workspace,
		size_t	token_workspace_bytes,
		const int *	segment_lengths,
		int	num_segments
	)

static

Definition at line 7854 of file attention_kernels.c.

{
    if (!q || !k_cache || !v_cache || !output || num_heads <= 0 ||
        num_kv_heads <= 0 || q_tokens <= 0 || past_tokens < 0 ||
        cache_capacity <= 0 || past_tokens > cache_capacity ||
        q_tokens > cache_capacity - past_tokens || head_dim <= 0 ||
        aligned_head_dim < head_dim || sliding_window < 0) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    int one_segment = q_tokens;
    if (!segment_lengths && num_segments == 0) {
        segment_lengths = &one_segment;
        num_segments = 1;
    }
    if (!segment_lengths || num_segments <= 0 ||
        (num_segments > q_tokens && num_segments > 3)) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    int planned_rows = 0;
    for (int s = 0; s < num_segments; ++s) {
        if (segment_lengths[s] < 0 || segment_lengths[s] > q_tokens - planned_rows) {
            return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
        }
        planned_rows += segment_lengths[s];
    }
    if (planned_rows != q_tokens) return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    if ((size_t) num_heads > SIZE_MAX / (size_t) aligned_head_dim) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    const size_t token_elems = (size_t) num_heads * (size_t) aligned_head_dim;
    if (token_elems > SIZE_MAX / (2 * sizeof(float)) || !token_workspace ||
        token_workspace_bytes < 2 * token_elems * sizeof(float)) {
        return CK_ATTENTION_STATUS_INVALID_ARGUMENT;
    }
    float *q_token = token_workspace;
    float *out_token = token_workspace + token_elems;
    int row_offset = 0;
    for (int s = 0; s < num_segments; ++s) {
        const int rows = segment_lengths[s];
        if (rows == 0) continue;
        if ((reduction == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64 ||
             reduction == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64_REGULAR_TAIL_NOSCALE) &&
            rows >= CK_GGML_FA_TILE_Q &&
            (sliding_window == 0 || past_tokens + row_offset + rows <= sliding_window)) {
            const ck_attention_status_t status = ck_attention_f16_prefill_qtile64_dispatch(
                q + (size_t) row_offset * (size_t) aligned_head_dim,
                k_cache, v_cache,
                output + (size_t) row_offset * (size_t) aligned_head_dim,
                num_heads, num_kv_heads, rows, past_tokens + row_offset,
                cache_capacity, head_dim, aligned_head_dim, 0, 1.0f,
                (size_t) q_tokens * (size_t) aligned_head_dim,
                (size_t) q_tokens * (size_t) aligned_head_dim,
                CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES);
            if (status != CK_ATTENTION_STATUS_OK) return status;
            row_offset += rows;
            continue;
        }
        const ck_attention_reduction_t selected =
            reduction == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64
                ? CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE :
            reduction == CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64_REGULAR_TAIL_NOSCALE
                ? CK_ATTN_REDUCTION_F16_REGULAR_NOSCALE : reduction;
        for (int local_row = 0; local_row < rows; ++local_row) {
            const int t = row_offset + local_row;
            for (int h = 0; h < num_heads; ++h) {
                const float *src = q +
                    ((size_t) h * (size_t) q_tokens + (size_t) t) *
                    (size_t) aligned_head_dim;
                memcpy(q_token + (size_t) h * (size_t) aligned_head_dim, src,
                       (size_t) aligned_head_dim * sizeof(float));
            }
            const ck_attention_status_t status = ck_attention_forward_decode_f16cache_gemma4(
                q_token, k_cache, v_cache, out_token, num_heads, num_kv_heads,
                past_tokens + t + 1, cache_capacity, head_dim, aligned_head_dim,
                sliding_window, selected);
            if (status != CK_ATTENTION_STATUS_OK) {
                return status;
            }
            for (int h = 0; h < num_heads; ++h) {
                float *dst = output +
                    ((size_t) h * (size_t) q_tokens + (size_t) t) *
                    (size_t) aligned_head_dim;
                memcpy(dst, out_token + (size_t) h * (size_t) aligned_head_dim,
                       (size_t) aligned_head_dim * sizeof(float));
            }
        }
        row_offset += rows;
    }
    return CK_ATTENTION_STATUS_OK;
}

References ck_attention_f16_prefill_qtile64_dispatch(), ck_attention_forward_decode_f16cache_gemma4(), CK_ATTENTION_STATUS_INVALID_ARGUMENT, CK_ATTENTION_STATUS_OK, CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES, CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64, CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64_REGULAR_TAIL_NOSCALE, CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE, CK_ATTN_REDUCTION_F16_REGULAR_NOSCALE, and CK_GGML_FA_TILE_Q.

Referenced by attention_forward_causal_head_major_gqa_prefill_append_f16cache_gemma4_workspace(), attention_forward_causal_head_major_gqa_prefill_append_f16cache_sliding_gemma4_workspace(), attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_gemma4_workspace(), and attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_sliding_gemma4_workspace().

◆ ck_attention_pytorch_sdpa_scale_f32()

float ck_attention_pytorch_sdpa_scale_f32 ( int head_dim )

Definition at line 4592 of file attention_kernels.c.

{
    /* PyTorch CPU flash calculate_scale evaluates in FP64, then narrows. */
    return head_dim > 0 ? (float)(1.0 / sqrt((double)head_dim)) : 0.0f;
}

Referenced by ck_attention_full_bf16_sdpa_amx_range(), and ck_attention_full_bf16_sdpa_tiled_range().

◆ ck_attention_query_key_f32_head_work()

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

static

Definition at line 240 of file attention_kernels.c.

{
    ck_attention_query_key_f32_head_args_t *args =
        (ck_attention_query_key_f32_head_args_t *)opaque;
    const size_t query_stride =
        (size_t)args->query_tokens * (size_t)args->head_dim;
    const size_t kv_stride =
        (size_t)args->key_tokens * (size_t)args->head_dim;
    for (int head = ith; head < args->num_heads; head += nth) {
        ck_attention_forward_query_key_head_major_f32_run(
            args->query + (size_t)head * query_stride,
            args->key + (size_t)head * kv_stride,
            args->value + (size_t)head * kv_stride,
            args->output + (size_t)head * query_stride,
            args->score_scratch + (size_t)head * args->key_tokens,
            NULL, 1, args->query_tokens, args->key_tokens,
            args->head_dim, args->scale);
    }
}

References ck_attention_forward_query_key_head_major_f32_run().

Referenced by attention_forward_query_key_head_major_f32_decode_heads().

◆ ck_attention_query_key_f32_transpose_work()

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

static

Definition at line 59 of file attention_kernels.c.

{
    ck_attention_query_key_f32_args_t *args =
        (ck_attention_query_key_f32_args_t *)opaque;
    const int rows = args->num_heads * args->head_dim;
    for (int row = ith; row < rows; row += nth) {
        const int head = row / args->head_dim;
        const int dim = row % args->head_dim;
        const float *key_head = args->key +
            (size_t)head * args->key_tokens * args->head_dim;
        float *packed = args->key_transpose_scratch +
            (size_t)row * args->key_tokens;
        for (int key_token = 0; key_token < args->key_tokens; ++key_token) {
            packed[key_token] =
                key_head[(size_t)key_token * args->head_dim + dim];
        }
    }
}

Referenced by ck_attention_forward_query_key_head_major_f32_run().

◆ ck_attention_query_key_f32_work()

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

static

Definition at line 78 of file attention_kernels.c.

{
    ck_attention_query_key_f32_args_t *args =
        (ck_attention_query_key_f32_args_t *)opaque;
    for (int q_token = ith; q_token < args->query_tokens; q_token += nth) {
        float *scores = args->score_scratch +
            (size_t)q_token * args->key_tokens;
        for (int head = 0; head < args->num_heads; ++head) {
            const float *q_head = args->query +
                (size_t)head * args->query_tokens * args->head_dim;
            const float *k_head = args->key +
                (size_t)head * args->key_tokens * args->head_dim;
            const float *v_head = args->value +
                (size_t)head * args->key_tokens * args->head_dim;
            float *out_head = args->output +
                (size_t)head * args->query_tokens * args->head_dim;
            const float *q_row = q_head + (size_t)q_token * args->head_dim;
            float maximum = -FLT_MAX;
            int k_token = 0;
#if defined(__AVX2__) && defined(__FMA__)
            for (; args->key_transpose_scratch != NULL &&
                   k_token + 7 < args->key_tokens; k_token += 8) {
                __m256 dots = _mm256_setzero_ps();
                for (int dim = 0; dim < args->head_dim; ++dim) {
                    const float *packed = args->key_transpose_scratch +
                        ((size_t)head * args->head_dim + dim) *
                            args->key_tokens + k_token;
                    dots = _mm256_fmadd_ps(
                        _mm256_set1_ps(q_row[dim]),
                        _mm256_loadu_ps(packed), dots);
                }
                float dot_values[8];
                _mm256_storeu_ps(dot_values, dots);
                for (int lane = 0; lane < 8; ++lane) {
                    const float score = dot_values[lane] * args->scale;
                    scores[k_token + lane] = score;
                    maximum = fmaxf(maximum, score);
                }
            }
#endif
            for (; k_token < args->key_tokens; ++k_token) {
                const float *k_row = k_head +
                    (size_t)k_token * args->head_dim;
                float dot = 0.0f;
                for (int dim = 0; dim < args->head_dim; ++dim) {
                    dot = fmaf(q_row[dim], k_row[dim], dot);
                }
                const float score = dot * args->scale;
                scores[k_token] = score;
                maximum = fmaxf(maximum, score);
            }
            double denominator = 0.0;
            for (int k_token = 0; k_token < args->key_tokens; ++k_token) {
                const float probability = expf(scores[k_token] - maximum);
                scores[k_token] = probability;
                denominator += (double)probability;
            }
            const float inverse = denominator > 0.0 ? (float)(1.0 / denominator) : 0.0f;
            float *out_row = out_head +
                (size_t)q_token * args->head_dim;
            for (int dim = 0; dim < args->head_dim; ++dim) {
                out_row[dim] = 0.0f;
            }
            for (int k_token = 0; k_token < args->key_tokens; ++k_token) {
                const float probability = scores[k_token] * inverse;
                const float *v_row = v_head +
                    (size_t)k_token * args->head_dim;
                for (int dim = 0; dim < args->head_dim; ++dim) {
                    out_row[dim] = fmaf(probability, v_row[dim], out_row[dim]);
                }
            }
        }
    }
}

References score.

Referenced by ck_attention_forward_query_key_head_major_f32_run().

◆ ck_attention_reference_expf()

static float ck_attention_reference_expf ( float value )

static

Definition at line 1243 of file attention_kernels.c.

{
#if defined(__linux__)
    static void *libm_handle = NULL;
    static ck_attention_math_f32_fn fn = NULL;
    static int resolved = 0;
    if (!resolved) {
        libm_handle = dlopen("libm.so.6", RTLD_NOW | RTLD_LOCAL);
        if (libm_handle) {
            fn = (ck_attention_math_f32_fn) dlsym(libm_handle, "expf");
        }
        resolved = 1;
    }
    if (fn) {
        return fn(value);
    }
#endif
    return expf(value);
}

Referenced by attention_forward_decode_head_major_gqa_f16cache_regular_noscale(), ck_attention_f16_prefill_gqa_reuse_work(), ck_attention_f16_prefill_qtile64_work(), ck_attention_f16_reduce_expf(), ck_attention_f16_split_work(), and ck_attention_full_tiled_f16kv_fp32_range().

◆ ck_attention_reverse_out_dot_enabled()

static int ck_attention_reverse_out_dot_enabled ( void )

static

Definition at line 487 of file attention_kernels.c.

{
    const char *v = getenv("CK_STRICT_ATTN_REVERSE_OUT_DOT");
    return v && v[0] && strcmp(v, "0") != 0;
}

◆ ck_attention_scale_f16_llama()

static void ck_attention_scale_f16_llama	(	uint16_t *	y,
		float	scale,
		int	n
	)

inlinestatic

Definition at line 6509 of file attention_kernels.c.

{
    int i = 0;
#if defined(__AVX512F__)
    const __m512 factor = _mm512_set1_ps(scale);
    for (; i + 63 < n; i += 64) {
        for (int lane = 0; lane < 4; ++lane) {
            const int offset = i + lane * 16;
            __m512 value = _mm512_cvtph_ps(
                _mm256_loadu_si256((const __m256i *) (y + offset)));
            value = _mm512_mul_ps(value, factor);
            _mm256_storeu_si256(
                (__m256i *) (y + offset), _mm512_cvtps_ph(value, 0));
        }
    }
#elif defined(__AVX2__) && defined(__F16C__)
    const __m256 factor = _mm256_set1_ps(scale);
    for (; i + 31 < n; i += 32) {
        for (int lane = 0; lane < 4; ++lane) {
            const int offset = i + lane * 8;
            __m256 value = _mm256_cvtph_ps(
                _mm_loadu_si128((const __m128i *) (y + offset)));
            value = _mm256_mul_ps(value, factor);
            _mm_storeu_si128(
                (__m128i *) (y + offset), _mm256_cvtps_ph(value, 0));
        }
    }
#endif
    for (; i < n; ++i) {
        const float value = CK_FP16_TO_FP32(y[i]);
        y[i] = CK_FP32_TO_FP16(value * scale);
    }
}

References CK_FP16_TO_FP32, and CK_FP32_TO_FP16.

Referenced by ck_attention_f16_split_work().

◆ ck_attention_sparse_bf16_pytorch_gqa_available()

int ck_attention_sparse_bf16_pytorch_gqa_available ( void )

Definition at line 9016 of file attention_kernels.c.

{
#if defined(__AVX512F__)
    return ck_attention_bf16_pytorch_gqa_available();
#else
    return 0;
#endif
}

References ck_attention_bf16_pytorch_gqa_available().

◆ ck_attention_strict_scale_f32()

static float ck_attention_strict_scale_f32 ( int head_dim )

static

Definition at line 1228 of file attention_kernels.c.

{
    // Keep strict parity on the precise libm sqrtf path. icx -O3 on AVX2 was
    // lowering 1/sqrtf(d) to a slightly smaller effective scale, which is
    // enough to move layer-0 vision softmax by ~2e-7 and snowball later.
    volatile float hd = (float) head_dim;
    float (*sqrtf_fn)(float) = sqrtf;
    volatile float root = sqrtf_fn(hd);
    volatile float one = 1.0f;
    volatile float scale = one / root;
    return scale;
}

Referenced by attention_forward_decode_head_major_gqa_flash_f16cache_split(), attention_forward_decode_head_major_gqa_regular(), attention_forward_full_head_major_gqa_exact_strided(), attention_forward_full_head_major_gqa_ggml_strided_workspace(), attention_forward_head_major_gqa_flash_impl(), attention_forward_query_key_head_major_tiled_f16kv_fp32(), ck_attention_f16_prefill_gqa_reuse_work(), ck_attention_f16_prefill_qtile64_work(), and ck_attention_full_tiled_f16kv_fp32_range().

◆ ck_attention_strict_unfused_f16_enabled()

static int ck_attention_strict_unfused_f16_enabled ( void )

static

Definition at line 4089 of file attention_kernels.c.

{
    const char *value = getenv("CK_STRICT_ATTN_F16_UNFUSED");
    return !value || !value[0] || strcmp(value, "0") != 0;
}

Referenced by attention_forward_head_major_gqa_flash_impl().

◆ ck_attention_trace()

static void ck_attention_trace	(	const char *	branch,
		int	layer_id,
		int	head_id
	)

static

Definition at line 596 of file attention_kernels.c.

{
    const char *enabled = getenv("CK_STRICT_ATTN_TRACE");
    const char *dir = getenv("CK_PARITY_DIR");
    if (!enabled || !enabled[0] || strcmp(enabled, "0") == 0 || !dir || !dir[0] || !branch || !branch[0]) {
        return;
    }
    char path[4096];
    snprintf(path, sizeof(path), "%s/%s", dir, "strict_trace.txt");
    FILE *f = fopen(path, "a");
    if (!f) {
        return;
    }
    fprintf(f, "layer=%d head=%d branch=%s\n", layer_id, head_id, branch);
    fclose(f);
}

Referenced by attention_forward_full_head_major_gqa_exact_strided(), and attention_forward_full_head_major_gqa_ggml_strided_workspace().

◆ ck_attention_trace_float()

static void ck_attention_trace_float	(	const char *	tag,
		int	layer_id,
		int	head_id,
		float	value
	)

static

Definition at line 634 of file attention_kernels.c.

{
    const char *enabled = getenv("CK_STRICT_ATTN_TRACE");
    const char *dir = getenv("CK_PARITY_DIR");
    if (!enabled || !enabled[0] || strcmp(enabled, "0") == 0 || !dir || !dir[0] || !tag || !tag[0]) {
        return;
    }
    char path[4096];
    snprintf(path, sizeof(path), "%s/%s", dir, "strict_trace.txt");
    FILE *f = fopen(path, "a");
    if (!f) {
        return;
    }
    fprintf(f, "layer=%d head=%d tag=%s float=%.17g\n", layer_id, head_id, tag, (double) value);
    fclose(f);
}

Referenced by attention_forward_full_head_major_gqa_ggml_strided_workspace().

◆ ck_attention_trace_query()

static void ck_attention_trace_query	(	const char *	tag,
		int	layer_id,
		int	head_id,
		int	query_id,
		int	value
	)

static

Definition at line 613 of file attention_kernels.c.

{
    const char *enabled = getenv("CK_STRICT_ATTN_TRACE");
    const char *dir = getenv("CK_PARITY_DIR");
    if (!enabled || !enabled[0] || strcmp(enabled, "0") == 0 || !dir || !dir[0] || !tag || !tag[0]) {
        return;
    }
    char path[4096];
    snprintf(path, sizeof(path), "%s/%s", dir, "strict_trace.txt");
    FILE *f = fopen(path, "a");
    if (!f) {
        return;
    }
    fprintf(f, "layer=%d head=%d query=%d tag=%s value=%d\n", layer_id, head_id, query_id, tag, value);
    fclose(f);
}

Referenced by ck_attention_vec_dump_selected_query(), and ck_attention_vec_dump_should_emit().

◆ ck_attention_u16_cache_to_f32()

static float ck_attention_u16_cache_to_f32	(	uint16_t	value,
		int	cache_is_bf16
	)

inlinestatic

Definition at line 7108 of file attention_kernels.c.

{
    return cache_is_bf16 ? bf16_to_float(value) : CK_FP16_TO_FP32(value);
}

References bf16_to_float(), and CK_FP16_TO_FP32.

Referenced by ck_attention_f16_prefill_gqa_reuse_work(), and ck_attention_f16_prefill_qtile64_work().

◆ ck_attention_vec_dump_enabled()

static int ck_attention_vec_dump_enabled ( void )

static

Definition at line 475 of file attention_kernels.c.

{
    const char *v = getenv("CK_STRICT_ATTN_VEC_DUMP");
    return v && v[0] && strcmp(v, "0") != 0;
}

Referenced by attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace(), attention_forward_head_major_gqa_flash_impl(), attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl(), ck_attention_llama_regular_impl(), and ck_attention_vec_dump_should_emit().

◆ ck_attention_vec_dump_exact_query()

static void ck_attention_vec_dump_exact_query	(	const float *	q_vec,
		const float *	k_head,
		const float *	out_vec,
		int	kv_tokens,
		int	head_dim,
		int	aligned_head_dim,
		float	scale,
		int	layer_id,
		int	head_id,
		int	query_id
	)

static

Definition at line 2141 of file attention_kernels.c.

{
    if (!ck_attention_vec_dump_should_emit(layer_id, head_id, query_id)) {
        return;
    }
 
    float *raw_scores = (float *) alloca((size_t) kv_tokens * sizeof(float));
    float *probs = (float *) alloca((size_t) kv_tokens * sizeof(float));
    float max_score = -INFINITY;
    for (int j = 0; j < kv_tokens; ++j) {
        const float *k_vec = k_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];
        }
        raw_scores[j] = dot;
        const float scaled = dot * scale;
        probs[j] = scaled;
        if (scaled > max_score) {
            max_score = scaled;
        }
    }
 
    float sum = 0.0f;
    for (int j = 0; j < kv_tokens; ++j) {
        probs[j] = expf(probs[j] - max_score);
        sum += probs[j];
    }
    if (sum > 0.0f) {
        const float inv_sum = 1.0f / sum;
        for (int j = 0; j < kv_tokens; ++j) {
            probs[j] *= inv_sum;
        }
    } else {
        memset(probs, 0, (size_t) kv_tokens * sizeof(float));
    }
 
    ck_attention_vec_dump_selected_query(raw_scores, probs, out_vec, NULL,
                                         kv_tokens, head_dim,
                                         layer_id, head_id, query_id);
}

References ck_attention_vec_dump_selected_query(), and ck_attention_vec_dump_should_emit().

Referenced by attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl().

◆ ck_attention_vec_dump_next_layer_id()

static int ck_attention_vec_dump_next_layer_id ( void )

static

Definition at line 554 of file attention_kernels.c.

{
    const int layer_id = ck_attention_vec_dump_layer_seq;
    ck_attention_vec_dump_layer_seq += 1;
    return layer_id;
}

References ck_attention_vec_dump_layer_seq.

Referenced by attention_forward_full_head_major_gqa_exact_strided(), attention_forward_full_head_major_gqa_ggml_strided_workspace(), attention_forward_head_major_gqa_flash_impl(), attention_forward_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl(), and ck_attention_llama_regular_impl().

◆ ck_attention_vec_dump_parse_env_int()

static int ck_attention_vec_dump_parse_env_int	(	const char *	name,
		int *	out
	)

static

Definition at line 506 of file attention_kernels.c.

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

References end.

Referenced by ck_attention_vec_dump_should_emit().

◆ ck_attention_vec_dump_selected_query()

static void ck_attention_vec_dump_selected_query	(	const float *	raw_scores,
		const float *	probs,
		const float *	out_vec,
		const float *	v_cols,
		int	kv_tokens,
		int	head_dim,
		int	layer_id,
		int	head_id,
		int	query_id
	)

static

Definition at line 654 of file attention_kernels.c.

{
    if (!ck_attention_vec_dump_should_emit(layer_id, head_id, query_id)) {
        return;
    }
    ck_attention_trace_query("vec_dump_selected_query", layer_id, head_id, query_id, 1);
    char name[32];
    snprintf(name, sizeof(name), "kq_scores_h%d_q%d", head_id, query_id);
    ck_attention_vec_dump_tensor(name, layer_id, query_id, raw_scores, (size_t) kv_tokens);
    snprintf(name, sizeof(name), "kq_soft_h%d_q%d", head_id, query_id);
    ck_attention_vec_dump_tensor(name, layer_id, query_id, probs, (size_t) kv_tokens);
    snprintf(name, sizeof(name), "kqv_out_h%d_q%d", head_id, query_id);
    ck_attention_vec_dump_tensor(name, layer_id, query_id, out_vec, (size_t) head_dim);
    if (v_cols && ck_attention_vec_dump_vcols_enabled()) {
        snprintf(name, sizeof(name), "vcols_h%d_q%d", head_id, query_id);
        ck_attention_vec_dump_tensor(name, layer_id, query_id, v_cols, (size_t) kv_tokens * (size_t) head_dim);
    }
}

References ck_attention_trace_query(), ck_attention_vec_dump_should_emit(), ck_attention_vec_dump_tensor(), and ck_attention_vec_dump_vcols_enabled().

Referenced by attention_forward_head_major_gqa_flash_impl(), attention_forward_head_major_gqa_unfused_f16_strict(), attention_query_full_exact_regular(), attention_query_full_ggml_regular(), ck_attention_llama_regular_impl(), and ck_attention_vec_dump_exact_query().

◆ ck_attention_vec_dump_should_emit()

static int ck_attention_vec_dump_should_emit	(	int	layer_id,
		int	head_id,
		int	query_id
	)

static

Definition at line 523 of file attention_kernels.c.

{
    if (!ck_attention_vec_dump_enabled()) {
        return 0;
    }
    int want_layer = -1;
    int want_head = -1;
    int want_query = -1;
    if (!ck_attention_vec_dump_parse_env_int("CK_STRICT_ATTN_DUMP_LAYER", &want_layer)) {
        return 0;
    }
    const int have_head = ck_attention_vec_dump_parse_env_int("CK_STRICT_ATTN_DUMP_HEAD", &want_head);
    const int have_query = ck_attention_vec_dump_parse_env_int("CK_STRICT_ATTN_DUMP_QUERY", &want_query);
    const int trace_target = layer_id == want_layer &&
                             (!have_head || head_id == want_head) &&
                             (!have_query || query_id == want_query);
    if (layer_id != want_layer) {
        return 0;
    }
    if (have_head && head_id != want_head) {
        return 0;
    }
    if (have_query && query_id != want_query) {
        return 0;
    }
    if (trace_target) {
        ck_attention_trace_query("vec_dump_should_emit", layer_id, head_id, query_id, 1);
    }
    return 1;
}

References ck_attention_trace_query(), ck_attention_vec_dump_enabled(), and ck_attention_vec_dump_parse_env_int().

Referenced by attention_forward_head_major_gqa_flash_impl(), attention_query_full_exact_regular(), attention_query_full_ggml_regular(), ck_attention_llama_regular_impl(), ck_attention_vec_dump_exact_query(), and ck_attention_vec_dump_selected_query().

◆ ck_attention_vec_dump_tensor()

static void ck_attention_vec_dump_tensor	(	const char *	name,
		int	layer_id,
		int	query_id,
		const float *	data,
		size_t	elem_count
	)

static

Definition at line 561 of file attention_kernels.c.

{
    const char *dir = getenv("CK_PARITY_DIR");
    if (!dir || !dir[0] || !name || !name[0] || !data || elem_count == 0) {
        return;
    }
 
    char path[4096];
    snprintf(path, sizeof(path), "%s/%s", dir, "strict_internal.bin");
    FILE *f = fopen(path, "ab");
    if (!f) {
        return;
    }
 
    ck_attention_vec_dump_header_t h;
    memset(&h, 0, sizeof(h));
    memcpy(h.magic, ck_attention_vec_dump_magic, sizeof(ck_attention_vec_dump_magic));
    h.version = ck_attention_vec_dump_version;
    h.layer_id = layer_id;
    strncpy(h.op_name, name, sizeof(h.op_name) - 1);
    h.dtype = 0u;
    h.rank = 1u;
    h.shape[0] = (int64_t) elem_count;
    h.elem_count = (uint32_t) elem_count;
    h.token_id = query_id;
 
    fwrite(&h, sizeof(h), 1, f);
    fwrite(data, sizeof(float), elem_count, f);
    fclose(f);
}

References ck_attention_vec_dump_magic, and ck_attention_vec_dump_version.

Referenced by ck_attention_vec_dump_selected_query().

◆ ck_attention_vec_dump_vcols_enabled()

static int ck_attention_vec_dump_vcols_enabled ( void )

static

Definition at line 481 of file attention_kernels.c.

{
    const char *v = getenv("CK_STRICT_ATTN_DUMP_VCOLS");
    return v && v[0] && strcmp(v, "0") != 0;
}

Referenced by ck_attention_vec_dump_selected_query().

◆ ck_bf16_dot_contract()

static float ck_bf16_dot_contract	(	const float *	a,
		const float *	b,
		int	count
	)

static

Definition at line 4564 of file attention_kernels.c.

{
#if defined(__AVX512BF16__)
    __m512 acc = _mm512_setzero_ps();
    int i = 0;
    for (; i + 32 <= count; i += 32) {
        const __m512bh av = _mm512_cvtne2ps_pbh(_mm512_loadu_ps(a + i + 16), _mm512_loadu_ps(a + i));
        const __m512bh bv = _mm512_cvtne2ps_pbh(_mm512_loadu_ps(b + i + 16), _mm512_loadu_ps(b + i));
        acc = _mm512_dpbf16_ps(acc, av, bv);
    }
    float sum = _mm512_reduce_add_ps(acc);
    for (; i < count; ++i) {
        const float av = bf16_to_float(float_to_bf16(a[i]));
        const float bv = bf16_to_float(float_to_bf16(b[i]));
        sum += av * bv;
    }
    return sum;
#else
    float sum = 0.0f;
    for (int i = 0; i < count; ++i) {
        const float av = bf16_to_float(float_to_bf16(a[i]));
        const float bv = bf16_to_float(float_to_bf16(b[i]));
        sum += av * bv;
    }
    return sum;
#endif
}

References bf16_to_float(), count, and float_to_bf16().

Referenced by ck_attention_full_bf16_sdpa_tiled_range().

◆ ck_ggml_vec_dot_f32_contig()

static float ck_ggml_vec_dot_f32_contig	(	const float *	x,
		const float *	y,
		int	n
	)

static

Definition at line 1090 of file attention_kernels.c.

{
    // Keep this a literal port of ggml_vec_dot_f32 so strict parity can match
    // llama.cpp's CPU attention path instead of merely approximating it.
#if defined(__AVX__)
    float sumf = 0.0f;
    const int np = (n & ~31);
    __m256 sum[4] = {
        _mm256_setzero_ps(),
        _mm256_setzero_ps(),
        _mm256_setzero_ps(),
        _mm256_setzero_ps(),
    };
 
    for (int i = 0; i < np; i += 32) {
        for (int j = 0; j < 4; ++j) {
            const __m256 ax = _mm256_loadu_ps(x + i + j * 8);
            const __m256 ay = _mm256_loadu_ps(y + i + j * 8);
#if defined(__FMA__)
            sum[j] = _mm256_fmadd_ps(ax, ay, sum[j]);
#else
            sum[j] = _mm256_add_ps(_mm256_mul_ps(ax, ay), sum[j]);
#endif
        }
    }
 
    sum[0] = _mm256_add_ps(sum[0], sum[2]);
    sum[1] = _mm256_add_ps(sum[1], sum[3]);
    sum[0] = _mm256_add_ps(sum[0], sum[1]);
    const __m128 t0 = _mm_add_ps(_mm256_castps256_ps128(sum[0]),
                                 _mm256_extractf128_ps(sum[0], 1));
    const __m128 t1 = _mm_hadd_ps(t0, t0);
    sumf = _mm_cvtss_f32(_mm_hadd_ps(t1, t1));
 
    for (int i = np; i < n; ++i) {
        sumf += x[i] * y[i];
    }
    return sumf;
#elif defined(__SSE2__)
    float sumf = 0.0f;
    const int np = (n & ~15);
    __m128 sum[4] = {
        _mm_setzero_ps(),
        _mm_setzero_ps(),
        _mm_setzero_ps(),
        _mm_setzero_ps(),
    };
 
    for (int i = 0; i < np; i += 16) {
        for (int j = 0; j < 4; ++j) {
            const __m128 ax = _mm_loadu_ps(x + i + j * 4);
            const __m128 ay = _mm_loadu_ps(y + i + j * 4);
#if defined(__FMA__)
            sum[j] = _mm_fmadd_ps(ax, ay, sum[j]);
#else
            sum[j] = _mm_add_ps(_mm_mul_ps(ax, ay), sum[j]);
#endif
        }
    }
 
    sum[0] = _mm_add_ps(sum[0], sum[2]);
    sum[1] = _mm_add_ps(sum[1], sum[3]);
    sum[0] = _mm_add_ps(sum[0], sum[1]);
#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
    sum[0] = _mm_add_ps(sum[0], _mm_movehl_ps(sum[0], sum[0]));
    sum[0] = _mm_add_ss(sum[0], _mm_movehdup_ps(sum[0]));
#else
    __m128 tmp = _mm_shuffle_ps(sum[0], sum[0], _MM_SHUFFLE(2, 3, 0, 1));
    sum[0] = _mm_add_ps(sum[0], tmp);
    tmp = _mm_movehl_ps(tmp, sum[0]);
    sum[0] = _mm_add_ss(sum[0], tmp);
#endif
    sumf = _mm_cvtss_f32(sum[0]);
 
    for (int i = np; i < n; ++i) {
        sumf += x[i] * y[i];
    }
    return sumf;
#else
    double sumf = 0.0;
    for (int i = 0; i < n; ++i) {
        sumf += (double) (x[i] * y[i]);
    }
    return (float) sumf;
#endif
}

Referenced by attention_forward_head_major_gqa_flash_impl(), and attention_query_full_ggml_regular().

◆ ck_ggml_vec_soft_max_row()

static double ck_ggml_vec_soft_max_row	(	int	n,
		float *	y,
		const float *	x,
		float	max
	)

static

Definition at line 1278 of file attention_kernels.c.

{
    int i = 0;
    double sum = 0.0;
 
#if defined(__AVX512F__) && defined(__AVX512DQ__)
    for (; i + 15 < n; i += 16) {
        const __m512 val = ck_ggml_v_expf512(
            _mm512_sub_ps(_mm512_loadu_ps(x + i), _mm512_set1_ps(max)));
        _mm512_storeu_ps(y + i, val);
        sum += (double) _mm512_reduce_add_ps(val);
    }
#elif defined(__AVX2__) && defined(__FMA__)
    for (; i + 7 < n; i += 8) {
        const __m256 val = ck_ggml_v_expf256(
            _mm256_sub_ps(_mm256_loadu_ps(x + i), _mm256_set1_ps(max)));
        _mm256_storeu_ps(y + i, val);
        __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1),
                                 _mm256_castps256_ps128(val));
        val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2));
        val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2));
        sum += (double) _mm_cvtss_f32(val2);
    }
#elif defined(__SSE2__)
    for (; i + 3 < n; i += 4) {
        const __m128 val = ck_ggml_v_expf128(
            _mm_sub_ps(_mm_loadu_ps(x + i), _mm_set1_ps(max)));
        _mm_storeu_ps(y + i, val);
#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
        __m128 acc = _mm_add_ps(val, _mm_movehl_ps(val, val));
        acc = _mm_add_ss(acc, _mm_movehdup_ps(acc));
#else
        __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1));
        __m128 acc = _mm_add_ps(val, tmp);
        tmp = _mm_movehl_ps(tmp, acc);
        acc = _mm_add_ss(acc, tmp);
#endif
        sum += (double) _mm_cvtss_f32(acc);
    }
#endif
 
    for (; i < n; ++i) {
        const float val = expf(x[i] - max);
        y[i] = val;
        sum += (double) val;
    }
 
    return sum;
}

Referenced by attention_forward_head_major_gqa_flash_impl(), attention_forward_head_major_gqa_unfused_f16_strict(), attention_query_full_ggml_regular(), ck_attention_f16_prefill_gqa_reuse_work(), ck_attention_f16_prefill_qtile64_work(), ck_attention_full_tiled_f16kv_fp32_range(), and ck_attention_llama_regular_query().

◆ ck_llama_kv_pad_256()

static int ck_llama_kv_pad_256	(	int	live_tokens,
		int	capacity
	)

inlinestatic

Definition at line 5802 of file attention_kernels.c.

{
    if (live_tokens <= 0 || capacity < live_tokens) return 0;
    int padded = live_tokens < 256 ? 256 : ((live_tokens + 255) / 256) * 256;
    return padded < capacity ? padded : capacity;
}

Referenced by ck_attention_llama_regular_impl().

◆ ck_llama_regular_dot_f16()

static float ck_llama_regular_dot_f16	(	const float *	a,
		const float *	b,
		int	count
	)

inlinestatic

Definition at line 5816 of file attention_kernels.c.

{
    int i = 0;
#if defined(__AVX512F__)
    __m512 acc[4] = {
        _mm512_setzero_ps(), _mm512_setzero_ps(),
        _mm512_setzero_ps(), _mm512_setzero_ps()
    };
    for (; i + 64 <= count; i += 64) {
        for (int lane = 0; lane < 4; ++lane) {
            acc[lane] = _mm512_fmadd_ps(
                ck_llama_regular_round_f16x16(_mm512_loadu_ps(a + i + lane * 16)),
                ck_llama_regular_round_f16x16(_mm512_loadu_ps(b + i + lane * 16)),
                acc[lane]);
        }
    }
    acc[0] = _mm512_add_ps(acc[0], acc[2]);
    acc[1] = _mm512_add_ps(acc[1], acc[3]);
    float sum = _mm512_reduce_add_ps(_mm512_add_ps(acc[0], acc[1]));
#elif defined(__AVX__)
    __m256 acc[4] = {
        _mm256_setzero_ps(), _mm256_setzero_ps(),
        _mm256_setzero_ps(), _mm256_setzero_ps()
    };
    for (; i + 32 <= count; i += 32) {
        for (int lane = 0; lane < 4; ++lane) {
#if defined(__FMA__)
            acc[lane] = _mm256_fmadd_ps(
                _mm256_cvtph_ps(_mm256_cvtps_ph(
                    _mm256_loadu_ps(a + i + lane * 8), _MM_FROUND_TO_NEAREST_INT)),
                _mm256_cvtph_ps(_mm256_cvtps_ph(
                    _mm256_loadu_ps(b + i + lane * 8), _MM_FROUND_TO_NEAREST_INT)),
                acc[lane]);
#else
            acc[lane] = _mm256_add_ps(acc[lane], _mm256_mul_ps(
                _mm256_cvtph_ps(_mm256_cvtps_ph(
                    _mm256_loadu_ps(a + i + lane * 8), _MM_FROUND_TO_NEAREST_INT)),
                _mm256_cvtph_ps(_mm256_cvtps_ph(
                    _mm256_loadu_ps(b + i + lane * 8), _MM_FROUND_TO_NEAREST_INT))));
#endif
        }
    }
    acc[0] = _mm256_add_ps(acc[0], acc[2]);
    acc[1] = _mm256_add_ps(acc[1], acc[3]);
    const __m256 merged = _mm256_add_ps(acc[0], acc[1]);
    __m128 halves = _mm_add_ps(
        _mm256_extractf128_ps(merged, 1), _mm256_castps256_ps128(merged));
    halves = _mm_add_ps(halves, _mm_movehl_ps(halves, halves));
    halves = _mm_add_ss(halves, _mm_movehdup_ps(halves));
    float sum = _mm_cvtss_f32(halves);
#else
    float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f};
    for (; i + 4 <= count; i += 4) {
        for (int lane = 0; lane < 4; ++lane) {
            const float av = CK_FP16_TO_FP32(CK_FP32_TO_FP16(a[i + lane]));
            const float bv = CK_FP16_TO_FP32(CK_FP32_TO_FP16(b[i + lane]));
            sums[lane] = fmaf(av, bv, sums[lane]);
        }
    }
    float sum = (sums[0] + sums[2]) + (sums[1] + sums[3]);
#endif
    for (; i < count; ++i) {
        const float av = CK_FP16_TO_FP32(CK_FP32_TO_FP16(a[i]));
        const float bv = CK_FP16_TO_FP32(CK_FP32_TO_FP16(b[i]));
        sum = fmaf(av, bv, sum);
    }
    return sum;
}

References CK_FP16_TO_FP32, CK_FP32_TO_FP16, and count.

Referenced by ck_attention_llama_regular_query().

◆ ck_llama_regular_gemm_f16()

static float ck_llama_regular_gemm_f16	(	const float *	probability,
		const float *	value_column,
		int	count
	)

inlinestatic

Definition at line 5885 of file attention_kernels.c.

{
    int i = 0;
#if defined(__AVX512F__)
    __m512 acc = _mm512_setzero_ps();
    for (; i + 16 <= count; i += 16) {
        acc = _mm512_fmadd_ps(
            ck_llama_regular_round_f16x16(_mm512_loadu_ps(value_column + i)),
            ck_llama_regular_round_f16x16(_mm512_loadu_ps(probability + i)), acc);
    }
    float sum = _mm512_reduce_add_ps(acc);
#elif defined(__AVX__)
    __m256 acc = _mm256_setzero_ps();
    for (; i + 8 <= count; i += 8) {
#if defined(__FMA__)
        acc = _mm256_fmadd_ps(
            _mm256_cvtph_ps(_mm256_cvtps_ph(
                _mm256_loadu_ps(value_column + i), _MM_FROUND_TO_NEAREST_INT)),
            _mm256_cvtph_ps(_mm256_cvtps_ph(
                _mm256_loadu_ps(probability + i), _MM_FROUND_TO_NEAREST_INT)), acc);
#else
        acc = _mm256_add_ps(acc, _mm256_mul_ps(
            _mm256_cvtph_ps(_mm256_cvtps_ph(
                _mm256_loadu_ps(value_column + i), _MM_FROUND_TO_NEAREST_INT)),
            _mm256_cvtph_ps(_mm256_cvtps_ph(
                _mm256_loadu_ps(probability + i), _MM_FROUND_TO_NEAREST_INT))));
#endif
    }
    __m128 halves = _mm_add_ps(
        _mm256_extractf128_ps(acc, 1), _mm256_castps256_ps128(acc));
    halves = _mm_add_ps(halves, _mm_movehl_ps(halves, halves));
    halves = _mm_add_ss(halves, _mm_movehdup_ps(halves));
    float sum = _mm_cvtss_f32(halves);
#else
    float sum = 0.0f;
#endif
    for (; i < count; ++i) {
        const float value = CK_FP16_TO_FP32(CK_FP32_TO_FP16(value_column[i]));
        const float probability_value = CK_FP16_TO_FP32(CK_FP32_TO_FP16(probability[i]));
        sum = fmaf(value, probability_value, sum);
    }
    return sum;
}

References CK_FP16_TO_FP32, CK_FP32_TO_FP16, and count.

Referenced by ck_attention_llama_regular_query().

◆ ck_llama_regular_prefill_work()

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

static

Definition at line 6052 of file attention_kernels.c.

{
    ck_llama_regular_prefill_args_t *a =
        (ck_llama_regular_prefill_args_t *)opaque;
    float *worker_scores = a->scores + (size_t)ith * (size_t)a->padded;
    float *worker_scaled =
        a->scaled_scores + (size_t)ith * (size_t)a->padded;
    const size_t worker_kv_stride =
        (size_t)a->kv_stride_tokens * (size_t)a->aligned_head_dim;
    const int jobs = a->num_heads * a->num_tokens;
    (void)nth;
    /* Query rows are independent, but their cost grows with the causal
     * prefix. Dynamic claims keep asymmetric workers balanced while each
     * worker retains exclusive ownership of its score scratch. */
    for (;;) {
        const int job = atomic_fetch_add_explicit(
            &a->next_job, 1, memory_order_relaxed);
        if (job >= jobs) break;
        const int head = job / a->num_tokens;
        const int query = job - head * a->num_tokens;
        const int kv_head =
            (int)((long long)head * a->num_kv_heads / a->num_heads);
        ck_attention_llama_regular_query(
            a->q + ((size_t)head * a->num_tokens + query) * a->aligned_head_dim,
            a->k + (size_t)kv_head * worker_kv_stride,
            a->value_columns +
                (size_t)kv_head * (size_t)a->head_dim * (size_t)a->padded,
            a->output +
                ((size_t)head * a->num_tokens + query) * a->aligned_head_dim,
            worker_scores, worker_scaled, a->num_tokens, a->padded, query,
            a->head_dim, a->aligned_head_dim, a->sliding_window, 1);
    }
}

References ck_attention_llama_regular_query().

Referenced by attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace().

◆ ck_local_fp16_to_fp32_2d()

static void ck_local_fp16_to_fp32_2d	(	const uint16_t *	src,
		float *	dst,
		int	rows,
		int	cols,
		int	src_stride,
		int	dst_stride
	)

inlinestatic

Definition at line 382 of file attention_kernels.c.

{
    if (!src || !dst || rows <= 0 || cols <= 0) {
        return;
    }
    for (int r = 0; r < rows; ++r) {
        ck_local_fp16_to_fp32_row(src + (size_t)r * (size_t)src_stride,
                                  dst + (size_t)r * (size_t)dst_stride,
                                  cols);
    }
}

References ck_local_fp16_to_fp32_row().

Referenced by attention_forward_decode_head_major_gqa_flash_f16cache().

◆ ck_local_fp16_to_fp32_row()

static void ck_local_fp16_to_fp32_row	(	const uint16_t *	src,
		float *	dst,
		int	n
	)

inlinestatic

Definition at line 372 of file attention_kernels.c.

{
    if (!src || !dst || n <= 0) {
        return;
    }
    for (int i = 0; i < n; ++i) {
        dst[i] = CK_FP16_TO_FP32(src[i]);
    }
}

References CK_FP16_TO_FP32.

Referenced by ck_local_fp16_to_fp32_2d().

◆ ck_muse_attention_checked_product()

static size_t ck_muse_attention_checked_product	(	size_t	left,
		size_t	right,
		const char *	label
	)

static

Definition at line 9036 of file attention_kernels.c.

{
    if (left != 0 && right > SIZE_MAX / left) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: Muse eager attention %s "
                "size overflow\n", label);
        abort();
    }
    return left * right;
}

References left, and right.

Referenced by ck_attention_muse_eager_bf16_storage_impl().

◆ ck_round_fp16_buffer()

static void ck_round_fp16_buffer	(	const float *	src,
		float *	dst,
		size_t	count
	)

static

Definition at line 356 of file attention_kernels.c.

{
    size_t i = 0;
#if defined(__AVX2__) && defined(__F16C__)
    for (; i + 8 <= count; i += 8) {
        const __m256 value = _mm256_loadu_ps(src + i);
        const __m128i half = _mm256_cvtps_ph(
            value, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
        _mm256_storeu_ps(dst + i, _mm256_cvtph_ps(half));
    }
#endif
    for (; i < count; ++i) {
        dst[i] = ck_round_fp16_scalar(src[i]);
    }
}

References ck_round_fp16_scalar(), and count.

Referenced by attention_flash_query_causal_exact_prerounded_f16kv(), attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace(), and attention_forward_head_major_gqa_flash_impl().

◆ ck_round_fp16_scalar()

static float ck_round_fp16_scalar ( float x )

inlinestatic

Definition at line 352 of file attention_kernels.c.

                                                  {
    return CK_FP16_TO_FP32(CK_FP32_TO_FP16(x));
}

References CK_FP16_TO_FP32, and CK_FP32_TO_FP16.

Referenced by attention_flash_query_causal_exact_f16kv(), attention_flash_query_causal_exact_prerounded_f16kv(), ck_attention_full_tiled_f16kv_fp32_range(), and ck_round_fp16_buffer().

◆ ck_vec_dot_f32_reverse_strict()

static float ck_vec_dot_f32_reverse_strict	(	const float *	x,
		const float *	y,
		int	n
	)

static

Definition at line 439 of file attention_kernels.c.

{
    float sumf = 0.0f;
    for (int i = n - 1; i >= 0; --i) {
        volatile float prod = x[i] * y[i];
        volatile float next = sumf + prod;
        sumf = next;
    }
    return sumf;
}

◆ ck_vec_dot_f32_strict()

static float ck_vec_dot_f32_strict	(	const float *	x,
		const float *	y,
		int	n
	)

static

Definition at line 413 of file attention_kernels.c.

{
    float sumf = 0.0f;
    for (int i = 0; i < n; ++i) {
        volatile float prod = x[i] * y[i];
        volatile float next = sumf + prod;
        sumf = next;
    }
    return sumf;
}

Referenced by attention_query_full_exact_regular().

◆ ck_vec_dot_f32x_f32_to_f32_via_f64()

static float ck_vec_dot_f32x_f32_to_f32_via_f64	(	const float *	x,
		const float *	y,
		int	n
	)

static

Definition at line 426 of file attention_kernels.c.

{
    double sum = 0.0;
    for (int i = 0; i < n; ++i) {
        volatile double prod = (double) x[i] * (double) y[i];
        volatile double next = sum + prod;
        sum = next;
    }
    return (float) sum;
}

Referenced by attention_query_full_exact_regular().

◆ ck_vec_max_f32_contig()

static float ck_vec_max_f32_contig	(	const float *	x,
		int	n
	)

inlinestatic

Definition at line 2900 of file attention_kernels.c.

{
    float max_val = -INFINITY;
    for (int i = 0; i < n; ++i) {
        if (x[i] > max_val) {
            max_val = x[i];
        }
    }
    return max_val;
}

Referenced by attention_forward_head_major_gqa_flash_impl(), attention_forward_head_major_gqa_unfused_f16_strict(), attention_query_full_ggml_regular(), ck_attention_f16_prefill_gqa_reuse_work(), ck_attention_f16_prefill_qtile64_work(), ck_attention_full_tiled_f16kv_fp32_range(), and ck_attention_llama_regular_query().

◆ ck_vec_scale_f32_inplace()

static void ck_vec_scale_f32_inplace	(	float *	x,
		int	n,
		float	scale
	)

inlinestatic

Definition at line 2893 of file attention_kernels.c.

{
    for (int i = 0; i < n; ++i) {
        x[i] *= scale;
    }
}

Referenced by attention_forward_head_major_gqa_flash_impl(), attention_query_full_ggml_regular(), ck_attention_f16_prefill_gqa_reuse_work(), ck_attention_f16_prefill_qtile64_work(), ck_attention_full_tiled_f16kv_fp32_range(), and ck_attention_llama_regular_query().

◆ convert_bf16_tensor_to_buf()

static void convert_bf16_tensor_to_buf	(	const uint16_t *	src,
		float *	dst,
		size_t	count
	)

static

Definition at line 321 of file attention_kernels.c.

{
    if (!dst || !src) return;
    bf16_tensor_to_float(src, dst, count);
}

References bf16_tensor_to_float(), and count.

Referenced by attention_backward_causal_head_major_gqa_bf16(), and attention_forward_causal_head_major_gqa_bf16().

◆ qkv_index()

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

inlinestatic

Definition at line 329 of file attention_kernels.c.

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

◆ score_index()

static size_t score_index	(	int	h,
		int	i,
		int	j,
		int	aligned_context_window
	)

inlinestatic

Definition at line 1333 of file attention_kernels.c.

{
    return ((size_t)h * (size_t)aligned_context_window * (size_t)aligned_context_window)
         + (size_t)i * (size_t)aligned_context_window
         + (size_t)j;
}

Referenced by attention_backward_causal_head_major_gqa(), attention_forward_causal_head_major(), attention_forward_causal_head_major_exact(), attention_forward_causal_head_major_gqa(), and attention_forward_causal_head_major_gqa_exact().

Variable Documentation

◆ ck_attention_vec_dump_layer_seq

int ck_attention_vec_dump_layer_seq = 0

static

Definition at line 467 of file attention_kernels.c.

Referenced by ck_attention_vec_dump_next_layer_id().

◆ ck_attention_vec_dump_magic

const char ck_attention_vec_dump_magic[8] = {'C', 'K', 'D', 'M', 'P', '\0', '\0', '\0'}

static

Definition at line 465 of file attention_kernels.c.

465{'C', 'K', 'D', 'M', 'P', '\0', '\0', '\0'};

Referenced by ck_attention_vec_dump_tensor().

◆ ck_attention_vec_dump_version

const uint32_t ck_attention_vec_dump_version = 1u

static

Definition at line 466 of file attention_kernels.c.

Referenced by ck_attention_vec_dump_tensor().

Macros

Typedefs

Functions

Variables

Detailed Description

CK-ENGINE KERNEL RULES:

Macro Definition Documentation

◆ CK_ENABLE_LLAMA_CPP_PARITY

◆ CK_GGML_FA_TILE_KV

◆ CK_GGML_FA_TILE_Q

◆ CK_GGML_FA_TILE_Q_LARGE

◆ CK_GGML_FA_TILE_Q_LARGE_MIN_TOKENS

◆ CK_MUSE_PREFILL_QUERY_TILE

◆ CK_NOINLINE

◆ CK_OPTNONE

◆ FLASH_QUERY_IMPL [1/2]

◆ FLASH_QUERY_IMPL [2/2]

◆ FLASH_QUERY_IMPL_DECODE

◆ RTLD_DEFAULT

Typedef Documentation

◆ ck_attention_math_f32_fn

Function Documentation

◆ attention_backward_causal_head_major()

◆ attention_backward_causal_head_major_gqa()

◆ attention_backward_causal_head_major_gqa_bf16()

◆ attention_flash_query_causal()

◆ attention_flash_query_causal_exact()

◆ attention_flash_query_causal_exact_f16kv()

◆ attention_flash_query_causal_exact_prerounded_f16kv()

◆ attention_forward_causal_head_major()

◆ attention_forward_causal_head_major_exact()

◆ attention_forward_causal_head_major_gqa()

◆ attention_forward_causal_head_major_gqa_bf16()

◆ attention_forward_causal_head_major_gqa_exact()

◆ attention_forward_causal_head_major_gqa_flash()

◆ attention_forward_causal_head_major_gqa_flash_strided()

◆ attention_forward_causal_head_major_gqa_flash_strided_f16kv()

◆ attention_forward_causal_head_major_gqa_flash_strided_f16kv_serial()

◆ attention_forward_causal_head_major_gqa_flash_strided_f16kv_workspace()

◆ attention_forward_causal_head_major_gqa_flash_strided_gemma4()

◆ attention_forward_causal_head_major_gqa_flash_strided_gemma4_token_output()

◆ attention_forward_causal_head_major_gqa_flash_strided_token_output()

◆ attention_forward_causal_head_major_gqa_llama_regular_strided_sliding_workspace()

◆ attention_forward_causal_head_major_gqa_muse_eager_bf16_storage()

◆ attention_forward_causal_head_major_gqa_muse_eager_bf16_storage_sliding()

◆ attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract()

◆ attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace()

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_auto_workspace()

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract()

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_contract_workspace()

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_gemma4_workspace()

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config()

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_workspace_bytes()

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_qtile64_schedule()

◆ attention_forward_causal_head_major_gqa_prefill_append_f16cache_sliding_gemma4_workspace()

◆ attention_forward_causal_head_major_gqa_prefill_full_bf16cache_pytorch_contract()

◆ attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_contract_workspace()

◆ attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_gemma4_workspace()

◆ attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_sliding_gemma4_workspace()

◆ attention_forward_causal_head_major_shared_kv_gemma4()

◆ attention_forward_chunk_head_major_gqa_flash_gemma4()

◆ attention_forward_decode_head_major_gqa_bf16cache_pytorch_contract()

◆ attention_forward_decode_head_major_gqa_f16cache_gemma4_contract()

◆ attention_forward_decode_head_major_gqa_f16cache_regular_noscale()

◆ attention_forward_decode_head_major_gqa_f16cache_sliding_gemma4_contract()

◆ attention_forward_decode_head_major_gqa_flash()

◆ attention_forward_decode_head_major_gqa_flash_f16cache()

◆ attention_forward_decode_head_major_gqa_flash_f16cache_contract()

◆ attention_forward_decode_head_major_gqa_flash_f16cache_split()

◆ attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned()

◆ attention_forward_decode_head_major_gqa_flash_f16kv()

◆ attention_forward_decode_head_major_gqa_flash_gemma4()

◆ attention_forward_decode_head_major_gqa_llama_regular_sliding_workspace()

◆ attention_forward_decode_head_major_gqa_muse_eager_bf16_storage()

◆ attention_forward_decode_head_major_gqa_muse_eager_bf16_storage_sliding()

◆ attention_forward_decode_head_major_gqa_regular()

◆ attention_forward_decode_head_major_shared_kv_gemma4()

◆ attention_forward_full_head_major_gqa_exact_strided()

◆ attention_forward_full_head_major_gqa_flash()

◆ attention_forward_full_head_major_gqa_flash_strided()