#include <stddef.h>
#include <stdint.h>
#include "ckernel_audio.h"
#include "cpu_features.h"
#include "ckernel_quant.h"
#include "mega_fused_attention.h"

Data Structures
struct	CKMathBackend

Enumerations
enum	ck_attention_prefill_schedule_t { CK_ATTN_PREFILL_SCHEDULE_KV_HEADS = 0 , CK_ATTN_PREFILL_SCHEDULE_QUERY_HEADS = 1 , CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES = 2 , CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES = 3 , CK_ATTN_PREFILL_SCHEDULE_GQA_SHARED_KV_TILES = 4 }

enum	ck_attention_reduction_t { CK_ATTN_REDUCTION_FP32_ONLINE = 0 , CK_ATTN_REDUCTION_F16_ONLINE_FP32_MERGE = 1 , CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE = 2 , CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64 = 3 , CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA = 4 }

enum	ck_attention_status_t { CK_ATTENTION_STATUS_OK = 0 , CK_ATTENTION_STATUS_INVALID_ARGUMENT = -1 , CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT = -2 , CK_ATTENTION_STATUS_INSUFFICIENT_WORKSPACE = -3 }

Functions
void	adamw_clip_update_multi_f32 (float const grads, float const weights, float const m_states, float const v_states, const size_t *numels, int tensor_count, float lr, float beta1, float beta2, float eps, float weight_decay, float max_grad_norm, int step)

void	adamw_update_f32 (const float grad, float weight, float m, float v, size_t numel, float lr, float beta1, float beta2, float eps, float weight_decay, int step)

void	add_backward_bf16 (const uint16_t d_y, uint16_t d_a, uint16_t *d_b, size_t n)

void	add_forward_2d_bf16 (const uint16_t a, const uint16_t b, uint16_t *y, int tokens, int dim, int aligned_dim)

void	add_forward_bf16 (const uint16_t a, const uint16_t b, uint16_t *y, size_t n)

void	add_forward_f32 (const float a, const float b, float *y, size_t n)

void	add_inplace_bf16 (uint16_t a, const uint16_t b, size_t n)

void	add_inplace_f32 (float a, const float b, size_t n)

void	add_scaled_forward_bf16 (const uint16_t a, const uint16_t b, uint16_t *y, float alpha, size_t n)

void	add_scaled_inplace_bf16 (uint16_t a, const uint16_t b, float alpha, size_t n)

void	add_stream_inplace (float a, const float b, size_t n)

void	add_stream_reorder_2d (float main_inout, float aux_scratch, int grid_h, int grid_w, int embed_dim, int merge_size)

int	argmax_f32 (const float *scores, int n)
	Find index of maximum value.

void	assistant_layer_scale_forward (float hidden, const float scale, int tokens, int embed_dim)

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)

void	attention_flash_decode (float out, const float q, const float k, const float v, int T_q, int T_k, int H, int D_h, float scale)
	Main flash attention function with SIMD dispatch.

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_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_sliding (const float q, const float k, const float v, float output, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)

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

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

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

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_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_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)

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_causal_head_major_shared_kv_sliding_gemma4 (const float q, float output, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int kv_stride_tokens, int sliding_window)

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

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)

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_flash_sliding (const float q_token, const float k_cache, const float v_cache, float out_token, int num_heads, int num_kv_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window)

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

void	attention_forward_decode_head_major_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_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_decode_head_major_shared_kv_sliding_gemma4 (const float q_token, const float k_cache, const float v_cache, float out_token, int num_heads, int kv_tokens, int cache_capacity, int head_dim, int aligned_head_dim, int sliding_window)

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)

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)

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_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)

void	attn_gate_sigmoid_mul_backward (const float d_out, const float x, const float gate, float d_x, float *d_gate, int rows, int num_heads, int state_dim)

void	attn_gate_sigmoid_mul_forward (const float x, const float gate, float *out, int rows, int num_heads, int state_dim)

void	attn_gate_sigmoid_mul_pytorch_bf16_storage (const float x, const float gate, float *out, int rows, int num_heads, int state_dim)

void	attn_gate_softplus_mul_forward (const float x, const float gate, float *out, int rows, int num_heads, int state_dim)

void	axpy_2d_f32 (float Y, const float X, float alpha, int num_tokens, int dim, int y_stride, int x_stride)
	Batched AXPY for 2D tensors: Y[t,:] += alpha * X[t,:].

void	axpy_f32 (float y, const float x, float alpha, int n)
	In-place AXPY: y += alpha * x.

void	axpy_zero_f32 (float y, const float x, float alpha, int n)
	Zero output then accumulate: y = 0; y += alpha * x.

void	backward_causal_softmax_head_major (float d_scores, const float weights, int num_heads, int num_tokens, int aligned_context_window)

void	backward_causal_softmax_head_major_bf16 (uint16_t d_scores, const uint16_t weights, int num_heads, int num_tokens, int aligned_context_window, float scratch_d_scores, float scratch_weights)

void	causal_softmax_head_major (float *scores, int num_heads, int num_tokens, int aligned_context_window)

void	causal_softmax_head_major_bf16 (uint16_t scores, int num_heads, int num_tokens, int aligned_context_window, float scratch)

void	causal_softmax_head_major_exact (float *scores, int num_heads, int num_tokens, int aligned_context_window)

int	ck_attention_bf16_pytorch_gqa_available (void)

void	ck_attention_flash_decode_wrapper (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)
	Wrapper to call TRUE flash attention from orchestration layer.

float	ck_attention_pytorch_sdpa_scale_f32 (int head_dim)

int	ck_attention_sparse_bf16_pytorch_gqa_available (void)

int	ck_flash_attn_choose_tile_k (int D_h)

int	ck_flash_attn_fast_exp_kind (void)

int	ck_gemm_bf16_amx_available (void)

int	ck_gemm_bf16_fp32out_amx_raw (const uint16_t A, const uint16_t B, float *C, int M, int N, int K, int accumulate)

int	ck_gemm_nt_f16_ggml_oracle (const float A, const void B, const float bias, float C, int M, int N, int K)

int	ck_gemm_nt_f16_simd_lanes (void)

void	ck_gemm_nt_head_major_q5_0 (const float attn_out, const void wo, const float bias, float output, int tokens, int embed_dim, int num_heads, int head_dim)
	Output projection from head-major attention (auto-dispatch)

void	ck_gemm_nt_head_major_q8_0 (const float attn_out, const void wo, const float bias, float output, int tokens, int embed_dim, int num_heads, int head_dim)
	Output projection from head-major attention (Q8_0 weights)

int	ck_get_num_threads (void)

int	ck_get_physical_cores (void)

void	ck_layout_head_to_token_f32 (const float src, float dst, int heads, int tokens, int head_dim)

void	ck_layout_token_to_head_f32 (const float src, float dst, int tokens, int heads, int head_dim)

void *	ck_memcpy_parallel_dispatch (void dst, const void src, size_t size)

int	ck_multimodal_mrope_positions_2d (int32_t *positions, int total_tokens, int prefix_start, int position_base, int prefix_tokens, int grid_x, int grid_y, int text_pos)

int	ck_multimodal_prefix_insert_f32 (const float source_rows, int32_t token_ids, float *decoder_rows, int row_count, int source_row_stride, int decoder_row_stride, int copy_dim, int start_row, int decoder_capacity)

const char *	ck_q6_k_prepared_provider_name (void)

const char *	ck_q6_k_q8_k_provider_name (void)

void	ck_residual_add_backward (const float d_out, float d_a, float *d_b, int tokens, int aligned_embed_dim)

void	ck_residual_add_token_major (const float a, const float b, float *out, int tokens, int aligned_embed_dim)

void	ck_residual_add_token_major_bf16_storage (const float a, const float b, float *out, int tokens, int aligned_embed_dim)

void	ck_residual_add_token_major_parallel_dispatch (const float a, const float b, float *out, int tokens, int aligned_embed_dim)

void	ck_set_num_threads (int num_threads)

void	ck_set_strict_parity (int enabled)

const float *	ck_strict_consume_next_gemm_a (size_t elems)

int	ck_strict_mtmd_clip_encode_planar_f32 (const float planar, int channels, int height, int width, float out, size_t out_elems)

int	ck_strict_parity_enabled (void)

void	ck_strict_store_next_gemm_a (const float *data, size_t elems)

CKMathBackend	ckernel_backend_native (void)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

void	dequant_q4_0_row (const void src, float dst, size_t n_elements)
	Dequantize Q4_0 row (multiple blocks)

void	dequant_q4_1_row (const void src, float dst, size_t n_elements)
	Dequantize Q4_1 row (multiple blocks)

void	dequant_q4_k_row (const void src, float dst, size_t n_elements)
	Dequantize Q4_K row (multiple blocks)

void	dequant_q5_0_row (const void src, float dst, size_t n_elements)
	Dequantize Q5_0 row (multiple blocks)

void	dequant_q5_1_row (const void src, float dst, size_t n_elements)
	Dequantize Q5_1 row (multiple blocks)

void	dequant_q6_k_row (const void src, float dst, size_t n_elements)
	Dequantize Q6_K row (multiple blocks)

void	dequant_q8_0_row (const void src, float dst, size_t n_elements)
	Dequantize Q8_0 row (multiple blocks)

void	embedding_backward (const int32_t token_ids, int token_count, const float d_output, float d_token_embeddings, float d_pos_embeddings, int vocab_size, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	embedding_backward_bf16 (const int32_t token_ids, int token_count, const uint16_t d_output, uint16_t d_token_embeddings, uint16_t d_pos_embeddings, int vocab_size, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	embedding_backward_bf16_mixed (const int32_t token_ids, int token_count, const uint16_t d_output, float d_token_embeddings, float d_pos_embeddings, int vocab_size, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	embedding_forward (const int32_t token_ids, int token_count, int vocab_size, const float token_embeddings, const float pos_embeddings, float output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	embedding_forward_bf16 (const int32_t token_ids, int token_count, int vocab_size, const uint16_t token_embeddings, const uint16_t pos_embeddings, uint16_t output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	embedding_forward_bf16_fp32 (const int32_t token_ids, int token_count, int vocab_size, const uint16_t token_embeddings, const float pos_embeddings, float output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	embedding_forward_q4_k (const int32_t token_ids, int token_count, int vocab_size, const void token_embeddings, const float pos_embeddings, float output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	embedding_forward_q5_0 (const int32_t token_ids, int token_count, int vocab_size, const void token_embeddings, const float pos_embeddings, float output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	embedding_forward_q6_k (const int32_t token_ids, int token_count, int vocab_size, const void token_embeddings, const float pos_embeddings, float output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	embedding_forward_q8_0 (const int32_t token_ids, int token_count, int vocab_size, const void token_embeddings, const float pos_embeddings, float output, int embed_dim, int aligned_embed_dim, int context_window, int add_pos)

void	farskip_swiglu_shared_combine_bf16 (const float hidden, const float routed, const float post_attn_residual, const uint16_t shared_gate, const uint16_t shared_up, const uint16_t shared_down, float main_output, float routed_free_output, int rows, int hidden_dim, int intermediate_dim)

void	farskip_swiglu_shared_combine_bf16_parallel_dispatch (const float hidden, const float routed, const float post_attn_residual, const uint16_t shared_gate, const uint16_t shared_up, const uint16_t shared_down, float main_output, float routed_free_output, int rows, int hidden_dim, int intermediate_dim)

void	farskip_swiglu_shared_combine_bf16_row_range (const float hidden, const float routed, const float post_attn_residual, const uint16_t shared_gate, const uint16_t shared_up, const uint16_t shared_down, float main_output, float routed_free_output, int rows, int hidden_dim, int intermediate_dim, int row_begin, int row_end)

void	fc1_backward_kernel (const float d_output, const float fc1_input, const float W_fc1, float d_input, float d_W_fc1, float d_b_fc1, int T, int aligned_in, int aligned_out, int num_threads)

void	fc2_backward_kernel (const float d_output, const float fc2_input, const float W_fc2, float d_input, float d_W_fc2, float d_b_fc2, int T, int aligned_in, int aligned_out, int num_threads)

void	feature_concat (const float main_input, const float branch_input, float *output, int rows, int main_dim, int branch_slice_dim, int num_branch_slices)

void	feature_concat_2way (const float main_input, const float branch_input, float *output, int rows, int main_dim, int branch_slice_dim, int num_branch_slices)

void	feature_slice_copy (const float src, float dst, int rows, int src_dim, int dst_dim, int dst_feature_offset)

void	final_logit_scale_f32 (float *logits, int tokens, int vocab_size, float scale)

void	fused_mlp_swiglu_decode (const float x, const float W_gate, const float W_up, const float W_down, const float b_gate, const float b_up, const float b_down, float output, int D, int Hff)

void	fused_mlp_swiglu_decode_tiled (const float x, const float W_gate, const float W_up, const float W_down, const float b_gate, const float b_up, const float b_down, float output, int D, int Hff)

void	fused_mlp_swiglu_decode_v2 (const float x, const float W_gate, const float W_up, const float W_down, const float b_gate, const float b_up, const float b_down, float output, int D, int Hff)

void	fused_mlp_swiglu_prefill (const float x, const float W_gate, const float W_up, const float W_down, float output, int seq_len, int hidden, int intermediate, float scratch)
	Fused MLP (Gate + Up + SwiGLU + Down) for prefill.

void	fused_mlp_swiglu_prefill_bias (const float x, const float W_gate, const float W_up, const float W_down, const float B_gate, const float B_up, const float B_down, float output, int seq_len, int hidden, int intermediate, float *scratch)
	Fused MLP (Gate + Up + SwiGLU + Down) for prefill with biases.

void	fused_mlp_swiglu_prefill_w1w2_quant (const float x, const void W1, const float B1, CKDataType w1_dt, const void W2, const float B2, CKDataType w2_dt, float output, int seq_len, int embed_dim, int aligned_embed_dim, int intermediate_dim, int aligned_intermediate_dim, void *scratch)
	Quantized fused MLP for prefill (W1=gate+up, W2=down)

size_t	fused_mlp_swiglu_prefill_w1w2_quant_scratch_size (int aligned_embed_dim, int aligned_intermediate_dim)
	Get scratch buffer size for fused_mlp_swiglu_prefill_w1w2_quant.

size_t	fused_mlp_swiglu_scratch_size (int intermediate)
	Get scratch buffer size for fused_mlp_swiglu_prefill.

void	fused_rmsnorm_qkv_prefill (const float x, const float gamma, const float Wq, const float Wk, const float Wv, float Q, float K, float V, int seq_len, int hidden, int q_dim, int kv_dim, float eps, float *scratch)
	Fused RMSNorm + QKV projection for prefill.

void	fused_rmsnorm_qkv_prefill_head_major (const float x, const float gamma, const float Wq, const float Bq, const float Wk, const float Bk, const float Wv, const float Bv, float Q, float K, float V, int seq_len, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int kv_stride_tokens, float eps, float scratch)
	Fused RMSNorm + QKV projection for prefill (head-major outputs)

void	fused_rmsnorm_qkv_prefill_head_major_quant (const float x, const float gamma, const void Wq, const float Bq, CKDataType wq_dt, const void Wk, const float Bk, CKDataType wk_dt, const void Wv, const float Bv, CKDataType wv_dt, float Q, float K, float V, int seq_len, int embed_dim, int aligned_embed_dim, int num_heads, int num_kv_heads, int head_dim, int aligned_head_dim, int kv_stride_tokens, float eps, void scratch)
	Fused RMSNorm + QKV projection for prefill (head-major, Q8 activations)

size_t	fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size (int aligned_embed_dim)
	Get scratch buffer size for fused_rmsnorm_qkv_prefill_head_major_quant.

size_t	fused_rmsnorm_qkv_scratch_size (int hidden)
	Get scratch buffer size for fused_rmsnorm_qkv_prefill.

void	gated_deltanet_autoregressive_backward (const float d_out, const float d_state_out, const float q, const float k, const float v, const float g, const float beta, const float state_in, const float state_out, float d_q, float d_k, float d_v, float d_g, float d_beta, float *d_state_in, int num_heads, int state_dim, float norm_eps)

void	gated_deltanet_autoregressive_forward (const float q, const float k, const float v, const float g, const float beta, const float state_in, float state_out, float out, int num_heads, int state_dim, float norm_eps)

void	gated_deltanet_llama_avx2_forward (const float q, const float k, const float v, const float g, const float beta, const float state_in, float state_out, float out, int num_heads, int group_count, int state_dim, float norm_eps)

void	gated_deltanet_llama_avx2_prefill_forward (const float q, const float k, const float v, const float g, const float beta, const float state_in, float state_out, float out, int rows, int num_heads, int group_count, int state_dim, float norm_eps)

void	gated_deltanet_llama_chunk64_head_forward (const float q, const float k, const float v, const float g, const float beta, const float state_in, float state_out, float out, int rows, int num_heads, int group_count, int head, int state_dim)

void	gated_deltanet_llama_chunk64_prefill_forward (const float q, const float k, const float v, const float g, const float beta, const float state_in, float state_out, float out, int rows, int num_heads, int group_count, int state_dim, float norm_eps)

void	gated_deltanet_prefill_forward (const float q, const float k, const float v, const float g, const float beta, const float state_in, float state_out, float out, int rows, int num_heads, int state_dim, float norm_eps)

void	gated_deltanet_pytorch_grouped_bf16_forward (const float q, const float k, const float v, const float g, const float beta, const float state_in, float state_out, float out, int num_heads, int group_count, int state_dim, float norm_eps)

void	gated_deltanet_pytorch_grouped_bf16_prefill_forward (const float q, const float k, const float v, const float g, const float beta, const float state_in, float state_out, float out, int rows, int num_heads, int group_count, int state_dim, float norm_eps)

void	geglu_backward_bf16_mixed (const uint16_t x, const uint16_t d_out, float *d_x, int tokens, int dim)

void	geglu_backward_fp32 (const float x, const float d_out, float *d_x, int tokens, int dim)

void	geglu_forward_bf16 (const uint16_t x, uint16_t out, int tokens, int dim, float *scratch)

void	geglu_forward_exact (const float x, float out, int tokens, int dim)

void	geglu_forward_fp32 (const float x, float out, int tokens, int dim)

void	geglu_forward_ggml_native (const float x, float out, int tokens, int dim)

void	gelu_backward_exact (const float input, const float d_output, float *d_input, size_t n)

void	gelu_backward_exact_bf16 (const uint16_t input, const uint16_t d_output, uint16_t d_input, size_t n, float scratch_input, float scratch_d_output, float scratch_d_input)

void	gelu_backward_fast (const float input, const float d_output, float *d_input, size_t n)

void	gelu_backward_fast_bf16 (const uint16_t input, const uint16_t d_output, uint16_t d_input, size_t n, float scratch_input, float scratch_d_output, float scratch_d_input)

void	gelu_backward_scalar (const float input, const float d_output, float *d_input, size_t n)

void	gelu_erf_bf16_storage (float *data, size_t n)

void	gelu_erf_fp64_f32_inplace (float *data, size_t n)

void	gelu_exact_inplace (float *data, size_t n)

void	gelu_fast_inplace (float *data, size_t n)

void	gelu_fast_inplace_bf16 (uint16_t data, size_t n, float scratch)

void	gelu_ggml_inplace (float *data, size_t n)

void	gelu_ggml_native_inplace (float *data, size_t n)

void	gelu_pytorch_erf_f32_inplace (float *data, size_t n)

void	gelu_pytorch_erf_sleef_bf16_storage (float *data, size_t n)

void	gelu_pytorch_tanh_bf16_storage (float *data, size_t n)

void	gemm_avx512_parallel (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_backward_bf16_mixed (const uint16_t d_output, const uint16_t input, const uint16_t weight, float d_input, float d_weight, float d_bias, int tokens, int in_dim, int out_dim)

void	gemm_backward_f32_train_parallel_dispatch (const float d_output, const float input, const float W, float d_input, float d_W, float d_b, int T, int aligned_in, int aligned_out, int num_threads)

void	gemm_backward_f32_train_parallel_dispatch_v2 (const float d_output, const float input, const float W, float d_input, float d_W, float d_b, int T, int aligned_in, int aligned_out, int num_threads)

void	gemm_bias_gelu_fused (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_bias_relu_fused (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_bias_silu_fused (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_blocked_serial (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_blocked_serial_bf16 (const uint16_t A, const uint16_t B, const uint16_t bias, uint16_t C, int M, int N, int K)

void	gemm_blocked_serial_train_parallel_dispatch (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_fine_grained_parallel (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_microkernel (const float A, const float B, float *C, int M, int N, int K, int B_transposed)

void	gemm_microkernel_blocked (const float A, const float B, float *C, int M, int N, int K)

void	gemm_microkernel_blocked_bt (const float A, const float B, float *C, int M, int N, int K)

void	gemm_microkernel_packed (const float A, const float B, float *C, int M, int N, int K)

void	gemm_naive_parallel (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_nn_avx512 (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_nn_blocked (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_nn_parallel (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_nn_simd (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16 (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16_amx_bf16_storage (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16_amx_bf16_storage_workspace (const float A, const void B, const float bias, float C, int M, int N, int K, uint16_t *a_bf16, size_t a_bf16_bytes)

void	gemm_nt_bf16_bf16_storage (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16_bf16_storage_parallel_dispatch (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16_bf16_storage_row_range (const float A, const void B, const float bias, float C, int M, int N, int K, int row_begin, int row_end)

void	gemm_nt_bf16_native_bf16_storage (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16_parallel_dispatch (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16_prefill_shape_safe_bf16_storage (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace (const float A, const void B, const float bias, float C, int M, int N, int K, uint16_t *a_bf16, size_t a_bf16_bytes)

void	gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_bf16_row_range (const float A, const void B, const float bias, float C, int M, int N, int K, int row_begin, int row_end)

void	gemm_nt_f16 (const float A, const void B, const float bias, float C, int M, int N, int K)
	NT GEMM wrapper for FP16 weights with the engine's standard ABI.

void	gemm_nt_f16_clipped (const float A, const void B, const float bias, const float input_min, const float input_max, const float output_min, const float output_max, float C, int M, int N, int K)

void	gemm_nt_f32_llama_production (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_nt_f32_llama_production_output_range (const float A, const float B, const float bias, float C, int M, int N, int K, int output_begin, int output_end)

void	gemm_nt_fp32_exact_parallel_dispatch (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q4_0 (const float A, const void B, const float bias, float C, int M, int N, int K)
	Matrix-matrix multiply: C[M,N] = A[M,K] @ B[N,K]^T + bias.

void	gemm_nt_q4_1 (const float A, const void B, const float bias, float C, int M, int N, int K)
	GEMM with transposed Q4_1 weights: C = A @ B^T.

void	gemm_nt_q4_k (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q4_k_q8_k (const void A_q8, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q5_0 (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q5_1 (const float A, const void B, const float bias, float C, int M, int N, int K)
	GEMM with transposed Q5_1 weights: C = A @ B^T.

void	gemm_nt_q5_1_q8_1 (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q5_1_q8_1_ref (const void A_q8, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q5_k (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q5_k_q8_k (const void A_q8, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q6_k (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q6_k_q8_k (const void A_q8, const void B, const float bias, float C, int M, int N, int K)
	NT GEMM: C = A @ B^T where A is Q8_K and B is Q6_K.

void	gemm_nt_q6_k_q8_k_m4_tile (const void A_q8, const void B, const float bias, float C, int M, int N, int K, int m0, int m1, int n0, int n1)

void	gemm_nt_q6_k_q8_k_tile (const void A_q8, const void B, const float bias, float C, int M, int N, int K, int m0, int m1, int n0, int n1)
	Compute one C[m0:m1, n0:n1] tile for Q8_K activations x Q6_K weights.

void	gemm_nt_q6_k_q8_k_tiled (const void A_q8, const void B, const float bias, float C, int M, int N, int K)
	Experimental single-thread tiled NT GEMM wrapper.

void	gemm_nt_q8_0 (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q8_0_q8_0 (const void A_q8, const void B, const float bias, float C, int M, int N, int K)
	gemm_nt_q8_0_q8_0 with optional bias (matches header signature)

void	gemm_nt_q8_0_q8_0_contract (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q8_0_q8_0_m2n4 (const void A_q8, const void B, const float bias, float C, int M, int N, int K)

void	gemm_nt_q8_0_q8_0_m2n4_tile (const void A_q8, const void B, const float bias, float C, int M, int N, int K, int ldc)

void	gemm_q4_k (float Y, const void W, const float *X, int M, int N, int K)
	Auto-dispatch GEMM based on available SIMD.

void	gemm_q4_k_q8_k (float Y, const void W, const void *X_q8, int M, int N, int K)

void	gemm_q6_k (float Y, const void W, const float *X, int M, int N, int K)

void	gemm_q6_k_q8_k (float Y, const void W, const void *X_q8, int M, int N, int K)
	GEMM: Y = W @ X^T where W is Q6_K and X is Q8_K.

void	gemm_swiglu_fused (const float x, const float W_gate, const float W_up, const float b_gate, const float b_up, float output, int M, int N, int K)

void	gemm_tn_avx512 (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_tn_blocked (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemm_tn_parallel (const float A, const float B, const float bias, float C, int M, int N, int K)

void	gemma4_final_logit_softcap_forward (float *logits, int tokens, int vocab_size, float cap)

void	gemma4_per_layer_embed_forward (float hidden, const float per_layer_input, const float inp_gate, const float proj, const float post_norm, const float out_scale, int tokens, int layer, int num_layers, int embed_dim, int per_layer_dim, float eps)

void	gemma4_per_layer_prepare_bf16_forward (float per_layer_input, const float hidden, const int32_t token_ids, const uint16_t per_layer_token_emb, const uint16_t per_layer_model_proj, const float per_layer_proj_norm, int tokens, int num_layers, int embed_dim, int per_layer_dim, int vocab_size, float eps)

void	gemma4_per_layer_prepare_forward (float per_layer_input, const float hidden, const int32_t token_ids, const void per_layer_token_emb, const uint16_t per_layer_model_proj, const float per_layer_proj_norm, int tokens, int num_layers, int embed_dim, int per_layer_dim, int vocab_size, float eps)

void	gemma4_v_norm_forward (const float input, float output, float *rstd_cache, int tokens, int num_kv_heads, int head_dim, float eps)

void	gemma4_v_norm_forward_parallel_dispatch (const float input, float output, float *rstd_cache, int tokens, int num_kv_heads, int head_dim, float eps)

void	gemma4_vision_projector_prep_forward (const float input, float output, int tokens, int dim, float scale, float eps)

void	gemv_bf16 (float y, const void W, const float *x, int M, int K)

void	gemv_bf16_bf16_storage (float y, const void W, const float *x, int M, int K)

void	gemv_bf16_bf16_storage_parallel_dispatch (float y, const void W, const float *x, int M, int K)

void	gemv_bf16_parallel_dispatch (float y, const void W, const float *x, int M, int K)

void	gemv_fused_q5_0_bias_dispatch (float y, const void W, const float x, const float bias, int M, int K)

void	gemv_fused_q8_0_bias_dispatch (float y, const void W, const float x, const float bias, int M, int K)

void	gemv_q4_0 (float y, const void W, const float *x, int M, int K)
	Auto-dispatch GEMV.

void	gemv_q4_k (float y, const void W, const float *x, int M, int K)
	Auto-dispatch GEMV based on available SIMD.

void	gemv_q4_k_q8_k (float y, const void W, const void *x_q8, int M, int K)

void	gemv_q4_k_q8_k_parallel (float y, const void W, const void *x_q8, int M, int K, int ith, int nth)

void	gemv_q4_k_q8_k_parallel_simd (float y, const void W, const void *x_q8, int M, int K, int ith, int nth)

void	gemv_q4_k_q8_k_ref (float y, const void W, const void *x_q8, int M, int K)

void	gemv_q5_0 (float y, const void W, const float *x, int M, int K)
	Auto-dispatch GEMV for Q5_0 weights based on CPU features.

void	gemv_q5_0_parallel (float y, const void W, const float *x, int M, int K, int ith, int nth)
	Parallel reference GEMV for Q5_0 × FP32.

void	gemv_q5_0_parallel_simd (float y, const void W, const float *x, int M, int K, int ith, int nth)
	Parallel SIMD GEMV for Q5_0 × FP32 with prefetching.

void	gemv_q5_0_q8_0 (float y, const void W, const void *x_q8, int M, int K)
	Matrix-vector multiply with Q5_0 weights and Q8_0 input.

void	gemv_q5_1 (float y, const void W, const float *x, int M, int K)
	Auto-dispatch GEMV.

void	gemv_q5_1_q8_1 (float y, const void W, const float *x, int M, int K)

void	gemv_q5_1_q8_1_ref (float y, const void W, const void *x_q8, int M, int K)

void	gemv_q5_k (float y, const void W, const float *x, int M, int K)

void	gemv_q5_k_q8_k (float y, const void W, const void *x_q8, int M, int K)

void	gemv_q6_k (float y, const void W, const float *x, int M, int K)

void	gemv_q6_k_q8_k (float y, const void W, const void *x_q8, int M, int K)
	GEMV: y = W @ x where W is Q6_K and x is Q8_K.

void	gemv_q6_k_q8_k_parallel (float y, const void W, const void *x_q8, int M, int K, int ith, int nth)
	Parallel reference GEMV for Q6_K × Q8_K.

void	gemv_q6_k_q8_k_parallel_simd (float y, const void W, const void *x_q8, int M, int K, int ith, int nth)
	Parallel SIMD GEMV for Q6_K × Q8_K.

void	gemv_q8_0 (float y, const void W, const float *x, int M, int K)
	Auto-dispatch GEMV for Q8_0 weights based on CPU features.

void	gemv_q8_0_q8_0 (float y, const void W, const void *x_q8, int M, int K)
	Matrix-vector multiply with Q8_0 weights and Q8_0 input.

void	gemv_q8_0_q8_0_contract (float y, const void W, const float *x, int M, int K)

void	gemv_q8_0_q8_0_x4 (float y, const void W, const void *x_q8, int M, int K)

void	gradient_accumulate_f32 (float dst, const float src, size_t numel)

void	gradient_accumulate_multi_f32 (float const dsts, const float const srcs, const size_t *numels, int tensor_count)

float	gradient_clip_norm_f32 (float *grad, size_t numel, float max_norm)
	Clip gradient norm (fp32)

float	gradient_global_norm_multi_f32 (const float const grads, const size_t *numels, int tensor_count)

void	gradient_scale_f32 (float *grad, size_t numel, float scale)

void	group_limited_topk_router_sigmoid_f32 (const float logits, const float correction_bias, int indices, float weights, int rows, int n_experts, int top_k, int n_group, int topk_group, int norm_topk_prob, float routed_scaling_factor)

void	hyper_connection_mix_bf16 (const float hyper_input, const float norm_weight, const uint16_t mix_down_weight, const uint16_t mix_up_weight, const uint16_t inject_weight, float mixed_output, float injection_output, float normalized_scratch, float dynamic_scratch, float mix_scratch, int rows, int streams, int hidden_dim, int dynamic_dim, float eps, int emit_injection)

void	hyper_connection_mix_q4k_q5_0_q4k (const float hyper_input, const float norm_weight, const void mix_down_weight, const void mix_up_weight, const void inject_weight, float mixed_output, float injection_output, float normalized_scratch, float dynamic_scratch, float mix_scratch, int rows, int streams, int hidden_dim, int dynamic_dim, float eps, int emit_injection)

void	hyper_connection_mix_q6k_q5_0_q4k (const float hyper_input, const float norm_weight, const void mix_down_weight, const void mix_up_weight, const void inject_weight, float mixed_output, float injection_output, float normalized_scratch, float dynamic_scratch, float mix_scratch, int rows, int streams, int hidden_dim, int dynamic_dim, float eps, int emit_injection)

void	hyper_stream_expand_bf16 (const float input, float output, int rows, int streams, int hidden_dim)

void	hyper_stream_expand_f32 (const float input, float output, int rows, int streams, int hidden_dim)

void	hyper_stream_inject_bf16 (const float hyper_input, const float block_output, const float injection_weight, float output, int rows, int streams, int hidden_dim)

void	hyper_stream_inject_f32 (const float hyper_input, const float block_output, const float injection_weight, float output, int rows, int streams, int hidden_dim)

void	im2patch (const float image, float patches, int C, int H, int W, int P)

void	im2patch_bf16 (const uint16_t image, uint16_t patches, int C, int H, int W, int P)

void	kv_cache_repack_head_major_inplace (float *buf, int num_heads, int tokens, int cache_capacity, int aligned_head_dim)

void	kv_cache_store (float __restrict kv_cache_k, float __restrict kv_cache_v, const float __restrict k, const float __restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)

void	kv_cache_store_batch_bf16 (uint16_t __restrict kv_cache_k, uint16_t __restrict kv_cache_v, const float __restrict k, const float __restrict v, int start_pos, int num_tokens, int num_kv_heads, int head_dim, int max_seq_len)

void	kv_cache_store_batch_f16 (uint16_t __restrict kv_cache_k, uint16_t __restrict kv_cache_v, const float __restrict k, const float __restrict v, int start_pos, int num_tokens, int num_kv_heads, int head_dim, int max_seq_len)

void	kv_cache_store_batch_f32 (float __restrict kv_cache_k, float __restrict kv_cache_v, const float __restrict k, const float __restrict v, int start_pos, int num_tokens, int num_kv_heads, int head_dim, int max_seq_len)

void	kv_cache_store_bf16 (uint16_t __restrict kv_cache_k, uint16_t __restrict kv_cache_v, const float __restrict k, const float __restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)

void	kv_cache_store_f16 (uint16_t __restrict kv_cache_k, uint16_t __restrict kv_cache_v, const float __restrict k, const float __restrict v, int layer, int pos, int num_kv_heads, int head_dim, int max_seq_len)

void	kv_cache_store_shared_q (float __restrict kv_cache_k, float __restrict kv_cache_v, const float *__restrict q, int layer, int pos, int num_heads, int head_dim, int max_seq_len)

void	kv_cache_write_head_major (const float __restrict k_token, const float __restrict v_token, float __restrict k_cache, float __restrict v_cache, int num_kv_heads, int token_index, int cache_capacity, int head_dim, int aligned_head_dim)

void	layernorm_backward_kernel (const float d_output, const float input, const float gamma, const float mean, const float rstd, float d_input, float d_gamma, float d_beta, int tokens, int d_model, int aligned_embed_dim)

void	layernorm_backward_kernel_bf16 (const uint16_t d_output, const uint16_t input, const float gamma, const float mean, const float rstd, uint16_t d_input, float d_gamma, float d_beta, int tokens, int d_model, int aligned_embed_dim, float scratch_d_output, float scratch_input, float *scratch_d_input)

void	layernorm_forward_rolled_slice (const float __restrict input_slice_base, const float __restrict gamma, const float __restrict beta, float __restrict output_slice_base, float __restrict mean_cache_slice, float __restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, int aligned_embed_dim, float eps)

void	layernorm_forward_rolled_slice_bf16 (const uint16_t __restrict input_slice_base, const float __restrict gamma, const float __restrict beta, uint16_t __restrict output_slice_base, float __restrict mean_cache_slice, float __restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, int aligned_embed_dim, float eps, float scratch_input, float scratch_output)

void	layernorm_forward_unrolled_slice (const float __restrict input_slice_base, const float __restrict gamma, const float __restrict beta, float __restrict output_slice_base, float __restrict mean_cache_slice, float __restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, float eps)

void	layernorm_forward_unrolled_slice_bf16 (const uint16_t __restrict input_slice_base, const float __restrict gamma, const float __restrict beta, uint16_t __restrict output_slice_base, float __restrict mean_cache_slice, float __restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, float eps, float scratch_input, float scratch_output)

void	layernorm_naive_serial (const float input, const float gamma, const float beta, float output, float mean_cache, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	layernorm_naive_serial_bf16_storage (const float input, const float gamma, const float beta, float output, float mean_cache, float rstd_cache, int tokens, int d_model, float eps)

void	layernorm_naive_serial_matched_precision (const float input, const float gamma, const float beta, float output, float mean_cache, float rstd_cache, int tokens, int d_model, float eps)

void	layernorm_pytorch_welford_bf16_storage (const float input, const float gamma, const float beta, float output, float mean_cache, float rstd_cache, int tokens, int d_model, float eps)

void	mamba2_conv1d_decode_f32 (const float state_in, const float x, const float weight, const float bias, float conv_out, float state_out, int rows, int conv_dim, int kernel_size)

void	mamba2_conv1d_f32_channel_range (const float state_in, const float x, const float weight, const float bias, float conv_out, float state_out, int rows, int conv_dim, int kernel_size, int channel_begin, int channel_end)

void	mamba2_conv1d_f32_parallel_dispatch (const float state_in, const float x, const float weight, const float bias, float conv_out, float state_out, int rows, int conv_dim, int kernel_size)

void	mamba2_dt_softplus_f32 (const float dt, const float dt_bias, float *dt_out, int rows, int num_heads, float dt_min, float dt_max)

void	mamba2_in_proj_split_f32 (const float projected, float gate, float hidden_bc, float dt, int rows, int d_mlp, int intermediate_dim, int conv_dim, int num_heads)

void	mamba2_rmsnorm_gate_f32 (const float x, const float gate, const float weight, float out, int rows, int inner_dim, int group_size, float eps)

void	mamba2_selective_scan_f32 (const float state_init, const float x, const float dt, const float a, const float b, const float c, const float d, float state_out, float *y, int batch, int seq_len, int num_heads, int head_dim, int state_dim, int num_groups)

void	mamba2_selective_scan_f32_head_range (const float state_init, const float x, const float dt, const float a, const float b, const float c, const float d, float state_out, float *y, int batch, int seq_len, int num_heads, int head_dim, int state_dim, int num_groups, int head_begin, int head_end)

void	mamba2_selective_scan_f32_parallel_dispatch (const float state_init, const float x, const float dt, const float a, const float b, const float c, const float d, float state_out, float *y, int batch, int seq_len, int num_heads, int head_dim, int state_dim, int num_groups)

void	mamba2_selective_state_update_decode_f32 (const float state_in, const float x, const float dt, const float a, const float b, const float c, const float d, float state_out, float *y, int rows, int num_heads, int head_dim, int state_dim, int num_groups)

void	mlp_token_parallel (const float input, const float W_fc1, const float b_fc1, const float W_fc2, const float b_fc2, float fc1_output, float *output, int T, int aligned_dim, int num_threads)

void	mlp_token_parallel_bf16 (const uint16_t input, const uint16_t W_fc1, const uint16_t b_fc1, const uint16_t W_fc2, const uint16_t b_fc2, float fc1_output, float output, int T, int aligned_dim, int num_threads, float scratch_bias1_f, float scratch_bias2_f, uint16_t scratch_fc1_bf16)

void	mlp_token_parallel_bf16_backward_mixed (const uint16_t input, const uint16_t W_fc1, const uint16_t b_fc1, const uint16_t W_fc2, const uint16_t d_output, float d_input, float d_W_fc1, float d_b_fc1, float d_W_fc2, float d_b_fc2, int T, int aligned_dim, int num_threads, float scratch_fc1_pre, uint16_t scratch_fc1_act_bf16, float *scratch_d_fc1)

void	mlp_token_parallel_bf16_fp32act (const uint16_t input, const uint16_t W_fc1, const uint16_t b_fc1, const uint16_t W_fc2, const uint16_t b_fc2, float fc1_output, float output, int T, int aligned_dim, int num_threads, float scratch_input_f, float scratch_bias1_f, float scratch_bias2_f, uint16_t *scratch_fc1_bf16)

void	mlp_token_parallel_exact (const float input, const float W_fc1, const float b_fc1, const float W_fc2, const float b_fc2, float fc1_output, float *output, int T, int aligned_dim, int num_threads)

void	moe_accumulate_expert_f32 (float output, const float expert_output, float routing_weight, int hidden_dim)
	Accumulate expert output: output += routing_weight * expert_output.

void	moe_relu2_expert_backward_f32 (const float d_output, const float hidden, const int indices, const float routing_weights, const float expert_up, const float expert_down, float d_hidden, float d_routing_weights, float d_expert_up, float d_expert_down, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

void	moe_relu2_expert_forward_f32 (const float hidden, const int indices, const float routing_weights, const float expert_up, const float expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

void	moe_relu2_expert_forward_q5_0_q5_0 (const float hidden, const int indices, const float routing_weights, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

void	moe_relu2_expert_forward_q5_0_q8_0 (const float hidden, const int indices, const float routing_weights, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

void	moe_relu2_shared_forward_q5_1_q8_0 (const float hidden, const float routed, const void shared_up, const void shared_down, float *output, int rows, int hidden_dim, int intermediate_dim)

int	moe_softmax_topk_router_llama_f32_workspace (const float logits, int indices, float weights, int rows, int n_experts, int top_k, float routed_scaling_factor, void workspace, size_t workspace_bytes)

int	moe_softmax_topk_router_pytorch_bf16_workspace (const float logits, int indices, float weights, int rows, int n_experts, int top_k, float routed_scaling_factor, void workspace, size_t workspace_bytes)

size_t	moe_softmax_topk_router_workspace_bytes (int n_experts)

void	moe_swiglu_expert_forward_bf16 (const float hidden, const int indices, const float routing_weights, const uint16_t expert_gate, const uint16_t expert_up, const uint16_t expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

void	moe_swiglu_expert_forward_bf16_parallel_dispatch (const float hidden, const int indices, const float routing_weights, const uint16_t expert_gate, const uint16_t expert_up, const uint16_t expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

void	moe_swiglu_expert_forward_bf16_row_range (const float hidden, const int indices, const float routing_weights, const uint16_t expert_gate, const uint16_t expert_up, const uint16_t expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, int row_begin, int row_end)

void	moe_swiglu_expert_forward_f32 (const float hidden, const int indices, const float routing_weights, const float expert_gate, const float expert_up, const float expert_down, float *output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

int	moe_swiglu_expert_forward_q4k_q4k_parallel_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q4k_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q5_0_parallel_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q5_0_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q5k_auto_prepared_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q5k_auto_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q5k_bucketed_prepared_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, const void expert_gate_packed, const void expert_up_packed, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q5k_parallel_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q5k_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q6k_parallel_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q6k_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q8_0_parallel_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

int	moe_swiglu_expert_forward_q4k_q8_0_workspace (const float hidden, const int indices, const float routing_weights, const void expert_gate, const void expert_up, const void expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes)

size_t	moe_swiglu_expert_q4k_q5k_workspace_bytes (int hidden_dim, int intermediate_dim)

size_t	moe_swiglu_expert_q4k_q8_0_workspace_bytes (int hidden_dim, int intermediate_dim)

void	moe_swiglu_packed_expert_forward_bf16 (const float hidden, const int indices, const float routing_weights, const uint16_t expert_gate_up, const uint16_t expert_down, float output, int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

void	moe_swiglu_shared_forward_bf16 (const float hidden, const float routed, const uint16_t shared_gate, const uint16_t shared_up, const uint16_t shared_down, float output, int rows, int hidden_dim, int intermediate_dim)

void	moe_swiglu_shared_forward_bf16_gated (const float hidden, const float routed, const uint16_t shared_gate, const uint16_t shared_up, const uint16_t shared_down, const uint16_t shared_router, float *output, int rows, int hidden_dim, int intermediate_dim)

void	moe_swiglu_shared_forward_bf16_gated_parallel_dispatch (const float hidden, const float routed, const uint16_t shared_gate, const uint16_t shared_up, const uint16_t shared_down, const uint16_t shared_router, float *output, int rows, int hidden_dim, int intermediate_dim)

void	moe_swiglu_shared_forward_bf16_gated_row_range (const float hidden, const float routed, const uint16_t shared_gate, const uint16_t shared_up, const uint16_t shared_down, const uint16_t shared_router, float *output, int rows, int hidden_dim, int intermediate_dim, int row_begin, int row_end)

void	moe_swiglu_shared_forward_bf16_parallel_dispatch (const float hidden, const float routed, const uint16_t shared_gate, const uint16_t shared_up, const uint16_t shared_down, float output, int rows, int hidden_dim, int intermediate_dim)

void	moe_swiglu_shared_forward_bf16_row_range (const float hidden, const float routed, const uint16_t shared_gate, const uint16_t shared_up, const uint16_t shared_down, float output, int rows, int hidden_dim, int intermediate_dim, int row_begin, int row_end)

void	moe_swiglu_shared_forward_f32 (const float hidden, const float routed, const float shared_gate, const float shared_up, const float shared_down, float output, int rows, int hidden_dim, int intermediate_dim)

int	moe_swiglu_shared_forward_q4k_q4k_parallel_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, float output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)

int	moe_swiglu_shared_forward_q4k_q4k_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, float output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)

int	moe_swiglu_shared_forward_q4k_q5_0_gated_parallel_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, const float shared_gate_input, float output, int rows, int hidden_dim, int intermediate_dim, void workspace, size_t workspace_bytes)

int	moe_swiglu_shared_forward_q4k_q5_0_gated_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, const float shared_gate_input, float output, int rows, int hidden_dim, int intermediate_dim, void workspace, size_t workspace_bytes)

int	moe_swiglu_shared_forward_q4k_q6k_parallel_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, float output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)

int	moe_swiglu_shared_forward_q4k_q6k_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, float output, int rows, int hidden_dim, int intermediate_dim, void *workspace, size_t workspace_bytes)

int	moe_swiglu_shared_forward_q4k_q8_0_gated_parallel_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, const float shared_gate_input, float output, int rows, int hidden_dim, int intermediate_dim, void workspace, size_t workspace_bytes)

int	moe_swiglu_shared_forward_q4k_q8_0_gated_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, const float shared_gate_input, float output, int rows, int hidden_dim, int intermediate_dim, void workspace, size_t workspace_bytes)

int	moe_swiglu_shared_forward_q8_0_gated_parallel_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, const float shared_gate_input, float output, int rows, int hidden_dim, int intermediate_dim, void workspace, size_t workspace_bytes)

int	moe_swiglu_shared_forward_q8_0_gated_workspace (const float hidden, const float routed, const void shared_gate, const void shared_up, const void shared_down, const float shared_gate_input, float output, int rows, int hidden_dim, int intermediate_dim, void workspace, size_t workspace_bytes)

size_t	moe_swiglu_shared_q4k_q8_0_gated_workspace_bytes (int hidden_dim, int intermediate_dim)

size_t	moe_swiglu_shared_q8_0_gated_workspace_bytes (int hidden_dim, int intermediate_dim)

void	mrope_qk_imrope_positions (float q, float k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void	mrope_qk_text (float q, float k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void	mrope_qk_text_imrope (float q, float k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void	mrope_qk_text_imrope_bf16_pytorch_storage (float q, float k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void	mrope_qk_text_imrope_positions_bf16_pytorch_storage (float q, float k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void	mrope_qk_vision (float q, float k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void	mrope_qk_vision_bf16_pytorch_storage (float q, float k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void	mrope_qk_vision_bf16_storage (float q, float k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void	mrope_qk_vision_fp16_storage (float q, float k, const int32_t *positions, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int n_dims, int section_0, int section_1, int section_2, int section_3, int n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow)

void	nemotron_group_limited_topk_router_f32 (const float scores, const float correction_bias, int indices, float weights, int rows, int n_experts, int top_k, int n_group, int topk_group, int norm_topk_prob, float routed_scaling_factor)

void	patch2im (const float d_patches, float d_image, int C, int H, int W, int P)

void	patch2im_bf16 (const uint16_t d_patches, uint16_t d_image, int C, int H, int W, int P)

void	patch_projection_bf16_pytorch_onednn_conv3d_storage (const float input, const void weights, const float bias, float output, int batch, int out_channels, int in_channels, int temporal, int patch_h, int patch_w)

void	patch_projection_image_bf16_native_storage (const float image, const void weights_t0, const void weights_t1, const float bias, float *output, int channels, int image_h, int image_w, int patch_size, int out_channels, int merge_size)

void	patch_projection_image_bf16_pytorch_onednn_conv3d_storage (const float image, const void weights_t0, const void weights_t1, const float bias, float *output, int channels, int image_h, int image_w, int patch_size, int out_channels, int merge_size)

void	position_embeddings_add (float x, const float position_embd, int num_tokens, int embed_dim, int num_positions)

void	position_embeddings_add_at_offset (float x, const float position_embd, int num_tokens, int embed_dim, int num_positions, int start_position)

void	position_embeddings_add_gemma4v_xy (float x, const float position_embd, int grid_h, int grid_w, int embed_dim, int source_grid_size)

void	position_embeddings_add_tiled_2d (float x, const float position_embd, int grid_h, int grid_w, int embed_dim, int merge_size, int source_grid_size)

void	position_embeddings_add_tiled_2d_align_corners (float x, const float position_embd, int grid_h, int grid_w, int embed_dim, int merge_size, int source_grid_size)

void	position_embeddings_add_tiled_2d_align_corners_bf16 (float x, const float position_embd, int grid_h, int grid_w, int embed_dim, int merge_size, int source_grid_size)

void	position_embeddings_add_tiled_2d_align_corners_fp32_interp_bf16 (float x, const float position_embd, int grid_h, int grid_w, int embed_dim, int merge_size, int source_grid_size)

void	q_norm_forward (float q, const float q_gamma, int num_heads, int num_tokens, int head_dim, float eps)

void	qk_norm_backward (const float d_q_out, const float d_k_out, const float q_in, const float k_in, const float q_gamma, const float k_gamma, float d_q_in, float d_k_in, float d_q_gamma, float d_k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

int	qk_norm_backward_last_isa (void)

void	qk_norm_forward (float q, float k, const float q_gamma, const float k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

void	qk_norm_forward_decode_exact (float q, float k, const float q_gamma, const float k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

void	qk_norm_forward_fp64_sum (float q, float k, const float q_gamma, const float k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

void	qk_norm_forward_llama_production (float q, float k, const float q_gamma, const float k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

void	qk_norm_forward_parallel_dispatch (float q, float k, const float q_gamma, const float k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

void	qk_norm_forward_prefill_exact (float q, float k, const float q_gamma, const float k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

void	qk_norm_forward_pytorch_bf16_storage (float q, float k, const float q_gamma, const float k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

void	qk_norm_forward_qwen4_pytorch_bf16_storage (float q, float k, const float q_gamma, const float k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)

void	quantize_batch_q8_0 (const float x, void y, int num_rows, int k)
	Batch quantize FP32 to Q8_0 format (row-major output)

void	quantize_batch_q8_k (const float x, void y, int num_rows, int k)
	Batch quantize FP32 to Q8_K format (row-major output)

void	quantize_batch_q8_k_4row_nearest_even (const float x, void y, int num_rows, int k)

void	quantize_row_q8_0 (const float x, void y, int k)
	Quantize FP32 to Q8_0 format (scalar reference)

void	quantize_row_q8_k (const float x, void y, int k)

void	qwen4_ple_gate_conv_inject_bf16 (const float hyper_input, const float key_projected, const float value_projected, const float norm_key_weight, const float norm_query_weight, const float norm_conv_weight, const uint16_t conv_weight, float hyper_output, float key_norm_scratch, float query_norm_scratch, float gated_scratch, float conv_norm_scratch, const float conv_state_in, float conv_state_out, int rows, int streams, int hidden_dim, int kernel_size, int dilation, float eps)

void	qwen4_ple_gate_conv_inject_fp16 (const float hyper_input, const float key_projected, const float value_projected, const float norm_key_weight, const float norm_query_weight, const float norm_conv_weight, const uint16_t conv_weight, float hyper_output, float key_norm_scratch, float query_norm_scratch, float gated_scratch, float conv_norm_scratch, const float conv_state_in, float conv_state_out, int rows, int streams, int hidden_dim, int kernel_size, int dilation, float eps)

void	qwen4_ple_gate_conv_inject_llama_fp16 (const float hyper_input, const float key_projected, const float value_projected, const float norm_key_weight, const float norm_query_weight, const float norm_conv_weight, const uint16_t conv_weight, float hyper_output, float key_norm_scratch, float query_norm_scratch, float gated_scratch, float conv_norm_scratch, const float conv_state_in, float conv_state_out, int rows, int streams, int hidden_dim, int kernel_size, int dilation, float eps)

void	qwen4_ple_ngram_embed_bf16 (const int32_t token_ids, const uint16_t embedding, const int64_t layer_multipliers, const int64_t head_offsets, const int64_t head_vocab_sizes, float output, const float token_state_in, float token_state_out, int rows, int ngram_size, int heads_per_ngram, int head_dim, int eos_token_id, int position)

void	qwen4_ple_ngram_embed_q5_0 (const int32_t token_ids, const void embedding, const int64_t layer_multipliers, const int64_t head_offsets, const int64_t head_vocab_sizes, float output, const float token_state_in, float token_state_out, int rows, int ngram_size, int heads_per_ngram, int head_dim, int eos_token_id, int position)

void	qwen4_qsa_index_select_bf16 (const float projected_qk, const float index_key_cache_in, const float q_norm_weight, const float k_norm_weight, float selected_indices, float index_key_cache_out, float q_norm_scratch, float pooled_key_scratch, float block_score_scratch, int32_t block_index_scratch, int rows, int query_heads, int index_head_dim, int token_budget, int compress_ratio, int rotary_dim, int context_length, int position, float rope_theta, float eps)

void	recurrent_conv_state_update_backward (const float d_conv_x, const float d_state_out, float d_state_in, float d_q, float d_k, float d_v, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)

void	recurrent_conv_state_update_backward_workspace (const float d_conv_x, const float d_state_out, float d_state_in, float d_q, float d_k, float d_v, float *d_conv_total, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)

void	recurrent_conv_state_update_forward (const float state_in, const float q, const float k, const float v, float conv_x, float state_out, int history_len, int num_seqs, int num_tokens, int q_dim, int k_dim, int v_dim)

void	recurrent_dt_gate_backward (const float d_gate, const float alpha, const float dt_bias, const float a, float d_alpha, float d_dt_bias, float *d_a, int rows, int dim)

void	recurrent_dt_gate_expanded_forward (const float alpha, const float dt_bias, const float a, float gate, int rows, int num_heads, int state_dim)

void	recurrent_dt_gate_forward (const float alpha, const float dt_bias, const float a, float gate, int rows, int num_heads, int state_dim)

void	recurrent_dt_gate_forward_pytorch_fp32 (const float alpha, const float dt_bias, const float a, float gate, int rows, int num_heads, int state_dim)

void	recurrent_norm_gate_backward (const float d_out, const float x, const float gate, const float weight, float d_x, float d_gate, float *d_weight, int rows, int num_heads, int head_dim, float eps)

void	recurrent_norm_gate_forward (const float x, const float gate, const float weight, float out, int rows, int num_heads, int head_dim, float eps)

void	recurrent_norm_gate_llama_avx2_forward (const float x, const float gate, const float weight, float out, int rows, int num_heads, int head_dim, float eps)

void	recurrent_norm_gate_pytorch_bf16_storage (const float x, const float gate, const float weight, float out, int rows, int num_heads, int head_dim, float eps)

void	recurrent_norm_sigmoid_gate_llama_avx2_forward (const float x, const float gate, const float weight, float out, int rows, int num_heads, int head_dim, float eps)

void	recurrent_norm_sigmoid_gate_pytorch_bf16_storage (const float x, const float gate, const float weight, float out, int rows, int num_heads, int head_dim, float eps)

void	recurrent_qk_l2_norm_backward (const float d_q_out, const float d_k_out, const float q, const float k, float d_q, float d_k, int rows, int q_dim, int k_dim, int head_dim, float eps)

void	recurrent_qk_l2_norm_forward (float q, float k, int rows, int q_dim, int k_dim, int head_dim, float eps)

void	recurrent_qk_l2_norm_pytorch_bf16_storage (float q, float k, int rows, int q_dim, int k_dim, int expanded_heads, int head_dim, float eps)

void	recurrent_qk_l2_norm_pytorch_fp32_output (float q, float k, int rows, int q_dim, int k_dim, int head_dim, float eps)

void	recurrent_sigmoid_forward_ggml (const float x, float out, int rows, int dim)

void	recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output (const float x, float out, int rows, int dim)

void	recurrent_silu_backward (const float d_out, const float x, float *d_x, int rows, int dim)

void	recurrent_silu_forward (const float x, float out, int rows, int dim)

void	recurrent_silu_forward_ggml (const float x, float out, int rows, int dim)

void	recurrent_silu_forward_pytorch_bf16_input_fp32_output (const float x, float out, int rows, int dim)

void	recurrent_silu_forward_pytorch_bf16_storage (const float x, float out, int rows, int dim)

void	recurrent_split_conv_qkv_backward (const float d_q, const float d_k, const float d_v, float d_packed_qkv, int rows, int q_dim, int k_dim, int v_dim)

void	recurrent_split_conv_qkv_forward (const float packed_qkv, float q, float k, float v, int rows, int q_dim, int k_dim, int v_dim)

void	recurrent_split_qkv_backward (const float d_q, const float d_k, const float d_v, float d_packed_qkv, int rows, int q_dim, int k_dim, int v_dim)

void	recurrent_split_qkv_forward (const float packed_qkv, float q, float k, float v, int rows, int q_dim, int k_dim, int v_dim)

void	relu2_backward (const float input, const float d_output, float *d_input, size_t n)

void	relu2_forward (const float input, float output, size_t n)

void	relu_backward (const float input, const float d_output, float *d_input, size_t n)

void	relu_backward_bf16 (const uint16_t input, const uint16_t d_output, uint16_t *d_input, size_t n)

void	relu_forward (const float input, float output, size_t n)

void	relu_forward_bf16 (const uint16_t input, uint16_t output, size_t n)

void	relu_forward_inplace (float *data, size_t n)

void	relu_forward_inplace_bf16 (uint16_t *data, size_t n)

void	rmsnorm_backward (const float d_output, const float input, const float gamma, const float rstd_cache, float d_input, float d_gamma, int tokens, int d_model, int aligned_embed_dim)

void	rmsnorm_backward_bf16 (const uint16_t d_output, const uint16_t input, const float gamma, const float rstd_cache, uint16_t d_input, float d_gamma, int tokens, int d_model, int aligned_embed_dim)

void	rmsnorm_backward_int4 (const uint8_t d_output, const uint8_t input, const float gamma, const float rstd_cache, uint8_t d_input, float d_gamma, int tokens, int d_model, int aligned_embed_dim, float scratch_d_output, float scratch_input, float *scratch_d_input)

void	rmsnorm_backward_int8 (const int8_t d_output, const int8_t input, const float gamma, const float rstd_cache, int8_t d_input, float d_gamma, int tokens, int d_model, int aligned_embed_dim, float scratch_d_output, float scratch_input, float *scratch_d_input)

void	rmsnorm_forward (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_bf16 (const uint16_t input, const float gamma, uint16_t output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_fp32_square_fp64_sum (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_fp64_sum (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_int4 (const uint8_t input, const float gamma, uint8_t output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps, float scratch_input, float scratch_output)

void	rmsnorm_forward_int8 (const int8_t input, const float gamma, int8_t output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps, float scratch_input, float scratch_output)

void	rmsnorm_forward_kv_lora (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_llama_production (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_no_weight (const float input, float output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_parallel_dispatch (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_pytorch_bf16_storage (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_qwen3next_pytorch_bf16_storage (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_strided_f32 (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)

void	rmsnorm_forward_strided_pytorch_bf16_storage (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)

void	rope_backward (const float d_out, float d_x, const float cos_cache, const float sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)

void	rope_backward_bf16 (const uint16_t d_out, uint16_t d_x, const float cos_cache, const float sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, float scratch_d_out, float scratch_d_x)

void	rope_backward_inplace (float d_x, const float cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)

void	rope_backward_qk (const float d_q_out, const float d_k_out, float d_q, float d_k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)

void	rope_backward_qk_bf16 (const uint16_t d_q_out, const uint16_t d_k_out, uint16_t d_q, uint16_t d_k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, float scratch_dq_out, float scratch_dq, float scratch_dk_out, float scratch_dk)

void	rope_backward_qk_pairwise_with_rotary_dim (const float d_q_out, const float d_k_out, float d_q, float d_k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)

void	rope_forward (float x, const float cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)

void	rope_forward_bf16 (uint16_t x, const float cos_cache, const float sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, float scratch)

void	rope_forward_bf16_with_rotary_dim (uint16_t x, const float cos_cache, const float sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float scratch)

void	rope_forward_q_split_direct_f32 (float q, const float freq_factors, int use_freq_factors, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base)

void	rope_forward_qk (float q, float k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset)

void	rope_forward_qk_bf16 (uint16_t q, uint16_t k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, float scratch_q, float scratch_k)

void	rope_forward_qk_bf16_with_rotary_dim (uint16_t q, uint16_t k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float scratch_q, float scratch_k)

void	rope_forward_qk_gemma4_direct (float q, float k, const float *freq_factors, int use_freq_factors, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base)

void	rope_forward_qk_gemma4v_vision_xy (float q, float k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int grid_w, int rotary_dim, float freq_base)

void	rope_forward_qk_pairwise_llama_cpu (float q, float k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)

void	rope_forward_qk_pairwise_with_rotary_dim (float q, float k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)

void	rope_forward_qk_split_direct_f32 (float q, float k, const float *freq_factors, int use_freq_factors, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base)

void	rope_forward_qk_split_direct_token_range_f32 (float q, float k, const float *freq_factors, int use_freq_factors, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base, int token_begin, int token_end)

void	rope_forward_qk_split_llama_token_range_f32 (float q, float k, const float *freq_factors, int use_freq_factors, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, float freq_base, int token_begin, int token_end)

void	rope_forward_qk_strided (float q, float k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int q_stride_tokens, int k_stride_tokens)

void	rope_forward_qk_strided_with_rotary_dim (float q, float k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int q_stride_tokens, int k_stride_tokens, int rotary_dim)

void	rope_forward_qk_with_rotary_dim (float q, float k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)

void	rope_forward_qk_with_rotary_dim_cache_stride (float q, float k, const float cos_cache, const float sin_cache, int num_heads, int num_kv_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim, int cache_rotary_dim)

void	rope_forward_strided (float x, const float cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int head_stride_tokens)

void	rope_forward_strided_with_rotary_dim (float x, const float cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int head_stride_tokens, int rotary_dim)

void	rope_forward_with_rotary_dim (float x, const float cos_cache, const float *sin_cache, int num_heads, int num_tokens, int head_dim, int aligned_head_dim, int pos_offset, int rotary_dim)

void	rope_precompute_cache (float cos_cache, float sin_cache, int max_seq_len, int head_dim, float base, int rotary_dim, const char *scaling_type, float scaling_factor)

void	rope_precompute_cache_llama_cpu (float cos_cache, float sin_cache, int max_seq_len, int head_dim, float base, int rotary_dim, const char *scaling_type, float scaling_factor)

void	rope_precompute_cache_split (float cos_cache, float sin_cache, int max_seq_len, int head_dim, float base)

void	rowwise_bias_add (float x, const float bias, int rows, int dim)

void	scal_copy_f32 (float y, const float x, float alpha, int n)
	Scaled copy: y = alpha * x.

void	sigmoid_backward (const float input, const float d_output, float *d_input, size_t n)

void	sigmoid_backward_bf16 (const uint16_t input, const uint16_t d_output, uint16_t d_input, size_t n, float scratch_input, float scratch_d_output, float scratch_d_input)

void	sigmoid_forward (const float input, float output, size_t n)

void	sigmoid_forward_bf16 (const uint16_t input, uint16_t output, size_t n, float scratch_input, float scratch_output)

float	sigmoid_scalar (float x)

void	softmax_cross_entropy_loss (const float logits, const int32_t targets, int tokens, int vocab_size, float d_logits, float loss_out)

void	softmax_cross_entropy_loss_bf16 (const uint16_t logits, const int32_t targets, int tokens, int vocab_size, uint16_t d_logits, float loss_out, float scratch_logits, float scratch_d_logits)

void	softmax_cross_entropy_loss_ptref (const float logits, const int32_t targets, int tokens, int vocab_size, float d_logits, float loss_out)

void	spatial_average_pool_contiguous (const float input, float output, int grid_h, int grid_w, int embed_dim, int merge_size)

void	spatial_merge_2x2 (const float input, float output, int grid_h, int grid_w, int embed_dim)

void	spatial_merge_contiguous_tiled (const float input, float output, int grid_h, int grid_w, int embed_dim, int merge_size)

void	speculative_commit_one_i32 (int accepted, int verified_token, int token_buffer, int token_count, int max_tokens, int target_position, int draft_position, int accepted_count, int rejected_count)
	Commit one verified speculative token and update decode counters.

void	speculative_verify_greedy_f32 (const float target_logits, int vocab_size, int draft_token, int accepted, int *verified_token)
	Greedy one-token speculative verification.

void	split_q_gate_backward (const float d_q, const float d_gate, float *d_packed_qg, int rows, int q_dim, int gate_dim, int group_dim)

void	split_q_gate_forward (const float packed_qg, float q, float *gate, int rows, int q_dim, int gate_dim, int group_dim)

void	split_qkv_packed_head_major_forward (const float packed_qkv, float q, float k, float v, int rows, int q_dim, int k_dim, int v_dim, int num_heads, int num_kv_heads)

void	ssm_conv1d_backward (const float d_out, const float conv_x, const float kernel, float d_conv_x, float *d_kernel, int kernel_size, int num_channels, int num_tokens, int num_seqs)

void	ssm_conv1d_forward (const float conv_x, const float kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)

void	ssm_conv1d_forward_llama_fma (const float conv_x, const float kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)

void	ssm_conv1d_forward_llama_production (const float conv_x, const float kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)

void	ssm_conv1d_forward_llama_production_serial (const float conv_x, const float kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)

void	ssm_conv1d_forward_pytorch_bf16_storage (const float conv_x, const float kernel, float *out, int kernel_size, int num_channels, int num_tokens, int num_seqs)

void	swiglu_backward (const float input, const float d_output, float *d_input, int tokens, int dim)

void	swiglu_backward_bf16 (const uint16_t input, const uint16_t d_output, uint16_t *d_input, int tokens, int dim)

void	swiglu_backward_exact (const float input, const float d_output, float *d_input, int tokens, int dim)

void	swiglu_forward (const float input, float output, int tokens, int dim)

void	swiglu_forward_bf16 (const uint16_t input, uint16_t output, int tokens, int dim)

void	swiglu_forward_exact (const float input, float output, int tokens, int dim)

void	swiglu_forward_ggml (const float input, float output, int tokens, int dim)

void	swiglu_forward_ggml_split (const float gate, const float up, float *output, int tokens, int dim)

void	swiglu_forward_pytorch_bf16_storage (const float input, float output, int tokens, int dim)

void	swiglu_forward_q8_k (const float input, void output_q8, int tokens, int dim)

void	topk_batched_f32 (const float scores, int num_tokens, int n_experts, int k, int indices, float *weights)
	Batched top-K selection for multiple tokens.

void	topk_f32 (const float scores, int n, int k, int indices, float *values)
	Find top-K indices and values from a score vector.

void	topk_softmax_backward_f32 (const int indices, const float weights, const float d_weights, float d_scores, int num_tokens, int n_experts_or_keys, int k)
	Backward for hard top-k followed by softmax over selected values.

void	topk_softmax_f32 (const float scores, int n, int k, int indices, float *weights)
	Find top-K indices with softmax-normalized weights.

void	unfused_rmsnorm_qkv_prefill (const float x, const float gamma, const float Wq, const float Wk, const float Wv, float x_norm, float Q, float K, float *V, int seq_len, int hidden, int q_dim, int kv_dim, float eps)
	Unfused version for benchmarking comparison.

void	vec_dot_q6_k_q8_k (int n, float s, const void vx, const void *vy)
	Q6_K x Q8_K dot product (single row)

void	vision_position_ids_2d_merge (int32_t *positions, int grid_h, int grid_w, int merge_size)

void	weighted_sum_f32 (float y, const float vectors, const float weights, int k, int n)
	Weighted sum of k vectors: y = sum_i(weights[i] * vectors[i])

void	yarn_rope_cache_contiguous_positions_f32 (float cos_cache, float sin_cache, int num_tokens, int rotary_dim, float freq_base, float factor, int original_context, float beta_fast, float beta_slow, float mscale, float mscale_all_dim)

void	yarn_rope_cache_explicit_positions_bf16 (uint16_t cos_cache, uint16_t sin_cache, const int32_t *positions, int num_tokens, int rotary_dim, float freq_base, float factor, int original_context, float beta_fast, float beta_slow, float mscale, float mscale_all_dim)

void	yarn_rope_cache_explicit_positions_f32 (float cos_cache, float sin_cache, const int32_t *positions, int num_tokens, int rotary_dim, float freq_base, float factor, int original_context, float beta_fast, float beta_slow, float mscale, float mscale_all_dim)

Enumeration Type Documentation

◆ ck_attention_prefill_schedule_t

enum ck_attention_prefill_schedule_t

Enumerator
CK_ATTN_PREFILL_SCHEDULE_KV_HEADS
CK_ATTN_PREFILL_SCHEDULE_QUERY_HEADS
CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES
CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES
CK_ATTN_PREFILL_SCHEDULE_GQA_SHARED_KV_TILES

Definition at line 2124 of file ckernel_engine.h.

             {
    CK_ATTN_PREFILL_SCHEDULE_KV_HEADS = 0,
    CK_ATTN_PREFILL_SCHEDULE_QUERY_HEADS = 1,
    CK_ATTN_PREFILL_SCHEDULE_QUERY_TILES = 2,
    CK_ATTN_PREFILL_SCHEDULE_KV_GROUP_QUERY_TILES = 3,
    CK_ATTN_PREFILL_SCHEDULE_GQA_SHARED_KV_TILES = 4,
} ck_attention_prefill_schedule_t;

◆ ck_attention_reduction_t

enum ck_attention_reduction_t

Enumerator
CK_ATTN_REDUCTION_FP32_ONLINE
CK_ATTN_REDUCTION_F16_ONLINE_FP32_MERGE
CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE
CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64
CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA

Definition at line 2114 of file ckernel_engine.h.

             {
    CK_ATTN_REDUCTION_FP32_ONLINE = 0,
    CK_ATTN_REDUCTION_F16_ONLINE_FP32_MERGE = 1,
    CK_ATTN_REDUCTION_F16_ONLINE_SINGLE_RANGE = 2,
    CK_ATTN_REDUCTION_F16_FLASH_AUTO_QTILE64 = 3,
    CK_ATTN_REDUCTION_BF16_PYTORCH_SDPA = 4,
} ck_attention_reduction_t;

◆ ck_attention_status_t

enum ck_attention_status_t

Enumerator
CK_ATTENTION_STATUS_OK
CK_ATTENTION_STATUS_INVALID_ARGUMENT
CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT
CK_ATTENTION_STATUS_INSUFFICIENT_WORKSPACE

Definition at line 2132 of file ckernel_engine.h.

             {
    CK_ATTENTION_STATUS_OK = 0,
    CK_ATTENTION_STATUS_INVALID_ARGUMENT = -1,
    CK_ATTENTION_STATUS_UNSUPPORTED_CONTRACT = -2,
    CK_ATTENTION_STATUS_INSUFFICIENT_WORKSPACE = -3,
} ck_attention_status_t;

Function Documentation

◆ adamw_clip_update_multi_f32()

void adamw_clip_update_multi_f32	(	float const	grads,
		float const	weights,
		float const	m_states,
		float const	v_states,
		const size_t *	numels,
		int	tensor_count,
		float	lr,
		float	beta1,
		float	beta2,
		float	eps,
		float	weight_decay,
		float	max_grad_norm,
		int	step
	)

Definition at line 638 of file optimizer_kernels.c.

{
    if (!grads || !weights || !m_states || !v_states || !numels || tensor_count <= 0) {
        return;
    }
 
    size_t total_numel = 0;
    int valid_tensors = 0;
    for (int i = 0; i < tensor_count; ++i) {
        if (grads[i] && weights[i] && m_states[i] && v_states[i] && numels[i] > 0) {
            total_numel += numels[i];
            valid_tensors += 1;
        }
    }
    if (valid_tensors == 0 || total_numel == 0) {
        return;
    }
 
    float grad_scale = 1.0f;
    if (max_grad_norm > 0.0f) {
        float global_norm = gradient_global_norm_multi_f32((const float *const *)grads, numels, tensor_count);
        if (global_norm > max_grad_norm) {
            grad_scale = max_grad_norm / global_norm;
        }
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
 
    if (!pool || nth <= 1 || nth > CK_OPT_PAR_MAX_THREADS ||
        total_numel < CK_OPT_PAR_MIN_NUMEL || valid_tensors < 2) {
        for (int i = 0; i < tensor_count; ++i) {
            float *g = grads[i];
            float *w = weights[i];
            float *m = m_states[i];
            float *v = v_states[i];
            size_t n = numels[i];
            if (!g || !w || !m || !v || n == 0) {
                continue;
            }
            if (grad_scale != 1.0f) {
                gradient_scale_f32_impl(g, n, grad_scale);
            }
            adamw_update_f32_impl(g, w, m, v, n, lr, beta1, beta2, eps, weight_decay, step);
        }
        return;
    }
    int active_nth = ck_opt_pick_active_threads(nth, total_numel, CK_OPT_PAR_MIN_NUMEL);
    if (active_nth > valid_tensors) {
        active_nth = valid_tensors;
    }
    if (active_nth <= 1) {
        for (int i = 0; i < tensor_count; ++i) {
            float *g = grads[i];
            float *w = weights[i];
            float *m = m_states[i];
            float *v = v_states[i];
            size_t n = numels[i];
            if (!g || !w || !m || !v || n == 0) {
                continue;
            }
            if (grad_scale != 1.0f) {
                gradient_scale_f32_impl(g, n, grad_scale);
            }
            adamw_update_f32_impl(g, w, m, v, n, lr, beta1, beta2, eps, weight_decay, step);
        }
        return;
    }
 
    ck_adamw_multi_parallel_args_t args = {
        .grads = grads,
        .weights = weights,
        .m_states = m_states,
        .v_states = v_states,
        .numels = numels,
        .tensor_count = tensor_count,
        .lr = lr,
        .beta1 = beta1,
        .beta2 = beta2,
        .eps = eps,
        .weight_decay = weight_decay,
        .grad_scale = grad_scale,
        .step = step,
    };
    ck_threadpool_dispatch_n(pool, active_nth, ck_adamw_multi_parallel_work, &args);
}

References adamw_update_f32_impl(), ck_adamw_multi_parallel_work(), CK_OPT_PAR_MAX_THREADS, CK_OPT_PAR_MIN_NUMEL, ck_opt_pick_active_threads(), ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), gradient_global_norm_multi_f32(), and gradient_scale_f32_impl().

◆ adamw_update_f32()

void adamw_update_f32	(	const float *	grad,
		float *	weight,
		float *	m,
		float *	v,
		size_t	numel,
		float	lr,
		float	beta1,
		float	beta2,
		float	eps,
		float	weight_decay,
		int	step
	)

Definition at line 592 of file optimizer_kernels.c.

{
    if (!grad || !weight || !m || !v || numel == 0) {
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
    if (!pool || nth <= 1 || nth > CK_OPT_PAR_MAX_THREADS || numel < CK_OPT_PAR_MIN_NUMEL) {
        adamw_update_f32_impl(grad, weight, m, v, numel, lr, beta1, beta2, eps, weight_decay, step);
        return;
    }
    int active_nth = ck_opt_pick_active_threads(nth, numel, CK_OPT_PAR_MIN_NUMEL);
    if (active_nth <= 1) {
        adamw_update_f32_impl(grad, weight, m, v, numel, lr, beta1, beta2, eps, weight_decay, step);
        return;
    }
 
    ck_adamw_parallel_args_t args = {
        .grad = grad,
        .weight = weight,
        .m = m,
        .v = v,
        .numel = numel,
        .lr = lr,
        .beta1 = beta1,
        .beta2 = beta2,
        .eps = eps,
        .weight_decay = weight_decay,
        .step = step,
    };
    ck_threadpool_dispatch_n(pool, active_nth, ck_adamw_parallel_work, &args);
}

References adamw_update_f32_impl(), ck_adamw_parallel_work(), CK_OPT_PAR_MAX_THREADS, CK_OPT_PAR_MIN_NUMEL, ck_opt_pick_active_threads(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

◆ add_backward_bf16()

void add_backward_bf16	(	const uint16_t *	d_y,
		uint16_t *	d_a,
		uint16_t *	d_b,
		size_t	n
	)

Definition at line 187 of file add_kernels_bf16.c.

{
    if (!d_y || n == 0) {
        return;
    }
 
    size_t i = 0;
 
    /* Copy to d_a if not in-place */
    if (d_a && d_a != d_y) {
#if defined(__AVX512F__)
        for (; i + 32 <= n; i += 32) {
            __m512i v0 = _mm512_loadu_si512((const __m512i*)&d_y[i]);
            __m512i v1 = _mm512_loadu_si512((const __m512i*)&d_y[i + 32]);
            _mm512_storeu_si512((__m512i*)&d_a[i], v0);
            _mm512_storeu_si512((__m512i*)&d_a[i + 32], v1);
        }
#endif
        for (; i < n; ++i) {
            d_a[i] = d_y[i];
        }
    }
 
    /* Copy to d_b if not in-place */
    i = 0;
    if (d_b && d_b != d_y) {
#if defined(__AVX512F__)
        for (; i + 32 <= n; i += 32) {
            __m512i v0 = _mm512_loadu_si512((const __m512i*)&d_y[i]);
            __m512i v1 = _mm512_loadu_si512((const __m512i*)&d_y[i + 32]);
            _mm512_storeu_si512((__m512i*)&d_b[i], v0);
            _mm512_storeu_si512((__m512i*)&d_b[i + 32], v1);
        }
#endif
        for (; i < n; ++i) {
            d_b[i] = d_y[i];
        }
    }
}

◆ add_forward_2d_bf16()

void add_forward_2d_bf16	(	const uint16_t *	a,
		const uint16_t *	b,
		uint16_t *	y,
		int	tokens,
		int	dim,
		int	aligned_dim
	)

Definition at line 235 of file add_kernels_bf16.c.

{
    if (!a || !b || !y || tokens <= 0 || dim <= 0) {
        return;
    }
 
    for (int t = 0; t < tokens; ++t) {
        const uint16_t *a_row = a + (size_t)t * aligned_dim;
        const uint16_t *b_row = b + (size_t)t * aligned_dim;
        uint16_t *y_row = y + (size_t)t * aligned_dim;
 
        int d = 0;
 
#if defined(__AVX512F__)
        for (; d + 16 <= dim; d += 16) {
            __m512 av = bf16_loadu_cvt_fp32(&a_row[d]);
            __m512 bv = bf16_loadu_cvt_fp32(&b_row[d]);
            __m512 yv = _mm512_add_ps(av, bv);
            fp32_cvt_storeu_bf16(&y_row[d], yv);
        }
#endif
 
        for (; d < dim; ++d) {
            float af = bf16_to_float(a_row[d]);
            float bf = bf16_to_float(b_row[d]);
            y_row[d] = float_to_bf16(af + bf);
        }
    }
}

References bf16_to_float(), and float_to_bf16().

◆ add_forward_bf16()

void add_forward_bf16	(	const uint16_t *	a,
		const uint16_t *	b,
		uint16_t *	y,
		size_t	n
	)

Definition at line 38 of file add_kernels_bf16.c.

{
    if (!a || !b || !y || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    /* AVX-512: Process 16 bf16 elements at a time */
    for (; i + 16 <= n; i += 16) {
        __m512 av = bf16_loadu_cvt_fp32(&a[i]);
        __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
        __m512 yv = _mm512_add_ps(av, bv);
        fp32_cvt_storeu_bf16(&y[i], yv);
    }
#endif
 
    /* Scalar fallback */
    for (; i < n; ++i) {
        float af = bf16_to_float(a[i]);
        float bf = bf16_to_float(b[i]);
        y[i] = float_to_bf16(af + bf);
    }
}

References bf16_to_float(), and float_to_bf16().

◆ add_forward_f32()

void add_forward_f32	(	const float *	a,
		const float *	b,
		float *	y,
		size_t	n
	)

Element-wise add: y = a + b

Test:

test_add.py::TestAddForward::test_add_forward_f32

test_add.py::TestAddForward::test_add_inplace_f32

test_multi_layer_parity.py::TestMultiLayerParity::test_residual_add

Element-wise addition of two vectors.

After changes: make test

Definition at line 284 of file add_kernels_bf16.c.

{
    if (!a || !b || !y || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    for (; i + 16 <= n; i += 16) {
        __m512 av = _mm512_loadu_ps(&a[i]);
        __m512 bv = _mm512_loadu_ps(&b[i]);
        __m512 yv = _mm512_add_ps(av, bv);
        _mm512_storeu_ps(&y[i], yv);
    }
#endif
 
#if defined(__AVX2__)
    for (; i + 8 <= n; i += 8) {
        __m256 av = _mm256_loadu_ps(&a[i]);
        __m256 bv = _mm256_loadu_ps(&b[i]);
        __m256 yv = _mm256_add_ps(av, bv);
        _mm256_storeu_ps(&y[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        y[i] = a[i] + b[i];
    }
}

◆ add_inplace_bf16()

void add_inplace_bf16	(	uint16_t *	a,
		const uint16_t *	b,
		size_t	n
	)

Definition at line 119 of file add_kernels_bf16.c.

{
    if (!a || !b || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    for (; i + 16 <= n; i += 16) {
        __m512 av = bf16_loadu_cvt_fp32(&a[i]);
        __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
        __m512 yv = _mm512_add_ps(av, bv);
        fp32_cvt_storeu_bf16(&a[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        float af = bf16_to_float(a[i]);
        float bf = bf16_to_float(b[i]);
        a[i] = float_to_bf16(af + bf);
    }
}

References bf16_to_float(), and float_to_bf16().

◆ add_inplace_f32()

void add_inplace_f32	(	float *	a,
		const float *	b,
		size_t	n
	)

Definition at line 318 of file add_kernels_bf16.c.

{
    if (!a || !b || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    for (; i + 16 <= n; i += 16) {
        __m512 av = _mm512_loadu_ps(&a[i]);
        __m512 bv = _mm512_loadu_ps(&b[i]);
        __m512 yv = _mm512_add_ps(av, bv);
        _mm512_storeu_ps(&a[i], yv);
    }
#endif
 
#if defined(__AVX2__)
    for (; i + 8 <= n; i += 8) {
        __m256 av = _mm256_loadu_ps(&a[i]);
        __m256 bv = _mm256_loadu_ps(&b[i]);
        __m256 yv = _mm256_add_ps(av, bv);
        _mm256_storeu_ps(&a[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        a[i] = a[i] + b[i];
    }
}

Referenced by mega_fused_outproj_mlp_prefill().

◆ add_scaled_forward_bf16()

void add_scaled_forward_bf16	(	const uint16_t *	a,
		const uint16_t *	b,
		uint16_t *	y,
		float	alpha,
		size_t	n
	)

Definition at line 86 of file add_kernels_bf16.c.

{
    if (!a || !b || !y || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 alpha_v = _mm512_set1_ps(alpha);
    for (; i + 16 <= n; i += 16) {
        __m512 av = bf16_loadu_cvt_fp32(&a[i]);
        __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
        __m512 yv = _mm512_fmadd_ps(bv, alpha_v, av);  /* a + alpha * b */
        fp32_cvt_storeu_bf16(&y[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        float af = bf16_to_float(a[i]);
        float bf = bf16_to_float(b[i]);
        y[i] = float_to_bf16(af + alpha * bf);
    }
}

References bf16_to_float(), and float_to_bf16().

◆ add_scaled_inplace_bf16()

void add_scaled_inplace_bf16	(	uint16_t *	a,
		const uint16_t *	b,
		float	alpha,
		size_t	n
	)

Definition at line 149 of file add_kernels_bf16.c.

{
    if (!a || !b || n == 0) {
        return;
    }
 
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 alpha_v = _mm512_set1_ps(alpha);
    for (; i + 16 <= n; i += 16) {
        __m512 av = bf16_loadu_cvt_fp32(&a[i]);
        __m512 bv = bf16_loadu_cvt_fp32(&b[i]);
        __m512 yv = _mm512_fmadd_ps(bv, alpha_v, av);
        fp32_cvt_storeu_bf16(&a[i], yv);
    }
#endif
 
    for (; i < n; ++i) {
        float af = bf16_to_float(a[i]);
        float bf = bf16_to_float(b[i]);
        a[i] = float_to_bf16(af + alpha * bf);
    }
}

References bf16_to_float(), and float_to_bf16().

◆ add_stream_inplace()

void add_stream_inplace	(	float *	a,
		const float *	b,
		size_t	n
	)

Definition at line 713 of file vision_kernels.c.

{
    if (!a || !b || n == 0) {
        return;
    }
 
    for (size_t i = 0; i < n; ++i) {
        a[i] += b[i];
    }
}

◆ add_stream_reorder_2d()

void add_stream_reorder_2d	(	float *	main_inout,
		float *	aux_scratch,
		int	grid_h,
		int	grid_w,
		int	embed_dim,
		int	merge_size
	)

Definition at line 726 of file vision_kernels.c.

{
    if (!main_inout || !aux_scratch || grid_h <= 0 || grid_w <= 0 || embed_dim <= 0 || merge_size <= 0) {
        return;
    }
 
    const int num_tokens = grid_h * grid_w;
    const size_t total_elems = (size_t) num_tokens * (size_t) embed_dim;
 
    for (size_t i = 0; i < total_elems; ++i) {
        main_inout[i] += aux_scratch[i];
    }
 
    for (int tok = 0; tok < num_tokens; ++tok) {
        const int src_tok = tile_order_index_2d(tok, grid_h, grid_w, merge_size);
        const float *src_main = main_inout + (size_t) src_tok * (size_t) embed_dim;
        float *dst = aux_scratch + (size_t) tok * (size_t) embed_dim;
        for (int d = 0; d < embed_dim; ++d) {
            dst[d] = src_main[d];
        }
    }
 
    memcpy(main_inout, aux_scratch, total_elems * sizeof(float));
}

References tile_order_index_2d().

◆ argmax_f32()

int argmax_f32	(	const float *	scores,
		int	n
	)

Find index of maximum value.

Parameters

scores	Input scores [n]
n	Number of scores

Returns: Index of maximum value

Definition at line 500 of file topk_kernels.c.

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

Referenced by speculative_verify_greedy_f32().

◆ assistant_layer_scale_forward()

void assistant_layer_scale_forward	(	float *	hidden,
		const float *	scale,
		int	tokens,
		int	embed_dim
	)

Definition at line 389 of file gemma4_per_layer_embed.c.

{
    if (!hidden || !scale || tokens <= 0 || embed_dim <= 0) {
        return;
    }
 
    const float s = scale[0];
    const size_t n = (size_t)tokens * (size_t)embed_dim;
    for (size_t i = 0; i < n; ++i) {
        hidden[i] *= s;
    }
}

◆ 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 9175 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 9036 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 8983 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_decode()

void attention_flash_decode	(	float *	out,
		const float *	q,
		const float *	k,
		const float *	v,
		int	T_q,
		int	T_k,
		int	H,
		int	D_h,
		float	scale
	)

Main flash attention function with SIMD dispatch.

Parameters

out	Output [T_q, H, D_h]
q	Query [T_q, H, D_h]
k	Key [T_k, H, D_h]
v	Value [T_k, H, D_h]
T_q	Number of query tokens (1 for decode)
T_k	Number of key/value tokens (context length)
H	Number of heads
D_h	Head dimension
scale	1/sqrt(D_h)

Definition at line 696 of file attention_flash_true.c.

{
    if (!out || !q || !k || !v) {
        return;
    }
    if (T_q <= 0 || T_k <= 0 || H <= 0 || D_h <= 0) {
        return;
    }
 
    // Dispatch based on CPU features
#if defined(__AVX512F__)
    attention_flash_decode_avx512(out, q, k, v, T_q, T_k, H, D_h, scale);
#elif defined(__AVX__) && !defined(__AVX512F__)
    attention_flash_decode_avx(out, q, k, v, T_q, T_k, H, D_h, scale);
#else
    attention_flash_decode_scalar(out, q, k, v, T_q, T_k, H, D_h, scale);
#endif
}

References attention_flash_decode_scalar().

Referenced by attention_forward_chunk_head_major_gqa_flash_gemma4(), attention_forward_decode_head_major_gqa_flash(), attention_forward_decode_head_major_gqa_flash_f16cache(), attention_forward_decode_head_major_gqa_flash_gemma4(), ck_attention_flash_decode_wrapper(), mega_fused_attention_prefill(), and mega_fused_attention_prefill_q8_0().

◆ 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 1281 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 1357 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 1435 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 1577 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 1505 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 4370 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 4422 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 5343 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_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 5303 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 5024 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 5047 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_sliding()

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

Flash attention forward with sliding window (prefill)

Test:: test_attention.py::TestAttentionForward::test_sliding_window_prefill

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

After changes: make test

Definition at line 566 of file attention_kernels_sliding.c.

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

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

◆ attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4()

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

Definition at line 763 of file attention_kernels_sliding.c.

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

References attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl().

Referenced by attention_forward_causal_head_major_shared_kv_sliding_gemma4().

◆ attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_token_output()

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

Definition at line 782 of file attention_kernels_sliding.c.

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

References attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4_impl().

◆ attention_forward_causal_head_major_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 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,
                                                /*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 5960 of file attention_kernels.c.

{
    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);
}

References ck_attention_llama_regular_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 8673 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 8596 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(), 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 7335 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 7641 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 7419 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,
            (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_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 7234 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 6716 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 7392 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,
        (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_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 8713 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 7616 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_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 5070 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_causal_head_major_shared_kv_sliding_gemma4()

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

Definition at line 801 of file attention_kernels_sliding.c.

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

References attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4().

◆ attention_forward_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 6089 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 8522 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_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 6002 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 8750 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 6586 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 6565 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);
}

References attention_forward_decode_head_major_gqa_flash_f16cache_split_partitioned().

Referenced by attention_forward_decode_head_major_gqa_flash_f16cache_contract().

◆ 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 6151 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 6046 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_flash_sliding()

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

Definition at line 817 of file attention_kernels_sliding.c.

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

References attention_forward_decode_head_major_gqa_flash(), and SLIDING_DECODE_IMPL.

◆ attention_forward_decode_head_major_gqa_flash_sliding_gemma4()

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

Definition at line 903 of file attention_kernels_sliding.c.

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

References attention_forward_decode_head_major_gqa_flash_gemma4(), and SLIDING_DECODE_IMPL_GEMMA4.

Referenced by attention_forward_decode_head_major_shared_kv_sliding_gemma4().

◆ attention_forward_decode_head_major_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 5975 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_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 8873 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 6135 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_decode_head_major_shared_kv_sliding_gemma4()

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

Definition at line 976 of file attention_kernels_sliding.c.

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

References attention_forward_decode_head_major_gqa_flash_sliding_gemma4().

◆ attention_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 5372 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 4391 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 4467 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 4920 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(), 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 5084 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 5690 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 5449 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 4973 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 4998 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 4948 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(), 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 3564 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 3547 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 3527 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_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 5216 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_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 5228 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_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 227 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 3581 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 8452 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.

◆ attn_gate_sigmoid_mul_backward()

void attn_gate_sigmoid_mul_backward	(	const float *	d_out,
		const float *	x,
		const float *	gate,
		float *	d_x,
		float *	d_gate,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 187 of file hybrid_attention_kernels.c.

                                                   {
    const int dim = num_heads * state_dim;
    for (int row = 0; row < rows; ++row) {
        const float *d_out_row = d_out + (size_t) row * (size_t) dim;
        const float *x_row = x + (size_t) row * (size_t) dim;
        const float *gate_row = gate + (size_t) row * (size_t) dim;
        float *d_x_row = d_x + (size_t) row * (size_t) dim;
        float *d_gate_row = d_gate + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float sig = hybrid_sigmoid(gate_row[col]);
            d_x_row[col] = d_out_row[col] * sig;
            d_gate_row[col] = d_out_row[col] * x_row[col] * sig * (1.0f - sig);
        }
    }
}

References hybrid_sigmoid().

◆ attn_gate_sigmoid_mul_forward()

void attn_gate_sigmoid_mul_forward	(	const float *	x,
		const float *	gate,
		float *	out,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 116 of file hybrid_attention_kernels.c.

                                                  {
    const int dim = num_heads * state_dim;
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) dim;
        const float *gate_row = gate + (size_t) row * (size_t) dim;
        float *out_row = out + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            out_row[col] = x_row[col] * hybrid_sigmoid(gate_row[col]);
        }
    }
}

References hybrid_sigmoid().

◆ attn_gate_sigmoid_mul_pytorch_bf16_storage()

void attn_gate_sigmoid_mul_pytorch_bf16_storage	(	const float *	x,
		const float *	gate,
		float *	out,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 133 of file hybrid_attention_kernels.c.

                                                               {
    if (!x || !gate || !out || rows <= 0 || num_heads <= 0 || state_dim <= 0) {
        return;
    }
    const int dim = num_heads * state_dim;
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t)row * (size_t)dim;
        const float *gate_row = gate + (size_t)row * (size_t)dim;
        float *out_row = out + (size_t)row * (size_t)dim;
        for (int col = 0; col < dim; col += 16) {
            const int width = dim - col < 16 ? dim - col : 16;
            float sigmoid[16];
            recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output(
                gate_row + col, sigmoid, 1, width);
            for (int lane = 0; lane < width; ++lane) {
                const float x_bf16 = bf16_to_float(float_to_bf16(x_row[col + lane]));
                const float sigmoid_bf16 = bf16_to_float(float_to_bf16(sigmoid[lane]));
                out_row[col + lane] = bf16_to_float(float_to_bf16(
                    x_bf16 * sigmoid_bf16));
            }
        }
    }
}

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

◆ attn_gate_softplus_mul_forward()

void attn_gate_softplus_mul_forward	(	const float *	x,
		const float *	gate,
		float *	out,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 162 of file hybrid_attention_kernels.c.

                                                   {
    const int dim = num_heads * state_dim;
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) dim;
        const float *gate_row = gate + (size_t) row * (size_t) num_heads;
        float *out_row = out + (size_t) row * (size_t) dim;
        for (int head = 0; head < num_heads; ++head) {
            const float value = gate_row[head];
            const float scale = value > 20.0f
                ? value
                : log1pf(expf(value));
            const size_t base = (size_t) head * (size_t) state_dim;
            for (int col = 0; col < state_dim; ++col) {
                out_row[base + (size_t) col] =
                    x_row[base + (size_t) col] * scale;
            }
        }
    }
}

◆ axpy_2d_f32()

void axpy_2d_f32	(	float *	Y,
		const float *	X,
		float	alpha,
		int	num_tokens,
		int	dim,
		int	y_stride,
		int	x_stride
	)

Batched AXPY for 2D tensors: Y[t,:] += alpha * X[t,:].

Parameters

Y	Output tensor [num_tokens, dim]
X	Input tensor [num_tokens, dim]
alpha	Scalar multiplier
num_tokens	Number of tokens
dim	Hidden dimension
y_stride	Stride between Y rows (for alignment)
x_stride	Stride between X rows

Definition at line 272 of file axpy_kernels.c.

{
    if (!Y || !X || num_tokens <= 0 || dim <= 0) {
        return;
    }
 
    /* Default strides if not specified */
    if (y_stride <= 0) y_stride = dim;
    if (x_stride <= 0) x_stride = dim;
 
    for (int t = 0; t < num_tokens; t++) {
        axpy_f32(Y + t * y_stride, X + t * x_stride, alpha, dim);
    }
}

References axpy_f32().

◆ axpy_f32()

void axpy_f32	(	float *	y,
		const float *	x,
		float	alpha,
		int	n
	)

In-place AXPY: y += alpha * x.

Test:

test_axpy.py::TestAXPY::test_axpy_f32

test_axpy.py::TestAXPY::test_axpy_vs_naive

In-place scaled vector addition: y += alpha * x BLAS-like axpy operation.

After changes: make test

Definition at line 105 of file axpy_kernels.c.

{
    if (!y || !x || n <= 0) {
        return;
    }
 
    int i = 0;
 
#ifdef __AVX512F__
    __m512 valpha = _mm512_set1_ps(alpha);
    for (; i + 16 <= n; i += 16) {
        __m512 vy = _mm512_loadu_ps(&y[i]);
        __m512 vx = _mm512_loadu_ps(&x[i]);
        vy = _mm512_fmadd_ps(vx, valpha, vy);  /* y = y + alpha * x */
        _mm512_storeu_ps(&y[i], vy);
    }
#endif
 
#ifdef __AVX2__
    __m256 valpha256 = _mm256_set1_ps(alpha);
    for (; i + 8 <= n; i += 8) {
        __m256 vy = _mm256_loadu_ps(&y[i]);
        __m256 vx = _mm256_loadu_ps(&x[i]);
        vy = _mm256_fmadd_ps(vx, valpha256, vy);
        _mm256_storeu_ps(&y[i], vy);
    }
#endif
 
    /* Scalar remainder */
    for (; i < n; i++) {
        y[i] += alpha * x[i];
    }
}

Referenced by axpy_2d_f32(), axpy_zero_f32(), ck_moe_q4k_mixed_route_parallel(), ck_moe_q4k_q5k_bucket_work(), ck_moe_q4k_q5k_route_parallel(), moe_accumulate_expert_f32(), moe_swiglu_expert_forward_q4k_q4k_workspace(), moe_swiglu_expert_forward_q4k_q5k_workspace(), moe_swiglu_expert_forward_q4k_q6k_workspace(), and weighted_sum_f32().

◆ axpy_zero_f32()

void axpy_zero_f32	(	float *	y,
		const float *	x,
		float	alpha,
		int	n
	)

Zero output then accumulate: y = 0; y += alpha * x.

Parameters

y	Output vector [n], zeroed then accumulated
x	Input vector [n]
alpha	Scalar multiplier
n	Vector length

Definition at line 239 of file axpy_kernels.c.

{
    if (!y || n <= 0) {
        return;
    }
 
    memset(y, 0, n * sizeof(float));
 
    if (x) {
        axpy_f32(y, x, alpha, n);
    }
}

References axpy_f32().

◆ backward_causal_softmax_head_major()

void backward_causal_softmax_head_major	(	float *	d_scores,
		const float *	weights,
		int	num_heads,
		int	num_tokens,
		int	aligned_context_window
	)

Definition at line 382 of file softmax_kernels.c.

{
    int H = num_heads;
    int T = num_tokens;
 
    for (int h = 0; h < H; ++h) {
        for (int i = 0; i < T; ++i) {
            int base = h * aligned_context_window * aligned_context_window
                     + i * aligned_context_window;
            float *drow = &d_scores[base];
            const float *wrow = &weights[base];
            int len = i + 1;
 
#if defined(__AVX512F__)
            // Compute dot product (vectorized)
            __m512 dot_vec = _mm512_setzero_ps();
            int j = 0;
            for (; j + 16 <= len; j += 16) {
                __m512 w = _mm512_loadu_ps(&wrow[j]);
                __m512 dw = _mm512_loadu_ps(&drow[j]);
                dot_vec = _mm512_fmadd_ps(w, dw, dot_vec);
            }
            float dot_product = _mm512_reduce_add_ps(dot_vec);
            for (; j < len; ++j) {
                dot_product += wrow[j] * drow[j];
            }
 
            // Compute gradient: d_scores = w * (dw - dot_product)
            __m512 dot_broadcast = _mm512_set1_ps(dot_product);
            j = 0;
            for (; j + 16 <= len; j += 16) {
                __m512 w = _mm512_loadu_ps(&wrow[j]);
                __m512 dw = _mm512_loadu_ps(&drow[j]);
                __m512 diff = _mm512_sub_ps(dw, dot_broadcast);
                __m512 result = _mm512_mul_ps(w, diff);
                _mm512_storeu_ps(&drow[j], result);
            }
            for (; j < len; ++j) {
                drow[j] = wrow[j] * (drow[j] - dot_product);
            }
 
            // Zero out future tokens
            __m512 zero = _mm512_setzero_ps();
            for (; j + 16 <= T; j += 16) {
                _mm512_storeu_ps(&drow[j], zero);
            }
            for (; j < T; ++j) {
                drow[j] = 0.0f;
            }
 
#elif defined(__AVX__)
            // Compute dot product (vectorized)
            __m256 dot_vec = _mm256_setzero_ps();
            int j = 0;
            for (; j + 8 <= len; j += 8) {
                __m256 w = _mm256_loadu_ps(&wrow[j]);
                __m256 dw = _mm256_loadu_ps(&drow[j]);
                // No FMA in AVX1: use mul + add
                __m256 prod = _mm256_mul_ps(w, dw);
                dot_vec = _mm256_add_ps(dot_vec, prod);
            }
            float dot_product = hsum256_ps_softmax(dot_vec);
            for (; j < len; ++j) {
                dot_product += wrow[j] * drow[j];
            }
 
            // Compute gradient: d_scores = w * (dw - dot_product)
            __m256 dot_broadcast = _mm256_set1_ps(dot_product);
            j = 0;
            for (; j + 8 <= len; j += 8) {
                __m256 w = _mm256_loadu_ps(&wrow[j]);
                __m256 dw = _mm256_loadu_ps(&drow[j]);
                __m256 diff = _mm256_sub_ps(dw, dot_broadcast);
                __m256 result = _mm256_mul_ps(w, diff);
                _mm256_storeu_ps(&drow[j], result);
            }
            for (; j < len; ++j) {
                drow[j] = wrow[j] * (drow[j] - dot_product);
            }
 
            // Zero out future tokens
            __m256 zero = _mm256_setzero_ps();
            for (; j + 8 <= T; j += 8) {
                _mm256_storeu_ps(&drow[j], zero);
            }
            for (; j < T; ++j) {
                drow[j] = 0.0f;
            }
 
#else
            // Scalar fallback
            float dot_product = 0.0f;
            for (int j = 0; j < len; ++j) {
                dot_product += wrow[j] * drow[j];
            }
 
            for (int j = 0; j < len; ++j) {
                drow[j] = wrow[j] * (drow[j] - dot_product);
            }
 
            for (int j = len; j < T; ++j) {
                drow[j] = 0.0f;
            }
#endif
        }
    }
}

Referenced by backward_causal_softmax_head_major_bf16().

◆ backward_causal_softmax_head_major_bf16()

void backward_causal_softmax_head_major_bf16	(	uint16_t *	d_scores,
		const uint16_t *	weights,
		int	num_heads,
		int	num_tokens,
		int	aligned_context_window,
		float *	scratch_d_scores,
		float *	scratch_weights
	)

Definition at line 53 of file softmax_kernels_bf16.c.

{
    if (!d_scores || !weights || num_heads <= 0 || num_tokens <= 0 || aligned_context_window <= 0) return;
    if (!scratch_d_scores || !scratch_weights) return;
 
    const size_t total = (size_t)num_heads *
                         (size_t)aligned_context_window *
                         (size_t)aligned_context_window;
 
    bf16_tensor_to_float(d_scores, scratch_d_scores, total);
    bf16_tensor_to_float(weights, scratch_weights, total);
    backward_causal_softmax_head_major(scratch_d_scores, scratch_weights, num_heads, num_tokens, aligned_context_window);
    float_tensor_to_bf16(scratch_d_scores, d_scores, total);
}

References backward_causal_softmax_head_major(), bf16_tensor_to_float(), and float_tensor_to_bf16().

◆ causal_softmax_head_major()

void causal_softmax_head_major	(	float *	scores,
		int	num_heads,
		int	num_tokens,
		int	aligned_context_window
	)

Causal softmax (in-place, row-wise)

Test:

test_softmax.py::TestSoftmaxForward::test_causal_softmax

test_softmax.py::TestSoftmaxForward::test_causal_vs_softmax

test_attention.py::TestAttentionForward::test_softmax_correctness

Applies causal mask (j > i => 0) and softmax to scores matrix. In-place on [num_heads, T, T] scores matrix.

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

Definition at line 144 of file softmax_kernels.c.

{
    for (int h = 0; h < num_heads; ++h) {
        for (int i = 0; i < num_tokens; ++i) {
            int base = h * aligned_context_window * aligned_context_window
                     + i * aligned_context_window;
            float *row = &scores[base];
            int len = i + 1;  // Number of valid elements (0..i inclusive)
 
#if defined(__AVX512F__)
            // Find max (vectorized)
            __m512 max_vec = _mm512_set1_ps(-INFINITY);
            int j = 0;
            for (; j + 16 <= len; j += 16) {
                __m512 v = _mm512_loadu_ps(&row[j]);
                max_vec = _mm512_max_ps(max_vec, v);
            }
            float max_val = _mm512_reduce_max_ps(max_vec);
            for (; j < len; ++j) {
                if (row[j] > max_val) max_val = row[j];
            }
 
            // Compute exp and sum (vectorized)
            __m512 max_broadcast = _mm512_set1_ps(max_val);
            __m512 sum_vec = _mm512_setzero_ps();
            j = 0;
            for (; j + 16 <= len; j += 16) {
                __m512 v = _mm512_loadu_ps(&row[j]);
                __m512 e = exp512_approx(_mm512_sub_ps(v, max_broadcast));
                _mm512_storeu_ps(&row[j], e);
                sum_vec = _mm512_add_ps(sum_vec, e);
            }
            float sum = _mm512_reduce_add_ps(sum_vec);
            for (; j < len; ++j) {
                float e = expf(row[j] - max_val);
                row[j] = e;
                sum += e;
            }
 
            // Normalize (vectorized)
            float inv_sum = 1.0f / sum;
            __m512 inv_sum_vec = _mm512_set1_ps(inv_sum);
            j = 0;
            for (; j + 16 <= len; j += 16) {
                __m512 v = _mm512_loadu_ps(&row[j]);
                _mm512_storeu_ps(&row[j], _mm512_mul_ps(v, inv_sum_vec));
            }
            for (; j < len; ++j) {
                row[j] *= inv_sum;
            }
 
            // Zero out future tokens (vectorized)
            __m512 zero = _mm512_setzero_ps();
            for (; j + 16 <= num_tokens; j += 16) {
                _mm512_storeu_ps(&row[j], zero);
            }
            for (; j < num_tokens; ++j) {
                row[j] = 0.0f;
            }
 
#elif defined(__AVX2__)
            // AVX2: Find max (vectorized)
            __m256 max_vec = _mm256_set1_ps(-INFINITY);
            int j = 0;
            for (; j + 8 <= len; j += 8) {
                __m256 v = _mm256_loadu_ps(&row[j]);
                max_vec = _mm256_max_ps(max_vec, v);
            }
            float max_val = hmax256_ps(max_vec);
            for (; j < len; ++j) {
                if (row[j] > max_val) max_val = row[j];
            }
 
            // Compute exp and sum (vectorized with fast exp)
            __m256 max_broadcast = _mm256_set1_ps(max_val);
            __m256 sum_vec = _mm256_setzero_ps();
            j = 0;
            for (; j + 8 <= len; j += 8) {
                __m256 v = _mm256_loadu_ps(&row[j]);
                __m256 e = exp256_approx(_mm256_sub_ps(v, max_broadcast));
                _mm256_storeu_ps(&row[j], e);
                sum_vec = _mm256_add_ps(sum_vec, e);
            }
            float sum = hsum256_ps_softmax(sum_vec);
            for (; j < len; ++j) {
                float e = expf(row[j] - max_val);
                row[j] = e;
                sum += e;
            }
 
            // Normalize (vectorized)
            float inv_sum = 1.0f / sum;
            __m256 inv_sum_vec = _mm256_set1_ps(inv_sum);
            j = 0;
            for (; j + 8 <= len; j += 8) {
                __m256 v = _mm256_loadu_ps(&row[j]);
                _mm256_storeu_ps(&row[j], _mm256_mul_ps(v, inv_sum_vec));
            }
            for (; j < len; ++j) {
                row[j] *= inv_sum;
            }
 
            // Zero out future tokens (vectorized)
            __m256 zero = _mm256_setzero_ps();
            for (; j + 8 <= num_tokens; j += 8) {
                _mm256_storeu_ps(&row[j], zero);
            }
            for (; j < num_tokens; ++j) {
                row[j] = 0.0f;
            }
 
#elif defined(__AVX__)
            // AVX1: vectorized max/sum/normalize, scalar exp
            __m256 max_vec = _mm256_set1_ps(-INFINITY);
            int j = 0;
            for (; j + 8 <= len; j += 8) {
                __m256 v = _mm256_loadu_ps(&row[j]);
                max_vec = _mm256_max_ps(max_vec, v);
            }
            float max_val = hmax256_ps(max_vec);
            for (; j < len; ++j) {
                if (row[j] > max_val) max_val = row[j];
            }
 
            // Compute exp and sum (scalar exp, no fast approx for AVX1)
            float sum = 0.0f;
            for (j = 0; j < len; ++j) {
                float e = expf(row[j] - max_val);
                row[j] = e;
                sum += e;
            }
 
            // Normalize (vectorized)
            float inv_sum = 1.0f / sum;
            __m256 inv_sum_vec = _mm256_set1_ps(inv_sum);
            j = 0;
            for (; j + 8 <= len; j += 8) {
                __m256 v = _mm256_loadu_ps(&row[j]);
                _mm256_storeu_ps(&row[j], _mm256_mul_ps(v, inv_sum_vec));
            }
            for (; j < len; ++j) {
                row[j] *= inv_sum;
            }
 
            // Zero out future tokens (vectorized)
            __m256 zero = _mm256_setzero_ps();
            for (; j + 8 <= num_tokens; j += 8) {
                _mm256_storeu_ps(&row[j], zero);
            }
            for (; j < num_tokens; ++j) {
                row[j] = 0.0f;
            }
 
#else
            // Scalar fallback
            float max_val = row[0];
            for (int j = 1; j < len; ++j) {
                if (row[j] > max_val) max_val = row[j];
            }
 
            float sum = 0.0f;
            for (int j = 0; j < len; ++j) {
                float e = expf(row[j] - max_val);
                row[j] = e;
                sum += e;
            }
 
            float inv_sum = 1.0f / sum;
            for (int j = 0; j < len; ++j) {
                row[j] *= inv_sum;
            }
 
            for (int j = len; j < num_tokens; ++j) {
                row[j] = 0.0f;
            }
#endif
        }
    }
}

Referenced by attention_forward_causal_head_major(), attention_forward_causal_head_major_gqa(), and causal_softmax_head_major_bf16().

◆ causal_softmax_head_major_bf16()

void causal_softmax_head_major_bf16	(	uint16_t *	scores,
		int	num_heads,
		int	num_tokens,
		int	aligned_context_window,
		float *	scratch
	)

Definition at line 31 of file softmax_kernels_bf16.c.

{
    if (!scores || num_heads <= 0 || num_tokens <= 0 || aligned_context_window <= 0) return;
    if (!scratch) return;
 
    const size_t total = (size_t)num_heads *
                         (size_t)aligned_context_window *
                         (size_t)aligned_context_window;
 
    bf16_tensor_to_float(scores, scratch, total);
    causal_softmax_head_major(scratch, num_heads, num_tokens, aligned_context_window);
    float_tensor_to_bf16(scratch, scores, total);
}

References bf16_tensor_to_float(), causal_softmax_head_major(), and float_tensor_to_bf16().

◆ causal_softmax_head_major_exact()

void causal_softmax_head_major_exact	(	float *	scores,
		int	num_heads,
		int	num_tokens,
		int	aligned_context_window
	)

Causal softmax (exact version using stdlib expf)

Test:

test_softmax.py::TestSoftmaxForward::test_causal_softmax_exact

test_softmax.py::TestSoftmaxForward::test_exact_vs_fast

Exact causal softmax using standard library expf for numerical accuracy reference.

After changes: make test

Definition at line 339 of file softmax_kernels.c.

{
    for (int h = 0; h < num_heads; ++h) {
        for (int i = 0; i < num_tokens; ++i) {
            int base = h * aligned_context_window * aligned_context_window
                     + i * aligned_context_window;
            float *row = &scores[base];
            int len = i + 1;
 
            // Find max
            float max_val = -INFINITY;
            for (int j = 0; j < len; ++j) {
                if (row[j] > max_val) max_val = row[j];
            }
 
            // Compute exp and sum using standard library expf
            float sum = 0.0f;
            for (int j = 0; j < len; ++j) {
                float e = expf(row[j] - max_val);
                row[j] = e;
                sum += e;
            }
 
            // Normalize
            float inv_sum = 1.0f / sum;
            for (int j = 0; j < len; ++j) {
                row[j] *= inv_sum;
            }
 
            // Zero out future tokens
            for (int j = len; j < num_tokens; ++j) {
                row[j] = 0.0f;
            }
        }
    }
}

Referenced by attention_forward_causal_head_major_exact(), and attention_forward_causal_head_major_gqa_exact().

◆ ck_attention_bf16_pytorch_gqa_available()

int ck_attention_bf16_pytorch_gqa_available ( void )

Definition at line 8511 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_flash_decode_wrapper()

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

Wrapper to call TRUE flash attention from orchestration layer.

Parameters

q_token	Query token [H, D_h]
k_cache	Cached keys [T_k, H, D_h]
v_cache	Cached values [T_k, H, D_h]
out_token	Output [H, D_h]
num_heads	Number of heads
num_kv_heads	Number of KV heads (for GQA)
kv_tokens	Number of tokens in KV cache
cache_capacity	Cache capacity
head_dim	Head dimension
aligned_head_dim	Aligned head dimension

Definition at line 72 of file ckernel_orchestration.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;
    }
 
    static int use_strict = -1;
    if (use_strict < 0) {
        const char *env = getenv("CK_FLASH_ATTN_STRICT");
        use_strict = (env && env[0] && env[0] != '0') ? 1 : 0;
    }
 
    if (use_strict) {
        attention_forward_decode_head_major_gqa_regular(q_token,
                                                        k_cache,
                                                        v_cache,
                                                        out_token,
                                                        num_heads,
                                                        num_kv_heads,
                                                        kv_tokens,
                                                        cache_capacity,
                                                        head_dim,
                                                        aligned_head_dim);
        return;
    }
 
    // Scale factor: 1/sqrt(head_dim)
    const float scale = 1.0f / sqrtf((float)head_dim);
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
 
#pragma omp parallel for schedule(static) if(num_heads > 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 *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;
 
        // Use aligned_head_dim as D_h so per-token stride matches the cache layout.
        attention_flash_decode(out_head,
                               q_head,
                               k_head,
                               v_head,
                               1,
                               kv_tokens,
                               1,
                               aligned_head_dim,
                               scale);
    }
}

References attention_flash_decode(), and attention_forward_decode_head_major_gqa_regular().

Referenced by ck_layer_forward_rmsnorm_swiglu_decode(), ck_layer_forward_rmsnorm_swiglu_decode_fused(), ck_layer_forward_rmsnorm_swiglu_decode_q4_k(), and ck_layer_forward_rmsnorm_swiglu_decode_quant().

◆ ck_attention_pytorch_sdpa_scale_f32()

float ck_attention_pytorch_sdpa_scale_f32 ( int head_dim )

Definition at line 4518 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_sparse_bf16_pytorch_gqa_available()

int ck_attention_sparse_bf16_pytorch_gqa_available ( void )

Definition at line 8502 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_flash_attn_choose_tile_k()

int ck_flash_attn_choose_tile_k ( int D_h )

Definition at line 108 of file attention_flash_true.c.

                                         {
    return ck_flash_attn_tile_k(D_h);
}

References ck_flash_attn_tile_k().

◆ ck_flash_attn_fast_exp_kind()

int ck_flash_attn_fast_exp_kind ( void )

Definition at line 112 of file attention_flash_true.c.

                                      {
#if CK_FLASH_ATTN_FAST_EXP
#if defined(__AVX512F__)
    return 512;
#elif defined(__AVX__)
    return 256;
#else
    return 0;
#endif
#else
    return 0;
#endif
}

◆ ck_gemm_bf16_amx_available()

int ck_gemm_bf16_amx_available ( void )

Definition at line 1196 of file gemm_kernels_bf16.c.

{
#if HAVE_AMX_BF16
    return ck_amx_request_xtile_data();
#else
    return 0;
#endif
}

Referenced by ck_attention_full_bf16_sdpa_amx_range(), gemm_nt_bf16_prefill_shape_safe_bf16_storage(), and gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace().

◆ ck_gemm_bf16_fp32out_amx_raw()

int ck_gemm_bf16_fp32out_amx_raw	(	const uint16_t *	A,
		const uint16_t *	B,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	accumulate
	)

Definition at line 1205 of file gemm_kernels_bf16.c.

{
#if HAVE_AMX_BF16
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0 ||
        (M % 16) != 0 || (N % 16) != 0 || (K % 2) != 0 ||
        !ck_amx_request_xtile_data()) {
        return 0;
    }
    /* Match oneDNN BRGEMM: largest even divisor of K no greater than 32. */
    int k_block = K < 32 ? K : 32;
    while (k_block > 2 && K % k_block != 0) k_block -= 2;
    ck_amx_config_bf16_16x16_kblock(k_block);
    uint16_t b_tile[16 * 32];
    for (int i = 0; i < M; i += 16) {
        for (int j = 0; j < N; j += 16) {
            if (accumulate) {
                _tile_loadd(2, C + (size_t)i * (size_t)N + (size_t)j,
                            N * (int)sizeof(float));
            } else {
                _tile_zero(2);
            }
            for (int k = 0; k < K; k += k_block) {
                memset(b_tile, 0, sizeof(b_tile));
                for (int kp = 0; kp < k_block / 2; ++kp) {
                    const int k0 = k + kp * 2;
                    for (int nn = 0; nn < 16; ++nn) {
                        b_tile[(size_t)kp * 32u + (size_t)nn * 2u] =
                            B[(size_t)(j + nn) * (size_t)K + (size_t)k0];
                        b_tile[(size_t)kp * 32u + (size_t)nn * 2u + 1u] =
                            B[(size_t)(j + nn) * (size_t)K + (size_t)k0 + 1u];
                    }
                }
                _tile_loadd(0, A + (size_t)i * (size_t)K + (size_t)k,
                            K * (int)sizeof(uint16_t));
                _tile_loadd(1, b_tile, 32 * (int)sizeof(uint16_t));
                _tile_dpbf16ps(2, 0, 1);
            }
            _tile_stored(2, C + (size_t)i * (size_t)N + (size_t)j,
                         N * (int)sizeof(float));
        }
    }
    _tile_release();
    return 1;
#else
    (void)A; (void)B; (void)C; (void)M; (void)N; (void)K; (void)accumulate;
    return 0;
#endif
}

References C.

Referenced by ck_attention_full_bf16_sdpa_amx_range().

◆ ck_gemm_nt_f16_ggml_oracle()

int ck_gemm_nt_f16_ggml_oracle	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 262 of file gemm_kernels_f16.c.

{
    return gemm_nt_f16_ggml_strict(A, B, bias, C, M, N, K);
}

References C, and gemm_nt_f16_ggml_strict().

◆ ck_gemm_nt_f16_simd_lanes()

int ck_gemm_nt_f16_simd_lanes ( void )

Definition at line 583 of file gemm_kernels_f16.c.

{
#if defined(__AVX512F__) && defined(__F16C__)
    return 16;
#elif defined(__F16C__) && defined(__AVX__)
    return 8;
#else
    return 1;
#endif
}

◆ ck_gemm_nt_head_major_q5_0()

void ck_gemm_nt_head_major_q5_0	(	const float *	attn_out,
		const void *	wo,
		const float *	bias,
		float *	output,
		int	tokens,
		int	embed_dim,
		int	num_heads,
		int	head_dim
	)

Output projection from head-major attention (auto-dispatch)

This replaces flatten_head_major() + ck_gemm_nt_quant() with a single strided-access kernel that reads head-major attention output directly.

Definition at line 290 of file gemm_head_major_output.c.

{
#if defined(__AVX__)
    gemv_nt_q5_0_head_major_output_avx(output, attn_out, wo, bias,
                                       tokens, embed_dim, num_heads, head_dim);
#else
    gemv_nt_q5_0_head_major_output(output, attn_out, wo, bias,
                                   tokens, embed_dim, num_heads, head_dim);
#endif
}

References gemv_nt_q5_0_head_major_output().

Referenced by mega_fused_attention_prefill().

◆ ck_gemm_nt_head_major_q8_0()

void ck_gemm_nt_head_major_q8_0	(	const float *	attn_out,
		const void *	wo,
		const float *	bias,
		float *	output,
		int	tokens,
		int	embed_dim,
		int	num_heads,
		int	head_dim
	)

Output projection from head-major attention (Q8_0 weights)

Definition at line 315 of file gemm_head_major_output.c.

{
    if (!output || !attn_out || !wo) return;
    if (tokens <= 0 || embed_dim <= 0 || num_heads <= 0 || head_dim <= 0) return;
 
    const int blocks_per_head = head_dim / QK8_0;
    const int blocks_per_row = embed_dim / QK8_0;
    const block_q8_0 *weights = (const block_q8_0 *)wo;
 
    const size_t token_stride = head_dim;
    const size_t head_stride = (size_t)tokens * token_stride;
 
    /* Initialize output */
    if (bias) {
        for (int t = 0; t < tokens; t++) {
            float *out_row = output + (size_t)t * embed_dim;
            for (int n = 0; n < embed_dim; n++) {
                out_row[n] = bias[n];
            }
        }
    } else {
        memset(output, 0, (size_t)tokens * embed_dim * sizeof(float));
    }
 
    /* Accumulate from each head */
    for (int h = 0; h < num_heads; h++) {
        const float *head_data = attn_out + (size_t)h * head_stride;
        const int head_offset = h * blocks_per_head;
 
        for (int n_block = 0; n_block < blocks_per_head; n_block++) {
            for (int n = 0; n < embed_dim; n++) {
                const block_q8_0 *w_row = weights + (size_t)n * blocks_per_row + head_offset + n_block;
                const float d = CK_FP16_TO_FP32(w_row->d);
 
                for (int t = 0; t < tokens; t++) {
                    const float *token_vec = head_data + (size_t)t * token_stride + (size_t)n_block * QK8_0;
                    float sum = 0.0f;
 
                    for (int j = 0; j < QK8_0; j++) {
                        sum += d * (float)w_row->qs[j] * token_vec[j];
                    }
 
                    output[(size_t)t * embed_dim + n] += sum;
                }
            }
        }
    }
}

References CK_FP16_TO_FP32, block_q8_0::d, QK8_0, and block_q8_0::qs.

Referenced by mega_fused_attention_prefill().

◆ ck_get_num_threads()

int ck_get_num_threads ( void )

Definition at line 310 of file ckernel_strict.c.

{
    // Auto-initialize if not set
    if (!g_threads_initialized) {
        ck_set_num_threads(0);  // Auto-detect
    }
    return g_num_threads;
}

References ck_set_num_threads(), g_num_threads, and g_threads_initialized.

Referenced by attention_forward_decode_head_major_gqa_flash_f16cache_contract(), and gemm_blocked_serial().

◆ ck_get_physical_cores()

int ck_get_physical_cores ( void )

Definition at line 170 of file ckernel_strict.c.

{
    int physical_cores = 0;
    int logical_cores = (int)sysconf(_SC_NPROCESSORS_ONLN);
    if (logical_cores <= 0) {
        logical_cores = 1;
    }
 
    int cpu_cores_hint = 0;
    int siblings_hint = 0;
 
    // Read from /proc/cpuinfo (Linux) and count unique (physical id, core id) pairs.
    FILE *f = fopen("/proc/cpuinfo", "r");
    if (f) {
        char line[256];
        int physical_id = -1;
        int core_id = -1;
 
        struct {
            int physical_id;
            int core_id;
        } seen[8192];
        int seen_count = 0;
 
        const int seen_cap = (int)(sizeof(seen) / sizeof(seen[0]));
 
        // Helper: add (pid,cid) to set if not present.
        #define CK_ADD_PAIR(pid, cid)                                            \
            do {                                                                 \
                if ((pid) >= 0 && (cid) >= 0) {                                  \
                    int exists = 0;                                              \
                    for (int ii = 0; ii < seen_count; ++ii) {                    \
                        if (seen[ii].physical_id == (pid) &&                     \
                            seen[ii].core_id == (cid)) {                         \
                            exists = 1;                                          \
                            break;                                               \
                        }                                                        \
                    }                                                            \
                    if (!exists && seen_count < seen_cap) {                      \
                        seen[seen_count].physical_id = (pid);                    \
                        seen[seen_count].core_id = (cid);                        \
                        ++seen_count;                                            \
                    }                                                            \
                }                                                                \
            } while (0)
 
        while (fgets(line, sizeof(line), f)) {
            int val;
 
            // Blank line separates processor blocks.
            if (line[0] == '\n' || line[0] == '\0') {
                CK_ADD_PAIR(physical_id, core_id);
                physical_id = -1;
                core_id = -1;
                continue;
            }
 
            if (sscanf(line, "physical id : %d", &val) == 1) {
                physical_id = val;
                continue;
            }
            if (sscanf(line, "core id : %d", &val) == 1) {
                core_id = val;
                continue;
            }
            if (sscanf(line, "cpu cores : %d", &val) == 1) {
                if (val > cpu_cores_hint) cpu_cores_hint = val;
                continue;
            }
            if (sscanf(line, "siblings : %d", &val) == 1) {
                if (val > siblings_hint) siblings_hint = val;
                continue;
            }
        }
        fclose(f);
 
        // Handle file without trailing blank line.
        CK_ADD_PAIR(physical_id, core_id);
 
        #undef CK_ADD_PAIR
 
        physical_cores = seen_count;
    }
 
    // Fallback: infer threads-per-core from siblings/cpu cores when pair data
    // is missing (common in containers/VMs).
    if (physical_cores <= 1 && logical_cores > 1) {
        int threads_per_core = 0;
        if (siblings_hint > 0 && cpu_cores_hint > 0 && siblings_hint >= cpu_cores_hint) {
            threads_per_core = siblings_hint / cpu_cores_hint;
        }
        if (threads_per_core > 1) {
            int inferred_physical = logical_cores / threads_per_core;
            if (inferred_physical > 1) {
                return inferred_physical;
            }
        }
        if (cpu_cores_hint > 1 && cpu_cores_hint <= logical_cores) {
            return cpu_cores_hint;
        }
        return logical_cores;
    }
 
    if (physical_cores > 1) {
        return physical_cores;
    }
 
    return logical_cores;
}

References CK_ADD_PAIR.

◆ ck_layout_head_to_token_f32()

void ck_layout_head_to_token_f32	(	const float *	src,
		float *	dst,
		int	heads,
		int	tokens,
		int	head_dim
	)

Definition at line 19 of file layout_kernels.c.

{
    if (!src || !dst || heads <= 0 || tokens <= 0 || head_dim <= 0) return;
    for (int h = 0; h < heads; ++h) {
        for (int t = 0; t < tokens; ++t) {
            memcpy(dst + ((size_t)t * (size_t)heads + (size_t)h) * (size_t)head_dim,
                   src + ((size_t)h * (size_t)tokens + (size_t)t) * (size_t)head_dim,
                   (size_t)head_dim * sizeof(float));
        }
    }
}

◆ ck_layout_token_to_head_f32()

void ck_layout_token_to_head_f32	(	const float *	src,
		float *	dst,
		int	tokens,
		int	heads,
		int	head_dim
	)

Definition at line 6 of file layout_kernels.c.

{
    if (!src || !dst || tokens <= 0 || heads <= 0 || head_dim <= 0) return;
    for (int t = 0; t < tokens; ++t) {
        for (int h = 0; h < heads; ++h) {
            memcpy(dst + ((size_t)h * (size_t)tokens + (size_t)t) * (size_t)head_dim,
                   src + ((size_t)t * (size_t)heads + (size_t)h) * (size_t)head_dim,
                   (size_t)head_dim * sizeof(float));
        }
    }
}

◆ ck_memcpy_parallel_dispatch()

void * ck_memcpy_parallel_dispatch	(	void *	dst,
		const void *	src,
		size_t	size
	)

◆ ck_multimodal_mrope_positions_2d()

int ck_multimodal_mrope_positions_2d	(	int32_t *	positions,
		int	total_tokens,
		int	prefix_start,
		int	position_base,
		int	prefix_tokens,
		int	grid_x,
		int	grid_y,
		int	text_pos
	)

Definition at line 2892 of file rope_kernels.c.

{
    if (!positions || total_tokens <= 0 || prefix_tokens <= 0) {
        return -1;
    }
    if (grid_x <= 0 || grid_y <= 0 || grid_x * grid_y != prefix_tokens) {
        return -2;
    }
    if (prefix_start < 0 || prefix_start > total_tokens) {
        return -3;
    }
    if (prefix_tokens > total_tokens - prefix_start) {
        return -4;
    }
 
    const int prefix_end = prefix_start + prefix_tokens;
    const int grid_extent = grid_x > grid_y ? grid_x : grid_y;
    const int resolved_text_pos = text_pos > 0
        ? text_pos
        : prefix_start + grid_extent;
 
    for (int token = 0; token < total_tokens; ++token) {
        int32_t pos0;
        int32_t pos1;
        int32_t pos2;
        if (token < prefix_start) {
            pos0 = pos1 = pos2 = (int32_t)token;
        } else if (token < prefix_end) {
            const int local_token = token - prefix_start;
            pos0 = (int32_t)position_base;
            pos1 = (int32_t)(position_base + local_token / grid_x);
            pos2 = (int32_t)(position_base + local_token % grid_x);
        } else {
            pos0 = pos1 = pos2 = (int32_t)(resolved_text_pos + token - prefix_end);
        }
        positions[token] = pos0;
        positions[token + total_tokens] = pos1;
        positions[token + 2 * total_tokens] = pos2;
        positions[token + 3 * total_tokens] = 0;
    }
    return resolved_text_pos;
}

References token.

◆ ck_multimodal_prefix_insert_f32()

int ck_multimodal_prefix_insert_f32	(	const float *	source_rows,
		int32_t *	token_ids,
		float *	decoder_rows,
		int	row_count,
		int	source_row_stride,
		int	decoder_row_stride,
		int	copy_dim,
		int	start_row,
		int	decoder_capacity
	)

Definition at line 342 of file embedding_kernels.c.

{
    if (!source_rows || !token_ids || !decoder_rows || row_count <= 0) {
        return -1;
    }
    if (source_row_stride < copy_dim || decoder_row_stride < copy_dim || copy_dim <= 0) {
        return -2;
    }
    if (start_row < 0 || start_row >= decoder_capacity) {
        return -3;
    }
    if (row_count > decoder_capacity - start_row) {
        row_count = decoder_capacity - start_row;
    }
 
    for (int row = 0; row < row_count; ++row) {
        const float *src = source_rows + (size_t)row * (size_t)source_row_stride;
        float *dst = decoder_rows
                   + (size_t)(start_row + row) * (size_t)decoder_row_stride;
        memcpy(dst, src, (size_t)copy_dim * sizeof(float));
        token_ids[start_row + row] = 0;
    }
    return row_count;
}

◆ ck_q6_k_prepared_provider_name()

const char * ck_q6_k_prepared_provider_name ( void )

Definition at line 62 of file gemm_kernels_q6k_q8k.c.

{
#if defined(__AVX512F__) && defined(__AVX512BW__) && \
    defined(__AVX512VNNI__)
    return "q6_k_prepared_avx512_vnni_exact";
#elif defined(__AVX2__)
    return "q6_k_prepared_avx2_exact";
#else
    return "q6_k_prepared_unavailable";
#endif
}

◆ ck_q6_k_q8_k_provider_name()

const char * ck_q6_k_q8_k_provider_name ( void )

Definition at line 1367 of file gemm_kernels_q6k_q8k.c.

{
    if (ck_strict_parity_enabled() || ck_q6k_q8k_force_ref()) {
        return "q6_k_q8_k_ref";
    }
#if defined(__AVX2__)
    return "q6_k_q8_k_avx2";
#elif defined(__AVX__)
    return "q6_k_q8_k_avx";
#elif defined(__SSE4_1__)
    return "q6_k_q8_k_sse";
#else
    return "q6_k_q8_k_ref";
#endif
}

References ck_q6k_q8k_force_ref(), and ck_strict_parity_enabled().

◆ ck_residual_add_backward()

void ck_residual_add_backward	(	const float *	d_out,
		float *	d_a,
		float *	d_b,
		int	tokens,
		int	aligned_embed_dim
	)

Definition at line 151 of file ckernel_orchestration.c.

{
    if (!d_out || !d_a || !d_b) {
        return;
    }
    size_t total = (size_t)tokens * (size_t)aligned_embed_dim;
    for (size_t i = 0; i < total; ++i) {
        float v = d_out[i];
        d_a[i] = v;
        d_b[i] = v;
    }
}

◆ ck_residual_add_token_major()

void ck_residual_add_token_major	(	const float *	a,
		const float *	b,
		float *	out,
		int	tokens,
		int	aligned_embed_dim
	)

Definition at line 139 of file ckernel_orchestration.c.

{
    size_t total = (size_t)tokens * (size_t)aligned_embed_dim;
    for (size_t i = 0; i < total; ++i) {
        out[i] = a[i] + b[i];
    }
}

Referenced by ck_layer_forward_rmsnorm_swiglu_decode_fused_attn_impl(), mega_fused_attention_prefill(), and mega_fused_attention_prefill_q8_0().

◆ ck_residual_add_token_major_bf16_storage()

void ck_residual_add_token_major_bf16_storage	(	const float *	a,
		const float *	b,
		float *	out,
		int	tokens,
		int	aligned_embed_dim
	)

Definition at line 67 of file add_kernels_bf16.c.

{
    const size_t count = (size_t)tokens * (size_t)aligned_embed_dim;
    for (size_t 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]));
        out[i] = bf16_to_float(float_to_bf16(av + bv));
    }
}

References bf16_to_float(), and float_to_bf16().

◆ ck_residual_add_token_major_parallel_dispatch()

void ck_residual_add_token_major_parallel_dispatch	(	const float *	a,
		const float *	b,
		float *	out,
		int	tokens,
		int	aligned_embed_dim
	)

◆ ck_set_num_threads()

void ck_set_num_threads ( int num_threads )

Definition at line 280 of file ckernel_strict.c.

{
    // 0 = auto-detect
    if (num_threads <= 0) {
        // Prefer explicit env controls when present:
        // - CK_NUM_THREADS: engine-level override
        // - OMP_NUM_THREADS: standard OpenMP control (set by `ck run --threads`)
        int env_threads = ck_parse_env_int("CK_NUM_THREADS");
        if (env_threads <= 0) {
            env_threads = ck_parse_env_int("OMP_NUM_THREADS");
        }
        num_threads = env_threads > 0 ? env_threads : ck_get_physical_cores();
    }
 
    g_num_threads = num_threads;
    g_threads_initialized = 1;
 
#ifdef _OPENMP
    omp_set_dynamic(0);  // Disable dynamic adjustment
    omp_set_num_threads(num_threads);
#endif
 
#if defined(USE_MKL)
    mkl_set_num_threads(num_threads);
#endif
 
    fprintf(stderr, "[CK] Set %d threads (auto=%d)\n",
            num_threads, ck_get_physical_cores());
}

References ck_get_physical_cores(), ck_parse_env_int(), g_num_threads, and g_threads_initialized.

Referenced by ck_get_num_threads().

◆ ck_set_strict_parity()

void ck_set_strict_parity ( int enabled )

Definition at line 28 of file ckernel_strict.c.

{
    ck_strict_parity = enabled ? 1 : 0;
    if (!ck_strict_parity) {
        ck_strict_next_gemm_a_valid = 0;
        ck_strict_next_gemm_a_size = 0;
    }
#ifdef _OPENMP
    if (ck_strict_parity) {
        omp_set_dynamic(0);
        omp_set_num_threads(1);
    }
#endif
}

References ck_strict_next_gemm_a_size, ck_strict_next_gemm_a_valid, and ck_strict_parity.

◆ ck_strict_consume_next_gemm_a()

const float * ck_strict_consume_next_gemm_a ( size_t elems )

Definition at line 70 of file ckernel_strict.c.

{
    if (!ck_strict_parity || !ck_strict_next_gemm_a_valid || ck_strict_next_gemm_a_size != elems) {
        return NULL;
    }
    ck_strict_next_gemm_a_valid = 0;
    return ck_strict_next_gemm_a;
}

References ck_strict_next_gemm_a, ck_strict_next_gemm_a_size, ck_strict_next_gemm_a_valid, and ck_strict_parity.

Referenced by gemm_nt_q8_0_q8_0_contract().

◆ ck_strict_mtmd_clip_encode_planar_f32()

int ck_strict_mtmd_clip_encode_planar_f32	(	const float *	planar,
		int	channels,
		int	height,
		int	width,
		float *	out,
		size_t	out_elems
	)

Definition at line 84 of file ckernel_strict.c.

{
    const char *gguf_path = getenv("CK_STRICT_GGUF_PATH");
    const char *shim_path = getenv("CK_STRICT_MTMD_SHIM_SO");
    if (!gguf_path || !gguf_path[0] || !shim_path || !shim_path[0]) {
        return 0;
    }
    if (!planar || !out || channels != 3 || height <= 0 || width <= 0) {
        return 0;
    }
 
    void *shim = dlopen(shim_path, RTLD_LAZY | RTLD_LOCAL);
    if (!shim) {
        return 0;
    }
 
    ck_strict_mtmd_clip_init_fn init_fn =
        (ck_strict_mtmd_clip_init_fn) dlsym(shim, "ck_mtmd_clip_init");
    ck_strict_mtmd_clip_free_fn free_fn =
        (ck_strict_mtmd_clip_free_fn) dlsym(shim, "ck_mtmd_clip_free");
    ck_strict_mtmd_clip_embd_nbytes_by_img_fn embd_nbytes_fn =
        (ck_strict_mtmd_clip_embd_nbytes_by_img_fn) dlsym(shim, "ck_mtmd_clip_embd_nbytes_by_img");
    ck_strict_mtmd_clip_encode_float_image_fn encode_fn =
        (ck_strict_mtmd_clip_encode_float_image_fn) dlsym(shim, "ck_mtmd_clip_encode_float_image");
 
    if (!init_fn || !free_fn || !embd_nbytes_fn || !encode_fn) {
        dlclose(shim);
        return 0;
    }
 
    const size_t pixel_count = (size_t) height * (size_t) width;
    float *interleaved = (float *) malloc(pixel_count * (size_t) channels * sizeof(float));
    if (!interleaved) {
        dlclose(shim);
        return 0;
    }
    for (size_t idx = 0; idx < pixel_count; ++idx) {
        interleaved[idx * 3 + 0] = planar[idx];
        interleaved[idx * 3 + 1] = planar[pixel_count + idx];
        interleaved[idx * 3 + 2] = planar[2 * pixel_count + idx];
    }
 
    int ok = 0;
    void *handle = init_fn(gguf_path, 0, 0, 0, 0, 0);
    if (handle) {
        const size_t needed_bytes = embd_nbytes_fn(handle, width, height);
        if (needed_bytes > 0 && needed_bytes <= out_elems * sizeof(float)) {
            ok = encode_fn(handle, 1, interleaved, height, width, out) ? 1 : 0;
        }
        free_fn(handle);
    }
 
    free(interleaved);
    dlclose(shim);
    return ok;
}

◆ ck_strict_parity_enabled()

int ck_strict_parity_enabled ( void )

Definition at line 43 of file ckernel_strict.c.

{
    return ck_strict_parity;
}

References ck_strict_parity.

Referenced by adamw_update_f32_impl(), 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_mixed_visual_chunk_head_major_gqa_flash_strided_gemma4_impl(), ck_dot_q6_k_q8_k_fast_or_ref(), ck_layer_forward_rmsnorm_swiglu(), ck_layer_forward_rmsnorm_swiglu_ref(), ck_mlp_swiglu_forward(), ck_mlp_swiglu_forward_ref(), ck_q6_k_q8_k_provider_name(), ck_q8k_activations_enabled(), gated_deltanet_autoregressive_forward(), gated_deltanet_impl_name(), geglu_forward_fp32(), gemm_avx512_parallel(), gemm_blocked_serial(), gemm_fine_grained_parallel(), gemm_naive_parallel(), gemm_nn_avx512(), gemm_nn_avx512_probe(), gemm_nn_blocked(), gemm_nn_parallel(), gemm_nt_f16(), gemm_nt_fp32_exact_parallel_dispatch(), gemm_nt_q6_k_q8_k_m4_tile(), gemm_nt_q8_0_q8_0_contract(), gemm_tn_avx512(), gemm_tn_blocked(), gemm_tn_parallel(), gemv_q4_k_q8_k_vnni(), gemv_q6_k_q8_k(), gemv_q6_k_q8_k_parallel_simd(), gemv_q8_0_q8_0_contract(), layernorm_forward_rolled_slice(), layernorm_forward_unrolled_slice(), mrope_qk_imrope_positions(), mrope_qk_vision(), rmsnorm_backward(), rmsnorm_forward_strided_f32(), softmax_cross_entropy_loss_index_mean_impl(), softmax_cross_entropy_loss_legacy_mean_impl(), swiglu_backward(), swiglu_forward(), and swiglu_forward_q8_k().

◆ ck_strict_store_next_gemm_a()

void ck_strict_store_next_gemm_a	(	const float *	data,
		size_t	elems
	)

Definition at line 48 of file ckernel_strict.c.

{
    if (!ck_strict_parity || !data || elems == 0) {
        ck_strict_next_gemm_a_valid = 0;
        ck_strict_next_gemm_a_size = 0;
        return;
    }
    if (elems > ck_strict_next_gemm_a_cap) {
        float *next = (float *) realloc(ck_strict_next_gemm_a, elems * sizeof(float));
        if (!next) {
            ck_strict_next_gemm_a_valid = 0;
            ck_strict_next_gemm_a_size = 0;
            return;
        }
        ck_strict_next_gemm_a = next;
        ck_strict_next_gemm_a_cap = elems;
    }
    memcpy(ck_strict_next_gemm_a, data, elems * sizeof(float));
    ck_strict_next_gemm_a_size = elems;
    ck_strict_next_gemm_a_valid = 1;
}

References ck_strict_next_gemm_a, ck_strict_next_gemm_a_cap, ck_strict_next_gemm_a_size, ck_strict_next_gemm_a_valid, and ck_strict_parity.

Referenced by ck_attention_full_ggml_graph_oracle_multihead().

◆ ckernel_backend_native()

CKMathBackend ckernel_backend_native ( void )

Obtain the built-in native backend (single-node CPU, C + intrinsics).

Definition at line 39 of file backend_native.c.

{
    CKMathBackend b;
    b.sgemm = &ckernel_sgemm_native;
    return b;
}

References ckernel_sgemm_native(), and CKMathBackend::sgemm.

◆ deepseek_csa_attention_backward_f32()

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

Definition at line 621 of file deepseek_kernels.c.

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

References ds_qkv_idx().

◆ deepseek_csa_attention_f32()

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

Definition at line 562 of file deepseek_kernels.c.

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

References ds_qkv_idx(), and ds_softmax().

Referenced by deepseek_hybrid_attention_f32().

◆ deepseek_dsa_topk_softmax_backward_f32()

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

Definition at line 153 of file deepseek_kernels.c.

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

References topk_softmax_backward_f32().

◆ deepseek_dsa_topk_softmax_f32()

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

Definition at line 86 of file deepseek_kernels.c.

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

◆ deepseek_hybrid_attention_f32()

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

Definition at line 683 of file deepseek_kernels.c.

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

References deepseek_csa_attention_f32().

◆ deepseek_mhc_mix_backward_f32()

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

Definition at line 53 of file deepseek_kernels.c.

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

References ds_mhc_idx(), and ds_mix_idx().

◆ deepseek_mhc_mix_f32()

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

Definition at line 30 of file deepseek_kernels.c.

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

References ds_mhc_idx(), and ds_mix_idx().

◆ deepseek_mla_attention_decode_f32()

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

Definition at line 1044 of file deepseek_kernels.c.

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

References deepseek_mla_attention_decode_f32_workspace().

◆ deepseek_mla_attention_decode_f32_workspace()

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

Definition at line 978 of file deepseek_kernels.c.

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

References score.

Referenced by deepseek_mla_attention_decode_f32().

◆ deepseek_mla_attention_f32()

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

Definition at line 879 of file deepseek_kernels.c.

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

References deepseek_mla_attention_f32_workspace().

◆ deepseek_mla_attention_f32_parallel_dispatch()

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

Definition at line 839 of file deepseek_kernels.c.

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

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

◆ deepseek_mla_attention_f32_workspace()

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

Definition at line 768 of file deepseek_kernels.c.

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

References ds_mla_attention_f32_query_range().

Referenced by deepseek_mla_attention_f32(), and deepseek_mla_attention_f32_parallel_dispatch().

◆ deepseek_mla_kv_cache_batch_store_f32()

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

Definition at line 900 of file deepseek_kernels.c.

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

◆ deepseek_mla_kv_cache_store_f32()

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

Definition at line 937 of file deepseek_kernels.c.

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

◆ deepseek_mla_kv_decompress_bf16()

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

Definition at line 270 of file deepseek_kernels.c.

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

References deepseek_mla_kv_decompress_bf16_token_range().

Referenced by deepseek_mla_kv_decompress_bf16_parallel_dispatch().

◆ deepseek_mla_kv_decompress_bf16_parallel_dispatch()

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

Definition at line 308 of file deepseek_kernels.c.

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

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

◆ deepseek_mla_kv_decompress_bf16_token_range()

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

Definition at line 225 of file deepseek_kernels.c.

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

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

Referenced by deepseek_mla_kv_decompress_bf16(), and ds_mla_kv_decompress_bf16_rows().

◆ deepseek_mla_kv_decompress_f32()

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

Definition at line 185 of file deepseek_kernels.c.

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

References ds_mla_thd_idx(), and ds_mla_tok_idx().

◆ deepseek_mla_partial_rope_concat_f32()

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

Definition at line 367 of file deepseek_kernels.c.

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

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

◆ deepseek_mla_partial_rope_concat_packed_bf16_storage()

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

Definition at line 473 of file deepseek_kernels.c.

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

References ds_mla_bf16_round(), and ds_mla_thd_idx().

◆ deepseek_mla_partial_rope_concat_packed_f32()

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

Definition at line 406 of file deepseek_kernels.c.

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

References ds_mla_thd_idx().

◆ dequant_q4_0_row()

void dequant_q4_0_row	(	const void *	src,
		float *	dst,
		size_t	n_elements
	)

Dequantize Q4_0 row (multiple blocks)

Parameters

src	Q4_0 data
dst	FP32 output
n_elements	Number of elements to dequantize

Definition at line 63 of file dequant_kernels.c.

{
    const block_q4_0 *blocks = (const block_q4_0 *)src;
    const size_t n_blocks = n_elements / QK4_0;
 
    for (size_t b = 0; b < n_blocks; b++) {
        dequant_q4_0_block(&blocks[b], &dst[b * QK4_0]);
    }
}

References dequant_q4_0_block(), and QK4_0.

◆ dequant_q4_1_row()

void dequant_q4_1_row	(	const void *	src,
		float *	dst,
		size_t	n_elements
	)

Dequantize Q4_1 row (multiple blocks)

Definition at line 141 of file dequant_kernels.c.

{
    const block_q4_1 *blocks = (const block_q4_1 *)src;
    const size_t n_blocks = n_elements / QK4_1;
 
    for (size_t b = 0; b < n_blocks; b++) {
        dequant_q4_1_block(&blocks[b], &dst[b * QK4_1]);
    }
}

References dequant_q4_1_block(), and QK4_1.

Referenced by dequant_row().

◆ dequant_q4_k_row()

void dequant_q4_k_row	(	const void *	src,
		float *	dst,
		size_t	n_elements
	)

Dequantize Q4_K row (multiple blocks)

Definition at line 372 of file dequant_kernels.c.

{
    const block_q4_K *blocks = (const block_q4_K *)src;
    const size_t n_blocks = n_elements / QK_K;
 
    for (size_t b = 0; b < n_blocks; b++) {
        dequant_q4_k_block(&blocks[b], &dst[b * QK_K]);
    }
}

References dequant_q4_k_block(), and QK_K.

Referenced by embedding_forward_q4_k().

◆ dequant_q5_0_row()

void dequant_q5_0_row	(	const void *	src,
		float *	dst,
		size_t	n_elements
	)

Dequantize Q5_0 row (multiple blocks)

Definition at line 198 of file dequant_kernels.c.

{
    const block_q5_0 *blocks = (const block_q5_0 *)src;
    const size_t n_blocks = n_elements / QK5_0;
 
    for (size_t b = 0; b < n_blocks; b++) {
        dequant_q5_0_block(&blocks[b], &dst[b * QK5_0]);
    }
}

References dequant_q5_0_block(), and QK5_0.

Referenced by embedding_forward_q5_0(), moe_relu2_expert_forward_q5_0_q5_0(), moe_relu2_expert_forward_q5_0_q8_0(), and qwen4_ple_ngram_embed_impl().

◆ dequant_q5_1_row()

void dequant_q5_1_row	(	const void *	src,
		float *	dst,
		size_t	n_elements
	)

Dequantize Q5_1 row (multiple blocks)

Definition at line 257 of file dequant_kernels.c.

{
    const block_q5_1 *blocks = (const block_q5_1 *)src;
    const size_t n_blocks = n_elements / QK5_1;
 
    for (size_t b = 0; b < n_blocks; b++) {
        dequant_q5_1_block(&blocks[b], &dst[b * QK5_1]);
    }
}

References dequant_q5_1_block(), and QK5_1.

Referenced by dequant_row(), and moe_relu2_shared_forward_q5_1_q8_0().

◆ dequant_q6_k_row()

void dequant_q6_k_row	(	const void *	src,
		float *	dst,
		size_t	n_elements
	)

Dequantize Q6_K row (multiple blocks)

Definition at line 422 of file dequant_kernels.c.

{
    const block_q6_K *blocks = (const block_q6_K *)src;
    const size_t n_blocks = n_elements / QK_K;
 
    for (size_t b = 0; b < n_blocks; b++) {
        dequant_q6_k_block(&blocks[b], &dst[b * QK_K]);
    }
}

References dequant_q6_k_block(), and QK_K.

Referenced by embedding_forward_q6_k().

◆ dequant_q8_0_row()

void dequant_q8_0_row	(	const void *	src,
		float *	dst,
		size_t	n_elements
	)

Dequantize Q8_0 row (multiple blocks)

Definition at line 288 of file dequant_kernels.c.

{
    const block_q8_0 *blocks = (const block_q8_0 *)src;
    const size_t n_blocks = n_elements / QK8_0;
 
    for (size_t b = 0; b < n_blocks; b++) {
        dequant_q8_0_block(&blocks[b], &dst[b * QK8_0]);
    }
}

References dequant_q8_0_block(), and QK8_0.

Referenced by dequant_row(), embedding_forward_q8_0(), moe_relu2_expert_forward_q5_0_q8_0(), and moe_relu2_shared_forward_q5_1_q8_0().

◆ embedding_backward()

void embedding_backward	(	const int32_t *	token_ids,
		int	token_count,
		const float *	d_output,
		float *	d_token_embeddings,
		float *	d_pos_embeddings,
		int	vocab_size,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 296 of file embedding_kernels.c.

{
    if (!token_ids || !d_output || !d_token_embeddings) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) {
        tokens = 0;
    }
    if (tokens > context_window) {
        tokens = context_window;
    }
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const float *d_out = d_output + (size_t)t * (size_t)aligned_embed_dim;
        float *d_tok = d_token_embeddings + (size_t)id * (size_t)aligned_embed_dim;
        float *d_pos = d_pos_embeddings ? (d_pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
 
        for (int d = 0; d < embed_dim; ++d) {
            float grad = d_out[d];
            d_tok[d] += grad;
            if (add_pos && d_pos) {
                d_pos[d] += grad;
            }
        }
    }
}

References vocab_size.

◆ embedding_backward_bf16()

void embedding_backward_bf16	(	const int32_t *	token_ids,
		int	token_count,
		const uint16_t *	d_output,
		uint16_t *	d_token_embeddings,
		uint16_t *	d_pos_embeddings,
		int	vocab_size,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 123 of file embedding_kernels_bf16.c.

{
    if (!token_ids || !d_output || !d_token_embeddings) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) tokens = 0;
    if (tokens > context_window) tokens = context_window;
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const uint16_t *d_out = d_output + (size_t)t * (size_t)aligned_embed_dim;
        uint16_t *d_tok = d_token_embeddings + (size_t)id * (size_t)aligned_embed_dim;
        uint16_t *d_pos = d_pos_embeddings ? (d_pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
 
        for (int d = 0; d < embed_dim; ++d) {
            float grad = bf16_to_float(d_out[d]);
 
            float cur_tok = bf16_to_float(d_tok[d]);
            d_tok[d] = float_to_bf16(cur_tok + grad);
 
            if (add_pos && d_pos) {
                float cur_pos = bf16_to_float(d_pos[d]);
                d_pos[d] = float_to_bf16(cur_pos + grad);
            }
        }
    }
}

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

◆ embedding_backward_bf16_mixed()

void embedding_backward_bf16_mixed	(	const int32_t *	token_ids,
		int	token_count,
		const uint16_t *	d_output,
		float *	d_token_embeddings,
		float *	d_pos_embeddings,
		int	vocab_size,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 166 of file embedding_kernels_bf16.c.

{
    if (!token_ids || !d_output || !d_token_embeddings) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) tokens = 0;
    if (tokens > context_window) tokens = context_window;
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const uint16_t *d_out = d_output + (size_t)t * (size_t)aligned_embed_dim;
        float *d_tok = d_token_embeddings + (size_t)id * (size_t)aligned_embed_dim;
        float *d_pos = d_pos_embeddings ? (d_pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
 
        for (int d = 0; d < embed_dim; ++d) {
            const float grad = bf16_to_float(d_out[d]);
            d_tok[d] += grad;
            if (add_pos && d_pos) {
                d_pos[d] += grad;
            }
        }
    }
}

References bf16_to_float(), and vocab_size.

◆ embedding_forward()

void embedding_forward	(	const int32_t *	token_ids,
		int	token_count,
		int	vocab_size,
		const float *	token_embeddings,
		const float *	pos_embeddings,
		float *	output,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 22 of file embedding_kernels.c.

{
    if (!token_ids || !token_embeddings || !output) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) {
        tokens = 0;
    }
    if (tokens > context_window) {
        tokens = context_window;
    }
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const float *tok = token_embeddings + (size_t)id * (size_t)aligned_embed_dim;
        const float *pos = pos_embeddings ? (pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
 
        if (add_pos && pos) {
            for (int d = 0; d < embed_dim; ++d) {
                out[d] = tok[d] + pos[d];
            }
        } else {
            for (int d = 0; d < embed_dim; ++d) {
                out[d] = tok[d];
            }
        }
 
        for (int d = embed_dim; d < aligned_embed_dim; ++d) {
            out[d] = 0.0f;
        }
    }
 
    for (int t = tokens; t < context_window; ++t) {
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
        memset(out, 0, (size_t)aligned_embed_dim * sizeof(float));
    }
}

References vocab_size.

◆ embedding_forward_bf16()

void embedding_forward_bf16	(	const int32_t *	token_ids,
		int	token_count,
		int	vocab_size,
		const uint16_t *	token_embeddings,
		const uint16_t *	pos_embeddings,
		uint16_t *	output,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 21 of file embedding_kernels_bf16.c.

{
    if (!token_ids || !token_embeddings || !output) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) tokens = 0;
    if (tokens > context_window) tokens = context_window;
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const uint16_t *tok = token_embeddings + (size_t)id * (size_t)aligned_embed_dim;
        const uint16_t *pos = pos_embeddings ? (pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
        uint16_t *out = output + (size_t)t * (size_t)aligned_embed_dim;
 
        if (add_pos && pos) {
            for (int d = 0; d < embed_dim; ++d) {
                float v = bf16_to_float(tok[d]) + bf16_to_float(pos[d]);
                out[d] = float_to_bf16(v);
            }
        } else {
            for (int d = 0; d < embed_dim; ++d) {
                out[d] = tok[d];
            }
        }
 
        for (int d = embed_dim; d < aligned_embed_dim; ++d) {
            out[d] = 0;
        }
    }
 
    for (int t = tokens; t < context_window; ++t) {
        uint16_t *out = output + (size_t)t * (size_t)aligned_embed_dim;
        memset(out, 0, (size_t)aligned_embed_dim * sizeof(uint16_t));
    }
}

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

◆ embedding_forward_bf16_fp32()

void embedding_forward_bf16_fp32	(	const int32_t *	token_ids,
		int	token_count,
		int	vocab_size,
		const uint16_t *	token_embeddings,
		const float *	pos_embeddings,
		float *	output,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 73 of file embedding_kernels_bf16.c.

{
    if (!token_ids || !token_embeddings || !output) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) tokens = 0;
    if (tokens > context_window) tokens = context_window;
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const uint16_t *tok = token_embeddings + (size_t)id * (size_t)aligned_embed_dim;
        const float *pos = pos_embeddings ? (pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
 
        if (add_pos && pos) {
            for (int d = 0; d < embed_dim; ++d) {
                out[d] = bf16_to_float(tok[d]) + pos[d];
            }
        } else {
            for (int d = 0; d < embed_dim; ++d) {
                out[d] = bf16_to_float(tok[d]);
            }
        }
 
        for (int d = embed_dim; d < aligned_embed_dim; ++d) {
            out[d] = 0.0f;
        }
    }
 
    for (int t = tokens; t < context_window; ++t) {
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
        memset(out, 0, (size_t)aligned_embed_dim * sizeof(float));
    }
}

References bf16_to_float(), and vocab_size.

◆ embedding_forward_q4_k()

void embedding_forward_q4_k	(	const int32_t *	token_ids,
		int	token_count,
		int	vocab_size,
		const void *	token_embeddings,
		const float *	pos_embeddings,
		float *	output,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 76 of file embedding_kernels.c.

{
    if (!token_ids || !token_embeddings || !output) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) {
        tokens = 0;
    }
    if (tokens > context_window) {
        tokens = context_window;
    }
 
    const size_t row_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, (size_t)aligned_embed_dim);
    const uint8_t *base = (const uint8_t *)token_embeddings;
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const void *tok = base + (size_t)id * row_bytes;
        const float *pos = pos_embeddings ? (pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
 
        dequant_q4_k_row(tok, out, (size_t)aligned_embed_dim);
 
        if (add_pos && pos) {
            for (int d = 0; d < embed_dim; ++d) {
                out[d] += pos[d];
            }
        }
 
        for (int d = embed_dim; d < aligned_embed_dim; ++d) {
            out[d] = 0.0f;
        }
    }
 
    for (int t = tokens; t < context_window; ++t) {
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
        memset(out, 0, (size_t)aligned_embed_dim * sizeof(float));
    }
}

References CK_DT_Q4_K, ck_dtype_row_bytes(), dequant_q4_k_row(), and vocab_size.

◆ embedding_forward_q5_0()

void embedding_forward_q5_0	(	const int32_t *	token_ids,
		int	token_count,
		int	vocab_size,
		const void *	token_embeddings,
		const float *	pos_embeddings,
		float *	output,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 131 of file embedding_kernels.c.

{
    if (!token_ids || !token_embeddings || !output) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) {
        tokens = 0;
    }
    if (tokens > context_window) {
        tokens = context_window;
    }
 
    const size_t row_bytes = ck_dtype_row_bytes(CK_DT_Q5_0, (size_t)aligned_embed_dim);
    const uint8_t *base = (const uint8_t *)token_embeddings;
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const void *tok = base + (size_t)id * row_bytes;
        const float *pos = pos_embeddings ? (pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
 
        dequant_q5_0_row(tok, out, (size_t)aligned_embed_dim);
 
        if (add_pos && pos) {
            for (int d = 0; d < embed_dim; ++d) {
                out[d] += pos[d];
            }
        }
 
        for (int d = embed_dim; d < aligned_embed_dim; ++d) {
            out[d] = 0.0f;
        }
    }
 
    for (int t = tokens; t < context_window; ++t) {
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
        memset(out, 0, (size_t)aligned_embed_dim * sizeof(float));
    }
}

References CK_DT_Q5_0, ck_dtype_row_bytes(), dequant_q5_0_row(), and vocab_size.

◆ embedding_forward_q6_k()

void embedding_forward_q6_k	(	const int32_t *	token_ids,
		int	token_count,
		int	vocab_size,
		const void *	token_embeddings,
		const float *	pos_embeddings,
		float *	output,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 241 of file embedding_kernels.c.

{
    if (!token_ids || !token_embeddings || !output) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) {
        tokens = 0;
    }
    if (tokens > context_window) {
        tokens = context_window;
    }
 
    const size_t row_bytes = ck_dtype_row_bytes(CK_DT_Q6_K, (size_t)aligned_embed_dim);
    const uint8_t *base = (const uint8_t *)token_embeddings;
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const void *tok = base + (size_t)id * row_bytes;
        const float *pos = pos_embeddings ? (pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
 
        dequant_q6_k_row(tok, out, (size_t)aligned_embed_dim);
 
        if (add_pos && pos) {
            for (int d = 0; d < embed_dim; ++d) {
                out[d] += pos[d];
            }
        }
 
        for (int d = embed_dim; d < aligned_embed_dim; ++d) {
            out[d] = 0.0f;
        }
    }
 
    for (int t = tokens; t < context_window; ++t) {
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
        memset(out, 0, (size_t)aligned_embed_dim * sizeof(float));
    }
}

References CK_DT_Q6_K, ck_dtype_row_bytes(), dequant_q6_k_row(), and vocab_size.

◆ embedding_forward_q8_0()

void embedding_forward_q8_0	(	const int32_t *	token_ids,
		int	token_count,
		int	vocab_size,
		const void *	token_embeddings,
		const float *	pos_embeddings,
		float *	output,
		int	embed_dim,
		int	aligned_embed_dim,
		int	context_window,
		int	add_pos
	)

Definition at line 186 of file embedding_kernels.c.

{
    if (!token_ids || !token_embeddings || !output) {
        return;
    }
 
    int tokens = token_count;
    if (tokens < 0) {
        tokens = 0;
    }
    if (tokens > context_window) {
        tokens = context_window;
    }
 
    const size_t row_bytes = ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)aligned_embed_dim);
    const uint8_t *base = (const uint8_t *)token_embeddings;
 
    for (int t = 0; t < tokens; ++t) {
        int id = token_ids[t];
        if (id < 0 || id >= vocab_size) {
            id = 0;
        }
 
        const void *tok = base + (size_t)id * row_bytes;
        const float *pos = pos_embeddings ? (pos_embeddings + (size_t)t * (size_t)aligned_embed_dim) : NULL;
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
 
        dequant_q8_0_row(tok, out, (size_t)aligned_embed_dim);
 
        if (add_pos && pos) {
            for (int d = 0; d < embed_dim; ++d) {
                out[d] += pos[d];
            }
        }
 
        for (int d = embed_dim; d < aligned_embed_dim; ++d) {
            out[d] = 0.0f;
        }
    }
 
    for (int t = tokens; t < context_window; ++t) {
        float *out = output + (size_t)t * (size_t)aligned_embed_dim;
        memset(out, 0, (size_t)aligned_embed_dim * sizeof(float));
    }
}

References CK_DT_Q8_0, ck_dtype_row_bytes(), dequant_q8_0_row(), and vocab_size.

Referenced by qwen2_0_5b_decode_decode_token(), and qwen2_0_5b_decode_forward_prefill_impl().

◆ farskip_swiglu_shared_combine_bf16()

void farskip_swiglu_shared_combine_bf16	(	const float *	hidden,
		const float *	routed,
		const float *	post_attn_residual,
		const uint16_t *	shared_gate,
		const uint16_t *	shared_up,
		const uint16_t *	shared_down,
		float *	main_output,
		float *	routed_free_output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 3381 of file axpy_kernels.c.

{
    if (rows <= 0) return;
    farskip_swiglu_shared_combine_bf16_row_range(
        hidden, routed, post_attn_residual, shared_gate, shared_up, shared_down,
        main_output, routed_free_output, rows, hidden_dim, intermediate_dim,
        0, rows);
}

References farskip_swiglu_shared_combine_bf16_row_range().

◆ farskip_swiglu_shared_combine_bf16_parallel_dispatch()

void farskip_swiglu_shared_combine_bf16_parallel_dispatch	(	const float *	hidden,
		const float *	routed,
		const float *	post_attn_residual,
		const uint16_t *	shared_gate,
		const uint16_t *	shared_up,
		const uint16_t *	shared_down,
		float *	main_output,
		float *	routed_free_output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim
	)

◆ farskip_swiglu_shared_combine_bf16_row_range()

void farskip_swiglu_shared_combine_bf16_row_range	(	const float *	hidden,
		const float *	routed,
		const float *	post_attn_residual,
		const uint16_t *	shared_gate,
		const uint16_t *	shared_up,
		const uint16_t *	shared_down,
		float *	main_output,
		float *	routed_free_output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	row_begin,
		int	row_end
	)

Definition at line 3328 of file axpy_kernels.c.

{
    if (!hidden || !routed || !post_attn_residual || !shared_gate || !shared_up ||
        !shared_down || !main_output || !routed_free_output || rows <= 0 ||
        hidden_dim <= 0 || intermediate_dim <= 0 || row_begin < 0 ||
        row_begin >= row_end || row_end > rows) {
        return;
    }
 
    float act[intermediate_dim];
 
    for (int r = row_begin; r < row_end; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        const float *route = routed + (size_t)r * (size_t)hidden_dim;
        const float *residual = post_attn_residual + (size_t)r * (size_t)hidden_dim;
        float *main = main_output + (size_t)r * (size_t)hidden_dim;
        float *routed_free = routed_free_output + (size_t)r * (size_t)hidden_dim;
 
        for (int i = 0; i < intermediate_dim; ++i) {
            float gv = 0.0f;
            float uv = 0.0f;
            for (int h = 0; h < hidden_dim; ++h) {
                gv += bf16_to_float(shared_gate[(size_t)i * (size_t)hidden_dim + (size_t)h]) * x[h];
                uv += bf16_to_float(shared_up[(size_t)i * (size_t)hidden_dim + (size_t)h]) * x[h];
            }
            act[i] = ck_moe_silu_f32(gv) * uv;
        }
 
        for (int h = 0; h < hidden_dim; ++h) {
            float shared = 0.0f;
            for (int i = 0; i < intermediate_dim; ++i) {
                shared += bf16_to_float(shared_down[(size_t)h * (size_t)intermediate_dim + (size_t)i]) * act[i];
            }
            const float mlp_output = route[h] + shared;
            routed_free[h] = residual[h] + shared;
            main[h] = residual[h] + mlp_output;
        }
    }
}

References bf16_to_float(), ck_moe_silu_f32(), and main().

Referenced by farskip_swiglu_shared_combine_bf16().

◆ fc1_backward_kernel()

void fc1_backward_kernel	(	const float *	d_output,
		const float *	fc1_input,
		const float *	W_fc1,
		float *	d_input,
		float *	d_W_fc1,
		float *	d_b_fc1,
		int	T,
		int	aligned_in,
		int	aligned_out,
		int	num_threads
	)

Definition at line 174 of file mlp_kernels.c.

{
    (void)num_threads;  // Threading handled by GEMM kernels
 
    // 1. d_input[T, in] = d_output[T, out] @ W[out, in]
    // Using gemm_nn: C[M,N] = A[M,K] @ B[K,N]
    // A = d_output [T, out], B = W [out, in], C = d_input [T, in]
    // M = T, N = aligned_in, K = aligned_out
    gemm_nn_simd(d_output, W_fc1, NULL, d_input,
                   T, aligned_in, aligned_out);
 
    // 2. d_W[out, in] = d_output[T, out].T @ fc1_input[T, in]
    // Using gemm_tn: C[M,N] = A[K,M].T @ B[K,N]
    // A = d_output [T, out] (stored as [K=T, M=out]), B = fc1_input [T, in]
    // C = d_W [out, in], M = aligned_out, N = aligned_in, K = T
    gemm_tn_parallel(d_output, fc1_input, NULL, d_W_fc1,
                   aligned_out, aligned_in, T);
 
    // 3. d_b_fc1 = sum_over_T(d_output)
    for (int out_idx = 0; out_idx < aligned_out; ++out_idx) {
        float bias_grad = 0.0f;
        for (int t = 0; t < T; ++t) {
            bias_grad += d_output[(size_t)t * aligned_out + out_idx];
        }
        d_b_fc1[out_idx] += bias_grad;
    }
}

References gemm_nn_simd(), and gemm_tn_parallel().

Referenced by ck_layer_backward_rmsnorm_swiglu().

◆ fc2_backward_kernel()

void fc2_backward_kernel	(	const float *	d_output,
		const float *	fc2_input,
		const float *	W_fc2,
		float *	d_input,
		float *	d_W_fc2,
		float *	d_b_fc2,
		int	T,
		int	aligned_in,
		int	aligned_out,
		int	num_threads
	)

Definition at line 126 of file mlp_kernels.c.

{
    (void)num_threads;  // Threading handled by GEMM kernels
 
    // 1. d_input[T, in] = d_output[T, out] @ W[out, in]
    // Using gemm_nn: C[M,N] = A[M,K] @ B[K,N]
    // A = d_output [T, out], B = W [out, in], C = d_input [T, in]
    // M = T, N = aligned_in, K = aligned_out
    gemm_nn_simd(d_output, W_fc2, NULL, d_input,
                   T, aligned_in, aligned_out);
 
    // 2. d_W[out, in] = d_output[T, out].T @ fc2_input[T, in]
    // Using gemm_tn: C[M,N] = A[K,M].T @ B[K,N]
    // A = d_output [T, out] (stored as [K=T, M=out]), B = fc2_input [T, in]
    // C = d_W [out, in], M = aligned_out, N = aligned_in, K = T
    // Note: gemm_tn overwrites, so we need to save and add if accumulating
    // For now, assume d_W starts zeroed (gradient accumulation handled at higher level)
    gemm_tn_parallel(d_output, fc2_input, NULL, d_W_fc2,
                   aligned_out, aligned_in, T);
 
    // 3. d_b_fc2 = sum_over_T(d_output)
    for (int out_idx = 0; out_idx < aligned_out; ++out_idx) {
        float bias_grad = 0.0f;
        for (int t = 0; t < T; ++t) {
            bias_grad += d_output[(size_t)t * aligned_out + out_idx];
        }
        d_b_fc2[out_idx] += bias_grad;
    }
}

References gemm_nn_simd(), and gemm_tn_parallel().

Referenced by ck_attention_project_head_major_backward(), ck_layer_backward_rmsnorm_swiglu(), and ck_qkv_project_head_major_backward().

◆ feature_concat()

void feature_concat	(	const float *	main_input,
		const float *	branch_input,
		float *	output,
		int	rows,
		int	main_dim,
		int	branch_slice_dim,
		int	num_branch_slices
	)

Definition at line 804 of file vision_kernels.c.

{
    if (!main_input || !output || rows <= 0 || main_dim < 0 || branch_slice_dim < 0 || num_branch_slices < 0) {
        return;
    }
 
    const int branch_total_dim = branch_slice_dim * num_branch_slices;
    const int out_dim = main_dim + branch_total_dim;
 
    const int in_place_expand = (main_input == output) && (out_dim > main_dim);
    const int row_start = in_place_expand ? rows - 1 : 0;
    const int row_end = in_place_expand ? -1 : rows;
    const int row_step = in_place_expand ? -1 : 1;
 
    for (int row = row_start; row != row_end; row += row_step) {
        const float *src_main = main_input + (size_t) row * (size_t) main_dim;
        float *dst_row = output + (size_t) row * (size_t) out_dim;
 
        if (main_dim > 0) {
            memmove(dst_row, src_main, (size_t) main_dim * sizeof(float));
        }
 
        for (int slice = 0; slice < num_branch_slices; ++slice) {
            const float *src_branch = branch_input
                + (size_t) slice * (size_t) rows * (size_t) branch_slice_dim
                + (size_t) row * (size_t) branch_slice_dim;
            float *dst_branch = dst_row + (size_t) main_dim + (size_t) slice * (size_t) branch_slice_dim;
            if (branch_slice_dim > 0) {
                memcpy(dst_branch, src_branch, (size_t) branch_slice_dim * sizeof(float));
            }
        }
    }
}

◆ feature_concat_2way()

void feature_concat_2way	(	const float *	main_input,
		const float *	branch_input,
		float *	output,
		int	rows,
		int	main_dim,
		int	branch_slice_dim,
		int	num_branch_slices
	)

Definition at line 756 of file vision_kernels.c.

{
    if (!main_input || !branch_input || !output || rows <= 0 || main_dim <= 0 ||
        branch_slice_dim < 0 || num_branch_slices < 0) {
        return;
    }
 
    const int branch_dim = branch_slice_dim * num_branch_slices;
    const size_t main_bytes = (size_t) main_dim * sizeof(float);
    const size_t branch_bytes = (size_t) branch_dim * sizeof(float);
    const size_t out_stride = (size_t) (main_dim + branch_dim);
 
    for (int r = 0; r < rows; ++r) {
        const float *src_main = main_input + ((size_t) r * (size_t) main_dim);
        const float *src_branch = branch_input + ((size_t) r * (size_t) branch_dim);
        float *dst = output + ((size_t) r * out_stride);
        memcpy(dst, src_main, main_bytes);
        memcpy(dst + main_dim, src_branch, branch_bytes);
    }
}

◆ feature_slice_copy()

void feature_slice_copy	(	const float *	src,
		float *	dst,
		int	rows,
		int	src_dim,
		int	dst_dim,
		int	dst_feature_offset
	)

Definition at line 783 of file vision_kernels.c.

{
    if (!src || !dst || rows <= 0 || src_dim <= 0 || dst_dim <= 0 || dst_feature_offset < 0) {
        return;
    }
    if (dst_feature_offset + src_dim > dst_dim) {
        return;
    }
 
    for (int row = 0; row < rows; ++row) {
        const float *src_row = src + (size_t) row * (size_t) src_dim;
        float *dst_row = dst + (size_t) row * (size_t) dst_dim + (size_t) dst_feature_offset;
        memcpy(dst_row, src_row, (size_t) src_dim * sizeof(float));
    }
}

◆ final_logit_scale_f32()

void final_logit_scale_f32	(	float *	logits,
		int	tokens,
		int	vocab_size,
		float	scale
	)

Definition at line 7 of file logit_kernels.c.

{
    if (!logits || tokens <= 0 || vocab_size <= 0) {
        return;
    }
 
    const size_t total = (size_t) tokens * (size_t) vocab_size;
    for (size_t i = 0; i < total; ++i) {
        logits[i] *= scale;
    }
}

References vocab_size.

◆ fused_mlp_swiglu_decode()

void fused_mlp_swiglu_decode	(	const float *	x,
		const float *	W_gate,
		const float *	W_up,
		const float *	W_down,
		const float *	b_gate,
		const float *	b_up,
		const float *	b_down,
		float *	output,
		int	D,
		int	Hff
	)

Definition at line 154 of file mlp_fused_decode.c.

{
#if defined(__AVX512F__)
    // Initialize output with bias or zero
    if (b_down) {
        memcpy(output, b_down, D * sizeof(float));
    } else {
        memset(output, 0, D * sizeof(float));
    }
 
    // Process intermediate dimension in tiles
    // Each tile computes MLP_TILE_SIZE swiglu values and immediately
    // accumulates them into the output
 
    /* Bounds check for stack allocation */
    if (D > 4096) return;
 
    #pragma omp parallel
    {
        /* Thread-local accumulator on stack (no malloc!) */
        float local_output[4096] __attribute__((aligned(64)));
        memset(local_output, 0, D * sizeof(float));
 
        #pragma omp for schedule(static)
        for (int t = 0; t < Hff; t += MLP_TILE_SIZE) {
            int tile_end = (t + MLP_TILE_SIZE < Hff) ? t + MLP_TILE_SIZE : Hff;
            int tile_size = tile_end - t;
 
            // Compute SwiGLU for this tile (stays in L1 cache)
            float swiglu_tile[MLP_TILE_SIZE] __attribute__((aligned(64)));
 
            for (int j = t; j < tile_end; j++) {
                const float *wg_row = &W_gate[j * D];
                const float *wu_row = &W_up[j * D];
 
                // Compute gate = x @ W_gate[j] using AVX-512
                __m512 gate_acc = _mm512_setzero_ps();
                __m512 up_acc = _mm512_setzero_ps();
 
                int k = 0;
                for (; k <= D - 16; k += 16) {
                    __m512 x_vec = _mm512_loadu_ps(&x[k]);
                    __m512 wg_vec = _mm512_loadu_ps(&wg_row[k]);
                    __m512 wu_vec = _mm512_loadu_ps(&wu_row[k]);
 
                    gate_acc = _mm512_fmadd_ps(x_vec, wg_vec, gate_acc);
                    up_acc = _mm512_fmadd_ps(x_vec, wu_vec, up_acc);
                }
 
                float gate = hsum512_ps(gate_acc);
                float up = hsum512_ps(up_acc);
 
                // Scalar remainder
                for (; k < D; k++) {
                    gate += x[k] * wg_row[k];
                    up += x[k] * wu_row[k];
                }
 
                // Add biases
                if (b_gate) gate += b_gate[j];
                if (b_up) up += b_up[j];
 
                // SwiGLU: SiLU(gate) * up
                swiglu_tile[j - t] = silu_scalar(gate) * up;
            }
 
            // Accumulate into output via W_down
            // output[i] += sum_j(swiglu_tile[j] * W_down[i, t+j])
            for (int i = 0; i < D; i++) {
                const float *wd_row = &W_down[i * Hff + t];
 
                __m512 acc = _mm512_setzero_ps();
                int j = 0;
                for (; j <= tile_size - 16; j += 16) {
                    __m512 sw_vec = _mm512_loadu_ps(&swiglu_tile[j]);
                    __m512 wd_vec = _mm512_loadu_ps(&wd_row[j]);
                    acc = _mm512_fmadd_ps(sw_vec, wd_vec, acc);
                }
 
                float sum = hsum512_ps(acc);
                for (; j < tile_size; j++) {
                    sum += swiglu_tile[j] * wd_row[j];
                }
 
                local_output[i] += sum;
            }
        }
 
        // Reduce thread-local outputs
        #pragma omp critical
        {
            for (int i = 0; i < D; i++) {
                output[i] += local_output[i];
            }
        }
        /* No free - stack buffer auto-deallocates */
    }
 
#else
    // Scalar fallback (same algorithm, no SIMD)
    if (b_down) {
        memcpy(output, b_down, D * sizeof(float));
    } else {
        memset(output, 0, D * sizeof(float));
    }
 
    for (int t = 0; t < Hff; t += MLP_TILE_SIZE) {
        int tile_end = (t + MLP_TILE_SIZE < Hff) ? t + MLP_TILE_SIZE : Hff;
        int tile_size = tile_end - t;
 
        float swiglu_tile[MLP_TILE_SIZE];
 
        for (int j = t; j < tile_end; j++) {
            float gate = 0.0f;
            float up = 0.0f;
 
            for (int k = 0; k < D; k++) {
                gate += x[k] * W_gate[j * D + k];
                up += x[k] * W_up[j * D + k];
            }
 
            if (b_gate) gate += b_gate[j];
            if (b_up) up += b_up[j];
 
            swiglu_tile[j - t] = silu_scalar(gate) * up;
        }
 
        for (int i = 0; i < D; i++) {
            for (int j = 0; j < tile_size; j++) {
                output[i] += swiglu_tile[j] * W_down[i * Hff + t + j];
            }
        }
    }
#endif
}

References __attribute__(), MLP_TILE_SIZE, and silu_scalar().

◆ fused_mlp_swiglu_decode_tiled()

void fused_mlp_swiglu_decode_tiled	(	const float *	x,
		const float *	W_gate,
		const float *	W_up,
		const float *	W_down,
		const float *	b_gate,
		const float *	b_up,
		const float *	b_down,
		float *	output,
		int	D,
		int	Hff
	)

Definition at line 429 of file mlp_fused_decode.c.

{
    // Tile size chosen to fit in L2 with W_down tile
    // Tile of swiglu: 256 floats = 1KB
    // Tile of W_down: 256 * D floats = 256 * 896 * 4 = 896KB
    // Fits in 2MB L2 with room for x and prefetch
    const int TILE = 256;
 
#if defined(__AVX512F__)
    // Initialize output
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < D; i++) {
        output[i] = b_down ? b_down[i] : 0.0f;
    }
 
    // Process tiles of intermediate dimension
    for (int t = 0; t < Hff; t += TILE) {
        int tile_end = (t + TILE < Hff) ? t + TILE : Hff;
        int tile_size = tile_end - t;
 
        // Compute swiglu tile
        float swiglu_tile[256] __attribute__((aligned(64)));
 
        #pragma omp parallel for schedule(static)
        for (int jj = 0; jj < tile_size; jj++) {
            int j = t + jj;
            const float *wg_row = &W_gate[j * D];
            const float *wu_row = &W_up[j * D];
 
            __m512 gate_acc = _mm512_setzero_ps();
            __m512 up_acc = _mm512_setzero_ps();
 
            int k = 0;
            for (; k <= D - 16; k += 16) {
                __m512 x_vec = _mm512_loadu_ps(&x[k]);
                __m512 wg_vec = _mm512_loadu_ps(&wg_row[k]);
                __m512 wu_vec = _mm512_loadu_ps(&wu_row[k]);
 
                gate_acc = _mm512_fmadd_ps(x_vec, wg_vec, gate_acc);
                up_acc = _mm512_fmadd_ps(x_vec, wu_vec, up_acc);
            }
 
            float gate = hsum512_ps(gate_acc);
            float up = hsum512_ps(up_acc);
 
            for (; k < D; k++) {
                gate += x[k] * wg_row[k];
                up += x[k] * wu_row[k];
            }
 
            if (b_gate) gate += b_gate[j];
            if (b_up) up += b_up[j];
 
            swiglu_tile[jj] = silu_scalar(gate) * up;
        }
 
        // Accumulate into output (parallelize over D)
        #pragma omp parallel for schedule(static)
        for (int i = 0; i < D; i++) {
            const float *wd_row = &W_down[i * Hff + t];
 
            __m512 acc = _mm512_setzero_ps();
            int j = 0;
            for (; j <= tile_size - 16; j += 16) {
                __m512 sw_vec = _mm512_loadu_ps(&swiglu_tile[j]);
                __m512 wd_vec = _mm512_loadu_ps(&wd_row[j]);
                acc = _mm512_fmadd_ps(sw_vec, wd_vec, acc);
            }
 
            float sum = hsum512_ps(acc);
            for (; j < tile_size; j++) {
                sum += swiglu_tile[j] * wd_row[j];
            }
 
            // Atomic add (or use thread-local buffers for better perf)
            #pragma omp atomic
            output[i] += sum;
        }
    }
 
#else
    // Scalar fallback
    for (int i = 0; i < D; i++) {
        output[i] = b_down ? b_down[i] : 0.0f;
    }
 
    for (int t = 0; t < Hff; t += TILE) {
        int tile_end = (t + TILE < Hff) ? t + TILE : Hff;
 
        float swiglu_tile[256];
 
        for (int j = t; j < tile_end; j++) {
            float gate = 0.0f, up = 0.0f;
            for (int k = 0; k < D; k++) {
                gate += x[k] * W_gate[j * D + k];
                up += x[k] * W_up[j * D + k];
            }
            if (b_gate) gate += b_gate[j];
            if (b_up) up += b_up[j];
            swiglu_tile[j - t] = silu_scalar(gate) * up;
        }
 
        for (int i = 0; i < D; i++) {
            for (int j = t; j < tile_end; j++) {
                output[i] += swiglu_tile[j - t] * W_down[i * Hff + j];
            }
        }
    }
#endif
}

References __attribute__(), and silu_scalar().

Referenced by fused_mlp_swiglu_decode_v2().

◆ fused_mlp_swiglu_decode_v2()

void fused_mlp_swiglu_decode_v2	(	const float *	x,
		const float *	W_gate,
		const float *	W_up,
		const float *	W_down,
		const float *	b_gate,
		const float *	b_up,
		const float *	b_down,
		float *	output,
		int	D,
		int	Hff
	)

Definition at line 318 of file mlp_fused_decode.c.

{
    // For large Hff, use tiled version to avoid stack overflow
    if (Hff > MAX_SWIGLU_STACK) {
        fused_mlp_swiglu_decode_tiled(x, W_gate, W_up, W_down,
                                      b_gate, b_up, b_down, output, D, Hff);
        return;
    }
 
#if defined(__AVX512F__)
    // Stack-allocated swiglu buffer (max 32KB)
    float swiglu[MAX_SWIGLU_STACK] __attribute__((aligned(64)));
 
    // Phase 1: Compute all swiglu values (parallelize over Hff)
    #pragma omp parallel for schedule(static)
    for (int j = 0; j < Hff; j++) {
        const float *wg_row = &W_gate[j * D];
        const float *wu_row = &W_up[j * D];
 
        __m512 gate_acc = _mm512_setzero_ps();
        __m512 up_acc = _mm512_setzero_ps();
 
        int k = 0;
        for (; k <= D - 16; k += 16) {
            __m512 x_vec = _mm512_loadu_ps(&x[k]);
            __m512 wg_vec = _mm512_loadu_ps(&wg_row[k]);
            __m512 wu_vec = _mm512_loadu_ps(&wu_row[k]);
 
            gate_acc = _mm512_fmadd_ps(x_vec, wg_vec, gate_acc);
            up_acc = _mm512_fmadd_ps(x_vec, wu_vec, up_acc);
        }
 
        float gate = hsum512_ps(gate_acc);
        float up = hsum512_ps(up_acc);
 
        for (; k < D; k++) {
            gate += x[k] * wg_row[k];
            up += x[k] * wu_row[k];
        }
 
        if (b_gate) gate += b_gate[j];
        if (b_up) up += b_up[j];
 
        swiglu[j] = silu_scalar(gate) * up;
    }
 
    // Phase 2: Down projection (parallelize over D)
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < D; i++) {
        const float *wd_row = &W_down[i * Hff];
 
        __m512 acc = _mm512_setzero_ps();
        int j = 0;
        for (; j <= Hff - 16; j += 16) {
            __m512 sw_vec = _mm512_loadu_ps(&swiglu[j]);
            __m512 wd_vec = _mm512_loadu_ps(&wd_row[j]);
            acc = _mm512_fmadd_ps(sw_vec, wd_vec, acc);
        }
 
        float sum = hsum512_ps(acc);
        for (; j < Hff; j++) {
            sum += swiglu[j] * wd_row[j];
        }
 
        output[i] = sum + (b_down ? b_down[i] : 0.0f);
    }
 
#else
    // Scalar fallback with stack buffer
    float swiglu[MAX_SWIGLU_STACK];
 
    for (int j = 0; j < Hff; j++) {
        float gate = 0.0f, up = 0.0f;
        for (int k = 0; k < D; k++) {
            gate += x[k] * W_gate[j * D + k];
            up += x[k] * W_up[j * D + k];
        }
        if (b_gate) gate += b_gate[j];
        if (b_up) up += b_up[j];
        swiglu[j] = silu_scalar(gate) * up;
    }
 
    for (int i = 0; i < D; i++) {
        float sum = 0.0f;
        for (int j = 0; j < Hff; j++) {
            sum += swiglu[j] * W_down[i * Hff + j];
        }
        output[i] = sum + (b_down ? b_down[i] : 0.0f);
    }
#endif
}

References __attribute__(), fused_mlp_swiglu_decode_tiled(), MAX_SWIGLU_STACK, and silu_scalar().

Referenced by ck_mlp_swiglu_forward_fully_fused_token().

◆ fused_mlp_swiglu_prefill()

void fused_mlp_swiglu_prefill	(	const float *	x,
		const float *	W_gate,
		const float *	W_up,
		const float *	W_down,
		float *	output,
		int	seq_len,
		int	hidden,
		int	intermediate,
		float *	scratch
	)

Fused MLP (Gate + Up + SwiGLU + Down) for prefill.

Tiles along token dimension to keep gate/up/hidden in L3 cache.

Parameters

scratch Temporary buffer from fused_mlp_swiglu_scratch_size()

Definition at line 878 of file prefill_fused_gemm.c.

{
    fused_mlp_swiglu_prefill_bias(x, W_gate, W_up, W_down,
                                  NULL, NULL, NULL,
                                  output, seq_len, hidden, intermediate,
                                  scratch);
}

References fused_mlp_swiglu_prefill_bias().

◆ fused_mlp_swiglu_prefill_bias()

void fused_mlp_swiglu_prefill_bias	(	const float *	x,
		const float *	W_gate,
		const float *	W_up,
		const float *	W_down,
		const float *	B_gate,
		const float *	B_up,
		const float *	B_down,
		float *	output,
		int	seq_len,
		int	hidden,
		int	intermediate,
		float *	scratch
	)

Fused MLP (Gate + Up + SwiGLU + Down) for prefill with biases.

Definition at line 746 of file prefill_fused_gemm.c.

{
    /* MLP is more complex because we have:
     * gate = x @ W_gate
     * up = x @ W_up
     * hidden = silu(gate) * up
     * out = hidden @ W_down
     *
     * The intermediate (gate, up, hidden) is large: seq_len × intermediate
     * For Qwen2-0.5B: 1024 × 4864 × 4 = 19.4MB (way bigger than L3!)
     *
     * Strategy: Tile along intermediate dimension for gate/up,
     * then fuse SwiGLU, then tile down projection.
     */
 
    /* scratch layout:
     * [gate_tile: TILE_M × TILE_N_INTER]
     * [up_tile: TILE_M × TILE_N_INTER]
     */
    const int TILE_N_INTER = 512;  /* Intermediate tile size */
    float *gate_tile = scratch;
    float *up_tile = scratch + (size_t)PREFILL_TILE_M * TILE_N_INTER;
    float *hidden_tile = gate_tile;  /* Reuse gate_tile for hidden after SwiGLU */
 
    /* For each chunk of intermediate dimension */
    for (int inter_start = 0; inter_start < intermediate; inter_start += TILE_N_INTER) {
        int tile_inter = (inter_start + TILE_N_INTER <= intermediate)
                             ? TILE_N_INTER : (intermediate - inter_start);
 
        const float *W_gate_tile = W_gate + (size_t)inter_start * hidden;
        const float *W_up_tile = W_up + (size_t)inter_start * hidden;
 
        /* For each chunk of tokens */
        for (int m_start = 0; m_start < seq_len; m_start += PREFILL_TILE_M) {
            int tile_m = (m_start + PREFILL_TILE_M <= seq_len)
                             ? PREFILL_TILE_M : (seq_len - m_start);
 
            const float *x_tile = x + (size_t)m_start * hidden;
 
            /* Compute gate and up projections for this tile */
            gemm_tile_nt_strided(x_tile, W_gate_tile, gate_tile,
                                  tile_m, tile_inter, hidden, tile_inter);
            gemm_tile_nt_strided(x_tile, W_up_tile, up_tile,
                                  tile_m, tile_inter, hidden, tile_inter);
            if (B_gate) {
                add_bias_tile(gate_tile, B_gate + inter_start, tile_m, tile_inter);
            }
            if (B_up) {
                add_bias_tile(up_tile, B_up + inter_start, tile_m, tile_inter);
            }
 
            /* Fused SwiGLU: hidden = silu(gate) * up */
            for (int i = 0; i < tile_m; ++i) {
                float *g = gate_tile + (size_t)i * tile_inter;
                float *u = up_tile + (size_t)i * tile_inter;
                for (int j = 0; j < tile_inter; ++j) {
                    float gv = g[j];
                    float silu = gv / (1.0f + expf(-gv));
                    g[j] = silu * u[j];  /* hidden_tile = gate_tile */
                }
            }
 
            /* Down projection: accumulate into output
             * out[m_start:, :] += hidden_tile @ W_down[inter_start:, :]^T
             */
            float *out_tile = output + (size_t)m_start * hidden;
 
            /* This is trickier - W_down is [hidden × intermediate]
             * We have hidden_tile[tile_m × tile_inter]
             * We want out[tile_m × hidden] += hidden_tile × W_down[:, inter_start:inter_start+tile_inter]^T
             *
             * For proper accumulation, need to handle this carefully.
             * For now, use a simpler approach: accumulate partial results.
             */
            for (int i = 0; i < tile_m; ++i) {
                float *h = hidden_tile + (size_t)i * tile_inter;
                float *o = out_tile + (size_t)i * hidden;
 
                for (int d = 0; d < hidden; ++d) {
                    const float *w_row = W_down + (size_t)d * intermediate + inter_start;
                    float sum = (inter_start == 0)
                        ? (B_down ? B_down[d] : 0.0f)
                        : o[d];
 
#if defined(__AVX512F__)
                    __m512 acc = _mm512_setzero_ps();
                    int j = 0;
                    for (; j + 16 <= tile_inter; j += 16) {
                        __m512 hv = _mm512_loadu_ps(h + j);
                        __m512 wv = _mm512_loadu_ps(w_row + j);
                        acc = _mm512_fmadd_ps(hv, wv, acc);
                    }
                    sum += _mm512_reduce_add_ps(acc);
                    for (; j < tile_inter; ++j) {
                        sum += h[j] * w_row[j];
                    }
#elif defined(__AVX__)
                    __m256 acc = _mm256_setzero_ps();
                    int j = 0;
                    for (; j + 8 <= tile_inter; j += 8) {
                        __m256 hv = _mm256_loadu_ps(h + j);
                        __m256 wv = _mm256_loadu_ps(w_row + j);
                        acc = _mm256_add_ps(acc, _mm256_mul_ps(hv, wv));
                    }
                    sum += hsum256_prefill(acc);
                    for (; j < tile_inter; ++j) {
                        sum += h[j] * w_row[j];
                    }
#else
                    for (int j = 0; j < tile_inter; ++j) {
                        sum += h[j] * w_row[j];
                    }
#endif
                    o[d] = sum;
                }
            }
        }
    }
}

References add_bias_tile(), gemm_tile_nt_strided(), PREFILL_TILE_M, and silu().

Referenced by fused_mlp_swiglu_prefill().

◆ fused_mlp_swiglu_prefill_w1w2_quant()

void fused_mlp_swiglu_prefill_w1w2_quant	(	const float *	x,
		const void *	W1,
		const float *	B1,
		CKDataType	w1_dt,
		const void *	W2,
		const float *	B2,
		CKDataType	w2_dt,
		float *	output,
		int	seq_len,
		int	embed_dim,
		int	aligned_embed_dim,
		int	intermediate_dim,
		int	aligned_intermediate_dim,
		void *	scratch
	)

Quantized fused MLP for prefill (W1=gate+up, W2=down)

W1 uses Q8_0 activations (Q5_0/Q8_0 weights), W2 uses Q8_K activations (Q4_K/Q6_K weights).

Uses Q8_0 activations for W1 (Q5_0/Q8_0 weights) and Q8_K activations for W2 (Q4_K/Q6_K weights).

Definition at line 964 of file prefill_fused_gemm.c.

{
    if (!x || !W1 || !W2 || !output || !scratch) {
        return;
    }
    if (seq_len <= 0 || embed_dim <= 0 || aligned_embed_dim <= 0 ||
        intermediate_dim <= 0 || aligned_intermediate_dim <= 0) {
        return;
    }
    if (aligned_embed_dim < embed_dim || aligned_intermediate_dim < intermediate_dim) {
        return;
    }
    if ((aligned_embed_dim % 32) != 0 || (aligned_intermediate_dim % 256) != 0) {
        return;
    }
    if (!mlp_q8_0_dtype_supported(w1_dt) || !mlp_q8_k_dtype_supported(w2_dt)) {
        return;
    }
 
    const int tile_m_max = PREFILL_TILE_M;
    const int inter = aligned_intermediate_dim;
    const size_t q8_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)aligned_embed_dim);
    const size_t q8k_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)aligned_intermediate_dim);
    const size_t w1_row_bytes = ck_dtype_row_bytes(w1_dt, (size_t)aligned_embed_dim);
 
    uint8_t *scratch_bytes = (uint8_t *)scratch;
    size_t q8_bytes = (size_t)tile_m_max * q8_row_bytes;
    size_t gate_bytes = (size_t)tile_m_max * (size_t)inter * sizeof(float);
    size_t up_bytes = gate_bytes;
    size_t gate_offset = align_up_size(q8_bytes, 64);
    size_t up_offset = gate_offset + align_up_size(gate_bytes, 64);
    size_t q8k_offset = up_offset + align_up_size(up_bytes, 64);
 
    uint8_t *q8_tile = scratch_bytes;
    float *gate_tile = (float *)(scratch_bytes + gate_offset);
    float *up_tile = (float *)(scratch_bytes + up_offset);
    uint8_t *q8k_tile = scratch_bytes + q8k_offset;
 
    const uint8_t *w1_base = (const uint8_t *)W1;
    const uint8_t *w_gate = w1_base;
    const uint8_t *w_up = w1_base + (size_t)inter * w1_row_bytes;
 
    const float *b_gate = B1;
    const float *b_up = B1 ? (B1 + (size_t)inter) : NULL;
 
    for (int m_start = 0; m_start < seq_len; m_start += tile_m_max) {
        int tile_m = (m_start + tile_m_max <= seq_len)
                         ? tile_m_max : (seq_len - m_start);
 
        const float *x_tile = x + (size_t)m_start * (size_t)aligned_embed_dim;
        float *out_tile = output + (size_t)m_start * (size_t)aligned_embed_dim;
 
        for (int t = 0; t < tile_m; ++t) {
            const float *row = x_tile + (size_t)t * (size_t)aligned_embed_dim;
            quantize_row_q8_0(row,
                              q8_tile + (size_t)t * q8_row_bytes,
                              aligned_embed_dim);
        }
 
        gemm_nt_q8_0_mlp_dispatch(q8_tile, w_gate, b_gate, gate_tile,
                                 tile_m, inter, aligned_embed_dim, w1_dt);
        gemm_nt_q8_0_mlp_dispatch(q8_tile, w_up, b_up, up_tile,
                                 tile_m, inter, aligned_embed_dim, w1_dt);
 
        for (int i = 0; i < tile_m; ++i) {
            float *g = gate_tile + (size_t)i * (size_t)inter;
            float *u = up_tile + (size_t)i * (size_t)inter;
            for (int j = 0; j < inter; ++j) {
                g[j] = silu_prefill(g[j]) * u[j];
            }
        }
 
        for (int i = 0; i < tile_m; ++i) {
            const float *row = gate_tile + (size_t)i * (size_t)inter;
            quantize_row_q8_k(row,
                              q8k_tile + (size_t)i * q8k_row_bytes,
                              aligned_intermediate_dim);
        }
 
        gemm_nt_q8_k_mlp_dispatch(q8k_tile, W2, B2, out_tile,
                                  tile_m, aligned_embed_dim, aligned_intermediate_dim, w2_dt);
    }
}

References align_up_size(), CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), gemm_nt_q8_0_mlp_dispatch(), gemm_nt_q8_k_mlp_dispatch(), mlp_q8_0_dtype_supported(), mlp_q8_k_dtype_supported(), PREFILL_TILE_M, quantize_row_q8_0(), quantize_row_q8_k(), and silu_prefill().

Referenced by mega_fused_outproj_mlp_prefill().

◆ fused_mlp_swiglu_prefill_w1w2_quant_scratch_size()

size_t fused_mlp_swiglu_prefill_w1w2_quant_scratch_size	(	int	aligned_embed_dim,
		int	aligned_intermediate_dim
	)

Get scratch buffer size for fused_mlp_swiglu_prefill_w1w2_quant.

Definition at line 1062 of file prefill_fused_gemm.c.

{
    if (aligned_embed_dim <= 0 || aligned_intermediate_dim <= 0) {
        return 0;
    }
    const size_t q8_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)aligned_embed_dim);
    const size_t q8k_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)aligned_intermediate_dim);
    const size_t q8_bytes = (size_t)PREFILL_TILE_M * q8_row_bytes;
    const size_t gate_bytes = (size_t)PREFILL_TILE_M * (size_t)aligned_intermediate_dim * sizeof(float);
    const size_t up_bytes = gate_bytes;
    const size_t q8k_bytes = (size_t)PREFILL_TILE_M * q8k_row_bytes;
 
    return align_up_size(q8_bytes, 64) +
           align_up_size(gate_bytes, 64) +
           align_up_size(up_bytes, 64) +
           align_up_size(q8k_bytes, 64);
}

References align_up_size(), CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), and PREFILL_TILE_M.

Referenced by mega_fused_outproj_mlp_prefill_scratch_size().

◆ fused_mlp_swiglu_scratch_size()

size_t fused_mlp_swiglu_scratch_size ( int intermediate )

Get scratch buffer size for fused_mlp_swiglu_prefill.

Definition at line 898 of file prefill_fused_gemm.c.

                                                       {
    const int TILE_N_INTER = 512;
    /* gate_tile + up_tile */
    return 2 * (size_t)PREFILL_TILE_M * TILE_N_INTER * sizeof(float);
}

References PREFILL_TILE_M.

◆ fused_rmsnorm_qkv_prefill()

void fused_rmsnorm_qkv_prefill	(	const float *	x,
		const float *	gamma,
		const float *	Wq,
		const float *	Wk,
		const float *	Wv,
		float *	Q,
		float *	K,
		float *	V,
		int	seq_len,
		int	hidden,
		int	q_dim,
		int	kv_dim,
		float	eps,
		float *	scratch
	)

Fused RMSNorm + QKV projection for prefill.

Tiles along token dimension to keep intermediate x_norm in L2 cache. Avoids ~7MB DRAM traffic per layer for seq_len=1024, hidden=896.

Parameters

scratch Temporary buffer from fused_rmsnorm_qkv_scratch_size()

Fused RMSNorm + QKV projection for prefill.

KEY INSIGHT: For Qwen2-0.5B, all QKV weights fit in L3: Wq (896×896) + Wk (128×896) + Wv (128×896) = 4.1MB < 6MB L3

So we use M-tiling (tokens) only:

For each token tile: a. Compute RMSNorm ONCE into scratch (x_norm stays in L2) b. Do all three GEMMs (Q, K, V) against cached x_norm c. Weights stay hot in L3 across all token tiles

This avoids both:

Large x_norm intermediate buffer (only TILE_M × hidden in L2)
RMSNorm recomputation (done once per token tile, used 3×)

Definition at line 393 of file prefill_fused_gemm.c.

{
    /* scratch is x_norm tile: [TILE_M × hidden] fits in L2 */
 
    /* Process token tiles - weights stay in L3 across all tiles */
    for (int m_start = 0; m_start < seq_len; m_start += PREFILL_TILE_M) {
        int tile_m = (m_start + PREFILL_TILE_M <= seq_len)
                         ? PREFILL_TILE_M : (seq_len - m_start);
 
        const float *x_tile = x + (size_t)m_start * hidden;
 
        /* Step 1: RMSNorm for this token tile (computed ONCE, used 3×) */
        rmsnorm_tile(x_tile, gamma, scratch, tile_m, hidden, hidden, eps);
 
        /* Step 2: Q projection - x_norm is hot in L2, Wq hot in L3 */
        float *Q_tile = Q + (size_t)m_start * q_dim;
        gemm_tile_nt_strided(scratch, Wq, Q_tile, tile_m, q_dim, hidden, q_dim);
 
        /* Step 3: K projection - x_norm still hot, Wk displaces some Wq */
        float *K_tile = K + (size_t)m_start * kv_dim;
        gemm_tile_nt_strided(scratch, Wk, K_tile, tile_m, kv_dim, hidden, kv_dim);
 
        /* Step 4: V projection - x_norm still hot, Wv displaces Wk */
        float *V_tile = V + (size_t)m_start * kv_dim;
        gemm_tile_nt_strided(scratch, Wv, V_tile, tile_m, kv_dim, hidden, kv_dim);
    }
}

References gemm_tile_nt_strided(), PREFILL_TILE_M, and rmsnorm_tile().

◆ fused_rmsnorm_qkv_prefill_head_major()

void fused_rmsnorm_qkv_prefill_head_major	(	const float *	x,
		const float *	gamma,
		const float *	Wq,
		const float *	Bq,
		const float *	Wk,
		const float *	Bk,
		const float *	Wv,
		const float *	Bv,
		float *	Q,
		float *	K,
		float *	V,
		int	seq_len,
		int	embed_dim,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	aligned_head_dim,
		int	kv_stride_tokens,
		float	eps,
		float *	scratch
	)

Fused RMSNorm + QKV projection for prefill (head-major outputs)

Writes Q as [num_heads, seq_len, aligned_head_dim] and K/V with stride kv_stride_tokens for KV-cache compatibility.

Q is written as [num_heads, seq_len, aligned_head_dim]. K/V are written with kv_stride_tokens for KV-cache compatibility.

Definition at line 441 of file prefill_fused_gemm.c.

{
    if (!x || !gamma || !Wq || !Wk || !Wv || !Q || !K || !V || !scratch) {
        return;
    }
    if (seq_len <= 0 || embed_dim <= 0 || aligned_embed_dim <= 0 ||
        head_dim <= 0 || aligned_head_dim <= 0 ||
        num_heads <= 0 || num_kv_heads <= 0) {
        return;
    }
    if (kv_stride_tokens < seq_len) {
        return;
    }
 
    const size_t q_head_stride = (size_t)seq_len * (size_t)aligned_head_dim;
    const size_t kv_head_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
    const size_t head_w_stride = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
 
    for (int m_start = 0; m_start < seq_len; m_start += PREFILL_TILE_M) {
        int tile_m = (m_start + PREFILL_TILE_M <= seq_len)
                         ? PREFILL_TILE_M : (seq_len - m_start);
 
        const float *x_tile = x + (size_t)m_start * (size_t)aligned_embed_dim;
        rmsnorm_tile(x_tile, gamma, scratch, tile_m, embed_dim, aligned_embed_dim, eps);
 
        for (int h = 0; h < num_heads; ++h) {
            const float *wq_h = Wq + (size_t)h * head_w_stride;
            const float *bq_h = Bq ? (Bq + (size_t)h * (size_t)aligned_head_dim) : NULL;
            float *q_h = Q + (size_t)h * q_head_stride + (size_t)m_start * (size_t)aligned_head_dim;
 
            gemm_tile_nt_strided(scratch, wq_h, q_h,
                                 tile_m, aligned_head_dim, aligned_embed_dim, aligned_head_dim);
            add_bias_tile(q_h, bq_h, tile_m, aligned_head_dim);
        }
 
        for (int h = 0; h < num_kv_heads; ++h) {
            const float *wk_h = Wk + (size_t)h * head_w_stride;
            const float *wv_h = Wv + (size_t)h * head_w_stride;
            const float *bk_h = Bk ? (Bk + (size_t)h * (size_t)aligned_head_dim) : NULL;
            const float *bv_h = Bv ? (Bv + (size_t)h * (size_t)aligned_head_dim) : NULL;
            float *k_h = K + (size_t)h * kv_head_stride + (size_t)m_start * (size_t)aligned_head_dim;
            float *v_h = V + (size_t)h * kv_head_stride + (size_t)m_start * (size_t)aligned_head_dim;
 
            gemm_tile_nt_strided(scratch, wk_h, k_h,
                                 tile_m, aligned_head_dim, aligned_embed_dim, aligned_head_dim);
            add_bias_tile(k_h, bk_h, tile_m, aligned_head_dim);
 
            gemm_tile_nt_strided(scratch, wv_h, v_h,
                                 tile_m, aligned_head_dim, aligned_embed_dim, aligned_head_dim);
            add_bias_tile(v_h, bv_h, tile_m, aligned_head_dim);
        }
    }
}

References add_bias_tile(), gemm_tile_nt_strided(), PREFILL_TILE_M, and rmsnorm_tile().

Referenced by mega_fused_attention_prefill(), and mega_fused_attention_prefill_q8_0().

◆ fused_rmsnorm_qkv_prefill_head_major_quant()

void fused_rmsnorm_qkv_prefill_head_major_quant	(	const float *	x,
		const float *	gamma,
		const void *	Wq,
		const float *	Bq,
		CKDataType	wq_dt,
		const void *	Wk,
		const float *	Bk,
		CKDataType	wk_dt,
		const void *	Wv,
		const float *	Bv,
		CKDataType	wv_dt,
		float *	Q,
		float *	K,
		float *	V,
		int	seq_len,
		int	embed_dim,
		int	aligned_embed_dim,
		int	num_heads,
		int	num_kv_heads,
		int	head_dim,
		int	aligned_head_dim,
		int	kv_stride_tokens,
		float	eps,
		void *	scratch
	)

Fused RMSNorm + QKV projection for prefill (head-major, Q8 activations)

Supports Q5_0 or Q8_0 weights with Q8_0 activations.

Supports Q5_0 or Q8_0 weights with Q8_0 activations. Writes K/V directly into KV cache layout (kv_stride_tokens).

Definition at line 519 of file prefill_fused_gemm.c.

{
    if (!x || !gamma || !Wq || !Wk || !Wv || !Q || !K || !V || !scratch) {
        return;
    }
    if (seq_len <= 0 || embed_dim <= 0 || aligned_embed_dim <= 0 ||
        head_dim <= 0 || aligned_head_dim <= 0 ||
        num_heads <= 0 || num_kv_heads <= 0) {
        return;
    }
    if (aligned_embed_dim % 32 != 0) {
        return;
    }
    if (kv_stride_tokens < seq_len) {
        return;
    }
    /* Determine quantization path: Q8_0 activations for Q5_0/Q8_0 weights,
     * Q8_K activations for Q4_K/Q6_K weights. All QKV weights must use
     * the same quantization family. */
    int use_q8_k_path = qkv_q8_k_dtype_supported(wq_dt);
    int use_q8_0_path = qkv_q8_0_dtype_supported(wq_dt);
 
    if (!use_q8_k_path && !use_q8_0_path) {
        /* Unsupported dtype for wq */
        return;
    }
 
    /* Verify all dtypes are from the same family */
    if (use_q8_k_path) {
        if (!qkv_q8_k_dtype_supported(wk_dt) || !qkv_q8_k_dtype_supported(wv_dt)) {
            return;  /* Mixed Q8_K and Q8_0 paths not supported */
        }
    } else {
        if (!qkv_q8_0_dtype_supported(wk_dt) || !qkv_q8_0_dtype_supported(wv_dt)) {
            return;
        }
    }
 
    const size_t float_bytes = (size_t)PREFILL_TILE_M * (size_t)aligned_embed_dim * sizeof(float);
    /* Q8_K has larger blocks (256) than Q8_0 (32), so use appropriate size */
    const CKDataType act_quant_type = use_q8_k_path ? CK_DT_Q8_K : CK_DT_Q8_0;
    const size_t q8_row_bytes = ck_dtype_row_bytes(act_quant_type, (size_t)aligned_embed_dim);
    const size_t q8_bytes = (size_t)PREFILL_TILE_M * q8_row_bytes;
    const size_t q8_offset = align_up_size(float_bytes, 64);
 
    float *normed = (float *)scratch;
    uint8_t *q8_tile = (uint8_t *)scratch + q8_offset;
    (void)q8_bytes;
 
    const size_t q_head_stride = (size_t)seq_len * (size_t)aligned_head_dim;
    const size_t kv_head_stride = (size_t)kv_stride_tokens * (size_t)aligned_head_dim;
    const size_t head_w_elems = (size_t)aligned_head_dim * (size_t)aligned_embed_dim;
    const size_t wq_head_bytes = ck_dtype_row_bytes(wq_dt, head_w_elems);
    const size_t wk_head_bytes = ck_dtype_row_bytes(wk_dt, head_w_elems);
    const size_t wv_head_bytes = ck_dtype_row_bytes(wv_dt, head_w_elems);
 
    for (int m_start = 0; m_start < seq_len; m_start += PREFILL_TILE_M) {
        int tile_m = (m_start + PREFILL_TILE_M <= seq_len)
                         ? PREFILL_TILE_M : (seq_len - m_start);
 
        const float *x_tile = x + (size_t)m_start * (size_t)aligned_embed_dim;
        rmsnorm_tile(x_tile, gamma, normed, tile_m, embed_dim, aligned_embed_dim, eps);
 
        /* Quantize activations to appropriate format */
        for (int t = 0; t < tile_m; ++t) {
            const float *row = normed + (size_t)t * (size_t)aligned_embed_dim;
            if (use_q8_k_path) {
                quantize_row_q8_k(row,
                                  q8_tile + (size_t)t * q8_row_bytes,
                                  aligned_embed_dim);
            } else {
                quantize_row_q8_0(row,
                                  q8_tile + (size_t)t * q8_row_bytes,
                                  aligned_embed_dim);
            }
        }
 
        for (int h = 0; h < num_heads; ++h) {
            const uint8_t *wq_h = (const uint8_t *)Wq + (size_t)h * wq_head_bytes;
            const float *bq_h = Bq ? (Bq + (size_t)h * (size_t)aligned_head_dim) : NULL;
            float *q_h = Q + (size_t)h * q_head_stride + (size_t)m_start * (size_t)aligned_head_dim;
 
            if (use_q8_k_path) {
                gemm_nt_q8_k_qkv_dispatch(q8_tile, wq_h, bq_h, q_h,
                                          tile_m, aligned_head_dim, aligned_embed_dim, wq_dt);
            } else {
                gemm_nt_q8_0_dispatch(q8_tile, wq_h, bq_h, q_h,
                                      tile_m, aligned_head_dim, aligned_embed_dim, wq_dt);
            }
        }
 
        for (int h = 0; h < num_kv_heads; ++h) {
            const uint8_t *wk_h = (const uint8_t *)Wk + (size_t)h * wk_head_bytes;
            const uint8_t *wv_h = (const uint8_t *)Wv + (size_t)h * wv_head_bytes;
            const float *bk_h = Bk ? (Bk + (size_t)h * (size_t)aligned_head_dim) : NULL;
            const float *bv_h = Bv ? (Bv + (size_t)h * (size_t)aligned_head_dim) : NULL;
            float *k_h = K + (size_t)h * kv_head_stride + (size_t)m_start * (size_t)aligned_head_dim;
            float *v_h = V + (size_t)h * kv_head_stride + (size_t)m_start * (size_t)aligned_head_dim;
 
            if (use_q8_k_path) {
                gemm_nt_q8_k_qkv_dispatch(q8_tile, wk_h, bk_h, k_h,
                                          tile_m, aligned_head_dim, aligned_embed_dim, wk_dt);
                gemm_nt_q8_k_qkv_dispatch(q8_tile, wv_h, bv_h, v_h,
                                          tile_m, aligned_head_dim, aligned_embed_dim, wv_dt);
            } else {
                gemm_nt_q8_0_dispatch(q8_tile, wk_h, bk_h, k_h,
                                      tile_m, aligned_head_dim, aligned_embed_dim, wk_dt);
                gemm_nt_q8_0_dispatch(q8_tile, wv_h, bv_h, v_h,
                                      tile_m, aligned_head_dim, aligned_embed_dim, wv_dt);
            }
        }
    }
}

References align_up_size(), CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), gemm_nt_q8_0_dispatch(), gemm_nt_q8_k_qkv_dispatch(), PREFILL_TILE_M, qkv_q8_0_dtype_supported(), qkv_q8_k_dtype_supported(), quantize_row_q8_0(), quantize_row_q8_k(), and rmsnorm_tile().

Referenced by mega_fused_attention_prefill(), and mega_fused_attention_prefill_q8_0().

◆ fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size()

size_t fused_rmsnorm_qkv_prefill_head_major_quant_scratch_size ( int aligned_embed_dim )

Get scratch buffer size for fused_rmsnorm_qkv_prefill_head_major_quant.

Definition at line 651 of file prefill_fused_gemm.c.

                                                                                      {
    if (aligned_embed_dim <= 0) {
        return 0;
    }
    const size_t float_bytes = (size_t)PREFILL_TILE_M * (size_t)aligned_embed_dim * sizeof(float);
    /* Use max of Q8_0 and Q8_K sizes to support both paths */
    const size_t q8_0_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)aligned_embed_dim);
    const size_t q8_k_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)aligned_embed_dim);
    const size_t q8_row_bytes = (q8_k_row_bytes > q8_0_row_bytes) ? q8_k_row_bytes : q8_0_row_bytes;
    const size_t q8_bytes = (size_t)PREFILL_TILE_M * q8_row_bytes;
    return align_up_size(float_bytes, 64) + q8_bytes;
}

References align_up_size(), CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), and PREFILL_TILE_M.

Referenced by mega_fused_attention_prefill(), mega_fused_attention_prefill_q8_0(), mega_fused_attention_prefill_q8_0_scratch_size(), and mega_fused_attention_prefill_scratch_size().

◆ fused_rmsnorm_qkv_scratch_size()

size_t fused_rmsnorm_qkv_scratch_size ( int hidden )

Get scratch buffer size for fused_rmsnorm_qkv_prefill.

Definition at line 739 of file prefill_fused_gemm.c.

                                                  {
    return (size_t)PREFILL_TILE_M * hidden * sizeof(float);
}

References PREFILL_TILE_M.

◆ gated_deltanet_autoregressive_backward()

void gated_deltanet_autoregressive_backward	(	const float *	d_out,
		const float *	d_state_out,
		const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		const float *	state_out,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		float *	d_g,
		float *	d_beta,
		float *	d_state_in,
		int	num_heads,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1988 of file deltanet_kernels.c.

{
    if (!d_out || !d_state_out || !q || !k || !v || !g || !beta || !state_in || !state_out ||
        !d_q || !d_k || !d_v || !d_g || !d_beta || !d_state_in) {
        return;
    }
    if (num_heads <= 0 || state_dim <= 0 || state_dim > CK_DELTANET_MAX_STACK_DIM) {
        return;
    }
 
    gated_deltanet_autoregressive_backward_ref(
        d_out,
        d_state_out,
        q,
        k,
        v,
        g,
        beta,
        state_in,
        state_out,
        d_q,
        d_k,
        d_v,
        d_g,
        d_beta,
        d_state_in,
        num_heads,
        state_dim,
        norm_eps);
}

References CK_DELTANET_MAX_STACK_DIM, and gated_deltanet_autoregressive_backward_ref().

◆ gated_deltanet_autoregressive_forward()

void gated_deltanet_autoregressive_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	num_heads,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1906 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out) {
        return;
    }
    if (num_heads <= 0 || state_dim <= 0) {
        return;
    }
 
    /*
     * q and k arrive pre-normalized by recurrent_qk_l2_norm, so the
     * ISA-specialized kernels can follow the same contract as the scalar ref.
     */
    if (ck_strict_parity_enabled() || ck_deltanet_force_ref()) {
        gated_deltanet_autoregressive_forward_ref(
            q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
        return;
    }
#if defined(__AVX512F__)
    gated_deltanet_autoregressive_forward_avx512(
        q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
#elif defined(__AVX2__)
    gated_deltanet_autoregressive_forward_avx2(
        q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
#elif defined(__AVX__)
    gated_deltanet_autoregressive_forward_avx(
        q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
#else
    gated_deltanet_autoregressive_forward_ref(
        q, k, v, g, beta, state_in, state_out, out, num_heads, state_dim, norm_eps);
#endif
}

References ck_deltanet_force_ref(), ck_strict_parity_enabled(), and gated_deltanet_autoregressive_forward_ref().

◆ gated_deltanet_llama_avx2_forward()

void gated_deltanet_llama_avx2_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 767 of file deltanet_kernels.c.

{
    gated_deltanet_llama_avx2_grouped_forward_transposed_impl(
        q, k, v, g, beta, state_in, state_out, out,
        num_heads, group_count, state_dim, norm_eps, 0, num_heads);
}

References gated_deltanet_llama_avx2_grouped_forward_transposed_impl().

Referenced by gated_deltanet_llama_avx2_prefill_forward().

◆ gated_deltanet_llama_avx2_prefill_forward()

void gated_deltanet_llama_avx2_prefill_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1131 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        rows <= 0 || num_heads <= 0 || group_count <= 0 ||
        num_heads % group_count != 0 || state_dim <= 0) {
        return;
    }
    const size_t qk_stride = (size_t) group_count * (size_t) state_dim;
    const size_t value_stride = (size_t) num_heads * (size_t) state_dim;
    const size_t gate_stride = (size_t) num_heads;
    for (int row = 0; row < rows; ++row) {
        gated_deltanet_llama_avx2_forward(
            q + (size_t) row * qk_stride,
            k + (size_t) row * qk_stride,
            v + (size_t) row * value_stride,
            g + (size_t) row * gate_stride,
            beta + (size_t) row * gate_stride,
            row == 0 ? state_in : state_out,
            state_out,
            out + (size_t) row * value_stride,
            num_heads, group_count, state_dim, norm_eps);
    }
}

References gated_deltanet_llama_avx2_forward().

Referenced by gated_deltanet_llama_chunk64_prefill_forward().

◆ gated_deltanet_llama_chunk64_head_forward()

void gated_deltanet_llama_chunk64_head_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	group_count,
		int	head,
		int	state_dim
	)

Definition at line 1202 of file deltanet_kernels.c.

{
#if defined(__AVX2__)
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        rows <= 0 || num_heads <= 0 || group_count <= 0 ||
        num_heads % group_count != 0 || head < 0 || head >= num_heads ||
        state_dim <= 0 || state_dim > CK_DELTANET_LLAMA_CHUNK_MAX_DIM) {
        return;
    }
    gated_deltanet_llama_chunk64_head(
        q, k, v, g, beta, state_in, state_out, out,
        rows, num_heads, group_count, head, state_dim);
#else
    (void)q;
    (void)k;
    (void)v;
    (void)g;
    (void)beta;
    (void)state_in;
    (void)state_out;
    (void)out;
    (void)rows;
    (void)num_heads;
    (void)group_count;
    (void)head;
    (void)state_dim;
#endif
}

References CK_DELTANET_LLAMA_CHUNK_MAX_DIM.

◆ gated_deltanet_llama_chunk64_prefill_forward()

void gated_deltanet_llama_chunk64_prefill_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1167 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        rows <= 0 || num_heads <= 0 || group_count <= 0 ||
        num_heads % group_count != 0 || state_dim <= 0) {
        return;
    }
#if defined(__AVX2__)
    (void)norm_eps;
    if (state_dim <= CK_DELTANET_LLAMA_CHUNK_MAX_DIM) {
        for (int head = 0; head < num_heads; ++head) {
            gated_deltanet_llama_chunk64_head(
                q, k, v, g, beta, state_in, state_out, out,
                rows, num_heads, group_count, head, state_dim);
        }
        return;
    }
#endif
    gated_deltanet_llama_avx2_prefill_forward(
        q, k, v, g, beta, state_in, state_out, out,
        rows, num_heads, group_count, state_dim, norm_eps);
}

References CK_DELTANET_LLAMA_CHUNK_MAX_DIM, and gated_deltanet_llama_avx2_prefill_forward().

◆ gated_deltanet_prefill_forward()

void gated_deltanet_prefill_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1949 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out) {
        return;
    }
    if (rows <= 0 || num_heads <= 0 || state_dim <= 0) {
        return;
    }
 
    const size_t vector_stride = (size_t)num_heads * (size_t)state_dim;
    const size_t gate_stride = (size_t)num_heads;
    for (int row = 0; row < rows; ++row) {
        const float *row_state_in = row == 0 ? state_in : state_out;
        gated_deltanet_autoregressive_forward(
            q + (size_t)row * vector_stride,
            k + (size_t)row * vector_stride,
            v + (size_t)row * vector_stride,
            g + (size_t)row * gate_stride,
            beta + (size_t)row * gate_stride,
            row_state_in,
            state_out,
            out + (size_t)row * vector_stride,
            num_heads,
            state_dim,
            norm_eps);
    }
}

References gated_deltanet_autoregressive_forward().

◆ gated_deltanet_pytorch_grouped_bf16_forward()

void gated_deltanet_pytorch_grouped_bf16_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1089 of file deltanet_kernels.c.

{
    gated_deltanet_pytorch_grouped_bf16_forward_impl(
        q, k, v, g, beta, state_in, state_out, out,
        NULL, NULL, NULL,
        num_heads, group_count, state_dim, norm_eps);
}

References gated_deltanet_pytorch_grouped_bf16_forward_impl().

Referenced by gated_deltanet_pytorch_grouped_bf16_prefill_forward().

◆ gated_deltanet_pytorch_grouped_bf16_prefill_forward()

void gated_deltanet_pytorch_grouped_bf16_prefill_forward	(	const float *	q,
		const float *	k,
		const float *	v,
		const float *	g,
		const float *	beta,
		const float *	state_in,
		float *	state_out,
		float *	out,
		int	rows,
		int	num_heads,
		int	group_count,
		int	state_dim,
		float	norm_eps
	)

Definition at line 1243 of file deltanet_kernels.c.

{
    if (!q || !k || !v || !g || !beta || !state_in || !state_out || !out ||
        rows <= 0 || num_heads <= 0 || group_count <= 0 ||
        num_heads % group_count != 0 || state_dim <= 0) {
        return;
    }
    const size_t qk_stride = (size_t)group_count * (size_t)state_dim;
    const size_t value_stride = (size_t)num_heads * (size_t)state_dim;
    const size_t gate_stride = (size_t)num_heads;
    for (int row = 0; row < rows; ++row) {
        gated_deltanet_pytorch_grouped_bf16_forward(
            q + (size_t)row * qk_stride,
            k + (size_t)row * qk_stride,
            v + (size_t)row * value_stride,
            g + (size_t)row * gate_stride,
            beta + (size_t)row * gate_stride,
            row == 0 ? state_in : state_out,
            state_out,
            out + (size_t)row * value_stride,
            num_heads, group_count, state_dim, norm_eps);
    }
}

References gated_deltanet_pytorch_grouped_bf16_forward().

◆ geglu_backward_bf16_mixed()

void geglu_backward_bf16_mixed	(	const uint16_t *	x,
		const uint16_t *	d_out,
		float *	d_x,
		int	tokens,
		int	dim
	)

Definition at line 139 of file geglu_kernels.c.

{
    if (!x || !d_out || !d_x || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
    const int inner_dim = dim * 2;
 
    for (int t = 0; t < tokens; ++t) {
        const uint16_t *x_ptr = x + (size_t)t * (size_t)inner_dim;
        const uint16_t *d_out_ptr = d_out + (size_t)t * (size_t)dim;
        float *d_x_ptr = d_x + (size_t)t * (size_t)inner_dim;
 
        for (int d = 0; d < dim; ++d) {
            const float a = bf16_to_float(x_ptr[d]);
            const float b = bf16_to_float(x_ptr[dim + d]);
            const float dout = bf16_to_float(d_out_ptr[d]);
 
            const float a2 = a * a;
            const float a3 = a2 * a;
            const float g = sqrt_2_over_pi * (a + coeff * a3);
            const float tanh_g = tanhf(g);
            const float sech2_g = 1.0f - tanh_g * tanh_g;
            const float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * a2);
 
            const float d_gelu = 0.5f * (1.0f + tanh_g) + 0.5f * a * sech2_g * g_prime;
            d_x_ptr[d] = dout * d_gelu * b;
 
            const float gelu_a = 0.5f * a * (1.0f + tanh_g);
            d_x_ptr[dim + d] = dout * gelu_a;
        }
    }
}

References bf16_to_float().

◆ geglu_backward_fp32()

void geglu_backward_fp32	(	const float *	x,
		const float *	d_out,
		float *	d_x,
		int	tokens,
		int	dim
	)

Definition at line 99 of file geglu_kernels.c.

{
    if (!x || !d_out || !d_x || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
    const int inner_dim = dim * 2;
 
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * (size_t)inner_dim;
        const float *d_out_ptr = d_out + (size_t)t * (size_t)dim;
        float *d_x_ptr = d_x + (size_t)t * (size_t)inner_dim;
 
        for (int d = 0; d < dim; ++d) {
            float a = x_ptr[d];
            float b = x_ptr[dim + d];
            float dout = d_out_ptr[d];
 
            float a2 = a * a;
            float a3 = a2 * a;
            float g = sqrt_2_over_pi * (a + coeff * a3);
            float tanh_g = tanhf(g);
            float sech2_g = 1.0f - tanh_g * tanh_g;
            float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * a2);
 
            float d_gelu = 0.5f * (1.0f + tanh_g) + 0.5f * a * sech2_g * g_prime;
            d_x_ptr[d] = dout * d_gelu * b;
 
            float gelu_a = 0.5f * a * (1.0f + tanh_g);
            d_x_ptr[dim + d] = dout * gelu_a;
        }
    }
}

◆ geglu_forward_bf16()

void geglu_forward_bf16	(	const uint16_t *	x,
		uint16_t *	out,
		int	tokens,
		int	dim,
		float *	scratch
	)

Definition at line 83 of file geglu_kernels.c.

{
    if (!x || !out || !scratch || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const size_t fp32_size = (size_t)tokens * (size_t)dim;
    const size_t input_size = fp32_size * 2;
    float *fp32_input = scratch;
    float *fp32_output = scratch + input_size;
 
    bf16_tensor_to_float(x, fp32_input, input_size);
    geglu_forward_fp32(fp32_input, fp32_output, tokens, dim);
    float_tensor_to_bf16(fp32_output, out, fp32_size);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and geglu_forward_fp32().

◆ geglu_forward_exact()

void geglu_forward_exact	(	const float *	x,
		float *	out,
		int	tokens,
		int	dim
	)

Definition at line 42 of file geglu_kernels.c.

{
    const int inner_dim = dim * 2;
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * (size_t)inner_dim;
        float *out_ptr = out + (size_t)t * (size_t)dim;
 
        for (int d = 0; d < dim; ++d) {
            out_ptr[d] = ck_gelu_tanh_parity_f32(x_ptr[d]) * x_ptr[dim + d];
        }
    }
}

References ck_gelu_tanh_parity_f32().

Referenced by geglu_forward_fp32().

◆ geglu_forward_fp32()

void geglu_forward_fp32	(	const float *	x,
		float *	out,
		int	tokens,
		int	dim
	)

Definition at line 55 of file geglu_kernels.c.

{
    if (!x || !out || tokens <= 0 || dim <= 0) {
        return;
    }
 
    if (ck_strict_parity_enabled()) {
        geglu_forward_exact(x, out, tokens, dim);
        return;
    }
 
    const int inner_dim = dim * 2;
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * (size_t)inner_dim;
        float *out_ptr = out + (size_t)t * (size_t)dim;
 
        for (int d = 0; d < dim; ++d) {
            out_ptr[d] = x_ptr[d];
        }
 
        gelu_fast_inplace(out_ptr, (size_t)dim);
 
        for (int d = 0; d < dim; ++d) {
            out_ptr[d] *= x_ptr[dim + d];
        }
    }
}

References ck_strict_parity_enabled(), geglu_forward_exact(), and gelu_fast_inplace().

Referenced by geglu_forward_bf16().

◆ geglu_forward_ggml_native()

void geglu_forward_ggml_native	(	const float *	x,
		float *	out,
		int	tokens,
		int	dim
	)

Definition at line 17 of file geglu_kernels.c.

{
    if (!x || !out || tokens <= 0 || dim <= 0) return;
    /* Preserve ascending compaction when output aliases the gate/up input. */
    for (int t = 0; t < tokens; ++t) {
        const float *row = x + (size_t)t * (size_t)dim * 2;
        float *dst = out + (size_t)t * dim;
        for (int d = 0; d < dim; ++d) {
            const float gate = row[d], up = row[dim + d];
            float value = gate;
            gelu_ggml_native_inplace(&value, 1);
            dst[d] = gate <= -10.0f ? 0.0f : value * up;
        }
    }
}

References gelu_ggml_native_inplace().

◆ gelu_backward_exact()

void gelu_backward_exact	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		size_t	n
	)

Definition at line 368 of file gelu_kernels.c.

{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
#if defined(__AVX512F__)
    const __m512 sqrt_2_pi_vec = _mm512_set1_ps(sqrt_2_over_pi);
    const __m512 coeff_vec = _mm512_set1_ps(coeff);
    const __m512 coeff3_vec = _mm512_set1_ps(3.0f * coeff);
    const __m512 half_vec = _mm512_set1_ps(0.5f);
    const __m512 one_vec = _mm512_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 x = _mm512_loadu_ps(&input[i]);
        __m512 dy = _mm512_loadu_ps(&d_output[i]);
 
        __m512 x2 = _mm512_mul_ps(x, x);
        __m512 x3 = _mm512_mul_ps(x2, x);
 
        // g = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m512 g = _mm512_fmadd_ps(coeff_vec, x3, x);
        g = _mm512_mul_ps(sqrt_2_pi_vec, g);
 
        __m512 tanh_g = tanh512_fast(g);
 
        // g' = sqrt(2/pi) * (1 + 3 * 0.044715 * x^2)
        __m512 g_prime = _mm512_fmadd_ps(coeff3_vec, x2, one_vec);
        g_prime = _mm512_mul_ps(sqrt_2_pi_vec, g_prime);
 
        // sech^2(g) = 1 - tanh^2(g)
        __m512 sech2_g = _mm512_fnmadd_ps(tanh_g, tanh_g, one_vec);
 
        // gelu_derivative = 0.5 * (1 + tanh_g) + 0.5 * x * sech2_g * g_prime
        __m512 term1 = _mm512_mul_ps(half_vec, _mm512_add_ps(one_vec, tanh_g));
        __m512 term2 = _mm512_mul_ps(half_vec, _mm512_mul_ps(x, _mm512_mul_ps(sech2_g, g_prime)));
        __m512 gelu_deriv = _mm512_add_ps(term1, term2);
 
        __m512 result = _mm512_mul_ps(dy, gelu_deriv);
        _mm512_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative = 0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#elif defined(__AVX2__)
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 coeff3_vec = _mm256_set1_ps(3.0f * coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&input[i]);
        __m256 dy = _mm256_loadu_ps(&d_output[i]);
 
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // g = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 g = _mm256_fmadd_ps(coeff_vec, x3, x);
        g = _mm256_mul_ps(sqrt_2_pi_vec, g);
 
        __m256 tanh_g = tanh256_fast(g);
 
        // g' = sqrt(2/pi) * (1 + 3 * 0.044715 * x^2)
        __m256 g_prime = _mm256_fmadd_ps(coeff3_vec, x2, one_vec);
        g_prime = _mm256_mul_ps(sqrt_2_pi_vec, g_prime);
 
        // sech^2(g) = 1 - tanh^2(g)
        __m256 sech2_g = _mm256_fnmadd_ps(tanh_g, tanh_g, one_vec);
 
        // gelu_derivative = 0.5 * (1 + tanh_g) + 0.5 * x * sech2_g * g_prime
        __m256 term1 = _mm256_mul_ps(half_vec, _mm256_add_ps(one_vec, tanh_g));
        __m256 term2 = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, _mm256_mul_ps(sech2_g, g_prime)));
        __m256 gelu_deriv = _mm256_add_ps(term1, term2);
 
        __m256 result = _mm256_mul_ps(dy, gelu_deriv);
        _mm256_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative = 0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#elif defined(__AVX__)
    // AVX1: Vectorize arithmetic, use scalar tanh
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 coeff3_vec = _mm256_set1_ps(3.0f * coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    float g_arr[8] __attribute__((aligned(32)));
    float tanh_arr[8] __attribute__((aligned(32)));
 
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&input[i]);
        __m256 dy = _mm256_loadu_ps(&d_output[i]);
 
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // g = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 coeff_x3 = _mm256_mul_ps(coeff_vec, x3);
        __m256 g = _mm256_mul_ps(sqrt_2_pi_vec, _mm256_add_ps(x, coeff_x3));
 
        // Compute tanh scalarly
        _mm256_store_ps(g_arr, g);
        for (int j = 0; j < 8; ++j) {
            tanh_arr[j] = tanhf(g_arr[j]);
        }
        __m256 tanh_g = _mm256_load_ps(tanh_arr);
 
        // g' = sqrt(2/pi) * (1 + 3 * 0.044715 * x^2)
        __m256 coeff3_x2 = _mm256_mul_ps(coeff3_vec, x2);
        __m256 g_prime = _mm256_mul_ps(sqrt_2_pi_vec, _mm256_add_ps(one_vec, coeff3_x2));
 
        // sech^2(g) = 1 - tanh^2(g)
        __m256 tanh_g_sq = _mm256_mul_ps(tanh_g, tanh_g);
        __m256 sech2_g = _mm256_sub_ps(one_vec, tanh_g_sq);
 
        // gelu_derivative = 0.5 * (1 + tanh_g) + 0.5 * x * sech2_g * g_prime
        __m256 term1 = _mm256_mul_ps(half_vec, _mm256_add_ps(one_vec, tanh_g));
        __m256 term2 = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, _mm256_mul_ps(sech2_g, g_prime)));
        __m256 gelu_deriv = _mm256_add_ps(term1, term2);
 
        __m256 result = _mm256_mul_ps(dy, gelu_deriv);
        _mm256_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative = 0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#else
    // Scalar fallback
    for (size_t i = 0; i < n; ++i) {
        float x = input[i];
 
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
 
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
 
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative =
            0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
 
        d_input[i] = d_output[i] * gelu_derivative;
    }
#endif
}

References __attribute__().

◆ gelu_backward_exact_bf16()

void gelu_backward_exact_bf16	(	const uint16_t *	input,
		const uint16_t *	d_output,
		uint16_t *	d_input,
		size_t	n,
		float *	scratch_input,
		float *	scratch_d_output,
		float *	scratch_d_input
	)

Definition at line 46 of file gelu_kernels_bf16.c.

{
    if (!scratch_input || !scratch_d_output || !scratch_d_input) return;
 
    bf16_tensor_to_float(input, scratch_input, n);
    bf16_tensor_to_float(d_output, scratch_d_output, n);
 
    // Use scalar exact version to avoid fast tanh approximation error
    // accumulating with BF16 precision loss.
    gelu_backward_scalar(scratch_input, scratch_d_output, scratch_d_input, n);
 
    float_tensor_to_bf16(scratch_d_input, d_input, n);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and gelu_backward_scalar().

◆ gelu_backward_fast()

void gelu_backward_fast	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		size_t	n
	)

Definition at line 818 of file gelu_kernels.c.

{
    const float beta = 1.702f;
 
#if defined(__AVX512F__)
    const __m512 beta_vec = _mm512_set1_ps(beta);
    const __m512 one_vec = _mm512_set1_ps(1.0f);
    const __m512 neg_beta_vec = _mm512_set1_ps(-beta);
 
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 x = _mm512_loadu_ps(&input[i]);
        __m512 dy = _mm512_loadu_ps(&d_output[i]);
 
        // s = sigmoid(beta * x) = 1 / (1 + exp(-beta * x))
        __m512 neg_beta_x = _mm512_mul_ps(neg_beta_vec, x);
        __m512 exp_neg = exp512_fast(neg_beta_x);
        __m512 s = _mm512_div_ps(one_vec, _mm512_add_ps(one_vec, exp_neg));
 
        // gelu_derivative = s * (1 + x * (1 - s) * beta)
        __m512 one_minus_s = _mm512_sub_ps(one_vec, s);
        __m512 inner = _mm512_fmadd_ps(_mm512_mul_ps(x, one_minus_s), beta_vec, one_vec);
        __m512 gelu_deriv = _mm512_mul_ps(s, inner);
 
        __m512 result = _mm512_mul_ps(dy, gelu_deriv);
        _mm512_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float s = 1.0f / (1.0f + expf(-beta * x));
        float gelu_derivative = s * (1.0f + x * (1.0f - s) * beta);
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#elif defined(__AVX2__)
    const __m256 beta_vec = _mm256_set1_ps(beta);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
    const __m256 neg_beta_vec = _mm256_set1_ps(-beta);
 
    size_t i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&input[i]);
        __m256 dy = _mm256_loadu_ps(&d_output[i]);
 
        // s = sigmoid(beta * x) = 1 / (1 + exp(-beta * x))
        __m256 neg_beta_x = _mm256_mul_ps(neg_beta_vec, x);
        __m256 exp_neg = exp256_fast(neg_beta_x);
        __m256 s = _mm256_div_ps(one_vec, _mm256_add_ps(one_vec, exp_neg));
 
        // gelu_derivative = s * (1 + x * (1 - s) * beta)
        __m256 one_minus_s = _mm256_sub_ps(one_vec, s);
        __m256 inner = _mm256_fmadd_ps(_mm256_mul_ps(x, one_minus_s), beta_vec, one_vec);
        __m256 gelu_deriv = _mm256_mul_ps(s, inner);
 
        __m256 result = _mm256_mul_ps(dy, gelu_deriv);
        _mm256_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float s = 1.0f / (1.0f + expf(-beta * x));
        float gelu_derivative = s * (1.0f + x * (1.0f - s) * beta);
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#elif defined(__AVX__)
    // AVX1: Vectorize arithmetic, use scalar exp
    const __m256 beta_vec = _mm256_set1_ps(beta);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
    const __m256 neg_beta_vec = _mm256_set1_ps(-beta);
 
    size_t i = 0;
    float neg_beta_x_arr[8] __attribute__((aligned(32)));
    float exp_arr[8] __attribute__((aligned(32)));
 
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&input[i]);
        __m256 dy = _mm256_loadu_ps(&d_output[i]);
 
        // s = sigmoid(beta * x) = 1 / (1 + exp(-beta * x))
        __m256 neg_beta_x = _mm256_mul_ps(neg_beta_vec, x);
 
        // Compute exp scalarly
        _mm256_store_ps(neg_beta_x_arr, neg_beta_x);
        for (int j = 0; j < 8; ++j) {
            exp_arr[j] = expf(neg_beta_x_arr[j]);
        }
        __m256 exp_neg = _mm256_load_ps(exp_arr);
 
        __m256 s = _mm256_div_ps(one_vec, _mm256_add_ps(one_vec, exp_neg));
 
        // gelu_derivative = s * (1 + x * (1 - s) * beta)
        __m256 one_minus_s = _mm256_sub_ps(one_vec, s);
        __m256 x_one_minus_s = _mm256_mul_ps(x, one_minus_s);
        __m256 x_one_minus_s_beta = _mm256_mul_ps(x_one_minus_s, beta_vec);
        __m256 inner = _mm256_add_ps(one_vec, x_one_minus_s_beta);
        __m256 gelu_deriv = _mm256_mul_ps(s, inner);
 
        __m256 result = _mm256_mul_ps(dy, gelu_deriv);
        _mm256_storeu_ps(&d_input[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = input[i];
        float s = 1.0f / (1.0f + expf(-beta * x));
        float gelu_derivative = s * (1.0f + x * (1.0f - s) * beta);
        d_input[i] = d_output[i] * gelu_derivative;
    }
 
#else
    // Scalar fallback
    for (size_t i = 0; i < n; ++i) {
        float x = input[i];
        float s = 1.0f / (1.0f + expf(-beta * x));
        float gelu_derivative = s * (1.0f + x * (1.0f - s) * beta);
        d_input[i] = d_output[i] * gelu_derivative;
    }
#endif
}

References __attribute__().

Referenced by gelu_backward_fast_bf16().

◆ gelu_backward_fast_bf16()

void gelu_backward_fast_bf16	(	const uint16_t *	input,
		const uint16_t *	d_output,
		uint16_t *	d_input,
		size_t	n,
		float *	scratch_input,
		float *	scratch_d_output,
		float *	scratch_d_input
	)

Definition at line 69 of file gelu_kernels_bf16.c.

{
    if (!scratch_input || !scratch_d_output || !scratch_d_input) return;
 
    bf16_tensor_to_float(input, scratch_input, n);
    bf16_tensor_to_float(d_output, scratch_d_output, n);
 
    gelu_backward_fast(scratch_input, scratch_d_output, scratch_d_input, n);
 
    float_tensor_to_bf16(scratch_d_input, d_input, n);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and gelu_backward_fast().

◆ gelu_backward_scalar()

void gelu_backward_scalar	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		size_t	n
	)

Definition at line 794 of file gelu_kernels.c.

{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
    for (size_t i = 0; i < n; ++i) {
        float x = input[i];
        float x3 = x * x * x;
        float g = sqrt_2_over_pi * (x + coeff * x3);
        float tanh_g = tanhf(g);
        float x2 = x * x;
        float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * x2);
        float sech2_g = 1.0f - tanh_g * tanh_g;
        float gelu_derivative = 0.5f * (1.0f + tanh_g) + 0.5f * x * sech2_g * g_prime;
        d_input[i] = d_output[i] * gelu_derivative;
    }
}

Referenced by gelu_backward_exact_bf16().

◆ gelu_erf_bf16_storage()

void gelu_erf_bf16_storage	(	float *	data,
		size_t	n
	)

Definition at line 609 of file gelu_kernels.c.

{
    const double inv_sqrt_2 = 0.707106781186547524400844362104849039;
    ck_gelu_math_f64_fn reference_erf = ck_gelu_system_erf();
    for (size_t i = 0; i < n; ++i) {
        const float x = bf16_to_float(float_to_bf16(data[i]));
        const double scaled = (double)x * inv_sqrt_2;
        const double erf_value = reference_erf ? reference_erf(scaled) : erf(scaled);
        const float output = (float)(0.5 * (double)x * (1.0 + erf_value));
        data[i] = bf16_to_float(float_to_bf16(output));
    }
}

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

◆ gelu_erf_fp64_f32_inplace()

void gelu_erf_fp64_f32_inplace	(	float *	data,
		size_t	n
	)

Definition at line 566 of file gelu_kernels.c.

{
    const double inv_sqrt_2 = 0.707106781186547524400844362104849039;
    ck_gelu_math_f64_fn reference_erf = ck_gelu_system_erf();
    for (size_t i = 0; i < n; ++i) {
        const float x = data[i];
        const double scaled = (double)x * inv_sqrt_2;
        const double erf_value = reference_erf ? reference_erf(scaled) : erf(scaled);
        data[i] = (float)(0.5 * (double)x * (1.0 + erf_value));
    }
}

References ck_gelu_system_erf().

Referenced by gelu_pytorch_erf_f32_inplace().

◆ gelu_exact_inplace()

void gelu_exact_inplace	(	float *	data,
		size_t	n
	)

Definition at line 557 of file gelu_kernels.c.

{
    for (size_t i = 0; i < n; ++i) {
        data[i] = ck_gelu_tanh_f32(data[i]);
    }
}

References ck_gelu_tanh_f32().

Referenced by gelu_fast_inplace_bf16(), and mlp_token_parallel_exact().

◆ gelu_fast_inplace()

void gelu_fast_inplace	(	float *	data,
		size_t	n
	)

GELU activation forward (fast approximation, in-place)

Test:

test_gelu.py::TestGELUForward::test_gelu_fast_inplace

test_gelu.py::TestGELUForward::test_gelu_vs_exact

test_parity.py::test_gelu_parity

Fast GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) In-place on contiguous buffer.

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

Definition at line 243 of file gelu_kernels.c.

{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
#if defined(__AVX512F__)
    const __m512 sqrt_2_pi_vec = _mm512_set1_ps(sqrt_2_over_pi);
    const __m512 coeff_vec = _mm512_set1_ps(coeff);
    const __m512 half_vec = _mm512_set1_ps(0.5f);
    const __m512 one_vec = _mm512_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 x = _mm512_loadu_ps(&data[i]);
        __m512 x2 = _mm512_mul_ps(x, x);
        __m512 x3 = _mm512_mul_ps(x2, x);
 
        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m512 inner = _mm512_fmadd_ps(coeff_vec, x3, x);
        inner = _mm512_mul_ps(sqrt_2_pi_vec, inner);
 
        // result = 0.5 * x * (1 + tanh(inner))
        __m512 tanh_val = tanh512_fast(inner);
        __m512 one_plus_tanh = _mm512_add_ps(one_vec, tanh_val);
        __m512 result = _mm512_mul_ps(half_vec, _mm512_mul_ps(x, one_plus_tanh));
 
        _mm512_storeu_ps(&data[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
 
#elif defined(__AVX2__)
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&data[i]);
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 inner = _mm256_fmadd_ps(coeff_vec, x3, x);
        inner = _mm256_mul_ps(sqrt_2_pi_vec, inner);
 
        // result = 0.5 * x * (1 + tanh(inner))
        __m256 tanh_val = tanh256_fast(inner);
        __m256 one_plus_tanh = _mm256_add_ps(one_vec, tanh_val);
        __m256 result = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, one_plus_tanh));
 
        _mm256_storeu_ps(&data[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
 
#elif defined(__AVX__)
    // AVX1: Vectorize arithmetic, use scalar tanh
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    float inner_arr[8] __attribute__((aligned(32)));
    float tanh_arr[8] __attribute__((aligned(32)));
 
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&data[i]);
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 coeff_x3 = _mm256_mul_ps(coeff_vec, x3);
        __m256 inner = _mm256_mul_ps(sqrt_2_pi_vec, _mm256_add_ps(x, coeff_x3));
 
        // Compute tanh scalarly
        _mm256_store_ps(inner_arr, inner);
        for (int j = 0; j < 8; ++j) {
            tanh_arr[j] = tanhf(inner_arr[j]);
        }
        __m256 tanh_val = _mm256_load_ps(tanh_arr);
 
        // result = 0.5 * x * (1 + tanh(inner))
        __m256 one_plus_tanh = _mm256_add_ps(one_vec, tanh_val);
        __m256 result = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, one_plus_tanh));
 
        _mm256_storeu_ps(&data[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
 
#else
    // Scalar fallback
    for (size_t i = 0; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
#endif
}

References __attribute__().

Referenced by mlp_token_parallel().

◆ gelu_fast_inplace_bf16()

void gelu_fast_inplace_bf16	(	uint16_t *	data,
		size_t	n,
		float *	scratch
	)

Definition at line 31 of file gelu_kernels_bf16.c.

{
    if (!scratch) return;
 
    bf16_tensor_to_float(data, scratch, n);
    // Use exact version to avoid fast tanh approximation error accumulating
    // with BF16 precision loss. Conversion overhead dominates anyway.
    gelu_exact_inplace(scratch, n);
    float_tensor_to_bf16(scratch, data, n);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and gelu_exact_inplace().

◆ gelu_ggml_inplace()

void gelu_ggml_inplace	(	float *	data,
		size_t	n
	)

Definition at line 753 of file gelu_kernels.c.

{
    pthread_once(&ck_gelu_ggml_runtime_once, ck_gelu_ggml_runtime_init);
    if (ck_gelu_runtime_ready) {
        for (size_t i = 0; i < n; ++i) {
            const float x = data[i];
            if (x <= -10.0f) {
                data[i] = 0.0f;
                continue;
            }
            if (x >= 10.0f) {
                data[i] = x;
                continue;
            }
            const ck_half x_fp16 = ck_gelu_runtime_fp32_to_fp16(x);
            const ck_half y_fp16 = ck_gelu_runtime_table_f16[(uint16_t) x_fp16];
            data[i] = ck_gelu_runtime_fp16_to_fp32(y_fp16);
        }
        return;
    }
 
    pthread_once(&ck_gelu_ggml_table_once, ck_gelu_ggml_table_init);
    for (size_t i = 0; i < n; ++i) {
        const float x = data[i];
        if (x <= -10.0f) {
            data[i] = 0.0f;
            continue;
        }
        if (x >= 10.0f) {
            data[i] = x;
            continue;
        }
        const ck_half x_fp16 = ggml_fp32_to_fp16(x);
        const ck_half y_fp16 = ck_gelu_ggml_table_f16[(uint16_t) x_fp16];
        data[i] = ggml_fp16_to_fp32(y_fp16);
    }
}

References ck_gelu_ggml_runtime_init(), ck_gelu_ggml_runtime_once, ck_gelu_ggml_table_f16, ck_gelu_ggml_table_init(), ck_gelu_ggml_table_once, ck_gelu_runtime_fp16_to_fp32, ck_gelu_runtime_fp32_to_fp16, ck_gelu_runtime_ready, ck_gelu_runtime_table_f16, ggml_fp16_to_fp32, and ggml_fp32_to_fp16.

◆ gelu_ggml_native_inplace()

void gelu_ggml_native_inplace	(	float *	data,
		size_t	n
	)

Definition at line 741 of file gelu_kernels.c.

{
    pthread_once(&ck_gelu_ggml_table_once, ck_gelu_ggml_table_init);
    for (size_t i = 0; i < n; ++i) {
        const float x = data[i];
        if (x <= -10.0f) data[i] = 0.0f;
        else if (x >= 10.0f) data[i] = x;
        else data[i] = ggml_fp16_to_fp32(
            ck_gelu_ggml_table_f16[(uint16_t)ggml_fp32_to_fp16(x)]);
    }
}

References ck_gelu_ggml_table_f16, ck_gelu_ggml_table_init(), ck_gelu_ggml_table_once, ggml_fp16_to_fp32, and ggml_fp32_to_fp16.

◆ gelu_pytorch_erf_f32_inplace()

void gelu_pytorch_erf_f32_inplace	(	float *	data,
		size_t	n
	)

Definition at line 580 of file gelu_kernels.c.

{
    gelu_erf_fp64_f32_inplace(data, n);
}

References gelu_erf_fp64_f32_inplace().

◆ gelu_pytorch_erf_sleef_bf16_storage()

void gelu_pytorch_erf_sleef_bf16_storage	(	float *	data,
		size_t	n
	)

Definition at line 682 of file gelu_kernels.c.

{
#if defined(__AVX512F__)
    pthread_once(&ck_pytorch_sleef_once, ck_bind_pytorch_sleef);
    if (!ck_pytorch_sleef_expf16) {
        fprintf(stderr,
                "[CK] PyTorch-exact BF16 GELU requires Sleef_expf16_u10; "
                "set CK_SLEEF_LIBRARY to libtorch_cpu.so or libsleef.so\n");
        abort();
    }
 
    const __m512 alpha = _mm512_set1_ps(0.70710678118654752440f);
    const __m512 half = _mm512_set1_ps(0.5f);
    const __m512 one = _mm512_set1_ps(1.0f);
    const __m512 neg_zero = _mm512_set1_ps(-0.0f);
    const __m512 p = _mm512_set1_ps(0.3275911f);
    const __m512 p1 = _mm512_set1_ps(0.254829592f);
    const __m512 p2 = _mm512_set1_ps(-0.284496736f);
    const __m512 p3 = _mm512_set1_ps(1.421413741f);
    const __m512 p4 = _mm512_set1_ps(-1.453152027f);
    const __m512 p5 = _mm512_set1_ps(1.061405429f);
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 x = _mm512_loadu_ps(data + i);
        __m512 erf_arg = _mm512_mul_ps(x, alpha);
        __m512 sign = _mm512_and_ps(neg_zero, erf_arg);
        __m512 abs_arg = _mm512_abs_ps(erf_arg);
        __m512 t = _mm512_div_ps(one, _mm512_fmadd_ps(p, abs_arg, one));
        __m512 r = _mm512_fmadd_ps(p5, t, p4);
        r = _mm512_fmadd_ps(r, t, p3);
        r = _mm512_fmadd_ps(r, t, p2);
        r = _mm512_fmadd_ps(r, t, p1);
        __m512 arg_sq = _mm512_mul_ps(erf_arg, erf_arg);
        __m512 exp_neg_sq = ck_pytorch_sleef_expf16(_mm512_xor_ps(neg_zero, arg_sq));
        __m512 neg_exp_t = _mm512_mul_ps(_mm512_xor_ps(neg_zero, exp_neg_sq), t);
        __m512 erf_x = _mm512_xor_ps(sign, _mm512_fmadd_ps(neg_exp_t, r, one));
        __m512 y = _mm512_mul_ps(_mm512_mul_ps(x, half), _mm512_add_ps(one, erf_x));
        float lanes[16];
        _mm512_storeu_ps(lanes, y);
        for (size_t lane = 0; lane < 16; ++lane) {
            const uint16_t input_code = float_to_bf16(data[i + lane]);
            const uint16_t output_code = ck_pytorch_gelu_erf_bf16_edge(
                input_code, float_to_bf16(lanes[lane]));
            data[i + lane] = bf16_to_float(output_code);
        }
    }
    for (; i < n; ++i) {
        const float x = bf16_to_float(float_to_bf16(data[i]));
        const float y = (x * 0.5f) * (1.0f + erff(x * 0.70710678118654752440f));
        data[i] = bf16_to_float(float_to_bf16(y));
    }
#else
    (void)data;
    (void)n;
    fprintf(stderr, "[CK] PyTorch-exact BF16 GELU requires an AVX-512 build\n");
    abort();
#endif
}

References bf16_to_float(), and float_to_bf16().

◆ gelu_pytorch_tanh_bf16_storage()

void gelu_pytorch_tanh_bf16_storage	(	float *	data,
		size_t	n
	)

Definition at line 590 of file gelu_kernels.c.

{
    /* PyTorch's x86 BF16 kernel widens to FP32 and evaluates tanh through
     * SLEEF's vector u10 provider.  That provider saturates at this exact
     * inner-argument boundary; libc tanhf retains a small tail and can round
     * to a different BF16 code before the following projection. */
    const float sleef_tanh_saturation = 8.664339742f;
    for (size_t i = 0; i < n; ++i) {
        const float x = bf16_to_float(float_to_bf16(data[i]));
        const float x3 = x * x * x;
        const float inner = 0.7978845608f * (x + 0.044715f * x3);
        const float tanh_inner = fabsf(inner) > sleef_tanh_saturation
            ? copysignf(1.0f, inner)
            : tanhf(inner);
        const float output = 0.5f * x * (1.0f + tanh_inner);
        data[i] = bf16_to_float(float_to_bf16(output));
    }
}

References bf16_to_float(), and float_to_bf16().

◆ gemm_avx512_parallel()

void gemm_avx512_parallel	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 230 of file gemm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        gemm_naive_serial_float(A, B, bias, C, M, N, K);
        return;
    }
#if defined(__AVX512F__)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            __m512 sum_vec = _mm512_setzero_ps();
            int k;
            for (k = 0; k <= K - 16; k += 16) {
                __m512 a_vec = _mm512_loadu_ps(&A[i * K + k]);
                __m512 b_vec = _mm512_loadu_ps(&B[j * K + k]);
                sum_vec = _mm512_fmadd_ps(a_vec, b_vec, sum_vec);
            }
            float sum = _mm512_reduce_add_ps(sum_vec);
            for (; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            float bias_val = bias ? bias[j] : 0.0f;
            C[i * N + j] = sum + bias_val;
        }
    }
#elif defined(__AVX__)
    // AVX1 path: 256-bit vectors, no FMA (use mul + add)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            __m256 sum_vec = _mm256_setzero_ps();
            int k;
            for (k = 0; k <= K - 8; k += 8) {
                __m256 a_vec = _mm256_loadu_ps(&A[i * K + k]);
                __m256 b_vec = _mm256_loadu_ps(&B[j * K + k]);
                __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, prod);
            }
            float sum = hsum256_ps(sum_vec);
            for (; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            float bias_val = bias ? bias[j] : 0.0f;
            C[i * N + j] = sum + bias_val;
        }
    }
#else
    gemm_naive_parallel(A, B, bias, C, M, N, K);
#endif
}

References C, ck_strict_parity_enabled(), gemm_naive_parallel(), and gemm_naive_serial_float().

◆ gemm_backward_bf16_mixed()

void gemm_backward_bf16_mixed	(	const uint16_t *	d_output,
		const uint16_t *	input,
		const uint16_t *	weight,
		float *	d_input,
		float *	d_weight,
		float *	d_bias,
		int	tokens,
		int	in_dim,
		int	out_dim
	)

Definition at line 2043 of file gemm_kernels_bf16.c.

{
    if (!d_output || !input || !weight || tokens <= 0 || in_dim <= 0 || out_dim <= 0) {
        return;
    }
 
    if (d_input) {
        for (int t = 0; t < tokens; ++t) {
            for (int i = 0; i < in_dim; ++i) {
                float sum = 0.0f;
                for (int o = 0; o < out_dim; ++o) {
                    const float dy = bf16_to_float(d_output[(size_t)t * (size_t)out_dim + (size_t)o]);
                    const float w = bf16_to_float(weight[(size_t)o * (size_t)in_dim + (size_t)i]);
                    sum += dy * w;
                }
                d_input[(size_t)t * (size_t)in_dim + (size_t)i] = sum;
            }
        }
    }
 
    if (d_weight) {
        for (int o = 0; o < out_dim; ++o) {
            for (int i = 0; i < in_dim; ++i) {
                float sum = 0.0f;
                for (int t = 0; t < tokens; ++t) {
                    const float dy = bf16_to_float(d_output[(size_t)t * (size_t)out_dim + (size_t)o]);
                    const float x = bf16_to_float(input[(size_t)t * (size_t)in_dim + (size_t)i]);
                    sum += dy * x;
                }
                d_weight[(size_t)o * (size_t)in_dim + (size_t)i] = sum;
            }
        }
    }
 
    if (d_bias) {
        for (int o = 0; o < out_dim; ++o) {
            float sum = 0.0f;
            for (int t = 0; t < tokens; ++t) {
                sum += bf16_to_float(d_output[(size_t)t * (size_t)out_dim + (size_t)o]);
            }
            d_bias[o] = sum;
        }
    }
}

References bf16_to_float().

◆ gemm_backward_f32_train_parallel_dispatch()

void gemm_backward_f32_train_parallel_dispatch	(	const float *	d_output,
		const float *	input,
		const float *	W,
		float *	d_input,
		float *	d_W,
		float *	d_b,
		int	T,
		int	aligned_in,
		int	aligned_out,
		int	num_threads
	)

Definition at line 610 of file ck_parallel_train.c.

                                                                {
    if (!d_output || !input || !W || !d_input || !d_W) {
        return;
    }
    if (T <= 0 || aligned_in <= 0 || aligned_out <= 0) {
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
    if (num_threads > 0 && num_threads < nth) {
        nth = num_threads;
    }
 
    if (!pool || nth <= 1) {
        ck_train_gemm_backward_serial(d_output, input, W, d_input, d_W, d_b, T, aligned_in, aligned_out);
        return;
    }
 
    /* T=1 is the dominant runtime shape in generated train microsteps.
     * Use vectorized d_input and parallel outer-product for dW/db. */
    if (T == 1) {
        gemm_nn_simd(d_output, W, NULL, d_input, 1, aligned_in, aligned_out);
 
        const size_t outer_work = (size_t)aligned_out * (size_t)aligned_in;
        if (aligned_out < nth * 2 || aligned_in < 64 || outer_work < (size_t)524288) {
            ck_train_outer_t1_compute_range(d_output, input, d_W, d_b, 0, aligned_out, aligned_in);
        } else {
            const int active_outer = ck_train_pick_active_threads(nth, (size_t)aligned_out, (size_t)64);
            ck_train_outer_t1_args_t t1_args = {
                .d_output = d_output,
                .input = input,
                .d_W = d_W,
                .d_b = d_b,
                .aligned_in = aligned_in,
                .aligned_out = aligned_out,
            };
            if (active_outer <= 1) {
                ck_train_outer_t1_compute_range(d_output, input, d_W, d_b, 0, aligned_out, aligned_in);
            } else {
                ck_threadpool_dispatch_n(pool, active_outer, ck_train_outer_t1_work, &t1_args);
            }
        }
        return;
    }
 
    const size_t nn_work = (size_t)T * (size_t)aligned_in * (size_t)aligned_out;
    const size_t tn_work = (size_t)aligned_out * (size_t)aligned_in * (size_t)T;
    if ((T < 2 && aligned_out < nth * 2) || (nn_work < (size_t)131072 && tn_work < (size_t)131072)) {
        ck_train_gemm_backward_serial(d_output, input, W, d_input, d_W, d_b, T, aligned_in, aligned_out);
        return;
    }
 
    ck_train_gemm_backward_args_t bw_args = {
        .d_output = d_output,
        .input = input,
        .W = W,
        .d_input = d_input,
        .d_W = d_W,
        .d_b = d_b,
        .T = T,
        .aligned_in = aligned_in,
        .aligned_out = aligned_out,
    };
    {
        const size_t bw_rows = (size_t)aligned_out > (size_t)T ? (size_t)aligned_out : (size_t)T;
        const int active_bw = ck_train_pick_active_threads(nth, bw_rows, (size_t)64);
        if (active_bw <= 1) {
            ck_train_gemm_backward_serial(d_output, input, W, d_input, d_W, d_b, T, aligned_in, aligned_out);
        } else {
            ck_threadpool_dispatch_n(pool, active_bw, ck_train_gemm_backward_work, &bw_args);
        }
    }
}

References ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_train_gemm_backward_serial(), ck_train_gemm_backward_work(), ck_train_outer_t1_compute_range(), ck_train_outer_t1_work(), ck_train_pick_active_threads(), and gemm_nn_simd().

◆ gemm_backward_f32_train_parallel_dispatch_v2()

void gemm_backward_f32_train_parallel_dispatch_v2	(	const float *	d_output,
		const float *	input,
		const float *	W,
		float *	d_input,
		float *	d_W,
		float *	d_b,
		int	T,
		int	aligned_in,
		int	aligned_out,
		int	num_threads
	)

Definition at line 699 of file ck_parallel_train.c.

                                                                   {
    if (!d_output || !input || !W || !d_input || !d_W) {
        return;
    }
    if (T <= 0 || aligned_in <= 0 || aligned_out <= 0) {
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
    if (num_threads > 0 && num_threads < nth) {
        nth = num_threads;
    }
 
    if (!pool || nth <= 1) {
        ck_train_gemm_backward_serial(d_output, input, W, d_input, d_W, d_b, T, aligned_in, aligned_out);
        return;
    }
 
    if (T == 1) {
        const size_t nn_work = (size_t)aligned_in * (size_t)aligned_out;
        if (aligned_in >= nth * 32 && nn_work >= (size_t)131072) {
            const int active_nn = ck_train_pick_active_threads(nth, (size_t)aligned_in, (size_t)64);
            ck_train_gemm_nn_args_t nn_args = {
                .A = d_output,
                .B = W,
                .bias = NULL,
                .C = d_input,
                .M = 1,
                .N = aligned_in,
                .K = aligned_out,
                .split_n = 1,
            };
            if (active_nn <= 1) {
                gemm_nn_simd(d_output, W, NULL, d_input, 1, aligned_in, aligned_out);
            } else {
                ck_threadpool_dispatch_n(pool, active_nn, ck_train_gemm_nn_work, &nn_args);
            }
        } else {
            gemm_nn_simd(d_output, W, NULL, d_input, 1, aligned_in, aligned_out);
        }
 
        const size_t outer_work = (size_t)aligned_out * (size_t)aligned_in;
        if (aligned_out >= nth && outer_work >= (size_t)131072) {
            const int active_outer = ck_train_pick_active_threads(nth, (size_t)aligned_out, (size_t)64);
            ck_train_outer_t1_args_t t1_args = {
                .d_output = d_output,
                .input = input,
                .d_W = d_W,
                .d_b = d_b,
                .aligned_in = aligned_in,
                .aligned_out = aligned_out,
            };
            if (active_outer <= 1) {
                ck_train_outer_t1_compute_range(d_output, input, d_W, d_b, 0, aligned_out, aligned_in);
            } else {
                ck_threadpool_dispatch_n(pool, active_outer, ck_train_outer_t1_work, &t1_args);
            }
        } else {
            ck_train_outer_t1_compute_range(d_output, input, d_W, d_b, 0, aligned_out, aligned_in);
        }
        return;
    }
 
    const size_t nn_work = (size_t)T * (size_t)aligned_in * (size_t)aligned_out;
    const size_t tn_work = (size_t)aligned_out * (size_t)aligned_in * (size_t)T;
    if ((T < nth && aligned_out < nth) || (nn_work < (size_t)131072 && tn_work < (size_t)131072)) {
        ck_train_gemm_backward_serial(d_output, input, W, d_input, d_W, d_b, T, aligned_in, aligned_out);
        return;
    }
 
    ck_train_gemm_backward_args_t bw_args = {
        .d_output = d_output,
        .input = input,
        .W = W,
        .d_input = d_input,
        .d_W = d_W,
        .d_b = d_b,
        .T = T,
        .aligned_in = aligned_in,
        .aligned_out = aligned_out,
    };
    {
        const size_t bw_rows = (size_t)aligned_out > (size_t)T ? (size_t)aligned_out : (size_t)T;
        const int active_bw = ck_train_pick_active_threads(nth, bw_rows, (size_t)64);
        if (active_bw <= 1) {
            ck_train_gemm_backward_serial(d_output, input, W, d_input, d_W, d_b, T, aligned_in, aligned_out);
        } else {
            ck_threadpool_dispatch_n(pool, active_bw, ck_train_gemm_backward_work, &bw_args);
        }
    }
}

References ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_train_gemm_backward_serial(), ck_train_gemm_backward_work(), ck_train_gemm_nn_work(), ck_train_outer_t1_compute_range(), ck_train_outer_t1_work(), ck_train_pick_active_threads(), and gemm_nn_simd().

◆ gemm_bias_gelu_fused()

void gemm_bias_gelu_fused	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 131 of file gemm_fused_kernels.c.

{
#if defined(__AVX__)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            __m256 sum_vec = _mm256_setzero_ps();
            int k;
            for (k = 0; k <= K - 8; k += 8) {
                __m256 a_vec = _mm256_loadu_ps(&A[i * K + k]);
                __m256 b_vec = _mm256_loadu_ps(&B[j * K + k]);
                __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, prod);
            }
            float sum = hsum256_ps_fused(sum_vec);
            for (; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            C[i * N + j] = fast_gelu_scalar(sum);
        }
    }
#else
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            C[i * N + j] = fast_gelu_scalar(sum);
        }
    }
#endif
}

References C, and fast_gelu_scalar().

◆ gemm_bias_relu_fused()

void gemm_bias_relu_fused	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 84 of file gemm_fused_kernels.c.

{
#if defined(__AVX__)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            __m256 sum_vec = _mm256_setzero_ps();
            int k;
            for (k = 0; k <= K - 8; k += 8) {
                __m256 a_vec = _mm256_loadu_ps(&A[i * K + k]);
                __m256 b_vec = _mm256_loadu_ps(&B[j * K + k]);
                __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, prod);
            }
            float sum = hsum256_ps_fused(sum_vec);
            for (; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            // Fused: add bias and ReLU while still in register
            sum += bias[j];
            C[i * N + j] = sum > 0.0f ? sum : 0.0f;
        }
    }
#else
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            C[i * N + j] = sum > 0.0f ? sum : 0.0f;
        }
    }
#endif
}

References C.

◆ gemm_bias_silu_fused()

void gemm_bias_silu_fused	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 177 of file gemm_fused_kernels.c.

{
#if defined(__AVX__)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            __m256 sum_vec = _mm256_setzero_ps();
            int k;
            for (k = 0; k <= K - 8; k += 8) {
                __m256 a_vec = _mm256_loadu_ps(&A[i * K + k]);
                __m256 b_vec = _mm256_loadu_ps(&B[j * K + k]);
                __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, prod);
            }
            float sum = hsum256_ps_fused(sum_vec);
            for (; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            // SiLU: x * sigmoid(x)
            float sig = 1.0f / (1.0f + expf(-sum));
            C[i * N + j] = sum * sig;
        }
    }
#else
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            sum += bias[j];
            float sig = 1.0f / (1.0f + expf(-sum));
            C[i * N + j] = sum * sig;
        }
    }
#endif
}

References C.

◆ gemm_blocked_serial()

void gemm_blocked_serial	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 849 of file gemm_kernels.c.

{
    // Ensure threads are initialized (auto-detects on first call)
    (void)ck_get_num_threads();
 
    if (ck_strict_parity_enabled()) {
        gemm_naive_serial_float(A, B, bias, C, M, N, K);
        return;
    }
 
    // Decode-time matvec (M=1) is extremely common and benefits from parallelism over N.
    // Lower threshold to parallelize more ops; OpenMP overhead is ~1-2μs per barrier.
    // For N*K >= 64K elements, parallel is worthwhile.
    if (M == 1 && (size_t)N * (size_t)K >= 65536) {
        gemm_nt_matvec_parallel(A, B, bias, C, N, K);
        return;
    }
 
    /*
     * Use gemm_microkernel for large matrices - it uses MKL/oneDNN when available,
     * which is substantially faster than our hand-written SIMD kernels.
     * B is stored as [N x K] (transposed), so we pass B_transposed=1.
     * Note: Use threshold of 32 to avoid numerical precision issues with small matrices.
     */
    if (M >= 32 && N >= 32 && K >= 32) {
        gemm_microkernel(A, B, C, M, N, K, 1);  // B_transposed=1
        ck_gemm_add_bias(C, bias, M, N);
        return;
    }
#if defined(__AVX512F__)
    const int block_size = 64;
#elif defined(__AVX__)
    const int block_size = 32;
#else
    const int block_size = 32;
#endif
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            C[i * N + j] = bias ? bias[j] : 0.0f;
        }
    }
    for (int ii = 0; ii < M; ii += block_size) {
        for (int jj = 0; jj < N; jj += block_size) {
            for (int kk = 0; kk < K; kk += block_size) {
                int i_end = ck_min(ii + block_size, M);
                int j_end = ck_min(jj + block_size, N);
                int k_end = ck_min(kk + block_size, K);
 
                for (int i = ii; i < i_end; i++) {
                    for (int j = jj; j < j_end; j++) {
#if defined(__AVX512F__)
                        __m512 sum_vec = _mm512_setzero_ps();
                        int k;
                        for (k = kk; k <= k_end - 16; k += 16) {
                            __m512 a_vec = _mm512_loadu_ps(&A[i * K + k]);
                            __m512 b_vec = _mm512_loadu_ps(&B[j * K + k]);
                            sum_vec = _mm512_fmadd_ps(a_vec, b_vec, sum_vec);
                        }
                        float partial_sum = _mm512_reduce_add_ps(sum_vec);
                        for (; k < k_end; k++) {
                            partial_sum += A[i * K + k] * B[j * K + k];
                        }
#elif defined(__AVX__)
                        __m256 sum_vec = _mm256_setzero_ps();
                        int k;
                        for (k = kk; k <= k_end - 8; k += 8) {
                            __m256 a_vec = _mm256_loadu_ps(&A[i * K + k]);
                            __m256 b_vec = _mm256_loadu_ps(&B[j * K + k]);
                            __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                            sum_vec = _mm256_add_ps(sum_vec, prod);
                        }
                        float partial_sum = hsum256_ps(sum_vec);
                        for (; k < k_end; k++) {
                            partial_sum += A[i * K + k] * B[j * K + k];
                        }
#else
                        float partial_sum = 0.0f;
                        for (int k = kk; k < k_end; k++) {
                            partial_sum += A[i * K + k] * B[j * K + k];
                        }
#endif
                        C[i * N + j] += partial_sum;
                    }
                }
            }
        }
    }
}

References C, ck_gemm_add_bias(), ck_get_num_threads(), ck_min(), ck_strict_parity_enabled(), gemm_microkernel(), gemm_naive_serial_float(), and gemm_nt_matvec_parallel().

Referenced by ck_attention_project_head_major(), ck_gemm_nt_quant(), ck_mlp_swiglu_forward(), ck_mlp_swiglu_forward_fused_token(), ck_qkv_project_head_major(), ck_qkv_project_head_major_token(), ck_train_gemm_work(), gemm_blocked_serial_train_parallel_dispatch(), mlp_token_parallel(), and mlp_token_parallel_exact().

◆ gemm_blocked_serial_bf16()

void gemm_blocked_serial_bf16	(	const uint16_t *	A,
		const uint16_t *	B,
		const uint16_t *	bias,
		uint16_t *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 424 of file gemm_kernels_bf16.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
#if HAVE_NATIVE_BF16
    /* Native BF16 instructions available (Ice Lake / Sapphire Rapids+) */
    gemm_bf16_native(A, B, bias, C, M, N, K);
#elif defined(__AVX512F__)
    /* Use AVX-512F with software BF16 conversion */
    if (M * N > 4096) {
        gemm_bf16_blocked_avx512(A, B, bias, C, M, N, K);
    } else {
        gemm_bf16_avx512(A, B, bias, C, M, N, K);
    }
#else
    /* Scalar fallback */
    gemm_bf16_scalar(A, B, bias, C, M, N, K);
#endif
}

References C.

◆ gemm_blocked_serial_train_parallel_dispatch()

void gemm_blocked_serial_train_parallel_dispatch	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 167 of file ck_parallel_train.c.

                                                        {
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    const int nth = pool ? ck_threadpool_n_threads(pool) : 1;
 
    /* Keep small shapes on serial path to avoid dispatch overhead. */
    const size_t work = (size_t)M * (size_t)N * (size_t)K;
    if (!pool || nth <= 1 || work < (size_t)131072) {
        gemm_blocked_serial(A, B, bias, C, M, N, K);
        return;
    }
 
    int split_n = 0;
    int active_nth = nth;
    if (M == 1) {
        /* Tensor-parallel decode-style split only when each worker has enough columns. */
        const int cols_per_worker = (N + nth - 1) / nth;
        if (cols_per_worker >= 256 && K >= 512 && work >= (size_t)2097152) {
            split_n = 1;
        } else {
            gemm_blocked_serial(A, B, bias, C, M, N, K);
            return;
        }
    } else {
        /* Row split for prefill/train batches with coarser chunks. */
        active_nth = M / 2;
        if (active_nth > nth) {
            active_nth = nth;
        }
        if (active_nth <= 1) {
            gemm_blocked_serial(A, B, bias, C, M, N, K);
            return;
        }
    }
 
    ck_train_gemm_args_t args = {
        .A = A,
        .B = B,
        .bias = bias,
        .C = C,
        .M = M,
        .N = N,
        .K = K,
        .split_n = split_n,
    };
 
    ck_threadpool_dispatch_n(pool, active_nth, ck_train_gemm_work, &args);
}

References C, ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_train_gemm_work(), and gemm_blocked_serial().

◆ gemm_fine_grained_parallel()

void gemm_fine_grained_parallel	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 286 of file gemm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        gemm_naive_serial_float(A, B, bias, C, M, N, K);
        return;
    }
#if defined(__AVX512F__)
    const int block_size = 64;
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            C[i * N + j] = bias ? bias[j] : 0.0f;
        }
    }
#pragma omp parallel for collapse(3)
    for (int ii = 0; ii < M; ii += block_size) {
        for (int jj = 0; jj < N; jj += block_size) {
            for (int kk = 0; kk < K; kk += block_size) {
                int i_end = ck_min(ii + block_size, M);
                int j_end = ck_min(jj + block_size, N);
                int k_end = ck_min(kk + block_size, K);
 
                for (int i = ii; i < i_end; i++) {
                    for (int j = jj; j < j_end; j++) {
                        __m512 sum_vec = _mm512_setzero_ps();
                        int k;
                        for (k = kk; k <= k_end - 16; k += 16) {
                            __m512 a_vec = _mm512_loadu_ps(&A[i * K + k]);
                            __m512 b_vec = _mm512_loadu_ps(&B[j * K + k]);
                            sum_vec = _mm512_fmadd_ps(a_vec, b_vec, sum_vec);
                        }
                        float partial_sum = _mm512_reduce_add_ps(sum_vec);
                        for (; k < k_end; k++) {
                            partial_sum += A[i * K + k] * B[j * K + k];
                        }
#pragma omp atomic
                        C[i * N + j] += partial_sum;
                    }
                }
            }
        }
    }
#elif defined(__AVX__)
    // AVX1 cache-blocked version
    const int block_size = 32;  // Smaller block for L1 cache
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            C[i * N + j] = bias ? bias[j] : 0.0f;
        }
    }
#pragma omp parallel for collapse(3)
    for (int ii = 0; ii < M; ii += block_size) {
        for (int jj = 0; jj < N; jj += block_size) {
            for (int kk = 0; kk < K; kk += block_size) {
                int i_end = ck_min(ii + block_size, M);
                int j_end = ck_min(jj + block_size, N);
                int k_end = ck_min(kk + block_size, K);
 
                for (int i = ii; i < i_end; i++) {
                    for (int j = jj; j < j_end; j++) {
                        __m256 sum_vec = _mm256_setzero_ps();
                        int k;
                        for (k = kk; k <= k_end - 8; k += 8) {
                            __m256 a_vec = _mm256_loadu_ps(&A[i * K + k]);
                            __m256 b_vec = _mm256_loadu_ps(&B[j * K + k]);
                            __m256 prod = _mm256_mul_ps(a_vec, b_vec);
                            sum_vec = _mm256_add_ps(sum_vec, prod);
                        }
                        float partial_sum = hsum256_ps(sum_vec);
                        for (; k < k_end; k++) {
                            partial_sum += A[i * K + k] * B[j * K + k];
                        }
#pragma omp atomic
                        C[i * N + j] += partial_sum;
                    }
                }
            }
        }
    }
#else
    gemm_naive_parallel(A, B, bias, C, M, N, K);
#endif
}

References C, ck_min(), ck_strict_parity_enabled(), gemm_naive_parallel(), and gemm_naive_serial_float().

◆ gemm_microkernel()

void gemm_microkernel	(	const float *	A,
		const float *	B,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	B_transposed
	)

Definition at line 1134 of file gemm_microkernel.c.

{
    if (B_transposed) {
        gemm_microkernel_blocked_bt(A, B, C, M, N, K);
    } else {
        // Use packed version for large matrices
        if (M >= PACK_THRESHOLD && N >= PACK_THRESHOLD && K >= PACK_THRESHOLD) {
            gemm_microkernel_packed(A, B, C, M, N, K);
        } else {
            gemm_microkernel_blocked(A, B, C, M, N, K);
        }
    }
}

References C, gemm_microkernel_blocked(), gemm_microkernel_blocked_bt(), gemm_microkernel_packed(), and PACK_THRESHOLD.

Referenced by gemm_blocked_serial().

◆ gemm_microkernel_blocked()

void gemm_microkernel_blocked	(	const float *	A,
		const float *	B,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 934 of file gemm_microkernel.c.

{
    const int mr = MR;
    const int nr = NR;
 
    // Use sequential version for small matrices to avoid OpenMP overhead
    // Threshold tuned for typical 4-8 core systems
    if ((size_t)M * N * K <= 512ULL * 512 * 512) {
        gemm_microkernel_sequential(A, B, C, M, N, K);
        return;
    }
 
    // Initialize thread count to physical cores (once)
    gemm_init_threads();
 
    // Zero output first
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < M; i++) {
        memset(&C[i * N], 0, N * sizeof(float));
    }
 
    // Block over K (outermost - for accumulation across all threads)
    for (int k0 = 0; k0 < K; k0 += KC) {
        int kb = (k0 + KC <= K) ? KC : (K - k0);
        int first_k = (k0 == 0);
 
        // Parallelize over M rows - each thread gets a chunk of M
        // This gives better cache locality than tile-level parallelism
        #pragma omp parallel for schedule(static)
        for (int m0 = 0; m0 < M; m0 += mr) {
            int mr_actual = (m0 + mr <= M) ? mr : (M - m0);
 
            // Each thread processes all N tiles for its M rows
            for (int n0 = 0; n0 < N; n0 += nr) {
                int nr_actual = (n0 + nr <= N) ? nr : (N - n0);
 
                const float *A_tile = &A[m0 * K + k0];
                const float *B_tile = &B[k0 * N + n0];
                float *C_tile = &C[m0 * N + n0];
 
                if (mr_actual == mr && nr_actual == nr) {
#if defined(__AVX512F__)
                    gemm_microkernel_6x32_avx512(kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#elif defined(__FMA__)
                    gemm_microkernel_6x16_avx(kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#elif defined(__AVX__)
                    gemm_microkernel_4x16_avx(kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#else
                    gemm_microkernel_edge(mr_actual, nr_actual, kb, A_tile, K, B_tile, N, C_tile, N, first_k);
#endif
                } else {
                    gemm_microkernel_edge(mr_actual, nr_actual, kb, A_tile, K, B_tile, N, C_tile, N, first_k);
                }
            }
        }
    }
}

References C, gemm_init_threads(), gemm_microkernel_edge(), gemm_microkernel_sequential(), KC, MR, and NR.

Referenced by gemm_microkernel(), and gemm_microkernel_packed().

◆ gemm_microkernel_blocked_bt()

void gemm_microkernel_blocked_bt	(	const float *	A,
		const float *	B,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 1058 of file gemm_microkernel.c.

{
    // Zero output first
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < M; i++) {
        memset(&C[i * N], 0, N * sizeof(float));
    }
 
    const int mr = MR;
    const int nr = NR;
 
    #pragma omp parallel for schedule(dynamic) collapse(2)
    for (int m0 = 0; m0 < M; m0 += MC) {
        for (int n0 = 0; n0 < N; n0 += NC) {
            int mb = (m0 + MC <= M) ? MC : (M - m0);
            int nb = (n0 + NC <= N) ? NC : (N - n0);
 
            for (int k0 = 0; k0 < K; k0 += KC) {
                int kb = (k0 + KC <= K) ? KC : (K - k0);
                int first_k = (k0 == 0);
 
                for (int m1 = 0; m1 < mb; m1 += mr) {
                    int mr_actual = (m1 + mr <= mb) ? mr : (mb - m1);
 
                    for (int n1 = 0; n1 < nb; n1 += nr) {
                        int nr_actual = (n1 + nr <= nb) ? nr : (nb - n1);
 
                        const float *A_tile = &A[(m0 + m1) * K + k0];
                        const float *B_tile = &B[(n0 + n1) * K + k0];
                        float *C_tile = &C[(m0 + m1) * N + (n0 + n1)];
 
                        if (mr_actual == mr && nr_actual == nr) {
#if defined(__AVX512F__)
                            gemm_microkernel_6x32_bt_avx512(kb, A_tile, K, B_tile, K, C_tile, N, first_k);
#else
                            // Scalar fallback for B-transposed
                            for (int i = 0; i < mr; i++) {
                                for (int j = 0; j < nr; j++) {
                                    float sum = first_k ? 0.0f : C_tile[i * N + j];
                                    for (int kk = 0; kk < kb; kk++) {
                                        sum += A_tile[i * K + kk] * B_tile[j * K + kk];
                                    }
                                    C_tile[i * N + j] = sum;
                                }
                            }
#endif
                        } else {
                            // Edge case
                            for (int i = 0; i < mr_actual; i++) {
                                for (int j = 0; j < nr_actual; j++) {
                                    float sum = first_k ? 0.0f : C_tile[i * N + j];
                                    for (int kk = 0; kk < kb; kk++) {
                                        sum += A_tile[i * K + kk] * B_tile[j * K + kk];
                                    }
                                    C_tile[i * N + j] = sum;
                                }
                            }
                        }
                    }
                }
            }
        }
    }
}

References C, KC, MC, MR, NC, and NR.

Referenced by gemm_microkernel().

◆ gemm_microkernel_packed()

void gemm_microkernel_packed	(	const float *	A,
		const float *	B,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 840 of file gemm_microkernel.c.

{
    // Use tile-parallel blocked version - scales better on many-core systems
    gemm_microkernel_blocked(A, B, C, M, N, K);
}

References C, and gemm_microkernel_blocked().

Referenced by gemm_microkernel().

◆ gemm_naive_parallel()

void gemm_naive_parallel	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 206 of file gemm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        gemm_naive_serial_float(A, B, bias, C, M, N, K);
        return;
    }
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[j * K + k];
            }
            float bias_val = bias ? bias[j] : 0.0f;
            C[i * N + j] = sum + bias_val;
        }
    }
}

References C, ck_strict_parity_enabled(), and gemm_naive_serial_float().

Referenced by ck_attention_project_head_major_ref(), ck_mlp_swiglu_forward_ref(), ck_qkv_project_head_major_ref(), gemm_avx512_parallel(), and gemm_fine_grained_parallel().

◆ gemm_nn_avx512()

void gemm_nn_avx512	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 420 of file gemm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        gemm_nn_serial_double(A, B, bias, C, M, N, K);
        return;
    }
#if defined(__AVX512F__)
    // For gemm_nn, we can't vectorize over K easily since B[k,j] has stride N.
    // Instead, vectorize over N (output columns) when N >= 16.
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        int j = 0;
        // Process 16 output columns at a time
        for (; j <= N - 16; j += 16) {
            __m512 sum_vec = bias ? _mm512_loadu_ps(&bias[j]) : _mm512_setzero_ps();
            for (int k = 0; k < K; k++) {
                __m512 a_broadcast = _mm512_set1_ps(A[i * K + k]);
                __m512 b_vec = _mm512_loadu_ps(&B[k * N + j]);
                sum_vec = _mm512_fmadd_ps(a_broadcast, b_vec, sum_vec);
            }
            _mm512_storeu_ps(&C[i * N + j], sum_vec);
        }
        // Handle remaining columns
        for (; j < N; j++) {
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[k * N + j];
            }
            C[i * N + j] = sum;
        }
    }
#elif defined(__AVX__)
    // AVX1: vectorize over N (8 columns at a time)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        int j = 0;
        for (; j <= N - 8; j += 8) {
            __m256 sum_vec = bias ? _mm256_loadu_ps(&bias[j]) : _mm256_setzero_ps();
            for (int k = 0; k < K; k++) {
                __m256 a_broadcast = _mm256_set1_ps(A[i * K + k]);
                __m256 b_vec = _mm256_loadu_ps(&B[k * N + j]);
                __m256 prod = _mm256_mul_ps(a_broadcast, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, prod);
            }
            _mm256_storeu_ps(&C[i * N + j], sum_vec);
        }
        for (; j < N; j++) {
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[k * N + j];
            }
            C[i * N + j] = sum;
        }
    }
#else
    gemm_nn_parallel(A, B, bias, C, M, N, K);
#endif
}

References C, ck_strict_parity_enabled(), gemm_nn_parallel(), and gemm_nn_serial_double().

Referenced by gemm_nn_simd().

◆ gemm_nn_blocked()

void gemm_nn_blocked	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 590 of file gemm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        gemm_nn_serial_double(A, B, bias, C, M, N, K);
        return;
    }
#if defined(__AVX512F__)
    const int block_size = 64;
#elif defined(__AVX__)
    const int block_size = 32;
#else
    const int block_size = 32;
#endif
    // Initialize C with bias (parallelized)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            C[i * N + j] = bias ? bias[j] : 0.0f;
        }
    }
    // Blocked multiply-accumulate (parallelized over M blocks)
#pragma omp parallel for
    for (int ii = 0; ii < M; ii += block_size) {
        for (int kk = 0; kk < K; kk += block_size) {
            for (int jj = 0; jj < N; jj += block_size) {
                int i_end = ck_min(ii + block_size, M);
                int k_end = ck_min(kk + block_size, K);
                int j_end = ck_min(jj + block_size, N);
 
                for (int i = ii; i < i_end; i++) {
                    for (int k = kk; k < k_end; k++) {
                        float a_val = A[i * K + k];
#if defined(__AVX512F__)
                        __m512 a_broadcast = _mm512_set1_ps(a_val);
                        int j;
                        for (j = jj; j <= j_end - 16; j += 16) {
                            __m512 b_vec = _mm512_loadu_ps(&B[k * N + j]);
                            __m512 c_vec = _mm512_loadu_ps(&C[i * N + j]);
                            c_vec = _mm512_fmadd_ps(a_broadcast, b_vec, c_vec);
                            _mm512_storeu_ps(&C[i * N + j], c_vec);
                        }
                        for (; j < j_end; j++) {
                            C[i * N + j] += a_val * B[k * N + j];
                        }
#elif defined(__AVX__)
                        __m256 a_broadcast = _mm256_set1_ps(a_val);
                        int j;
                        for (j = jj; j <= j_end - 8; j += 8) {
                            __m256 b_vec = _mm256_loadu_ps(&B[k * N + j]);
                            __m256 c_vec = _mm256_loadu_ps(&C[i * N + j]);
                            __m256 prod = _mm256_mul_ps(a_broadcast, b_vec);
                            c_vec = _mm256_add_ps(c_vec, prod);
                            _mm256_storeu_ps(&C[i * N + j], c_vec);
                        }
                        for (; j < j_end; j++) {
                            C[i * N + j] += a_val * B[k * N + j];
                        }
#else
                        for (int j = jj; j < j_end; j++) {
                            C[i * N + j] += a_val * B[k * N + j];
                        }
#endif
                    }
                }
            }
        }
    }
}

References C, ck_min(), ck_strict_parity_enabled(), and gemm_nn_serial_double().

◆ gemm_nn_parallel()

void gemm_nn_parallel	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 398 of file gemm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        gemm_nn_serial_double(A, B, bias, C, M, N, K);
        return;
    }
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[k * N + j];
            }
            C[i * N + j] = sum;
        }
    }
}

References C, ck_strict_parity_enabled(), and gemm_nn_serial_double().

Referenced by gemm_nn_avx512(), and gemm_nn_avx512_probe().

◆ gemm_nn_simd()

void gemm_nn_simd	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 577 of file gemm_kernels.c.

{
    if (ck_gemm_nn_impl_probe_enabled()) {
        gemm_nn_avx512_probe(A, B, bias, C, M, N, K);
        return;
    }
    gemm_nn_avx512(A, B, bias, C, M, N, K);
}

References C, ck_gemm_nn_impl_probe_enabled(), gemm_nn_avx512(), and gemm_nn_avx512_probe().

Referenced by fc1_backward_kernel(), fc2_backward_kernel(), gemm_backward_f32_train_parallel_dispatch(), and gemm_backward_f32_train_parallel_dispatch_v2().

◆ gemm_nt_bf16()

void gemm_nt_bf16	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 748 of file gemm_kernels_bf16.c.

{
    if (M <= 0) return;
    gemm_nt_bf16_row_range(A, B, bias, C, M, N, K, 0, M);
}

References C, and gemm_nt_bf16_row_range().

Referenced by gemm_nt_bf16_parallel_dispatch().

◆ gemm_nt_bf16_amx_bf16_storage()

void gemm_nt_bf16_amx_bf16_storage	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 1311 of file gemm_kernels_bf16.c.

{
    size_t input_count = 0;
    if (M > 0 && K > 0 && (size_t)M <= SIZE_MAX / (size_t)K) {
        input_count = (size_t)M * (size_t)K;
    }
    uint16_t *workspace = input_count > 0 && input_count <= SIZE_MAX / sizeof(uint16_t)
        ? (uint16_t *)malloc(input_count * sizeof(uint16_t))
        : NULL;
    if (!workspace) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: AMX BF16 compatibility workspace allocation failed\n");
        abort();
    }
    gemm_nt_bf16_amx_bf16_storage_workspace(
        A, B, bias, C, M, N, K, workspace, input_count * sizeof(uint16_t));
    free(workspace);
}

References C, and gemm_nt_bf16_amx_bf16_storage_workspace().

Referenced by gemm_nt_bf16_prefill_shape_safe_bf16_storage().

◆ gemm_nt_bf16_amx_bf16_storage_workspace()

void gemm_nt_bf16_amx_bf16_storage_workspace	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		uint16_t *	a_bf16,
		size_t	a_bf16_bytes
	)

Definition at line 1258 of file gemm_kernels_bf16.c.

{
#if HAVE_AMX_BF16
    if (!A || !B || !C || M < 16 || N < 16 || K < 32 ||
        (M % 16) != 0 || (N % 16) != 0 || (K % 32) != 0) {
        fprintf(stderr,
            "HARD KERNEL CONTRACT FAULT: AMX BF16 GEMM requires non-null buffers "
            "and M%%16=N%%16=K%%32=0 (M=%d N=%d K=%d)\n",
            M, N, K);
        abort();
    }
    const size_t input_count = (size_t)M * K;
    if (!a_bf16 || input_count > SIZE_MAX / sizeof(uint16_t) ||
        a_bf16_bytes < input_count * sizeof(uint16_t)) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: AMX BF16 activation workspace is too small\n");
        abort();
    }
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > 24) active = 24;
    ck_bf16_convert_args_t convert = {.src=A, .dst=a_bf16, .count=input_count};
    if (pool && active > 1) ck_threadpool_dispatch_n(pool, active, ck_bf16_convert_work, &convert);
    else ck_bf16_convert_work(0, 1, &convert);
    ck_gemm_bf16_amx_args_t gemm = {
        .A=a_bf16, .B=(const uint16_t *)B, .bias=bias, .C=C,
        .M=M, .N=N, .K=K, .failed=0
    };
    if (pool && active > 1) ck_threadpool_dispatch_n(pool, active, ck_gemm_bf16_amx_work, &gemm);
    else ck_gemm_bf16_amx_work(0, 1, &gemm);
    if (__atomic_load_n(&gemm.failed, __ATOMIC_RELAXED)) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: AMX tile permission request failed\n");
        abort();
    }
    ck_bf16_round_args_t round = {.values=C, .count=(size_t)M * N};
    if (pool && active > 1) ck_threadpool_dispatch_n(pool, active, ck_bf16_round_work, &round);
    else ck_bf16_round_work(0, 1, &round);
    return;
#else
    (void)A; (void)B; (void)bias; (void)C; (void)M; (void)N; (void)K;
    (void)a_bf16; (void)a_bf16_bytes;
    fprintf(stderr,
            "HARD KERNEL CONTRACT FAULT: gemm_nt_bf16_amx_bf16_storage was selected "
            "without AMX BF16 support\n");
    abort();
#endif
}

References C, ck_bf16_convert_work(), ck_bf16_round_work(), ck_gemm_bf16_amx_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by gemm_nt_bf16_amx_bf16_storage(), and gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace().

◆ gemm_nt_bf16_bf16_storage()

void gemm_nt_bf16_bf16_storage	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 2034 of file gemm_kernels_bf16.c.

{
    gemm_nt_bf16_bf16_storage_parallel_dispatch(A, B, bias, C, M, N, K);
}

References C, and gemm_nt_bf16_bf16_storage_parallel_dispatch().

◆ gemm_nt_bf16_bf16_storage_parallel_dispatch()

void gemm_nt_bf16_bf16_storage_parallel_dispatch	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 2010 of file gemm_kernels_bf16.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    if (!pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)N <= 4096) {
        gemm_nt_bf16_bf16_storage_row_range(A, B, bias, C, M, N, K, 0, M);
        return;
    }
 
    ck_gemm_nt_bf16_exact_args_t args = {
        .A = A, .B = B, .bias = bias, .C = C, .M = M, .N = N, .K = K,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemm_nt_bf16_storage_exact_rows, &args);
}

References C, ck_gemm_nt_bf16_storage_exact_rows(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and gemm_nt_bf16_bf16_storage_row_range().

Referenced by gemm_nt_bf16_bf16_storage().

◆ gemm_nt_bf16_bf16_storage_row_range()

void gemm_nt_bf16_bf16_storage_row_range	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	row_begin,
		int	row_end
	)

Definition at line 1985 of file gemm_kernels_bf16.c.

{
    gemm_nt_bf16_row_range(A, B, bias, C, M, N, K, row_begin, row_end);
    for (int row = row_begin; row < row_end; ++row) {
        float *dst = C + (size_t)row * (size_t)N;
        for (int col = 0; col < N; ++col) {
            dst[col] = bf16_to_float(float_to_bf16(dst[col]));
        }
    }
}

References bf16_to_float(), C, float_to_bf16(), and gemm_nt_bf16_row_range().

Referenced by ck_gemm_nt_bf16_storage_exact_rows(), and gemm_nt_bf16_bf16_storage_parallel_dispatch().

◆ gemm_nt_bf16_native_bf16_storage()

void gemm_nt_bf16_native_bf16_storage	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 1071 of file gemm_kernels_bf16.c.

{
    const uint16_t *weights = (const uint16_t *)B;
    if (!A || !weights || !C || M <= 0 || N <= 0 || K <= 0) return;
 
    ck_gemm_bf16_native_args_t args = {
        .A = A, .B = weights, .bias = bias, .C = C, .M = M, .N = N, .K = K
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > M) active = M;
    if (active > 24) active = 24;
    if (!pool || active <= 1 || (size_t)M * (size_t)N <= 4096) {
        ck_gemm_bf16_native_work(0, 1, &args);
        return;
    }
    ck_threadpool_dispatch_n(pool, active, ck_gemm_bf16_native_work, &args);
}

References C, ck_gemm_bf16_native_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), and ck_threadpool_n_threads().

Referenced by gemm_nt_bf16_prefill_shape_safe_bf16_storage(), and gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace().

◆ gemm_nt_bf16_parallel_dispatch()

void gemm_nt_bf16_parallel_dispatch	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 777 of file gemm_kernels_bf16.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    const char *disabled = getenv("CK_DISABLE_BF16_GEMM_PARALLEL_PREFILL");
    if ((disabled && disabled[0] && strcmp(disabled, "0") != 0) ||
        !pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)N <= 4096) {
        gemm_nt_bf16(A, B, bias, C, M, N, K);
        return;
    }
 
    ck_gemm_nt_bf16_exact_args_t args = {
        .A = A, .B = B, .bias = bias, .C = C, .M = M, .N = N, .K = K,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemm_nt_bf16_exact_rows, &args);
}

References C, ck_gemm_nt_bf16_exact_rows(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and gemm_nt_bf16().

◆ gemm_nt_bf16_prefill_shape_safe_bf16_storage()

void gemm_nt_bf16_prefill_shape_safe_bf16_storage	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 1347 of file gemm_kernels_bf16.c.

{
    const int amx_shape = M >= 16 && N >= 16 && K >= 32 &&
                          (M % 16) == 0 && (N % 16) == 0 && (K % 32) == 0;
    if (amx_shape && ck_gemm_bf16_amx_available()) {
        gemm_nt_bf16_amx_bf16_storage(A, B, bias, C, M, N, K);
        return;
    }
    gemm_nt_bf16_native_bf16_storage(A, B, bias, C, M, N, K);
}

References C, ck_gemm_bf16_amx_available(), gemm_nt_bf16_amx_bf16_storage(), and gemm_nt_bf16_native_bf16_storage().

◆ gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace()

void gemm_nt_bf16_prefill_shape_safe_bf16_storage_workspace	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		uint16_t *	a_bf16,
		size_t	a_bf16_bytes
	)

Definition at line 1333 of file gemm_kernels_bf16.c.

{
    const int amx_shape = M >= 16 && N >= 16 && K >= 32 &&
                          (M % 16) == 0 && (N % 16) == 0 && (K % 32) == 0;
    if (amx_shape && ck_gemm_bf16_amx_available()) {
        gemm_nt_bf16_amx_bf16_storage_workspace(
            A, B, bias, C, M, N, K, a_bf16, a_bf16_bytes);
        return;
    }
    gemm_nt_bf16_native_bf16_storage(A, B, bias, C, M, N, K);
}

References C, ck_gemm_bf16_amx_available(), gemm_nt_bf16_amx_bf16_storage_workspace(), and gemm_nt_bf16_native_bf16_storage().

◆ gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage()

void gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 1551 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    ck_pytorch_brgemm_require_version(
        3, 12, 0, "80afa71049cd69a3df32adcccb623b12cd7baa22",
        "gemm_nt_bf16_pytorch_onednn_3_12_brgemm_bf16_storage", M, N, K);
#endif
    gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl(
        A, B, bias, C, M, N, K);
}

References C, and gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl().

◆ gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage()

void gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 1536 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    ck_pytorch_brgemm_require_version(
        3, 7, 1, "8d263e693366ef8db40acc569cc7d8edf644556d",
        "gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage", M, N, K);
#endif
    gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl(
        A, B, bias, C, M, N, K);
}

References C, and gemm_nt_bf16_pytorch_onednn_brgemm_bf16_storage_impl().

◆ gemm_nt_bf16_row_range()

void gemm_nt_bf16_row_range	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	row_begin,
		int	row_end
	)

Definition at line 720 of file gemm_kernels_bf16.c.

{
    const uint16_t *w = (const uint16_t *)B;
    if (!A || !w || !C || M <= 0 || N <= 0 || K <= 0 ||
        row_begin < 0 || row_begin >= row_end || row_end > M) {
        return;
    }
 
    for (int i = row_begin; i < row_end; ++i) {
        const float *a_row = A + (size_t)i * (size_t)K;
        float *c_row = C + (size_t)i * (size_t)N;
        for (int j = 0; j < N; ++j) {
            const uint16_t *w_row = w + (size_t)j * (size_t)K;
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; ++k) {
                const float ab = bf16_to_float(float_to_bf16(a_row[k]));
                sum += ab * bf16_to_float(w_row[k]);
            }
            c_row[j] = sum;
        }
    }
}

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

Referenced by ck_gemm_nt_bf16_exact_rows(), gemm_nt_bf16(), and gemm_nt_bf16_bf16_storage_row_range().

◆ gemm_nt_f16()

void gemm_nt_f16	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

NT GEMM wrapper for FP16 weights with the engine's standard ABI.

Contract: A: [M, K] fp32 activation matrix B: [N, K] fp16 weight matrix stored row-major (transposed layout) C: [M, N] fp32 output matrix

This wrapper follows llama.cpp's CPU F16 mul_mat contract: activation rows are rounded to FP16 first, then the dot runs as F16 x F16 with FP32 output accumulation. The lower-level gemm_f16() helper remains the direct F16-weight x FP32-activation operator for generic use.

Definition at line 970 of file gemm_kernels_f16.c.

{
    if (ck_strict_parity_enabled() &&
        gemm_nt_f16_ggml_strict(A, B, bias, C, M, N, K)) {
        return;
    }
 
    if (!gemm_f16_input_fp16_threadpool(C, (const uint16_t *)B, A, N, M, K)) {
        gemm_f16_input_fp16_ref(C, (const uint16_t *)B, A, N, M, K);
    }
 
    if (!bias) {
        return;
    }
 
#pragma omp parallel for schedule(static) if(M > 1)
    for (int i = 0; i < M; ++i) {
        float *c_row = C + (size_t)i * (size_t)N;
        for (int j = 0; j < N; ++j) {
            c_row[j] += bias[j];
        }
    }
}

References C, ck_strict_parity_enabled(), gemm_f16_input_fp16_ref(), gemm_f16_input_fp16_threadpool(), and gemm_nt_f16_ggml_strict().

◆ gemm_nt_f16_clipped()

void gemm_nt_f16_clipped	(	const float *	A,
		const void *	B,
		const float *	bias,
		const float *	input_min,
		const float *	input_max,
		const float *	output_min,
		const float *	output_max,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 998 of file gemm_kernels_f16.c.

{
    const float in_min = input_min ? input_min[0] : -3.4028234663852886e38f;
    const float in_max = input_max ? input_max[0] : 3.4028234663852886e38f;
    const float out_min = output_min ? output_min[0] : -3.4028234663852886e38f;
    const float out_max = output_max ? output_max[0] : 3.4028234663852886e38f;
    const uint16_t *W = (const uint16_t *)B;
 
#pragma omp parallel for schedule(static) if(M > 1)
    for (int m = 0; m < M; ++m) {
        const float *a_row = A + (size_t)m * (size_t)K;
        uint16_t a_f16[K];
 
        for (int k = 0; k < K; ++k) {
            float x = a_row[k];
            if (x < in_min) x = in_min;
            if (x > in_max) x = in_max;
            a_f16[k] = fp32_to_fp16(x);
        }
 
        float *c_row = C + (size_t)m * (size_t)N;
        for (int n = 0; n < N; ++n) {
            const uint16_t *w_row = W + (size_t)n * (size_t)K;
            float sum = bias ? bias[n] : 0.0f;
            for (int k = 0; k < K; ++k) {
                sum += fp16_to_fp32(w_row[k]) * fp16_to_fp32(a_f16[k]);
            }
            if (sum < out_min) sum = out_min;
            if (sum > out_max) sum = out_max;
            c_row[n] = sum;
        }
    }
}

References C, fp16_to_fp32, and fp32_to_fp16.

◆ gemm_nt_f32_llama_production()

void gemm_nt_f32_llama_production	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 1062 of file gemm_kernels.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
#pragma omp parallel for schedule(static) if ((size_t)M * (size_t)N >= 96)
    for (int index = 0; index < M * N; ++index) {
        ck_gemm_nt_f32_llama_production_output(A, B, bias, C, M, N, K, index);
    }
}

References C, and ck_gemm_nt_f32_llama_production_output().

Referenced by ck_moe_shared_q4k_gated_workspace(), and moe_swiglu_shared_forward_q8_0_gated_workspace().

◆ gemm_nt_f32_llama_production_output_range()

void gemm_nt_f32_llama_production_output_range	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	output_begin,
		int	output_end
	)

Definition at line 1049 of file gemm_kernels.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) return;
    const int total = M * N;
    if (output_begin < 0) output_begin = 0;
    if (output_end > total) output_end = total;
    for (int index = output_begin; index < output_end; ++index) {
        ck_gemm_nt_f32_llama_production_output(A, B, bias, C, M, N, K, index);
    }
}

References C, and ck_gemm_nt_f32_llama_production_output().

◆ gemm_nt_fp32_exact_parallel_dispatch()

void gemm_nt_fp32_exact_parallel_dispatch	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 179 of file gemm_kernels.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) return;
    ck_gemm_nt_fp32_exact_args_t args = {
        .A = A, .B = B, .bias = bias, .C = C,
        .M = M, .N = N, .K = K,
        .bias_before_reduction = ck_strict_parity_enabled(),
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    if (!pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)N <= 4096) {
        ck_gemm_nt_fp32_exact_rows(0, M, &args);
        return;
    }
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemm_nt_fp32_exact_rows, &args);
}

References C, ck_gemm_nt_fp32_exact_rows(), ck_strict_parity_enabled(), ck_threadpool_global(), ck_threadpool_n_threads(), and ck_threadpool_parallel_for_n().

◆ gemm_nt_q4_0()

void gemm_nt_q4_0	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Matrix-matrix multiply: C[M,N] = A[M,K] @ B[N,K]^T + bias.

Parameters

A	Input matrix [M x K], row-major FP32
B	Weight matrix in Q4_0 format, [N x K] stored row-major
bias	Optional bias [N], NULL if not used
C	Output [M x N], row-major FP32
M	Batch size (number of tokens)
N	Output dimension (number of rows in B)
K	Input dimension

Definition at line 176 of file gemm_kernels_q4_0.c.

{
    const block_q4_0 *blocks = (const block_q4_0 *)B;
    const int blocks_per_row = K / QK4_0;
 
    for (int m = 0; m < M; m++) {
        const float *a_row = &A[m * K];
 
        for (int n = 0; n < N; n++) {
            float sum = 0.0f;
 
            for (int b = 0; b < blocks_per_row; b++) {
                const block_q4_0 *block = &blocks[n * blocks_per_row + b];
                const float d = CK_FP16_TO_FP32(block->d);
                const float *ap = &a_row[b * QK4_0];
 
                for (int i = 0; i < QK4_0 / 2; i++) {
                    const uint8_t packed = block->qs[i];
                    const int q0 = (packed & 0x0F) - 8;
                    const int q1 = (packed >> 4) - 8;
 
                    sum += d * (float)q0 * ap[2 * i + 0];
                    sum += d * (float)q1 * ap[2 * i + 1];
                }
            }
 
            C[m * N + n] = sum + (bias ? bias[n] : 0.0f);
        }
    }
}

References C, CK_FP16_TO_FP32, block_q4_0::d, QK4_0, and block_q4_0::qs.

Referenced by ck_gemm_nt_quant().

◆ gemm_nt_q4_1()

void gemm_nt_q4_1	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

GEMM with transposed Q4_1 weights: C = A @ B^T.

Parameters

A	Input activations [M x K], row-major FP32
B	Weight matrix in Q4_1 format [N x K], row-major quantized
bias	Optional bias [N], NULL if not used
C	Output [M x N], row-major FP32
M	Batch size (number of tokens)
N	Output dimension
K	Input dimension

Definition at line 256 of file gemm_kernels_q4_1.c.

{
    const block_q4_1 *blocks = (const block_q4_1 *)B;
    const int blocks_per_row = K / QK4_1;
 
    for (int m = 0; m < M; m++) {
        const float *a_row = &A[m * K];
 
        for (int n = 0; n < N; n++) {
            float sum = 0.0f;
 
            for (int b = 0; b < blocks_per_row; b++) {
                const block_q4_1 *block = &blocks[n * blocks_per_row + b];
                const float d = CK_FP16_TO_FP32(block->d);
                const float min = CK_FP16_TO_FP32(block->m);
                const float *ap = &a_row[b * QK4_1];
 
                for (int i = 0; i < QK4_1 / 2; i++) {
                    const uint8_t packed = block->qs[i];
                    const int q0 = (packed & 0x0F);
                    const int q1 = (packed >> 4);
 
                    const float w0 = d * (float)q0 + min;
                    const float w1 = d * (float)q1 + min;
 
                    sum += w0 * ap[2 * i + 0];
                    sum += w1 * ap[2 * i + 1];
                }
            }
 
            C[m * N + n] = sum + (bias ? bias[n] : 0.0f);
        }
    }
}

References C, CK_FP16_TO_FP32, block_q4_1::d, block_q4_1::m, QK4_1, and block_q4_1::qs.

Referenced by ck_gemm_nt_quant().

◆ gemm_nt_q4_k()

void gemm_nt_q4_k	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 708 of file gemm_kernels_q4k.c.

{
    if (!A || !B || !C) {
        return;
    }
    if (M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    /* gemm_q4_k produces Y as [batch x M_out]. Here:
     *   batch = M (tokens)
     *   M_out = N (output channels) */
    gemm_q4_k(C, B, A, /*M_out=*/N, /*N_batch=*/M, K);
 
    if (!bias) {
        return;
    }
 
    for (int i = 0; i < M; ++i) {
        float *row = C + (size_t)i * (size_t)N;
        for (int j = 0; j < N; ++j) {
            row[j] += bias[j];
        }
    }
}

References C, and gemm_q4_k().

◆ gemm_nt_q4_k_q8_k()

void gemm_nt_q4_k_q8_k	(	const void *	A_q8,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 397 of file gemm_kernels_q4k_q8k.c.

{
    if (!A_q8 || !B || !C) {
        return;
    }
    if (M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    gemm_q4_k_q8_k(C, B, A_q8, /*M_out=*/N, /*N_batch=*/M, K);
 
    if (!bias) {
        return;
    }
 
    for (int i = 0; i < M; ++i) {
        float *row = C + (size_t)i * (size_t)N;
        for (int j = 0; j < N; ++j) {
            row[j] += bias[j];
        }
    }
}

References C, and gemm_q4_k_q8_k().

Referenced by ck_attention_project_head_major_q4_k_q8_k(), ck_layer_forward_rmsnorm_swiglu_decode_q4_k(), ck_mlp_swiglu_forward_q4_k_q8_k(), ck_mlp_swiglu_forward_q4_k_q8_k_prefill(), ck_qkv_project_head_major_token_q4_k_q8_k(), gemm_nt_q8_k_mlp_dispatch(), gemm_nt_q8_k_qkv_dispatch(), model_forward_prefill_impl(), model_forward_prefill_impl(), model_forward_prefill_impl(), model_forward_prefill_impl(), qwen2_0_5b_decode_forward_prefill_impl(), qwen2_0_5b_decode_forward_prefill_impl(), qwen2_0_5b_decode_forward_prefill_impl(), and qwen2_0_5b_decode_forward_prefill_impl().

◆ gemm_nt_q5_0()

void gemm_nt_q5_0	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 833 of file gemm_kernels_q5_0.c.

{
    /* For decode (M=1), use direct GEMV which has AVX optimization */
    if (M == 1) {
        /* gemm_q5_0 expects column-major output, but we need row-major
         * So we call gemv_q5_0 directly for each output element */
        gemv_q5_0(C, B, A, N, K);
        if (bias) {
            for (int n = 0; n < N; n++) {
                C[n] += bias[n];
            }
        }
        return;
    }
 
    /* For prefill (M>1), use GEMM which dispatches to GEMV with AVX/AVX512 */
    /* gemm_q5_0 produces Y as [batch x M_out]. Here:
     *   batch = M (tokens)
     *   M_out = N (output channels) */
    gemm_q5_0(C, B, A, /*M_out=*/N, /*N_batch=*/M, K);
 
    if (bias) {
        for (int m = 0; m < M; m++) {
            float *row = C + (size_t)m * (size_t)N;
            for (int n = 0; n < N; n++) {
                row[n] += bias[n];
            }
        }
    }
}

References C, gemm_q5_0(), and gemv_q5_0().

◆ gemm_nt_q5_1()

void gemm_nt_q5_1	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

GEMM with transposed Q5_1 weights: C = A @ B^T.

Parameters

A	Input activations [M x K], row-major FP32
B	Weight matrix in Q5_1 format [N x K], row-major quantized
bias	Optional bias [N], NULL if not used
C	Output [M x N], row-major FP32
M	Batch size (number of tokens)
N	Output dimension
K	Input dimension

Definition at line 314 of file gemm_kernels_q5_1.c.

{
    const block_q5_1 *blocks = (const block_q5_1 *)B;
    const int blocks_per_row = K / QK5_1;
 
    for (int m = 0; m < M; m++) {
        const float *a_row = &A[m * K];
 
        for (int n = 0; n < N; n++) {
            float sum = 0.0f;
 
            for (int b = 0; b < blocks_per_row; b++) {
                const block_q5_1 *block = &blocks[n * blocks_per_row + b];
                const float d = CK_FP16_TO_FP32(block->d);
                const float min = CK_FP16_TO_FP32(block->m);
                const float *ap = &a_row[b * QK5_1];
 
                uint32_t qh;
                memcpy(&qh, block->qh, sizeof(qh));
 
                for (int i = 0; i < QK5_1 / 2; i++) {
                    const uint8_t packed = block->qs[i];
                    const int lo0 = (packed & 0x0F);
                    const int lo1 = (packed >> 4);
                    const int hi0 = ((qh >> (2 * i + 0)) & 1) << 4;
                    const int hi1 = ((qh >> (2 * i + 1)) & 1) << 4;
                    const int q0 = lo0 | hi0;
                    const int q1 = lo1 | hi1;
 
                    const float w0 = d * (float)q0 + min;
                    const float w1 = d * (float)q1 + min;
 
                    sum += w0 * ap[2 * i + 0];
                    sum += w1 * ap[2 * i + 1];
                }
            }
 
            C[m * N + n] = sum + (bias ? bias[n] : 0.0f);
        }
    }
}

References C, CK_FP16_TO_FP32, block_q5_1::d, block_q5_1::m, block_q5_1::qh, QK5_1, and block_q5_1::qs.

Referenced by ck_gemm_nt_quant().

◆ gemm_nt_q5_1_q8_1()

void gemm_nt_q5_1_q8_1	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 289 of file gemm_kernels_q5_1_q8_1.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0 || (K % QK5_1) != 0) {
        return;
    }
 
    const int blocks_per_row = K / QK5_1;
    if (blocks_per_row > CK_Q51_STACK_Q8_BLOCKS) {
        return;
    }
 
    const block_q5_1 *W = (const block_q5_1 *)B;
 
    for (int m = 0; m < M; ++m) {
        block_q8_1 a_q8[CK_Q51_STACK_Q8_BLOCKS];
        quantize_row_q8_1_scalar(&A[m * K], a_q8, K);
        float *c_row = &C[(size_t)m * (size_t)N];
 
        int n = 0;
        for (; n + 7 < N; n += 8) {
            const block_q5_1 *w0 = &W[(size_t)(n + 0) * (size_t)blocks_per_row];
            const block_q5_1 *w1 = &W[(size_t)(n + 1) * (size_t)blocks_per_row];
            const block_q5_1 *w2 = &W[(size_t)(n + 2) * (size_t)blocks_per_row];
            const block_q5_1 *w3 = &W[(size_t)(n + 3) * (size_t)blocks_per_row];
            const block_q5_1 *w4 = &W[(size_t)(n + 4) * (size_t)blocks_per_row];
            const block_q5_1 *w5 = &W[(size_t)(n + 5) * (size_t)blocks_per_row];
            const block_q5_1 *w6 = &W[(size_t)(n + 6) * (size_t)blocks_per_row];
            const block_q5_1 *w7 = &W[(size_t)(n + 7) * (size_t)blocks_per_row];
            float s0 = 0.0f;
            float s1 = 0.0f;
            float s2 = 0.0f;
            float s3 = 0.0f;
            float s4 = 0.0f;
            float s5 = 0.0f;
            float s6 = 0.0f;
            float s7 = 0.0f;
 
            for (int b = 0; b < blocks_per_row; ++b) {
                const block_q8_1 *x = &a_q8[b];
                s0 += dot_q5_1_q8_1_block(&w0[b], x);
                s1 += dot_q5_1_q8_1_block(&w1[b], x);
                s2 += dot_q5_1_q8_1_block(&w2[b], x);
                s3 += dot_q5_1_q8_1_block(&w3[b], x);
                s4 += dot_q5_1_q8_1_block(&w4[b], x);
                s5 += dot_q5_1_q8_1_block(&w5[b], x);
                s6 += dot_q5_1_q8_1_block(&w6[b], x);
                s7 += dot_q5_1_q8_1_block(&w7[b], x);
            }
 
            c_row[n + 0] = s0 + (bias ? bias[n + 0] : 0.0f);
            c_row[n + 1] = s1 + (bias ? bias[n + 1] : 0.0f);
            c_row[n + 2] = s2 + (bias ? bias[n + 2] : 0.0f);
            c_row[n + 3] = s3 + (bias ? bias[n + 3] : 0.0f);
            c_row[n + 4] = s4 + (bias ? bias[n + 4] : 0.0f);
            c_row[n + 5] = s5 + (bias ? bias[n + 5] : 0.0f);
            c_row[n + 6] = s6 + (bias ? bias[n + 6] : 0.0f);
            c_row[n + 7] = s7 + (bias ? bias[n + 7] : 0.0f);
        }
 
        for (; n < N; ++n) {
            const block_q5_1 *w_row = &W[(size_t)n * (size_t)blocks_per_row];
            float sum = 0.0f;
            for (int b = 0; b < blocks_per_row; ++b) {
                sum += dot_q5_1_q8_1_block(&w_row[b], &a_q8[b]);
            }
            c_row[n] = sum + (bias ? bias[n] : 0.0f);
        }
    }
}

References C, CK_Q51_STACK_Q8_BLOCKS, dot_q5_1_q8_1_block(), QK5_1, and quantize_row_q8_1_scalar().

Referenced by gemm_nt_q5_1_q8_1_m4().

◆ gemm_nt_q5_1_q8_1_ref()

void gemm_nt_q5_1_q8_1_ref	(	const void *	A_q8,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 242 of file gemm_kernels_q5_1_q8_1.c.

{
    if (!A_q8 || !B || !C || M <= 0 || N <= 0 || K <= 0 || (K % QK5_1) != 0) {
        return;
    }
 
    const block_q8_1 *A = (const block_q8_1 *)A_q8;
    const block_q5_1 *W = (const block_q5_1 *)B;
    const int blocks_per_row = K / QK5_1;
 
    for (int m = 0; m < M; ++m) {
        const block_q8_1 *a_row = &A[m * blocks_per_row];
        for (int n = 0; n < N; ++n) {
            const block_q5_1 *w_row = &W[n * blocks_per_row];
            float sum = 0.0f;
            for (int b = 0; b < blocks_per_row; ++b) {
                sum += dot_q5_1_q8_1_block(&w_row[b], &a_row[b]);
            }
            C[m * N + n] = sum + (bias ? bias[n] : 0.0f);
        }
    }
}

References C, dot_q5_1_q8_1_block(), and QK5_1.

◆ gemm_nt_q5_k()

void gemm_nt_q5_k	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 1001 of file gemm_kernels_q5_k.c.

{
#if defined(__AVX512F__)
    /* TODO: AVX-512 implementation */
    gemm_nt_q5_k_ref(A, B, bias, C, M, N, K);
#elif defined(__AVX2__)
    /* TODO: AVX-2 implementation */
    gemm_nt_q5_k_ref(A, B, bias, C, M, N, K);
#elif defined(__AVX__)
    /* TODO: AVX implementation */
    gemm_nt_q5_k_ref(A, B, bias, C, M, N, K);
#elif defined(__SSE4_1__)
    /* TODO: SSE4.1 implementation */
    gemm_nt_q5_k_ref(A, B, bias, C, M, N, K);
#else
    gemm_nt_q5_k_ref(A, B, bias, C, M, N, K);
#endif
}

References C, and gemm_nt_q5_k_ref().

◆ gemm_nt_q5_k_q8_k()

void gemm_nt_q5_k_q8_k	(	const void *	A_q8,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 914 of file gemm_kernels_q5_k.c.

{
#if defined(__AVX512F__)
    /* TODO: AVX-512 implementation */
    gemm_nt_q5_k_q8_k_ref(A_q8, B, bias, C, M, N, K);
#elif defined(__AVX2__)
    /* TODO: AVX-2 implementation */
    gemm_nt_q5_k_q8_k_ref(A_q8, B, bias, C, M, N, K);
#elif defined(__AVX__)
    /* TODO: AVX implementation */
    gemm_nt_q5_k_q8_k_ref(A_q8, B, bias, C, M, N, K);
#elif defined(__SSE4_1__)
    /* TODO: SSE4.1 implementation */
    gemm_nt_q5_k_q8_k_ref(A_q8, B, bias, C, M, N, K);
#else
    gemm_nt_q5_k_q8_k_ref(A_q8, B, bias, C, M, N, K);
#endif
}

References C, and gemm_nt_q5_k_q8_k_ref().

◆ gemm_nt_q6_k()

void gemm_nt_q6_k	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 212 of file gemm_kernels_q6k.c.

{
    if (!A || !B || !C) {
        return;
    }
    if (M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    /* gemm_q6_k produces Y as [batch x M_out] where:
     *   batch = M (tokens)
     *   M_out = N (output channels) */
    gemm_q6_k(C, B, A, /*M_out=*/N, /*N_batch=*/M, K);
 
    if (!bias) {
        return;
    }
 
    for (int i = 0; i < M; ++i) {
        float *row = C + (size_t)i * (size_t)N;
        for (int j = 0; j < N; ++j) {
            row[j] += bias[j];
        }
    }
}

References C, and gemm_q6_k().

Referenced by ck_gemm_nt_quant(), gemm_nt_q6_k_ref(), qwen2_0_5b_decode_layer_0_decode(), qwen2_0_5b_decode_layer_10_decode(), qwen2_0_5b_decode_layer_13_decode(), qwen2_0_5b_decode_layer_16_decode(), qwen2_0_5b_decode_layer_19_decode(), qwen2_0_5b_decode_layer_1_decode(), qwen2_0_5b_decode_layer_21_decode(), qwen2_0_5b_decode_layer_3_decode(), qwen2_0_5b_decode_layer_6_decode(), qwen2_0_5b_decode_layer_7_decode(), qwen2_0_5b_decode_layer_8_decode(), and qwen2_0_5b_decode_layer_9_decode().

◆ gemm_nt_q6_k_q8_k()

void gemm_nt_q6_k_q8_k	(	const void *	A_q8,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

NT GEMM: C = A @ B^T where A is Q8_K and B is Q6_K.

This is the typical inference pattern:

A: Activations in Q8_K format [M x K]
B: Weights in Q6_K format [N x K]
C: Output [M x N]

Parameters

A_q8	Input activations in Q8_K format
B	Weight matrix in Q6_K format
bias	Optional bias vector [N]
C	Output matrix
M	Batch size (number of tokens)
N	Output dimension
K	Input dimension

Definition at line 1515 of file gemm_kernels_q6k_q8k.c.

{
    if (!A_q8 || !B || !C) {
        return;
    }
    if (M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    /* Prefill GEMM is the hot Qwen2/Qwen3.5 MLP-down path. Keep decode
     * gemv_q6_k_q8_k() conservative, but allow GEMM/prefill to use the
     * parity-gated SIMD dot helper by default. CK strict parity and
     * CK_DEBUG_Q6K_Q8K_REF=1 still force the scalar reference reduction. */
    const block_q8_K *A = (const block_q8_K *)A_q8;
    const block_q6_K *W = (const block_q6_K *)B;
    const int blocks_per_vec = K / QK_K;
    const int blocks_per_row = K / QK_K;
 
    for (int m = 0; m < M; ++m) {
        const block_q8_K *a_row = A + (size_t)m * (size_t)blocks_per_vec;
        float *c_row = C + (size_t)m * (size_t)N;
        for (int n = 0; n < N; ++n) {
            const block_q6_K *w_row = W + (size_t)n * (size_t)blocks_per_row;
            const float b = bias ? bias[n] : 0.0f;
            c_row[n] = ck_dot_q6_k_q8_k_fast_or_ref(w_row, a_row, K) + b;
        }
    }
}

References C, ck_dot_q6_k_q8_k_fast_or_ref(), and QK_K.

Referenced by gemm_nt_q8_k_mlp_dispatch(), and gemm_nt_q8_k_qkv_dispatch().

◆ gemm_nt_q6_k_q8_k_m4_tile()

void gemm_nt_q6_k_q8_k_m4_tile	(	const void *	A_q8,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	m0,
		int	m1,
		int	n0,
		int	n1
	)

Definition at line 1702 of file gemm_kernels_q6k_q8k.c.

{
    if (!A_q8 || !B || !C || M <= 0 || N <= 0 || K <= 0 ||
        K % QK_K != 0) {
        return;
    }
    if (m0 < 0) m0 = 0;
    if (n0 < 0) n0 = 0;
    if (m1 > M) m1 = M;
    if (n1 > N) n1 = N;
    if (m0 >= m1 || n0 >= n1) return;
 
#if defined(__AVX2__)
    if (!ck_strict_parity_enabled() && !ck_q6k_q8k_force_ref()) {
        const block_q8_K *A = (const block_q8_K *)A_q8;
        const block_q6_K *W = (const block_q6_K *)B;
        const int blocks_per_vec = K / QK_K;
        for (int n = n0; n < n1; ++n) {
            const block_q6_K *w_row =
                W + (size_t)n * (size_t)blocks_per_vec;
            const float b = bias ? bias[n] : 0.0f;
            int m = m0;
            for (; m + 4 <= m1; m += 4) {
                float values[4];
                dot_q6_k_q8_k_avx2_m4(
                    w_row, A + (size_t)m * (size_t)blocks_per_vec,
                    blocks_per_vec, 4, K, values);
                for (int r = 0; r < 4; ++r) {
                    C[(size_t)(m + r) * (size_t)N + (size_t)n] = values[r] + b;
                }
            }
            if (m < m1) {
                float values[4];
                const int rows = m1 - m;
                dot_q6_k_q8_k_avx2_m4(
                    w_row, A + (size_t)m * (size_t)blocks_per_vec,
                    blocks_per_vec, rows, K, values);
                for (int r = 0; r < rows; ++r) {
                    C[(size_t)(m + r) * (size_t)N + (size_t)n] = values[r] + b;
                }
            }
        }
        return;
    }
#endif
    gemm_nt_q6_k_q8_k_tile(
        A_q8, B, bias, C, M, N, K, m0, m1, n0, n1);
}

References C, ck_q6k_q8k_force_ref(), ck_strict_parity_enabled(), gemm_nt_q6_k_q8_k_tile(), and QK_K.

◆ gemm_nt_q6_k_q8_k_tile()

void gemm_nt_q6_k_q8_k_tile	(	const void *	A_q8,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	m0,
		int	m1,
		int	n0,
		int	n1
	)

Compute one C[m0:m1, n0:n1] tile for Q8_K activations x Q6_K weights.

Pure tile math only: no threadpool, no global scheduling, no allocation. The orchestrator decides how to split tile jobs across cores.

Definition at line 1665 of file gemm_kernels_q6k_q8k.c.

{
    if (!A_q8 || !B || !C) {
        return;
    }
    if (M <= 0 || N <= 0 || K <= 0 || K % QK_K != 0) {
        return;
    }
    if (m0 < 0) m0 = 0;
    if (n0 < 0) n0 = 0;
    if (m1 > M) m1 = M;
    if (n1 > N) n1 = N;
    if (m0 >= m1 || n0 >= n1) {
        return;
    }
 
    const block_q8_K *A = (const block_q8_K *)A_q8;
    const block_q6_K *W = (const block_q6_K *)B;
    const int blocks_per_vec = K / QK_K;
    const int blocks_per_row = K / QK_K;
 
    for (int n = n0; n < n1; ++n) {
        const block_q6_K *w_row = W + (size_t)n * (size_t)blocks_per_row;
        const float b = bias ? bias[n] : 0.0f;
        for (int m = m0; m < m1; ++m) {
            const block_q8_K *a_row = A + (size_t)m * (size_t)blocks_per_vec;
            C[(size_t)m * (size_t)N + (size_t)n] = ck_dot_q6_k_q8_k_fast_or_ref(w_row, a_row, K) + b;
        }
    }
}

References C, ck_dot_q6_k_q8_k_fast_or_ref(), and QK_K.

Referenced by gemm_nt_q6_k_q8_k_m4_tile(), and gemm_nt_q6_k_q8_k_tiled().

◆ gemm_nt_q6_k_q8_k_tiled()

void gemm_nt_q6_k_q8_k_tiled	(	const void *	A_q8,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Experimental single-thread tiled NT GEMM wrapper.

Kept as a separate symbol from gemm_nt_q6_k_q8_k for benchmarks and parity. Production prefill should prefer the v8 2D tile scheduler when enabled.

Definition at line 1763 of file gemm_kernels_q6k_q8k.c.

{
    enum { TILE_M = 8, TILE_N = 16 };
    for (int n0 = 0; n0 < N; n0 += TILE_N) {
        const int n1 = (n0 + TILE_N < N) ? (n0 + TILE_N) : N;
        for (int m0 = 0; m0 < M; m0 += TILE_M) {
            const int m1 = (m0 + TILE_M < M) ? (m0 + TILE_M) : M;
            gemm_nt_q6_k_q8_k_tile(A_q8, B, bias, C, M, N, K, m0, m1, n0, n1);
        }
    }
}

References C, and gemm_nt_q6_k_q8_k_tile().

◆ gemm_nt_q8_0()

void gemm_nt_q8_0	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 879 of file gemm_kernels_q8_0.c.

{
#if defined(__AVX512F__)
    if (ck_q8_0_fp32_m4n4_enabled() && M >= 4 && N >= 4 && K % QK8_0 == 0) {
        gemm_nt_q8_0_m4n4_avx512(A, B, bias, C, M, N, K);
        return;
    }
#endif
    gemm_nt_q8_0_rowloop(A, B, bias, C, M, N, K);
}

References C, ck_q8_0_fp32_m4n4_enabled(), gemm_nt_q8_0_rowloop(), and QK8_0.

◆ gemm_nt_q8_0_q8_0()

void gemm_nt_q8_0_q8_0	(	const void *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

gemm_nt_q8_0_q8_0 with optional bias (matches header signature)

C[m,n] = A[m,K] @ B[n,K]^T + bias[n]

Definition at line 552 of file gemm_batch_int8.c.

{
    /* First compute GEMM */
#if defined(__AVX2__)
    /*
     * The production contract is bit-exact with llama.cpp's eight-lane
     * FP32 accumulation tree.  AVX-512/VNNI changes the integer dot, but it
     * must not silently replace that FP32 reduction with the scalar-per-block
     * candidate above.  The AVX2-named entry point delegates each activation
     * row to the certified x4 provider, which also uses VNNI instructions when
     * they are available while preserving the declared reduction order.
     */
    gemm_nt_q8_0_q8_0_avx2(A, B, C, M, N, K);
#elif defined(__AVX__)
    gemm_nt_q8_0_q8_0_avx(A, B, C, M, N, K);
#else
    gemm_nt_q8_0_q8_0_ref(A, B, C, M, N, K);
#endif
 
    /* Add bias if provided */
    if (bias != NULL) {
        for (int m = 0; m < M; m++) {
            for (int n = 0; n < N; n++) {
                C[(size_t)m * N + n] += bias[n];
            }
        }
    }
}

References C, and gemm_nt_q8_0_q8_0_ref().

Referenced by gemm_nt_q8_0_dispatch(), and gemm_nt_q8_0_mlp_dispatch().

◆ gemm_nt_q8_0_q8_0_contract()

void gemm_nt_q8_0_q8_0_contract	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 414 of file gemm_kernels_q8_0_q8_0_contract.c.

{
    if (!A || !B || !C || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    const float *A_use = A;
    const int strict = ck_strict_parity_enabled();
    const int dump_enabled = strict && ck_q80_contract_dump_enabled();
    int strict_cached_layer = -1;
    int strict_dump_layer = -1;
    if (ck_strict_parity_enabled() && ck_q80_contract_cached_input_enabled()) {
        const float *cached = ck_strict_consume_next_gemm_a((size_t) M * (size_t) K);
        if (cached) {
            A_use = cached;
            strict_cached_layer = ck_q80_contract_cached_gemm_seq++;
        }
    }
    if (dump_enabled) {
        strict_dump_layer = strict_cached_layer >= 0
            ? strict_cached_layer
            : ck_q80_contract_cached_gemm_seq++;
    }
 
    if (dump_enabled && strict_dump_layer >= 0) {
        ck_q80_contract_dump_tensor(strict_cached_layer >= 0
                                        ? "strict_out_proj_input_cached"
                                        : "strict_out_proj_input_live",
                                    strict_dump_layer,
                                    A_use,
                                    (size_t) M * (size_t) K);
    }
 
    if (strict &&
        gemm_nt_q8_0_q8_0_ggml_strict(A_use, B, bias, C, M, N, K)) {
        if (dump_enabled && strict_dump_layer >= 0) {
            ck_q80_contract_dump_tensor("strict_out_proj_output_ggml",
                                        strict_dump_layer,
                                        C,
                                        (size_t) M * (size_t) N);
        }
        return;
    }
 
    if (!strict) {
#pragma omp parallel for schedule(static) if(M > 1)
        for (int m = 0; m < M; ++m) {
            gemv_q8_0_q8_0_contract(&C[m * N], B, &A_use[m * K], N, K);
            if (bias) {
                for (int n = 0; n < N; ++n) {
                    C[m * N + n] += bias[n];
                }
            }
        }
        return;
    }
 
    for (int m = 0; m < M; ++m) {
        gemv_q8_0_q8_0_contract(&C[m * N], B, &A_use[m * K], N, K);
        if (bias) {
            for (int n = 0; n < N; ++n) {
                C[m * N + n] += bias[n];
            }
        }
    }
 
    if (dump_enabled && strict_dump_layer >= 0) {
        ck_q80_contract_dump_tensor("strict_out_proj_output_fallback",
                                    strict_dump_layer,
                                    C,
                                    (size_t) M * (size_t) N);
    }
}

References C, ck_q80_contract_cached_gemm_seq, ck_q80_contract_cached_input_enabled(), ck_q80_contract_dump_enabled(), ck_q80_contract_dump_tensor(), ck_strict_consume_next_gemm_a(), ck_strict_parity_enabled(), gemm_nt_q8_0_q8_0_ggml_strict(), and gemv_q8_0_q8_0_contract().

◆ gemm_nt_q8_0_q8_0_m2n4()

void gemm_nt_q8_0_q8_0_m2n4	(	const void *	A_q8,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 586 of file gemm_batch_int8.c.

{
    gemm_q8_0_q8_0_m2n4(C, B, A, M, N, K);
    if (bias != NULL) {
        for (int m = 0; m < M; ++m) {
            for (int n = 0; n < N; ++n) {
                C[(size_t)m * (size_t)N + n] += bias[n];
            }
        }
    }
}

References C, and gemm_q8_0_q8_0_m2n4().

◆ gemm_nt_q8_0_q8_0_m2n4_tile()

void gemm_nt_q8_0_q8_0_m2n4_tile	(	const void *	A_q8,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K,
		int	ldc
	)

Definition at line 603 of file gemm_batch_int8.c.

{
    gemm_q8_0_q8_0_m2n4_strided(C, ldc, B, A, M, N, K);
    if (bias != NULL) {
        for (int m = 0; m < M; ++m) {
            for (int n = 0; n < N; ++n) {
                C[(size_t)m * (size_t)ldc + n] += bias[n];
            }
        }
    }
}

References C, and gemm_q8_0_q8_0_m2n4_strided().

◆ gemm_q4_k()

void gemm_q4_k	(	float *	Y,
		const void *	W,
		const float *	X,
		int	M,
		int	N,
		int	K
	)

Auto-dispatch GEMM based on available SIMD.

Definition at line 486 of file gemm_kernels_q4k.c.

{
    /* Use reference implementation for correctness
     * TODO: Fix AVX-512 version to match llama.cpp layout */
    gemm_q4_k_ref(Y, W, X, M, N, K);
}

References gemm_q4_k_ref().

Referenced by gemm_nt_q4_k().

◆ gemm_q4_k_q8_k()

void gemm_q4_k_q8_k	(	float *	Y,
		const void *	W,
		const void *	X_q8,
		int	M,
		int	N,
		int	K
	)

Definition at line 354 of file gemm_kernels_q4k_q8k.c.

{
    if (!Y || !W || !X_q8 || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    const block_q8_K *X = (const block_q8_K *)X_q8;
    const int blocks_per_vec = K / QK_K;
    const int blocks_per_row = K / QK_K;
    const size_t work_items = (size_t)M * (size_t)N;
 
    if (work_items >= 4096u && M >= 512 && N > 1) {
        ck_threadpool_t *pool = ck_threadpool_global();
        const int pool_threads = pool ? ck_threadpool_n_threads(pool) : 1;
        int active_threads = pool_threads;
        if (active_threads > M) {
            active_threads = M;
        }
        if (active_threads > 1) {
            gemm_q4_k_q8_k_work_t work = {
                .Y = Y,
                .W = W,
                .X = X,
                .M_out = M,
                .N_batch = N,
                .K = K,
                .blocks_per_vec = blocks_per_vec,
                .blocks_per_row = blocks_per_row,
            };
            ck_threadpool_dispatch_n(pool, active_threads, gemm_q4_k_q8_k_thread_fn, &work);
            return;
        }
    }
 
    for (int n = 0; n < N; ++n) {
        const block_q8_K *x_row = X + (size_t)n * (size_t)blocks_per_vec;
        gemv_q4_k_q8_k(&Y[n * M], W, x_row, M, K);
    }
}

References ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), gemm_q4_k_q8_k_thread_fn(), gemv_q4_k_q8_k(), and QK_K.

Referenced by gemm_nt_q4_k_q8_k().

◆ gemm_q6_k()

void gemm_q6_k	(	float *	Y,
		const void *	W,
		const float *	X,
		int	M,
		int	N,
		int	K
	)

Definition at line 195 of file gemm_kernels_q6k.c.

{
    if (!Y || !W || !X) {
        return;
    }
    if (M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    for (int n = 0; n < N; ++n) {
        gemv_q6_k(&Y[n * M], W, &X[n * K], M, K);
    }
}

References gemv_q6_k().

Referenced by gemm_nt_q6_k().

◆ gemm_q6_k_q8_k()

void gemm_q6_k_q8_k	(	float *	Y,
		const void *	W,
		const void *	X_q8,
		int	M,
		int	N,
		int	K
	)

GEMM: Y = W @ X^T where W is Q6_K and X is Q8_K.

Parameters

Y	Output matrix [N x M] in row-major
W	Weight matrix in Q6_K format [M x K]
X_q8	Input matrix in Q8_K format [N x K]
M	Number of output rows (output dim)
N	Number of input vectors (batch size)
K	Input dimension

Definition at line 1481 of file gemm_kernels_q6k_q8k.c.

{
    if (!Y || !W || !X_q8 || M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    const block_q8_K *X = (const block_q8_K *)X_q8;
    const int blocks_per_vec = K / QK_K;
 
    for (int n = 0; n < N; ++n) {
        const block_q8_K *x_row = X + (size_t)n * (size_t)blocks_per_vec;
        gemv_q6_k_q8_k(&Y[n * M], W, x_row, M, K);
    }
}

References gemv_q6_k_q8_k(), and QK_K.

◆ gemm_swiglu_fused()

void gemm_swiglu_fused	(	const float *	x,
		const float *	W_gate,
		const float *	W_up,
		const float *	b_gate,
		const float *	b_up,
		float *	output,
		int	M,
		int	N,
		int	K
	)

Definition at line 241 of file gemm_fused_kernels.c.

{
#if defined(__AVX__)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        const float *x_row = &x[i * K];
        float *out_row = &output[i * N];
 
        for (int j = 0; j < N; j++) {
            const float *w_gate_row = &W_gate[j * K];
            const float *w_up_row = &W_up[j * K];
 
            // Compute both dot products in parallel using SIMD
            __m256 gate_vec = _mm256_setzero_ps();
            __m256 up_vec = _mm256_setzero_ps();
 
            int k;
            for (k = 0; k <= K - 8; k += 8) {
                __m256 x_vec = _mm256_loadu_ps(&x_row[k]);
                __m256 wg_vec = _mm256_loadu_ps(&w_gate_row[k]);
                __m256 wu_vec = _mm256_loadu_ps(&w_up_row[k]);
 
                // gate += x * W_gate
                gate_vec = _mm256_add_ps(gate_vec, _mm256_mul_ps(x_vec, wg_vec));
                // up += x * W_up
                up_vec = _mm256_add_ps(up_vec, _mm256_mul_ps(x_vec, wu_vec));
            }
 
            // Horizontal sum
            float gate = hsum256_ps_fused(gate_vec);
            float up = hsum256_ps_fused(up_vec);
 
            // Scalar remainder
            for (; k < K; k++) {
                gate += x_row[k] * w_gate_row[k];
                up += x_row[k] * w_up_row[k];
            }
 
            // Add biases
            if (b_gate) gate += b_gate[j];
            if (b_up) up += b_up[j];
 
            // SwiGLU: SiLU(gate) * up = gate * sigmoid(gate) * up
            float sig = 1.0f / (1.0f + expf(-gate));
            out_row[j] = gate * sig * up;
        }
    }
#else
    // Scalar fallback
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float gate = 0.0f;
            float up = 0.0f;
 
            for (int k = 0; k < K; k++) {
                gate += x[i * K + k] * W_gate[j * K + k];
                up += x[i * K + k] * W_up[j * K + k];
            }
 
            if (b_gate) gate += b_gate[j];
            if (b_up) up += b_up[j];
 
            // SwiGLU: SiLU(gate) * up
            float sig = 1.0f / (1.0f + expf(-gate));
            output[i * N + j] = gate * sig * up;
        }
    }
#endif
}

Referenced by ck_mlp_swiglu_forward_fused_token().

◆ gemm_tn_avx512()

void gemm_tn_avx512	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 709 of file gemm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        gemm_tn_serial_double(A, B, bias, C, M, N, K);
        return;
    }
#if defined(__AVX512F__)
    // Vectorize over N (output columns)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        int j = 0;
        for (; j <= N - 16; j += 16) {
            __m512 sum_vec = bias ? _mm512_loadu_ps(&bias[j]) : _mm512_setzero_ps();
            for (int k = 0; k < K; k++) {
                __m512 a_broadcast = _mm512_set1_ps(A[k * M + i]);
                __m512 b_vec = _mm512_loadu_ps(&B[k * N + j]);
                sum_vec = _mm512_fmadd_ps(a_broadcast, b_vec, sum_vec);
            }
            _mm512_storeu_ps(&C[i * N + j], sum_vec);
        }
        for (; j < N; j++) {
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[k * M + i] * B[k * N + j];
            }
            C[i * N + j] = sum;
        }
    }
#elif defined(__AVX__)
    // AVX1: vectorize over N (8 columns at a time)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        int j = 0;
        for (; j <= N - 8; j += 8) {
            __m256 sum_vec = bias ? _mm256_loadu_ps(&bias[j]) : _mm256_setzero_ps();
            for (int k = 0; k < K; k++) {
                __m256 a_broadcast = _mm256_set1_ps(A[k * M + i]);
                __m256 b_vec = _mm256_loadu_ps(&B[k * N + j]);
                __m256 prod = _mm256_mul_ps(a_broadcast, b_vec);
                sum_vec = _mm256_add_ps(sum_vec, prod);
            }
            _mm256_storeu_ps(&C[i * N + j], sum_vec);
        }
        for (; j < N; j++) {
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[k * M + i] * B[k * N + j];
            }
            C[i * N + j] = sum;
        }
    }
#else
    gemm_tn_parallel(A, B, bias, C, M, N, K);
#endif
}

References C, ck_strict_parity_enabled(), gemm_tn_parallel(), and gemm_tn_serial_double().

◆ gemm_tn_blocked()

void gemm_tn_blocked	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 769 of file gemm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        gemm_tn_serial_double(A, B, bias, C, M, N, K);
        return;
    }
#if defined(__AVX512F__)
    const int block_size = 64;
#elif defined(__AVX__)
    const int block_size = 32;
#else
    const int block_size = 32;
#endif
    // Initialize C with bias (parallelized)
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            C[i * N + j] = bias ? bias[j] : 0.0f;
        }
    }
    // Blocked multiply-accumulate (parallelized over M blocks)
#pragma omp parallel for
    for (int ii = 0; ii < M; ii += block_size) {
        for (int kk = 0; kk < K; kk += block_size) {
            for (int jj = 0; jj < N; jj += block_size) {
                int i_end = ck_min(ii + block_size, M);
                int k_end = ck_min(kk + block_size, K);
                int j_end = ck_min(jj + block_size, N);
 
                for (int k = kk; k < k_end; k++) {
                    for (int i = ii; i < i_end; i++) {
                        float a_val = A[k * M + i];
#if defined(__AVX512F__)
                        __m512 a_broadcast = _mm512_set1_ps(a_val);
                        int j;
                        for (j = jj; j <= j_end - 16; j += 16) {
                            __m512 b_vec = _mm512_loadu_ps(&B[k * N + j]);
                            __m512 c_vec = _mm512_loadu_ps(&C[i * N + j]);
                            c_vec = _mm512_fmadd_ps(a_broadcast, b_vec, c_vec);
                            _mm512_storeu_ps(&C[i * N + j], c_vec);
                        }
                        for (; j < j_end; j++) {
                            C[i * N + j] += a_val * B[k * N + j];
                        }
#elif defined(__AVX__)
                        __m256 a_broadcast = _mm256_set1_ps(a_val);
                        int j;
                        for (j = jj; j <= j_end - 8; j += 8) {
                            __m256 b_vec = _mm256_loadu_ps(&B[k * N + j]);
                            __m256 c_vec = _mm256_loadu_ps(&C[i * N + j]);
                            __m256 prod = _mm256_mul_ps(a_broadcast, b_vec);
                            c_vec = _mm256_add_ps(c_vec, prod);
                            _mm256_storeu_ps(&C[i * N + j], c_vec);
                        }
                        for (; j < j_end; j++) {
                            C[i * N + j] += a_val * B[k * N + j];
                        }
#else
                        for (int j = jj; j < j_end; j++) {
                            C[i * N + j] += a_val * B[k * N + j];
                        }
#endif
                    }
                }
            }
        }
    }
}

References C, ck_min(), ck_strict_parity_enabled(), and gemm_tn_serial_double().

◆ gemm_tn_parallel()

void gemm_tn_parallel	(	const float *	A,
		const float *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 687 of file gemm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        gemm_tn_serial_double(A, B, bias, C, M, N, K);
        return;
    }
#pragma omp parallel for
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = bias ? bias[j] : 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[k * M + i] * B[k * N + j];
            }
            C[i * N + j] = sum;
        }
    }
}

References C, ck_strict_parity_enabled(), and gemm_tn_serial_double().

Referenced by fc1_backward_kernel(), fc2_backward_kernel(), and gemm_tn_avx512().

◆ gemma4_final_logit_softcap_forward()

void gemma4_final_logit_softcap_forward	(	float *	logits,
		int	tokens,
		int	vocab_size,
		float	cap
	)

Definition at line 405 of file gemma4_per_layer_embed.c.

{
    if (!logits || tokens <= 0 || vocab_size <= 0 || cap <= 0.0f) {
        return;
    }
    const float inv_cap = 1.0f / cap;
    const size_t total = (size_t)tokens * (size_t)vocab_size;
    for (size_t i = 0; i < total; ++i) {
        logits[i] = tanhf(logits[i] * inv_cap) * cap;
    }
}

References vocab_size.

◆ gemma4_per_layer_embed_forward()

void gemma4_per_layer_embed_forward	(	float *	hidden,
		const float *	per_layer_input,
		const float *	inp_gate,
		const float *	proj,
		const float *	post_norm,
		const float *	out_scale,
		int	tokens,
		int	layer,
		int	num_layers,
		int	embed_dim,
		int	per_layer_dim,
		float	eps
	)

Definition at line 348 of file gemma4_per_layer_embed.c.

{
    if (!hidden || !per_layer_input || !inp_gate || !proj || !post_norm ||
        tokens <= 0 || layer < 0 || layer >= num_layers || embed_dim <= 0 ||
        per_layer_dim != QK_K || embed_dim > 4096) {
        return;
    }
 
    ck_gemma4_embed_args_t args = {
        .hidden = hidden,
        .per_layer_input = per_layer_input,
        .inp_gate = inp_gate,
        .proj = proj,
        .post_norm = post_norm,
        .out_scale = out_scale,
        .layer = layer,
        .num_layers = num_layers,
        .embed_dim = embed_dim,
        .per_layer_dim = per_layer_dim,
        .eps = eps,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    const char *disabled = getenv("CK_DISABLE_GEMMA4_EMBED_PARALLEL");
    if (disabled && disabled[0] && strcmp(disabled, "0") != 0) active = 1;
    if (active > tokens) active = tokens;
    ck_threadpool_parallel_for_n(
        pool, active, 0, tokens, 1, ck_gemma4_embed_range, &args);
}

References ck_gemma4_embed_range(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and QK_K.

◆ gemma4_per_layer_prepare_bf16_forward()

void gemma4_per_layer_prepare_bf16_forward	(	float *	per_layer_input,
		const float *	hidden,
		const int32_t *	token_ids,
		const uint16_t *	per_layer_token_emb,
		const uint16_t *	per_layer_model_proj,
		const float *	per_layer_proj_norm,
		int	tokens,
		int	num_layers,
		int	embed_dim,
		int	per_layer_dim,
		int	vocab_size,
		float	eps
	)

Definition at line 254 of file gemma4_per_layer_embed.c.

{
    if (!per_layer_input || !hidden || !token_ids || !per_layer_token_emb ||
        !per_layer_model_proj || !per_layer_proj_norm || tokens <= 0 ||
        num_layers <= 0 || embed_dim <= 0 || per_layer_dim <= 0 || vocab_size <= 0) {
        return;
    }
 
    if (per_layer_dim > QK_K) {
        return;
    }
    ck_gemma4_prepare_args_t args = {
        .per_layer_input = per_layer_input,
        .hidden = hidden,
        .token_ids = token_ids,
        .per_layer_token_emb = per_layer_token_emb,
        .per_layer_model_proj = per_layer_model_proj,
        .per_layer_proj_norm = per_layer_proj_norm,
        .num_layers = num_layers,
        .embed_dim = embed_dim,
        .per_layer_dim = per_layer_dim,
        .vocab_size = vocab_size,
        .eps = eps,
    };
    ck_gemma4_prepare_parallel(tokens, ck_gemma4_prepare_bf16_range, &args);
}

References ck_gemma4_prepare_bf16_range(), ck_gemma4_prepare_parallel(), QK_K, and vocab_size.

◆ gemma4_per_layer_prepare_forward()

void gemma4_per_layer_prepare_forward	(	float *	per_layer_input,
		const float *	hidden,
		const int32_t *	token_ids,
		const void *	per_layer_token_emb,
		const uint16_t *	per_layer_model_proj,
		const float *	per_layer_proj_norm,
		int	tokens,
		int	num_layers,
		int	embed_dim,
		int	per_layer_dim,
		int	vocab_size,
		float	eps
	)

Definition at line 218 of file gemma4_per_layer_embed.c.

{
    if (!per_layer_input || !hidden || !token_ids || !per_layer_token_emb ||
        !per_layer_model_proj || !per_layer_proj_norm || tokens <= 0 ||
        num_layers <= 0 || embed_dim <= 0 || per_layer_dim != QK_K || vocab_size <= 0) {
        return;
    }
 
    ck_gemma4_prepare_args_t args = {
        .per_layer_input = per_layer_input,
        .hidden = hidden,
        .token_ids = token_ids,
        .per_layer_token_emb = per_layer_token_emb,
        .per_layer_model_proj = per_layer_model_proj,
        .per_layer_proj_norm = per_layer_proj_norm,
        .num_layers = num_layers,
        .embed_dim = embed_dim,
        .per_layer_dim = per_layer_dim,
        .vocab_size = vocab_size,
        .eps = eps,
    };
    ck_gemma4_prepare_parallel(tokens, ck_gemma4_prepare_q5_range, &args);
}

References ck_gemma4_prepare_parallel(), ck_gemma4_prepare_q5_range(), QK_K, and vocab_size.

◆ gemma4_v_norm_forward()

void gemma4_v_norm_forward	(	const float *	input,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	num_kv_heads,
		int	head_dim,
		float	eps
	)

Definition at line 687 of file rmsnorm_kernels.c.

{
    if (!input || !output || tokens <= 0 || num_kv_heads <= 0 || head_dim <= 0) {
        return;
    }
    rmsnorm_forward_no_weight(input, output, rstd_cache,
                              tokens * num_kv_heads, head_dim, head_dim, eps);
}

References rmsnorm_forward_no_weight().

◆ gemma4_v_norm_forward_parallel_dispatch()

void gemma4_v_norm_forward_parallel_dispatch	(	const float *	input,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	num_kv_heads,
		int	head_dim,
		float	eps
	)

◆ gemma4_vision_projector_prep_forward()

void gemma4_vision_projector_prep_forward	(	const float *	input,
		float *	output,
		int	tokens,
		int	dim,
		float	scale,
		float	eps
	)

Definition at line 633 of file vision_kernels.c.

{
    if (input == NULL || output == NULL || tokens <= 0 || dim <= 0) return;
    if (eps <= 0.0f) eps = 1.0e-6f;
    for (int t = 0; t < tokens; ++t) {
        const float *src = input + (size_t)t * (size_t)dim;
        float *dst = output + (size_t)t * (size_t)dim;
        double ss = 0.0;
        for (int i = 0; i < dim; ++i) {
            const float v = src[i] * scale;
            ss += (double)v * (double)v;
        }
        const float inv_rms = 1.0f / sqrtf((float)(ss / (double)dim) + eps);
        for (int i = 0; i < dim; ++i) {
            dst[i] = (src[i] * scale) * inv_rms;
        }
    }
}

◆ gemv_bf16()

void gemv_bf16	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 622 of file gemm_kernels_bf16.c.

{
    gemv_bf16_row_range(
        y, (const uint16_t *)W, x, M, K, 0, M);
}

References gemv_bf16_row_range().

Referenced by gemv_bf16_parallel_dispatch().

◆ gemv_bf16_bf16_storage()

void gemv_bf16_bf16_storage	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 712 of file gemm_kernels_bf16.c.

{
    gemv_bf16_bf16_storage_parallel_dispatch(y, W, x, M, K);
}

References gemv_bf16_bf16_storage_parallel_dispatch().

◆ gemv_bf16_bf16_storage_parallel_dispatch()

void gemv_bf16_bf16_storage_parallel_dispatch	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 688 of file gemm_kernels_bf16.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    if (!pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)K <= 65536) {
        gemv_bf16_bf16_storage_row_range(
            y, (const uint16_t *)W, x, M, K, 0, M);
        return;
    }
 
    ck_gemv_bf16_args_t args = {
        .y = y, .w = (const uint16_t *)W, .x = x, .M = M, .K = K,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemv_bf16_storage_rows, &args);
}

References ck_gemv_bf16_storage_rows(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and gemv_bf16_bf16_storage_row_range().

Referenced by gemv_bf16_bf16_storage().

◆ gemv_bf16_parallel_dispatch()

void gemv_bf16_parallel_dispatch	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 646 of file gemm_kernels_bf16.c.

{
    ck_threadpool_t *pool = ck_threadpool_global();
    if (!pool || ck_threadpool_n_threads(pool) <= 1 || M < 2 ||
        (size_t)M * (size_t)K <= 65536) {
        gemv_bf16(y, W, x, M, K);
        return;
    }
 
    ck_gemv_bf16_args_t args = {
        .y = y, .w = (const uint16_t *)W, .x = x, .M = M, .K = K,
    };
    int active = ck_threadpool_n_threads(pool);
    if (active > M) active = M;
    int grain = M / (active * 4);
    if (grain < 1) grain = 1;
    ck_threadpool_parallel_for_n(
        pool, active, 0, M, grain, ck_gemv_bf16_rows, &args);
}

References ck_gemv_bf16_rows(), ck_threadpool_global(), ck_threadpool_n_threads(), ck_threadpool_parallel_for_n(), and gemv_bf16().

◆ gemv_fused_q5_0_bias_dispatch()

void gemv_fused_q5_0_bias_dispatch	(	float *	y,
		const void *	W,
		const float *	x,
		const float *	bias,
		int	M,
		int	K
	)

Definition at line 510 of file gemv_fused_quant_bias.c.

{
#if defined(__AVX__)
    gemv_fused_q5_0_bias_avx(y, W, x, bias, M, K);
#else
    gemv_fused_q5_0_bias(y, W, x, bias, M, K);
#endif
}

References gemv_fused_q5_0_bias().

◆ gemv_fused_q8_0_bias_dispatch()

void gemv_fused_q8_0_bias_dispatch	(	float *	y,
		const void *	W,
		const float *	x,
		const float *	bias,
		int	M,
		int	K
	)

Definition at line 525 of file gemv_fused_quant_bias.c.

{
#if defined(__AVX__)
    gemv_fused_q8_0_bias_avx(y, W, x, bias, M, K);
#else
    gemv_fused_q8_0_bias(y, W, x, bias, M, K);
#endif
}

References gemv_fused_q8_0_bias().

◆ gemv_q4_0()

void gemv_q4_0	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Auto-dispatch GEMV.

Definition at line 132 of file gemm_kernels_q4_0.c.

{
#ifdef __AVX512F__
    gemv_q4_0_avx512(y, W, x, M, K);
#else
    gemv_q4_0_ref(y, W, x, M, K);
#endif
}

References gemv_q4_0_ref().

Referenced by dot_q4_0(), and gemm_q4_0().

◆ gemv_q4_k()

void gemv_q4_k	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Auto-dispatch GEMV based on available SIMD.

Definition at line 301 of file gemm_kernels_q4k.c.

{
    if (ck_q4k_debug_q8_contract() && K > 0 && (K % QK_K) == 0) {
        const int nb = K / QK_K;
        if (nb <= CK_Q4K_STACK_Q8_BLOCKS) {
            block_q8_K x_q8[CK_Q4K_STACK_Q8_BLOCKS];
            quantize_row_q8_k(x, x_q8, K);
            gemv_q4_k_q8_k(y, W, x_q8, M, K);
            return;
        }
    }
#ifdef __AVX512F__
    gemv_q4_k_avx512(y, W, x, M, K);
#elif defined(__AVX__)
    gemv_q4_k_avx(y, W, x, M, K);
#else
    gemv_q4_k_ref(y, W, x, M, K);
#endif
}

References ck_q4k_debug_q8_contract(), CK_Q4K_STACK_Q8_BLOCKS, gemv_q4_k_q8_k(), gemv_q4_k_ref(), QK_K, and quantize_row_q8_k().

Referenced by attention_mlp_fused_q4k(), dot_q4_k(), gemm_q4_k_ref(), layer_fused_attn_mlp_qkv_q4k(), and rmsnorm_qkv_q4k_fused().

◆ gemv_q4_k_q8_k()

void gemv_q4_k_q8_k	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K
	)

Definition at line 273 of file gemm_kernels_q4k_q8k.c.

{
    if (ck_q4k_q8k_force_ref()) {
        gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
        return;
    }
#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && !defined(CK_NO_AVX512_VNNI)
    /* VNNI: Best for decode (single token) - INT8 dot product acceleration */
    gemv_q4_k_q8_k_vnni(y, W, x_q8, M, K);
#elif defined(__AVX2__)
    gemv_q4_k_q8_k_avx2(y, W, x_q8, M, K);
#elif defined(__AVX__)
    /* AVX version uses maddubs_epi16 (more efficient than SSE) */
    gemv_q4_k_q8_k_avx(y, W, x_q8, M, K);
#elif defined(__SSE4_1__)
    gemv_q4_k_q8_k_sse(y, W, x_q8, M, K);
#else
    gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
#endif
}

References ck_q4k_q8k_force_ref(), gemv_q4_k_q8_k_avx(), gemv_q4_k_q8_k_avx2(), gemv_q4_k_q8_k_ref(), gemv_q4_k_q8_k_sse(), and gemv_q4_k_q8_k_vnni().

◆ gemv_q4_k_q8_k_parallel()

void gemv_q4_k_q8_k_parallel	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K,
		int	ith,
		int	nth
	)

Definition at line 240 of file gemm_kernels_q4k_q8k.c.

{
    if (!y || !W || !x_q8 || M <= 0 || K <= 0) {
        return;
    }
    if (ith < 0 || nth <= 0 || ith >= nth) {
        return;
    }
 
    /* Compute row range for this thread */
    const int dr = (M + nth - 1) / nth;
    const int r0 = dr * ith;
    const int r1 = (r0 + dr < M) ? (r0 + dr) : M;
 
    if (r0 >= M) {
        return;  /* This thread has no work */
    }
 
    const block_q4_K *blocks = (const block_q4_K *)W;
    const block_q8_K *x = (const block_q8_K *)x_q8;
    const int blocks_per_row = K / QK_K;
 
    /* Only process rows [r0, r1) */
    for (int row = r0; row < r1; ++row) {
        const block_q4_K *w_row = blocks + (size_t)row * (size_t)blocks_per_row;
        y[row] = dot_q4_k_q8_k_ref(w_row, x, K);
    }
}

References dot_q4_k_q8_k_ref(), and QK_K.

◆ gemv_q4_k_q8_k_parallel_simd()

void gemv_q4_k_q8_k_parallel_simd	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K,
		int	ith,
		int	nth
	)

Definition at line 263 of file gemm_kernels_q4k_avx.c.

{
    /* Fall back to reference parallel version */
    gemv_q4_k_q8_k_parallel(y, W, x_q8, M, K, ith, nth);
}

References gemv_q4_k_q8_k_parallel().

Referenced by decode_layer_parallel(), mlp_parallel(), and qkv_projection_parallel().

◆ gemv_q4_k_q8_k_ref()

void gemv_q4_k_q8_k_ref	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K
	)

Definition at line 201 of file gemm_kernels_q4k_q8k.c.

{
    if (!y || !W || !x_q8 || M <= 0 || K <= 0) {
        return;
    }
 
    const block_q4_K *blocks = (const block_q4_K *)W;
    const block_q8_K *x = (const block_q8_K *)x_q8;
    const int blocks_per_row = K / QK_K;
 
    for (int row = 0; row < M; ++row) {
        const block_q4_K *w_row = blocks + (size_t)row * (size_t)blocks_per_row;
        y[row] = dot_q4_k_q8_k_ref(w_row, x, K);
    }
}

References dot_q4_k_q8_k_ref(), and QK_K.

◆ gemv_q5_0()

void gemv_q5_0	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Auto-dispatch GEMV for Q5_0 weights based on CPU features.

Dispatch priority (best available):

AVX-512 (512-bit vectors) - Intel Skylake-X+
AVX2+FMA (256-bit vectors) - Intel Haswell+
AVX (256-bit vectors) - Intel Sandy Bridge+
SSE4.1 (128-bit vectors) - Intel Nehalem+
Reference (scalar) - Fallback

Uses ck_features.h for standardized feature detection.

Parameters

y	Output vector [M]
W	Weight matrix in Q5_0 format [M x K]
x	Input vector [K]
M	Number of output rows
K	Number of input columns (hidden dimension)

Definition at line 550 of file gemm_kernels_q5_0.c.

{
// Dispatch order: AVX512 > AVX2 > AVX > SSE > ref
#if defined(__AVX512F__)
    gemv_q5_0_avx512(y, W, x, M, K);
#elif defined(__AVX2__)
    gemv_q5_0_avx2(y, W, x, M, K);
#elif defined(__AVX__)
    gemv_q5_0_avx(y, W, x, M, K);
#elif defined(__SSE4_1__)
    gemv_q5_0_ref(y, W, x, M, K);
#else
    gemv_q5_0_ref(y, W, x, M, K);
#endif
}

References gemv_q5_0_ref().

◆ gemv_q5_0_parallel()

void gemv_q5_0_parallel	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K,
		int	ith,
		int	nth
	)

Parallel reference GEMV for Q5_0 × FP32.

Definition at line 579 of file gemm_kernels_q5_0.c.

{
    if (!y || !W || !x || M <= 0 || K <= 0) return;
    if (ith < 0 || nth <= 0 || ith >= nth) return;
 
    const int dr = (M + nth - 1) / nth;
    const int r0 = dr * ith;
    const int r1 = (r0 + dr < M) ? (r0 + dr) : M;
 
    if (r0 >= M) return;
 
    const block_q5_0 *blocks = (const block_q5_0 *)W;
    const int blocks_per_row = K / QK5_0;
 
    for (int row = r0; row < r1; row++) {
        float sum = 0.0f;
        for (int b = 0; b < blocks_per_row; b++) {
            const block_q5_0 *block = &blocks[row * blocks_per_row + b];
            const float d = CK_FP16_TO_FP32(block->d);
            const float *xp = &x[b * QK5_0];
 
            uint32_t qh;
            memcpy(&qh, block->qh, sizeof(qh));
 
            for (int j = 0; j < QK5_0 / 2; j++) {
                const uint8_t packed = block->qs[j];
                const int lo = (packed & 0x0F);
                const int hi = (packed >> 4);
                const int xh_0 = ((qh >> (j + 0)) << 4) & 0x10;
                const int xh_1 = ((qh >> (j + 12))) & 0x10;
                const int w0 = (lo | xh_0) - 16;
                const int w1 = (hi | xh_1) - 16;
                sum += d * (w0 * xp[j] + w1 * xp[j + QK5_0/2]);
            }
        }
        y[row] = sum;
    }
}

References CK_FP16_TO_FP32, block_q5_0::d, block_q5_0::qh, QK5_0, and block_q5_0::qs.

Referenced by gemv_q5_0_parallel_simd().

◆ gemv_q5_0_parallel_simd()

void gemv_q5_0_parallel_simd	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K,
		int	ith,
		int	nth
	)

Parallel SIMD GEMV for Q5_0 × FP32 with prefetching.

Definition at line 625 of file gemm_kernels_q5_0.c.

{
    if (!y || !W || !x || M <= 0 || K <= 0) return;
    if (ith < 0 || nth <= 0 || ith >= nth) return;
 
    const int dr = (M + nth - 1) / nth;
    const int r0 = dr * ith;
    const int r1 = (r0 + dr < M) ? (r0 + dr) : M;
 
    if (r0 >= M) return;
 
#if defined(__AVX__) || defined(__SSE4_1__)
    const block_q5_0 *blocks = (const block_q5_0 *)W;
    const int blocks_per_row = K / QK5_0;
    /* Prefetch first few rows */
    const int PREFETCH_ROWS = 4;
    for (int p = 0; p < PREFETCH_ROWS && r0 + p < r1; ++p) {
        const char *row_ptr = (const char *)(blocks + (r0 + p) * blocks_per_row);
        _mm_prefetch(row_ptr, _MM_HINT_T0);
        _mm_prefetch(row_ptr + 64, _MM_HINT_T0);
    }
 
    for (int row = r0; row < r1; ++row) {
        /* Prefetch rows ahead */
        if (row + PREFETCH_ROWS < r1) {
            const char *prefetch_ptr = (const char *)(blocks + (row + PREFETCH_ROWS) * blocks_per_row);
            _mm_prefetch(prefetch_ptr, _MM_HINT_T0);
            _mm_prefetch(prefetch_ptr + 64, _MM_HINT_T0);
        }
 
        /* Use SIMD dot product for this row */
#if defined(__AVX512F__)
        /* Call single-row AVX512 implementation */
        gemv_q5_0_avx512(&y[row], (const char *)blocks + row * blocks_per_row * sizeof(block_q5_0), x, 1, K);
#elif defined(__AVX2__)
        gemv_q5_0_avx2(&y[row], (const char *)blocks + row * blocks_per_row * sizeof(block_q5_0), x, 1, K);
#elif defined(__AVX__)
        gemv_q5_0_avx(&y[row], (const char *)blocks + row * blocks_per_row * sizeof(block_q5_0), x, 1, K);
#else
        gemv_q5_0_ref(&y[row], (const char *)blocks + row * blocks_per_row * sizeof(block_q5_0), x, 1, K);
#endif
    }
#else
    /* Fallback to reference parallel */
    gemv_q5_0_parallel(y, W, x, M, K, ith, nth);
#endif
}

References gemv_q5_0_parallel(), gemv_q5_0_ref(), and QK5_0.

◆ gemv_q5_0_q8_0()

void gemv_q5_0_q8_0	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K
	)

Matrix-vector multiply with Q5_0 weights and Q8_0 input.

Parameters

y	Output vector [M]
W	Weight matrix in Q5_0 format [M x K]
x_q8	Input vector in Q8_0 format [K]
M	Number of output rows
K	Number of columns (must be multiple of 32)

Definition at line 1640 of file gemm_kernels_q5_0.c.

{
    const block_q5_0 *w_blocks = (const block_q5_0 *)W;
    const block_q8_0 *x_blocks = (const block_q8_0 *)x_q8;
    const int blocks_per_row = K / QK5_0;
 
    for (int row = 0; row < M; row++) {
        vec_dot_q5_0_q8_0(K, &y[row],
                          &w_blocks[row * blocks_per_row],
                          x_blocks);
    }
}

References QK5_0, and vec_dot_q5_0_q8_0().

Referenced by ck_moe_q4k_mixed_route_work(), moe_swiglu_expert_forward_q4k_q5_0_workspace(), and moe_swiglu_shared_forward_q4k_q5_0_gated_workspace().

◆ gemv_q5_1()

void gemv_q5_1	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Auto-dispatch GEMV.

Definition at line 187 of file gemm_kernels_q5_1.c.

{
#ifdef __AVX512F__
    gemv_q5_1_avx512(y, W, x, M, K);
#else
    gemv_q5_1_ref(y, W, x, M, K);
#endif
}

References gemv_q5_1_ref().

Referenced by dot_q5_1(), and gemm_q5_1().

◆ gemv_q5_1_q8_1()

void gemv_q5_1_q8_1	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 269 of file gemm_kernels_q5_1_q8_1.c.

{
    if (!y || !W || !x || M <= 0 || K <= 0 || (K % QK5_1) != 0) {
        return;
    }
 
    const int blocks_per_row = K / QK5_1;
    if (blocks_per_row > CK_Q51_STACK_Q8_BLOCKS) {
        return;
    }
 
    block_q8_1 x_q8[CK_Q51_STACK_Q8_BLOCKS];
    quantize_row_q8_1_scalar(x, x_q8, K);
    gemv_q5_1_q8_1_ref(y, W, x_q8, M, K);
}

References CK_Q51_STACK_Q8_BLOCKS, gemv_q5_1_q8_1_ref(), QK5_1, and quantize_row_q8_1_scalar().

◆ gemv_q5_1_q8_1_ref()

void gemv_q5_1_q8_1_ref	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K
	)

Definition at line 219 of file gemm_kernels_q5_1_q8_1.c.

{
    if (!y || !W || !x_q8 || M <= 0 || K <= 0 || (K % QK5_1) != 0) {
        return;
    }
 
    const block_q5_1 *blocks = (const block_q5_1 *)W;
    const block_q8_1 *x = (const block_q8_1 *)x_q8;
    const int blocks_per_row = K / QK5_1;
 
    for (int row = 0; row < M; ++row) {
        const block_q5_1 *w_row = &blocks[row * blocks_per_row];
        float sum = 0.0f;
        for (int b = 0; b < blocks_per_row; ++b) {
            sum += dot_q5_1_q8_1_block(&w_row[b], &x[b]);
        }
        y[row] = sum;
    }
}

References dot_q5_1_q8_1_block(), and QK5_1.

Referenced by gemv_q5_1_q8_1().

◆ gemv_q5_k()

void gemv_q5_k	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 982 of file gemm_kernels_q5_k.c.

{
#if defined(__AVX512F__)
    /* TODO: AVX-512 implementation */
    gemv_q5_k_ref(y, W, x, M, K);
#elif defined(__AVX2__)
    /* TODO: AVX-2 implementation */
    gemv_q5_k_ref(y, W, x, M, K);
#elif defined(__AVX__)
    /* TODO: AVX implementation */
    gemv_q5_k_ref(y, W, x, M, K);
#elif defined(__SSE4_1__)
    /* TODO: SSE4.1 implementation */
    gemv_q5_k_ref(y, W, x, M, K);
#else
    gemv_q5_k_ref(y, W, x, M, K);
#endif
}

References gemv_q5_k_ref().

◆ gemv_q5_k_q8_k()

void gemv_q5_k_q8_k	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K
	)

Definition at line 892 of file gemm_kernels_q5_k.c.

{
#if defined(__AVX512F__)
    /* TODO: AVX-512 implementation */
    gemv_q5_k_q8_k_ref(y, W, x_q8, M, K);
#elif defined(__AVX2__)
    /* TODO: AVX-2 implementation */
    gemv_q5_k_q8_k_ref(y, W, x_q8, M, K);
#elif defined(__AVX__)
    /* TODO: AVX implementation */
    gemv_q5_k_q8_k_ref(y, W, x_q8, M, K);
#elif defined(__SSE4_1__)
    /* TODO: SSE4.1 implementation */
    gemv_q5_k_q8_k_ref(y, W, x_q8, M, K);
#else
    gemv_q5_k_q8_k_ref(y, W, x_q8, M, K);
#endif
}

References gemv_q5_k_q8_k_ref().

Referenced by ck_moe_q4k_q5k_route_work(), gemm_q5_k_q8_k_compact_rows4(), and moe_swiglu_expert_forward_q4k_q5k_workspace().

◆ gemv_q6_k()

void gemv_q6_k	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 169 of file gemm_kernels_q6k.c.

{
    if (!y || !W || !x) {
        return;
    }
    if (M <= 0 || K <= 0) {
        return;
    }
    // TEMPORARILY DISABLE NEW AVX KERNELS - USE REFERENCE ONLY
 
    const block_q6_K *blocks = (const block_q6_K *)W;
    const int blocks_per_row = K / QK_K;
 
    for (int row = 0; row < M; ++row) {
        const block_q6_K *w_row = blocks + (size_t)row * (size_t)blocks_per_row;
#if defined(__AVX__) && !defined(__AVX512F__)
        y[row] = dot_q6_k_avx(w_row, x, K);
#else
        y[row] = dot_q6_k_ref(w_row, x, K);
#endif
    }
}

References dot_q6_k_ref(), and QK_K.

Referenced by gemm_q6_k().

◆ gemv_q6_k_q8_k()

void gemv_q6_k_q8_k	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K
	)

GEMV: y = W @ x where W is Q6_K and x is Q8_K.

Definition at line 1341 of file gemm_kernels_q6k_q8k.c.

{
    if (ck_strict_parity_enabled() || ck_q6k_q8k_force_ref()) {
        gemv_q6_k_q8_k_ref(y, W, x_q8, M, K);
        return;
    }
 
#if defined(__AVX2__)
    /* llama.cpp's x86 Q6_K production graph keeps the AVX2 reduction order
     * even when AVX-512 is available. Wider ISA availability is not a license
     * to change this numerical contract. */
    gemv_q6_k_q8_k_avx2(y, W, x_q8, M, K);
    return;
#elif defined(__AVX__)
    gemv_q6_k_q8_k_avx(y, W, x_q8, M, K);
    return;
#elif defined(__SSE4_1__)
    gemv_q6_k_q8_k_sse(y, W, x_q8, M, K);
    return;
#endif
    gemv_q6_k_q8_k_ref(y, W, x_q8, M, K);
}

References ck_q6k_q8k_force_ref(), ck_strict_parity_enabled(), gemv_q6_k_q8_k_avx(), gemv_q6_k_q8_k_avx2(), gemv_q6_k_q8_k_ref(), and gemv_q6_k_q8_k_sse().

Referenced by ck_moe_q4k_mixed_route_work(), moe_swiglu_expert_forward_q4k_q6k_workspace(), and moe_swiglu_shared_forward_q4k_q6k_workspace().

◆ gemv_q6_k_q8_k_parallel()

void gemv_q6_k_q8_k_parallel	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K,
		int	ith,
		int	nth
	)

Parallel reference GEMV for Q6_K × Q8_K.

Caller provides ith (thread index) and nth (total threads). Each thread processes rows [r0, r1).

Definition at line 1398 of file gemm_kernels_q6k_q8k.c.

{
    if (!y || !W || !x_q8 || M <= 0 || K <= 0) return;
    if (ith < 0 || nth <= 0 || ith >= nth) return;
 
    /* Compute row range for this thread */
    const int dr = (M + nth - 1) / nth;
    const int r0 = dr * ith;
    const int r1 = (r0 + dr < M) ? (r0 + dr) : M;
 
    if (r0 >= M) return;
 
    const block_q6_K *blocks = (const block_q6_K *)W;
    const block_q8_K *x = (const block_q8_K *)x_q8;
    const int blocks_per_row = K / QK_K;
 
    for (int row = r0; row < r1; ++row) {
        const block_q6_K *w_row = blocks + (size_t)row * (size_t)blocks_per_row;
        y[row] = dot_q6_k_q8_k_ref(w_row, x, K);
    }
}

References dot_q6_k_q8_k_ref(), and QK_K.

◆ gemv_q6_k_q8_k_parallel_simd()

void gemv_q6_k_q8_k_parallel_simd	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K,
		int	ith,
		int	nth
	)

Parallel SIMD GEMV for Q6_K × Q8_K.

Uses best available SIMD (AVX/SSE) with row prefetching. Caller provides ith/nth from OpenMP region.

Definition at line 1430 of file gemm_kernels_q6k_q8k.c.

{
    if (!y || !W || !x_q8 || M <= 0 || K <= 0) return;
    if (ith < 0 || nth <= 0 || ith >= nth) return;
 
    const int dr = (M + nth - 1) / nth;
    const int r0 = dr * ith;
    const int r1 = (r0 + dr < M) ? (r0 + dr) : M;
 
    if (r0 >= M) return;
 
    const block_q6_K *blocks = (const block_q6_K *)W;
    const block_q8_K *x = (const block_q8_K *)x_q8;
    const int blocks_per_row = K / QK_K;
    const int strict = ck_strict_parity_enabled() || ck_q6k_q8k_force_ref();
 
    for (int row = r0; row < r1; ++row) {
        const block_q6_K *w_row = blocks + (size_t)row * (size_t)blocks_per_row;
#if defined(__AVX2__)
        y[row] = strict ? dot_q6_k_q8_k_ref(w_row, x, K)
                        : dot_q6_k_q8_k_avx2(w_row, x, K);
#elif defined(__AVX__)
        y[row] = strict ? dot_q6_k_q8_k_ref(w_row, x, K)
                        : dot_q6_k_q8_k_avx(w_row, x, K);
#elif defined(__SSE4_1__)
        y[row] = strict ? dot_q6_k_q8_k_ref(w_row, x, K)
                        : dot_q6_k_q8_k_sse(w_row, x, K);
#else
        y[row] = dot_q6_k_q8_k_ref(w_row, x, K);
#endif
    }
}

References ck_q6k_q8k_force_ref(), ck_strict_parity_enabled(), dot_q6_k_q8_k_ref(), and QK_K.

◆ gemv_q8_0()

void gemv_q8_0	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Auto-dispatch GEMV for Q8_0 weights based on CPU features.

Dispatch priority (best available):

AVX-512 (512-bit vectors) - Intel Skylake-X+
AVX2+FMA (256-bit vectors) - Intel Haswell+
AVX (256-bit vectors) - Intel Sandy Bridge+
SSE4.1 (128-bit vectors) - Intel Nehalem+
Reference (scalar) - Fallback

Uses ck_features.h for standardized feature detection.

Parameters

y	Output vector [M]
W	Weight matrix in Q8_0 format [M x K]
x	Input vector [K]
M	Number of output rows
K	Number of input columns (hidden dimension)

Definition at line 694 of file gemm_kernels_q8_0.c.

{
    if (ck_q8_0_debug_ref()) {
        gemv_q8_0_ref(y, W, x, M, K);
        return;
    }
 
// Dispatch order: AVX512 > AVX2 > AVX > SSE > ref
#if defined(__AVX512F__)
    gemv_q8_0_avx512(y, W, x, M, K);
#elif defined(__AVX2__)
    gemv_q8_0_avx2(y, W, x, M, K);
#elif defined(__AVX__)
    gemv_q8_0_avx(y, W, x, M, K);
#elif defined(__SSE4_1__)
    gemv_q8_0_sse(y, W, x, M, K);
#else
    gemv_q8_0_ref(y, W, x, M, K);
#endif
}

References ck_q8_0_debug_ref(), and gemv_q8_0_ref().

Referenced by gemv_q8_0_q8_0_contract().

◆ gemv_q8_0_q8_0()

void gemv_q8_0_q8_0	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K
	)

Matrix-vector multiply with Q8_0 weights and Q8_0 input.

Parameters

y	Output vector [M]
W	Weight matrix in Q8_0 format [M x K]
x_q8	Input vector in Q8_0 format [K]
M	Number of output rows
K	Number of columns (must be multiple of 32)

Definition at line 1405 of file gemm_kernels_q8_0.c.

{
    const block_q8_0 *w_blocks = (const block_q8_0 *)W;
    const block_q8_0 *x_blocks = (const block_q8_0 *)x_q8;
    const int blocks_per_row = K / QK8_0;
 
    for (int row = 0; row < M; row++) {
        vec_dot_q8_0_q8_0(K, &y[row],
                          &w_blocks[row * blocks_per_row],
                          x_blocks);
    }
}

References QK8_0, and vec_dot_q8_0_q8_0().

Referenced by ck_moe_q4k_mixed_route_work(), moe_swiglu_expert_forward_q4k_q8_0_workspace(), moe_swiglu_shared_forward_q4k_q8_0_gated_workspace(), and moe_swiglu_shared_forward_q8_0_gated_workspace().

◆ gemv_q8_0_q8_0_contract()

void gemv_q8_0_q8_0_contract	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Definition at line 383 of file gemm_kernels_q8_0_q8_0_contract.c.

{
    if (!y || !W || !x || M <= 0 || K <= 0) {
        return;
    }
 
    if ((K % QK8_0) != 0) {
        gemv_q8_0(y, W, x, M, K);
        return;
    }
 
    const int blocks_per_row = K / QK8_0;
    if (blocks_per_row > CK_Q80_STACK_Q8_BLOCKS) {
        gemv_q8_0(y, W, x, M, K);
        return;
    }
 
    block_q8_0 x_q8[CK_Q80_STACK_Q8_BLOCKS];
    if (ck_strict_parity_enabled()) {
        quantize_row_q8_0_ref_local(x, x_q8, K);
        gemv_q8_0_q8_0_ref_rows(y, W, x_q8, M, K);
        return;
    }
    quantize_row_q8_0(x, x_q8, K);
    gemv_q8_0_q8_0_x4(y, W, x_q8, M, K);
}

References CK_Q80_STACK_Q8_BLOCKS, ck_strict_parity_enabled(), gemv_q8_0(), gemv_q8_0_q8_0_ref_rows(), gemv_q8_0_q8_0_x4(), QK8_0, quantize_row_q8_0(), and quantize_row_q8_0_ref_local().

Referenced by gemm_nt_q8_0_q8_0_contract().

◆ gemv_q8_0_q8_0_x4()

void gemv_q8_0_q8_0_x4	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K
	)

Definition at line 1428 of file gemm_kernels_q8_0.c.

{
#if defined(__AVX2__) || defined(__AVX512F__)
    if (ck_q8_0_q8_0_debug_ref() || (K % QK8_0) != 0) {
        gemv_q8_0_q8_0(y, W, x_q8, M, K);
        return;
    }
 
    const block_q8_0 *w = (const block_q8_0 *)W;
    const block_q8_0 *x = (const block_q8_0 *)x_q8;
    const int nb = K / QK8_0;
    int row = 0;
    for (; row + 3 < M; row += 4) {
        __m256 acc0 = _mm256_setzero_ps();
        __m256 acc1 = _mm256_setzero_ps();
        __m256 acc2 = _mm256_setzero_ps();
        __m256 acc3 = _mm256_setzero_ps();
        const block_q8_0 *w0 = w + (size_t)(row + 0) * (size_t)nb;
        const block_q8_0 *w1 = w + (size_t)(row + 1) * (size_t)nb;
        const block_q8_0 *w2 = w + (size_t)(row + 2) * (size_t)nb;
        const block_q8_0 *w3 = w + (size_t)(row + 3) * (size_t)nb;
 
        for (int ib = 0; ib < nb; ++ib) {
            const __m256i qx = _mm256_loadu_si256((const __m256i *)x[ib].qs);
            const float dx = CK_FP16_TO_FP32(x[ib].d);
            const __m256i qw0 = _mm256_loadu_si256((const __m256i *)w0[ib].qs);
            const __m256i qw1 = _mm256_loadu_si256((const __m256i *)w1[ib].qs);
            const __m256i qw2 = _mm256_loadu_si256((const __m256i *)w2[ib].qs);
            const __m256i qw3 = _mm256_loadu_si256((const __m256i *)w3[ib].qs);
            const __m256 p0 = mul_sum_i8_pairs_float_q8_0_avx2(qw0, qx);
            const __m256 p1 = mul_sum_i8_pairs_float_q8_0_avx2(qw1, qx);
            const __m256 p2 = mul_sum_i8_pairs_float_q8_0_avx2(qw2, qx);
            const __m256 p3 = mul_sum_i8_pairs_float_q8_0_avx2(qw3, qx);
            const __m256 d0 = _mm256_set1_ps(CK_FP16_TO_FP32(w0[ib].d) * dx);
            const __m256 d1 = _mm256_set1_ps(CK_FP16_TO_FP32(w1[ib].d) * dx);
            const __m256 d2 = _mm256_set1_ps(CK_FP16_TO_FP32(w2[ib].d) * dx);
            const __m256 d3 = _mm256_set1_ps(CK_FP16_TO_FP32(w3[ib].d) * dx);
#if defined(__FMA__)
            acc0 = _mm256_fmadd_ps(d0, p0, acc0);
            acc1 = _mm256_fmadd_ps(d1, p1, acc1);
            acc2 = _mm256_fmadd_ps(d2, p2, acc2);
            acc3 = _mm256_fmadd_ps(d3, p3, acc3);
#else
            acc0 = _mm256_add_ps(_mm256_mul_ps(d0, p0), acc0);
            acc1 = _mm256_add_ps(_mm256_mul_ps(d1, p1), acc1);
            acc2 = _mm256_add_ps(_mm256_mul_ps(d2, p2), acc2);
            acc3 = _mm256_add_ps(_mm256_mul_ps(d3, p3), acc3);
#endif
        }
        y[row + 0] = hsum_float_8_q8_0(acc0);
        y[row + 1] = hsum_float_8_q8_0(acc1);
        y[row + 2] = hsum_float_8_q8_0(acc2);
        y[row + 3] = hsum_float_8_q8_0(acc3);
    }
    if (row < M) {
        gemv_q8_0_q8_0(y + row, w + (size_t)row * (size_t)nb,
                       x, M - row, K);
    }
#else
    gemv_q8_0_q8_0(y, W, x_q8, M, K);
#endif
}

References CK_FP16_TO_FP32, ck_q8_0_q8_0_debug_ref(), gemv_q8_0_q8_0(), and QK8_0.

Referenced by gemv_q8_0_q8_0_contract().

◆ gradient_accumulate_f32()

void gradient_accumulate_f32	(	float *	dst,
		const float *	src,
		size_t	numel
	)

Definition at line 925 of file optimizer_kernels.c.

{
    if (!dst || !src || numel == 0) {
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
    if (!pool || nth <= 1 || nth > CK_OPT_PAR_MAX_THREADS || numel < CK_OPT_PAR_MIN_NUMEL) {
        gradient_accumulate_f32_impl(dst, src, numel);
        return;
    }
    int active_nth = ck_opt_pick_active_threads(nth, numel, CK_OPT_PAR_MIN_NUMEL);
    if (active_nth <= 1) {
        gradient_accumulate_f32_impl(dst, src, numel);
        return;
    }
 
    ck_accum_parallel_args_t args = {
        .dst = dst,
        .src = src,
        .numel = numel,
    };
    ck_threadpool_dispatch_n(pool, active_nth, ck_accum_parallel_work, &args);
}

References ck_accum_parallel_work(), CK_OPT_PAR_MAX_THREADS, CK_OPT_PAR_MIN_NUMEL, ck_opt_pick_active_threads(), ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), and gradient_accumulate_f32_impl().

◆ gradient_accumulate_multi_f32()

void gradient_accumulate_multi_f32	(	float const	dsts,
		const float const	srcs,
		const size_t *	numels,
		int	tensor_count
	)

Definition at line 951 of file optimizer_kernels.c.

{
    if (!dsts || !srcs || !numels || tensor_count <= 0) {
        return;
    }
 
    size_t total_numel = 0;
    int valid_tensors = 0;
    for (int i = 0; i < tensor_count; ++i) {
        float *dst = dsts[i];
        const float *src = srcs[i];
        size_t n = numels[i];
        if (!dst || !src || n == 0) {
            continue;
        }
        total_numel += n;
        valid_tensors += 1;
    }
    if (total_numel == 0 || valid_tensors == 0) {
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
    if (!pool || nth <= 1 || nth > CK_OPT_PAR_MAX_THREADS || total_numel < CK_OPT_PAR_MIN_NUMEL) {
        for (int i = 0; i < tensor_count; ++i) {
            float *dst = dsts[i];
            const float *src = srcs[i];
            size_t n = numels[i];
            if (!dst || !src || n == 0) {
                continue;
            }
            gradient_accumulate_f32_impl(dst, src, n);
        }
        return;
    }
 
    int active_nth = ck_opt_pick_active_threads(nth, total_numel, CK_OPT_PAR_MIN_NUMEL);
    if (active_nth <= 1) {
        for (int i = 0; i < tensor_count; ++i) {
            float *dst = dsts[i];
            const float *src = srcs[i];
            size_t n = numels[i];
            if (!dst || !src || n == 0) {
                continue;
            }
            gradient_accumulate_f32_impl(dst, src, n);
        }
        return;
    }
 
    ck_accum_multi_parallel_args_t args = {
        .dsts = dsts,
        .srcs = srcs,
        .numels = numels,
        .tensor_count = tensor_count,
        .total_numel = total_numel,
    };
    ck_threadpool_dispatch_n(pool, active_nth, ck_accum_multi_parallel_work, &args);
}

References ck_accum_multi_parallel_work(), CK_OPT_PAR_MAX_THREADS, CK_OPT_PAR_MIN_NUMEL, ck_opt_pick_active_threads(), ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), and gradient_accumulate_f32_impl().

◆ gradient_clip_norm_f32()

float gradient_clip_norm_f32	(	float *	grad,
		size_t	numel,
		float	max_norm
	)

Clip gradient norm (fp32)

If ||grad||_2 > max_norm, scale grad so that ||grad||_2 = max_norm

Parameters

grad	Gradient tensor to clip (in-place) [numel]
numel	Number of elements
max_norm	Maximum allowed L2 norm

Returns: The original L2 norm before clipping

Definition at line 1169 of file optimizer_kernels.c.

{
    if (!grad || numel == 0 || max_norm <= 0.0f) {
        return 0.0f;
    }
 
    double sum_sq = 0.0;
    ck_threadpool_t *pool = ck_threadpool_global();
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
 
    if (pool && nth > 1 && nth <= CK_OPT_PAR_MAX_THREADS && numel >= CK_OPT_PAR_MIN_NUMEL) {
        int active_nth = ck_opt_pick_active_threads(nth, numel, CK_OPT_PAR_MIN_NUMEL);
        if (active_nth <= 1) {
            sum_sq = gradient_sum_sq_f32_impl(grad, numel);
        } else {
            ck_sum_sq_parallel_args_t args;
            args.grad = grad;
            args.numel = numel;
            for (int i = 0; i < CK_OPT_PAR_MAX_THREADS; ++i) {
                args.partial[i] = 0.0;
            }
            ck_threadpool_dispatch_n(pool, active_nth, ck_sum_sq_parallel_work, &args);
            for (int i = 0; i < active_nth; ++i) {
                sum_sq += args.partial[i];
            }
        }
    } else {
        sum_sq = gradient_sum_sq_f32_impl(grad, numel);
    }
 
    float norm = sqrtf((float)sum_sq);
    if (norm > max_norm) {
        float scale = max_norm / norm;
        gradient_scale_f32(grad, numel, scale);
    }
    return norm;
}

References CK_OPT_PAR_MAX_THREADS, CK_OPT_PAR_MIN_NUMEL, ck_opt_pick_active_threads(), ck_sum_sq_parallel_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), gradient_scale_f32(), and gradient_sum_sq_f32_impl().

◆ gradient_global_norm_multi_f32()

float gradient_global_norm_multi_f32	(	const float const	grads,
		const size_t *	numels,
		int	tensor_count
	)

Definition at line 1207 of file optimizer_kernels.c.

{
    if (!grads || !numels || tensor_count <= 0) {
        return 0.0f;
    }
 
    size_t total_numel = 0;
    int valid_tensors = 0;
    for (int i = 0; i < tensor_count; ++i) {
        if (!grads[i] || numels[i] == 0) {
            continue;
        }
        total_numel += numels[i];
        valid_tensors += 1;
    }
    if (total_numel == 0 || valid_tensors == 0) {
        return 0.0f;
    }
 
    double sum_sq = 0.0;
    ck_threadpool_t *pool = ck_threadpool_global();
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
 
    if (pool && nth > 1 && nth <= CK_OPT_PAR_MAX_THREADS &&
        total_numel >= CK_OPT_PAR_MIN_NUMEL && valid_tensors > 1) {
        int active_nth = ck_opt_pick_active_threads(nth, total_numel, CK_OPT_PAR_MIN_NUMEL);
        if (active_nth > valid_tensors) {
            active_nth = valid_tensors;
        }
        if (active_nth <= 1) {
            for (int i = 0; i < tensor_count; ++i) {
                const float *g = grads[i];
                size_t n = numels[i];
                if (!g || n == 0) {
                    continue;
                }
                sum_sq += gradient_sum_sq_f32_impl(g, n);
            }
        } else {
            ck_sum_sq_multi_parallel_args_t args;
            args.grads = grads;
            args.numels = numels;
            args.tensor_count = tensor_count;
            for (int i = 0; i < CK_OPT_PAR_MAX_THREADS; ++i) {
                args.partial[i] = 0.0;
            }
            ck_threadpool_dispatch_n(pool, active_nth, ck_sum_sq_multi_parallel_work, &args);
            for (int i = 0; i < active_nth; ++i) {
                sum_sq += args.partial[i];
            }
        }
    } else {
        for (int i = 0; i < tensor_count; ++i) {
            const float *g = grads[i];
            size_t n = numels[i];
            if (!g || n == 0) {
                continue;
            }
            sum_sq += gradient_sum_sq_f32_impl(g, n);
        }
    }
 
    if (!(sum_sq > 0.0)) {
        return 0.0f;
    }
    return sqrtf((float)sum_sq);
}

References CK_OPT_PAR_MAX_THREADS, CK_OPT_PAR_MIN_NUMEL, ck_opt_pick_active_threads(), ck_sum_sq_multi_parallel_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), and gradient_sum_sq_f32_impl().

Referenced by adamw_clip_update_multi_f32().

◆ gradient_scale_f32()

void gradient_scale_f32	(	float *	grad,
		size_t	numel,
		float	scale
	)

Definition at line 1073 of file optimizer_kernels.c.

{
    if (!grad || numel == 0) {
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int nth = pool ? ck_threadpool_n_threads(pool) : 1;
    if (!pool || nth <= 1 || nth > CK_OPT_PAR_MAX_THREADS || numel < CK_OPT_PAR_MIN_NUMEL) {
        gradient_scale_f32_impl(grad, numel, scale);
        return;
    }
    int active_nth = ck_opt_pick_active_threads(nth, numel, CK_OPT_PAR_MIN_NUMEL);
    if (active_nth <= 1) {
        gradient_scale_f32_impl(grad, numel, scale);
        return;
    }
 
    ck_scale_parallel_args_t args = {
        .grad = grad,
        .numel = numel,
        .scale = scale,
    };
    ck_threadpool_dispatch_n(pool, active_nth, ck_scale_parallel_work, &args);
}

References CK_OPT_PAR_MAX_THREADS, CK_OPT_PAR_MIN_NUMEL, ck_opt_pick_active_threads(), ck_scale_parallel_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), ck_threadpool_n_threads(), and gradient_scale_f32_impl().

◆ group_limited_topk_router_sigmoid_f32()

void group_limited_topk_router_sigmoid_f32	(	const float *	logits,
		const float *	correction_bias,
		int *	indices,
		float *	weights,
		int	rows,
		int	n_experts,
		int	top_k,
		int	n_group,
		int	topk_group,
		int	norm_topk_prob,
		float	routed_scaling_factor
	)

Definition at line 798 of file topk_kernels.c.

{
    group_limited_topk_router_f32_impl(
        logits, correction_bias, indices, weights, rows, n_experts, top_k,
        n_group, topk_group, norm_topk_prob, routed_scaling_factor, 1
    );
}

References group_limited_topk_router_f32_impl().

◆ hyper_connection_mix_bf16()

void hyper_connection_mix_bf16	(	const float *	hyper_input,
		const float *	norm_weight,
		const uint16_t *	mix_down_weight,
		const uint16_t *	mix_up_weight,
		const uint16_t *	inject_weight,
		float *	mixed_output,
		float *	injection_output,
		float *	normalized_scratch,
		float *	dynamic_scratch,
		float *	mix_scratch,
		int	rows,
		int	streams,
		int	hidden_dim,
		int	dynamic_dim,
		float	eps,
		int	emit_injection
	)

Definition at line 71 of file hyper_connection_kernels.c.

                                                   {
    if (!hyper_input || !norm_weight || !mix_down_weight || !mix_up_weight ||
        !mixed_output || !normalized_scratch || !dynamic_scratch || !mix_scratch ||
        rows <= 0 || streams <= 0 || hidden_dim <= 0 || dynamic_dim <= 0) {
        return;
    }
    if (emit_injection && (!inject_weight || !injection_output)) {
        return;
    }
 
    const int hyper_dim = streams * hidden_dim;
    const float inv_streams = 1.0f / (float)streams;
 
    for (int row = 0; row < rows; ++row) {
        const float *input_row =
            hyper_input + (size_t)row * (size_t)hyper_dim;
        float *norm_row =
            normalized_scratch + (size_t)row * (size_t)hyper_dim;
        float *dynamic_row =
            dynamic_scratch + (size_t)row * (size_t)dynamic_dim;
        float *mix_row = mix_scratch + (size_t)row * (size_t)hyper_dim;
 
        for (int stream = 0; stream < streams; ++stream) {
            const int base = stream * hidden_dim;
            rmsnorm_forward_qwen3next_pytorch_bf16_storage(
                input_row + base,
                norm_weight + base,
                norm_row + base,
                NULL,
                1,
                hidden_dim,
                hidden_dim,
                eps
            );
        }
 
        for (int out = 0; out < dynamic_dim; ++out) {
            const uint16_t *weight_row =
                mix_down_weight + (size_t)out * (size_t)hyper_dim;
            float sum = 0.0f;
            for (int col = 0; col < hyper_dim; ++col) {
                sum += norm_row[col] * bf16_to_float(weight_row[col]);
            }
            const float projected = ck_bf16_round(sum * inv_streams);
            dynamic_row[out] = ck_bf16_round(
                projected / (1.0f + expf(-projected)));
        }
 
        for (int out = 0; out < hyper_dim; ++out) {
            const uint16_t *weight_row =
                mix_up_weight + (size_t)out * (size_t)dynamic_dim;
            float sum = 0.0f;
            for (int col = 0; col < dynamic_dim; ++col) {
                sum += dynamic_row[col] * bf16_to_float(weight_row[col]);
            }
            mix_row[out] = ck_sigmoid_bf16(ck_bf16_round(sum));
        }
 
        float *mixed_row =
            mixed_output + (size_t)row * (size_t)hidden_dim;
        for (int col = 0; col < hidden_dim; ++col) {
            float sum = 0.0f;
            for (int stream = 0; stream < streams; ++stream) {
                const int index = stream * hidden_dim + col;
                sum += ck_bf16_round(norm_row[index] * mix_row[index]);
            }
            mixed_row[col] = ck_bf16_round(sum * inv_streams);
        }
 
        if (emit_injection) {
            float *injection_row =
                injection_output + (size_t)row * (size_t)streams;
            for (int stream = 0; stream < streams; ++stream) {
                const uint16_t *weight_row =
                    inject_weight + (size_t)stream * (size_t)hyper_dim;
                float sum = 0.0f;
                for (int col = 0; col < hyper_dim; ++col) {
                    sum += norm_row[col] * bf16_to_float(weight_row[col]);
                }
                injection_row[stream] = ck_bf16_round(
                    2.0f * ck_sigmoid_bf16(ck_bf16_round(sum * inv_streams)));
            }
        }
    }
}

References bf16_to_float(), ck_bf16_round(), ck_sigmoid_bf16(), and rmsnorm_forward_qwen3next_pytorch_bf16_storage().

◆ hyper_connection_mix_q4k_q5_0_q4k()

void hyper_connection_mix_q4k_q5_0_q4k	(	const float *	hyper_input,
		const float *	norm_weight,
		const void *	mix_down_weight,
		const void *	mix_up_weight,
		const void *	inject_weight,
		float *	mixed_output,
		float *	injection_output,
		float *	normalized_scratch,
		float *	dynamic_scratch,
		float *	mix_scratch,
		int	rows,
		int	streams,
		int	hidden_dim,
		int	dynamic_dim,
		float	eps,
		int	emit_injection
	)

Definition at line 344 of file hyper_connection_kernels.c.

                                                           {
    hyper_connection_mix_quantized(
        hyper_input, norm_weight, mix_down_weight, mix_up_weight, inject_weight,
        mixed_output, injection_output, normalized_scratch, dynamic_scratch,
        mix_scratch, rows, streams, hidden_dim, dynamic_dim, eps,
        emit_injection, hyper_injection_q4k_q8k_llama_dispatch,
        gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch);
}

References gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch(), hyper_connection_mix_quantized(), and hyper_injection_q4k_q8k_llama_dispatch().

◆ hyper_connection_mix_q6k_q5_0_q4k()

void hyper_connection_mix_q6k_q5_0_q4k	(	const float *	hyper_input,
		const float *	norm_weight,
		const void *	mix_down_weight,
		const void *	mix_up_weight,
		const void *	inject_weight,
		float *	mixed_output,
		float *	injection_output,
		float *	normalized_scratch,
		float *	dynamic_scratch,
		float *	mix_scratch,
		int	rows,
		int	streams,
		int	hidden_dim,
		int	dynamic_dim,
		float	eps,
		int	emit_injection
	)

Definition at line 368 of file hyper_connection_kernels.c.

                                                           {
    hyper_connection_mix_quantized(
        hyper_input, norm_weight, mix_down_weight, mix_up_weight, inject_weight,
        mixed_output, injection_output, normalized_scratch, dynamic_scratch,
        mix_scratch, rows, streams, hidden_dim, dynamic_dim, eps,
        emit_injection, hyper_injection_q4k_q8k_llama_dispatch,
        gemm_nt_q6_k_q8_k_parallel_dispatch);
}

References gemm_nt_q6_k_q8_k_parallel_dispatch(), hyper_connection_mix_quantized(), and hyper_injection_q4k_q8k_llama_dispatch().

◆ hyper_stream_expand_bf16()

void hyper_stream_expand_bf16	(	const float *	input,
		float *	output,
		int	rows,
		int	streams,
		int	hidden_dim
	)

Definition at line 51 of file hyper_connection_kernels.c.

                                              {
    if (!input || !output || rows <= 0 || streams <= 0 || hidden_dim <= 0) {
        return;
    }
    for (int row = 0; row < rows; ++row) {
        const float *src = input + (size_t)row * (size_t)hidden_dim;
        float *dst = output + (size_t)row * (size_t)streams * (size_t)hidden_dim;
        for (int stream = 0; stream < streams; ++stream) {
            for (int col = 0; col < hidden_dim; ++col) {
                dst[(size_t)stream * (size_t)hidden_dim + (size_t)col] =
                    ck_bf16_round(src[col]);
            }
        }
    }
}

References ck_bf16_round().

◆ hyper_stream_expand_f32()

void hyper_stream_expand_f32	(	const float *	input,
		float *	output,
		int	rows,
		int	streams,
		int	hidden_dim
	)

Definition at line 32 of file hyper_connection_kernels.c.

                                             {
    if (!input || !output || rows <= 0 || streams <= 0 || hidden_dim <= 0) {
        return;
    }
    for (int row = 0; row < rows; ++row) {
        const float *src = input + (size_t)row * (size_t)hidden_dim;
        float *dst = output + (size_t)row * (size_t)streams * (size_t)hidden_dim;
        for (int stream = 0; stream < streams; ++stream) {
            for (int col = 0; col < hidden_dim; ++col) {
                dst[(size_t)stream * (size_t)hidden_dim + (size_t)col] = src[col];
            }
        }
    }
}

◆ hyper_stream_inject_bf16()

void hyper_stream_inject_bf16	(	const float *	hyper_input,
		const float *	block_output,
		const float *	injection_weight,
		float *	output,
		int	rows,
		int	streams,
		int	hidden_dim
	)

Definition at line 392 of file hyper_connection_kernels.c.

                                              {
    if (!hyper_input || !block_output || !injection_weight || !output ||
        rows <= 0 || streams <= 0 || hidden_dim <= 0) {
        return;
    }
    const int hyper_dim = streams * hidden_dim;
    for (int row = 0; row < rows; ++row) {
        const float *hyper_row =
            hyper_input + (size_t)row * (size_t)hyper_dim;
        const float *block_row =
            block_output + (size_t)row * (size_t)hidden_dim;
        const float *inject_row =
            injection_weight + (size_t)row * (size_t)streams;
        float *output_row = output + (size_t)row * (size_t)hyper_dim;
        for (int stream = 0; stream < streams; ++stream) {
            for (int col = 0; col < hidden_dim; ++col) {
                const int index = stream * hidden_dim + col;
                output_row[index] = ck_bf16_round(
                    hyper_row[index] +
                    ck_bf16_round(block_row[col] * inject_row[stream]));
            }
        }
    }
}

References ck_bf16_round().

◆ hyper_stream_inject_f32()

void hyper_stream_inject_f32	(	const float *	hyper_input,
		const float *	block_output,
		const float *	injection_weight,
		float *	output,
		int	rows,
		int	streams,
		int	hidden_dim
	)

Definition at line 423 of file hyper_connection_kernels.c.

                                             {
    if (!hyper_input || !block_output || !injection_weight || !output ||
        rows <= 0 || streams <= 0 || hidden_dim <= 0) {
        return;
    }
    const int hyper_dim = streams * hidden_dim;
    for (int row = 0; row < rows; ++row) {
        const float *hyper_row =
            hyper_input + (size_t)row * (size_t)hyper_dim;
        const float *block_row =
            block_output + (size_t)row * (size_t)hidden_dim;
        const float *inject_row =
            injection_weight + (size_t)row * (size_t)streams;
        float *output_row = output + (size_t)row * (size_t)hyper_dim;
        for (int stream = 0; stream < streams; ++stream) {
            for (int col = 0; col < hidden_dim; ++col) {
                const int index = stream * hidden_dim + col;
                volatile const float weighted =
                    block_row[col] * inject_row[stream];
                output_row[index] = hyper_row[index] + weighted;
            }
        }
    }
}

◆ im2patch()

void im2patch	(	const float *	image,
		float *	patches,
		int	C,
		int	H,
		int	W,
		int	P
	)

im2patch: Transforms an image into a sequence of flattened patches.

Image Layout: [C, H, W] (Row-major: W is fastest moving) Output Layout: [num_patches, C * P * P]

num_patches = (H/P) * (W/P) P = patch_size

Definition at line 82 of file vision_kernels.c.

{
    int num_patches_h = H / P;
    int num_patches_w = W / P;
    int patch_dim = C * P * P;
 
    // ph, pw: patch grid coordinates
    for (int ph = 0; ph < num_patches_h; ++ph) {
        for (int pw = 0; pw < num_patches_w; ++pw) {
            
            int patch_idx = ph * num_patches_w + pw;
            float *dst_patch = patches + (size_t)patch_idx * patch_dim;
 
            // For each patch, grab pixels from all channels
            for (int c = 0; c < C; ++c) {
                for (int py = 0; py < P; ++py) {
                    int y = ph * P + py;
                    int x = pw * P;
                    
                    // Input row start in the image
                    const float *src_row = image + (size_t)c * H * W + (size_t)y * W + x;
                    
                    // Destination row in the flattened patch sequence
                    float *dst_row = dst_patch + (size_t)c * P * P + (size_t)py * P;
                    
                    // Copy P pixels (one row of the patch)
                    memcpy(dst_row, src_row, P * sizeof(float));
                }
            }
        }
    }
}

References C.

◆ im2patch_bf16()

void im2patch_bf16	(	const uint16_t *	image,
		uint16_t *	patches,
		int	C,
		int	H,
		int	W,
		int	P
	)

Definition at line 22 of file vision_kernels_bf16.c.

{
    if (!image || !patches || C <= 0 || H <= 0 || W <= 0 || P <= 0) {
        return;
    }
 
    int num_patches_h = H / P;
    int num_patches_w = W / P;
    int patch_dim = C * P * P;
 
    for (int ph = 0; ph < num_patches_h; ++ph) {
        for (int pw = 0; pw < num_patches_w; ++pw) {
            int patch_idx = ph * num_patches_w + pw;
            uint16_t *dst_patch = patches + (size_t)patch_idx * (size_t)patch_dim;
 
            for (int c = 0; c < C; ++c) {
                for (int py = 0; py < P; ++py) {
                    int y = ph * P + py;
                    int x = pw * P;
 
                    const uint16_t *src_row = image + (size_t)c * (size_t)H * (size_t)W + (size_t)y * (size_t)W + (size_t)x;
                    uint16_t *dst_row = dst_patch + (size_t)c * (size_t)P * (size_t)P + (size_t)py * (size_t)P;
 
                    memcpy(dst_row, src_row, (size_t)P * sizeof(uint16_t));
                }
            }
        }
    }
}

References C.

◆ kv_cache_repack_head_major_inplace()

void kv_cache_repack_head_major_inplace	(	float *	buf,
		int	num_heads,
		int	tokens,
		int	cache_capacity,
		int	aligned_head_dim
	)

Definition at line 49 of file kv_cache_kernels.c.

{
    if (!buf) {
        return;
    }
    if (num_heads <= 0 || tokens <= 0 || cache_capacity <= 0 || aligned_head_dim <= 0) {
        return;
    }
    if (tokens > cache_capacity) {
        tokens = cache_capacity;
    }
    if (tokens == cache_capacity) {
        return;
    }
 
    const size_t old_head_stride = (size_t)tokens * (size_t)aligned_head_dim;
    const size_t new_head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
    const size_t bytes = (size_t)tokens * (size_t)aligned_head_dim * sizeof(float);
 
    // Move head blocks from high to low to avoid overwriting source data
    // for heads that have not yet been moved.
    for (int h = num_heads - 1; h >= 0; --h) {
        float *src = buf + (size_t)h * old_head_stride;
        float *dst = buf + (size_t)h * new_head_stride;
        memmove(dst, src, bytes);
    }
}

Referenced by qwen2_0_5b_decode_forward_prefill_impl(), qwen2_0_5b_decode_forward_prefill_impl(), and qwen2_0_5b_decode_forward_prefill_impl().

◆ kv_cache_store()

void kv_cache_store	(	float *__restrict	kv_cache_k,
		float *__restrict	kv_cache_v,
		const float *__restrict	k,
		const float *__restrict	v,
		int	layer,
		int	pos,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 122 of file kv_cache_kernels.c.

{
    (void)layer;
    kv_cache_write_head_major(k, v,
                              kv_cache_k, kv_cache_v,
                              num_kv_heads,
                              pos,
                              max_seq_len,
                              head_dim,
                              head_dim);
}

References kv_cache_write_head_major().

◆ kv_cache_store_batch_bf16()

void kv_cache_store_batch_bf16	(	uint16_t *__restrict	kv_cache_k,
		uint16_t *__restrict	kv_cache_v,
		const float *__restrict	k,
		const float *__restrict	v,
		int	start_pos,
		int	num_tokens,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 291 of file kv_cache_kernels.c.

{
    if (!kv_cache_k || !kv_cache_v || !k || !v ||
        start_pos < 0 || num_tokens <= 0 || num_kv_heads <= 0 ||
        head_dim <= 0 || max_seq_len <= 0 ||
        start_pos > max_seq_len - num_tokens) {
        return;
    }
 
    const size_t compact_head_stride = (size_t)num_tokens * (size_t)head_dim;
    const size_t cache_head_stride = (size_t)max_seq_len * (size_t)head_dim;
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_head = k + (size_t)h * compact_head_stride;
        const float *v_head = v + (size_t)h * compact_head_stride;
        uint16_t *k_head_cache = kv_cache_k + (size_t)h * cache_head_stride;
        uint16_t *v_head_cache = kv_cache_v + (size_t)h * cache_head_stride;
        for (int t = 0; t < num_tokens; ++t) {
            const size_t src_offset = (size_t)t * (size_t)head_dim;
            const size_t dst_offset = (size_t)(start_pos + t) * (size_t)head_dim;
            ck_local_fp32_to_bf16_row(k_head + src_offset, k_head_cache + dst_offset, head_dim);
            ck_local_fp32_to_bf16_row(v_head + src_offset, v_head_cache + dst_offset, head_dim);
        }
    }
}

References ck_local_fp32_to_bf16_row().

◆ kv_cache_store_batch_f16()

void kv_cache_store_batch_f16	(	uint16_t *__restrict	kv_cache_k,
		uint16_t *__restrict	kv_cache_v,
		const float *__restrict	k,
		const float *__restrict	v,
		int	start_pos,
		int	num_tokens,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 255 of file kv_cache_kernels.c.

{
    if (!kv_cache_k || !kv_cache_v || !k || !v) {
        return;
    }
    if (start_pos < 0 || num_tokens <= 0 || num_kv_heads <= 0 ||
        head_dim <= 0 || max_seq_len <= 0 ||
        start_pos > max_seq_len - num_tokens) {
        return;
    }
 
    const size_t compact_head_stride = (size_t)num_tokens * (size_t)head_dim;
    const size_t cache_head_stride = (size_t)max_seq_len * (size_t)head_dim;
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_head = k + (size_t)h * compact_head_stride;
        const float *v_head = v + (size_t)h * compact_head_stride;
        uint16_t *k_head_cache = kv_cache_k + (size_t)h * cache_head_stride;
        uint16_t *v_head_cache = kv_cache_v + (size_t)h * cache_head_stride;
        for (int t = 0; t < num_tokens; ++t) {
            const size_t src_offset = (size_t)t * (size_t)head_dim;
            const size_t dst_offset = (size_t)(start_pos + t) * (size_t)head_dim;
            ck_local_fp32_to_fp16_row(k_head + src_offset, k_head_cache + dst_offset, head_dim);
            ck_local_fp32_to_fp16_row(v_head + src_offset, v_head_cache + dst_offset, head_dim);
        }
    }
}

References ck_local_fp32_to_fp16_row().

◆ kv_cache_store_batch_f32()

void kv_cache_store_batch_f32	(	float *__restrict	kv_cache_k,
		float *__restrict	kv_cache_v,
		const float *__restrict	k,
		const float *__restrict	v,
		int	start_pos,
		int	num_tokens,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 223 of file kv_cache_kernels.c.

{
    if (!kv_cache_k || !kv_cache_v || !k || !v ||
        start_pos < 0 || num_tokens <= 0 || num_kv_heads <= 0 ||
        head_dim <= 0 || max_seq_len <= 0 ||
        start_pos > max_seq_len - num_tokens) {
        return;
    }
 
    const size_t compact_head_stride = (size_t)num_tokens * (size_t)head_dim;
    const size_t cache_head_stride = (size_t)max_seq_len * (size_t)head_dim;
    const size_t token_bytes = (size_t)num_tokens * (size_t)head_dim * sizeof(float);
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_head = k + (size_t)h * compact_head_stride;
        const float *v_head = v + (size_t)h * compact_head_stride;
        float *k_head_cache = kv_cache_k + (size_t)h * cache_head_stride
                            + (size_t)start_pos * (size_t)head_dim;
        float *v_head_cache = kv_cache_v + (size_t)h * cache_head_stride
                            + (size_t)start_pos * (size_t)head_dim;
        memcpy(k_head_cache, k_head, token_bytes);
        memcpy(v_head_cache, v_head, token_bytes);
    }
}

◆ kv_cache_store_bf16()

void kv_cache_store_bf16	(	uint16_t *__restrict	kv_cache_k,
		uint16_t *__restrict	kv_cache_v,
		const float *__restrict	k,
		const float *__restrict	v,
		int	layer,
		int	pos,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 195 of file kv_cache_kernels.c.

{
    (void)layer;
    if (!kv_cache_k || !kv_cache_v || !k || !v ||
        num_kv_heads <= 0 || pos < 0 || pos >= max_seq_len ||
        head_dim <= 0 || max_seq_len <= 0) {
        return;
    }
 
    const size_t head_stride = (size_t)max_seq_len * (size_t)head_dim;
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_src = k + (size_t)h * (size_t)head_dim;
        const float *v_src = v + (size_t)h * (size_t)head_dim;
        uint16_t *k_dst = kv_cache_k + (size_t)h * head_stride + (size_t)pos * (size_t)head_dim;
        uint16_t *v_dst = kv_cache_v + (size_t)h * head_stride + (size_t)pos * (size_t)head_dim;
        ck_local_fp32_to_bf16_row(k_src, k_dst, head_dim);
        ck_local_fp32_to_bf16_row(v_src, v_dst, head_dim);
    }
}

References ck_local_fp32_to_bf16_row().

◆ kv_cache_store_f16()

void kv_cache_store_f16	(	uint16_t *__restrict	kv_cache_k,
		uint16_t *__restrict	kv_cache_v,
		const float *__restrict	k,
		const float *__restrict	v,
		int	layer,
		int	pos,
		int	num_kv_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 161 of file kv_cache_kernels.c.

{
    (void)layer;
    if (!kv_cache_k || !kv_cache_v || !k || !v) {
        return;
    }
    if (num_kv_heads <= 0 || pos < 0 || head_dim <= 0 || max_seq_len <= 0) {
        return;
    }
    if (pos >= max_seq_len) {
        return;
    }
 
    const size_t head_stride = (size_t)max_seq_len * (size_t)head_dim;
    const size_t token_stride = (size_t)head_dim;
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_src = k + (size_t)h * token_stride;
        const float *v_src = v + (size_t)h * token_stride;
        uint16_t *k_dst = kv_cache_k + (size_t)h * head_stride + (size_t)pos * token_stride;
        uint16_t *v_dst = kv_cache_v + (size_t)h * head_stride + (size_t)pos * token_stride;
        ck_local_fp32_to_fp16_row(k_src, k_dst, head_dim);
        ck_local_fp32_to_fp16_row(v_src, v_dst, head_dim);
    }
}

References ck_local_fp32_to_fp16_row().

◆ kv_cache_store_shared_q()

void kv_cache_store_shared_q	(	float *__restrict	kv_cache_k,
		float *__restrict	kv_cache_v,
		const float *__restrict	q,
		int	layer,
		int	pos,
		int	num_heads,
		int	head_dim,
		int	max_seq_len
	)

Definition at line 142 of file kv_cache_kernels.c.

{
    (void)layer;
    kv_cache_write_head_major(q, q,
                              kv_cache_k, kv_cache_v,
                              num_heads,
                              pos,
                              max_seq_len,
                              head_dim,
                              head_dim);
}

References kv_cache_write_head_major().

◆ kv_cache_write_head_major()

void kv_cache_write_head_major	(	const float *__restrict	k_token,
		const float *__restrict	v_token,
		float *__restrict	k_cache,
		float *__restrict	v_cache,
		int	num_kv_heads,
		int	token_index,
		int	cache_capacity,
		int	head_dim,
		int	aligned_head_dim
	)

Definition at line 81 of file kv_cache_kernels.c.

{
    if (!k_token || !v_token || !k_cache || !v_cache) {
        return;
    }
    if (num_kv_heads <= 0 || token_index < 0 || cache_capacity <= 0) {
        return;
    }
    if (token_index >= cache_capacity || head_dim <= 0 || aligned_head_dim <= 0) {
        return;
    }
 
    const size_t head_stride = (size_t)cache_capacity * (size_t)aligned_head_dim;
    const size_t token_stride = (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_kv_heads; ++h) {
        const float *k_src = k_token + (size_t)h * token_stride;
        const float *v_src = v_token + (size_t)h * token_stride;
 
        float *k_dst = k_cache + (size_t)h * head_stride + (size_t)token_index * token_stride;
        float *v_dst = v_cache + (size_t)h * head_stride + (size_t)token_index * token_stride;
 
        for (int d = 0; d < head_dim; ++d) {
            k_dst[d] = k_src[d];
            v_dst[d] = v_src[d];
        }
        for (int d = head_dim; d < aligned_head_dim; ++d) {
            k_dst[d] = 0.0f;
            v_dst[d] = 0.0f;
        }
    }
}

◆ layernorm_backward_kernel()

void layernorm_backward_kernel	(	const float *	d_output,
		const float *	input,
		const float *	gamma,
		const float *	mean,
		const float *	rstd,
		float *	d_input,
		float *	d_gamma,
		float *	d_beta,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim
	)

Definition at line 1007 of file layernorm_kernels.c.

{
    int T = tokens;
    int D = d_model;
    int aligned_D = aligned_embed_dim;
 
    // Per-token input gradients
    for (int t = 0; t < T; ++t) {
        float mean_t = mean[t];
        float rstd_t = rstd[t];
 
        float d_y_gamma_sum = 0.0f;
        float d_y_gamma_xhat_sum = 0.0f;
 
        // First pass: compute sums
        for (int d = 0; d < D; ++d) {
            float x = input[t * aligned_D + d];
            float x_hat = (x - mean_t) * rstd_t;
            float d_y = d_output[t * aligned_D + d];
            float d_y_gamma = d_y * gamma[d];
 
            d_y_gamma_sum += d_y_gamma;
            d_y_gamma_xhat_sum += d_y_gamma * x_hat;
        }
 
        // Second pass: compute input gradients
        float scale = rstd_t / (float)D;
        for (int d = 0; d < D; ++d) {
            float x = input[t * aligned_D + d];
            float x_hat = (x - mean_t) * rstd_t;
            float d_y = d_output[t * aligned_D + d];
 
            d_input[t * aligned_D + d] =
                scale * ((float)D * d_y * gamma[d] - d_y_gamma_sum - x_hat * d_y_gamma_xhat_sum);
        }
 
        // Zero padding for aligned dimension beyond D
        for (int d = D; d < aligned_D; ++d) {
            d_input[t * aligned_D + d] = 0.0f;
        }
    }
 
    // Parameter gradients (gamma, beta)
    for (int d = 0; d < D; ++d) {
        float gamma_grad = 0.0f;
        float beta_grad = 0.0f;
 
        for (int t = 0; t < T; ++t) {
            float x = input[t * aligned_D + d];
            float x_hat = (x - mean[t]) * rstd[t];
            float d_y = d_output[t * aligned_D + d];
 
            gamma_grad += d_y * x_hat;
            beta_grad += d_y;
        }
 
        d_gamma[d] += gamma_grad;
        d_beta[d] += beta_grad;
    }
}

Referenced by layernorm_backward_kernel_bf16().

◆ layernorm_backward_kernel_bf16()

void layernorm_backward_kernel_bf16	(	const uint16_t *	d_output,
		const uint16_t *	input,
		const float *	gamma,
		const float *	mean,
		const float *	rstd,
		uint16_t *	d_input,
		float *	d_gamma,
		float *	d_beta,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float *	scratch_d_output,
		float *	scratch_input,
		float *	scratch_d_input
	)

Definition at line 84 of file layernorm_kernels_bf16.c.

{
    if (!scratch_d_output || !scratch_input || !scratch_d_input) return;
 
    size_t total = (size_t)tokens * (size_t)aligned_embed_dim;
 
    bf16_tensor_to_float(d_output, scratch_d_output, total);
    bf16_tensor_to_float(input, scratch_input, total);
 
    layernorm_backward_kernel(scratch_d_output, scratch_input, gamma, mean, rstd,
                              scratch_d_input, d_gamma, d_beta,
                              tokens, d_model, aligned_embed_dim);
 
    float_tensor_to_bf16(scratch_d_input, d_input, total);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and layernorm_backward_kernel().

◆ layernorm_forward_rolled_slice()

void layernorm_forward_rolled_slice	(	const float *__restrict	input_slice_base,
		const float *__restrict	gamma,
		const float *__restrict	beta,
		float *__restrict	output_slice_base,
		float *__restrict	mean_cache_slice,
		float *__restrict	rstd_cache_slice,
		int	num_tokens_in_slice,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 398 of file layernorm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        layernorm_forward_ggml_exact(input_slice_base, gamma, beta,
                                     output_slice_base, mean_cache_slice, rstd_cache_slice,
                                     num_tokens_in_slice, d_model,
                                     aligned_embed_dim, aligned_embed_dim, aligned_embed_dim, eps);
        return;
    }
 
#if defined(__AVX512F__)
    layernorm_forward_rolled_slice_avx512(input_slice_base, gamma, beta,
                                           output_slice_base, mean_cache_slice, rstd_cache_slice,
                                           num_tokens_in_slice, d_model, aligned_embed_dim, eps);
#elif defined(__AVX2__) || defined(__AVX__)
    layernorm_forward_rolled_slice_avx256(input_slice_base, gamma, beta,
                                           output_slice_base, mean_cache_slice, rstd_cache_slice,
                                           num_tokens_in_slice, d_model, aligned_embed_dim, eps);
#else
    layernorm_naive_serial(input_slice_base, gamma, beta,
                           output_slice_base, mean_cache_slice, rstd_cache_slice,
                           num_tokens_in_slice, d_model, aligned_embed_dim, eps);
#endif
}

References ck_strict_parity_enabled(), layernorm_forward_ggml_exact(), and layernorm_naive_serial().

Referenced by layernorm_forward_rolled_slice_bf16().

◆ layernorm_forward_rolled_slice_bf16()

void layernorm_forward_rolled_slice_bf16	(	const uint16_t *__restrict	input_slice_base,
		const float *__restrict	gamma,
		const float *__restrict	beta,
		uint16_t *__restrict	output_slice_base,
		float *__restrict	mean_cache_slice,
		float *__restrict	rstd_cache_slice,
		int	num_tokens_in_slice,
		int	d_model,
		int	aligned_embed_dim,
		float	eps,
		float *	scratch_input,
		float *	scratch_output
	)

Definition at line 30 of file layernorm_kernels_bf16.c.

{
    if (!scratch_input || !scratch_output) return;
 
    size_t total = (size_t)num_tokens_in_slice * (size_t)aligned_embed_dim;
 
    bf16_tensor_to_float(input_slice_base, scratch_input, total);
    layernorm_forward_rolled_slice(scratch_input, gamma, beta,
                                   scratch_output, mean_cache_slice, rstd_cache_slice,
                                   num_tokens_in_slice, d_model, aligned_embed_dim, eps);
    float_tensor_to_bf16(scratch_output, output_slice_base, total);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and layernorm_forward_rolled_slice().

◆ layernorm_forward_unrolled_slice()

void layernorm_forward_unrolled_slice	(	const float *__restrict	input_slice_base,
		const float *__restrict	gamma,
		const float *__restrict	beta,
		float *__restrict	output_slice_base,
		float *__restrict	mean_cache_slice,
		float *__restrict	rstd_cache_slice,
		int	num_tokens_in_slice,
		int	d_model,
		float	eps
	)

Definition at line 730 of file layernorm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        layernorm_forward_ggml_exact(input_slice_base, gamma, beta,
                                     output_slice_base, mean_cache_slice, rstd_cache_slice,
                                     num_tokens_in_slice, d_model,
                                     d_model, d_model, d_model, eps);
        return;
    }
 
#if defined(__AVX512F__)
    layernorm_forward_unrolled_slice_avx512(input_slice_base, gamma, beta,
                                             output_slice_base, mean_cache_slice, rstd_cache_slice,
                                             num_tokens_in_slice, d_model, eps);
#elif defined(__AVX2__) || defined(__AVX__)
    layernorm_forward_unrolled_slice_avx256(input_slice_base, gamma, beta,
                                             output_slice_base, mean_cache_slice, rstd_cache_slice,
                                             num_tokens_in_slice, d_model, eps);
#else
    layernorm_forward_unrolled_slice_scalar(input_slice_base, gamma, beta,
                                            output_slice_base, mean_cache_slice, rstd_cache_slice,
                                            num_tokens_in_slice, d_model, eps);
#endif
}

References ck_strict_parity_enabled(), layernorm_forward_ggml_exact(), and layernorm_forward_unrolled_slice_scalar().

Referenced by layernorm_forward_unrolled_slice_bf16().

◆ layernorm_forward_unrolled_slice_bf16()

void layernorm_forward_unrolled_slice_bf16	(	const uint16_t *__restrict	input_slice_base,
		const float *__restrict	gamma,
		const float *__restrict	beta,
		uint16_t *__restrict	output_slice_base,
		float *__restrict	mean_cache_slice,
		float *__restrict	rstd_cache_slice,
		int	num_tokens_in_slice,
		int	d_model,
		float	eps,
		float *	scratch_input,
		float *	scratch_output
	)

Definition at line 57 of file layernorm_kernels_bf16.c.

{
    if (!scratch_input || !scratch_output) return;
 
    size_t total = (size_t)num_tokens_in_slice * (size_t)d_model;
 
    bf16_tensor_to_float(input_slice_base, scratch_input, total);
    layernorm_forward_unrolled_slice(scratch_input, gamma, beta,
                                     scratch_output, mean_cache_slice, rstd_cache_slice,
                                     num_tokens_in_slice, d_model, eps);
    float_tensor_to_bf16(scratch_output, output_slice_base, total);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and layernorm_forward_unrolled_slice().

◆ layernorm_naive_serial()

void layernorm_naive_serial	(	const float *	input,
		const float *	gamma,
		const float *	beta,
		float *	output,
		float *	mean_cache,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 175 of file layernorm_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *in_ptr = input + t * aligned_embed_dim;
        float *out_ptr = output + t * aligned_embed_dim;
 
        float sum_val = 0.0f;
        for (int i = 0; i < d_model; ++i) {
            sum_val += in_ptr[i];
        }
        float mean = sum_val / (float)d_model;
 
        float sum_sq_diff = 0.0f;
        for (int i = 0; i < d_model; ++i) {
            float diff = in_ptr[i] - mean;
            sum_sq_diff += diff * diff;
        }
        float variance = sum_sq_diff / (float)d_model + eps;
 
        double var_double = (double)variance;
        float inv_std = (float)(1.0 / sqrt(var_double));
 
        for (int i = 0; i < d_model; ++i) {
            float normalized_val = (in_ptr[i] - mean) * inv_std;
            out_ptr[i] = normalized_val * gamma[i] + beta[i];
        }
 
        if (mean_cache) {
            mean_cache[t] = mean;
        }
        if (rstd_cache) {
            rstd_cache[t] = inv_std;
        }
        /* Keep aligned padding quiet so future GEMMs see deterministic memory. */
        if (aligned_embed_dim > d_model) {
            /* Keep padded lanes zeroed so subsequent GEMMs never read stale data. */
            for (int i = d_model; i < aligned_embed_dim; ++i) {
                out_ptr[i] = 0.0f;
            }
        }
    }
}

Referenced by layernorm_forward_rolled_slice().

◆ layernorm_naive_serial_bf16_storage()

void layernorm_naive_serial_bf16_storage	(	const float *	input,
		const float *	gamma,
		const float *	beta,
		float *	output,
		float *	mean_cache,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		float	eps
	)

Definition at line 786 of file layernorm_kernels.c.

{
    const size_t count = (size_t)tokens * (size_t)d_model;
    for (size_t i = 0; i < count; ++i) {
        output[i] = bf16_to_float(float_to_bf16(input[i]));
    }
    layernorm_forward_ggml_exact(output, gamma, beta,
                                 output, mean_cache, rstd_cache,
                                 tokens, d_model,
                                 d_model, d_model, d_model, eps);
    for (size_t i = 0; i < count; ++i) {
        output[i] = bf16_to_float(float_to_bf16(output[i]));
    }
}

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

◆ layernorm_naive_serial_matched_precision()

void layernorm_naive_serial_matched_precision	(	const float *	input,
		const float *	gamma,
		const float *	beta,
		float *	output,
		float *	mean_cache,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		float	eps
	)

Definition at line 764 of file layernorm_kernels.c.

{
    layernorm_forward_ggml_exact(input, gamma, beta,
                                 output, mean_cache, rstd_cache,
                                 tokens, d_model,
                                 d_model, d_model, d_model, eps);
}

References layernorm_forward_ggml_exact().

Referenced by layernorm_forward_unrolled_slice_scalar().

◆ layernorm_pytorch_welford_bf16_storage()

void layernorm_pytorch_welford_bf16_storage	(	const float *	input,
		const float *	gamma,
		const float *	beta,
		float *	output,
		float *	mean_cache,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		float	eps
	)

Definition at line 952 of file layernorm_kernels.c.

{
#if !defined(__AVX2__) || !defined(__FMA__)
    (void)input; (void)gamma; (void)beta; (void)output;
    (void)mean_cache; (void)rstd_cache; (void)tokens; (void)d_model; (void)eps;
    abort();
#else
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)d_model;
        float *y = output + (size_t)t * (size_t)d_model;
        float mean;
        float variance;
        layernorm_pytorch_bf16_rowwise_moments_avx2(x, d_model, &mean, &variance);
        const float rstd = 1.0f / sqrtf(variance + eps);
        const float bias = -rstd * mean;
        int i = 0;
        for (; i + 7 < d_model; i += 8) {
            const __m256 x_vec = _mm256_loadu_ps(x + i);
            const __m256 gamma_vec = gamma ? _mm256_loadu_ps(gamma + i) : _mm256_set1_ps(1.0f);
            const __m256 beta_vec = beta ? _mm256_loadu_ps(beta + i) : _mm256_setzero_ps();
            const __m256 normalized = _mm256_fmadd_ps(
                x_vec, _mm256_set1_ps(rstd), _mm256_set1_ps(bias));
            const __m256 transformed = _mm256_fmadd_ps(normalized, gamma_vec, beta_vec);
            float lanes[8];
            _mm256_storeu_ps(lanes, transformed);
            for (int lane = 0; lane < 8; ++lane) {
                y[i + lane] = bf16_to_float(float_to_bf16(lanes[lane]));
            }
        }
        for (; i < d_model; ++i) {
            const float gamma_v = gamma ? gamma[i] : 1.0f;
            const float beta_v = beta ? beta[i] : 0.0f;
            const float value = fmaf(fmaf(x[i], rstd, bias), gamma_v, beta_v);
            y[i] = bf16_to_float(float_to_bf16(value));
        }
        if (mean_cache) {
            mean_cache[t] = mean;
        }
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
    }
#endif
}

References bf16_to_float(), and float_to_bf16().

◆ mamba2_conv1d_decode_f32()

void mamba2_conv1d_decode_f32	(	const float *	state_in,
		const float *	x,
		const float *	weight,
		const float *	bias,
		float *	conv_out,
		float *	state_out,
		int	rows,
		int	conv_dim,
		int	kernel_size
	)

Definition at line 148 of file mamba2_kernels.c.

                                               {
    if (!state_in || !x || !weight || !conv_out || !state_out ||
        rows <= 0 || conv_dim <= 0 || kernel_size <= 0) {
        return;
    }
 
    ck_mamba_debug_finite("conv.state_in", state_in, (size_t)conv_dim * (size_t)kernel_size);
    ck_mamba_debug_finite("conv.x[0]", x, (size_t)conv_dim);
    mamba2_conv1d_f32_channel_range(
        state_in, x, weight, bias, conv_out, state_out,
        rows, conv_dim, kernel_size, 0, conv_dim);
    ck_mamba_debug_finite("conv.out[0]", conv_out, (size_t)conv_dim);
    ck_mamba_debug_finite("conv.state_out", state_out, (size_t)conv_dim * (size_t)kernel_size);
}

References ck_mamba_debug_finite(), and mamba2_conv1d_f32_channel_range().

◆ mamba2_conv1d_f32_channel_range()

void mamba2_conv1d_f32_channel_range	(	const float *	state_in,
		const float *	x,
		const float *	weight,
		const float *	bias,
		float *	conv_out,
		float *	state_out,
		int	rows,
		int	conv_dim,
		int	kernel_size,
		int	channel_begin,
		int	channel_end
	)

Definition at line 102 of file mamba2_kernels.c.

                                                      {
    if (!state_in || !x || !weight || !conv_out || !state_out ||
        rows <= 0 || conv_dim <= 0 || kernel_size <= 0 ||
        channel_begin < 0 || channel_begin >= channel_end ||
        channel_end > conv_dim) {
        return;
    }
 
    /* Time rows are dependent, but channels are independent. Keeping one
     * channel's rolling state local preserves its exact update order while
     * allowing disjoint channel ranges to execute concurrently.
     */
    for (int ch = channel_begin; ch < channel_end; ++ch) {
        float state_work[(size_t)kernel_size];
        const size_t base = (size_t)ch * (size_t)kernel_size;
        memcpy(state_work, state_in + base,
               (size_t)kernel_size * sizeof(float));
        for (int row = 0; row < rows; ++row) {
            for (int k = 0; k < kernel_size - 1; ++k) {
                state_work[(size_t)k] = state_work[(size_t)k + 1u];
            }
            state_work[(size_t)kernel_size - 1u] =
                x[(size_t)row * (size_t)conv_dim + (size_t)ch];
 
            float acc = bias ? bias[ch] : 0.0f;
            for (int k = 0; k < kernel_size; ++k) {
                acc += state_work[(size_t)k] *
                       weight[(size_t)ch * (size_t)kernel_size + (size_t)k];
            }
            conv_out[(size_t)row * (size_t)conv_dim + (size_t)ch] = mamba2_silu_f32(acc);
        }
        memcpy(state_out + base, state_work,
               (size_t)kernel_size * sizeof(float));
    }
}

References mamba2_silu_f32().

Referenced by mamba2_conv1d_decode_f32().

◆ mamba2_conv1d_f32_parallel_dispatch()

void mamba2_conv1d_f32_parallel_dispatch	(	const float *	state_in,
		const float *	x,
		const float *	weight,
		const float *	bias,
		float *	conv_out,
		float *	state_out,
		int	rows,
		int	conv_dim,
		int	kernel_size
	)

◆ mamba2_dt_softplus_f32()

void mamba2_dt_softplus_f32	(	const float *	dt,
		const float *	dt_bias,
		float *	dt_out,
		int	rows,
		int	num_heads,
		float	dt_min,
		float	dt_max
	)

Definition at line 171 of file mamba2_kernels.c.

                                          {
    if (!dt || !dt_out || rows <= 0 || num_heads <= 0) {
        return;
    }
 
    ck_mamba_debug_finite("dt.in[0]", dt, (size_t)num_heads);
    ck_mamba_debug_finite("dt.bias", dt_bias, (size_t)num_heads);
    for (int row = 0; row < rows; ++row) {
        for (int h = 0; h < num_heads; ++h) {
            float v = dt[(size_t)row * (size_t)num_heads + (size_t)h];
            if (dt_bias) {
                v += dt_bias[h];
            }
            v = mamba2_softplus_f32(v);
            if (dt_min < dt_max) {
                if (v < dt_min) {
                    v = dt_min;
                } else if (v > dt_max) {
                    v = dt_max;
                }
            }
            dt_out[(size_t)row * (size_t)num_heads + (size_t)h] = v;
        }
    }
    ck_mamba_debug_finite("dt.out[0]", dt_out, (size_t)num_heads);
}

References ck_mamba_debug_finite(), and mamba2_softplus_f32().

◆ mamba2_in_proj_split_f32()

void mamba2_in_proj_split_f32	(	const float *	projected,
		float *	gate,
		float *	hidden_bc,
		float *	dt,
		int	rows,
		int	d_mlp,
		int	intermediate_dim,
		int	conv_dim,
		int	num_heads
	)

Definition at line 65 of file mamba2_kernels.c.

                                             {
    if (!projected || !gate || !hidden_bc || !dt ||
        rows <= 0 || d_mlp < 0 || intermediate_dim <= 0 || conv_dim <= 0 || num_heads <= 0) {
        return;
    }
 
    const int projection_dim = 2 * d_mlp + intermediate_dim + conv_dim + num_heads;
    const int gate_offset = 2 * d_mlp;
    const int hidden_bc_offset = gate_offset + intermediate_dim;
    const int dt_offset = hidden_bc_offset + conv_dim;
 
    ck_mamba_debug_finite("split.projected[0]", projected, (size_t)projection_dim);
    for (int row = 0; row < rows; ++row) {
        const float *src = projected + (size_t)row * (size_t)projection_dim;
        memcpy(gate + (size_t)row * (size_t)intermediate_dim,
               src + gate_offset,
               (size_t)intermediate_dim * sizeof(float));
        memcpy(hidden_bc + (size_t)row * (size_t)conv_dim,
               src + hidden_bc_offset,
               (size_t)conv_dim * sizeof(float));
        memcpy(dt + (size_t)row * (size_t)num_heads,
               src + dt_offset,
               (size_t)num_heads * sizeof(float));
    }
    ck_mamba_debug_finite("split.gate[0]", gate, (size_t)intermediate_dim);
    ck_mamba_debug_finite("split.hidden_bc[0]", hidden_bc, (size_t)conv_dim);
    ck_mamba_debug_finite("split.dt[0]", dt, (size_t)num_heads);
}

References ck_mamba_debug_finite().

◆ mamba2_rmsnorm_gate_f32()

void mamba2_rmsnorm_gate_f32	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	inner_dim,
		int	group_size,
		float	eps
	)

Definition at line 373 of file mamba2_kernels.c.

                                        {
    if (!x || !gate || !weight || !out || rows <= 0 || inner_dim <= 0 || group_size <= 0) {
        return;
    }
 
    ck_mamba_debug_finite("rmsgate.x[0]", x, (size_t)inner_dim);
    ck_mamba_debug_finite("rmsgate.gate[0]", gate, (size_t)inner_dim);
    ck_mamba_debug_finite("rmsgate.weight", weight, (size_t)inner_dim);
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t)row * (size_t)inner_dim;
        const float *gate_row = gate + (size_t)row * (size_t)inner_dim;
        float *out_row = out + (size_t)row * (size_t)inner_dim;
 
        for (int start = 0; start < inner_dim; start += group_size) {
            int end = start + group_size;
            if (end > inner_dim) {
                end = inner_dim;
            }
            const int count = end - start;
            float ms = 0.0f;
            for (int col = start; col < end; ++col) {
                const float gated = x_row[col] * mamba2_silu_f32(gate_row[col]);
                ms += gated * gated;
            }
            const float inv_rms = 1.0f / sqrtf(ms / (float)count + eps);
            for (int col = start; col < end; ++col) {
                const float gated = x_row[col] * mamba2_silu_f32(gate_row[col]);
                out_row[col] = gated * inv_rms * weight[col];
            }
        }
    }
    ck_mamba_debug_finite("rmsgate.out[0]", out, (size_t)inner_dim);
}

References ck_mamba_debug_finite(), end, mamba2_silu_f32(), and start.

◆ mamba2_selective_scan_f32()

void mamba2_selective_scan_f32	(	const float *	state_init,
		const float *	x,
		const float *	dt,
		const float *	a,
		const float *	b,
		const float *	c,
		const float *	d,
		float *	state_out,
		float *	y,
		int	batch,
		int	seq_len,
		int	num_heads,
		int	head_dim,
		int	state_dim,
		int	num_groups
	)

Definition at line 339 of file mamba2_kernels.c.

                                               {
    if (!state_init || !x || !dt || !a || !b || !c || !d || !state_out || !y ||
        batch <= 0 || seq_len <= 0 || num_heads <= 0 || head_dim <= 0 || state_dim <= 0 || num_groups <= 0) {
        return;
    }
 
    const size_t state_per_batch = (size_t)num_heads * (size_t)head_dim * (size_t)state_dim;
    ck_mamba_debug_finite("scan.state_init", state_init, state_per_batch);
    ck_mamba_debug_finite("scan.x[0]", x, (size_t)num_heads * (size_t)head_dim + 2u * (size_t)num_groups * (size_t)state_dim);
    ck_mamba_debug_finite("scan.dt[0]", dt, (size_t)num_heads);
    ck_mamba_debug_finite("scan.a", a, (size_t)num_heads);
    ck_mamba_debug_finite("scan.d", d, (size_t)num_heads);
    mamba2_selective_scan_f32_head_range(
        state_init, x, dt, a, b, c, d, state_out, y,
        batch, seq_len, num_heads, head_dim, state_dim, num_groups,
        0, num_heads);
    ck_mamba_debug_finite("scan.state_out", state_out, (size_t)batch * state_per_batch);
    ck_mamba_debug_finite("scan.y[0]", y, (size_t)num_heads * (size_t)head_dim);
}

References ck_mamba_debug_finite(), and mamba2_selective_scan_f32_head_range().

◆ mamba2_selective_scan_f32_head_range()

void mamba2_selective_scan_f32_head_range	(	const float *	state_init,
		const float *	x,
		const float *	dt,
		const float *	a,
		const float *	b,
		const float *	c,
		const float *	d,
		float *	state_out,
		float *	y,
		int	batch,
		int	seq_len,
		int	num_heads,
		int	head_dim,
		int	state_dim,
		int	num_groups,
		int	head_begin,
		int	head_end
	)

Definition at line 264 of file mamba2_kernels.c.

                                                        {
    if (!state_init || !x || !dt || !a || !b || !c || !d || !state_out || !y ||
        batch <= 0 || seq_len <= 0 || num_heads <= 0 || head_dim <= 0 ||
        state_dim <= 0 || num_groups <= 0 || head_begin < 0 ||
        head_begin >= head_end || head_end > num_heads) {
        return;
    }
 
    const size_t state_per_batch = (size_t)num_heads * (size_t)head_dim * (size_t)state_dim;
    const size_t head_state = (size_t)head_dim * (size_t)state_dim;
    const int packed_xbc = (x == b && b == c);
    const size_t inner_dim = (size_t)num_heads * (size_t)head_dim;
    const size_t bc_dim = (size_t)num_groups * (size_t)state_dim;
    const size_t packed_stride = inner_dim + 2u * bc_dim;
    const int heads_per_group = (num_heads + num_groups - 1) / num_groups;
 
    for (int bs = 0; bs < batch; ++bs) {
        float *state_batch = state_out + (size_t)bs * state_per_batch;
        memcpy(state_batch + (size_t)head_begin * head_state,
               state_init + (size_t)bs * state_per_batch + (size_t)head_begin * head_state,
               (size_t)(head_end - head_begin) * head_state * sizeof(float));
        for (int t = 0; t < seq_len; ++t) {
            for (int h = head_begin; h < head_end; ++h) {
                /* Nemotron-H prefill and decode both use the repeating
                 * B/C group map from repeat_interleave semantics.  Keeping
                 * this in sync is required for full-prefix and incremental
                 * decode equivalence.
                 */
                int group = h / heads_per_group;
                if (group >= num_groups) group = num_groups - 1;
                const float dt_h = dt[((size_t)bs * (size_t)seq_len + (size_t)t) * (size_t)num_heads + (size_t)h];
                const float d_a = expf(dt_h * a[h]);
                const float d_h = d[h];
                const float *packed_row = x + ((size_t)bs * (size_t)seq_len + (size_t)t) * packed_stride;
                const float *b_row = packed_xbc
                    ? (packed_row + inner_dim + (size_t)group * (size_t)state_dim)
                    : (b + (((size_t)bs * (size_t)seq_len + (size_t)t) * (size_t)num_groups + (size_t)group) * (size_t)state_dim);
                const float *c_row = packed_xbc
                    ? (packed_row + inner_dim + bc_dim + (size_t)group * (size_t)state_dim)
                    : (c + (((size_t)bs * (size_t)seq_len + (size_t)t) * (size_t)num_groups + (size_t)group) * (size_t)state_dim);
 
                for (int hd = 0; hd < head_dim; ++hd) {
                    const size_t x_idx = (((size_t)bs * (size_t)seq_len + (size_t)t) * (size_t)num_heads + (size_t)h) * (size_t)head_dim + (size_t)hd;
                    const float x_val = packed_xbc ? packed_row[(size_t)h * (size_t)head_dim + (size_t)hd] : x[x_idx];
                    const size_t state_base = ((size_t)h * (size_t)head_dim + (size_t)hd) * (size_t)state_dim;
                    float acc = 0.0f;
                    for (int st = 0; st < state_dim; ++st) {
                        const size_t si = state_base + (size_t)st;
                        const float new_state = state_batch[si] * d_a + dt_h * b_row[st] * x_val;
                        state_batch[si] = new_state;
                        acc += new_state * c_row[st];
                    }
                    y[x_idx] = acc + d_h * x_val;
                }
            }
        }
    }
}

Referenced by mamba2_selective_scan_f32().

◆ mamba2_selective_scan_f32_parallel_dispatch()

void mamba2_selective_scan_f32_parallel_dispatch	(	const float *	state_init,
		const float *	x,
		const float *	dt,
		const float *	a,
		const float *	b,
		const float *	c,
		const float *	d,
		float *	state_out,
		float *	y,
		int	batch,
		int	seq_len,
		int	num_heads,
		int	head_dim,
		int	state_dim,
		int	num_groups
	)

◆ mamba2_selective_state_update_decode_f32()

void mamba2_selective_state_update_decode_f32	(	const float *	state_in,
		const float *	x,
		const float *	dt,
		const float *	a,
		const float *	b,
		const float *	c,
		const float *	d,
		float *	state_out,
		float *	y,
		int	rows,
		int	num_heads,
		int	head_dim,
		int	state_dim,
		int	num_groups
	)

Definition at line 204 of file mamba2_kernels.c.

                                                              {
    if (!state_in || !x || !dt || !a || !b || !c || !d || !state_out || !y ||
        rows <= 0 || num_heads <= 0 || head_dim <= 0 || state_dim <= 0 || num_groups <= 0) {
        return;
    }
 
    const int packed_xbc = (x == b && b == c);
    const size_t inner_dim = (size_t)num_heads * (size_t)head_dim;
    const size_t bc_dim = (size_t)num_groups * (size_t)state_dim;
    const size_t packed_stride = inner_dim + 2u * bc_dim;
 
    for (int row = 0; row < rows; ++row) {
        const float *packed_row = packed_xbc ? x + (size_t)row * packed_stride : NULL;
        for (int h = 0; h < num_heads; ++h) {
            const int heads_per_group = (num_heads + num_groups - 1) / num_groups;
            int group = h / heads_per_group;
            if (group >= num_groups) group = num_groups - 1;
            const float dt_h = dt[(size_t)row * (size_t)num_heads + (size_t)h];
            const float d_a = expf(dt_h * a[h]);
            const float d_h = d[h];
            const float *b_row = packed_xbc
                ? (packed_row + inner_dim + (size_t)group * (size_t)state_dim)
                : (b + ((size_t)row * (size_t)num_groups + (size_t)group) * (size_t)state_dim);
            const float *c_row = packed_xbc
                ? (packed_row + inner_dim + bc_dim + (size_t)group * (size_t)state_dim)
                : (c + ((size_t)row * (size_t)num_groups + (size_t)group) * (size_t)state_dim);
 
            for (int hd = 0; hd < head_dim; ++hd) {
                const size_t x_idx = ((size_t)row * (size_t)num_heads + (size_t)h) * (size_t)head_dim + (size_t)hd;
                const float x_val = packed_xbc ? packed_row[(size_t)h * (size_t)head_dim + (size_t)hd] : x[x_idx];
                const size_t state_base =
                    (((size_t)row * (size_t)num_heads + (size_t)h) * (size_t)head_dim + (size_t)hd) *
                    (size_t)state_dim;
                float acc = 0.0f;
                for (int s = 0; s < state_dim; ++s) {
                    const size_t si = state_base + (size_t)s;
                    const float new_state = state_in[si] * d_a + dt_h * b_row[s] * x_val;
                    state_out[si] = new_state;
                    acc += new_state * c_row[s];
                }
                y[x_idx] = acc + d_h * x_val;
            }
        }
    }
}

◆ mlp_token_parallel()

void mlp_token_parallel	(	const float *	input,
		const float *	W_fc1,
		const float *	b_fc1,
		const float *	W_fc2,
		const float *	b_fc2,
		float *	fc1_output,
		float *	output,
		int	T,
		int	aligned_dim,
		int	num_threads
	)

Definition at line 49 of file mlp_kernels.c.

{
    int D = aligned_dim;
    int fourD = 4 * D;
 
    // FC1: [T × D] · [D × 4D] -> [T × 4D]
    // Our GEMM layout: A[M×K], B[N×K], so B is [4D × D].
    gemm_blocked_serial(input, W_fc1, b_fc1,
                        fc1_output,
                        T,      // M
                        fourD,  // N
                        D);     // K
 
    // GELU in-place on FC1 output
    gelu_fast_inplace(fc1_output, (size_t)T * (size_t)fourD);
 
    // FC2: [T × 4D] · [4D × D] -> [T × D]
    gemm_blocked_serial(fc1_output, W_fc2, b_fc2,
                        output,
                        T,  // M
                        D,  // N
                        fourD); // K
}

References gelu_fast_inplace(), and gemm_blocked_serial().

◆ mlp_token_parallel_bf16()

void mlp_token_parallel_bf16	(	const uint16_t *	input,
		const uint16_t *	W_fc1,
		const uint16_t *	b_fc1,
		const uint16_t *	W_fc2,
		const uint16_t *	b_fc2,
		float *	fc1_output,
		float *	output,
		int	T,
		int	aligned_dim,
		int	num_threads,
		float *	scratch_bias1_f,
		float *	scratch_bias2_f,
		uint16_t *	scratch_fc1_bf16
	)

Optimized MLP Forward (BF16 weights, FP32 activations)

Caller-provided scratch buffers: scratch_bias1_f: [4*D] floats scratch_bias2_f: [D] floats scratch_fc1_bf16: [T * 4*D] uint16_t (BF16)

Definition at line 104 of file mlp_kernels_bf16.c.

{
    if (!input || !W_fc1 || !b_fc1 || !W_fc2 || !b_fc2 || !fc1_output || !output) return;
    if (!scratch_bias1_f || !scratch_bias2_f || !scratch_fc1_bf16) return;
 
    (void)num_threads;
    const int D = aligned_dim;
    const int fourD = 4 * D;
 
    /* Convert biases to FP32 */
    for (int i = 0; i < fourD; ++i) {
        scratch_bias1_f[i] = bf16_to_float(b_fc1[i]);
    }
    for (int i = 0; i < D; ++i) {
        scratch_bias2_f[i] = bf16_to_float(b_fc2[i]);
    }
 
    /* FC1: [T, D] x [4D, D].T -> [T, 4D] */
    gemm_bf16_fp32out(input, W_fc1, scratch_bias1_f, fc1_output, T, fourD, D);
 
    /* GELU activation */
#if defined(__AVX512F__)
    #pragma omp parallel for
    for (int t = 0; t < T; ++t) {
        float *row = fc1_output + (size_t)t * fourD;
        int j = 0;
        for (; j <= fourD - 16; j += 16) {
            __m512 x = _mm512_loadu_ps(row + j);
            _mm512_storeu_ps(row + j, gelu_avx512(x));
        }
        for (; j < fourD; ++j) {
            row[j] = gelu_scalar(row[j]);
        }
    }
#else
    for (int t = 0; t < T; ++t) {
        for (int j = 0; j < fourD; ++j) {
            fc1_output[t * fourD + j] = gelu_scalar(fc1_output[t * fourD + j]);
        }
    }
#endif
 
    /* Convert FP32 activations to BF16 */
#if defined(__AVX512F__)
    #pragma omp parallel for
    for (int t = 0; t < T; ++t) {
        float *src = fc1_output + (size_t)t * fourD;
        uint16_t *dst = scratch_fc1_bf16 + (size_t)t * fourD;
        int j = 0;
        for (; j <= fourD - 16; j += 16) {
            __m512 fp32 = _mm512_loadu_ps(src + j);
            __m512i as_int = _mm512_castps_si512(fp32);
            __m512i lsb = _mm512_srli_epi32(as_int, 16);
            lsb = _mm512_and_si512(lsb, _mm512_set1_epi32(1));
            __m512i rounding = _mm512_add_epi32(_mm512_set1_epi32(0x7FFF), lsb);
            __m512i rounded = _mm512_add_epi32(as_int, rounding);
            __m512i shifted = _mm512_srli_epi32(rounded, 16);
            __m256i bf16 = _mm512_cvtepi32_epi16(shifted);
            _mm256_storeu_si256((__m256i *)(dst + j), bf16);
        }
        for (; j < fourD; ++j) {
            dst[j] = float_to_bf16(src[j]);
        }
    }
#else
    for (size_t i = 0; i < (size_t)T * fourD; ++i) {
        scratch_fc1_bf16[i] = float_to_bf16(fc1_output[i]);
    }
#endif
 
    /* FC2: BF16 GEMM with FP32 output */
    gemm_bf16_fp32out(scratch_fc1_bf16, W_fc2, scratch_bias2_f, output, T, D, fourD);
}

References bf16_to_float(), float_to_bf16(), gelu_scalar(), and gemm_bf16_fp32out().

◆ mlp_token_parallel_bf16_backward_mixed()

void mlp_token_parallel_bf16_backward_mixed	(	const uint16_t *	input,
		const uint16_t *	W_fc1,
		const uint16_t *	b_fc1,
		const uint16_t *	W_fc2,
		const uint16_t *	d_output,
		float *	d_input,
		float *	d_W_fc1,
		float *	d_b_fc1,
		float *	d_W_fc2,
		float *	d_b_fc2,
		int	T,
		int	aligned_dim,
		int	num_threads,
		float *	scratch_fc1_pre,
		uint16_t *	scratch_fc1_act_bf16,
		float *	scratch_d_fc1
	)

BF16 MLP backward with FP32 gradient accumulation.

Forward contract matched here: z1 = input_bf16 @ W_fc1_bf16.T + b_fc1_bf16 h = GELU(z1) hq = round_to_bf16(h) y = hq @ W_fc2_bf16.T + b_fc2_bf16

Gradients are accumulated and written as FP32. The BF16 activation cast is treated like PyTorch's mixed-precision cast: gradient flows through to h, while d_W_fc2 uses the rounded hq values that FC2 actually consumed.

Definition at line 268 of file mlp_kernels_bf16.c.

{
    if (!input || !W_fc1 || !b_fc1 || !W_fc2 || !d_output) return;
    if (!scratch_fc1_pre || !scratch_fc1_act_bf16 || !scratch_d_fc1) return;
    if (T <= 0 || aligned_dim <= 0) return;
 
    (void)num_threads;
    const int D = aligned_dim;
    const int fourD = 4 * D;
 
    /* Recompute FC1 pre-activation and the rounded activation consumed by FC2. */
    for (int t = 0; t < T; ++t) {
        for (int j = 0; j < fourD; ++j) {
            float sum = bf16_to_float(b_fc1[j]);
            for (int i = 0; i < D; ++i) {
                const float x = bf16_to_float(input[(size_t)t * (size_t)D + (size_t)i]);
                const float w = bf16_to_float(W_fc1[(size_t)j * (size_t)D + (size_t)i]);
                sum += x * w;
            }
            scratch_fc1_pre[(size_t)t * (size_t)fourD + (size_t)j] = sum;
            scratch_fc1_act_bf16[(size_t)t * (size_t)fourD + (size_t)j] = float_to_bf16(gelu_scalar(sum));
        }
    }
 
    if (d_input) {
        for (int t = 0; t < T; ++t) {
            for (int i = 0; i < D; ++i) {
                d_input[(size_t)t * (size_t)D + (size_t)i] = 0.0f;
            }
        }
    }
    if (d_W_fc1) {
        for (size_t i = 0; i < (size_t)fourD * (size_t)D; ++i) d_W_fc1[i] = 0.0f;
    }
    if (d_b_fc1) {
        for (int j = 0; j < fourD; ++j) d_b_fc1[j] = 0.0f;
    }
    if (d_W_fc2) {
        for (size_t i = 0; i < (size_t)D * (size_t)fourD; ++i) d_W_fc2[i] = 0.0f;
    }
    if (d_b_fc2) {
        for (int o = 0; o < D; ++o) d_b_fc2[o] = 0.0f;
    }
 
    /* d_W_fc2, d_b_fc2, and d_h = d_output @ W_fc2. */
    for (int t = 0; t < T; ++t) {
        for (int j = 0; j < fourD; ++j) {
            float dh = 0.0f;
            const float hq = bf16_to_float(scratch_fc1_act_bf16[(size_t)t * (size_t)fourD + (size_t)j]);
            for (int o = 0; o < D; ++o) {
                const float dy = bf16_to_float(d_output[(size_t)t * (size_t)D + (size_t)o]);
                const float w2 = bf16_to_float(W_fc2[(size_t)o * (size_t)fourD + (size_t)j]);
                dh += dy * w2;
                if (d_W_fc2) {
                    d_W_fc2[(size_t)o * (size_t)fourD + (size_t)j] += dy * hq;
                }
            }
            const float z = scratch_fc1_pre[(size_t)t * (size_t)fourD + (size_t)j];
            scratch_d_fc1[(size_t)t * (size_t)fourD + (size_t)j] = dh * gelu_derivative_scalar(z);
        }
        if (d_b_fc2) {
            for (int o = 0; o < D; ++o) {
                d_b_fc2[o] += bf16_to_float(d_output[(size_t)t * (size_t)D + (size_t)o]);
            }
        }
    }
 
    /* Backprop through FC1. */
    for (int t = 0; t < T; ++t) {
        for (int j = 0; j < fourD; ++j) {
            const float dz = scratch_d_fc1[(size_t)t * (size_t)fourD + (size_t)j];
            if (d_b_fc1) d_b_fc1[j] += dz;
            for (int i = 0; i < D; ++i) {
                const float x = bf16_to_float(input[(size_t)t * (size_t)D + (size_t)i]);
                const float w1 = bf16_to_float(W_fc1[(size_t)j * (size_t)D + (size_t)i]);
                if (d_W_fc1) d_W_fc1[(size_t)j * (size_t)D + (size_t)i] += dz * x;
                if (d_input) d_input[(size_t)t * (size_t)D + (size_t)i] += dz * w1;
            }
        }
    }
}

References bf16_to_float(), float_to_bf16(), gelu_derivative_scalar(), and gelu_scalar().

◆ mlp_token_parallel_bf16_fp32act()

void mlp_token_parallel_bf16_fp32act	(	const uint16_t *	input,
		const uint16_t *	W_fc1,
		const uint16_t *	b_fc1,
		const uint16_t *	W_fc2,
		const uint16_t *	b_fc2,
		float *	fc1_output,
		float *	output,
		int	T,
		int	aligned_dim,
		int	num_threads,
		float *	scratch_input_f,
		float *	scratch_bias1_f,
		float *	scratch_bias2_f,
		uint16_t *	scratch_fc1_bf16
	)

Alternative: Fully FP32 activations throughout

Caller-provided scratch buffers: scratch_input_f: [T * D] floats scratch_bias1_f: [4*D] floats scratch_bias2_f: [D] floats scratch_fc1_bf16: [T * 4*D] uint16_t (BF16)

Definition at line 199 of file mlp_kernels_bf16.c.

{
    if (!input || !W_fc1 || !b_fc1 || !W_fc2 || !b_fc2 || !fc1_output || !output) return;
    if (!scratch_input_f || !scratch_bias1_f || !scratch_bias2_f || !scratch_fc1_bf16) return;
 
    (void)num_threads;
    const int D = aligned_dim;
    const int fourD = 4 * D;
 
    /* Convert input and biases to FP32 */
    bf16_tensor_to_float(input, scratch_input_f, (size_t)T * D);
    bf16_tensor_to_float(b_fc1, scratch_bias1_f, fourD);
    bf16_tensor_to_float(b_fc2, scratch_bias2_f, D);
 
    /* FC1 */
    gemm_bf16_fp32out(input, W_fc1, scratch_bias1_f, fc1_output, T, fourD, D);
 
    /* GELU */
#if defined(__AVX512F__)
    #pragma omp parallel for
    for (int t = 0; t < T; ++t) {
        float *row = fc1_output + (size_t)t * fourD;
        int j = 0;
        for (; j <= fourD - 16; j += 16) {
            __m512 x = _mm512_loadu_ps(row + j);
            _mm512_storeu_ps(row + j, gelu_avx512(x));
        }
        for (; j < fourD; ++j) {
            row[j] = gelu_scalar(row[j]);
        }
    }
#else
    for (size_t i = 0; i < (size_t)T * fourD; ++i) {
        fc1_output[i] = gelu_scalar(fc1_output[i]);
    }
#endif
 
    /* Convert fc1_output to BF16 for FC2 */
    float_tensor_to_bf16(fc1_output, scratch_fc1_bf16, (size_t)T * fourD);
    gemm_bf16_fp32out(scratch_fc1_bf16, W_fc2, scratch_bias2_f, output, T, D, fourD);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), gelu_scalar(), and gemm_bf16_fp32out().

◆ mlp_token_parallel_exact()

void mlp_token_parallel_exact	(	const float *	input,
		const float *	W_fc1,
		const float *	b_fc1,
		const float *	W_fc2,
		const float *	b_fc2,
		float *	fc1_output,
		float *	output,
		int	T,
		int	aligned_dim,
		int	num_threads
	)

Definition at line 84 of file mlp_kernels.c.

{
    (void)num_threads;
    int D = aligned_dim;
    int fourD = 4 * D;
 
    // FC1: [T × D] · [D × 4D] -> [T × 4D]
    gemm_blocked_serial(input, W_fc1, b_fc1,
                        fc1_output,
                        T,      // M
                        fourD,  // N
                        D);     // K
 
    // Exact GELU using standard library tanhf
    gelu_exact_inplace(fc1_output, (size_t)T * (size_t)fourD);
 
    // FC2: [T × 4D] · [4D × D] -> [T × D]
    gemm_blocked_serial(fc1_output, W_fc2, b_fc2,
                        output,
                        T,  // M
                        D,  // N
                        fourD); // K
}

References gelu_exact_inplace(), and gemm_blocked_serial().

◆ moe_accumulate_expert_f32()

void moe_accumulate_expert_f32	(	float *	output,
		const float *	expert_output,
		float	routing_weight,
		int	hidden_dim
	)

Accumulate expert output: output += routing_weight * expert_output.

Parameters

output	Token output buffer [hidden_dim], accumulated in place
expert_output	Expert's output for this token [hidden_dim]
routing_weight	Softmax routing weight for this expert
hidden_dim	Hidden dimension

Definition at line 307 of file axpy_kernels.c.

{
    axpy_f32(output, expert_output, routing_weight, hidden_dim);
}

References axpy_f32().

◆ moe_relu2_expert_backward_f32()

void moe_relu2_expert_backward_f32	(	const float *	d_output,
		const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const float *	expert_up,
		const float *	expert_down,
		float *	d_hidden,
		float *	d_routing_weights,
		float *	d_expert_up,
		float *	d_expert_down,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k
	)

Definition at line 3706 of file axpy_kernels.c.

{
    if (!d_output || !hidden || !indices || !routing_weights || !expert_up || !expert_down ||
        !d_hidden || !d_routing_weights || !d_expert_up || !d_expert_down ||
        rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0 || n_experts <= 0 || top_k <= 0) {
        return;
    }
 
    for (size_t p = 0; p < (size_t)rows * (size_t)hidden_dim; ++p) d_hidden[p] = 0.0f;
    for (size_t p = 0; p < (size_t)rows * (size_t)top_k; ++p) d_routing_weights[p] = 0.0f;
    for (size_t p = 0; p < (size_t)n_experts * (size_t)intermediate_dim * (size_t)hidden_dim; ++p) d_expert_up[p] = 0.0f;
    for (size_t p = 0; p < (size_t)n_experts * (size_t)hidden_dim * (size_t)intermediate_dim; ++p) d_expert_down[p] = 0.0f;
 
    float pre[intermediate_dim];
    float act[intermediate_dim];
    float d_act[intermediate_dim];
    float d_pre[intermediate_dim];
    float expert_out[hidden_dim];
 
    for (int r = 0; r < rows; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        const float *dy = d_output + (size_t)r * (size_t)hidden_dim;
        float *dx = d_hidden + (size_t)r * (size_t)hidden_dim;
 
        for (int slot = 0; slot < top_k; ++slot) {
            const int e = indices[(size_t)r * (size_t)top_k + (size_t)slot];
            if (e < 0 || e >= n_experts) continue;
            const float route_w = routing_weights[(size_t)r * (size_t)top_k + (size_t)slot];
 
            for (int i = 0; i < intermediate_dim; ++i) {
                float v = 0.0f;
                for (int h = 0; h < hidden_dim; ++h) {
                    v += expert_up[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)] * x[h];
                }
                pre[i] = v;
                act[i] = (v > 0.0f) ? v * v : 0.0f;
                d_act[i] = 0.0f;
            }
 
            for (int h = 0; h < hidden_dim; ++h) {
                float v = 0.0f;
                for (int i = 0; i < intermediate_dim; ++i) {
                    v += expert_down[ck_moe_down_idx(e, h, i, hidden_dim, intermediate_dim)] * act[i];
                }
                expert_out[h] = v;
            }
 
            float d_route = 0.0f;
            for (int h = 0; h < hidden_dim; ++h) {
                const float d_expert_out = dy[h] * route_w;
                d_route += dy[h] * expert_out[h];
                for (int i = 0; i < intermediate_dim; ++i) {
                    d_expert_down[ck_moe_down_idx(e, h, i, hidden_dim, intermediate_dim)] += d_expert_out * act[i];
                    d_act[i] += d_expert_out * expert_down[ck_moe_down_idx(e, h, i, hidden_dim, intermediate_dim)];
                }
            }
            d_routing_weights[(size_t)r * (size_t)top_k + (size_t)slot] += d_route;
 
            for (int i = 0; i < intermediate_dim; ++i) {
                d_pre[i] = (pre[i] > 0.0f) ? d_act[i] * 2.0f * pre[i] : 0.0f;
            }
 
            for (int i = 0; i < intermediate_dim; ++i) {
                const float dpi = d_pre[i];
                for (int h = 0; h < hidden_dim; ++h) {
                    d_expert_up[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)] += dpi * x[h];
                    dx[h] += dpi * expert_up[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)];
                }
            }
        }
    }
}

References ck_moe_down_idx(), and ck_moe_up_idx().

◆ moe_relu2_expert_forward_f32()

void moe_relu2_expert_forward_f32	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const float *	expert_up,
		const float *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k
	)

Definition at line 374 of file axpy_kernels.c.

{
    if (!hidden || !indices || !routing_weights || !expert_up || !expert_down || !output ||
        rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0 || n_experts <= 0 || top_k <= 0) {
        return;
    }
 
    const size_t out_count = (size_t)rows * (size_t)hidden_dim;
    for (size_t p = 0; p < out_count; ++p) output[p] = 0.0f;
 
    float pre[intermediate_dim];
    float gate[intermediate_dim];
    float up[intermediate_dim];
    float act[intermediate_dim];
 
    for (int r = 0; r < rows; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        float *y = output + (size_t)r * (size_t)hidden_dim;
        for (int slot = 0; slot < top_k; ++slot) {
            const int e = indices[(size_t)r * (size_t)top_k + (size_t)slot];
            if (e < 0 || e >= n_experts) continue;
            const float route_w = routing_weights[(size_t)r * (size_t)top_k + (size_t)slot];
 
            for (int i = 0; i < intermediate_dim; ++i) {
                float v = 0.0f;
                for (int h = 0; h < hidden_dim; ++h) {
                    v += expert_up[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)] * x[h];
                }
                pre[i] = v;
                act[i] = (v > 0.0f) ? v * v : 0.0f;
            }
 
            for (int h = 0; h < hidden_dim; ++h) {
                float v = 0.0f;
                for (int i = 0; i < intermediate_dim; ++i) {
                    v += expert_down[ck_moe_down_idx(e, h, i, hidden_dim, intermediate_dim)] * act[i];
                }
                y[h] += route_w * v;
            }
        }
    }
}

References ck_moe_down_idx(), and ck_moe_up_idx().

◆ moe_relu2_expert_forward_q5_0_q5_0()

void moe_relu2_expert_forward_q5_0_q5_0	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k
	)

Definition at line 3557 of file axpy_kernels.c.

{
    if (!hidden || !indices || !routing_weights || !expert_up || !expert_down || !output ||
        rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0 || n_experts <= 0 || top_k <= 0) {
        return;
    }
 
    const size_t out_count = (size_t)rows * (size_t)hidden_dim;
    for (size_t p = 0; p < out_count; ++p) output[p] = 0.0f;
 
    const size_t up_row_bytes = ck_dtype_row_bytes(CK_DT_Q5_0, (size_t)hidden_dim);
    const size_t down_row_bytes = ck_dtype_row_bytes(CK_DT_Q5_0, (size_t)intermediate_dim);
    const uint8_t *up_base = (const uint8_t *)expert_up;
    const uint8_t *down_base = (const uint8_t *)expert_down;
 
    float up_row[hidden_dim];
    float down_row[intermediate_dim];
    float act[intermediate_dim];
 
    for (int r = 0; r < rows; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        float *y = output + (size_t)r * (size_t)hidden_dim;
        if (ck_moe_debug_enabled() && r == 0) {
            fprintf(stderr,
                    "[CK_DEBUG_MOE] routed_q5q5 rows=%d hidden=%d intermediate=%d experts=%d top_k=%d up_row_bytes=%zu down_row_bytes=%zu\n",
                    rows,
                    hidden_dim,
                    intermediate_dim,
                    n_experts,
                    top_k,
                    up_row_bytes,
                    down_row_bytes);
            fprintf(stderr, "[CK_DEBUG_MOE] routed slots:");
            for (int dbg_slot = 0; dbg_slot < top_k; ++dbg_slot) {
                fprintf(stderr,
                        " (%d,%g)",
                        indices[(size_t)r * (size_t)top_k + (size_t)dbg_slot],
                        routing_weights[(size_t)r * (size_t)top_k + (size_t)dbg_slot]);
            }
            fprintf(stderr, "\n");
            ck_moe_debug_finite("routed.hidden[0]", x, (size_t)hidden_dim);
            ck_moe_debug_finite("routed.hidden_all", hidden, (size_t)rows * (size_t)hidden_dim);
        }
        for (int slot = 0; slot < top_k; ++slot) {
            const int e = indices[(size_t)r * (size_t)top_k + (size_t)slot];
            if (e < 0 || e >= n_experts) continue;
            const float route_w = routing_weights[(size_t)r * (size_t)top_k + (size_t)slot];
            const uint8_t *expert_up_base = up_base + (size_t)e * (size_t)intermediate_dim * up_row_bytes;
            const uint8_t *expert_down_base = down_base + (size_t)e * (size_t)hidden_dim * down_row_bytes;
 
            for (int i = 0; i < intermediate_dim; ++i) {
                dequant_q5_0_row(expert_up_base + (size_t)i * up_row_bytes, up_row, (size_t)hidden_dim);
                float v = 0.0f;
                for (int h = 0; h < hidden_dim; ++h) v += up_row[h] * x[h];
                act[i] = (v > 0.0f) ? v * v : 0.0f;
            }
 
            for (int h = 0; h < hidden_dim; ++h) {
                dequant_q5_0_row(expert_down_base + (size_t)h * down_row_bytes, down_row, (size_t)intermediate_dim);
                float v = 0.0f;
                for (int i = 0; i < intermediate_dim; ++i) v += down_row[i] * act[i];
                y[h] += route_w * v;
            }
        }
        if (ck_moe_debug_enabled() && r == 0) {
            ck_moe_debug_finite("routed.output[0]", y, (size_t)hidden_dim);
            ck_moe_debug_finite("routed.output_all", output, (size_t)rows * (size_t)hidden_dim);
        }
    }
}

References CK_DT_Q5_0, ck_dtype_row_bytes(), ck_moe_debug_enabled(), ck_moe_debug_finite(), and dequant_q5_0_row().

◆ moe_relu2_expert_forward_q5_0_q8_0()

void moe_relu2_expert_forward_q5_0_q8_0	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k
	)

Definition at line 3475 of file axpy_kernels.c.

{
    if (!hidden || !indices || !routing_weights || !expert_up || !expert_down || !output ||
        rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0 || n_experts <= 0 || top_k <= 0) {
        return;
    }
 
    const size_t out_count = (size_t)rows * (size_t)hidden_dim;
    for (size_t p = 0; p < out_count; ++p) output[p] = 0.0f;
 
    const size_t up_row_bytes = ck_dtype_row_bytes(CK_DT_Q5_0, (size_t)hidden_dim);
    const size_t down_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)intermediate_dim);
    const uint8_t *up_base = (const uint8_t *)expert_up;
    const uint8_t *down_base = (const uint8_t *)expert_down;
 
    float up_row[hidden_dim];
    float down_row[intermediate_dim];
    float act[intermediate_dim];
 
    for (int r = 0; r < rows; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        float *y = output + (size_t)r * (size_t)hidden_dim;
        if (ck_moe_debug_enabled() && r == 0) {
            fprintf(stderr,
                    "[CK_DEBUG_MOE] routed_q5q8 rows=%d hidden=%d intermediate=%d experts=%d top_k=%d up_row_bytes=%zu down_row_bytes=%zu\n",
                    rows,
                    hidden_dim,
                    intermediate_dim,
                    n_experts,
                    top_k,
                    up_row_bytes,
                    down_row_bytes);
            fprintf(stderr, "[CK_DEBUG_MOE] routed slots:");
            for (int dbg_slot = 0; dbg_slot < top_k; ++dbg_slot) {
                fprintf(stderr,
                        " (%d,%g)",
                        indices[(size_t)r * (size_t)top_k + (size_t)dbg_slot],
                        routing_weights[(size_t)r * (size_t)top_k + (size_t)dbg_slot]);
            }
            fprintf(stderr, "\n");
            ck_moe_debug_finite("routed.hidden[0]", x, (size_t)hidden_dim);
            ck_moe_debug_finite("routed.hidden_all", hidden, (size_t)rows * (size_t)hidden_dim);
        }
        for (int slot = 0; slot < top_k; ++slot) {
            const int e = indices[(size_t)r * (size_t)top_k + (size_t)slot];
            if (e < 0 || e >= n_experts) continue;
            const float route_w = routing_weights[(size_t)r * (size_t)top_k + (size_t)slot];
            const uint8_t *expert_up_base = up_base + (size_t)e * (size_t)intermediate_dim * up_row_bytes;
            const uint8_t *expert_down_base = down_base + (size_t)e * (size_t)hidden_dim * down_row_bytes;
 
            for (int i = 0; i < intermediate_dim; ++i) {
                dequant_q5_0_row(expert_up_base + (size_t)i * up_row_bytes, up_row, (size_t)hidden_dim);
                float v = 0.0f;
                for (int h = 0; h < hidden_dim; ++h) v += up_row[h] * x[h];
                act[i] = (v > 0.0f) ? v * v : 0.0f;
            }
 
            for (int h = 0; h < hidden_dim; ++h) {
                dequant_q8_0_row(expert_down_base + (size_t)h * down_row_bytes, down_row, (size_t)intermediate_dim);
                float v = 0.0f;
                for (int i = 0; i < intermediate_dim; ++i) v += down_row[i] * act[i];
                y[h] += route_w * v;
            }
        }
        if (ck_moe_debug_enabled() && r == 0) {
            ck_moe_debug_finite("routed.output[0]", y, (size_t)hidden_dim);
            ck_moe_debug_finite("routed.output_all", output, (size_t)rows * (size_t)hidden_dim);
        }
    }
}

References CK_DT_Q5_0, CK_DT_Q8_0, ck_dtype_row_bytes(), ck_moe_debug_enabled(), ck_moe_debug_finite(), dequant_q5_0_row(), and dequant_q8_0_row().

◆ moe_relu2_shared_forward_q5_1_q8_0()

void moe_relu2_shared_forward_q5_1_q8_0	(	const float *	hidden,
		const float *	routed,
		const void *	shared_up,
		const void *	shared_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 3638 of file axpy_kernels.c.

{
    if (!hidden || !shared_up || !shared_down || !output || rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0) {
        return;
    }
 
    const size_t up_row_bytes = ck_dtype_row_bytes(CK_DT_Q5_1, (size_t)hidden_dim);
    const size_t down_row_bytes = ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)intermediate_dim);
    const uint8_t *up_base = (const uint8_t *)shared_up;
    const uint8_t *down_base = (const uint8_t *)shared_down;
 
    float up_row[hidden_dim];
    float down_row[intermediate_dim];
    float act[intermediate_dim];
 
    for (int r = 0; r < rows; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        const float *route = routed ? (routed + (size_t)r * (size_t)hidden_dim) : NULL;
        float *y = output + (size_t)r * (size_t)hidden_dim;
        float x_alias[hidden_dim];
        if (output == hidden) {
            memcpy(x_alias, x, (size_t)hidden_dim * sizeof(float));
            x = x_alias;
        }
 
        if (ck_moe_debug_enabled() && r == 0) {
            fprintf(stderr,
                    "[CK_DEBUG_MOE] shared_q5q8 rows=%d hidden=%d intermediate=%d up_row_bytes=%zu down_row_bytes=%zu alias=%d\n",
                    rows,
                    hidden_dim,
                    intermediate_dim,
                    up_row_bytes,
                    down_row_bytes,
                    output == hidden);
            ck_moe_debug_finite("shared.hidden[0]", x, (size_t)hidden_dim);
            ck_moe_debug_finite("shared.hidden_all", hidden, (size_t)rows * (size_t)hidden_dim);
            ck_moe_debug_finite("shared.routed[0]", route, (size_t)hidden_dim);
            if (routed) ck_moe_debug_finite("shared.routed_all", routed, (size_t)rows * (size_t)hidden_dim);
        }
 
        for (int i = 0; i < intermediate_dim; ++i) {
            dequant_q5_1_row(up_base + (size_t)i * up_row_bytes, up_row, (size_t)hidden_dim);
            float v = 0.0f;
            for (int h = 0; h < hidden_dim; ++h) v += up_row[h] * x[h];
            act[i] = (v > 0.0f) ? v * v : 0.0f;
        }
 
        for (int h = 0; h < hidden_dim; ++h) {
            dequant_q8_0_row(down_base + (size_t)h * down_row_bytes, down_row, (size_t)intermediate_dim);
            float v = route ? route[h] : 0.0f;
            for (int i = 0; i < intermediate_dim; ++i) v += down_row[i] * act[i];
            y[h] = v;
        }
 
        if (ck_moe_debug_enabled() && r == 0) {
            ck_moe_debug_finite("shared.output[0]", y, (size_t)hidden_dim);
            ck_moe_debug_finite("shared.output_all", output, (size_t)rows * (size_t)hidden_dim);
        }
    }
}

References CK_DT_Q5_1, CK_DT_Q8_0, ck_dtype_row_bytes(), ck_moe_debug_enabled(), ck_moe_debug_finite(), dequant_q5_1_row(), and dequant_q8_0_row().

◆ moe_softmax_topk_router_llama_f32_workspace()

int moe_softmax_topk_router_llama_f32_workspace	(	const float *	logits,
		int *	indices,
		float *	weights,
		int	rows,
		int	n_experts,
		int	top_k,
		float	routed_scaling_factor,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 316 of file topk_kernels.c.

{
    const size_t required = moe_softmax_topk_router_workspace_bytes(n_experts);
    if (!logits || !indices || !weights || !workspace || rows <= 0 ||
        n_experts <= 0 || top_k <= 0 || top_k > n_experts ||
        !isfinite(routed_scaling_factor) || required == 0 ||
        workspace_bytes < required) {
        return -1;
    }
 
    float *probabilities = (float *)workspace;
    for (int row = 0; row < rows; ++row) {
        const float *row_logits = logits + (size_t)row * (size_t)n_experts;
        int *row_indices = indices + (size_t)row * (size_t)top_k;
        float *row_weights = weights + (size_t)row * (size_t)top_k;
        float max_value = -INFINITY;
        for (int expert = 0; expert < n_experts; ++expert) {
            if (!isfinite(row_logits[expert])) {
                return -2;
            }
            if (row_logits[expert] > max_value) {
                max_value = row_logits[expert];
            }
        }
 
        const double softmax_sum = ck_moe_llama_softmax_row(
            probabilities, row_logits, n_experts, max_value);
        const float inverse_softmax_sum = (float)(1.0 / softmax_sum);
        for (int expert = 0; expert < n_experts; ++expert) {
            probabilities[expert] *= inverse_softmax_sum;
        }
 
        topk_f32(probabilities, n_experts, top_k, row_indices, NULL);
        double selected_sum_f64 = 0.0;
        for (int slot = 0; slot < top_k; ++slot) {
            row_weights[slot] = probabilities[row_indices[slot]];
            selected_sum_f64 += (double)row_weights[slot];
        }
        float selected_sum = (float)selected_sum_f64;
        if (selected_sum < 6.103515625e-5f) {
            selected_sum = 6.103515625e-5f;
        }
        for (int slot = 0; slot < top_k; ++slot) {
            row_weights[slot] =
                (row_weights[slot] / selected_sum) * routed_scaling_factor;
        }
    }
    return 0;
}

References ck_moe_llama_softmax_row(), moe_softmax_topk_router_workspace_bytes(), and topk_f32().

Referenced by moe_softmax_topk_router_pytorch_bf16_workspace().

◆ moe_softmax_topk_router_pytorch_bf16_workspace()

int moe_softmax_topk_router_pytorch_bf16_workspace	(	const float *	logits,
		int *	indices,
		float *	weights,
		int	rows,
		int	n_experts,
		int	top_k,
		float	routed_scaling_factor,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 375 of file topk_kernels.c.

{
    const int status = moe_softmax_topk_router_llama_f32_workspace(
        logits, indices, weights, rows, n_experts, top_k,
        routed_scaling_factor, workspace, workspace_bytes);
    if (status != 0) {
        return status;
    }
    for (size_t index = 0; index < (size_t)rows * (size_t)top_k; ++index) {
        weights[index] = bf16_to_float(float_to_bf16(weights[index]));
    }
    return 0;
}

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

◆ moe_softmax_topk_router_workspace_bytes()

size_t moe_softmax_topk_router_workspace_bytes ( int n_experts )

Definition at line 177 of file topk_kernels.c.

{
    if (n_experts <= 0) {
        return 0;
    }
    return ((size_t)n_experts * sizeof(float) + 63u) & ~(size_t)63u;
}

Referenced by moe_softmax_topk_router_llama_f32_workspace().

◆ moe_swiglu_expert_forward_bf16()

void moe_swiglu_expert_forward_bf16	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const uint16_t *	expert_gate,
		const uint16_t *	expert_up,
		const uint16_t *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k
	)

Definition at line 563 of file axpy_kernels.c.

{
    if (rows <= 0) return;
    moe_swiglu_expert_forward_bf16_row_range(
        hidden, indices, routing_weights, expert_gate, expert_up, expert_down,
        output, rows, hidden_dim, intermediate_dim, n_experts, top_k, 0, rows);
}

References moe_swiglu_expert_forward_bf16_row_range().

◆ moe_swiglu_expert_forward_bf16_parallel_dispatch()

void moe_swiglu_expert_forward_bf16_parallel_dispatch	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const uint16_t *	expert_gate,
		const uint16_t *	expert_up,
		const uint16_t *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k
	)

◆ moe_swiglu_expert_forward_bf16_row_range()

void moe_swiglu_expert_forward_bf16_row_range	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const uint16_t *	expert_gate,
		const uint16_t *	expert_up,
		const uint16_t *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		int	row_begin,
		int	row_end
	)

Definition at line 505 of file axpy_kernels.c.

{
    if (!hidden || !indices || !routing_weights || !expert_gate || !expert_up || !expert_down || !output ||
        rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0 || n_experts <= 0 || top_k <= 0 ||
        row_begin < 0 || row_begin >= row_end || row_end > rows) {
        return;
    }
 
    float gate[intermediate_dim];
    float up[intermediate_dim];
    float act[intermediate_dim];
 
    for (int r = row_begin; r < row_end; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        float *y = output + (size_t)r * (size_t)hidden_dim;
        for (int h = 0; h < hidden_dim; ++h) y[h] = 0.0f;
        for (int slot = 0; slot < top_k; ++slot) {
            const int e = indices[(size_t)r * (size_t)top_k + (size_t)slot];
            if (e < 0 || e >= n_experts) continue;
            const float route_w = routing_weights[(size_t)r * (size_t)top_k + (size_t)slot];
 
            for (int i = 0; i < intermediate_dim; ++i) {
                float gv = 0.0f;
                float uv = 0.0f;
                for (int h = 0; h < hidden_dim; ++h) {
                    gv += bf16_to_float(expert_gate[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)]) * x[h];
                    uv += bf16_to_float(expert_up[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)]) * x[h];
                }
                gate[i] = gv;
                up[i] = uv;
                act[i] = ck_moe_silu_f32(gv) * uv;
            }
 
            for (int h = 0; h < hidden_dim; ++h) {
                float v = 0.0f;
                for (int i = 0; i < intermediate_dim; ++i) {
                    v += bf16_to_float(expert_down[ck_moe_down_idx(e, h, i, hidden_dim, intermediate_dim)]) * act[i];
                }
                y[h] += route_w * v;
            }
        }
    }
}

References bf16_to_float(), ck_moe_down_idx(), ck_moe_silu_f32(), and ck_moe_up_idx().

Referenced by moe_swiglu_expert_forward_bf16().

◆ moe_swiglu_expert_forward_f32()

void moe_swiglu_expert_forward_f32	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const float *	expert_gate,
		const float *	expert_up,
		const float *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k
	)

Definition at line 450 of file axpy_kernels.c.

{
    if (!hidden || !indices || !routing_weights || !expert_gate || !expert_up || !expert_down || !output ||
        rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0 || n_experts <= 0 || top_k <= 0) {
        return;
    }
 
    for (size_t p = 0; p < (size_t)rows * (size_t)hidden_dim; ++p) output[p] = 0.0f;
 
    float gate[intermediate_dim];
    float up[intermediate_dim];
    float act[intermediate_dim];
 
    for (int r = 0; r < rows; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        float *y = output + (size_t)r * (size_t)hidden_dim;
        for (int slot = 0; slot < top_k; ++slot) {
            const int e = indices[(size_t)r * (size_t)top_k + (size_t)slot];
            if (e < 0 || e >= n_experts) continue;
            const float route_w = routing_weights[(size_t)r * (size_t)top_k + (size_t)slot];
 
            for (int i = 0; i < intermediate_dim; ++i) {
                float gv = 0.0f;
                float uv = 0.0f;
                for (int h = 0; h < hidden_dim; ++h) {
                    gv += expert_gate[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)] * x[h];
                    uv += expert_up[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)] * x[h];
                }
                gate[i] = gv;
                up[i] = uv;
                act[i] = ck_moe_silu_f32(gv) * uv;
            }
 
            for (int h = 0; h < hidden_dim; ++h) {
                float v = 0.0f;
                for (int i = 0; i < intermediate_dim; ++i) {
                    v += expert_down[ck_moe_down_idx(e, h, i, hidden_dim, intermediate_dim)] * act[i];
                }
                y[h] += route_w * v;
            }
        }
    }
}

References ck_moe_down_idx(), ck_moe_silu_f32(), and ck_moe_up_idx().

◆ moe_swiglu_expert_forward_q4k_q4k_parallel_workspace()

int moe_swiglu_expert_forward_q4k_q4k_parallel_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 1881 of file axpy_kernels.c.

{
    return ck_moe_q4k_mixed_parallel_workspace(
        hidden, indices, routing_weights, expert_gate, expert_up, expert_down,
        output, rows, hidden_dim, intermediate_dim, n_experts, top_k,
        workspace, workspace_bytes,
        moe_swiglu_expert_forward_q4k_q4k_workspace, CK_MOE_DOWN_Q4_K);
}

References CK_MOE_DOWN_Q4_K, ck_moe_q4k_mixed_parallel_workspace(), and moe_swiglu_expert_forward_q4k_q4k_workspace().

◆ moe_swiglu_expert_forward_q4k_q4k_workspace()

int moe_swiglu_expert_forward_q4k_q4k_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 1108 of file axpy_kernels.c.

{
    const size_t required = moe_swiglu_expert_q4k_q5k_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !indices || !routing_weights || !expert_gate ||
        !expert_up || !expert_down || !output || !workspace || required == 0 ||
        workspace_bytes < required || rows <= 0 || n_experts <= 0 ||
        top_k <= 0 || top_k > n_experts) {
        return -1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)intermediate_dim * sizeof(float));
    const size_t act_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)intermediate_dim));
    uint8_t *cursor = (uint8_t *)workspace;
    void *hidden_q8 = cursor;
    cursor += hidden_q8_bytes;
    float *gate_up = (float *)cursor;
    cursor += gate_up_bytes;
    void *act_q8 = cursor;
    cursor += act_q8_bytes;
    float *expert_output = (float *)cursor;
 
    const size_t gate_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q4_K, (size_t)hidden_dim);
    const size_t down_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q4_K, (size_t)intermediate_dim);
    const uint8_t *gate_base = (const uint8_t *)expert_gate;
    const uint8_t *up_base = (const uint8_t *)expert_up;
    const uint8_t *down_base = (const uint8_t *)expert_down;
 
    memset(output, 0, (size_t)rows * (size_t)hidden_dim * sizeof(float));
    for (int row = 0; row < rows; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        float *y = output + (size_t)row * (size_t)hidden_dim;
        quantize_row_q8_k(x, hidden_q8, hidden_dim);
        for (int slot = 0; slot < top_k; ++slot) {
            const size_t route_index =
                (size_t)row * (size_t)top_k + (size_t)slot;
            const int expert = indices[route_index];
            if (expert < 0 || expert >= n_experts) return -2;
            const size_t up_offset =
                (size_t)expert * (size_t)intermediate_dim * gate_row_bytes;
            const size_t down_offset =
                (size_t)expert * (size_t)hidden_dim * down_row_bytes;
            gemv_q4_k_q8_k(gate_up, gate_base + up_offset, hidden_q8,
                           intermediate_dim, hidden_dim);
            gemv_q4_k_q8_k(gate_up + intermediate_dim,
                           up_base + up_offset, hidden_q8,
                           intermediate_dim, hidden_dim);
            swiglu_forward_ggml(gate_up, gate_up, 1, intermediate_dim);
            quantize_row_q8_k(gate_up, act_q8, intermediate_dim);
            gemv_q4_k_q8_k(expert_output, down_base + down_offset, act_q8,
                           hidden_dim, intermediate_dim);
            axpy_f32(y, expert_output, routing_weights[route_index], hidden_dim);
        }
    }
    return 0;
}

References axpy_f32(), CK_DT_Q4_K, CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), gemv_q4_k_q8_k(), moe_swiglu_expert_q4k_q5k_workspace_bytes(), quantize_row_q8_k(), and swiglu_forward_ggml().

Referenced by moe_swiglu_expert_forward_q4k_q4k_parallel_workspace().

◆ moe_swiglu_expert_forward_q4k_q5_0_parallel_workspace()

int moe_swiglu_expert_forward_q4k_q5_0_parallel_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 1927 of file axpy_kernels.c.

{
    return ck_moe_q4k_mixed_parallel_workspace(
        hidden, indices, routing_weights, expert_gate, expert_up, expert_down,
        output, rows, hidden_dim, intermediate_dim, n_experts, top_k,
        workspace, workspace_bytes,
        moe_swiglu_expert_forward_q4k_q5_0_workspace, CK_MOE_DOWN_Q5_0);
}

References CK_MOE_DOWN_Q5_0, ck_moe_q4k_mixed_parallel_workspace(), and moe_swiglu_expert_forward_q4k_q5_0_workspace().

◆ moe_swiglu_expert_forward_q4k_q5_0_workspace()

int moe_swiglu_expert_forward_q4k_q5_0_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 1027 of file axpy_kernels.c.

{
    const size_t required = moe_swiglu_expert_q4k_q8_0_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !indices || !routing_weights || !expert_gate ||
        !expert_up || !expert_down || !output || !workspace || required == 0 ||
        workspace_bytes < required || rows <= 0 || n_experts <= 0 ||
        top_k <= 0 || top_k > n_experts) {
        return -1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)intermediate_dim * sizeof(float));
    const size_t act_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)intermediate_dim));
    uint8_t *cursor = (uint8_t *)workspace;
    void *hidden_q8 = cursor;
    cursor += hidden_q8_bytes;
    float *gate_up = (float *)cursor;
    cursor += gate_up_bytes;
    void *act_q8 = cursor;
    cursor += act_q8_bytes;
    float *expert_output = (float *)cursor;
    cursor += ck_moe_align64((size_t)hidden_dim * sizeof(float));
    void *projection_scratch = cursor;
 
    const size_t gate_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q4_K, (size_t)hidden_dim);
    const size_t down_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q5_0, (size_t)intermediate_dim);
    const uint8_t *gate_base = (const uint8_t *)expert_gate;
    const uint8_t *up_base = (const uint8_t *)expert_up;
    const uint8_t *down_base = (const uint8_t *)expert_down;
 
    memset(output, 0, (size_t)rows * (size_t)hidden_dim * sizeof(float));
    for (int row = 0; row < rows; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        float *y = output + (size_t)row * (size_t)hidden_dim;
        quantize_row_q8_k(x, hidden_q8, hidden_dim);
        for (int slot = 0; slot < top_k; ++slot) {
            const size_t route_index =
                (size_t)row * (size_t)top_k + (size_t)slot;
            const int expert = indices[route_index];
            if (expert < 0 || expert >= n_experts) return -2;
            const size_t gate_offset =
                (size_t)expert * (size_t)intermediate_dim * gate_row_bytes;
            const size_t down_offset =
                (size_t)expert * (size_t)hidden_dim * down_row_bytes;
            ck_moe_q4k_llama_projection(
                gate_up, gate_base + gate_offset, hidden_q8,
                intermediate_dim, hidden_dim, projection_scratch);
            ck_moe_q4k_llama_projection(
                gate_up + intermediate_dim, up_base + gate_offset, hidden_q8,
                intermediate_dim, hidden_dim, projection_scratch);
            swiglu_forward_ggml(gate_up, gate_up, 1, intermediate_dim);
            quantize_row_q8_0(gate_up, act_q8, intermediate_dim);
            gemv_q5_0_q8_0(expert_output, down_base + down_offset, act_q8,
                           hidden_dim, intermediate_dim);
            ck_moe_llama_weighted_accumulate(
                y, expert_output, routing_weights[route_index], hidden_dim);
        }
    }
    return 0;
}

References CK_DT_Q4_K, CK_DT_Q5_0, CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), ck_moe_llama_weighted_accumulate(), ck_moe_q4k_llama_projection(), gemv_q5_0_q8_0(), moe_swiglu_expert_q4k_q8_0_workspace_bytes(), quantize_row_q8_0(), quantize_row_q8_k(), and swiglu_forward_ggml().

Referenced by moe_swiglu_expert_forward_q4k_q5_0_parallel_workspace().

◆ moe_swiglu_expert_forward_q4k_q5k_auto_prepared_workspace()

int moe_swiglu_expert_forward_q4k_q5k_auto_prepared_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

◆ moe_swiglu_expert_forward_q4k_q5k_auto_workspace()

int moe_swiglu_expert_forward_q4k_q5k_auto_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2595 of file axpy_kernels.c.

{
    if (rows < 512) {
        return moe_swiglu_expert_forward_q4k_q5k_parallel_workspace(
            hidden, indices, routing_weights,
            expert_gate, expert_up, expert_down, output,
            rows, hidden_dim, intermediate_dim, n_experts, top_k,
            workspace, workspace_bytes);
    }
    return moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace(
        hidden, indices, routing_weights,
        expert_gate, expert_up, expert_down, output,
        rows, hidden_dim, intermediate_dim, n_experts, top_k,
        workspace, workspace_bytes);
}

References moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace(), and moe_swiglu_expert_forward_q4k_q5k_parallel_workspace().

◆ moe_swiglu_expert_forward_q4k_q5k_bucketed_prepared_workspace()

int moe_swiglu_expert_forward_q4k_q5k_bucketed_prepared_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		const void *	expert_gate_packed,
		const void *	expert_up_packed,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2569 of file axpy_kernels.c.

{
    return ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl(
        hidden, indices, routing_weights,
        expert_gate, expert_up, expert_down,
        expert_gate_packed, expert_up_packed, output,
        rows, hidden_dim, intermediate_dim, n_experts, top_k,
        workspace, workspace_bytes);
}

References ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl().

◆ moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace()

int moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2546 of file axpy_kernels.c.

{
    return ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl(
        hidden, indices, routing_weights,
        expert_gate, expert_up, expert_down, NULL, NULL, output,
        rows, hidden_dim, intermediate_dim, n_experts, top_k,
        workspace, workspace_bytes);
}

References ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl().

Referenced by moe_swiglu_expert_forward_q4k_q5k_auto_workspace().

◆ moe_swiglu_expert_forward_q4k_q5k_parallel_workspace()

int moe_swiglu_expert_forward_q4k_q5k_parallel_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 1470 of file axpy_kernels.c.

{
    const size_t stride = moe_swiglu_expert_q4k_q5k_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !indices || !routing_weights || !expert_gate || !expert_up ||
        !expert_down || !output || !workspace || stride == 0 || rows <= 0 ||
        n_experts <= 0 || top_k <= 0 || top_k > n_experts) {
        return -1;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    if (rows == 1) {
        const int route_status = ck_moe_q4k_q5k_route_parallel(
            hidden, indices, routing_weights,
            expert_gate, expert_up, expert_down, output,
            hidden_dim, intermediate_dim, n_experts, top_k,
            workspace, workspace_bytes, stride, pool);
        if (route_status <= 0) return route_status;
    }
 
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > rows) active = rows;
    if (active > CK_THREADPOOL_MAX_THREADS) active = CK_THREADPOOL_MAX_THREADS;
    const size_t workspace_workers = workspace_bytes / stride;
    if (workspace_workers == 0) return -1;
    if ((size_t)active > workspace_workers) active = (int)workspace_workers;
    if (active <= 1) {
        return moe_swiglu_expert_forward_q4k_q5k_workspace(
            hidden, indices, routing_weights, expert_gate, expert_up, expert_down,
            output, rows, hidden_dim, intermediate_dim, n_experts, top_k,
            workspace, stride);
    }
 
    ck_moe_q4k_q5k_parallel_args_t args = {
        .hidden = hidden,
        .indices = indices,
        .routing_weights = routing_weights,
        .expert_gate = expert_gate,
        .expert_up = expert_up,
        .expert_down = expert_down,
        .output = output,
        .rows = rows,
        .hidden_dim = hidden_dim,
        .intermediate_dim = intermediate_dim,
        .n_experts = n_experts,
        .top_k = top_k,
        .workspace = (uint8_t *)workspace,
        .workspace_stride = stride,
        .status = {0},
    };
    ck_threadpool_dispatch_n(pool, active, ck_moe_q4k_q5k_parallel_work, &args);
    for (int ith = 0; ith < active; ++ith) {
        if (args.status[ith] != 0) return args.status[ith];
    }
    return 0;
}

References ck_moe_q4k_q5k_parallel_work(), ck_moe_q4k_q5k_route_parallel(), ck_threadpool_dispatch_n(), ck_threadpool_global(), CK_THREADPOOL_MAX_THREADS, ck_threadpool_n_threads(), moe_swiglu_expert_forward_q4k_q5k_workspace(), and moe_swiglu_expert_q4k_q5k_workspace_bytes().

Referenced by moe_swiglu_expert_forward_q4k_q5k_auto_workspace().

◆ moe_swiglu_expert_forward_q4k_q5k_workspace()

int moe_swiglu_expert_forward_q4k_q5k_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 781 of file axpy_kernels.c.

{
    const size_t required = moe_swiglu_expert_q4k_q5k_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !indices || !routing_weights || !expert_gate || !expert_up ||
        !expert_down || !output || !workspace || required == 0 ||
        workspace_bytes < required || rows <= 0 || n_experts <= 0 ||
        top_k <= 0 || top_k > n_experts) {
        return -1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)intermediate_dim * sizeof(float));
    const size_t act_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)intermediate_dim));
    uint8_t *cursor = (uint8_t *)workspace;
    void *hidden_q8 = cursor;
    cursor += hidden_q8_bytes;
    float *gate_up = (float *)cursor;
    cursor += gate_up_bytes;
    void *act_q8 = cursor;
    cursor += act_q8_bytes;
    float *expert_output = (float *)cursor;
 
    const size_t q4_row_bytes = ck_dtype_row_bytes(CK_DT_Q4_K, (size_t)hidden_dim);
    const size_t q5_row_bytes = ck_dtype_row_bytes(CK_DT_Q5_K, (size_t)intermediate_dim);
    const uint8_t *gate_base = (const uint8_t *)expert_gate;
    const uint8_t *up_base = (const uint8_t *)expert_up;
    const uint8_t *down_base = (const uint8_t *)expert_down;
 
    memset(output, 0, (size_t)rows * (size_t)hidden_dim * sizeof(float));
    for (int row = 0; row < rows; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        float *y = output + (size_t)row * (size_t)hidden_dim;
        quantize_row_q8_k(x, hidden_q8, hidden_dim);
 
        for (int slot = 0; slot < top_k; ++slot) {
            const size_t route_index = (size_t)row * (size_t)top_k + (size_t)slot;
            const int expert = indices[route_index];
            if (expert < 0 || expert >= n_experts) {
                return -2;
            }
 
            const size_t up_expert_offset =
                (size_t)expert * (size_t)intermediate_dim * q4_row_bytes;
            const size_t down_expert_offset =
                (size_t)expert * (size_t)hidden_dim * q5_row_bytes;
            gemv_q4_k_q8_k(gate_up,
                           gate_base + up_expert_offset,
                           hidden_q8,
                           intermediate_dim,
                           hidden_dim);
            gemv_q4_k_q8_k(gate_up + intermediate_dim,
                           up_base + up_expert_offset,
                           hidden_q8,
                           intermediate_dim,
                           hidden_dim);
            swiglu_forward_ggml(gate_up, gate_up, 1, intermediate_dim);
            quantize_row_q8_k(gate_up, act_q8, intermediate_dim);
            gemv_q5_k_q8_k(expert_output,
                           down_base + down_expert_offset,
                           act_q8,
                           hidden_dim,
                           intermediate_dim);
 
            const float route_weight = routing_weights[route_index];
            axpy_f32(y, expert_output, route_weight, hidden_dim);
        }
    }
    return 0;
}

References axpy_f32(), CK_DT_Q4_K, CK_DT_Q5_K, CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), gemv_q4_k_q8_k(), gemv_q5_k_q8_k(), moe_swiglu_expert_q4k_q5k_workspace_bytes(), quantize_row_q8_k(), and swiglu_forward_ggml().

Referenced by ck_moe_q4k_q5k_parallel_work(), and moe_swiglu_expert_forward_q4k_q5k_parallel_workspace().

◆ moe_swiglu_expert_forward_q4k_q6k_parallel_workspace()

int moe_swiglu_expert_forward_q4k_q6k_parallel_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 1858 of file axpy_kernels.c.

{
    return ck_moe_q4k_mixed_parallel_workspace(
        hidden, indices, routing_weights, expert_gate, expert_up, expert_down,
        output, rows, hidden_dim, intermediate_dim, n_experts, top_k,
        workspace, workspace_bytes,
        moe_swiglu_expert_forward_q4k_q6k_workspace, CK_MOE_DOWN_Q6_K);
}

References CK_MOE_DOWN_Q6_K, ck_moe_q4k_mixed_parallel_workspace(), and moe_swiglu_expert_forward_q4k_q6k_workspace().

◆ moe_swiglu_expert_forward_q4k_q6k_workspace()

int moe_swiglu_expert_forward_q4k_q6k_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 869 of file axpy_kernels.c.

{
    const size_t required = moe_swiglu_expert_q4k_q5k_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !indices || !routing_weights || !expert_gate ||
        !expert_up || !expert_down || !output || !workspace || required == 0 ||
        workspace_bytes < required || rows <= 0 || n_experts <= 0 ||
        top_k <= 0 || top_k > n_experts) {
        return -1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)intermediate_dim * sizeof(float));
    const size_t act_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)intermediate_dim));
    uint8_t *cursor = (uint8_t *)workspace;
    void *hidden_q8 = cursor;
    cursor += hidden_q8_bytes;
    float *gate_up = (float *)cursor;
    cursor += gate_up_bytes;
    void *act_q8 = cursor;
    cursor += act_q8_bytes;
    float *expert_output = (float *)cursor;
 
    const size_t q4_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q4_K, (size_t)hidden_dim);
    const size_t q6_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q6_K, (size_t)intermediate_dim);
    const uint8_t *gate_base = (const uint8_t *)expert_gate;
    const uint8_t *up_base = (const uint8_t *)expert_up;
    const uint8_t *down_base = (const uint8_t *)expert_down;
 
    memset(output, 0, (size_t)rows * (size_t)hidden_dim * sizeof(float));
    for (int row = 0; row < rows; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        float *y = output + (size_t)row * (size_t)hidden_dim;
        quantize_row_q8_k(x, hidden_q8, hidden_dim);
        for (int slot = 0; slot < top_k; ++slot) {
            const size_t route_index =
                (size_t)row * (size_t)top_k + (size_t)slot;
            const int expert = indices[route_index];
            if (expert < 0 || expert >= n_experts) return -2;
            const size_t up_offset =
                (size_t)expert * (size_t)intermediate_dim * q4_row_bytes;
            const size_t down_offset =
                (size_t)expert * (size_t)hidden_dim * q6_row_bytes;
            gemv_q4_k_q8_k(gate_up, gate_base + up_offset, hidden_q8,
                           intermediate_dim, hidden_dim);
            gemv_q4_k_q8_k(gate_up + intermediate_dim,
                           up_base + up_offset, hidden_q8,
                           intermediate_dim, hidden_dim);
            swiglu_forward_ggml(gate_up, gate_up, 1, intermediate_dim);
            quantize_row_q8_k(gate_up, act_q8, intermediate_dim);
            gemv_q6_k_q8_k(expert_output, down_base + down_offset, act_q8,
                           hidden_dim, intermediate_dim);
            axpy_f32(y, expert_output, routing_weights[route_index], hidden_dim);
        }
    }
    return 0;
}

References axpy_f32(), CK_DT_Q4_K, CK_DT_Q6_K, CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), gemv_q4_k_q8_k(), gemv_q6_k_q8_k(), moe_swiglu_expert_q4k_q5k_workspace_bytes(), quantize_row_q8_k(), and swiglu_forward_ggml().

Referenced by moe_swiglu_expert_forward_q4k_q6k_parallel_workspace().

◆ moe_swiglu_expert_forward_q4k_q8_0_parallel_workspace()

int moe_swiglu_expert_forward_q4k_q8_0_parallel_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 1904 of file axpy_kernels.c.

{
    return ck_moe_q4k_mixed_parallel_workspace(
        hidden, indices, routing_weights, expert_gate, expert_up, expert_down,
        output, rows, hidden_dim, intermediate_dim, n_experts, top_k,
        workspace, workspace_bytes,
        moe_swiglu_expert_forward_q4k_q8_0_workspace, CK_MOE_DOWN_Q8_0);
}

References CK_MOE_DOWN_Q8_0, ck_moe_q4k_mixed_parallel_workspace(), and moe_swiglu_expert_forward_q4k_q8_0_workspace().

◆ moe_swiglu_expert_forward_q4k_q8_0_workspace()

int moe_swiglu_expert_forward_q4k_q8_0_workspace	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const void *	expert_gate,
		const void *	expert_up,
		const void *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 946 of file axpy_kernels.c.

{
    const size_t required = moe_swiglu_expert_q4k_q8_0_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !indices || !routing_weights || !expert_gate ||
        !expert_up || !expert_down || !output || !workspace || required == 0 ||
        workspace_bytes < required || rows <= 0 || n_experts <= 0 ||
        top_k <= 0 || top_k > n_experts) {
        return -1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)intermediate_dim * sizeof(float));
    const size_t act_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)intermediate_dim));
    uint8_t *cursor = (uint8_t *)workspace;
    void *hidden_q8 = cursor;
    cursor += hidden_q8_bytes;
    float *gate_up = (float *)cursor;
    cursor += gate_up_bytes;
    void *act_q8 = cursor;
    cursor += act_q8_bytes;
    float *expert_output = (float *)cursor;
    cursor += ck_moe_align64((size_t)hidden_dim * sizeof(float));
    void *projection_scratch = cursor;
 
    const size_t gate_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q4_K, (size_t)hidden_dim);
    const size_t down_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q8_0, (size_t)intermediate_dim);
    const uint8_t *gate_base = (const uint8_t *)expert_gate;
    const uint8_t *up_base = (const uint8_t *)expert_up;
    const uint8_t *down_base = (const uint8_t *)expert_down;
 
    memset(output, 0, (size_t)rows * (size_t)hidden_dim * sizeof(float));
    for (int row = 0; row < rows; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        float *y = output + (size_t)row * (size_t)hidden_dim;
        quantize_row_q8_k(x, hidden_q8, hidden_dim);
        for (int slot = 0; slot < top_k; ++slot) {
            const size_t route_index =
                (size_t)row * (size_t)top_k + (size_t)slot;
            const int expert = indices[route_index];
            if (expert < 0 || expert >= n_experts) return -2;
            const size_t gate_offset =
                (size_t)expert * (size_t)intermediate_dim * gate_row_bytes;
            const size_t down_offset =
                (size_t)expert * (size_t)hidden_dim * down_row_bytes;
            ck_moe_q4k_llama_projection(
                gate_up, gate_base + gate_offset, hidden_q8,
                intermediate_dim, hidden_dim, projection_scratch);
            ck_moe_q4k_llama_projection(
                gate_up + intermediate_dim, up_base + gate_offset, hidden_q8,
                intermediate_dim, hidden_dim, projection_scratch);
            swiglu_forward_ggml(gate_up, gate_up, 1, intermediate_dim);
            quantize_row_q8_0(gate_up, act_q8, intermediate_dim);
            gemv_q8_0_q8_0(expert_output, down_base + down_offset, act_q8,
                           hidden_dim, intermediate_dim);
            ck_moe_llama_weighted_accumulate(
                y, expert_output, routing_weights[route_index], hidden_dim);
        }
    }
    return 0;
}

References CK_DT_Q4_K, CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), ck_moe_llama_weighted_accumulate(), ck_moe_q4k_llama_projection(), gemv_q8_0_q8_0(), moe_swiglu_expert_q4k_q8_0_workspace_bytes(), quantize_row_q8_0(), quantize_row_q8_k(), and swiglu_forward_ggml().

Referenced by moe_swiglu_expert_forward_q4k_q8_0_parallel_workspace().

◆ moe_swiglu_expert_q4k_q5k_workspace_bytes()

size_t moe_swiglu_expert_q4k_q5k_workspace_bytes	(	int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 727 of file axpy_kernels.c.

{
    if (hidden_dim <= 0 || intermediate_dim <= 0 ||
        hidden_dim % 256 != 0 || intermediate_dim % 256 != 0) {
        return 0;
    }
 
    size_t bytes = ck_moe_align64(ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    bytes += ck_moe_align64(2u * (size_t)intermediate_dim * sizeof(float));
    bytes += ck_moe_align64(ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)intermediate_dim));
    bytes += ck_moe_align64((size_t)hidden_dim * sizeof(float));
    return bytes;
}

References CK_DT_Q8_K, ck_dtype_row_bytes(), and ck_moe_align64().

Referenced by ck_moe_q4k_mixed_parallel_workspace(), ck_moe_shared_q4k_parallel_workspace(), moe_swiglu_expert_forward_q4k_q4k_workspace(), moe_swiglu_expert_forward_q4k_q5k_parallel_workspace(), moe_swiglu_expert_forward_q4k_q5k_workspace(), moe_swiglu_expert_forward_q4k_q6k_workspace(), moe_swiglu_shared_forward_q4k_q4k_workspace(), and moe_swiglu_shared_forward_q4k_q6k_workspace().

◆ moe_swiglu_expert_q4k_q8_0_workspace_bytes()

size_t moe_swiglu_expert_q4k_q8_0_workspace_bytes	(	int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 742 of file axpy_kernels.c.

{
    if (hidden_dim <= 0 || intermediate_dim <= 0 ||
        hidden_dim % 256 != 0 || intermediate_dim % QK8_0 != 0) {
        return 0;
    }
 
    size_t bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    bytes += ck_moe_align64(2u * (size_t)intermediate_dim * sizeof(float));
    bytes += ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)intermediate_dim));
    bytes += ck_moe_align64((size_t)hidden_dim * sizeof(float));
    bytes += ck_moe_q4k_llama_projection_scratch_bytes(
        intermediate_dim, hidden_dim);
    return bytes;
}

References CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), ck_moe_q4k_llama_projection_scratch_bytes(), and QK8_0.

Referenced by ck_moe_q4k_mixed_parallel_workspace(), moe_swiglu_expert_forward_q4k_q5_0_workspace(), and moe_swiglu_expert_forward_q4k_q8_0_workspace().

◆ moe_swiglu_packed_expert_forward_bf16()

void moe_swiglu_packed_expert_forward_bf16	(	const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const uint16_t *	expert_gate_up,
		const uint16_t *	expert_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k
	)

Definition at line 582 of file axpy_kernels.c.

{
    if (!hidden || !indices || !routing_weights || !expert_gate_up ||
        !expert_down || !output || rows <= 0 || hidden_dim <= 0 ||
        intermediate_dim <= 0 || n_experts <= 0 || top_k <= 0) {
        return;
    }
 
    float projection[2 * intermediate_dim];
    float activation[intermediate_dim];
    const size_t packed_expert_stride =
        (size_t)2 * (size_t)intermediate_dim * (size_t)hidden_dim;
    const size_t packed_up_offset =
        (size_t)intermediate_dim * (size_t)hidden_dim;
    for (int row = 0; row < rows; ++row) {
        float *y = output + (size_t)row * (size_t)hidden_dim;
        for (int hidden_col = 0; hidden_col < hidden_dim; ++hidden_col) {
            y[hidden_col] = 0.0f;
        }
    }
    for (int row = 0; row < rows; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        float *y = output + (size_t)row * (size_t)hidden_dim;
        for (int slot = 0; slot < top_k; ++slot) {
                const size_t route_index =
                    (size_t)row * (size_t)top_k + (size_t)slot;
                const int expert = indices[route_index];
                if (expert < 0 || expert >= n_experts) continue;
                const uint16_t *packed =
                    expert_gate_up + (size_t)expert * packed_expert_stride;
                const uint16_t *gate = packed;
                const uint16_t *up = packed + packed_up_offset;
                for (int intermediate = 0; intermediate < intermediate_dim;
                     ++intermediate) {
                    float gate_value = 0.0f;
                    float up_value = 0.0f;
                    const size_t row_offset =
                        (size_t)intermediate * (size_t)hidden_dim;
                    for (int hidden_col = 0; hidden_col < hidden_dim;
                         ++hidden_col) {
                        const float input =
                            ck_moe_bf16_round(x[hidden_col]);
                        gate_value +=
                            bf16_to_float(gate[row_offset + hidden_col]) * input;
                        up_value +=
                            bf16_to_float(up[row_offset + hidden_col]) * input;
                    }
                    projection[intermediate] =
                        ck_moe_bf16_round(gate_value);
                    projection[intermediate_dim + intermediate] =
                        ck_moe_bf16_round(up_value);
                }
                swiglu_forward_pytorch_bf16_storage(
                    projection, activation, 1, intermediate_dim);
 
                const float route_weight =
                    ck_moe_bf16_round(routing_weights[route_index]);
                for (int hidden_col = 0; hidden_col < hidden_dim;
                     ++hidden_col) {
                    float value = 0.0f;
                    for (int intermediate = 0;
                         intermediate < intermediate_dim; ++intermediate) {
                        value += bf16_to_float(expert_down[
                            ck_moe_down_idx(expert, hidden_col, intermediate,
                                            hidden_dim, intermediate_dim)]) *
                                 activation[intermediate];
                    }
                    const float down_stored = ck_moe_bf16_round(value);
                    const float weighted = ck_moe_bf16_round(
                        down_stored * route_weight);
                    y[hidden_col] += weighted;
                }
        }
        for (int hidden_col = 0; hidden_col < hidden_dim; ++hidden_col) {
            y[hidden_col] = ck_moe_bf16_round(y[hidden_col]);
        }
    }
}

References bf16_to_float(), ck_moe_bf16_round(), ck_moe_down_idx(), and swiglu_forward_pytorch_bf16_storage().

◆ moe_swiglu_shared_forward_bf16()

void moe_swiglu_shared_forward_bf16	(	const float *	hidden,
		const float *	routed,
		const uint16_t *	shared_gate,
		const uint16_t *	shared_up,
		const uint16_t *	shared_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 3205 of file axpy_kernels.c.

{
    if (rows <= 0) return;
    moe_swiglu_shared_forward_bf16_row_range(
        hidden, routed, shared_gate, shared_up, shared_down, output,
        rows, hidden_dim, intermediate_dim, 0, rows);
}

References moe_swiglu_shared_forward_bf16_row_range().

◆ moe_swiglu_shared_forward_bf16_gated()

void moe_swiglu_shared_forward_bf16_gated	(	const float *	hidden,
		const float *	routed,
		const uint16_t *	shared_gate,
		const uint16_t *	shared_up,
		const uint16_t *	shared_down,
		const uint16_t *	shared_router,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 3298 of file axpy_kernels.c.

{
    if (rows <= 0) return;
    moe_swiglu_shared_forward_bf16_gated_row_range(
        hidden, routed, shared_gate, shared_up, shared_down, shared_router,
        output, rows, hidden_dim, intermediate_dim, 0, rows);
}

References moe_swiglu_shared_forward_bf16_gated_row_range().

◆ moe_swiglu_shared_forward_bf16_gated_parallel_dispatch()

void moe_swiglu_shared_forward_bf16_gated_parallel_dispatch	(	const float *	hidden,
		const float *	routed,
		const uint16_t *	shared_gate,
		const uint16_t *	shared_up,
		const uint16_t *	shared_down,
		const uint16_t *	shared_router,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim
	)

◆ moe_swiglu_shared_forward_bf16_gated_row_range()

void moe_swiglu_shared_forward_bf16_gated_row_range	(	const float *	hidden,
		const float *	routed,
		const uint16_t *	shared_gate,
		const uint16_t *	shared_up,
		const uint16_t *	shared_down,
		const uint16_t *	shared_router,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	row_begin,
		int	row_end
	)

Definition at line 3221 of file axpy_kernels.c.

{
    if (!hidden || !shared_gate || !shared_up || !shared_down ||
        !shared_router || !output || rows <= 0 || hidden_dim <= 0 ||
        intermediate_dim <= 0 || row_begin < 0 || row_begin >= row_end ||
        row_end > rows) {
        return;
    }
 
    float activation[intermediate_dim];
    for (int row = row_begin; row < row_end; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        const float *route = routed
            ? routed + (size_t)row * (size_t)hidden_dim
            : NULL;
        float *y = output + (size_t)row * (size_t)hidden_dim;
 
        for (int intermediate = 0; intermediate < intermediate_dim;
             ++intermediate) {
            float gate_sum = 0.0f;
            float up_sum = 0.0f;
            const size_t weight_base =
                (size_t)intermediate * (size_t)hidden_dim;
            for (int hidden_col = 0; hidden_col < hidden_dim; ++hidden_col) {
                const float value = ck_moe_bf16_round(x[hidden_col]);
                gate_sum += bf16_to_float(
                    shared_gate[weight_base + (size_t)hidden_col]) * value;
                up_sum += bf16_to_float(
                    shared_up[weight_base + (size_t)hidden_col]) * value;
            }
            const float gate_value = ck_moe_bf16_round(gate_sum);
            const float up_value = ck_moe_bf16_round(up_sum);
            const float silu = ck_moe_bf16_round(
                ck_moe_silu_f32(gate_value));
            activation[intermediate] = ck_moe_bf16_round(silu * up_value);
        }
 
        float router_sum = 0.0f;
        for (int hidden_col = 0; hidden_col < hidden_dim; ++hidden_col) {
            router_sum += bf16_to_float(shared_router[hidden_col]) *
                ck_moe_bf16_round(x[hidden_col]);
        }
        const float shared_scale = ck_moe_bf16_round(
            ck_moe_sigmoid_f32(ck_moe_bf16_round(router_sum)));
 
        for (int hidden_col = 0; hidden_col < hidden_dim; ++hidden_col) {
            float down_sum = 0.0f;
            const size_t weight_base =
                (size_t)hidden_col * (size_t)intermediate_dim;
            for (int intermediate = 0; intermediate < intermediate_dim;
                 ++intermediate) {
                down_sum += bf16_to_float(
                    shared_down[weight_base + (size_t)intermediate]) *
                    activation[intermediate];
            }
            const float shared = ck_moe_bf16_round(down_sum);
            const float gated_shared = ck_moe_bf16_round(shared * shared_scale);
            const float routed_value = route
                ? ck_moe_bf16_round(route[hidden_col])
                : 0.0f;
            y[hidden_col] = ck_moe_bf16_round(routed_value + gated_shared);
        }
    }
}

References bf16_to_float(), ck_moe_bf16_round(), ck_moe_sigmoid_f32(), ck_moe_silu_f32(), and silu().

Referenced by moe_swiglu_shared_forward_bf16_gated().

◆ moe_swiglu_shared_forward_bf16_parallel_dispatch()

void moe_swiglu_shared_forward_bf16_parallel_dispatch	(	const float *	hidden,
		const float *	routed,
		const uint16_t *	shared_gate,
		const uint16_t *	shared_up,
		const uint16_t *	shared_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim
	)

◆ moe_swiglu_shared_forward_bf16_row_range()

void moe_swiglu_shared_forward_bf16_row_range	(	const float *	hidden,
		const float *	routed,
		const uint16_t *	shared_gate,
		const uint16_t *	shared_up,
		const uint16_t *	shared_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	row_begin,
		int	row_end
	)

Definition at line 3157 of file axpy_kernels.c.

{
    if (!hidden || !shared_gate || !shared_up || !shared_down || !output ||
        rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0 ||
        row_begin < 0 || row_begin >= row_end || row_end > rows) {
        return;
    }
 
    float gate[intermediate_dim];
    float up[intermediate_dim];
    float act[intermediate_dim];
 
    for (int r = row_begin; r < row_end; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        const float *route = routed ? (routed + (size_t)r * (size_t)hidden_dim) : NULL;
        float *y = output + (size_t)r * (size_t)hidden_dim;
        for (int i = 0; i < intermediate_dim; ++i) {
            float gv = 0.0f;
            float uv = 0.0f;
            for (int h = 0; h < hidden_dim; ++h) {
                gv += bf16_to_float(shared_gate[(size_t)i * (size_t)hidden_dim + (size_t)h]) * x[h];
                uv += bf16_to_float(shared_up[(size_t)i * (size_t)hidden_dim + (size_t)h]) * x[h];
            }
            gate[i] = gv;
            up[i] = uv;
            act[i] = ck_moe_silu_f32(gv) * uv;
        }
        for (int h = 0; h < hidden_dim; ++h) {
            float v = route ? route[h] : 0.0f;
            for (int i = 0; i < intermediate_dim; ++i) {
                v += bf16_to_float(shared_down[(size_t)h * (size_t)intermediate_dim + (size_t)i]) * act[i];
            }
            y[h] = v;
        }
    }
}

References bf16_to_float(), and ck_moe_silu_f32().

Referenced by moe_swiglu_shared_forward_bf16().

◆ moe_swiglu_shared_forward_f32()

void moe_swiglu_shared_forward_f32	(	const float *	hidden,
		const float *	routed,
		const float *	shared_gate,
		const float *	shared_up,
		const float *	shared_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 3116 of file axpy_kernels.c.

{
    if (!hidden || !shared_gate || !shared_up || !shared_down || !output || rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0) {
        return;
    }
 
    float gate[intermediate_dim];
    float up[intermediate_dim];
    float act[intermediate_dim];
 
    for (int r = 0; r < rows; ++r) {
        const float *x = hidden + (size_t)r * (size_t)hidden_dim;
        const float *route = routed ? (routed + (size_t)r * (size_t)hidden_dim) : NULL;
        float *y = output + (size_t)r * (size_t)hidden_dim;
        for (int i = 0; i < intermediate_dim; ++i) {
            float gv = 0.0f;
            float uv = 0.0f;
            for (int h = 0; h < hidden_dim; ++h) {
                gv += shared_gate[(size_t)i * (size_t)hidden_dim + (size_t)h] * x[h];
                uv += shared_up[(size_t)i * (size_t)hidden_dim + (size_t)h] * x[h];
            }
            act[i] = ck_moe_silu_f32(gv) * uv;
        }
        for (int h = 0; h < hidden_dim; ++h) {
            float v = route ? route[h] : 0.0f;
            for (int i = 0; i < intermediate_dim; ++i) {
                v += shared_down[(size_t)h * (size_t)intermediate_dim + (size_t)i] * act[i];
            }
            y[h] = v;
        }
    }
}

References ck_moe_silu_f32().

◆ moe_swiglu_shared_forward_q4k_q4k_parallel_workspace()

int moe_swiglu_shared_forward_q4k_q4k_parallel_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2071 of file axpy_kernels.c.

{
    return ck_moe_shared_q4k_parallel_workspace(
        hidden, routed, shared_gate, shared_up, shared_down, output, rows,
        hidden_dim, intermediate_dim, workspace, workspace_bytes,
        moe_swiglu_shared_forward_q4k_q4k_workspace);
}

References ck_moe_shared_q4k_parallel_workspace(), and moe_swiglu_shared_forward_q4k_q4k_workspace().

◆ moe_swiglu_shared_forward_q4k_q4k_workspace()

int moe_swiglu_shared_forward_q4k_q4k_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 1242 of file axpy_kernels.c.

{
    const size_t required = moe_swiglu_expert_q4k_q5k_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !shared_gate || !shared_up || !shared_down || !output ||
        !workspace || required == 0 || workspace_bytes < required || rows <= 0) {
        return -1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)intermediate_dim * sizeof(float));
    const size_t act_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)intermediate_dim));
    uint8_t *cursor = (uint8_t *)workspace;
    void *hidden_q8 = cursor;
    cursor += hidden_q8_bytes;
    float *gate_up = (float *)cursor;
    cursor += gate_up_bytes;
    void *act_q8 = cursor;
    cursor += act_q8_bytes;
    float *shared_output = (float *)cursor;
 
    for (int row = 0; row < rows; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        float *y = output + (size_t)row * (size_t)hidden_dim;
        quantize_row_q8_k(x, hidden_q8, hidden_dim);
        gemv_q4_k_q8_k(gate_up, shared_gate, hidden_q8,
                       intermediate_dim, hidden_dim);
        gemv_q4_k_q8_k(gate_up + intermediate_dim, shared_up, hidden_q8,
                       intermediate_dim, hidden_dim);
        swiglu_forward_ggml(gate_up, gate_up, 1, intermediate_dim);
        quantize_row_q8_k(gate_up, act_q8, intermediate_dim);
        gemv_q4_k_q8_k(shared_output, shared_down, act_q8,
                       hidden_dim, intermediate_dim);
        const float *route = routed
            ? routed + (size_t)row * (size_t)hidden_dim
            : NULL;
        for (int col = 0; col < hidden_dim; ++col) {
            y[col] = shared_output[col] + (route ? route[col] : 0.0f);
        }
    }
    return 0;
}

References CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), gemv_q4_k_q8_k(), moe_swiglu_expert_q4k_q5k_workspace_bytes(), quantize_row_q8_k(), and swiglu_forward_ggml().

Referenced by moe_swiglu_shared_forward_q4k_q4k_parallel_workspace().

◆ moe_swiglu_shared_forward_q4k_q5_0_gated_parallel_workspace()

int moe_swiglu_shared_forward_q4k_q5_0_gated_parallel_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		const float *	shared_gate_input,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2999 of file axpy_kernels.c.

{
    /* See the Q8_0-down variant above: inner projection dispatch owns the
     * pool so nested row dispatch would deadlock. */
    return moe_swiglu_shared_forward_q4k_q5_0_gated_workspace(
        hidden, routed, shared_gate, shared_up, shared_down, shared_gate_input,
        output, rows, hidden_dim, intermediate_dim, workspace, workspace_bytes);
}

References moe_swiglu_shared_forward_q4k_q5_0_gated_workspace().

◆ moe_swiglu_shared_forward_q4k_q5_0_gated_workspace()

int moe_swiglu_shared_forward_q4k_q5_0_gated_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		const float *	shared_gate_input,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2828 of file axpy_kernels.c.

{
    return ck_moe_shared_q4k_gated_workspace(
        hidden, routed, shared_gate, shared_up, shared_down, shared_gate_input,
        output, rows, hidden_dim, intermediate_dim, workspace, workspace_bytes,
        gemv_q5_0_q8_0);
}

References ck_moe_shared_q4k_gated_workspace(), and gemv_q5_0_q8_0().

Referenced by moe_swiglu_shared_forward_q4k_q5_0_gated_parallel_workspace().

◆ moe_swiglu_shared_forward_q4k_q6k_parallel_workspace()

int moe_swiglu_shared_forward_q4k_q6k_parallel_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2052 of file axpy_kernels.c.

{
    return ck_moe_shared_q4k_parallel_workspace(
        hidden, routed, shared_gate, shared_up, shared_down, output, rows,
        hidden_dim, intermediate_dim, workspace, workspace_bytes,
        moe_swiglu_shared_forward_q4k_q6k_workspace);
}

References ck_moe_shared_q4k_parallel_workspace(), and moe_swiglu_shared_forward_q4k_q6k_workspace().

◆ moe_swiglu_shared_forward_q4k_q6k_workspace()

int moe_swiglu_shared_forward_q4k_q6k_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 1185 of file axpy_kernels.c.

{
    const size_t required = moe_swiglu_expert_q4k_q5k_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !shared_gate || !shared_up || !shared_down || !output ||
        !workspace || required == 0 || workspace_bytes < required || rows <= 0) {
        return -1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)intermediate_dim * sizeof(float));
    const size_t act_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)intermediate_dim));
    uint8_t *cursor = (uint8_t *)workspace;
    void *hidden_q8 = cursor;
    cursor += hidden_q8_bytes;
    float *gate_up = (float *)cursor;
    cursor += gate_up_bytes;
    void *act_q8 = cursor;
    cursor += act_q8_bytes;
    float *shared_output = (float *)cursor;
 
    for (int row = 0; row < rows; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        float *y = output + (size_t)row * (size_t)hidden_dim;
        quantize_row_q8_k(x, hidden_q8, hidden_dim);
        gemv_q4_k_q8_k(gate_up, shared_gate, hidden_q8,
                       intermediate_dim, hidden_dim);
        gemv_q4_k_q8_k(gate_up + intermediate_dim, shared_up, hidden_q8,
                       intermediate_dim, hidden_dim);
        swiglu_forward_ggml(gate_up, gate_up, 1, intermediate_dim);
        quantize_row_q8_k(gate_up, act_q8, intermediate_dim);
        gemv_q6_k_q8_k(shared_output, shared_down, act_q8,
                       hidden_dim, intermediate_dim);
        const float *route = routed
            ? routed + (size_t)row * (size_t)hidden_dim
            : NULL;
        for (int col = 0; col < hidden_dim; ++col) {
            y[col] = shared_output[col] + (route ? route[col] : 0.0f);
        }
    }
    return 0;
}

References CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), gemv_q4_k_q8_k(), gemv_q6_k_q8_k(), moe_swiglu_expert_q4k_q5k_workspace_bytes(), quantize_row_q8_k(), and swiglu_forward_ggml().

Referenced by moe_swiglu_shared_forward_q4k_q6k_parallel_workspace().

◆ moe_swiglu_shared_forward_q4k_q8_0_gated_parallel_workspace()

int moe_swiglu_shared_forward_q4k_q8_0_gated_parallel_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		const float *	shared_gate_input,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2978 of file axpy_kernels.c.

{
    /* Gate/up use the output-parallel pairwise Q4_K provider. Do not wrap
     * that provider in a second dispatch on the same persistent pool. */
    return moe_swiglu_shared_forward_q4k_q8_0_gated_workspace(
        hidden, routed, shared_gate, shared_up, shared_down, shared_gate_input,
        output, rows, hidden_dim, intermediate_dim, workspace, workspace_bytes);
}

References moe_swiglu_shared_forward_q4k_q8_0_gated_workspace().

◆ moe_swiglu_shared_forward_q4k_q8_0_gated_workspace()

int moe_swiglu_shared_forward_q4k_q8_0_gated_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		const float *	shared_gate_input,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2815 of file axpy_kernels.c.

{
    return ck_moe_shared_q4k_gated_workspace(
        hidden, routed, shared_gate, shared_up, shared_down, shared_gate_input,
        output, rows, hidden_dim, intermediate_dim, workspace, workspace_bytes,
        gemv_q8_0_q8_0);
}

References ck_moe_shared_q4k_gated_workspace(), and gemv_q8_0_q8_0().

Referenced by moe_swiglu_shared_forward_q4k_q8_0_gated_parallel_workspace().

◆ moe_swiglu_shared_forward_q8_0_gated_parallel_workspace()

int moe_swiglu_shared_forward_q8_0_gated_parallel_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		const float *	shared_gate_input,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2957 of file axpy_kernels.c.

{
    return ck_moe_shared_gated_parallel_workspace(
        hidden, routed, shared_gate, shared_up, shared_down, shared_gate_input,
        output, rows, hidden_dim, intermediate_dim, workspace, workspace_bytes,
        moe_swiglu_shared_q8_0_gated_workspace_bytes(hidden_dim, intermediate_dim),
        moe_swiglu_shared_forward_q8_0_gated_workspace);
}

References ck_moe_shared_gated_parallel_workspace(), moe_swiglu_shared_forward_q8_0_gated_workspace(), and moe_swiglu_shared_q8_0_gated_workspace_bytes().

◆ moe_swiglu_shared_forward_q8_0_gated_workspace()

int moe_swiglu_shared_forward_q8_0_gated_workspace	(	const float *	hidden,
		const float *	routed,
		const void *	shared_gate,
		const void *	shared_up,
		const void *	shared_down,
		const float *	shared_gate_input,
		float *	output,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		void *	workspace,
		size_t	workspace_bytes
	)

Definition at line 2642 of file axpy_kernels.c.

{
    const size_t required = moe_swiglu_shared_q8_0_gated_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !shared_gate || !shared_up || !shared_down ||
        !shared_gate_input || !output || !workspace || required == 0 ||
        workspace_bytes < required || rows <= 0) {
        return -1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)hidden_dim));
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)intermediate_dim * sizeof(float));
    const size_t activation_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)intermediate_dim));
    uint8_t *cursor = (uint8_t *)workspace;
    void *hidden_q8 = cursor;
    cursor += hidden_q8_bytes;
    float *gate_up = (float *)cursor;
    cursor += gate_up_bytes;
    void *activation_q8 = cursor;
    cursor += activation_q8_bytes;
    float *shared_output = (float *)cursor;
 
    for (int row = 0; row < rows; ++row) {
        const float *x = hidden + (size_t)row * (size_t)hidden_dim;
        const float *routed_row = routed
            ? routed + (size_t)row * (size_t)hidden_dim
            : NULL;
        float *output_row = output + (size_t)row * (size_t)hidden_dim;
 
        quantize_row_q8_0(x, hidden_q8, hidden_dim);
        gemv_q8_0_q8_0(gate_up, shared_gate, hidden_q8,
                       intermediate_dim, hidden_dim);
        gemv_q8_0_q8_0(gate_up + intermediate_dim, shared_up, hidden_q8,
                       intermediate_dim, hidden_dim);
        swiglu_forward_ggml(gate_up, gate_up, 1, intermediate_dim);
        quantize_row_q8_0(gate_up, activation_q8, intermediate_dim);
        gemv_q8_0_q8_0(shared_output, shared_down, activation_q8,
                       hidden_dim, intermediate_dim);
 
        float gate_value = 0.0f;
        gemm_nt_f32_llama_production(
            x, shared_gate_input, NULL, &gate_value, 1, 1, hidden_dim);
        const float gate_scale = 1.0f / (1.0f + expf(-gate_value));
        for (int h = 0; h < hidden_dim; ++h) {
            const float routed_value = routed_row ? routed_row[h] : 0.0f;
            volatile float gated_shared = shared_output[h] * gate_scale;
            output_row[h] = routed_value + gated_shared;
        }
    }
    return 0;
}

References CK_DT_Q8_0, ck_dtype_row_bytes(), ck_moe_align64(), gemm_nt_f32_llama_production(), gemv_q8_0_q8_0(), moe_swiglu_shared_q8_0_gated_workspace_bytes(), quantize_row_q8_0(), and swiglu_forward_ggml().

Referenced by moe_swiglu_shared_forward_q8_0_gated_parallel_workspace().

◆ moe_swiglu_shared_q4k_q8_0_gated_workspace_bytes()

size_t moe_swiglu_shared_q4k_q8_0_gated_workspace_bytes	(	int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 761 of file axpy_kernels.c.

{
    enum { batch_rows = 4 };
    if (hidden_dim <= 0 || intermediate_dim <= 0 ||
        hidden_dim % 256 != 0 || intermediate_dim % QK8_0 != 0) {
        return 0;
    }
 
    size_t bytes = ck_moe_align64(
        batch_rows * ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    bytes += 2u * ck_moe_align64(
        batch_rows * (size_t)intermediate_dim * sizeof(float));
    bytes += ck_moe_align64(
        batch_rows * ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)intermediate_dim));
    bytes += ck_moe_align64(
        batch_rows * (size_t)hidden_dim * sizeof(float));
    return bytes;
}

References CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), and QK8_0.

Referenced by ck_moe_shared_q4k_gated_workspace().

◆ moe_swiglu_shared_q8_0_gated_workspace_bytes()

size_t moe_swiglu_shared_q8_0_gated_workspace_bytes	(	int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 2625 of file axpy_kernels.c.

{
    if (hidden_dim <= 0 || intermediate_dim <= 0 ||
        hidden_dim % 32 != 0 || intermediate_dim % 32 != 0) {
        return 0;
    }
 
    size_t bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)hidden_dim));
    bytes += ck_moe_align64(2u * (size_t)intermediate_dim * sizeof(float));
    bytes += ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)intermediate_dim));
    bytes += ck_moe_align64((size_t)hidden_dim * sizeof(float));
    return bytes;
}

References CK_DT_Q8_0, ck_dtype_row_bytes(), and ck_moe_align64().

Referenced by moe_swiglu_shared_forward_q8_0_gated_parallel_workspace(), and moe_swiglu_shared_forward_q8_0_gated_workspace().

◆ mrope_qk_imrope_positions()

void mrope_qk_imrope_positions	(	float *	q,
		float *	k,
		const int32_t *	positions,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	n_dims,
		int	section_0,
		int	section_1,
		int	section_2,
		int	section_3,
		int	n_ctx_orig,
		float	freq_base,
		float	freq_scale,
		float	ext_factor,
		float	attn_factor,
		float	beta_fast,
		float	beta_slow
	)

Definition at line 2686 of file rope_kernels.c.

{
    if (!q || !k || !positions || num_heads <= 0 || num_kv_heads <= 0 || num_tokens <= 0) {
        return;
    }
 
    const int sections[4] = {section_0, section_1, section_2, section_3};
 
    if (ck_strict_parity_enabled()) {
        const int q_ok = explicit_mrope_apply_ggml_exact(
            q, positions, num_heads, num_tokens, head_dim, aligned_head_dim, n_dims, sections,
            n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, GGML_ROPE_TYPE_IMROPE);
        const int k_ok = explicit_mrope_apply_ggml_exact(
            k, positions, num_kv_heads, num_tokens, head_dim, aligned_head_dim, n_dims, sections,
            n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, GGML_ROPE_TYPE_IMROPE);
        if (q_ok && k_ok) {
            return;
        }
    }
 
    const size_t q_head_stride = (size_t) num_tokens * (size_t) aligned_head_dim;
    const size_t k_head_stride = (size_t) num_tokens * (size_t) aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        explicit_mrope_apply_head(
            q + (size_t) h * q_head_stride,
            positions,
            num_tokens,
            head_dim,
            aligned_head_dim,
            n_dims,
            sections,
            n_ctx_orig,
            freq_base,
            freq_scale,
            ext_factor,
            attn_factor,
            beta_fast,
            beta_slow,
            1
        );
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        explicit_mrope_apply_head(
            k + (size_t) h * k_head_stride,
            positions,
            num_tokens,
            head_dim,
            aligned_head_dim,
            n_dims,
            sections,
            n_ctx_orig,
            freq_base,
            freq_scale,
            ext_factor,
            attn_factor,
            beta_fast,
            beta_slow,
            1
        );
    }
}

References ck_strict_parity_enabled(), explicit_mrope_apply_ggml_exact(), explicit_mrope_apply_head(), and GGML_ROPE_TYPE_IMROPE.

◆ mrope_qk_text()

void mrope_qk_text	(	float *	q,
		float *	k,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	n_dims,
		int	section_0,
		int	section_1,
		int	section_2,
		int	section_3,
		int	n_ctx_orig,
		float	freq_base,
		float	freq_scale,
		float	ext_factor,
		float	attn_factor,
		float	beta_fast,
		float	beta_slow
	)

Definition at line 2203 of file rope_kernels.c.

{
    if (!q || !k || num_heads <= 0 || num_kv_heads <= 0 || num_tokens <= 0) {
        return;
    }
 
    const int sections[4] = {section_0, section_1, section_2, section_3};
    const size_t q_head_stride = (size_t) num_tokens * (size_t) aligned_head_dim;
    const size_t k_head_stride = (size_t) num_tokens * (size_t) aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        text_mrope_apply_head(
            q + (size_t) h * q_head_stride,
            num_tokens,
            head_dim,
            aligned_head_dim,
            pos_offset,
            n_dims,
            sections,
            n_ctx_orig,
            freq_base,
            freq_scale,
            ext_factor,
            attn_factor,
            beta_fast,
            beta_slow,
            0
        );
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        text_mrope_apply_head(
            k + (size_t) h * k_head_stride,
            num_tokens,
            head_dim,
            aligned_head_dim,
            pos_offset,
            n_dims,
            sections,
            n_ctx_orig,
            freq_base,
            freq_scale,
            ext_factor,
            attn_factor,
            beta_fast,
            beta_slow,
            0
        );
    }
}

References text_mrope_apply_head().

◆ mrope_qk_text_imrope()

void mrope_qk_text_imrope	(	float *	q,
		float *	k,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	n_dims,
		int	section_0,
		int	section_1,
		int	section_2,
		int	section_3,
		int	n_ctx_orig,
		float	freq_base,
		float	freq_scale,
		float	ext_factor,
		float	attn_factor,
		float	beta_fast,
		float	beta_slow
	)

Definition at line 2273 of file rope_kernels.c.

{
    if (!q || !k || num_heads <= 0 || num_kv_heads <= 0 || num_tokens <= 0) {
        return;
    }
 
    const int sections[4] = {section_0, section_1, section_2, section_3};
    const size_t q_head_stride = (size_t) num_tokens * (size_t) aligned_head_dim;
    const size_t k_head_stride = (size_t) num_tokens * (size_t) aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        text_mrope_apply_head(
            q + (size_t) h * q_head_stride,
            num_tokens,
            head_dim,
            aligned_head_dim,
            pos_offset,
            n_dims,
            sections,
            n_ctx_orig,
            freq_base,
            freq_scale,
            ext_factor,
            attn_factor,
            beta_fast,
            beta_slow,
            1
        );
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        text_mrope_apply_head(
            k + (size_t) h * k_head_stride,
            num_tokens,
            head_dim,
            aligned_head_dim,
            pos_offset,
            n_dims,
            sections,
            n_ctx_orig,
            freq_base,
            freq_scale,
            ext_factor,
            attn_factor,
            beta_fast,
            beta_slow,
            1
        );
    }
}

References text_mrope_apply_head().

◆ mrope_qk_text_imrope_bf16_pytorch_storage()

void mrope_qk_text_imrope_bf16_pytorch_storage	(	float *	q,
		float *	k,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	n_dims,
		int	section_0,
		int	section_1,
		int	section_2,
		int	section_3,
		int	n_ctx_orig,
		float	freq_base,
		float	freq_scale,
		float	ext_factor,
		float	attn_factor,
		float	beta_fast,
		float	beta_slow
	)

Definition at line 2496 of file rope_kernels.c.

{
    (void)section_0;
    (void)section_1;
    (void)section_2;
    (void)section_3;
    (void)n_ctx_orig;
    (void)ext_factor;
    (void)attn_factor;
    (void)beta_fast;
    (void)beta_slow;
    text_mrope_apply_pytorch_bf16_storage(q, num_heads, num_tokens, head_dim,
                                           aligned_head_dim, pos_offset, n_dims,
                                           freq_base, freq_scale);
    text_mrope_apply_pytorch_bf16_storage(k, num_kv_heads, num_tokens, head_dim,
                                           aligned_head_dim, pos_offset, n_dims,
                                           freq_base, freq_scale);
}

References text_mrope_apply_pytorch_bf16_storage().

◆ mrope_qk_text_imrope_positions_bf16_pytorch_storage()

void mrope_qk_text_imrope_positions_bf16_pytorch_storage	(	float *	q,
		float *	k,
		const int32_t *	positions,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	n_dims,
		int	section_0,
		int	section_1,
		int	section_2,
		int	section_3,
		int	n_ctx_orig,
		float	freq_base,
		float	freq_scale,
		float	ext_factor,
		float	attn_factor,
		float	beta_fast,
		float	beta_slow
	)

Definition at line 2460 of file rope_kernels.c.

{
    (void)n_ctx_orig;
    (void)ext_factor;
    (void)attn_factor;
    (void)beta_fast;
    (void)beta_slow;
    const int sections[4] = {section_0, section_1, section_2, section_3};
    text_mrope_apply_positions_pytorch_bf16_storage(
        q, positions, num_heads, num_tokens, head_dim, aligned_head_dim, n_dims,
        sections, freq_base, freq_scale);
    text_mrope_apply_positions_pytorch_bf16_storage(
        k, positions, num_kv_heads, num_tokens, head_dim, aligned_head_dim, n_dims,
        sections, freq_base, freq_scale);
}

References text_mrope_apply_positions_pytorch_bf16_storage().

◆ mrope_qk_vision()

void mrope_qk_vision	(	float *	q,
		float *	k,
		const int32_t *	positions,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	n_dims,
		int	section_0,
		int	section_1,
		int	section_2,
		int	section_3,
		int	n_ctx_orig,
		float	freq_base,
		float	freq_scale,
		float	ext_factor,
		float	attn_factor,
		float	beta_fast,
		float	beta_slow
	)

Definition at line 2534 of file rope_kernels.c.

{
    if (!q || !k || !positions || num_heads <= 0 || num_kv_heads <= 0 || num_tokens <= 0) {
        return;
    }
 
    const int sections[4] = {section_0, section_1, section_2, section_3};
 
    if (ck_strict_parity_enabled()) {
        if (explicit_mrope_apply_ggml_exact(
                q, positions, num_heads, num_tokens, head_dim, aligned_head_dim, n_dims, sections,
                n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, GGML_ROPE_TYPE_VISION) &&
            explicit_mrope_apply_ggml_exact(
                k, positions, num_kv_heads, num_tokens, head_dim, aligned_head_dim, n_dims, sections,
                n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, GGML_ROPE_TYPE_VISION)) {
            return;
        }
    }
 
    const size_t q_head_stride = (size_t) num_tokens * (size_t) aligned_head_dim;
    const size_t k_head_stride = (size_t) num_tokens * (size_t) aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        vision_mrope_apply_head(
            q + (size_t) h * q_head_stride,
            positions,
            num_tokens,
            head_dim,
            aligned_head_dim,
            n_dims,
            sections,
            n_ctx_orig,
            freq_base,
            freq_scale,
            ext_factor,
            attn_factor,
            beta_fast,
            beta_slow
        );
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        vision_mrope_apply_head(
            k + (size_t) h * k_head_stride,
            positions,
            num_tokens,
            head_dim,
            aligned_head_dim,
            n_dims,
            sections,
            n_ctx_orig,
            freq_base,
            freq_scale,
            ext_factor,
            attn_factor,
            beta_fast,
            beta_slow
        );
    }
}

References ck_strict_parity_enabled(), explicit_mrope_apply_ggml_exact(), GGML_ROPE_TYPE_VISION, and vision_mrope_apply_head().

◆ mrope_qk_vision_bf16_pytorch_storage()

void mrope_qk_vision_bf16_pytorch_storage	(	float *	q,
		float *	k,
		const int32_t *	positions,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	n_dims,
		int	section_0,
		int	section_1,
		int	section_2,
		int	section_3,
		int	n_ctx_orig,
		float	freq_base,
		float	freq_scale,
		float	ext_factor,
		float	attn_factor,
		float	beta_fast,
		float	beta_slow
	)

◆ mrope_qk_vision_bf16_storage()

void mrope_qk_vision_bf16_storage	(	float *	q,
		float *	k,
		const int32_t *	positions,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	n_dims,
		int	section_0,
		int	section_1,
		int	section_2,
		int	section_3,
		int	n_ctx_orig,
		float	freq_base,
		float	freq_scale,
		float	ext_factor,
		float	attn_factor,
		float	beta_fast,
		float	beta_slow
	)

Definition at line 2645 of file rope_kernels.c.

◆ mrope_qk_vision_fp16_storage()

void mrope_qk_vision_fp16_storage	(	float *	q,
		float *	k,
		const int32_t *	positions,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	n_dims,
		int	section_0,
		int	section_1,
		int	section_2,
		int	section_3,
		int	n_ctx_orig,
		float	freq_base,
		float	freq_scale,
		float	ext_factor,
		float	attn_factor,
		float	beta_fast,
		float	beta_slow
	)

Definition at line 2646 of file rope_kernels.c.

◆ nemotron_group_limited_topk_router_f32()

void nemotron_group_limited_topk_router_f32	(	const float *	scores,
		const float *	correction_bias,
		int *	indices,
		float *	weights,
		int	rows,
		int	n_experts,
		int	top_k,
		int	n_group,
		int	topk_group,
		int	norm_topk_prob,
		float	routed_scaling_factor
	)

Definition at line 780 of file topk_kernels.c.

{
    group_limited_topk_router_f32_impl(
        scores, correction_bias, indices, weights, rows, n_experts, top_k,
        n_group, topk_group, norm_topk_prob, routed_scaling_factor, 0
    );
}

References group_limited_topk_router_f32_impl().

◆ patch2im()

void patch2im	(	const float *	d_patches,
		float *	d_image,
		int	C,
		int	H,
		int	W,
		int	P
	)

patch2im: Accumulates gradients from patches back into the image. (Backward pass)

d_patches: [num_patches, C * P * P] d_image: [C, H, W] (Accumulated)

Definition at line 123 of file vision_kernels.c.

{
    int num_patches_h = H / P;
    int num_patches_w = W / P;
    int patch_dim = C * P * P;
 
    // Zero out the image first as we are accumulating gradients
    memset(d_image, 0, (size_t)C * H * W * sizeof(float));
 
    for (int ph = 0; ph < num_patches_h; ++ph) {
        for (int pw = 0; pw < num_patches_w; ++pw) {
            
            int patch_idx = ph * num_patches_w + pw;
            const float *src_patch = d_patches + (size_t)patch_idx * patch_dim;
 
            for (int c = 0; c < C; ++c) {
                for (int py = 0; py < P; ++py) {
                    int y = ph * P + py;
                    int x = pw * P;
                    
                    float *dst_row = d_image + (size_t)c * H * W + (size_t)y * W + x;
                    const float *src_row = src_patch + (size_t)c * P * P + (size_t)py * P;
                    
                    // Add the patch gradient to the image gradient
                    for (int px = 0; px < P; ++px) {
                        dst_row[px] += src_row[px];
                    }
                }
            }
        }
    }
}

References C.

◆ patch2im_bf16()

void patch2im_bf16	(	const uint16_t *	d_patches,
		uint16_t *	d_image,
		int	C,
		int	H,
		int	W,
		int	P
	)

Definition at line 57 of file vision_kernels_bf16.c.

{
    if (!d_patches || !d_image || C <= 0 || H <= 0 || W <= 0 || P <= 0) {
        return;
    }
 
    int num_patches_h = H / P;
    int num_patches_w = W / P;
    int patch_dim = C * P * P;
 
    memset(d_image, 0, (size_t)C * (size_t)H * (size_t)W * sizeof(uint16_t));
 
    for (int ph = 0; ph < num_patches_h; ++ph) {
        for (int pw = 0; pw < num_patches_w; ++pw) {
            int patch_idx = ph * num_patches_w + pw;
            const uint16_t *src_patch = d_patches + (size_t)patch_idx * (size_t)patch_dim;
 
            for (int c = 0; c < C; ++c) {
                for (int py = 0; py < P; ++py) {
                    int y = ph * P + py;
                    int x = pw * P;
 
                    uint16_t *dst_row = d_image + (size_t)c * (size_t)H * (size_t)W + (size_t)y * (size_t)W + (size_t)x;
                    const uint16_t *src_row = src_patch + (size_t)c * (size_t)P * (size_t)P + (size_t)py * (size_t)P;
 
                    for (int px = 0; px < P; ++px) {
                        float acc = bf16_to_float(dst_row[px]) + bf16_to_float(src_row[px]);
                        dst_row[px] = float_to_bf16(acc);
                    }
                }
            }
        }
    }
}

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

◆ patch_projection_bf16_pytorch_onednn_conv3d_storage()

void patch_projection_bf16_pytorch_onednn_conv3d_storage	(	const float *	input,
		const void *	weights,
		const float *	bias,
		float *	output,
		int	batch,
		int	out_channels,
		int	in_channels,
		int	temporal,
		int	patch_h,
		int	patch_w
	)

Definition at line 1564 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    if (!input || !weights || !bias || !output || batch <= 0 ||
        out_channels <= 0 || in_channels <= 0 || temporal <= 0 ||
        patch_h <= 0 || patch_w <= 0) {
        ck_pytorch_brgemm_fault("invalid Conv3D patch contract", batch,
                                out_channels, in_channels * temporal * patch_h * patch_w);
    }
    pthread_once(&ck_pytorch_brgemm_once, ck_pytorch_brgemm_init);
    if (ck_pytorch_brgemm_init_status != 0) {
        ck_pytorch_brgemm_fault("could not initialize oneDNN Conv3D", batch,
                                out_channels, in_channels * temporal * patch_h * patch_w);
    }
    ck_pytorch_brgemm_require_version(
        3, 7, 1, "8d263e693366ef8db40acc569cc7d8edf644556d",
        "patch_projection_bf16_pytorch_onednn_conv3d_storage",
        batch, out_channels, in_channels * temporal * patch_h * patch_w);
 
    const size_t input_count = (size_t)batch * (size_t)in_channels *
        (size_t)temporal * (size_t)patch_h * (size_t)patch_w;
    const size_t output_count = (size_t)batch * (size_t)out_channels;
    uint16_t *input_bf16 = (uint16_t *)malloc(input_count * sizeof(*input_bf16));
    uint16_t *bias_bf16 = (uint16_t *)malloc((size_t)out_channels * sizeof(*bias_bf16));
    uint16_t *output_bf16 = (uint16_t *)malloc(output_count * sizeof(*output_bf16));
    if (!input_bf16 || !bias_bf16 || !output_bf16) {
        free(output_bf16);
        free(bias_bf16);
        free(input_bf16);
        ck_pytorch_brgemm_fault("Conv3D workspace allocation failed", batch,
                                out_channels, in_channels * temporal * patch_h * patch_w);
    }
    for (size_t i = 0; i < input_count; ++i) input_bf16[i] = float_to_bf16(input[i]);
    for (int i = 0; i < out_channels; ++i) bias_bf16[i] = float_to_bf16(bias[i]);
 
    dnnl_dims_t src_dims = {batch, in_channels, temporal, patch_h, patch_w};
    dnnl_dims_t weight_dims = {
        out_channels, in_channels, temporal, patch_h, patch_w};
    dnnl_dims_t bias_dims = {out_channels};
    dnnl_dims_t dst_dims = {batch, out_channels, 1, 1, 1};
    dnnl_dims_t strides = {temporal, patch_h, patch_w};
    dnnl_dims_t dilates = {0, 0, 0};
    dnnl_dims_t padding = {0, 0, 0};
 
    dnnl_memory_desc_t user_src_md = NULL, user_weight_md = NULL;
    dnnl_memory_desc_t bias_md = NULL, user_dst_md = NULL;
    dnnl_memory_desc_t any_src_md = NULL, any_weight_md = NULL, any_dst_md = NULL;
    dnnl_primitive_desc_t conv_pd = NULL;
    dnnl_primitive_t conv = NULL;
    dnnl_primitive_desc_t reorder_pd = NULL;
    dnnl_primitive_t reorder = NULL;
    dnnl_exec_arg_t reorder_args[2];
    dnnl_memory_t user_src = NULL, user_weight = NULL, bias_mem = NULL, user_dst = NULL;
    dnnl_memory_t conv_src = NULL, conv_weight = NULL, conv_dst = NULL;
    dnnl_status_t status = dnnl_success;
 
#define CK_DNNL_CONV(call) do { status = (call); if (status != dnnl_success) goto cleanup_conv; } while (0)
    pthread_mutex_lock(&ck_pytorch_brgemm_lock);
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &user_src_md, 5, src_dims, dnnl_bf16, dnnl_ncdhw));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &user_weight_md, 5, weight_dims, dnnl_bf16, dnnl_oidhw));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &bias_md, 1, bias_dims, dnnl_bf16, dnnl_x));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &user_dst_md, 5, dst_dims, dnnl_bf16, dnnl_ncdhw));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &any_src_md, 5, src_dims, dnnl_bf16, dnnl_format_tag_any));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &any_weight_md, 5, weight_dims, dnnl_bf16, dnnl_format_tag_any));
    CK_DNNL_CONV(dnnl_memory_desc_create_with_tag(
        &any_dst_md, 5, dst_dims, dnnl_bf16, dnnl_format_tag_any));
    CK_DNNL_CONV(dnnl_convolution_forward_primitive_desc_create(
        &conv_pd, ck_pytorch_brgemm_engine, dnnl_forward_training,
        dnnl_convolution_direct, any_src_md, any_weight_md, bias_md, any_dst_md,
        strides, dilates, padding, padding, NULL));
 
    const_dnnl_memory_desc_t conv_src_md =
        dnnl_primitive_desc_query_md(conv_pd, dnnl_query_src_md, 0);
    const_dnnl_memory_desc_t conv_weight_md =
        dnnl_primitive_desc_query_md(conv_pd, dnnl_query_weights_md, 0);
    const_dnnl_memory_desc_t conv_dst_md =
        dnnl_primitive_desc_query_md(conv_pd, dnnl_query_dst_md, 0);
    CK_DNNL_CONV(dnnl_memory_create(
        &user_src, user_src_md, ck_pytorch_brgemm_engine, input_bf16));
    CK_DNNL_CONV(dnnl_memory_create(
        &user_weight, user_weight_md, ck_pytorch_brgemm_engine, (void *)weights));
    CK_DNNL_CONV(dnnl_memory_create(
        &bias_mem, bias_md, ck_pytorch_brgemm_engine, bias_bf16));
    CK_DNNL_CONV(dnnl_memory_create(
        &user_dst, user_dst_md, ck_pytorch_brgemm_engine, output_bf16));
    CK_DNNL_CONV(dnnl_memory_create(
        &conv_src, conv_src_md, ck_pytorch_brgemm_engine, DNNL_MEMORY_ALLOCATE));
    CK_DNNL_CONV(dnnl_memory_create(
        &conv_weight, conv_weight_md, ck_pytorch_brgemm_engine, DNNL_MEMORY_ALLOCATE));
    CK_DNNL_CONV(dnnl_memory_create(
        &conv_dst, conv_dst_md, ck_pytorch_brgemm_engine, DNNL_MEMORY_ALLOCATE));
 
    CK_DNNL_CONV(dnnl_reorder_primitive_desc_create(
        &reorder_pd, user_src_md, ck_pytorch_brgemm_engine,
        conv_src_md, ck_pytorch_brgemm_engine, NULL));
    CK_DNNL_CONV(dnnl_primitive_create(&reorder, reorder_pd));
    reorder_args[0] = (dnnl_exec_arg_t){DNNL_ARG_FROM, user_src};
    reorder_args[1] = (dnnl_exec_arg_t){DNNL_ARG_TO, conv_src};
    CK_DNNL_CONV(dnnl_primitive_execute(
        reorder, ck_pytorch_brgemm_stream, 2, reorder_args));
    dnnl_primitive_destroy(reorder); reorder = NULL;
    dnnl_primitive_desc_destroy(reorder_pd); reorder_pd = NULL;
 
    CK_DNNL_CONV(dnnl_reorder_primitive_desc_create(
        &reorder_pd, user_weight_md, ck_pytorch_brgemm_engine,
        conv_weight_md, ck_pytorch_brgemm_engine, NULL));
    CK_DNNL_CONV(dnnl_primitive_create(&reorder, reorder_pd));
    reorder_args[0] = (dnnl_exec_arg_t){DNNL_ARG_FROM, user_weight};
    reorder_args[1] = (dnnl_exec_arg_t){DNNL_ARG_TO, conv_weight};
    CK_DNNL_CONV(dnnl_primitive_execute(
        reorder, ck_pytorch_brgemm_stream, 2, reorder_args));
    dnnl_primitive_destroy(reorder); reorder = NULL;
    dnnl_primitive_desc_destroy(reorder_pd); reorder_pd = NULL;
 
    CK_DNNL_CONV(dnnl_primitive_create(&conv, conv_pd));
    dnnl_exec_arg_t conv_args[] = {
        {DNNL_ARG_SRC, conv_src},
        {DNNL_ARG_WEIGHTS, conv_weight},
        {DNNL_ARG_BIAS, bias_mem},
        {DNNL_ARG_DST, conv_dst},
    };
    CK_DNNL_CONV(dnnl_primitive_execute(
        conv, ck_pytorch_brgemm_stream,
        (int)(sizeof(conv_args) / sizeof(conv_args[0])), conv_args));
 
    CK_DNNL_CONV(dnnl_reorder_primitive_desc_create(
        &reorder_pd, conv_dst_md, ck_pytorch_brgemm_engine,
        user_dst_md, ck_pytorch_brgemm_engine, NULL));
    CK_DNNL_CONV(dnnl_primitive_create(&reorder, reorder_pd));
    reorder_args[0] = (dnnl_exec_arg_t){DNNL_ARG_FROM, conv_dst};
    reorder_args[1] = (dnnl_exec_arg_t){DNNL_ARG_TO, user_dst};
    CK_DNNL_CONV(dnnl_primitive_execute(
        reorder, ck_pytorch_brgemm_stream, 2, reorder_args));
    CK_DNNL_CONV(dnnl_stream_wait(ck_pytorch_brgemm_stream));
 
cleanup_conv:
    if (reorder) dnnl_primitive_destroy(reorder);
    if (reorder_pd) dnnl_primitive_desc_destroy(reorder_pd);
    if (conv) dnnl_primitive_destroy(conv);
    if (conv_dst) dnnl_memory_destroy(conv_dst);
    if (conv_weight) dnnl_memory_destroy(conv_weight);
    if (conv_src) dnnl_memory_destroy(conv_src);
    if (user_dst) dnnl_memory_destroy(user_dst);
    if (bias_mem) dnnl_memory_destroy(bias_mem);
    if (user_weight) dnnl_memory_destroy(user_weight);
    if (user_src) dnnl_memory_destroy(user_src);
    if (conv_pd) dnnl_primitive_desc_destroy(conv_pd);
    if (any_dst_md) dnnl_memory_desc_destroy(any_dst_md);
    if (any_weight_md) dnnl_memory_desc_destroy(any_weight_md);
    if (any_src_md) dnnl_memory_desc_destroy(any_src_md);
    if (user_dst_md) dnnl_memory_desc_destroy(user_dst_md);
    if (bias_md) dnnl_memory_desc_destroy(bias_md);
    if (user_weight_md) dnnl_memory_desc_destroy(user_weight_md);
    if (user_src_md) dnnl_memory_desc_destroy(user_src_md);
    pthread_mutex_unlock(&ck_pytorch_brgemm_lock);
#undef CK_DNNL_CONV
 
    if (status != dnnl_success) {
        free(output_bf16);
        free(bias_bf16);
        free(input_bf16);
        ck_pytorch_brgemm_fault("oneDNN Conv3D execution failed", batch,
                                out_channels, in_channels * temporal * patch_h * patch_w);
    }
    for (size_t i = 0; i < output_count; ++i) output[i] = bf16_to_float(output_bf16[i]);
    free(output_bf16);
    free(bias_bf16);
    free(input_bf16);
#else
    (void)input; (void)weights; (void)bias; (void)output; (void)batch;
    (void)out_channels; (void)in_channels; (void)temporal; (void)patch_h; (void)patch_w;
    fprintf(stderr, "HARD KERNEL CONTRACT FAULT: PyTorch oneDNN BF16 Conv3D "
                    "was selected without USE_ONEDNN=1\n");
    abort();
#endif
}

References bf16_to_float(), and float_to_bf16().

Referenced by patch_projection_image_bf16_pytorch_onednn_conv3d_storage().

◆ patch_projection_image_bf16_native_storage()

void patch_projection_image_bf16_native_storage	(	const float *	image,
		const void *	weights_t0,
		const void *	weights_t1,
		const float *	bias,
		float *	output,
		int	channels,
		int	image_h,
		int	image_w,
		int	patch_size,
		int	out_channels,
		int	merge_size
	)

Definition at line 1940 of file gemm_kernels_bf16.c.

{
    if (!image || !weights_t0 || !weights_t1 || !output || channels <= 0 ||
        image_h <= 0 || image_w <= 0 || patch_size <= 0 || out_channels <= 0 ||
        merge_size <= 0 || image_h % patch_size != 0 ||
        image_w % patch_size != 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid native BF16 image patch projection\n");
        abort();
    }
    const int grid_h = image_h / patch_size;
    const int grid_w = image_w / patch_size;
    if (grid_h % merge_size != 0 || grid_w % merge_size != 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: native BF16 patch grid is not merge aligned\n");
        abort();
    }
    ck_patch_projection_bf16_native_args_t args = {
        .image = image,
        .weights_t0 = (const uint16_t *)weights_t0,
        .weights_t1 = (const uint16_t *)weights_t1,
        .bias = bias,
        .output = output,
        .channels = channels,
        .image_h = image_h,
        .image_w = image_w,
        .patch_size = patch_size,
        .out_channels = out_channels,
        .merge_size = merge_size,
        .grid_w = grid_w,
        .batch = grid_h * grid_w,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    if (active > args.batch) active = args.batch;
    if (active > 24) active = 24;
    if (pool && active > 1) {
        ck_threadpool_dispatch_n(
            pool, active, ck_patch_projection_bf16_native_work, &args);
    } else {
        ck_patch_projection_bf16_native_work(0, 1, &args);
    }
}

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

◆ patch_projection_image_bf16_pytorch_onednn_conv3d_storage()

void patch_projection_image_bf16_pytorch_onednn_conv3d_storage	(	const float *	image,
		const void *	weights_t0,
		const void *	weights_t1,
		const float *	bias,
		float *	output,
		int	channels,
		int	image_h,
		int	image_w,
		int	patch_size,
		int	out_channels,
		int	merge_size
	)

Definition at line 1750 of file gemm_kernels_bf16.c.

{
#ifdef USE_ONEDNN
    if (!image || !weights_t0 || !weights_t1 || !bias || !output ||
        channels <= 0 || image_h <= 0 || image_w <= 0 || patch_size <= 0 ||
        out_channels <= 0 || merge_size <= 0 || image_h % patch_size != 0 ||
        image_w % patch_size != 0) {
        ck_pytorch_brgemm_fault("invalid image patch projection contract",
                                image_h, image_w, patch_size);
    }
    const int grid_h = image_h / patch_size;
    const int grid_w = image_w / patch_size;
    if (grid_h % merge_size != 0 || grid_w % merge_size != 0) {
        ck_pytorch_brgemm_fault("patch grid is not merge-tile aligned",
                                grid_h, grid_w, merge_size);
    }
    const int batch = grid_h * grid_w;
    const int temporal = 2;
    const int half_k = channels * patch_size * patch_size;
    const int full_k = temporal * half_k;
    float *patches = (float *)malloc((size_t)batch * (size_t)full_k * sizeof(*patches));
    uint16_t *weights = (uint16_t *)malloc(
        (size_t)out_channels * (size_t)full_k * sizeof(*weights));
    if (!patches || !weights) {
        free(weights);
        free(patches);
        ck_pytorch_brgemm_fault("image patch projection workspace allocation failed",
                                batch, out_channels, full_k);
    }
 
    for (int tok = 0; tok < batch; ++tok) {
        const int tiles_per_row = grid_w / merge_size;
        const int tile_area = merge_size * merge_size;
        const int tile = tok / tile_area;
        const int within = tok % tile_area;
        const int patch_y = (tile / tiles_per_row) * merge_size + within / merge_size;
        const int patch_x = (tile % tiles_per_row) * merge_size + within % merge_size;
        float *dst = patches + (size_t)tok * (size_t)full_k;
        for (int c = 0; c < channels; ++c) {
            for (int t = 0; t < temporal; ++t) {
                for (int py = 0; py < patch_size; ++py) {
                    const float *src = image +
                        ((size_t)c * (size_t)image_h +
                         (size_t)(patch_y * patch_size + py)) * (size_t)image_w +
                        (size_t)(patch_x * patch_size);
                    memcpy(dst, src, (size_t)patch_size * sizeof(*dst));
                    dst += patch_size;
                }
            }
        }
    }
 
    const uint16_t *w0 = (const uint16_t *)weights_t0;
    const uint16_t *w1 = (const uint16_t *)weights_t1;
    for (int n = 0; n < out_channels; ++n) {
        uint16_t *dst = weights + (size_t)n * (size_t)full_k;
        for (int c = 0; c < channels; ++c) {
            const size_t channel_offset =
                (size_t)n * (size_t)half_k +
                (size_t)c * (size_t)patch_size * (size_t)patch_size;
            const size_t plane_bytes =
                (size_t)patch_size * (size_t)patch_size * sizeof(*dst);
            memcpy(dst, w0 + channel_offset, plane_bytes);
            dst += patch_size * patch_size;
            memcpy(dst, w1 + channel_offset, plane_bytes);
            dst += patch_size * patch_size;
        }
    }
 
    patch_projection_bf16_pytorch_onednn_conv3d_storage(
        patches, weights, bias, output, batch, out_channels, channels,
        temporal, patch_size, patch_size);
    free(weights);
    free(patches);
#else
    (void)image; (void)weights_t0; (void)weights_t1; (void)bias; (void)output;
    (void)channels; (void)image_h; (void)image_w; (void)patch_size;
    (void)out_channels; (void)merge_size;
    fprintf(stderr, "HARD KERNEL CONTRACT FAULT: PyTorch oneDNN BF16 image "
                    "patch projection was selected without USE_ONEDNN=1\n");
    abort();
#endif
}

References patch_projection_bf16_pytorch_onednn_conv3d_storage().

◆ position_embeddings_add()

void position_embeddings_add	(	float *	x,
		const float *	position_embd,
		int	num_tokens,
		int	embed_dim,
		int	num_positions
	)

Add learned absolute position embeddings in-place.

x layout: [num_tokens, embed_dim] position_embd: [num_positions, embed_dim]

This first v8 vision path intentionally assumes native resized embeddings are already materialized in the weight tensor, so token i maps directly to position_embd[i]. This is the correct contract for fixed-size bring-up.

Definition at line 168 of file vision_kernels.c.

{
    if (x == NULL || position_embd == NULL || num_tokens <= 0 || embed_dim <= 0) {
        return;
    }
    const int limit = num_tokens < num_positions ? num_tokens : num_positions;
    for (int tok = 0; tok < limit; ++tok) {
        float *dst = x + (size_t)tok * embed_dim;
        const float *src = position_embd + (size_t)tok * embed_dim;
        for (int d = 0; d < embed_dim; ++d) {
            dst[d] += src[d];
        }
    }
}

◆ position_embeddings_add_at_offset()

void position_embeddings_add_at_offset	(	float *	x,
		const float *	position_embd,
		int	num_tokens,
		int	embed_dim,
		int	num_positions,
		int	start_position
	)

Definition at line 187 of file vision_kernels.c.

{
    if (x == NULL || position_embd == NULL || num_tokens <= 0 || embed_dim <= 0 ||
        num_positions <= 0 || start_position < 0 || start_position >= num_positions) {
        return;
    }
    const int available = num_positions - start_position;
    const int limit = num_tokens < available ? num_tokens : available;
    for (int tok = 0; tok < limit; ++tok) {
        float *dst = x + (size_t)tok * embed_dim;
        const float *src = position_embd + (size_t)(start_position + tok) * embed_dim;
        for (int d = 0; d < embed_dim; ++d) {
            dst[d] += src[d];
        }
    }
}

◆ position_embeddings_add_gemma4v_xy()

void position_embeddings_add_gemma4v_xy	(	float *	x,
		const float *	position_embd,
		int	grid_h,
		int	grid_w,
		int	embed_dim,
		int	source_grid_size
	)

Definition at line 209 of file vision_kernels.c.

{
    if (x == NULL || position_embd == NULL || grid_h <= 0 || grid_w <= 0 || embed_dim <= 0) {
        return;
    }
    if (source_grid_size <= 0) {
        source_grid_size = grid_w > grid_h ? grid_w : grid_h;
    }
 
    const size_t table_stride = (size_t) source_grid_size * (size_t) embed_dim;
    const float *table_x = position_embd;
    const float *table_y = position_embd + table_stride;
 
    for (int y = 0; y < grid_h; ++y) {
        const int yy = y < source_grid_size ? y : (source_grid_size - 1);
        const float *row_y = table_y + (size_t) yy * (size_t) embed_dim;
        for (int x_pos = 0; x_pos < grid_w; ++x_pos) {
            const int xx = x_pos < source_grid_size ? x_pos : (source_grid_size - 1);
            const float *row_x = table_x + (size_t) xx * (size_t) embed_dim;
            float *dst = x + ((size_t) y * (size_t) grid_w + (size_t) x_pos) * (size_t) embed_dim;
            for (int d = 0; d < embed_dim; ++d) {
                dst[d] += row_x[d] + row_y[d];
            }
        }
    }
}

◆ position_embeddings_add_tiled_2d()

void position_embeddings_add_tiled_2d	(	float *	x,
		const float *	position_embd,
		int	grid_h,
		int	grid_w,
		int	embed_dim,
		int	merge_size,
		int	source_grid_size
	)

Definition at line 248 of file vision_kernels.c.

{
    if (x == NULL || position_embd == NULL || grid_h <= 0 || grid_w <= 0 || embed_dim <= 0 || merge_size <= 0) {
        return;
    }
 
    if (source_grid_size <= 0) {
        source_grid_size = grid_h == grid_w ? grid_h : (grid_h > grid_w ? grid_h : grid_w);
    }
 
    const int num_tokens = grid_h * grid_w;
    const int source_tokens = source_grid_size * source_grid_size;
    const int needs_resize = source_grid_size != grid_h || source_grid_size != grid_w;
 
    if (!needs_resize) {
        for (int tok = 0; tok < num_tokens; ++tok) {
            const int flat = tile_order_index_2d(tok, grid_h, grid_w, merge_size);
            if (flat < 0 || flat >= source_tokens) {
                continue;
            }
            float *dst = x + (size_t) tok * (size_t) embed_dim;
            const float *src = position_embd + (size_t) flat * (size_t) embed_dim;
            for (int d = 0; d < embed_dim; ++d) {
                dst[d] += src[d];
            }
        }
        return;
    }
 
    const float sf_x = (float) grid_w / (float) source_grid_size;
    const float sf_y = (float) grid_h / (float) source_grid_size;
    const float pixel_offset = 0.5f;
    const float support_x = fmaxf(1.0f, 1.0f / sf_x);
    const float support_y = fmaxf(1.0f, 1.0f / sf_y);
    const float invscale_x = 1.0f / support_x;
    const float invscale_y = 1.0f / support_y;
 
    /* Preserve ggml's channel -> row -> column loop nesting. The destination
     * index converts the row-major interpolation result directly to CK's
     * merge-tiled token layout without allocating an intermediate tensor. */
    for (int d = 0; d < embed_dim; ++d) {
        for (int dst_y = 0; dst_y < grid_h; ++dst_y) {
            const float y_src = ((float) dst_y + pixel_offset) / sf_y;
            int y_min = (int) (y_src - support_y + pixel_offset);
            int y_max = (int) (y_src + support_y + pixel_offset);
            if (y_min < 0) y_min = 0;
            if (y_max > source_grid_size) y_max = source_grid_size;
 
            for (int dst_x = 0; dst_x < grid_w; ++dst_x) {
                const float x_src = ((float) dst_x + pixel_offset) / sf_x;
                int x_min = (int) (x_src - support_x + pixel_offset);
                int x_max = (int) (x_src + support_x + pixel_offset);
                if (x_min < 0) x_min = 0;
                if (x_max > source_grid_size) x_max = source_grid_size;
                float val = 0.0f;
                float total_weight = 0.0f;
                for (int sy = y_min; sy < y_max; ++sy) {
                    const float wy_arg = ((float) sy - y_src + pixel_offset) * invscale_y;
                    const float wy = fmaxf(1.0f - fabsf(wy_arg), 0.0f);
                    if (wy <= 0.0f) {
                        continue;
                    }
                    for (int sx = x_min; sx < x_max; ++sx) {
                        const float wx_arg = ((float) sx - x_src + pixel_offset) * invscale_x;
                        const float wx = fmaxf(1.0f - fabsf(wx_arg), 0.0f);
                        const float weight = wx * wy;
                        if (weight <= 0.0f) {
                            continue;
                        }
                        const float sample = position_embd[((size_t) sy * (size_t) source_grid_size + (size_t) sx) * (size_t) embed_dim + (size_t) d];
                        val = fmaf(sample, weight, val);
                        total_weight += weight;
                    }
                }
                if (total_weight > 0.0f) {
                    val /= total_weight;
                    const int tok = tile_order_linear_index_2d(
                        dst_y, dst_x, grid_h, grid_w, merge_size);
                    x[(size_t) tok * (size_t) embed_dim + (size_t) d] += val;
                }
            }
        }
    }
}

References tile_order_index_2d(), and tile_order_linear_index_2d().

◆ position_embeddings_add_tiled_2d_align_corners()

void position_embeddings_add_tiled_2d_align_corners	(	float *	x,
		const float *	position_embd,
		int	grid_h,
		int	grid_w,
		int	embed_dim,
		int	merge_size,
		int	source_grid_size
	)

Definition at line 342 of file vision_kernels.c.

{
    if (x == NULL || position_embd == NULL || grid_h <= 0 || grid_w <= 0 ||
        embed_dim <= 0 || merge_size <= 0 || source_grid_size <= 0) {
        return;
    }
 
    const float y_scale = grid_h > 1
        ? (float)(source_grid_size - 1) / (float)(grid_h - 1)
        : 0.0f;
    const float x_scale = grid_w > 1
        ? (float)(source_grid_size - 1) / (float)(grid_w - 1)
        : 0.0f;
 
    for (int tok = 0; tok < grid_h * grid_w; ++tok) {
        const int row_major = tile_order_index_2d(tok, grid_h, grid_w, merge_size);
        const int dst_y = row_major / grid_w;
        const int dst_x = row_major % grid_w;
        const float src_y = (float)dst_y * y_scale;
        const float src_x = (float)dst_x * x_scale;
        const int y0 = (int)src_y;
        const int x0 = (int)src_x;
        const int y1 = y0 + 1 < source_grid_size ? y0 + 1 : y0;
        const int x1 = x0 + 1 < source_grid_size ? x0 + 1 : x0;
        const float dy = src_y - (float)y0;
        const float dx = src_x - (float)x0;
        const float w00 = (1.0f - dy) * (1.0f - dx);
        const float w01 = (1.0f - dy) * dx;
        const float w10 = dy * (1.0f - dx);
        const float w11 = dy * dx;
        const float *p00 = position_embd + ((size_t)y0 * source_grid_size + x0) * embed_dim;
        const float *p01 = position_embd + ((size_t)y0 * source_grid_size + x1) * embed_dim;
        const float *p10 = position_embd + ((size_t)y1 * source_grid_size + x0) * embed_dim;
        const float *p11 = position_embd + ((size_t)y1 * source_grid_size + x1) * embed_dim;
        float *dst = x + (size_t)tok * embed_dim;
 
        for (int d = 0; d < embed_dim; ++d) {
            const float pos = p00[d] * w00 + p01[d] * w01 + p10[d] * w10 + p11[d] * w11;
            dst[d] += pos;
        }
    }
}

References tile_order_index_2d().

◆ position_embeddings_add_tiled_2d_align_corners_bf16()

void position_embeddings_add_tiled_2d_align_corners_bf16	(	float *	x,
		const float *	position_embd,
		int	grid_h,
		int	grid_w,
		int	embed_dim,
		int	merge_size,
		int	source_grid_size
	)

Definition at line 397 of file vision_kernels.c.

{
    if (x == NULL || position_embd == NULL || grid_h <= 0 || grid_w <= 0 ||
        embed_dim <= 0 || merge_size <= 0 || source_grid_size <= 0) {
        return;
    }
 
    for (int tok = 0; tok < grid_h * grid_w; ++tok) {
        const int row_major = tile_order_index_2d(tok, grid_h, grid_w, merge_size);
        const int dst_y = row_major / grid_w;
        const int dst_x = row_major % grid_w;
        /* torch.linspace computes each rational coordinate before FP32 storage. */
        const float src_y = grid_h > 1
            ? (float)((double)dst_y * (double)(source_grid_size - 1) / (double)(grid_h - 1))
            : 0.0f;
        const float src_x = grid_w > 1
            ? (float)((double)dst_x * (double)(source_grid_size - 1) / (double)(grid_w - 1))
            : 0.0f;
        const int y0 = (int)src_y;
        const int x0 = (int)src_x;
        const int y1 = y0 + 1 < source_grid_size ? y0 + 1 : y0;
        const int x1 = x0 + 1 < source_grid_size ? x0 + 1 : x0;
        const float dy = src_y - (float)y0;
        const float dx = src_x - (float)x0;
        const float w00 = bf16_to_float(float_to_bf16((1.0f - dy) * (1.0f - dx)));
        const float w01 = bf16_to_float(float_to_bf16((1.0f - dy) * dx));
        const float w10 = bf16_to_float(float_to_bf16(dy * (1.0f - dx)));
        const float w11 = bf16_to_float(float_to_bf16(dy * dx));
        const float *p00 = position_embd + ((size_t)y0 * source_grid_size + x0) * embed_dim;
        const float *p01 = position_embd + ((size_t)y0 * source_grid_size + x1) * embed_dim;
        const float *p10 = position_embd + ((size_t)y1 * source_grid_size + x0) * embed_dim;
        const float *p11 = position_embd + ((size_t)y1 * source_grid_size + x1) * embed_dim;
        float *dst = x + (size_t)tok * embed_dim;
 
        for (int d = 0; d < embed_dim; ++d) {
            const float v00 = bf16_to_float(float_to_bf16(p00[d] * w00));
            const float v01 = bf16_to_float(float_to_bf16(p01[d] * w01));
            const float v10 = bf16_to_float(float_to_bf16(p10[d] * w10));
            const float v11 = bf16_to_float(float_to_bf16(p11[d] * w11));
            float pos = bf16_to_float(float_to_bf16(v00 + v01));
            pos = bf16_to_float(float_to_bf16(pos + v10));
            pos = bf16_to_float(float_to_bf16(pos + v11));
            const float hidden = bf16_to_float(float_to_bf16(dst[d]));
            dst[d] = bf16_to_float(float_to_bf16(hidden + pos));
        }
    }
}

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

◆ position_embeddings_add_tiled_2d_align_corners_fp32_interp_bf16()

void position_embeddings_add_tiled_2d_align_corners_fp32_interp_bf16	(	float *	x,
		const float *	position_embd,
		int	grid_h,
		int	grid_w,
		int	embed_dim,
		int	merge_size,
		int	source_grid_size
	)

Definition at line 458 of file vision_kernels.c.

{
    if (x == NULL || position_embd == NULL || grid_h <= 0 || grid_w <= 0 ||
        embed_dim <= 0 || merge_size <= 0 || source_grid_size <= 0) {
        return;
    }
 
    for (int tok = 0; tok < grid_h * grid_w; ++tok) {
        const int row_major = tile_order_index_2d(tok, grid_h, grid_w, merge_size);
        const int dst_y = row_major / grid_w;
        const int dst_x = row_major % grid_w;
        const float src_y = grid_h > 1
            ? (float)dst_y * (float)(source_grid_size - 1) / (float)(grid_h - 1)
            : 0.0f;
        const float src_x = grid_w > 1
            ? (float)dst_x * (float)(source_grid_size - 1) / (float)(grid_w - 1)
            : 0.0f;
        const int y0 = (int)src_y;
        const int x0 = (int)src_x;
        const int y1 = y0 + 1 < source_grid_size ? y0 + 1 : y0;
        const int x1 = x0 + 1 < source_grid_size ? x0 + 1 : x0;
        const float y_distance0 = fabsf(src_y - (float)y0);
        const float y_distance1 = fabsf(src_y - (float)y0 - 1.0f);
        const float x_distance0 = fabsf(src_x - (float)x0);
        const float x_distance1 = fabsf(src_x - (float)x0 - 1.0f);
        const float wy0 = fmaxf(1.0f - y_distance0, 0.0f);
        const float wy1 = fmaxf(1.0f - y_distance1, 0.0f);
        const float wx0 = fmaxf(1.0f - x_distance0, 0.0f);
        const float wx1 = fmaxf(1.0f - x_distance1, 0.0f);
        const float w00 = wy0 * wx0;
        const float w01 = wy0 * wx1;
        const float w10 = wy1 * wx0;
        const float w11 = wy1 * wx1;
        const float *p00 = position_embd + ((size_t)y0 * source_grid_size + x0) * embed_dim;
        const float *p01 = position_embd + ((size_t)y0 * source_grid_size + x1) * embed_dim;
        const float *p10 = position_embd + ((size_t)y1 * source_grid_size + x0) * embed_dim;
        const float *p11 = position_embd + ((size_t)y1 * source_grid_size + x1) * embed_dim;
        float *dst = x + (size_t)tok * embed_dim;
 
        for (int d = 0; d < embed_dim; ++d) {
            /* Materialize products before the reduction, as torch.sum does. */
            volatile float v00 = p00[d] * w00;
            volatile float v01 = p01[d] * w01;
            volatile float v10 = p10[d] * w10;
            volatile float v11 = p11[d] * w11;
            volatile float pos01 = v00 + v01;
            volatile float pos012 = pos01 + v10;
            const float pos = pos012 + v11;
            const float pos_bf16 = bf16_to_float(float_to_bf16(pos));
            const float hidden = bf16_to_float(float_to_bf16(dst[d]));
            dst[d] = bf16_to_float(float_to_bf16(hidden + pos_bf16));
        }
    }
}

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

◆ q_norm_forward()

void q_norm_forward	(	float *	q,
		const float *	q_gamma,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Forward pass for Gemma4-assistant q-only per-head RMSNorm.

Some Gemma4 assistant/drafter checkpoints project only Q and then reuse Q as the shared K/V stream. This wrapper keeps that public kernel contract explicit while reusing the same row-wise RMSNorm implementation as qk_norm_forward.

Definition at line 405 of file qk_norm_kernels.c.

{
    rmsnorm_forward(q, q_gamma, q, NULL,
                    num_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward().

◆ qk_norm_backward()

void qk_norm_backward	(	const float *	d_q_out,
		const float *	d_k_out,
		const float *	q_in,
		const float *	k_in,
		const float *	q_gamma,
		const float *	k_gamma,
		float *	d_q_in,
		float *	d_k_in,
		float *	d_q_gamma,
		float *	d_k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Backward pass for per-head QK RMSNorm.

This computes:

d_q / d_k for the Q and K activations
d_q_gamma / d_k_gamma for shared per-head gamma vectors

Implementation is reference-first and deterministic: 1) recompute row rstd values from saved q/k inputs 2) call rmsnorm_backward on flattened [rows, head_dim] views

Definition at line 427 of file qk_norm_kernels.c.

{
    int q_rows = num_heads * num_tokens;
    int k_rows = num_kv_heads * num_tokens;
 
    if (q_rows > 0) {
        float q_rstd_cache[q_rows];
        qk_norm_compute_rstd(q_in, q_rstd_cache, q_rows, head_dim, eps);
        rmsnorm_backward(d_q_out, q_in, q_gamma, q_rstd_cache,
                         d_q_in, d_q_gamma, q_rows, head_dim, head_dim);
    }
 
    if (k_rows > 0) {
        float k_rstd_cache[k_rows];
        qk_norm_compute_rstd(k_in, k_rstd_cache, k_rows, head_dim, eps);
        rmsnorm_backward(d_k_out, k_in, k_gamma, k_rstd_cache,
                         d_k_in, d_k_gamma, k_rows, head_dim, head_dim);
    }
}

References qk_norm_compute_rstd(), and rmsnorm_backward().

◆ qk_norm_backward_last_isa()

int qk_norm_backward_last_isa ( void )

Definition at line 111 of file qk_norm_kernels.c.

{
    return g_qk_norm_last_isa;
}

References g_qk_norm_last_isa.

◆ qk_norm_forward()

void qk_norm_forward	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Per-head RMSNorm on Q and K.

Parameters

q	Q scratch buffer [num_heads * num_tokens * head_dim], in-place
k	K scratch buffer [num_kv_heads * num_tokens * head_dim], in-place
q_gamma	Q norm gamma weights [head_dim]
k_gamma	K norm gamma weights [head_dim]
num_heads	Number of query heads (e.g. 32 for Qwen3-8B)
num_kv_heads	Number of KV heads (e.g. 8 for Qwen3-8B with GQA)
num_tokens	Number of tokens (1 for decode, T for prefill)
head_dim	Dimension per head (e.g. 128)
eps	RMSNorm epsilon (e.g. 1e-6)

Test:: unittest/test_qk_norm.py

Definition at line 326 of file qk_norm_kernels.c.

{
    /* Q norm: [num_heads * num_tokens] rows of [head_dim]
     * Each row is one head's vector for one token. */
    rmsnorm_forward(q, q_gamma, q, NULL,
                    num_heads * num_tokens, head_dim, head_dim, eps);
 
    /* K norm: [num_kv_heads * num_tokens] rows of [head_dim]
     * Same logic, fewer rows when using GQA. */
    rmsnorm_forward(k, k_gamma, k, NULL,
                    num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward().

◆ qk_norm_forward_decode_exact()

void qk_norm_forward_decode_exact	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

◆ qk_norm_forward_fp64_sum()

void qk_norm_forward_fp64_sum	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Definition at line 342 of file qk_norm_kernels.c.

{
    rmsnorm_forward_fp64_sum(q, q_gamma, q, NULL,
                             num_heads * num_tokens, head_dim, head_dim, eps);
    rmsnorm_forward_fp64_sum(k, k_gamma, k, NULL,
                             num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward_fp64_sum().

◆ qk_norm_forward_llama_production()

void qk_norm_forward_llama_production	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Definition at line 353 of file qk_norm_kernels.c.

{
    rmsnorm_forward_llama_production(
        q, q_gamma, q, NULL,
        num_heads * num_tokens, head_dim, head_dim, eps);
    rmsnorm_forward_llama_production(
        k, k_gamma, k, NULL,
        num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward_llama_production().

◆ qk_norm_forward_parallel_dispatch()

void qk_norm_forward_parallel_dispatch	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

◆ qk_norm_forward_prefill_exact()

void qk_norm_forward_prefill_exact	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

◆ qk_norm_forward_pytorch_bf16_storage()

void qk_norm_forward_pytorch_bf16_storage	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Definition at line 366 of file qk_norm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage(
        q, q_gamma, q, NULL,
        num_heads * num_tokens, head_dim, head_dim, eps);
    rmsnorm_forward_pytorch_bf16_storage(
        k, k_gamma, k, NULL,
        num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward_pytorch_bf16_storage().

◆ qk_norm_forward_qwen4_pytorch_bf16_storage()

void qk_norm_forward_qwen4_pytorch_bf16_storage	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Definition at line 381 of file qk_norm_kernels.c.

{
    rmsnorm_forward_qwen3next_pytorch_bf16_storage(
        q, q_gamma, q, NULL,
        num_heads * num_tokens, head_dim, head_dim, eps);
    rmsnorm_forward_qwen3next_pytorch_bf16_storage(
        k, k_gamma, k, NULL,
        num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward_qwen3next_pytorch_bf16_storage().

◆ quantize_batch_q8_0()

void quantize_batch_q8_0	(	const float *	x,
		void *	vy,
		int	num_rows,
		int	k
	)

Batch quantize FP32 to Q8_0 format (row-major output)

Quantizes multiple rows of FP32 data to Q8_0 format, placing each row's Q8_0 output at the correct byte offset for GEMM compatibility.

Memory layout: Input: [num_rows, k] FP32, row-major (stride = k * sizeof(float)) Output: [num_rows, q8_row_bytes] Q8_0, row-major (stride = q8_row_bytes)

where q8_row_bytes = (k / 32) * sizeof(block_q8_0) = (k / 32) * 34

Parameters

x	Input FP32 values [num_rows * k]
vy	Output Q8_0 blocks [num_rows * (k/32) blocks]
num_rows	Number of rows (batch size / tokens)
k	Elements per row (must be multiple of 32)

Definition at line 256 of file gemm_kernels_q8_0.c.

{
    const size_t row_bytes_in = (size_t)k * sizeof(float);
    const size_t row_bytes_out = (size_t)(k / QK8_0) * sizeof(block_q8_0);
 
    uint8_t *out = (uint8_t *)vy;
    const uint8_t *in = (const uint8_t *)x;
 
    for (int row = 0; row < num_rows; ++row) {
        quantize_row_q8_0(
            (const float *)(in + row * row_bytes_in),
            (void *)(out + row * row_bytes_out),
            k
        );
    }
}

References QK8_0, and quantize_row_q8_0().

◆ quantize_batch_q8_k()

void quantize_batch_q8_k	(	const float *	x,
		void *	vy,
		int	num_rows,
		int	k
	)

Batch quantize FP32 to Q8_K format (row-major output)

Same as quantize_batch_q8_0 but for Q8_K format (super-blocks).

Parameters

x	Input FP32 values [num_rows * k]
vy	Output Q8_K blocks
num_rows	Number of rows (batch size / tokens)
k	Elements per row (must be multiple of 256)

Definition at line 283 of file gemm_kernels_q8_0.c.

{
    /* Q8_K: 256 elements per super-block, each block is larger */
    const size_t row_bytes_in = (size_t)k * sizeof(float);
    /* Q8_K block size = 2 (d) + 256 (qs) + 32 (bsums/2) = ~274 bytes for 256 elements */
    /* Actual: sizeof(block_q8_K) from ckernel_quant.h */
    const size_t row_bytes_out = (size_t)(k / 256) * sizeof(block_q8_K);
 
    uint8_t *out = (uint8_t *)vy;
    const uint8_t *in = (const uint8_t *)x;
 
    for (int row = 0; row < num_rows; ++row) {
        quantize_row_q8_k(
            (const float *)(in + row * row_bytes_in),
            (void *)(out + row * row_bytes_out),
            k
        );
    }
}

References quantize_row_q8_k().

◆ quantize_batch_q8_k_4row_nearest_even()

void quantize_batch_q8_k_4row_nearest_even	(	const float *	x,
		void *	y,
		int	num_rows,
		int	k
	)

Definition at line 140 of file gemm_kernels_q4k_q8k.c.

                                                                 {
    if (!x || !vy || num_rows <= 0 || k <= 0) {
        return;
    }
    assert(k % QK_K == 0);
 
    block_q8_K *y = (block_q8_K *)vy;
    const int blocks_per_row = k / QK_K;
 
    /* Q8_K bytes are a numerical ABI. The previous four-row AVX2 path used
     * a different max/scale evaluation order and changed real Qwen3-VL
     * visual-prefix scales by one FP32 ULP. Keep this public grouped ABI on
     * the canonical row provider until an optimized implementation is
     * byte-exact against llama.cpp on both synthetic and production inputs. */
    for (int row = 0; row < num_rows; ++row) {
        quantize_row_q8_k(
                x + (size_t)row * (size_t)k,
                y + (size_t)row * (size_t)blocks_per_row,
                k);
    }
}

References QK_K, and quantize_row_q8_k().

◆ quantize_row_q8_0()

void quantize_row_q8_0	(	const float *	x,
		void *	vy,
		int	k
	)

Quantize FP32 to Q8_0 format (scalar reference)

Parameters

x	Input FP32 values
vy	Output Q8_0 blocks
k	Number of elements (must be multiple of 32)

Definition at line 125 of file gemm_kernels_q8_0.c.

{
    block_q8_0 *y = (block_q8_0 *)vy;
    const int nb = k / QK8_0;  /* QK8_0 = 32 */
 
#if defined(__AVX__)
    const __m256 sign_bit = _mm256_set1_ps(-0.0f);
 
    for (int i = 0; i < nb; i++) {
        __m256 v0 = _mm256_loadu_ps(x + 0);
        __m256 v1 = _mm256_loadu_ps(x + 8);
        __m256 v2 = _mm256_loadu_ps(x + 16);
        __m256 v3 = _mm256_loadu_ps(x + 24);
        x += QK8_0;
 
        __m256 max_abs = _mm256_andnot_ps(sign_bit, v0);
        max_abs = _mm256_max_ps(max_abs, _mm256_andnot_ps(sign_bit, v1));
        max_abs = _mm256_max_ps(max_abs, _mm256_andnot_ps(sign_bit, v2));
        max_abs = _mm256_max_ps(max_abs, _mm256_andnot_ps(sign_bit, v3));
 
        __m128 max4 = _mm_max_ps(_mm256_extractf128_ps(max_abs, 1),
                                 _mm256_castps256_ps128(max_abs));
        max4 = _mm_max_ps(max4, _mm_movehl_ps(max4, max4));
        max4 = _mm_max_ss(max4, _mm_movehdup_ps(max4));
        const float max_scalar = _mm_cvtss_f32(max4);
 
#if defined(__INTEL_LLVM_COMPILER)
        const float d = ck_q8_0_div_rounded_f32(max_scalar, 127.0f);
        const float id = max_scalar != 0.0f
            ? ck_q8_0_div_rounded_f32(127.0f, max_scalar)
            : 0.0f;
#else
        const float d = max_scalar / 127.0f;
        const float id = max_scalar != 0.0f ? 127.0f / max_scalar : 0.0f;
#endif
        y[i].d = CK_FP32_TO_FP16(d);
 
        const __m256 mul = _mm256_set1_ps(id);
        v0 = _mm256_mul_ps(v0, mul);
        v1 = _mm256_mul_ps(v1, mul);
        v2 = _mm256_mul_ps(v2, mul);
        v3 = _mm256_mul_ps(v3, mul);
 
        /* Match llama.cpp x86 Q8 quantization: nearest-even rounding. */
        v0 = _mm256_round_ps(v0, _MM_ROUND_NEAREST);
        v1 = _mm256_round_ps(v1, _MM_ROUND_NEAREST);
        v2 = _mm256_round_ps(v2, _MM_ROUND_NEAREST);
        v3 = _mm256_round_ps(v3, _MM_ROUND_NEAREST);
 
        __m256i i0 = _mm256_cvtps_epi32(v0);
        __m256i i1 = _mm256_cvtps_epi32(v1);
        __m256i i2 = _mm256_cvtps_epi32(v2);
        __m256i i3 = _mm256_cvtps_epi32(v3);
 
#if defined(__AVX2__)
        i0 = _mm256_packs_epi32(i0, i1);
        i2 = _mm256_packs_epi32(i2, i3);
        i0 = _mm256_packs_epi16(i0, i2);
 
        const __m256i perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7);
        i0 = _mm256_permutevar8x32_epi32(i0, perm);
        _mm256_storeu_si256((__m256i *)y[i].qs, i0);
#else
        __m128i ni0 = _mm256_castsi256_si128(i0);
        __m128i ni1 = _mm256_extractf128_si256(i0, 1);
        __m128i ni2 = _mm256_castsi256_si128(i1);
        __m128i ni3 = _mm256_extractf128_si256(i1, 1);
        __m128i ni4 = _mm256_castsi256_si128(i2);
        __m128i ni5 = _mm256_extractf128_si256(i2, 1);
        __m128i ni6 = _mm256_castsi256_si128(i3);
        __m128i ni7 = _mm256_extractf128_si256(i3, 1);
 
        ni0 = _mm_packs_epi32(ni0, ni1);
        ni2 = _mm_packs_epi32(ni2, ni3);
        ni4 = _mm_packs_epi32(ni4, ni5);
        ni6 = _mm_packs_epi32(ni6, ni7);
 
        ni0 = _mm_packs_epi16(ni0, ni2);
        ni4 = _mm_packs_epi16(ni4, ni6);
 
        _mm_storeu_si128((__m128i *)(y[i].qs + 0), ni0);
        _mm_storeu_si128((__m128i *)(y[i].qs + 16), ni4);
#endif
    }
#else
    for (int i = 0; i < nb; i++) {
        const float *xb = x + i * QK8_0;
 
        /* Find max absolute value in block */
        float amax = 0.0f;
        for (int j = 0; j < QK8_0; j++) {
            float av = xb[j] >= 0 ? xb[j] : -xb[j];
            if (av > amax) amax = av;
        }
 
        /* Compute scale: d = max / 127 */
        float d = amax / 127.0f;
        float id = d != 0.0f ? 127.0f / amax : 0.0f;
 
        /* Store scale as FP16 */
        y[i].d = CK_FP32_TO_FP16(d);
 
        /* Quantize values */
        for (int j = 0; j < QK8_0; j++) {
            float v = xb[j] * id;
            int q = ck_nearest_int_q8_0(v);
            if (q > 127) q = 127;
            if (q < -127) q = -127;
            y[i].qs[j] = (int8_t)q;
        }
    }
#endif
}

References CK_FP32_TO_FP16, ck_nearest_int_q8_0(), block_q8_0::d, id, QK8_0, and block_q8_0::qs.

Referenced by ck_moe_q4k_mixed_route_work(), ck_moe_shared_q4k_gated_workspace(), ck_moe_swiglu_nvfp4_projection(), fused_mlp_swiglu_prefill_w1w2_quant(), fused_rmsnorm_qkv_prefill_head_major_quant(), gemv_q8_0_q8_0_contract(), hyper_connection_mix_quantized(), moe_swiglu_expert_forward_q4k_q5_0_workspace(), moe_swiglu_expert_forward_q4k_q8_0_workspace(), moe_swiglu_shared_forward_q8_0_gated_workspace(), quantize_attn_out_head_major_q8_0(), quantize_attn_out_head_major_q8_0(), and quantize_attn_out_head_major_q8_0().

◆ quantize_row_q8_k()

void quantize_row_q8_k	(	const float *	x,
		void *	y,
		int	k
	)

Definition at line 121 of file gemm_kernels_q4k_q8k.c.

                                                        {
    const char *ref_env = getenv("CK_DEBUG_Q8K_REF");
    if (ref_env && atoi(ref_env) != 0) {
        quantize_row_q8_k_ref(x, vy, k);
        return;
    }
#if defined(__AVX512F__) && defined(__AVX512BW__)
    quantize_row_q8_k_avx512(x, vy, k);
#elif defined(__AVX2__)
    quantize_row_q8_k_avx2(x, vy, k);
#elif defined(__AVX__)
    quantize_row_q8_k_avx(x, vy, k);
#elif defined(__SSE4_1__)
    quantize_row_q8_k_sse(x, vy, k);
#else
    quantize_row_q8_k_ref(x, vy, k);
#endif
}

References quantize_row_q8_k_avx(), quantize_row_q8_k_avx2(), quantize_row_q8_k_avx512(), quantize_row_q8_k_ref(), and quantize_row_q8_k_sse().

◆ qwen4_ple_gate_conv_inject_bf16()

void qwen4_ple_gate_conv_inject_bf16	(	const float *	hyper_input,
		const float *	key_projected,
		const float *	value_projected,
		const float *	norm_key_weight,
		const float *	norm_query_weight,
		const float *	norm_conv_weight,
		const uint16_t *	conv_weight,
		float *	hyper_output,
		float *	key_norm_scratch,
		float *	query_norm_scratch,
		float *	gated_scratch,
		float *	conv_norm_scratch,
		const float *	conv_state_in,
		float *	conv_state_out,
		int	rows,
		int	streams,
		int	hidden_dim,
		int	kernel_size,
		int	dilation,
		float	eps
	)

Definition at line 402 of file qwen4_exp_kernels.c.

                                              {
    qwen4_ple_gate_conv_inject_impl(
        hyper_input, key_projected, value_projected, norm_key_weight,
        norm_query_weight, norm_conv_weight, conv_weight, hyper_output,
        key_norm_scratch, query_norm_scratch, gated_scratch, conv_norm_scratch,
        conv_state_in, conv_state_out, rows, streams, hidden_dim, kernel_size,
        dilation, eps, 0, 0);
}

References qwen4_ple_gate_conv_inject_impl().

◆ qwen4_ple_gate_conv_inject_fp16()

void qwen4_ple_gate_conv_inject_fp16	(	const float *	hyper_input,
		const float *	key_projected,
		const float *	value_projected,
		const float *	norm_key_weight,
		const float *	norm_query_weight,
		const float *	norm_conv_weight,
		const uint16_t *	conv_weight,
		float *	hyper_output,
		float *	key_norm_scratch,
		float *	query_norm_scratch,
		float *	gated_scratch,
		float *	conv_norm_scratch,
		const float *	conv_state_in,
		float *	conv_state_out,
		int	rows,
		int	streams,
		int	hidden_dim,
		int	kernel_size,
		int	dilation,
		float	eps
	)

Definition at line 419 of file qwen4_exp_kernels.c.

                                              {
    qwen4_ple_gate_conv_inject_impl(
        hyper_input, key_projected, value_projected, norm_key_weight,
        norm_query_weight, norm_conv_weight, conv_weight, hyper_output,
        key_norm_scratch, query_norm_scratch, gated_scratch, conv_norm_scratch,
        conv_state_in, conv_state_out, rows, streams, hidden_dim, kernel_size,
        dilation, eps, 1, 0);
}

References qwen4_ple_gate_conv_inject_impl().

◆ qwen4_ple_gate_conv_inject_llama_fp16()

void qwen4_ple_gate_conv_inject_llama_fp16	(	const float *	hyper_input,
		const float *	key_projected,
		const float *	value_projected,
		const float *	norm_key_weight,
		const float *	norm_query_weight,
		const float *	norm_conv_weight,
		const uint16_t *	conv_weight,
		float *	hyper_output,
		float *	key_norm_scratch,
		float *	query_norm_scratch,
		float *	gated_scratch,
		float *	conv_norm_scratch,
		const float *	conv_state_in,
		float *	conv_state_out,
		int	rows,
		int	streams,
		int	hidden_dim,
		int	kernel_size,
		int	dilation,
		float	eps
	)

Definition at line 436 of file qwen4_exp_kernels.c.

                                              {
    qwen4_ple_gate_conv_inject_impl(
        hyper_input, key_projected, value_projected, norm_key_weight,
        norm_query_weight, norm_conv_weight, conv_weight, hyper_output,
        key_norm_scratch, query_norm_scratch, gated_scratch, conv_norm_scratch,
        conv_state_in, conv_state_out, rows, streams, hidden_dim, kernel_size,
        dilation, eps, 1, 1);
}

References qwen4_ple_gate_conv_inject_impl().

◆ qwen4_ple_ngram_embed_bf16()

void qwen4_ple_ngram_embed_bf16	(	const int32_t *	token_ids,
		const uint16_t *	embedding,
		const int64_t *	layer_multipliers,
		const int64_t *	head_offsets,
		const int64_t *	head_vocab_sizes,
		float *	output,
		const float *	token_state_in,
		float *	token_state_out,
		int	rows,
		int	ngram_size,
		int	heads_per_ngram,
		int	head_dim,
		int	eos_token_id,
		int	position
	)

Definition at line 121 of file qwen4_exp_kernels.c.

                                                     {
    qwen4_ple_ngram_embed_impl(
        token_ids, embedding, layer_multipliers, head_offsets,
        head_vocab_sizes, output, token_state_in, token_state_out, rows,
        ngram_size, heads_per_ngram, head_dim, eos_token_id, position_offset, 0);
}

References qwen4_ple_ngram_embed_impl().

◆ qwen4_ple_ngram_embed_q5_0()

void qwen4_ple_ngram_embed_q5_0	(	const int32_t *	token_ids,
		const void *	embedding,
		const int64_t *	layer_multipliers,
		const int64_t *	head_offsets,
		const int64_t *	head_vocab_sizes,
		float *	output,
		const float *	token_state_in,
		float *	token_state_out,
		int	rows,
		int	ngram_size,
		int	heads_per_ngram,
		int	head_dim,
		int	eos_token_id,
		int	position
	)

Definition at line 141 of file qwen4_exp_kernels.c.

                                                     {
    qwen4_ple_ngram_embed_impl(
        token_ids, embedding, layer_multipliers, head_offsets,
        head_vocab_sizes, output, token_state_in, token_state_out, rows,
        ngram_size, heads_per_ngram, head_dim, eos_token_id, position_offset, 1);
}

References qwen4_ple_ngram_embed_impl().

◆ qwen4_qsa_index_select_bf16()

void qwen4_qsa_index_select_bf16	(	const float *	projected_qk,
		const float *	index_key_cache_in,
		const float *	q_norm_weight,
		const float *	k_norm_weight,
		float *	selected_indices,
		float *	index_key_cache_out,
		float *	q_norm_scratch,
		float *	pooled_key_scratch,
		float *	block_score_scratch,
		int32_t *	block_index_scratch,
		int	rows,
		int	query_heads,
		int	index_head_dim,
		int	token_budget,
		int	compress_ratio,
		int	rotary_dim,
		int	context_length,
		int	position,
		float	rope_theta,
		float	eps
	)

Definition at line 488 of file qwen4_exp_kernels.c.

               {
    if (!projected_qk || !index_key_cache_in || !q_norm_weight || !k_norm_weight ||
        !selected_indices || !index_key_cache_out || !q_norm_scratch ||
        !pooled_key_scratch || !block_score_scratch || !block_index_scratch ||
        rows <= 0 || query_heads <= 0 || index_head_dim <= 0 || token_budget <= 0 ||
        compress_ratio <= 0 || rotary_dim <= 0 || rotary_dim > index_head_dim ||
        context_length <= 0 || position < 0 || rows > context_length - position) return;
 
    const int projected_dim = (query_heads + 1) * index_head_dim;
    const int selection_width = token_budget + compress_ratio - 1;
    const int block_topk = token_budget / compress_ratio;
    if (index_key_cache_out != index_key_cache_in && position > 0) {
        memcpy(index_key_cache_out, index_key_cache_in,
               (size_t)position * index_head_dim * sizeof(*index_key_cache_out));
    }
 
    for (int row = 0; row < rows; ++row) {
        const int absolute_position = position + row;
        const float *projected = projected_qk + (size_t)row * projected_dim;
        float *query_normed = q_norm_scratch;
        qwen4_shared_head_rmsnorm(projected, q_norm_weight, query_normed,
                                  query_heads, index_head_dim, eps);
        for (int head = 0; head < query_heads; ++head) {
            qwen4_rope_split_inplace(query_normed + (size_t)head * index_head_dim,
                                     rotary_dim, absolute_position, rope_theta);
        }
        float *raw_key = index_key_cache_out + (size_t)absolute_position * index_head_dim;
        const float *projected_key = projected + (size_t)query_heads * index_head_dim;
        for (int col = 0; col < index_head_dim; ++col) {
            raw_key[col] = qwen4_bf16_round(projected_key[col]);
        }
 
        float *selected_row = selected_indices + (size_t)row * selection_width;
        for (int slot = 0; slot < selection_width; ++slot) selected_row[slot] = -1.0f;
        const int visible = absolute_position + 1;
        const int complete_blocks = visible / compress_ratio;
        const int selected_blocks = complete_blocks < block_topk ? complete_blocks : block_topk;
        for (int slot = 0; slot < selected_blocks; ++slot) {
            block_score_scratch[slot] = -INFINITY;
            block_index_scratch[slot] = -1;
        }
 
        for (int block = 0; block < complete_blocks; ++block) {
            const int block_start = block * compress_ratio;
            for (int col = 0; col < index_head_dim; ++col) {
                float sum = 0.0f;
                for (int token = 0; token < compress_ratio; ++token) {
                    sum += index_key_cache_out[
                        (size_t)(block_start + token) * index_head_dim + col];
                }
                pooled_key_scratch[col] = qwen4_bf16_round(sum / (float)compress_ratio);
            }
            qwen4_shared_head_rmsnorm(pooled_key_scratch, k_norm_weight,
                                      pooled_key_scratch, 1, index_head_dim, eps);
            qwen4_rope_split_inplace(pooled_key_scratch, rotary_dim, block_start, rope_theta);
            float score = 0.0f;
            for (int head = 0; head < query_heads; ++head) {
                const float *query_head = query_normed + (size_t)head * index_head_dim;
                float dot = 0.0f;
                for (int col = 0; col < index_head_dim; ++col) {
                    dot += query_head[col] * pooled_key_scratch[col];
                }
                if (dot > 0.0f) score += dot;
            }
            score /= sqrtf((float)index_head_dim);
            int target = -1;
            for (int slot = 0; slot < selected_blocks; ++slot) {
                if (target < 0 || block_score_scratch[slot] < block_score_scratch[target]) target = slot;
            }
            if (target >= 0 && score > block_score_scratch[target]) {
                block_score_scratch[target] = score;
                block_index_scratch[target] = block;
            }
        }
 
        for (int outer = 1; outer < selected_blocks; ++outer) {
            const int32_t value = block_index_scratch[outer];
            int inner = outer - 1;
            while (inner >= 0 && block_index_scratch[inner] > value) {
                block_index_scratch[inner + 1] = block_index_scratch[inner];
                --inner;
            }
            block_index_scratch[inner + 1] = value;
        }
        int output_count = 0;
        for (int slot = 0; slot < selected_blocks; ++slot) {
            const int block = block_index_scratch[slot];
            if (block < 0) continue;
            for (int token = 0; token < compress_ratio; ++token) {
                selected_row[output_count++] = (float)(block * compress_ratio + token);
            }
        }
        for (int token = complete_blocks * compress_ratio;
             token < visible && output_count < selection_width; ++token) {
            selected_row[output_count++] = (float)token;
        }
    }
}

References qwen4_bf16_round(), qwen4_rope_split_inplace(), qwen4_shared_head_rmsnorm(), score, and token.

◆ recurrent_conv_state_update_backward()

void recurrent_conv_state_update_backward	(	const float *	d_conv_x,
		const float *	d_state_out,
		float *	d_state_in,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		int	history_len,
		int	num_seqs,
		int	num_tokens,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 160 of file recurrent_state_kernels.c.

                                                     {
    int channels = 0;
    int total_len = 0;
    size_t elements = 0;
    if (!recurrent_conv_backward_extents(
            history_len, num_seqs, num_tokens, q_dim, k_dim, v_dim,
            &channels, &total_len, &elements) ||
        elements > SIZE_MAX / sizeof(float)) {
        return;
    }
    (void)channels;
    (void)total_len;
    float *workspace = (float *)malloc(elements * sizeof(float));
    if (!workspace) {
        return;
    }
    recurrent_conv_state_update_backward_workspace(
        d_conv_x, d_state_out, d_state_in, d_q, d_k, d_v, workspace,
        history_len, num_seqs, num_tokens, q_dim, k_dim, v_dim);
    free(workspace);
}

References recurrent_conv_backward_extents(), and recurrent_conv_state_update_backward_workspace().

◆ recurrent_conv_state_update_backward_workspace()

void recurrent_conv_state_update_backward_workspace	(	const float *	d_conv_x,
		const float *	d_state_out,
		float *	d_state_in,
		float *	d_q,
		float *	d_k,
		float *	d_v,
		float *	d_conv_total,
		int	history_len,
		int	num_seqs,
		int	num_tokens,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 91 of file recurrent_state_kernels.c.

                                                                {
    int channels = 0;
    int total_len = 0;
    size_t elements = 0;
    if (!d_conv_x || !d_state_out || !d_state_in || !d_q || !d_k || !d_v ||
        !d_conv_total || !recurrent_conv_backward_extents(
            history_len, num_seqs, num_tokens, q_dim, k_dim, v_dim,
            &channels, &total_len, &elements)) {
        return;
    }
    if (elements > SIZE_MAX / sizeof(float)) {
        return;
    }
 
    memcpy(d_conv_total, d_conv_x, elements * sizeof(float));
 
    for (int seq = 0; seq < num_seqs; ++seq) {
        const float *d_state_out_seq = d_state_out + (size_t) seq * (size_t) channels * (size_t) history_len;
        float *d_conv_seq = d_conv_total + (size_t) seq * (size_t) channels * (size_t) total_len;
        for (int ch = 0; ch < channels; ++ch) {
            float *dst = d_conv_seq + (size_t) ch * (size_t) total_len + (size_t) num_tokens;
            const float *src = d_state_out_seq + (size_t) ch * (size_t) history_len;
            for (int idx = 0; idx < history_len; ++idx) {
                dst[idx] += src[idx];
            }
        }
    }
 
    for (int seq = 0; seq < num_seqs; ++seq) {
        const float *d_conv_seq = d_conv_total + (size_t) seq * (size_t) channels * (size_t) total_len;
        float *d_state_in_seq = d_state_in + (size_t) seq * (size_t) channels * (size_t) history_len;
 
        for (int ch = 0; ch < channels; ++ch) {
            memcpy(
                d_state_in_seq + (size_t) ch * (size_t) history_len,
                d_conv_seq + (size_t) ch * (size_t) total_len,
                (size_t) history_len * sizeof(float));
        }
 
        for (int tok = 0; tok < num_tokens; ++tok) {
            const int row = seq * num_tokens + tok;
            float *d_q_row = d_q + (size_t) row * (size_t) q_dim;
            float *d_k_row = d_k + (size_t) row * (size_t) k_dim;
            float *d_v_row = d_v + (size_t) row * (size_t) v_dim;
            for (int col = 0; col < q_dim; ++col) {
                d_q_row[col] = d_conv_seq[(size_t) col * (size_t) total_len + (size_t) (history_len + tok)];
            }
            for (int col = 0; col < k_dim; ++col) {
                d_k_row[col] = d_conv_seq[(size_t) (q_dim + col) * (size_t) total_len + (size_t) (history_len + tok)];
            }
            for (int col = 0; col < v_dim; ++col) {
                d_v_row[col] = d_conv_seq[(size_t) (q_dim + k_dim + col) * (size_t) total_len + (size_t) (history_len + tok)];
            }
        }
    }
}

References recurrent_conv_backward_extents().

Referenced by recurrent_conv_state_update_backward().

◆ recurrent_conv_state_update_forward()

void recurrent_conv_state_update_forward	(	const float *	state_in,
		const float *	q,
		const float *	k,
		const float *	v,
		float *	conv_x,
		float *	state_out,
		int	history_len,
		int	num_seqs,
		int	num_tokens,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 8 of file recurrent_state_kernels.c.

                                                    {
    const int channels = q_dim + k_dim + v_dim;
    const int total_len = history_len + num_tokens;
    for (int seq = 0; seq < num_seqs; ++seq) {
        const float *state_seq = state_in + (size_t) seq * (size_t) channels * (size_t) history_len;
        float *conv_seq = conv_x + (size_t) seq * (size_t) channels * (size_t) total_len;
        float *state_out_seq = state_out + (size_t) seq * (size_t) channels * (size_t) history_len;
        for (int ch = 0; ch < channels; ++ch) {
            memcpy(
                conv_seq + (size_t) ch * (size_t) total_len,
                state_seq + (size_t) ch * (size_t) history_len,
                (size_t) history_len * sizeof(float));
        }
 
        for (int tok = 0; tok < num_tokens; ++tok) {
            const int row = seq * num_tokens + tok;
            const float *q_row = q + (size_t) row * (size_t) q_dim;
            const float *k_row = k + (size_t) row * (size_t) k_dim;
            const float *v_row = v + (size_t) row * (size_t) v_dim;
            for (int col = 0; col < q_dim; ++col) {
                conv_seq[(size_t) col * (size_t) total_len + (size_t) (history_len + tok)] = q_row[col];
            }
            for (int col = 0; col < k_dim; ++col) {
                conv_seq[(size_t) (q_dim + col) * (size_t) total_len + (size_t) (history_len + tok)] = k_row[col];
            }
            for (int col = 0; col < v_dim; ++col) {
                conv_seq[(size_t) (q_dim + k_dim + col) * (size_t) total_len + (size_t) (history_len + tok)] = v_row[col];
            }
        }
 
        for (int ch = 0; ch < channels; ++ch) {
            memcpy(
                state_out_seq + (size_t) ch * (size_t) history_len,
                conv_seq + (size_t) ch * (size_t) total_len + (size_t) num_tokens,
                (size_t) history_len * sizeof(float));
        }
    }
}

◆ recurrent_dt_gate_backward()

void recurrent_dt_gate_backward	(	const float *	d_gate,
		const float *	alpha,
		const float *	dt_bias,
		const float *	a,
		float *	d_alpha,
		float *	d_dt_bias,
		float *	d_a,
		int	rows,
		int	dim
	)

Definition at line 104 of file recurrent_gate_kernels.c.

                                         {
    for (int col = 0; col < dim; ++col) {
        d_dt_bias[col] = 0.0f;
        d_a[col] = 0.0f;
    }
 
    for (int row = 0; row < rows; ++row) {
        const float *d_gate_row = d_gate + (size_t) row * (size_t) dim;
        const float *alpha_row = alpha + (size_t) row * (size_t) dim;
        float *d_alpha_row = d_alpha + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float x = alpha_row[col] + dt_bias[col];
            const float sp = recurrent_softplus(x);
            const float sig = recurrent_sigmoid(x);
            const float d_out = d_gate_row[col];
            d_a[col] += d_out * sp;
            {
                const float d_x = d_out * a[col] * sig;
                d_alpha_row[col] = d_x;
                d_dt_bias[col] += d_x;
            }
        }
    }
}

References recurrent_sigmoid(), and recurrent_softplus().

◆ recurrent_dt_gate_expanded_forward()

void recurrent_dt_gate_expanded_forward	(	const float *	alpha,
		const float *	dt_bias,
		const float *	a,
		float *	gate,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 83 of file recurrent_gate_kernels.c.

                                                       {
    for (int row = 0; row < rows; ++row) {
        const float *alpha_row = alpha + (size_t) row * (size_t) num_heads;
        float *gate_row = gate + (size_t) row * (size_t) num_heads * (size_t) state_dim;
        for (int h = 0; h < num_heads; ++h) {
            const float sp = recurrent_softplus(alpha_row[h] + dt_bias[h]);
            const float *a_head = a + (size_t) h * (size_t) state_dim;
            float *gate_head = gate_row + (size_t) h * (size_t) state_dim;
            for (int col = 0; col < state_dim; ++col) {
                gate_head[col] = sp * a_head[col];
            }
        }
    }
}

References recurrent_softplus().

◆ recurrent_dt_gate_forward()

void recurrent_dt_gate_forward	(	const float *	alpha,
		const float *	dt_bias,
		const float *	a,
		float *	gate,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 65 of file recurrent_gate_kernels.c.

                                              {
    const int dim = num_heads * state_dim;
    for (int row = 0; row < rows; ++row) {
        const float *alpha_row = alpha + (size_t) row * (size_t) dim;
        float *gate_row = gate + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float x = alpha_row[col] + dt_bias[col];
            gate_row[col] = recurrent_softplus(x) * a[col];
        }
    }
}

References recurrent_softplus().

◆ recurrent_dt_gate_forward_pytorch_fp32()

void recurrent_dt_gate_forward_pytorch_fp32	(	const float *	alpha,
		const float *	dt_bias,
		const float *	a,
		float *	gate,
		int	rows,
		int	num_heads,
		int	state_dim
	)

Definition at line 179 of file recurrent_gate_kernels.c.

{
    if (!alpha || !dt_bias || !a || !gate || rows < 0 || num_heads < 0 || state_dim != 1) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: invalid PyTorch FP32 recurrent dt-gate arguments\n");
        abort();
    }
#if defined(__AVX512F__)
    pthread_once(&ck_recurrent_sleef_once, ck_bind_recurrent_pytorch_sleef);
    if (!ck_recurrent_pytorch_expf16 || !ck_recurrent_pytorch_log1pf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch FP32 recurrent dt gate requires "
                "SLEEF Sleef_expf16_u10 and Sleef_log1pf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
    if ((num_heads & 15) != 0) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch AVX-512 recurrent dt gate requires "
                "a head count divisible by 16 (got %d)\n",
                num_heads);
        abort();
    }
    pthread_once(&ck_recurrent_libm_once, ck_bind_recurrent_llama_libm);
    if (!ck_recurrent_pytorch_log1pf) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch FP32 recurrent dt gate requires "
                "log1pf from libm.so.6\n");
        abort();
    }
    const __m512 threshold = _mm512_set1_ps(20.0f);
    const int count = rows * num_heads;
    int index = 0;
    for (; index + 32 <= count; index += 32) {
        for (int half = 0; half < 2; ++half) {
            float x_lanes[16] __attribute__((aligned(64)));
            float a_lanes[16] __attribute__((aligned(64)));
            const int base = index + half * 16;
            for (int lane = 0; lane < 16; ++lane) {
                const int head = (base + lane) % num_heads;
                x_lanes[lane] = alpha[base + lane] + dt_bias[head];
                a_lanes[lane] = a[head];
            }
            const __m512 x = _mm512_load_ps(x_lanes);
            const __m512 softplus = _mm512_mask_blend_ps(
                _mm512_cmp_ps_mask(x, threshold, _CMP_GT_OQ),
                ck_recurrent_pytorch_log1pf16(ck_recurrent_pytorch_expf16(x)),
                x);
            _mm512_storeu_ps(
                gate + base,
                _mm512_mul_ps(softplus, _mm512_load_ps(a_lanes)));
        }
    }
    for (; index < count; ++index) {
        const int head = index % num_heads;
        const float x = alpha[index] + dt_bias[head];
        const float softplus = x > 20.0f
            ? x
            : ck_recurrent_pytorch_log1pf(ck_recurrent_llama_expf(x));
        gate[index] = softplus * a[head];
    }
#else
    fprintf(stderr,
            "HARD KERNEL CONTRACT FAULT: PyTorch FP32 recurrent dt gate requires AVX-512\n");
    abort();
#endif
}

References __attribute__(), ck_bind_recurrent_llama_libm(), ck_recurrent_libm_once, ck_recurrent_llama_expf, and ck_recurrent_pytorch_log1pf.

◆ recurrent_norm_gate_backward()

void recurrent_norm_gate_backward	(	const float *	d_out,
		const float *	x,
		const float *	gate,
		const float *	weight,
		float *	d_x,
		float *	d_gate,
		float *	d_weight,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 182 of file recurrent_norm_kernels.c.

                                             {
    const int inner_dim = num_heads * head_dim;
    memset(d_weight, 0, (size_t) head_dim * sizeof(float));
 
    for (int row = 0; row < rows; ++row) {
        const float *d_out_row = d_out + (size_t) row * (size_t) inner_dim;
        const float *x_row = x + (size_t) row * (size_t) inner_dim;
        const float *gate_row = gate + (size_t) row * (size_t) inner_dim;
        float *d_x_row = d_x + (size_t) row * (size_t) inner_dim;
        float *d_gate_row = d_gate + (size_t) row * (size_t) inner_dim;
 
        for (int head = 0; head < num_heads; ++head) {
            const float *x_head = x_row + (size_t) head * (size_t) head_dim;
            const float *gate_head = gate_row + (size_t) head * (size_t) head_dim;
            const float *d_out_head = d_out_row + (size_t) head * (size_t) head_dim;
            float *d_x_head = d_x_row + (size_t) head * (size_t) head_dim;
            float *d_gate_head = d_gate_row + (size_t) head * (size_t) head_dim;
 
            float ms = 0.0f;
            for (int col = 0; col < head_dim; ++col) {
                ms += x_head[col] * x_head[col];
            }
            ms /= (float) head_dim;
            const float inv_rms = 1.0f / sqrtf(ms + eps);
            const float inv_rms3_over_dim = (inv_rms * inv_rms * inv_rms) / (float) head_dim;
 
            float dot = 0.0f;
            for (int col = 0; col < head_dim; ++col) {
                const float g = gate_head[col];
                const float sig = recurrent_sigmoid_local(g);
                const float silu = g * sig;
                dot += d_out_head[col] * weight[col] * silu * x_head[col];
            }
 
            for (int col = 0; col < head_dim; ++col) {
                const float g = gate_head[col];
                const float sig = recurrent_sigmoid_local(g);
                const float silu = g * sig;
                const float scaled = inv_rms * weight[col] * silu;
                d_x_head[col] = d_out_head[col] * scaled - x_head[col] * inv_rms3_over_dim * dot;
                d_weight[col] += d_out_head[col] * (x_head[col] * inv_rms) * silu;
                d_gate_head[col] = d_out_head[col] * (x_head[col] * inv_rms * weight[col]) *
                                   (sig + g * sig * (1.0f - sig));
            }
        }
    }
}

References recurrent_sigmoid_local(), and silu().

◆ recurrent_norm_gate_forward()

void recurrent_norm_gate_forward	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 18 of file recurrent_norm_kernels.c.

                                            {
    const int inner_dim = num_heads * head_dim;
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) inner_dim;
        const float *gate_row = gate + (size_t) row * (size_t) inner_dim;
        float *out_row = out + (size_t) row * (size_t) inner_dim;
 
        for (int head = 0; head < num_heads; ++head) {
            const float *x_head = x_row + (size_t) head * (size_t) head_dim;
            const float *gate_head = gate_row + (size_t) head * (size_t) head_dim;
            float *out_head = out_row + (size_t) head * (size_t) head_dim;
 
            float ms = 0.0f;
            for (int col = 0; col < head_dim; ++col) {
                ms += x_head[col] * x_head[col];
            }
            ms /= (float) head_dim;
            const float inv_rms = 1.0f / sqrtf(ms + eps);
 
            for (int col = 0; col < head_dim; ++col) {
                const float g = gate_head[col];
                const float silu = g * recurrent_sigmoid_local(g);
                out_head[col] = x_head[col] * inv_rms * weight[col] * silu;
            }
        }
    }
}

References recurrent_sigmoid_local(), and silu().

◆ recurrent_norm_gate_llama_avx2_forward()

void recurrent_norm_gate_llama_avx2_forward	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 53 of file recurrent_norm_kernels.c.

                                                        {
    if (!x || !gate || !weight || !out || rows <= 0 || num_heads <= 0 ||
        head_dim <= 0 || head_dim > 4096) {
        return;
    }
    const int inner_dim = num_heads * head_dim;
    float normalized[4096];
    float silu[4096];
    for (int row = 0; row < rows; ++row) {
        for (int head = 0; head < num_heads; ++head) {
            const size_t offset = (size_t) row * (size_t) inner_dim
                                + (size_t) head * (size_t) head_dim;
            rmsnorm_forward_llama_production(
                x + offset, weight, normalized, NULL, 1, head_dim, head_dim, eps);
            recurrent_silu_forward_ggml(gate + offset, silu, 1, head_dim);
            for (int col = 0; col < head_dim; ++col) {
                out[offset + (size_t) col] = normalized[col] * silu[col];
            }
        }
    }
}

References recurrent_silu_forward_ggml(), rmsnorm_forward_llama_production(), and silu().

◆ recurrent_norm_gate_pytorch_bf16_storage()

void recurrent_norm_gate_pytorch_bf16_storage	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 111 of file recurrent_norm_kernels.c.

{
    if (!x || !gate || !weight || !out || rows <= 0 || num_heads <= 0 ||
        head_dim <= 0 || head_dim > 4096) {
        return;
    }
    const int inner_dim = num_heads * head_dim;
    float normalized[4096];
    float silu[4096];
    for (int row = 0; row < rows; ++row) {
        for (int head = 0; head < num_heads; ++head) {
            const size_t offset = (size_t)row * (size_t)inner_dim +
                (size_t)head * (size_t)head_dim;
            /*
             * Qwen3Next order:
             * FP32 RMS statistics -> BF16 normalized value -> BF16 weight
             * multiply -> FP32 SiLU(gate) -> BF16 final output.
             */
            rmsnorm_forward_pytorch_bf16_storage(
                x + offset, weight, normalized, NULL,
                1, head_dim, head_dim, eps);
            recurrent_silu_forward_pytorch_bf16_input_fp32_output(
                gate + offset, silu, 1, head_dim);
            for (int col = 0; col < head_dim; ++col) {
                out[offset + (size_t)col] = bf16_to_float(float_to_bf16(
                    normalized[col] * silu[col]));
            }
        }
    }
}

References bf16_to_float(), float_to_bf16(), recurrent_silu_forward_pytorch_bf16_input_fp32_output(), rmsnorm_forward_pytorch_bf16_storage(), and silu().

◆ recurrent_norm_sigmoid_gate_llama_avx2_forward()

void recurrent_norm_sigmoid_gate_llama_avx2_forward	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 82 of file recurrent_norm_kernels.c.

                                                                {
    if (!x || !gate || !weight || !out || rows <= 0 || num_heads <= 0 ||
        head_dim <= 0 || head_dim > 4096) {
        return;
    }
    const int inner_dim = num_heads * head_dim;
    float normalized[4096];
    float sigmoid[4096];
    for (int row = 0; row < rows; ++row) {
        for (int head = 0; head < num_heads; ++head) {
            const size_t offset = (size_t) row * (size_t) inner_dim
                                + (size_t) head * (size_t) head_dim;
            rmsnorm_forward_llama_production(
                x + offset, weight, normalized, NULL, 1, head_dim, head_dim, eps);
            recurrent_sigmoid_forward_ggml(gate + offset, sigmoid, 1, head_dim);
            for (int col = 0; col < head_dim; ++col) {
                out[offset + (size_t) col] = normalized[col] * sigmoid[col];
            }
        }
    }
}

References recurrent_sigmoid_forward_ggml(), and rmsnorm_forward_llama_production().

◆ recurrent_norm_sigmoid_gate_pytorch_bf16_storage()

void recurrent_norm_sigmoid_gate_pytorch_bf16_storage	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 149 of file recurrent_norm_kernels.c.

{
    if (!x || !gate || !weight || !out || rows <= 0 || num_heads <= 0 ||
        head_dim <= 0 || head_dim > 4096) {
        return;
    }
    const int inner_dim = num_heads * head_dim;
    float normalized[4096];
    float sigmoid[4096];
    for (int row = 0; row < rows; ++row) {
        for (int head = 0; head < num_heads; ++head) {
            const size_t offset = (size_t)row * (size_t)inner_dim +
                (size_t)head * (size_t)head_dim;
            rmsnorm_forward_pytorch_bf16_storage(
                x + offset, weight, normalized, NULL,
                1, head_dim, head_dim, eps);
            recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output(
                gate + offset, sigmoid, 1, head_dim);
            for (int col = 0; col < head_dim; ++col) {
                out[offset + (size_t)col] = bf16_to_float(float_to_bf16(
                    normalized[col] * sigmoid[col]));
            }
        }
    }
}

References bf16_to_float(), float_to_bf16(), recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output(), and rmsnorm_forward_pytorch_bf16_storage().

◆ recurrent_qk_l2_norm_backward()

void recurrent_qk_l2_norm_backward	(	const float *	d_q_out,
		const float *	d_k_out,
		const float *	q,
		const float *	k,
		float *	d_q,
		float *	d_k,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	head_dim,
		float	eps
	)

Definition at line 383 of file recurrent_qk_norm_kernels.c.

                                              {
    recurrent_l2_norm_rows_backward_one(d_q_out, q, d_q, rows, q_dim, head_dim, eps);
    recurrent_l2_norm_rows_backward_one(d_k_out, k, d_k, rows, k_dim, head_dim, eps);
}

References recurrent_l2_norm_rows_backward_one().

◆ recurrent_qk_l2_norm_forward()

void recurrent_qk_l2_norm_forward	(	float *	q,
		float *	k,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	head_dim,
		float	eps
	)

Definition at line 86 of file recurrent_qk_norm_kernels.c.

                                             {
    recurrent_l2_norm_rows_forward_one(q, rows, q_dim, head_dim, eps);
    recurrent_l2_norm_rows_forward_one(k, rows, k_dim, head_dim, eps);
}

References recurrent_l2_norm_rows_forward_one().

◆ recurrent_qk_l2_norm_pytorch_bf16_storage()

void recurrent_qk_l2_norm_pytorch_bf16_storage	(	float *	q,
		float *	k,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	expanded_heads,
		int	head_dim,
		float	eps
	)

Definition at line 368 of file recurrent_qk_norm_kernels.c.

{
    recurrent_pytorch_bf16_l2_rows(
        q, rows, q_dim, expanded_heads, head_dim, eps);
    recurrent_pytorch_bf16_l2_rows(
        k, rows, k_dim, expanded_heads, head_dim, eps);
}

References recurrent_pytorch_bf16_l2_rows().

◆ recurrent_qk_l2_norm_pytorch_fp32_output()

void recurrent_qk_l2_norm_pytorch_fp32_output	(	float *	q,
		float *	k,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	head_dim,
		float	eps
	)

Definition at line 246 of file recurrent_qk_norm_kernels.c.

{
    recurrent_pytorch_fp32_l2_rows(q, rows, q_dim, head_dim, eps);
    recurrent_pytorch_fp32_l2_rows(k, rows, k_dim, head_dim, eps);
}

References recurrent_pytorch_fp32_l2_rows().

◆ recurrent_sigmoid_forward_ggml()

void recurrent_sigmoid_forward_ggml	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 501 of file recurrent_gate_kernels.c.

                                             {
    float (*volatile llama_expf)(float) = expf;
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) dim;
        float *out_row = out + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            out_row[col] = 1.0f / (1.0f + llama_expf(-x_row[col]));
        }
    }
}

Referenced by hyper_connection_mix_quantized(), and recurrent_norm_sigmoid_gate_llama_avx2_forward().

◆ recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output()

void recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 339 of file recurrent_gate_kernels.c.

{
    if (!x || !out || rows < 0 || dim < 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid BF16-input FP32-output sigmoid arguments\n");
        abort();
    }
#if defined(__AVX512F__)
    pthread_once(&ck_recurrent_sleef_once, ck_bind_recurrent_pytorch_sleef);
    if (!ck_recurrent_pytorch_expf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch BF16-input sigmoid requires "
                "SLEEF Sleef_expf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
#endif
    const int count = rows * dim;
    int i = 0;
#if defined(__AVX512F__)
    for (; i + 16 <= count; i += 16) {
        float lanes[16] __attribute__((aligned(64)));
        for (int lane = 0; lane < 16; ++lane) {
            lanes[lane] = bf16_to_float(float_to_bf16(x[i + lane]));
        }
        const __m512 values = _mm512_load_ps(lanes);
        const __m512 denominator = _mm512_add_ps(
            _mm512_set1_ps(1.0f),
            ck_recurrent_pytorch_expf16(
                _mm512_sub_ps(_mm512_setzero_ps(), values)));
        _mm512_storeu_ps(out + i,
                         _mm512_div_ps(_mm512_set1_ps(1.0f), denominator));
    }
#endif
    for (; i < count; ++i) {
        const float value = bf16_to_float(float_to_bf16(x[i]));
        out[i] = 1.0f / (1.0f + expf(-value));
    }
}

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

Referenced by attn_gate_sigmoid_mul_pytorch_bf16_storage(), and recurrent_norm_sigmoid_gate_pytorch_bf16_storage().

◆ recurrent_silu_backward()

void recurrent_silu_backward	(	const float *	d_out,
		const float *	x,
		float *	d_x,
		int	rows,
		int	dim
	)

Definition at line 515 of file recurrent_gate_kernels.c.

                                      {
    for (int row = 0; row < rows; ++row) {
        const float *d_out_row = d_out + (size_t) row * (size_t) dim;
        const float *x_row = x + (size_t) row * (size_t) dim;
        float *d_x_row = d_x + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float xv = x_row[col];
            const float sig = recurrent_sigmoid(xv);
            d_x_row[col] = d_out_row[col] * (sig + xv * sig * (1.0f - sig));
        }
    }
}

References recurrent_sigmoid().

◆ recurrent_silu_forward()

void recurrent_silu_forward	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 137 of file recurrent_gate_kernels.c.

                                     {
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) dim;
        float *out_row = out + (size_t) row * (size_t) dim;
        for (int col = 0; col < dim; ++col) {
            const float xv = x_row[col];
            out_row[col] = xv * recurrent_sigmoid(xv);
        }
    }
}

References recurrent_sigmoid().

◆ recurrent_silu_forward_ggml()

void recurrent_silu_forward_ggml	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 468 of file recurrent_gate_kernels.c.

                                          {
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) dim;
        float *out_row = out + (size_t) row * (size_t) dim;
        int col = 0;
#if defined(__AVX512F__) && defined(__AVX512DQ__)
        for (; col + 16 <= dim; col += 16) {
            const __m512 xv = _mm512_loadu_ps(x_row + col);
            const __m512 neg = _mm512_sub_ps(_mm512_setzero_ps(), xv);
            const __m512 denom = _mm512_add_ps(
                _mm512_set1_ps(1.0f), recurrent_ggml_expf_avx512(neg));
            _mm512_storeu_ps(out_row + col, _mm512_div_ps(xv, denom));
        }
#elif defined(__AVX2__) && defined(__FMA__)
        for (; col + 8 <= dim; col += 8) {
            const __m256 xv = _mm256_loadu_ps(x_row + col);
            const __m256 neg = _mm256_sub_ps(_mm256_setzero_ps(), xv);
            const __m256 denom = _mm256_add_ps(
                _mm256_set1_ps(1.0f), recurrent_ggml_expf_avx2(neg));
            _mm256_storeu_ps(out_row + col, _mm256_div_ps(xv, denom));
        }
#endif
        float (*volatile llama_expf)(float) = expf;
        for (; col < dim; ++col) {
            const float xv = x_row[col];
            out_row[col] = xv / (1.0f + llama_expf(-xv));
        }
    }
}

Referenced by hyper_connection_mix_quantized(), qwen4_ple_gate_conv_inject_impl(), and recurrent_norm_gate_llama_avx2_forward().

◆ recurrent_silu_forward_pytorch_bf16_input_fp32_output()

void recurrent_silu_forward_pytorch_bf16_input_fp32_output	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 299 of file recurrent_gate_kernels.c.

{
    if (!x || !out || rows < 0 || dim < 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid BF16-input FP32-output SiLU arguments\n");
        abort();
    }
#if defined(__AVX512F__)
    pthread_once(&ck_recurrent_sleef_once, ck_bind_recurrent_pytorch_sleef);
    if (!ck_recurrent_pytorch_expf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch BF16-input SiLU requires "
                "SLEEF Sleef_expf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
#endif
    const int count = rows * dim;
    int i = 0;
#if defined(__AVX512F__)
    for (; i + 16 <= count; i += 16) {
        float lanes[16] __attribute__((aligned(64)));
        for (int lane = 0; lane < 16; ++lane) {
            lanes[lane] = bf16_to_float(float_to_bf16(x[i + lane]));
        }
        const __m512 values = _mm512_load_ps(lanes);
        const __m512 denominator = _mm512_add_ps(
            _mm512_set1_ps(1.0f),
            ck_recurrent_pytorch_expf16(
                _mm512_sub_ps(_mm512_setzero_ps(), values)));
        _mm512_storeu_ps(out + i, _mm512_div_ps(values, denominator));
    }
#endif
    for (; i < count; ++i) {
        const float value = bf16_to_float(float_to_bf16(x[i]));
        out[i] = value / (1.0f + expf(-value));
    }
}

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

Referenced by recurrent_norm_gate_pytorch_bf16_storage().

◆ recurrent_silu_forward_pytorch_bf16_storage()

void recurrent_silu_forward_pytorch_bf16_storage	(	const float *	x,
		float *	out,
		int	rows,
		int	dim
	)

Definition at line 252 of file recurrent_gate_kernels.c.

{
    if (!x || !out || rows < 0 || dim < 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid PyTorch BF16 recurrent SiLU arguments\n");
        abort();
    }
#if defined(__AVX512F__)
    pthread_once(&ck_recurrent_sleef_once, ck_bind_recurrent_pytorch_sleef);
    if (!ck_recurrent_pytorch_expf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch BF16 recurrent SiLU requires "
                "SLEEF Sleef_expf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
#endif
    for (int row = 0; row < rows; ++row) {
        const float *src = x + (size_t)row * (size_t)dim;
        float *dst = out + (size_t)row * (size_t)dim;
        int col = 0;
#if defined(__AVX512F__)
        for (; col + 16 <= dim; col += 16) {
            float lanes[16] __attribute__((aligned(64)));
            for (int lane = 0; lane < 16; ++lane) {
                lanes[lane] = bf16_to_float(float_to_bf16(src[col + lane]));
            }
            const __m512 values = _mm512_load_ps(lanes);
            const __m512 denominator = _mm512_add_ps(
                _mm512_set1_ps(1.0f),
                ck_recurrent_pytorch_expf16(
                    _mm512_sub_ps(_mm512_setzero_ps(), values)));
            _mm512_store_ps(lanes, _mm512_div_ps(values, denominator));
            for (int lane = 0; lane < 16; ++lane) {
                dst[col + lane] = bf16_to_float(float_to_bf16(lanes[lane]));
            }
        }
#endif
        for (; col < dim; ++col) {
            const float value = bf16_to_float(float_to_bf16(src[col]));
            const float silu = value / (1.0f + expf(-value));
            dst[col] = bf16_to_float(float_to_bf16(silu));
        }
    }
}

References __attribute__(), bf16_to_float(), float_to_bf16(), and silu().

◆ recurrent_split_conv_qkv_backward()

void recurrent_split_conv_qkv_backward	(	const float *	d_q,
		const float *	d_k,
		const float *	d_v,
		float *	d_packed_qkv,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 108 of file recurrent_split_kernels.c.

                                                  {
    recurrent_split_qkv_backward(d_q, d_k, d_v, d_packed_qkv, rows, q_dim, k_dim, v_dim);
}

References recurrent_split_qkv_backward().

◆ recurrent_split_conv_qkv_forward()

void recurrent_split_conv_qkv_forward	(	const float *	packed_qkv,
		float *	q,
		float *	k,
		float *	v,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 97 of file recurrent_split_kernels.c.

                                                 {
    recurrent_split_qkv_forward(packed_qkv, q, k, v, rows, q_dim, k_dim, v_dim);
}

References recurrent_split_qkv_forward().

◆ recurrent_split_qkv_backward()

void recurrent_split_qkv_backward	(	const float *	d_q,
		const float *	d_k,
		const float *	d_v,
		float *	d_packed_qkv,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 77 of file recurrent_split_kernels.c.

                                             {
    const int packed_dim = q_dim + k_dim + v_dim;
    for (int row = 0; row < rows; ++row) {
        const float *dq_src = d_q + (size_t) row * (size_t) q_dim;
        const float *dk_src = d_k + (size_t) row * (size_t) k_dim;
        const float *dv_src = d_v + (size_t) row * (size_t) v_dim;
        float *dst = d_packed_qkv + (size_t) row * (size_t) packed_dim;
        memcpy(dst, dq_src, (size_t) q_dim * sizeof(float));
        memcpy(dst + q_dim, dk_src, (size_t) k_dim * sizeof(float));
        memcpy(dst + q_dim + k_dim, dv_src, (size_t) v_dim * sizeof(float));
    }
}

Referenced by recurrent_split_conv_qkv_backward().

◆ recurrent_split_qkv_forward()

void recurrent_split_qkv_forward	(	const float *	packed_qkv,
		float *	q,
		float *	k,
		float *	v,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim
	)

Definition at line 5 of file recurrent_split_kernels.c.

                                            {
    const int packed_dim = q_dim + k_dim + v_dim;
    for (int row = 0; row < rows; ++row) {
        const float *src = packed_qkv + (size_t) row * (size_t) packed_dim;
        float *q_dst = q + (size_t) row * (size_t) q_dim;
        float *k_dst = k + (size_t) row * (size_t) k_dim;
        float *v_dst = v + (size_t) row * (size_t) v_dim;
        memcpy(q_dst, src, (size_t) q_dim * sizeof(float));
        memcpy(k_dst, src + q_dim, (size_t) k_dim * sizeof(float));
        memcpy(v_dst, src + q_dim + k_dim, (size_t) v_dim * sizeof(float));
    }
}

◆ relu2_backward()

void relu2_backward	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		size_t	n
	)

Definition at line 147 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    const __m512 vzero = _mm512_setzero_ps();
    const __m512 vtwo = _mm512_set1_ps(2.0f);
    for (; i + 15 < n; i += 16) {
        const __m512 vx = _mm512_loadu_ps(input + i);
        const __m512 vdy = _mm512_loadu_ps(d_output + i);
        const __mmask16 mask = _mm512_cmp_ps_mask(vx, vzero, _CMP_GT_OQ);
        const __m512 vdx = _mm512_mul_ps(_mm512_mul_ps(vtwo, vx), vdy);
        _mm512_storeu_ps(d_input + i, _mm512_maskz_mov_ps(mask, vdx));
    }
#elif defined(__AVX2__) || defined(__AVX__)
    const __m256 vzero = _mm256_setzero_ps();
    const __m256 vtwo = _mm256_set1_ps(2.0f);
    for (; i + 7 < n; i += 8) {
        const __m256 vx = _mm256_loadu_ps(input + i);
        const __m256 vdy = _mm256_loadu_ps(d_output + i);
        const __m256 mask = _mm256_cmp_ps(vx, vzero, _CMP_GT_OQ);
        const __m256 vdx = _mm256_mul_ps(_mm256_mul_ps(vtwo, vx), vdy);
        _mm256_storeu_ps(d_input + i, _mm256_and_ps(mask, vdx));
    }
#endif
 
    for (; i < n; ++i) {
        const float x = input[i];
        d_input[i] = (x > 0.0f) ? d_output[i] * 2.0f * x : 0.0f;
    }
}

References mask.

◆ relu2_forward()

void relu2_forward	(	const float *	input,
		float *	output,
		size_t	n
	)

Definition at line 120 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    const __m512 vzero = _mm512_setzero_ps();
    for (; i + 15 < n; i += 16) {
        const __m512 vx = _mm512_loadu_ps(input + i);
        const __m512 vr = _mm512_max_ps(vx, vzero);
        _mm512_storeu_ps(output + i, _mm512_mul_ps(vr, vr));
    }
#elif defined(__AVX2__) || defined(__AVX__)
    const __m256 vzero = _mm256_setzero_ps();
    for (; i + 7 < n; i += 8) {
        const __m256 vx = _mm256_loadu_ps(input + i);
        const __m256 vr = _mm256_max_ps(vx, vzero);
        _mm256_storeu_ps(output + i, _mm256_mul_ps(vr, vr));
    }
#endif
 
    for (; i < n; ++i) {
        const float x = input[i];
        output[i] = (x > 0.0f) ? x * x : 0.0f;
    }
}

◆ relu_backward()

void relu_backward	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		size_t	n
	)

Definition at line 84 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 vzero = _mm512_setzero_ps();
    for (; i + 15 < n; i += 16) {
        __m512 vx = _mm512_loadu_ps(input + i);
        __m512 vdy = _mm512_loadu_ps(d_output + i);
        __mmask16 mask = _mm512_cmp_ps_mask(vx, vzero, _CMP_GT_OQ);
        __m512 vdx = _mm512_maskz_mov_ps(mask, vdy);
        _mm512_storeu_ps(d_input + i, vdx);
    }
#elif defined(__AVX2__) || defined(__AVX__)
    __m256 vzero = _mm256_setzero_ps();
    for (; i + 7 < n; i += 8) {
        __m256 vx = _mm256_loadu_ps(input + i);
        __m256 vdy = _mm256_loadu_ps(d_output + i);
        // Result is all 1s (0xFFFFFFFF) if true, 0 if false.
        __m256 mask = _mm256_cmp_ps(vx, vzero, _CMP_GT_OQ);
        __m256 vdx = _mm256_and_ps(mask, vdy);
        _mm256_storeu_ps(d_input + i, vdx);
    }
#endif
 
    // Scalar fallback
    for (; i < n; ++i) {
        d_input[i] = (input[i] > 0.0f) ? d_output[i] : 0.0f;
    }
}

References mask.

◆ relu_backward_bf16()

void relu_backward_bf16	(	const uint16_t *	input,
		const uint16_t *	d_output,
		uint16_t *	d_input,
		size_t	n
	)

Definition at line 69 of file relu_kernels_bf16.c.

{
    if (!input || !d_output || !d_input) {
        return;
    }
 
    size_t i = 0;
#if defined(__AVX2__)
    const __m256i zero = _mm256_setzero_si256();
    for (; i + 16 <= n; i += 16) {
        const __m256i x = _mm256_loadu_si256((const __m256i *)(input + i));
        const __m256i dy = _mm256_loadu_si256((const __m256i *)(d_output + i));
        const __m256i mask = _mm256_cmpgt_epi16(x, zero);
        const __m256i dx = _mm256_and_si256(dy, mask);
        _mm256_storeu_si256((__m256i *)(d_input + i), dx);
    }
#endif
    for (; i < n; ++i) {
        d_input[i] = ((input[i] & 0x8000u) == 0u && (input[i] & 0x7fffu) != 0u) ? d_output[i] : 0u;
    }
}

References mask.

◆ relu_forward()

void relu_forward	(	const float *	input,
		float *	output,
		size_t	n
	)

Definition at line 26 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 vzero = _mm512_setzero_ps();
    for (; i + 15 < n; i += 16) {
        __m512 vx = _mm512_loadu_ps(input + i);
        __m512 vy = _mm512_max_ps(vx, vzero);
        _mm512_storeu_ps(output + i, vy);
    }
#elif defined(__AVX2__) || defined(__AVX__)
    __m256 vzero = _mm256_setzero_ps();
    for (; i + 7 < n; i += 8) {
        __m256 vx = _mm256_loadu_ps(input + i);
        __m256 vy = _mm256_max_ps(vx, vzero);
        _mm256_storeu_ps(output + i, vy);
    }
#endif
 
    // Scalar fallback
    for (; i < n; ++i) {
        float x = input[i];
        output[i] = (x > 0.0f) ? x : 0.0f;
    }
}

◆ relu_forward_bf16()

void relu_forward_bf16	(	const uint16_t *	input,
		uint16_t *	output,
		size_t	n
	)

Definition at line 27 of file relu_kernels_bf16.c.

{
    if (!input || !output) {
        return;
    }
 
    size_t i = 0;
#if defined(__AVX2__)
    const __m256i zero = _mm256_setzero_si256();
    for (; i + 16 <= n; i += 16) {
        const __m256i x = _mm256_loadu_si256((const __m256i *)(input + i));
        const __m256i mask = _mm256_cmpgt_epi16(x, zero);
        const __m256i y = _mm256_and_si256(x, mask);
        _mm256_storeu_si256((__m256i *)(output + i), y);
    }
#endif
    for (; i < n; ++i) {
        output[i] = (input[i] & 0x8000u) ? 0u : input[i];
    }
}

References mask.

◆ relu_forward_inplace()

void relu_forward_inplace	(	float *	data,
		size_t	n
	)

Definition at line 54 of file relu_kernels.c.

{
    size_t i = 0;
 
#if defined(__AVX512F__)
    __m512 vzero = _mm512_setzero_ps();
    for (; i + 15 < n; i += 16) {
        __m512 vx = _mm512_loadu_ps(data + i);
        __m512 vy = _mm512_max_ps(vx, vzero);
        _mm512_storeu_ps(data + i, vy);
    }
#elif defined(__AVX2__) || defined(__AVX__)
    __m256 vzero = _mm256_setzero_ps();
    for (; i + 7 < n; i += 8) {
        __m256 vx = _mm256_loadu_ps(data + i);
        __m256 vy = _mm256_max_ps(vx, vzero);
        _mm256_storeu_ps(data + i, vy);
    }
#endif
 
    // Scalar fallback
    for (; i < n; ++i) {
        float x = data[i];
        if (x < 0.0f) {
            data[i] = 0.0f;
        }
    }
}

◆ relu_forward_inplace_bf16()

void relu_forward_inplace_bf16	(	uint16_t *	data,
		size_t	n
	)

Definition at line 48 of file relu_kernels_bf16.c.

{
    if (!data) {
        return;
    }
 
    size_t i = 0;
#if defined(__AVX2__)
    const __m256i zero = _mm256_setzero_si256();
    for (; i + 16 <= n; i += 16) {
        const __m256i x = _mm256_loadu_si256((const __m256i *)(data + i));
        const __m256i mask = _mm256_cmpgt_epi16(x, zero);
        const __m256i y = _mm256_and_si256(x, mask);
        _mm256_storeu_si256((__m256i *)(data + i), y);
    }
#endif
    for (; i < n; ++i) {
        data[i] = (data[i] & 0x8000u) ? 0u : data[i];
    }
}

References mask.

◆ rmsnorm_backward()

void rmsnorm_backward	(	const float *	d_output,
		const float *	input,
		const float *	gamma,
		const float *	rstd_cache,
		float *	d_input,
		float *	d_gamma,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim
	)

RMSNorm backward pass

Test:

test_rmsnorm.py::TestRMSNormBackward::test_backward_tokens

test_rmsnorm.py::TestRMSNormBackward::test_backward_single

test_parity.py::test_rmsnorm_backward_parity

Computes dX and dGamma given dY, X, gamma, and cached rstd. dX_i = rstd * (dY_i * gamma_i - x_hat_i * m) dGamma_i = sum_t (dY_i * x_hat_i)

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

Definition at line 715 of file rmsnorm_kernels.c.

{
    int T = tokens;
    int D = d_model;
    int aligned = aligned_embed_dim;
 
    if (ck_strict_parity_enabled()) {
        rmsnorm_backward_strict_scalar(d_output, input, gamma, rstd_cache, d_input, d_gamma, T, D, aligned);
        return;
    }
 
    // Zero parameter gradients
#if defined(__AVX512F__)
    {
        int d = 0;
        for (; d + 16 <= D; d += 16) {
            _mm512_storeu_ps(&d_gamma[d], _mm512_setzero_ps());
        }
        for (; d < D; ++d) {
            d_gamma[d] = 0.0f;
        }
    }
#elif defined(__AVX__)
    {
        int d = 0;
        for (; d + 8 <= D; d += 8) {
            _mm256_storeu_ps(&d_gamma[d], _mm256_setzero_ps());
        }
        for (; d < D; ++d) {
            d_gamma[d] = 0.0f;
        }
    }
#else
    for (int d = 0; d < D; ++d) {
        d_gamma[d] = 0.0f;
    }
#endif
 
    for (int t = 0; t < T; ++t) {
        const float *x = input + (size_t)t * aligned;
        const float *dY = d_output + (size_t)t * aligned;
        float *dX = d_input + (size_t)t * aligned;
 
        float rstd = rstd_cache[t];
 
#if defined(__AVX512F__)
        // Compute m = (1/D) * sum_j (dY_j * gamma_j * x_hat_j)
        __m512 rstd_vec = _mm512_set1_ps(rstd);
        __m512 sum_vec = _mm512_setzero_ps();
        int d = 0;
 
        for (; d + 16 <= D; d += 16) {
            __m512 xv = _mm512_loadu_ps(&x[d]);
            __m512 dyv = _mm512_loadu_ps(&dY[d]);
            __m512 gv = _mm512_loadu_ps(&gamma[d]);
            __m512 x_hat = _mm512_mul_ps(xv, rstd_vec);
            // sum += dY * gamma * x_hat
            __m512 prod = _mm512_mul_ps(dyv, gv);
            sum_vec = _mm512_fmadd_ps(prod, x_hat, sum_vec);
        }
        float sum_dY_g_xhat = _mm512_reduce_add_ps(sum_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            float x_hat = x[d] * rstd;
            sum_dY_g_xhat += dY[d] * gamma[d] * x_hat;
        }
        float m = sum_dY_g_xhat / (float)D;
 
        // Compute dX and accumulate dGamma (vectorized)
        __m512 m_vec = _mm512_set1_ps(m);
        d = 0;
        for (; d + 16 <= D; d += 16) {
            __m512 xv = _mm512_loadu_ps(&x[d]);
            __m512 dyv = _mm512_loadu_ps(&dY[d]);
            __m512 gv = _mm512_loadu_ps(&gamma[d]);
            __m512 dgv = _mm512_loadu_ps(&d_gamma[d]);
 
            __m512 x_hat = _mm512_mul_ps(xv, rstd_vec);
 
            // dX = rstd * (dY * gamma - x_hat * m)
            __m512 dy_g = _mm512_mul_ps(dyv, gv);
            __m512 xhat_m = _mm512_mul_ps(x_hat, m_vec);
            __m512 diff = _mm512_sub_ps(dy_g, xhat_m);
            __m512 dxv = _mm512_mul_ps(rstd_vec, diff);
            _mm512_storeu_ps(&dX[d], dxv);
 
            // d_gamma += dY * x_hat
            dgv = _mm512_fmadd_ps(dyv, x_hat, dgv);
            _mm512_storeu_ps(&d_gamma[d], dgv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            float x_hat = x[d] * rstd;
            float dy = dY[d];
            dX[d] = rstd * (dy * gamma[d] - x_hat * m);
            d_gamma[d] += dy * x_hat;
        }
 
#elif defined(__AVX__)
        // Compute m = (1/D) * sum_j (dY_j * gamma_j * x_hat_j)
        __m256 rstd_vec = _mm256_set1_ps(rstd);
        __m256 sum_vec = _mm256_setzero_ps();
        int d = 0;
 
        for (; d + 8 <= D; d += 8) {
            __m256 xv = _mm256_loadu_ps(&x[d]);
            __m256 dyv = _mm256_loadu_ps(&dY[d]);
            __m256 gv = _mm256_loadu_ps(&gamma[d]);
            __m256 x_hat = _mm256_mul_ps(xv, rstd_vec);
            // sum += dY * gamma * x_hat (no FMA, use mul + mul + add)
            __m256 prod = _mm256_mul_ps(dyv, gv);
            __m256 prod2 = _mm256_mul_ps(prod, x_hat);
            sum_vec = _mm256_add_ps(sum_vec, prod2);
        }
        float sum_dY_g_xhat = hsum256_ps_rmsnorm(sum_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            float x_hat = x[d] * rstd;
            sum_dY_g_xhat += dY[d] * gamma[d] * x_hat;
        }
        float m = sum_dY_g_xhat / (float)D;
 
        // Compute dX and accumulate dGamma (vectorized)
        __m256 m_vec = _mm256_set1_ps(m);
        d = 0;
        for (; d + 8 <= D; d += 8) {
            __m256 xv = _mm256_loadu_ps(&x[d]);
            __m256 dyv = _mm256_loadu_ps(&dY[d]);
            __m256 gv = _mm256_loadu_ps(&gamma[d]);
            __m256 dgv = _mm256_loadu_ps(&d_gamma[d]);
 
            __m256 x_hat = _mm256_mul_ps(xv, rstd_vec);
 
            // dX = rstd * (dY * gamma - x_hat * m)
            __m256 dy_g = _mm256_mul_ps(dyv, gv);
            __m256 xhat_m = _mm256_mul_ps(x_hat, m_vec);
            __m256 diff = _mm256_sub_ps(dy_g, xhat_m);
            __m256 dxv = _mm256_mul_ps(rstd_vec, diff);
            _mm256_storeu_ps(&dX[d], dxv);
 
            // d_gamma += dY * x_hat
            __m256 dy_xhat = _mm256_mul_ps(dyv, x_hat);
            dgv = _mm256_add_ps(dgv, dy_xhat);
            _mm256_storeu_ps(&d_gamma[d], dgv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            float x_hat = x[d] * rstd;
            float dy = dY[d];
            dX[d] = rstd * (dy * gamma[d] - x_hat * m);
            d_gamma[d] += dy * x_hat;
        }
 
#else
        // Scalar fallback
        // Compute m = (1/D) * sum_j (dY_j * gamma_j * x_hat_j)
        float sum_dY_g_xhat = 0.0f;
        for (int d = 0; d < D; ++d) {
            float x_hat = x[d] * rstd;
            sum_dY_g_xhat += dY[d] * gamma[d] * x_hat;
        }
        float m = sum_dY_g_xhat / (float)D;
 
        // Compute dX and accumulate dGamma
        for (int d = 0; d < D; ++d) {
            float x_hat = x[d] * rstd;
            float dy = dY[d];
            dX[d] = rstd * (dy * gamma[d] - x_hat * m);
            d_gamma[d] += dy * x_hat;
        }
#endif
 
        // Zero padding gradients (if any)
        for (int d = D; d < aligned; ++d) {
            dX[d] = 0.0f;
        }
    }
}

References ck_strict_parity_enabled(), and rmsnorm_backward_strict_scalar().

Referenced by ck_layer_backward_rmsnorm_swiglu(), rmsnorm_backward_int4(), and rmsnorm_backward_int8().

◆ rmsnorm_backward_bf16()

void rmsnorm_backward_bf16	(	const uint16_t *	d_output,
		const uint16_t *	input,
		const float *	gamma,
		const float *	rstd_cache,
		uint16_t *	d_input,
		float *	d_gamma,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim
	)

Definition at line 113 of file rmsnorm_kernels_bf16.c.

{
    int T = tokens;
    int D = d_model;
    int aligned = aligned_embed_dim;
 
    if (!d_output || !input || !gamma || !rstd_cache || !d_input || !d_gamma) {
        return;
    }
 
    // Zero parameter gradients
#if defined(__AVX512F__)
    {
        int d = 0;
        for (; d + 16 <= D; d += 16) {
            _mm512_storeu_ps(&d_gamma[d], _mm512_setzero_ps());
        }
        for (; d < D; ++d) {
            d_gamma[d] = 0.0f;
        }
    }
#else
    for (int d = 0; d < D; ++d) {
        d_gamma[d] = 0.0f;
    }
#endif
 
    for (int t = 0; t < T; ++t) {
        const uint16_t *x_bf16 = input + (size_t)t * aligned;
        const uint16_t *dY_bf16 = d_output + (size_t)t * aligned;
        uint16_t *dX_bf16 = d_input + (size_t)t * aligned;
        float rstd = rstd_cache[t];
 
#if defined(__AVX512F__)
        // Compute m = (1/D) * sum_j (dY_j * gamma_j * x_hat_j)
        __m512 rstd_vec = _mm512_set1_ps(rstd);
        __m512 sum_vec = _mm512_setzero_ps();
        int d = 0;
 
        for (; d + 16 <= D; d += 16) {
            __m512 xv = bf16_loadu_cvt_fp32(&x_bf16[d]);
            __m512 dyv = bf16_loadu_cvt_fp32(&dY_bf16[d]);
            __m512 gv = _mm512_loadu_ps(&gamma[d]);
            __m512 x_hat = _mm512_mul_ps(xv, rstd_vec);
            // sum += dY * gamma * x_hat
            __m512 prod = _mm512_mul_ps(dyv, gv);
            sum_vec = _mm512_fmadd_ps(prod, x_hat, sum_vec);
        }
        float sum_dY_g_xhat = _mm512_reduce_add_ps(sum_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            float x = bf16_to_float(x_bf16[d]);
            float x_hat = x * rstd;
            float dy = bf16_to_float(dY_bf16[d]);
            sum_dY_g_xhat += dy * gamma[d] * x_hat;
        }
        float m = sum_dY_g_xhat / (float)D;
 
        // Compute dX and accumulate dGamma (vectorized)
        __m512 m_vec = _mm512_set1_ps(m);
        d = 0;
        for (; d + 16 <= D; d += 16) {
            __m512 xv = bf16_loadu_cvt_fp32(&x_bf16[d]);
            __m512 dyv = bf16_loadu_cvt_fp32(&dY_bf16[d]);
            __m512 gv = _mm512_loadu_ps(&gamma[d]);
            __m512 dgv = _mm512_loadu_ps(&d_gamma[d]);
 
            __m512 x_hat = _mm512_mul_ps(xv, rstd_vec);
 
            // dX = rstd * (dY * gamma - x_hat * m)
            __m512 dy_g = _mm512_mul_ps(dyv, gv);
            __m512 xhat_m = _mm512_mul_ps(x_hat, m_vec);
            __m512 diff = _mm512_sub_ps(dy_g, xhat_m);
            __m512 dxv = _mm512_mul_ps(rstd_vec, diff);
            fp32_cvt_storeu_bf16(&dX_bf16[d], dxv);
 
            // d_gamma += dY * x_hat
            dgv = _mm512_fmadd_ps(dyv, x_hat, dgv);
            _mm512_storeu_ps(&d_gamma[d], dgv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            float x = bf16_to_float(x_bf16[d]);
            float x_hat = x * rstd;
            float dy = bf16_to_float(dY_bf16[d]);
            float dx = rstd * (dy * gamma[d] - x_hat * m);
            dX_bf16[d] = float_to_bf16(dx);
            d_gamma[d] += dy * x_hat;
        }
 
#else
        // Scalar fallback
        double sum_dY_g_xhat = 0.0;
        for (int d = 0; d < D; ++d) {
            float x = bf16_to_float(x_bf16[d]);
            float x_hat = x * rstd;
            float dy = bf16_to_float(dY_bf16[d]);
            sum_dY_g_xhat += (double)dy * (double)gamma[d] * (double)x_hat;
        }
        float m = (float)(sum_dY_g_xhat / (double)D);
 
        for (int d = 0; d < D; ++d) {
            float x = bf16_to_float(x_bf16[d]);
            float x_hat = x * rstd;
            float dy = bf16_to_float(dY_bf16[d]);
            float dx = rstd * (dy * gamma[d] - x_hat * m);
            dX_bf16[d] = float_to_bf16(dx);
            d_gamma[d] += dy * x_hat;
        }
#endif
 
        // Zero padding gradients
        for (int d = D; d < aligned; ++d) {
            dX_bf16[d] = 0;
        }
    }
}

References bf16_to_float(), and float_to_bf16().

◆ rmsnorm_backward_int4()

void rmsnorm_backward_int4	(	const uint8_t *	d_output,
		const uint8_t *	input,
		const float *	gamma,
		const float *	rstd_cache,
		uint8_t *	d_input,
		float *	d_gamma,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float *	scratch_d_output,
		float *	scratch_input,
		float *	scratch_d_input
	)

Definition at line 104 of file rmsnorm_kernels_int4.c.

{
    if (!d_output || !input || !gamma || !rstd_cache || !d_input || !d_gamma) return;
    if (!scratch_d_output || !scratch_input || !scratch_d_input) return;
 
    size_t total = (size_t)tokens * (size_t)aligned_embed_dim;
 
    convert_int4_to_float(d_output, scratch_d_output, total);
    convert_int4_to_float(input, scratch_input, total);
 
    for (int d = 0; d < d_model; ++d) {
        d_gamma[d] = 0.0f;
    }
 
    rmsnorm_backward(scratch_d_output,
                     scratch_input,
                     gamma,
                     rstd_cache,
                     scratch_d_input,
                     d_gamma,
                     tokens,
                     d_model,
                     aligned_embed_dim);
 
    convert_float_to_int4(scratch_d_input, d_input, total);
}

References convert_float_to_int4(), convert_int4_to_float(), and rmsnorm_backward().

◆ rmsnorm_backward_int8()

void rmsnorm_backward_int8	(	const int8_t *	d_output,
		const int8_t *	input,
		const float *	gamma,
		const float *	rstd_cache,
		int8_t *	d_input,
		float *	d_gamma,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float *	scratch_d_output,
		float *	scratch_input,
		float *	scratch_d_input
	)

Definition at line 84 of file rmsnorm_kernels_int8.c.

{
    if (!d_output || !input || !gamma || !rstd_cache || !d_input || !d_gamma) return;
    if (!scratch_d_output || !scratch_input || !scratch_d_input) return;
 
    size_t total = (size_t)tokens * (size_t)aligned_embed_dim;
 
    convert_int8_to_float(d_output, scratch_d_output, total);
    convert_int8_to_float(input, scratch_input, total);
 
    // Zero gamma gradient before accumulation.
    for (int d = 0; d < d_model; ++d) {
        d_gamma[d] = 0.0f;
    }
 
    rmsnorm_backward(scratch_d_output,
                     scratch_input,
                     gamma,
                     rstd_cache,
                     scratch_d_input,
                     d_gamma,
                     tokens,
                     d_model,
                     aligned_embed_dim);
 
    convert_float_to_int8(scratch_d_input, d_input, total);
}

References convert_float_to_int8(), convert_int8_to_float(), and rmsnorm_backward().

◆ rmsnorm_forward()

void rmsnorm_forward	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 621 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_strided_f32(
        input,
        gamma,
        output,
        rstd_cache,
        tokens,
        d_model,
        aligned_embed_dim,
        aligned_embed_dim,
        eps
    );
}

References rmsnorm_forward_strided_f32().

◆ rmsnorm_forward_bf16()

void rmsnorm_forward_bf16	(	const uint16_t *	input,
		const float *	gamma,
		uint16_t *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 24 of file rmsnorm_kernels_bf16.c.

{
    int T = tokens;
    int D = d_model;
    int aligned = aligned_embed_dim;
 
    for (int t = 0; t < T; ++t) {
        const uint16_t *x_bf16 = input + (size_t)t * aligned;
        float *rstd_ptr = rstd_cache ? (rstd_cache + t) : NULL;
        uint16_t *out_bf16 = output + (size_t)t * aligned;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        __m512 sum_sq_vec = _mm512_setzero_ps();
        int d = 0;
 
        // Vectorized sum of squares
        for (; d + 16 <= D; d += 16) {
            __m512 xv = bf16_loadu_cvt_fp32(&x_bf16[d]);
            sum_sq_vec = _mm512_fmadd_ps(xv, xv, sum_sq_vec);
        }
        float sum_sq = _mm512_reduce_add_ps(sum_sq_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            float x = bf16_to_float(x_bf16[d]);
            sum_sq += x * x;
        }
 
        float mean_sq = sum_sq / (float)D;
        float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_ptr) {
            *rstd_ptr = rstd;
        }
 
        // Apply normalization and scale (vectorized)
        __m512 rstd_vec = _mm512_set1_ps(rstd);
        d = 0;
        for (; d + 16 <= D; d += 16) {
            __m512 xv = bf16_loadu_cvt_fp32(&x_bf16[d]);
            __m512 gv = _mm512_loadu_ps(&gamma[d]);
            __m512 x_hat = _mm512_mul_ps(xv, rstd_vec);
            __m512 yv = _mm512_mul_ps(x_hat, gv);
            fp32_cvt_storeu_bf16(&out_bf16[d], yv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            float x = bf16_to_float(x_bf16[d]);
            float y = x * rstd * gamma[d];
            out_bf16[d] = float_to_bf16(y);
        }
 
#else
        // Scalar fallback
        double sum_sq = 0.0;
        for (int d = 0; d < D; ++d) {
            float x = bf16_to_float(x_bf16[d]);
            sum_sq += (double)x * (double)x;
        }
        double mean_sq = sum_sq / (double)D;
        double r = sqrt(mean_sq + (double)eps);
        float rstd = (float)(1.0 / r);
        if (rstd_ptr) {
            *rstd_ptr = rstd;
        }
 
        for (int d = 0; d < D; ++d) {
            float x = bf16_to_float(x_bf16[d]);
            float x_hat = x * rstd;
            float y = x_hat * gamma[d];
            out_bf16[d] = float_to_bf16(y);
        }
#endif
 
        // Zero padding
        for (int d = D; d < aligned; ++d) {
            out_bf16[d] = 0;
        }
    }
}

References bf16_to_float(), and float_to_bf16().

◆ rmsnorm_forward_fp32_square_fp64_sum()

void rmsnorm_forward_fp32_square_fp64_sum	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

◆ rmsnorm_forward_fp64_sum()

void rmsnorm_forward_fp64_sum	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 137 of file rmsnorm_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)aligned_embed_dim;
        float *y = output + (size_t)t * (size_t)aligned_embed_dim;
        /* This provider's contract requires an ascending scalar reduction.
         * Keep the accumulator volatile so whole-program optimization cannot
         * reassociate the sum or replace it with SIMD partial reductions. */
        volatile double sum_sq = 0.0;
        for (int d = 0; d < d_model; ++d) {
            const float square = x[d] * x[d];
            sum_sq = sum_sq + (double)square;
        }
        const float mean_sq = (float)(sum_sq / (double)d_model);
        const float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
        for (int d = 0; d < d_model; ++d) {
            const float normalized = x[d] * rstd;
            y[d] = normalized * gamma[d];
        }
        for (int d = d_model; d < aligned_embed_dim; ++d) {
            y[d] = 0.0f;
        }
    }
}

◆ rmsnorm_forward_int4()

void rmsnorm_forward_int4	(	const uint8_t *	input,
		const float *	gamma,
		uint8_t *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps,
		float *	scratch_input,
		float *	scratch_output
	)

Definition at line 78 of file rmsnorm_kernels_int4.c.

{
    if (!input || !gamma || !output) return;
    if (!scratch_input || !scratch_output) return;
 
    size_t total = (size_t)tokens * (size_t)aligned_embed_dim;
 
    convert_int4_to_float(input, scratch_input, total);
    rmsnorm_forward(scratch_input, gamma, scratch_output, rstd_cache,
                    tokens, d_model, aligned_embed_dim, eps);
    convert_float_to_int4(scratch_output, output, total);
}

References convert_float_to_int4(), convert_int4_to_float(), and rmsnorm_forward().

◆ rmsnorm_forward_int8()

void rmsnorm_forward_int8	(	const int8_t *	input,
		const float *	gamma,
		int8_t *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps,
		float *	scratch_input,
		float *	scratch_output
	)

Definition at line 58 of file rmsnorm_kernels_int8.c.

{
    if (!input || !gamma || !output) return;
    if (!scratch_input || !scratch_output) return;
 
    size_t total = (size_t)tokens * (size_t)aligned_embed_dim;
 
    convert_int8_to_float(input, scratch_input, total);
    rmsnorm_forward(scratch_input, gamma, scratch_output, rstd_cache,
                    tokens, d_model, aligned_embed_dim, eps);
    convert_float_to_int8(scratch_output, output, total);
}

References convert_float_to_int8(), convert_int8_to_float(), and rmsnorm_forward().

◆ rmsnorm_forward_kv_lora()

void rmsnorm_forward_kv_lora	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 643 of file rmsnorm_kernels.c.

{
    rmsnorm_forward(input, gamma, output, rstd_cache, tokens, d_model, aligned_embed_dim, eps);
}

References rmsnorm_forward().

◆ rmsnorm_forward_llama_production()

void rmsnorm_forward_llama_production	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 196 of file rmsnorm_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)aligned_embed_dim;
        float *y = output + (size_t)t * (size_t)aligned_embed_dim;
        volatile double sum_sq = 0.0;
        for (int d = 0; d < d_model; ++d) {
            const float square = x[d] * x[d];
            sum_sq = sum_sq + (double)square;
        }
        const float mean_sq = (float)(sum_sq / (double)d_model);
        const float rstd = rmsnorm_llama_production_rstd(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
        for (int d = 0; d < d_model; ++d) {
            /*
             * Keep the RMSNorm + scale expression fused at the source level.
             * llama.cpp's CPU graph fuses GGML_OP_RMS_NORM followed by
             * GGML_OP_MUL and evaluates this left-associative expression in
             * one kernel.  Materializing the normalized value as a named
             * float introduces a store/load rounding boundary under ICX and
             * differs by one ULP for otherwise identical inputs.
             */
            y[d] = x[d] * rstd * gamma[d];
        }
        for (int d = d_model; d < aligned_embed_dim; ++d) {
            y[d] = 0.0f;
        }
    }
}

References rmsnorm_llama_production_rstd().

Referenced by qwen4_group_rmsnorm_llama(), recurrent_norm_gate_llama_avx2_forward(), and recurrent_norm_sigmoid_gate_llama_avx2_forward().

◆ rmsnorm_forward_no_weight()

void rmsnorm_forward_no_weight	(	const float *	input,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 655 of file rmsnorm_kernels.c.

{
    if (!input || !output || tokens <= 0 || d_model <= 0 || aligned_embed_dim <= 0) {
        return;
    }
    const float inv_d = 1.0f / (float)d_model;
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)aligned_embed_dim;
        float *y = output + (size_t)t * (size_t)aligned_embed_dim;
        double sum_sq = 0.0;
        for (int d = 0; d < d_model; ++d) {
            sum_sq += (double)x[d] * (double)x[d];
        }
        const float rstd = 1.0f / sqrtf((float)(sum_sq * (double)inv_d) + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
        for (int d = 0; d < d_model; ++d) {
            y[d] = x[d] * rstd;
        }
        for (int d = d_model; d < aligned_embed_dim; ++d) {
            y[d] = 0.0f;
        }
    }
}

Referenced by gemma4_v_norm_forward().

◆ rmsnorm_forward_parallel_dispatch()

void rmsnorm_forward_parallel_dispatch	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

◆ rmsnorm_forward_pytorch_bf16_storage()

void rmsnorm_forward_pytorch_bf16_storage	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 395 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage_impl(
        input, gamma, output, rstd_cache, tokens, d_model,
        aligned_embed_dim, aligned_embed_dim, eps, 0);
}

References rmsnorm_forward_pytorch_bf16_storage_impl().

Referenced by recurrent_norm_gate_pytorch_bf16_storage(), and recurrent_norm_sigmoid_gate_pytorch_bf16_storage().

◆ rmsnorm_forward_qwen3next_pytorch_bf16_storage()

void rmsnorm_forward_qwen3next_pytorch_bf16_storage	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 424 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage_impl(
        input, gamma, output, rstd_cache, tokens, d_model,
        aligned_embed_dim, aligned_embed_dim, eps, 1);
}

References rmsnorm_forward_pytorch_bf16_storage_impl().

Referenced by hyper_connection_mix_bf16(), and qwen4_group_rmsnorm_pytorch_bf16().

◆ rmsnorm_forward_strided_f32()

void rmsnorm_forward_strided_f32	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	input_stride,
		int	output_stride,
		float	eps
	)

RMSNorm forward pass

Test:

test_rmsnorm.py::TestRMSNormForward::test_fp32_tokens

test_rmsnorm.py::TestRMSNormForward::test_fp32_single

test_rmsnorm.py::TestRMSNormForward::test_perf_rolled

test_layernorm.py::TestLayerNormForward::test_rmsnorm_compat

test_parity.py::test_rmsnorm_parity

RMSNorm: y[i] = gamma[i] * x[i] / sqrt(mean(x^2) + eps)

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

Definition at line 492 of file rmsnorm_kernels.c.

{
    int T = tokens;
    int D = d_model;
 
    const char *exact_env = getenv("CK_RMSNORM_EXACT");
    if (ck_strict_parity_enabled() || (exact_env && atoi(exact_env) != 0)) {
        rmsnorm_forward_strict_scalar(
            input, gamma, output, rstd_cache, T, D, input_stride, output_stride, eps
        );
        return;
    }
 
    for (int t = 0; t < T; ++t) {
        const float *x = input + (size_t)t * (size_t)input_stride;
        float *y = output + (size_t)t * (size_t)output_stride;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        __m512 sum_sq_vec = _mm512_setzero_ps();
        int d = 0;
 
        // Vectorized sum of squares
        for (; d + 16 <= D; d += 16) {
            __m512 xv = _mm512_loadu_ps(&x[d]);
            sum_sq_vec = _mm512_fmadd_ps(xv, xv, sum_sq_vec);
        }
        float sum_sq = _mm512_reduce_add_ps(sum_sq_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            sum_sq += x[d] * x[d];
        }
 
        float mean_sq = sum_sq / (float)D;
        float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
 
        // Apply normalization and scale (vectorized)
        __m512 rstd_vec = _mm512_set1_ps(rstd);
        d = 0;
        for (; d + 16 <= D; d += 16) {
            __m512 xv = _mm512_loadu_ps(&x[d]);
            __m512 gv = _mm512_loadu_ps(&gamma[d]);
            __m512 x_hat = _mm512_mul_ps(xv, rstd_vec);
            __m512 yv = _mm512_mul_ps(x_hat, gv);
            _mm512_storeu_ps(&y[d], yv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            y[d] = x[d] * rstd * gamma[d];
        }
 
#elif defined(__AVX__)
        // AVX: Process 8 floats at a time
        __m256 sum_sq_vec = _mm256_setzero_ps();
        int d = 0;
 
        // Vectorized sum of squares (no FMA in AVX1, use mul + add)
        for (; d + 8 <= D; d += 8) {
            __m256 xv = _mm256_loadu_ps(&x[d]);
            __m256 xv_sq = _mm256_mul_ps(xv, xv);
            sum_sq_vec = _mm256_add_ps(sum_sq_vec, xv_sq);
        }
        float sum_sq = hsum256_ps_rmsnorm(sum_sq_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            sum_sq += x[d] * x[d];
        }
 
        float mean_sq = sum_sq / (float)D;
        float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
 
        // Apply normalization and scale (vectorized)
        __m256 rstd_vec = _mm256_set1_ps(rstd);
        d = 0;
        for (; d + 8 <= D; d += 8) {
            __m256 xv = _mm256_loadu_ps(&x[d]);
            __m256 gv = _mm256_loadu_ps(&gamma[d]);
            __m256 x_hat = _mm256_mul_ps(xv, rstd_vec);
            __m256 yv = _mm256_mul_ps(x_hat, gv);
            _mm256_storeu_ps(&y[d], yv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            y[d] = x[d] * rstd * gamma[d];
        }
 
#else
        // Scalar fallback
        float sum_sq = 0.0f;
        for (int d = 0; d < D; ++d) {
            float v = x[d];
            sum_sq += v * v;
        }
        float mean_sq = sum_sq / (float)D;
        float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
 
        // Apply normalization and scale
        for (int d = 0; d < D; ++d) {
            float x_hat = x[d] * rstd;
            y[d] = x_hat * gamma[d];
        }
#endif
 
        // Zero padding (if any)
        for (int d = D; d < output_stride; ++d) {
            y[d] = 0.0f;
        }
    }
}

References ck_strict_parity_enabled(), and rmsnorm_forward_strict_scalar().

Referenced by rmsnorm_forward().

◆ rmsnorm_forward_strided_pytorch_bf16_storage()

void rmsnorm_forward_strided_pytorch_bf16_storage	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	input_stride,
		int	output_stride,
		float	eps
	)

Definition at line 409 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage_impl(
        input, gamma, output, rstd_cache, tokens, d_model,
        input_stride, output_stride, eps, 0);
}

References rmsnorm_forward_pytorch_bf16_storage_impl().

◆ rope_backward()

void rope_backward	(	const float *	d_out,
		float *	d_x,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset
	)

RoPE backward (inverse rotation)

Test:

test_rope.py::TestRoPEBackward::test_rope_backward

test_rope.py::TestRoPEBackward::test_rope_backward_vs_separate

RoPE backward: inverse rotation (rotate by -θ). Since cos(-θ) = cos(θ) and sin(-θ) = -sin(θ): d_x[2i] = d0 * c + d1 * s d_x[2i+1] = -d0 * s + d1 * c

After changes: make test

Definition at line 884 of file rope_kernels.c.

{
    size_t head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
    int half_dim = head_dim / 2;
 
    for (int h = 0; h < num_heads; ++h) {
        for (int t = 0; t < num_tokens; ++t) {
            int pos = pos_offset + t;
            const float *cos_row = cos_cache + pos * half_dim;
            const float *sin_row = sin_cache + pos * half_dim;
 
            size_t idx = h * head_stride + (size_t)t * (size_t)aligned_head_dim;
            const float *d_out_row = d_out + idx;
            float *d_x_row = d_x + idx;
 
#if defined(__AVX512F__)
            int i = 0;
            for (; i + 16 <= half_dim; i += 16) {
                __m512 d0 = _mm512_loadu_ps(&d_out_row[i]);
                __m512 d1 = _mm512_loadu_ps(&d_out_row[i + half_dim]);
                __m512 c = _mm512_loadu_ps(&cos_row[i]);
                __m512 s = _mm512_loadu_ps(&sin_row[i]);
 
                // Inverse: d_x[i] = d0 * c + d1 * s
                __m512 r0 = _mm512_fmadd_ps(d0, c, _mm512_mul_ps(d1, s));
                // Inverse: d_x[i+half] = -d0 * s + d1 * c
                __m512 r1 = _mm512_fmsub_ps(d1, c, _mm512_mul_ps(d0, s));
 
                _mm512_storeu_ps(&d_x_row[i], r0);
                _mm512_storeu_ps(&d_x_row[i + half_dim], r1);
            }
            for (; i < half_dim; ++i) {
                float d0 = d_out_row[i];
                float d1 = d_out_row[i + half_dim];
                float c = cos_row[i];
                float s = sin_row[i];
                d_x_row[i] = d0 * c + d1 * s;
                d_x_row[i + half_dim] = -d0 * s + d1 * c;
            }
 
#elif defined(__AVX__)
            int i = 0;
            for (; i + 8 <= half_dim; i += 8) {
                __m256 d0 = _mm256_loadu_ps(&d_out_row[i]);
                __m256 d1 = _mm256_loadu_ps(&d_out_row[i + half_dim]);
                __m256 c = _mm256_loadu_ps(&cos_row[i]);
                __m256 s = _mm256_loadu_ps(&sin_row[i]);
 
                // Inverse: d_x[i] = d0 * c + d1 * s
                __m256 d0c = _mm256_mul_ps(d0, c);
                __m256 d1s = _mm256_mul_ps(d1, s);
                __m256 r0 = _mm256_add_ps(d0c, d1s);
 
                // Inverse: d_x[i+half] = -d0 * s + d1 * c = d1 * c - d0 * s
                __m256 d1c = _mm256_mul_ps(d1, c);
                __m256 d0s = _mm256_mul_ps(d0, s);
                __m256 r1 = _mm256_sub_ps(d1c, d0s);
 
                _mm256_storeu_ps(&d_x_row[i], r0);
                _mm256_storeu_ps(&d_x_row[i + half_dim], r1);
            }
            for (; i < half_dim; ++i) {
                float d0 = d_out_row[i];
                float d1 = d_out_row[i + half_dim];
                float c = cos_row[i];
                float s = sin_row[i];
                d_x_row[i] = d0 * c + d1 * s;
                d_x_row[i + half_dim] = -d0 * s + d1 * c;
            }
 
#else
            for (int i = 0; i < half_dim; ++i) {
                float d0 = d_out_row[i];
                float d1 = d_out_row[i + half_dim];
                float c = cos_row[i];
                float s = sin_row[i];
 
                // Inverse rotation: rotate by -θ
                d_x_row[i] = d0 * c + d1 * s;
                d_x_row[i + half_dim] = -d0 * s + d1 * c;
            }
#endif
 
            for (int i = head_dim; i < aligned_head_dim; ++i) {
                d_x_row[i] = 0.0f;
            }
        }
    }
}

Referenced by rope_backward_bf16(), and rope_backward_qk().

◆ rope_backward_bf16()

void rope_backward_bf16	(	const uint16_t *	d_out,
		uint16_t *	d_x,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		float *	scratch_d_out,
		float *	scratch_d_x
	)

Definition at line 67 of file rope_kernels_bf16.c.

{
    if (!scratch_d_out || !scratch_d_x) return;
 
    size_t total = (size_t)num_heads * (size_t)num_tokens * (size_t)aligned_head_dim;
 
    bf16_tensor_to_float(d_out, scratch_d_out, total);
    rope_backward(scratch_d_out, scratch_d_x, cos_cache, sin_cache,
                  num_heads, num_tokens, head_dim, aligned_head_dim, pos_offset);
    float_tensor_to_bf16(scratch_d_x, d_x, total);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and rope_backward().

Referenced by rope_backward_qk_bf16().

◆ rope_backward_inplace()

void rope_backward_inplace	(	float *	d_x,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset
	)

RoPE backward in-place (overwrite with inverse rotation)

Test:: test_rope.py::TestRoPEBackward::test_rope_backward_inplace

In-place backward: overwrite d_out with inverse-rotated gradients. Useful when d_x == d_out is acceptable (saves memory).

After changes: make test

Definition at line 991 of file rope_kernels.c.

{
    size_t head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
    int half_dim = head_dim / 2;
 
    for (int h = 0; h < num_heads; ++h) {
        for (int t = 0; t < num_tokens; ++t) {
            int pos = pos_offset + t;
            const float *cos_row = cos_cache + pos * half_dim;
            const float *sin_row = sin_cache + pos * half_dim;
 
            float *d_row = d_x + h * head_stride + (size_t)t * (size_t)aligned_head_dim;
 
#if defined(__AVX512F__)
            int i = 0;
            for (; i + 16 <= half_dim; i += 16) {
                __m512 d0 = _mm512_loadu_ps(&d_row[i]);
                __m512 d1 = _mm512_loadu_ps(&d_row[i + half_dim]);
                __m512 c = _mm512_loadu_ps(&cos_row[i]);
                __m512 s = _mm512_loadu_ps(&sin_row[i]);
 
                __m512 r0 = _mm512_fmadd_ps(d0, c, _mm512_mul_ps(d1, s));
                __m512 r1 = _mm512_fmsub_ps(d1, c, _mm512_mul_ps(d0, s));
 
                _mm512_storeu_ps(&d_row[i], r0);
                _mm512_storeu_ps(&d_row[i + half_dim], r1);
            }
            for (; i < half_dim; ++i) {
                float d0 = d_row[i];
                float d1 = d_row[i + half_dim];
                float c = cos_row[i];
                float s = sin_row[i];
                d_row[i] = d0 * c + d1 * s;
                d_row[i + half_dim] = -d0 * s + d1 * c;
            }
 
#elif defined(__AVX__)
            int i = 0;
            for (; i + 8 <= half_dim; i += 8) {
                __m256 d0 = _mm256_loadu_ps(&d_row[i]);
                __m256 d1 = _mm256_loadu_ps(&d_row[i + half_dim]);
                __m256 c = _mm256_loadu_ps(&cos_row[i]);
                __m256 s = _mm256_loadu_ps(&sin_row[i]);
 
                __m256 d0c = _mm256_mul_ps(d0, c);
                __m256 d1s = _mm256_mul_ps(d1, s);
                __m256 r0 = _mm256_add_ps(d0c, d1s);
 
                __m256 d1c = _mm256_mul_ps(d1, c);
                __m256 d0s = _mm256_mul_ps(d0, s);
                __m256 r1 = _mm256_sub_ps(d1c, d0s);
 
                _mm256_storeu_ps(&d_row[i], r0);
                _mm256_storeu_ps(&d_row[i + half_dim], r1);
            }
            for (; i < half_dim; ++i) {
                float d0 = d_row[i];
                float d1 = d_row[i + half_dim];
                float c = cos_row[i];
                float s = sin_row[i];
                d_row[i] = d0 * c + d1 * s;
                d_row[i + half_dim] = -d0 * s + d1 * c;
            }
 
#else
            for (int i = 0; i < half_dim; ++i) {
                float d0 = d_row[i];
                float d1 = d_row[i + half_dim];
                float c = cos_row[i];
                float s = sin_row[i];
 
                // Inverse rotation: rotate by -θ
                d_row[i] = d0 * c + d1 * s;
                d_row[i + half_dim] = -d0 * s + d1 * c;
            }
#endif
 
            for (int i = head_dim; i < aligned_head_dim; ++i) {
                d_row[i] = 0.0f;
            }
        }
    }
}

◆ rope_backward_qk()

void rope_backward_qk	(	const float *	d_q_out,
		const float *	d_k_out,
		float *	d_q,
		float *	d_k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset
	)

RoPE backward for both dQ and dK

Test:: test_rope.py::TestRoPEBackward::test_rope_backward_qk

Combined RoPE backward for both dQ and dK gradients.

After changes: make test

Definition at line 2826 of file rope_kernels.c.

{
    rope_backward(d_q_out, d_q, cos_cache, sin_cache, num_heads, num_tokens, head_dim, aligned_head_dim, pos_offset);
    rope_backward(d_k_out, d_k, cos_cache, sin_cache, num_kv_heads, num_tokens, head_dim, aligned_head_dim, pos_offset);
}

References rope_backward().

Referenced by ck_layer_backward_rmsnorm_swiglu().

◆ rope_backward_qk_bf16()

void rope_backward_qk_bf16	(	const uint16_t *	d_q_out,
		const uint16_t *	d_k_out,
		uint16_t *	d_q,
		uint16_t *	d_k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		float *	scratch_dq_out,
		float *	scratch_dq,
		float *	scratch_dk_out,
		float *	scratch_dk
	)

Definition at line 139 of file rope_kernels_bf16.c.

{
    if (!d_q_out || !d_k_out || !d_q || !d_k) return;
 
    rope_backward_bf16(d_q_out, d_q, cos_cache, sin_cache,
                       num_heads, num_tokens, head_dim, aligned_head_dim, pos_offset,
                       scratch_dq_out, scratch_dq);
    rope_backward_bf16(d_k_out, d_k, cos_cache, sin_cache,
                       num_kv_heads, num_tokens, head_dim, aligned_head_dim, pos_offset,
                       scratch_dk_out, scratch_dk);
}

References rope_backward_bf16().

◆ rope_backward_qk_pairwise_with_rotary_dim()

void rope_backward_qk_pairwise_with_rotary_dim	(	const float *	d_q_out,
		const float *	d_k_out,
		float *	d_q,
		float *	d_k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim
	)

Definition at line 2843 of file rope_kernels.c.

{
    size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
    size_t k_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        rope_backward_apply_head_pairwise(
            d_q_out + (size_t)h * q_head_stride,
            d_q + (size_t)h * q_head_stride,
            cos_cache,
            sin_cache,
            num_tokens,
            head_dim,
            aligned_head_dim,
            pos_offset,
            rotary_dim
        );
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        rope_backward_apply_head_pairwise(
            d_k_out + (size_t)h * k_head_stride,
            d_k + (size_t)h * k_head_stride,
            cos_cache,
            sin_cache,
            num_tokens,
            head_dim,
            aligned_head_dim,
            pos_offset,
            rotary_dim
        );
    }
}

References rope_backward_apply_head_pairwise().

◆ rope_forward()

void rope_forward	(	float *	x,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset
	)

RoPE forward (head-major layout, in-place)

Test:

test_rope.py::TestRoPEForward::test_rope_forward

test_rope.py::TestRoPEForward::test_rope_vs_separate

test_parity.py::test_rope_parity

Applies rotary position embeddings in-place to Q or K tensor. x: [num_heads, num_tokens, head_dim] head-major

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

Definition at line 796 of file rope_kernels.c.

{
    rope_forward_with_rotary_dim(x, cos_cache, sin_cache, num_heads, num_tokens,
                                 head_dim, aligned_head_dim, pos_offset, head_dim);
}

References rope_forward_with_rotary_dim().

◆ rope_forward_bf16()

void rope_forward_bf16	(	uint16_t *	x,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		float *	scratch
	)

Definition at line 28 of file rope_kernels_bf16.c.

{
    rope_forward_bf16_with_rotary_dim(x, cos_cache, sin_cache, num_heads, num_tokens,
                                     head_dim, aligned_head_dim, pos_offset, head_dim, scratch);
}

References rope_forward_bf16_with_rotary_dim().

◆ rope_forward_bf16_with_rotary_dim()

void rope_forward_bf16_with_rotary_dim	(	uint16_t *	x,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim,
		float *	scratch
	)

Definition at line 42 of file rope_kernels_bf16.c.

{
    if (!scratch) return;
 
    size_t total = (size_t)num_heads * (size_t)num_tokens * (size_t)aligned_head_dim;
 
    bf16_tensor_to_float(x, scratch, total);
    rope_forward_with_rotary_dim(scratch, cos_cache, sin_cache,
                                num_heads, num_tokens, head_dim, aligned_head_dim, pos_offset, rotary_dim);
    float_tensor_to_bf16(scratch, x, total);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and rope_forward_with_rotary_dim().

Referenced by rope_forward_bf16(), and rope_forward_qk_bf16_with_rotary_dim().

◆ rope_forward_q_split_direct_f32()

void rope_forward_q_split_direct_f32	(	float *	q,
		const float *	freq_factors,
		int	use_freq_factors,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim,
		float	freq_base
	)

Definition at line 1301 of file rope_kernels.c.

{
    rope_forward_split_direct_one(q, freq_factors, use_freq_factors,
                                  num_heads, num_tokens, head_dim, aligned_head_dim,
                                  pos_offset, rotary_dim, freq_base);
}

References rope_forward_split_direct_one().

◆ rope_forward_qk()

void rope_forward_qk	(	float *	q,
		float *	k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset
	)

RoPE forward for both Q and K (common inference pattern)

Test:

test_rope.py::TestRoPEForward::test_rope_forward_qk

test_fused_attention_decode.py::TestFusedAttentionDecode::test_qk_rope

test_parity.py::test_rope_qk_parity

Combined RoPE forward for both Q and K in one call. q: [num_heads, num_tokens, head_dim] k: [num_kv_heads, num_tokens, head_dim]

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

Definition at line 1094 of file rope_kernels.c.

{
    rope_forward_qk_with_rotary_dim(q, k, cos_cache, sin_cache, num_heads, num_kv_heads,
                                    num_tokens, head_dim, aligned_head_dim, pos_offset, head_dim);
}

References rope_forward_qk_with_rotary_dim().

◆ rope_forward_qk_bf16()

void rope_forward_qk_bf16	(	uint16_t *	q,
		uint16_t *	k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		float *	scratch_q,
		float *	scratch_k
	)

Definition at line 94 of file rope_kernels_bf16.c.

{
    rope_forward_qk_bf16_with_rotary_dim(q, k, cos_cache, sin_cache, num_heads, num_kv_heads,
                                        num_tokens, head_dim, aligned_head_dim, pos_offset,
                                        head_dim, scratch_q, scratch_k);
}

References rope_forward_qk_bf16_with_rotary_dim().

◆ rope_forward_qk_bf16_with_rotary_dim()

void rope_forward_qk_bf16_with_rotary_dim	(	uint16_t *	q,
		uint16_t *	k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim,
		float *	scratch_q,
		float *	scratch_k
	)

Definition at line 112 of file rope_kernels_bf16.c.

{
    if (!q || !k) return;
 
    rope_forward_bf16_with_rotary_dim(q, cos_cache, sin_cache,
                                      num_heads, num_tokens, head_dim, aligned_head_dim, pos_offset,
                                      rotary_dim, scratch_q);
    rope_forward_bf16_with_rotary_dim(k, cos_cache, sin_cache,
                                      num_kv_heads, num_tokens, head_dim, aligned_head_dim, pos_offset,
                                      rotary_dim, scratch_k);
}

References rope_forward_bf16_with_rotary_dim().

Referenced by rope_forward_qk_bf16().

◆ rope_forward_qk_gemma4_direct()

void rope_forward_qk_gemma4_direct	(	float *	q,
		float *	k,
		const float *	freq_factors,
		int	use_freq_factors,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim,
		float	freq_base
	)

Definition at line 1317 of file rope_kernels.c.

{
    rope_forward_qk_split_direct_f32(q, k, freq_factors, use_freq_factors,
                                    num_heads, num_kv_heads, num_tokens,
                                    head_dim, aligned_head_dim, pos_offset,
                                    rotary_dim, freq_base);
}

References rope_forward_qk_split_direct_f32().

◆ rope_forward_qk_gemma4v_vision_xy()

void rope_forward_qk_gemma4v_vision_xy	(	float *	q,
		float *	k,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	grid_w,
		int	rotary_dim,
		float	freq_base
	)

Definition at line 1404 of file rope_kernels.c.

{
    rope_forward_gemma4v_vision_xy_one(q, num_heads, num_tokens, head_dim, aligned_head_dim, grid_w, rotary_dim, freq_base);
    rope_forward_gemma4v_vision_xy_one(k, num_kv_heads, num_tokens, head_dim, aligned_head_dim, grid_w, rotary_dim, freq_base);
}

References rope_forward_gemma4v_vision_xy_one().

◆ rope_forward_qk_pairwise_llama_cpu()

void rope_forward_qk_pairwise_llama_cpu	(	float *	q,
		float *	k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim
	)

Definition at line 1540 of file rope_kernels.c.

{
    if (!q || !k || !cos_cache || !sin_cache || num_tokens <= 0) {
        return;
    }
    if (rotary_dim <= 0 || rotary_dim > head_dim) {
        rotary_dim = head_dim;
    }
    const int cache_half = rotary_dim / 2;
    const size_t head_stride =
        (size_t)num_tokens * (size_t)aligned_head_dim;
    if (num_tokens == 1) {
        const float *cos_row =
            cos_cache + (size_t)pos_offset * (size_t)cache_half;
        const float *sin_row =
            sin_cache + (size_t)pos_offset * (size_t)cache_half;
        rope_apply_decode_pairwise_llama_cpu(
            q, cos_row, sin_row, num_heads, aligned_head_dim, rotary_dim);
        rope_apply_decode_pairwise_llama_cpu(
            k, cos_row, sin_row, num_kv_heads, aligned_head_dim, rotary_dim);
        return;
    }
    for (int token = 0; token < num_tokens; ++token) {
        const int pos = pos_offset + token;
        const float *cos_row =
            cos_cache + (size_t)pos * (size_t)cache_half;
        const float *sin_row =
            sin_cache + (size_t)pos * (size_t)cache_half;
        for (int head = 0; head < num_heads; ++head) {
            rope_apply_decode_pairwise_llama_cpu(
                q + (size_t)head * head_stride
                  + (size_t)token * (size_t)aligned_head_dim,
                cos_row,
                sin_row,
                1,
                aligned_head_dim,
                rotary_dim);
        }
        for (int head = 0; head < num_kv_heads; ++head) {
            rope_apply_decode_pairwise_llama_cpu(
                k + (size_t)head * head_stride
                  + (size_t)token * (size_t)aligned_head_dim,
                cos_row,
                sin_row,
                1,
                aligned_head_dim,
                rotary_dim);
        }
    }
}

References rope_apply_decode_pairwise_llama_cpu(), and token.

◆ rope_forward_qk_pairwise_with_rotary_dim()

void rope_forward_qk_pairwise_with_rotary_dim	(	float *	q,
		float *	k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim
	)

Definition at line 1483 of file rope_kernels.c.

{
    size_t q_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
    size_t k_head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        rope_apply_head_pairwise(q + (size_t) h * q_head_stride,
                                 cos_cache, sin_cache,
                                 num_tokens, head_dim, aligned_head_dim, pos_offset, rotary_dim);
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        rope_apply_head_pairwise(k + (size_t) h * k_head_stride,
                                 cos_cache, sin_cache,
                                 num_tokens, head_dim, aligned_head_dim, pos_offset, rotary_dim);
    }
}

References rope_apply_head_pairwise().

◆ rope_forward_qk_split_direct_f32()

void rope_forward_qk_split_direct_f32	(	float *	q,
		float *	k,
		const float *	freq_factors,
		int	use_freq_factors,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim,
		float	freq_base
	)

Definition at line 1180 of file rope_kernels.c.

{
    rope_forward_split_direct_one(q, freq_factors, use_freq_factors,
                                  num_heads, num_tokens, head_dim, aligned_head_dim,
                                  pos_offset, rotary_dim, freq_base);
    rope_forward_split_direct_one(k, freq_factors, use_freq_factors,
                                  num_kv_heads, num_tokens, head_dim, aligned_head_dim,
                                  pos_offset, rotary_dim, freq_base);
}

References rope_forward_split_direct_one().

Referenced by rope_forward_qk_gemma4_direct().

◆ rope_forward_qk_split_direct_token_range_f32()

void rope_forward_qk_split_direct_token_range_f32	(	float *	q,
		float *	k,
		const float *	freq_factors,
		int	use_freq_factors,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim,
		float	freq_base,
		int	token_begin,
		int	token_end
	)

Definition at line 1201 of file rope_kernels.c.

{
    if ((!q && !k) || num_tokens <= 0 || head_dim <= 0 ||
        aligned_head_dim <= 0 || token_begin < 0 ||
        token_begin >= token_end || token_end > num_tokens) {
        return;
    }
    if (rotary_dim <= 0 || rotary_dim > head_dim) {
        rotary_dim = head_dim;
    }
    if (freq_base <= 0.0f) {
        freq_base = 10000.0f;
    }
 
    const int rotary_half = rotary_dim / 2;
    const size_t head_stride =
        (size_t)num_tokens * (size_t)aligned_head_dim;
    const float theta_scale = powf(freq_base, -2.0f / (float)rotary_dim);
 
    for (int t = token_begin; t < token_end; ++t) {
        const float pos = (float)(pos_offset + t);
        for (int i = 0; i < rotary_half; ++i) {
            const int idx0 = i;
            const int idx1 = i + rotary_half;
            const float ff =
                (use_freq_factors && freq_factors) ? freq_factors[i] : 1.0f;
            const float theta = pos * powf(theta_scale, (float)i) / ff;
            const float c = cosf(theta);
            const float sv = sinf(theta);
 
            for (int h = 0; q && h < num_heads; ++h) {
                float *row = q + (size_t)h * head_stride +
                             (size_t)t * (size_t)aligned_head_dim;
                const float x0 = row[idx0];
                const float x1 = row[idx1];
                row[idx0] = x0 * c - x1 * sv;
                row[idx1] = x1 * c + x0 * sv;
            }
            for (int h = 0; k && h < num_kv_heads; ++h) {
                float *row = k + (size_t)h * head_stride +
                             (size_t)t * (size_t)aligned_head_dim;
                const float x0 = row[idx0];
                const float x1 = row[idx1];
                row[idx0] = x0 * c - x1 * sv;
                row[idx1] = x1 * c + x0 * sv;
            }
        }
    }
}

◆ rope_forward_qk_split_llama_token_range_f32()

void rope_forward_qk_split_llama_token_range_f32	(	float *	q,
		float *	k,
		const float *	freq_factors,
		int	use_freq_factors,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim,
		float	freq_base,
		int	token_begin,
		int	token_end
	)

Definition at line 1265 of file rope_kernels.c.

{
    if ((!q && !k) || num_tokens <= 0 || head_dim <= 0 ||
        aligned_head_dim < head_dim || token_begin < 0 ||
        token_begin >= token_end || token_end > num_tokens) return;
    if (rotary_dim <= 0 || rotary_dim > head_dim) rotary_dim = head_dim;
    if (freq_base <= 0.0f) freq_base = 10000.0f;
    const int half = rotary_dim / 2;
    const size_t stride = (size_t)num_tokens * aligned_head_dim;
    const float scale = ck_rope_reference_powf(freq_base, -2.0f / rotary_dim);
    for (int t = token_begin; t < token_end; ++t) {
        /* The rounded recurrence is part of the ggml CPU numerical contract. */
        volatile float theta = (float)(pos_offset + t);
        for (int i = 0; i < half; ++i) {
            const float ff = use_freq_factors && freq_factors ? freq_factors[i] : 1.0f;
            const float angle = theta / ff;
            const float c = ck_rope_reference_cosf(angle);
            const float s = ck_rope_reference_sinf(angle);
            for (int h = 0; h < num_heads + num_kv_heads; ++h) {
                float *base = h < num_heads ? q : k;
                if (!base) continue;
                const int head = h < num_heads ? h : h - num_heads;
                float *row = base + (size_t)head * stride + (size_t)t * aligned_head_dim;
                const float x0 = row[i], x1 = row[i + half];
                row[i] = fmaf(x0, c, -(x1 * s));
                row[i + half] = fmaf(x0, s, x1 * c);
            }
            theta *= scale;
        }
    }
}

References ck_rope_reference_cosf(), ck_rope_reference_powf(), and ck_rope_reference_sinf().

◆ rope_forward_qk_strided()

void rope_forward_qk_strided	(	float *	q,
		float *	k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	q_stride_tokens,
		int	k_stride_tokens
	)

RoPE forward for both Q and K with custom strides (KV cache layouts)

Test:

test_rope.py::TestRoPEForward::test_rope_forward_qk_strided

test_kv_cache_attention.py::TestKVCacheAttention::test_qk_rope_strided

Combined QK RoPE with configurable strides for KV cache layouts.

After changes: make test

Definition at line 2778 of file rope_kernels.c.

{
    rope_forward_qk_strided_with_rotary_dim(q, k, cos_cache, sin_cache, num_heads, num_kv_heads,
                                           num_tokens, head_dim, aligned_head_dim, pos_offset,
                                           q_stride_tokens, k_stride_tokens, head_dim);
}

References rope_forward_qk_strided_with_rotary_dim().

◆ rope_forward_qk_strided_with_rotary_dim()

void rope_forward_qk_strided_with_rotary_dim	(	float *	q,
		float *	k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	q_stride_tokens,
		int	k_stride_tokens,
		int	rotary_dim
	)

Definition at line 2796 of file rope_kernels.c.

{
    rope_forward_strided_with_rotary_dim(q, cos_cache, sin_cache, num_heads, num_tokens,
                                       head_dim, aligned_head_dim, pos_offset,
                                       q_stride_tokens, rotary_dim);
    rope_forward_strided_with_rotary_dim(k, cos_cache, sin_cache, num_kv_heads, num_tokens,
                                       head_dim, aligned_head_dim, pos_offset,
                                       k_stride_tokens, rotary_dim);
}

References rope_forward_strided_with_rotary_dim().

Referenced by rope_forward_qk_strided().

◆ rope_forward_qk_with_rotary_dim()

void rope_forward_qk_with_rotary_dim	(	float *	q,
		float *	k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim
	)

Definition at line 1109 of file rope_kernels.c.

{
    rope_forward_with_rotary_dim(q, cos_cache, sin_cache, num_heads, num_tokens,
                                head_dim, aligned_head_dim, pos_offset, rotary_dim);
    rope_forward_with_rotary_dim(k, cos_cache, sin_cache, num_kv_heads, num_tokens,
                                head_dim, aligned_head_dim, pos_offset, rotary_dim);
}

References rope_forward_with_rotary_dim().

Referenced by rope_forward_qk().

◆ rope_forward_qk_with_rotary_dim_cache_stride()

void rope_forward_qk_with_rotary_dim_cache_stride	(	float *	q,
		float *	k,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim,
		int	cache_rotary_dim
	)

Definition at line 1419 of file rope_kernels.c.

{
    if (rotary_dim <= 0 || rotary_dim > head_dim) {
        rotary_dim = head_dim;
    }
    if (cache_rotary_dim < rotary_dim) {
        cache_rotary_dim = rotary_dim;
    }
    const int rotary_half = rotary_dim / 2;
    const int cache_half = cache_rotary_dim / 2;
    const size_t head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        float *head = q + (size_t)h * head_stride;
        for (int t = 0; t < num_tokens; ++t) {
            const int pos = pos_offset + t;
            const float *cos_row = cos_cache + (size_t)pos * (size_t)cache_half;
            const float *sin_row = sin_cache + (size_t)pos * (size_t)cache_half;
            float *x_row = head + (size_t)t * (size_t)aligned_head_dim;
            for (int i = 0; i < rotary_half; ++i) {
                const int idx0 = 2 * i;
                const int idx1 = idx0 + 1;
                const float x0 = x_row[idx0];
                const float x1 = x_row[idx1];
                const float c = cos_row[i];
                const float sv = sin_row[i];
                x_row[idx0] = x0 * c - x1 * sv;
                x_row[idx1] = x0 * sv + x1 * c;
            }
        }
    }
 
    for (int h = 0; h < num_kv_heads; ++h) {
        float *head = k + (size_t)h * head_stride;
        for (int t = 0; t < num_tokens; ++t) {
            const int pos = pos_offset + t;
            const float *cos_row = cos_cache + (size_t)pos * (size_t)cache_half;
            const float *sin_row = sin_cache + (size_t)pos * (size_t)cache_half;
            float *x_row = head + (size_t)t * (size_t)aligned_head_dim;
            for (int i = 0; i < rotary_half; ++i) {
                const int idx0 = 2 * i;
                const int idx1 = idx0 + 1;
                const float x0 = x_row[idx0];
                const float x1 = x_row[idx1];
                const float c = cos_row[i];
                const float sv = sin_row[i];
                x_row[idx0] = x0 * c - x1 * sv;
                x_row[idx1] = x0 * sv + x1 * c;
            }
        }
    }
}

◆ rope_forward_strided()

void rope_forward_strided	(	float *	x,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	head_stride_tokens
	)

RoPE forward with custom head stride (for KV cache layouts)

Test:

test_rope.py::TestRoPEForward::test_rope_strided

test_kv_cache_attention.py::TestKVCacheAttention::test_rope_decode

Variant with configurable head_stride_tokens for non-contiguous head layouts.

After changes: make test

Definition at line 837 of file rope_kernels.c.

{
    rope_forward_strided_with_rotary_dim(x, cos_cache, sin_cache, num_heads, num_tokens,
                                         head_dim, aligned_head_dim, pos_offset,
                                         head_stride_tokens, head_dim);
}

References rope_forward_strided_with_rotary_dim().

◆ rope_forward_strided_with_rotary_dim()

void rope_forward_strided_with_rotary_dim	(	float *	x,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	head_stride_tokens,
		int	rotary_dim
	)

Definition at line 852 of file rope_kernels.c.

{
    size_t head_stride = (size_t)head_stride_tokens * (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        rope_apply_head(x + h * head_stride,
                        cos_cache, sin_cache,
                        num_tokens, head_dim, aligned_head_dim, pos_offset, rotary_dim);
    }
}

References rope_apply_head().

Referenced by rope_forward_qk_strided_with_rotary_dim(), and rope_forward_strided().

◆ rope_forward_with_rotary_dim()

void rope_forward_with_rotary_dim	(	float *	x,
		const float *	cos_cache,
		const float *	sin_cache,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		int	aligned_head_dim,
		int	pos_offset,
		int	rotary_dim
	)

Definition at line 809 of file rope_kernels.c.

{
    size_t head_stride = (size_t)num_tokens * (size_t)aligned_head_dim;
 
    for (int h = 0; h < num_heads; ++h) {
        rope_apply_head(x + h * head_stride,
                        cos_cache, sin_cache,
                        num_tokens, head_dim, aligned_head_dim, pos_offset, rotary_dim);
    }
}

References rope_apply_head().

Referenced by rope_forward(), rope_forward_bf16_with_rotary_dim(), and rope_forward_qk_with_rotary_dim().

◆ rope_precompute_cache()

void rope_precompute_cache	(	float *	cos_cache,
		float *	sin_cache,
		int	max_seq_len,
		int	head_dim,
		float	base,
		int	rotary_dim,
		const char *	scaling_type,
		float	scaling_factor
	)

Precompute RoPE cos/sin cache with rotary_dim and scaling support

Test:

test_rope.py::TestRoPECache::test_cache_computation

test_rope.py::TestRoPECache::test_cache_values

Precomputes cos(m * theta_i) and sin(m * theta_i) for positions 0..max_seq_len-1. Only computes for first rotary_dim channels; remaining head_dim - rotary_dim channels are NOT rotated (pass through unchanged).

Scaling types:

"none": Standard RoPE
"linear": Scale positions by 1/scaling_factor
"dynamic": NTK-aware dynamic scaling
"yarn": YaRN scaling (beta-based)

Parameters

cos_cache	Output: [max_seq_len, rotary_dim/2] cos values
sin_cache	Output: [max_seq_len, rotary_dim/2] sin values
max_seq_len	Maximum sequence length for cache
head_dim	Full head dimension (for frequency computation)
base	RoPE base frequency (theta)
rotary_dim	Number of dimensions to rotate (0 = use head_dim)
scaling_type	Scaling type string: "none", "linear", "dynamic", "yarn"
scaling_factor	Scaling factor (1.0 = no scaling)

After changes: make test

Definition at line 384 of file rope_kernels.c.

{
    // Use rotary_dim = head_dim if not specified
    if (rotary_dim <= 0 || rotary_dim > head_dim) {
        rotary_dim = head_dim;
    }
 
    // Use no scaling if not specified
    int is_linear_scaling = 0;
    if (scaling_type != NULL && strcmp(scaling_type, "linear") == 0 && scaling_factor > 0.0f && scaling_factor != 1.0f) {
        is_linear_scaling = 1;
    }
 
    int rotary_half = rotary_dim / 2;
 
    for (int pos = 0; pos < max_seq_len; ++pos) {
        // Apply linear scaling to position if needed
        float effective_pos = (float)pos;
        if (is_linear_scaling) {
            effective_pos = (float)pos / scaling_factor;
        }
 
        for (int i = 0; i < rotary_half; ++i) {
            // Match the FP32 reference contract directly. Computing this via
            // long-double log/exp makes the final float depend on the host
            // libm and long-double ABI.
            const float exponent = ((float)(2 * i)) / (float)rotary_dim;
            const float freq_f = 1.0f / powf(base, exponent);
            float angle_f = effective_pos * freq_f;
            cos_cache[pos * rotary_half + i] = cosf(angle_f);
            sin_cache[pos * rotary_half + i] = sinf(angle_f);
        }
    }
}

◆ rope_precompute_cache_llama_cpu()

void rope_precompute_cache_llama_cpu	(	float *	cos_cache,
		float *	sin_cache,
		int	max_seq_len,
		int	head_dim,
		float	base,
		int	rotary_dim,
		const char *	scaling_type,
		float	scaling_factor
	)

Definition at line 560 of file rope_kernels.c.

{
    if (!cos_cache || !sin_cache || max_seq_len <= 0 || head_dim <= 0) {
        return;
    }
    if (rotary_dim <= 0 || rotary_dim > head_dim) {
        rotary_dim = head_dim;
    }
    if (base <= 0.0f) {
        base = 10000.0f;
    }
 
    const int linear_scaling =
        scaling_type != NULL &&
        strcmp(scaling_type, "linear") == 0 &&
        scaling_factor > 0.0f &&
        scaling_factor != 1.0f;
    const int rotary_half = rotary_dim / 2;
    const float theta_scale =
        ck_rope_reference_powf(base, -2.0f / (float)rotary_dim);
 
    for (int pos = 0; pos < max_seq_len; ++pos) {
        float theta = (float)pos;
        if (linear_scaling) {
            theta /= scaling_factor;
        }
        for (int i = 0; i < rotary_half; ++i) {
            cos_cache[(size_t)pos * (size_t)rotary_half + (size_t)i] =
                ck_rope_reference_cosf(theta);
            sin_cache[(size_t)pos * (size_t)rotary_half + (size_t)i] =
                ck_rope_reference_sinf(theta);
            theta *= theta_scale;
        }
    }
}

References ck_rope_reference_cosf(), ck_rope_reference_powf(), and ck_rope_reference_sinf().

◆ rope_precompute_cache_split()

void rope_precompute_cache_split	(	float *	cos_cache,
		float *	sin_cache,
		int	max_seq_len,
		int	head_dim,
		float	base
	)

Precompute RoPE cos/sin cache (split layout: head_dim/2) Legacy layout used before rotary_dim/scaling support.

Parameters

cos_cache	Output: [max_seq_len, head_dim/2] cos values
sin_cache	Output: [max_seq_len, head_dim/2] sin values
max_seq_len	Maximum sequence length for cache
head_dim	Full head dimension
base	RoPE base frequency (theta)

Definition at line 340 of file rope_kernels.c.

{
    int half_dim = head_dim / 2;
    for (int pos = 0; pos < max_seq_len; ++pos) {
        for (int i = 0; i < half_dim; ++i) {
            const float exponent = ((float)(2 * i)) / (float)head_dim;
            const float freq_f = 1.0f / powf(base, exponent);
            const float angle_f = (float)pos * freq_f;
            cos_cache[pos * half_dim + i] = cosf(angle_f);
            sin_cache[pos * half_dim + i] = sinf(angle_f);
        }
    }
}

◆ rowwise_bias_add()

void rowwise_bias_add	(	float *	x,
		const float *	bias,
		int	rows,
		int	dim
	)

Definition at line 696 of file vision_kernels.c.

{
    if (!x || !bias || rows <= 0 || dim <= 0) {
        return;
    }
 
    for (int r = 0; r < rows; ++r) {
        float *row = x + ((size_t) r * (size_t) dim);
        for (int c = 0; c < dim; ++c) {
            row[c] += bias[c];
        }
    }
}

◆ scal_copy_f32()

void scal_copy_f32	(	float *	y,
		const float *	x,
		float	alpha,
		int	n
	)

Scaled copy: y = alpha * x.

Parameters

y	Output vector [n]
x	Input vector [n]
alpha	Scalar multiplier
n	Vector length

Definition at line 156 of file axpy_kernels.c.

{
    if (!y || !x || n <= 0) {
        return;
    }
 
    int i = 0;
 
#ifdef __AVX512F__
    __m512 valpha = _mm512_set1_ps(alpha);
    for (; i + 16 <= n; i += 16) {
        __m512 vx = _mm512_loadu_ps(&x[i]);
        __m512 vy = _mm512_mul_ps(vx, valpha);
        _mm512_storeu_ps(&y[i], vy);
    }
#endif
 
#ifdef __AVX2__
    __m256 valpha256 = _mm256_set1_ps(alpha);
    for (; i + 8 <= n; i += 8) {
        __m256 vx = _mm256_loadu_ps(&x[i]);
        __m256 vy = _mm256_mul_ps(vx, valpha256);
        _mm256_storeu_ps(&y[i], vy);
    }
#endif
 
    for (; i < n; i++) {
        y[i] = alpha * x[i];
    }
}

Referenced by weighted_sum_f32().

◆ sigmoid_backward()

void sigmoid_backward	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		size_t	n
	)

Definition at line 138 of file sigmoid_kernels.c.

{
#if defined(__AVX512F__)
    sigmoid_backward_avx512(input, d_output, d_input, n);
#else
    for (size_t i = 0; i < n; ++i) {
        float x = input[i];
        float s = sigmoid_scalar(x);
        float s_prime = s * (1.0f - s);
        d_input[i] = d_output[i] * s_prime;
    }
#endif
}

References sigmoid_scalar().

Referenced by sigmoid_backward_bf16().

◆ sigmoid_backward_bf16()

void sigmoid_backward_bf16	(	const uint16_t *	input,
		const uint16_t *	d_output,
		uint16_t *	d_input,
		size_t	n,
		float *	scratch_input,
		float *	scratch_d_output,
		float *	scratch_d_input
	)

Definition at line 45 of file sigmoid_kernels_bf16.c.

{
    if (!input || !d_output || !d_input || n == 0) return;
    if (!scratch_input || !scratch_d_output || !scratch_d_input) return;
 
    bf16_tensor_to_float(input, scratch_input, n);
    bf16_tensor_to_float(d_output, scratch_d_output, n);
    sigmoid_backward(scratch_input, scratch_d_output, scratch_d_input, n);
    float_tensor_to_bf16(scratch_d_input, d_input, n);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and sigmoid_backward().

◆ sigmoid_forward()

void sigmoid_forward	(	const float *	input,
		float *	output,
		size_t	n
	)

Definition at line 122 of file sigmoid_kernels.c.

{
#if defined(__AVX512F__)
    sigmoid_forward_avx512(input, output, n);
#else
    for (size_t i = 0; i < n; ++i) {
        output[i] = sigmoid_scalar(input[i]);
    }
#endif
}

References sigmoid_scalar().

Referenced by sigmoid_forward_bf16().

◆ sigmoid_forward_bf16()

void sigmoid_forward_bf16	(	const uint16_t *	input,
		uint16_t *	output,
		size_t	n,
		float *	scratch_input,
		float *	scratch_output
	)

Definition at line 27 of file sigmoid_kernels_bf16.c.

{
    if (!input || !output || n == 0) return;
    if (!scratch_input || !scratch_output) return;
 
    bf16_tensor_to_float(input, scratch_input, n);
    sigmoid_forward(scratch_input, scratch_output, n);
    float_tensor_to_bf16(scratch_output, output, n);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and sigmoid_forward().

◆ sigmoid_scalar()

float sigmoid_scalar ( float x )

Definition at line 26 of file sigmoid_kernels.c.

{
    return 1.0f / (1.0f + expf(-x));
}

Referenced by sigmoid_backward(), sigmoid_forward(), swiglu_backward(), swiglu_backward_bf16(), swiglu_forward(), swiglu_forward_bf16(), and swiglu_forward_q8_k().

◆ softmax_cross_entropy_loss()

void softmax_cross_entropy_loss	(	const float *	logits,
		const int32_t *	targets,
		int	tokens,
		int	vocab_size,
		float *	d_logits,
		float *	loss_out
	)

Definition at line 259 of file loss_kernels.c.

{
    if (!logits || !targets || !d_logits || tokens <= 0 || vocab_size <= 0) {
        if (loss_out) {
            *loss_out = 0.0f;
        }
        return;
    }
 
    if (ce_legacy_mode_enabled() && ce_targets_all_valid_no_ignore(targets, tokens, vocab_size)) {
        softmax_cross_entropy_loss_legacy_mean_impl(
            logits, targets, tokens, vocab_size, d_logits, loss_out);
        return;
    }
 
    softmax_cross_entropy_loss_index_mean_impl(
        logits, targets, tokens, vocab_size, d_logits, loss_out, -1);
}

References ce_legacy_mode_enabled(), ce_targets_all_valid_no_ignore(), softmax_cross_entropy_loss_index_mean_impl(), softmax_cross_entropy_loss_legacy_mean_impl(), and vocab_size.

Referenced by softmax_cross_entropy_loss_bf16().

◆ softmax_cross_entropy_loss_bf16()

void softmax_cross_entropy_loss_bf16	(	const uint16_t *	logits,
		const int32_t *	targets,
		int	tokens,
		int	vocab_size,
		uint16_t *	d_logits,
		float *	loss_out,
		float *	scratch_logits,
		float *	scratch_d_logits
	)

Definition at line 25 of file loss_kernels_bf16.c.

{
    if (!logits || !targets || !d_logits || tokens <= 0 || vocab_size <= 0) {
        if (loss_out) *loss_out = 0.0f;
        return;
    }
    if (!scratch_logits || !scratch_d_logits) {
        if (loss_out) *loss_out = 0.0f;
        return;
    }
 
    const size_t count = (size_t)tokens * (size_t)vocab_size;
 
    bf16_tensor_to_float(logits, scratch_logits, count);
    softmax_cross_entropy_loss(scratch_logits, targets, tokens, vocab_size, scratch_d_logits, loss_out);
    float_tensor_to_bf16(scratch_d_logits, d_logits, count);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), softmax_cross_entropy_loss(), and vocab_size.

◆ softmax_cross_entropy_loss_ptref()

void softmax_cross_entropy_loss_ptref	(	const float *	logits,
		const int32_t *	targets,
		int	tokens,
		int	vocab_size,
		float *	d_logits,
		float *	loss_out
	)

Definition at line 292 of file loss_kernels.c.

{
    /*
     * Keep a strict reference variant for parity experiments:
     * - always uses strict math path,
     * - same reduction / ignore semantics as default kernel.
     */
    softmax_cross_entropy_loss_index_mean_impl(
        logits, targets, tokens, vocab_size, d_logits, loss_out, 1);
}

References softmax_cross_entropy_loss_index_mean_impl(), and vocab_size.

◆ spatial_average_pool_contiguous()

void spatial_average_pool_contiguous	(	const float *	input,
		float *	output,
		int	grid_h,
		int	grid_w,
		int	embed_dim,
		int	merge_size
	)

Definition at line 657 of file vision_kernels.c.

{
    if (input == NULL || output == NULL || grid_h <= 0 || grid_w <= 0 || embed_dim <= 0 || merge_size <= 0) {
        return;
    }
 
    const int out_h = grid_h / merge_size;
    const int out_w = grid_w / merge_size;
    if (out_h <= 0 || out_w <= 0) {
        return;
    }
 
    const float inv_area = 1.0f / (float)(merge_size * merge_size);
    for (int oy = 0; oy < out_h; ++oy) {
        for (int ox = 0; ox < out_w; ++ox) {
            float *dst = output + ((size_t)oy * (size_t)out_w + (size_t)ox) * (size_t)embed_dim;
            memset(dst, 0, (size_t)embed_dim * sizeof(float));
            for (int dy = 0; dy < merge_size; ++dy) {
                const int iy = oy * merge_size + dy;
                for (int dx = 0; dx < merge_size; ++dx) {
                    const int ix = ox * merge_size + dx;
                    const float *src = input + ((size_t)iy * (size_t)grid_w + (size_t)ix) * (size_t)embed_dim;
                    for (int c = 0; c < embed_dim; ++c) {
                        dst[c] += src[c];
                    }
                }
            }
            for (int c = 0; c < embed_dim; ++c) {
                dst[c] *= inv_area;
            }
        }
    }
}

◆ spatial_merge_2x2()

void spatial_merge_2x2	(	const float *	input,
		float *	output,
		int	grid_h,
		int	grid_w,
		int	embed_dim
	)

Merge 2x2 neighboring tokens into a single wider token.

Input layout: [grid_h * grid_w, embed_dim] Output layout: [(grid_h/2) * (grid_w/2), embed_dim * 4]

Pack order within each merged token: top-left, top-right, bottom-left, bottom-right

Definition at line 577 of file vision_kernels.c.

{
    if (input == NULL || output == NULL || grid_h <= 0 || grid_w <= 0 || embed_dim <= 0) {
        return;
    }
 
    const int merged_h = grid_h / 2;
    const int merged_w = grid_w / 2;
    const size_t token_stride = (size_t)embed_dim;
    const size_t merged_stride = (size_t)embed_dim * 4;
 
    for (int mh = 0; mh < merged_h; ++mh) {
        for (int mw = 0; mw < merged_w; ++mw) {
            const int y0 = mh * 2;
            const int x0 = mw * 2;
            const int in_idx00 = y0 * grid_w + x0;
            const int in_idx01 = in_idx00 + 1;
            const int in_idx10 = in_idx00 + grid_w;
            const int in_idx11 = in_idx10 + 1;
            const int out_idx = mh * merged_w + mw;
 
            const float *src00 = input + (size_t)in_idx00 * token_stride;
            const float *src01 = input + (size_t)in_idx01 * token_stride;
            const float *src10 = input + (size_t)in_idx10 * token_stride;
            const float *src11 = input + (size_t)in_idx11 * token_stride;
            float *dst = output + (size_t)out_idx * merged_stride;
 
            memcpy(dst + 0 * embed_dim, src00, (size_t)embed_dim * sizeof(float));
            memcpy(dst + 1 * embed_dim, src01, (size_t)embed_dim * sizeof(float));
            memcpy(dst + 2 * embed_dim, src10, (size_t)embed_dim * sizeof(float));
            memcpy(dst + 3 * embed_dim, src11, (size_t)embed_dim * sizeof(float));
        }
    }
}

◆ spatial_merge_contiguous_tiled()

void spatial_merge_contiguous_tiled	(	const float *	input,
		float *	output,
		int	grid_h,
		int	grid_w,
		int	embed_dim,
		int	merge_size
	)

Definition at line 616 of file vision_kernels.c.

{
    if (input == NULL || output == NULL || grid_h <= 0 || grid_w <= 0 || embed_dim <= 0 || merge_size <= 0) {
        return;
    }
 
    const size_t num_tokens = (size_t) grid_h * (size_t) grid_w;
    const size_t merge_factor = (size_t) merge_size * (size_t) merge_size;
    const size_t merged_tokens = num_tokens / merge_factor;
    memcpy(output, input, merged_tokens * (size_t) embed_dim * merge_factor * sizeof(float));
}

◆ speculative_commit_one_i32()

void speculative_commit_one_i32	(	int	accepted,
		int	verified_token,
		int *	token_buffer,
		int *	token_count,
		int	max_tokens,
		int *	target_position,
		int *	draft_position,
		int *	accepted_count,
		int *	rejected_count
	)

Commit one verified speculative token and update decode counters.

This is the minimal state transition for the first Gemma4 assistant bridge: greedy, one draft token, target remains authoritative. For this milestone the draft cache is kept synchronized with the target position after each token. Multi-token speculative decoding can later replace this with prefix accept and partial draft-cache rollback.

Definition at line 611 of file topk_kernels.c.

{
    int next_count = token_count ? *token_count : 0;
    if (token_buffer && token_count && next_count >= 0 && next_count < max_tokens) {
        token_buffer[next_count] = verified_token;
        next_count += 1;
        *token_count = next_count;
    }
 
    if (target_position) {
        *target_position += 1;
        if (draft_position) {
            *draft_position = *target_position;
        }
    } else if (draft_position) {
        *draft_position += 1;
    }
 
    if (accepted) {
        if (accepted_count) {
            *accepted_count += 1;
        }
    } else {
        if (rejected_count) {
            *rejected_count += 1;
        }
    }
}

◆ speculative_verify_greedy_f32()

void speculative_verify_greedy_f32	(	const float *	target_logits,
		int	vocab_size,
		int	draft_token,
		int *	accepted,
		int *	verified_token
	)

Greedy one-token speculative verification.

The draft model proposes draft_token. The target model is authoritative: if draft_token equals argmax(target_logits), the candidate is accepted and emitted. Otherwise the target argmax is emitted and the draft path must be reset or rewound by the runtime loop.

Parameters

target_logits	Target/backbone logits [vocab_size]
vocab_size	Number of logits
draft_token	Candidate token from draft/assistant model
accepted	Output scalar: 1 if accepted, 0 otherwise
verified_token	Output scalar: accepted draft token or target argmax

Definition at line 585 of file topk_kernels.c.

{
    const int target_token = argmax_f32(target_logits, vocab_size);
    const int ok = (target_token >= 0 && draft_token == target_token) ? 1 : 0;
 
    if (accepted) {
        *accepted = ok;
    }
    if (verified_token) {
        *verified_token = ok ? draft_token : target_token;
    }
}

References argmax_f32(), and vocab_size.

◆ split_q_gate_backward()

void split_q_gate_backward	(	const float *	d_q,
		const float *	d_gate,
		float *	d_packed_qg,
		int	rows,
		int	q_dim,
		int	gate_dim,
		int	group_dim
	)

Definition at line 77 of file hybrid_attention_kernels.c.

                                          {
    const int packed_dim = q_dim + gate_dim;
    if (!d_q || !d_gate || !d_packed_qg || rows <= 0 || q_dim <= 0 || gate_dim <= 0) {
        return;
    }
    if (group_dim <= 0) {
        group_dim = q_dim;
    }
    const int q_groups = q_dim / group_dim;
    const int gate_group_dim = (q_groups > 0 && gate_dim % q_groups == 0) ? (gate_dim / q_groups) : gate_dim;
    for (int row = 0; row < rows; ++row) {
        const float *dq_src = d_q + (size_t) row * (size_t) q_dim;
        const float *dg_src = d_gate + (size_t) row * (size_t) gate_dim;
        float *dst = d_packed_qg + (size_t) row * (size_t) packed_dim;
        if (q_groups > 0 && q_groups * group_dim == q_dim && q_groups * gate_group_dim == gate_dim) {
            for (int group = 0; group < q_groups; ++group) {
                const size_t dst_group_off = (size_t) group * (size_t) (group_dim + gate_group_dim);
                memcpy(
                    dst + dst_group_off,
                    dq_src + (size_t) group * (size_t) group_dim,
                    (size_t) group_dim * sizeof(float));
                memcpy(
                    dst + dst_group_off + (size_t) group_dim,
                    dg_src + (size_t) group * (size_t) gate_group_dim,
                    (size_t) gate_group_dim * sizeof(float));
            }
        } else {
            memcpy(dst, dq_src, (size_t) q_dim * sizeof(float));
            memcpy(dst + q_dim, dg_src, (size_t) gate_dim * sizeof(float));
        }
    }
}

◆ split_q_gate_forward()

void split_q_gate_forward	(	const float *	packed_qg,
		float *	q,
		float *	gate,
		int	rows,
		int	q_dim,
		int	gate_dim,
		int	group_dim
	)

Definition at line 38 of file hybrid_attention_kernels.c.

                                         {
    const int packed_dim = q_dim + gate_dim;
    if (!packed_qg || !q || !gate || rows <= 0 || q_dim <= 0 || gate_dim <= 0) {
        return;
    }
    if (group_dim <= 0) {
        group_dim = q_dim;
    }
    const int q_groups = q_dim / group_dim;
    const int gate_group_dim = (q_groups > 0 && gate_dim % q_groups == 0) ? (gate_dim / q_groups) : gate_dim;
    for (int row = 0; row < rows; ++row) {
        const float *src = packed_qg + (size_t) row * (size_t) packed_dim;
        float *q_dst = q + (size_t) row * (size_t) q_dim;
        float *gate_dst = gate + (size_t) row * (size_t) gate_dim;
        if (q_groups > 0 && q_groups * group_dim == q_dim && q_groups * gate_group_dim == gate_dim) {
            for (int group = 0; group < q_groups; ++group) {
                const size_t src_group_off = (size_t) group * (size_t) (group_dim + gate_group_dim);
                memcpy(
                    q_dst + (size_t) group * (size_t) group_dim,
                    src + src_group_off,
                    (size_t) group_dim * sizeof(float));
                memcpy(
                    gate_dst + (size_t) group * (size_t) gate_group_dim,
                    src + src_group_off + (size_t) group_dim,
                    (size_t) gate_group_dim * sizeof(float));
            }
        } else {
            memcpy(q_dst, src, (size_t) q_dim * sizeof(float));
            memcpy(gate_dst, src + q_dim, (size_t) gate_dim * sizeof(float));
        }
    }
}

◆ split_qkv_packed_head_major_forward()

void split_qkv_packed_head_major_forward	(	const float *	packed_qkv,
		float *	q,
		float *	k,
		float *	v,
		int	rows,
		int	q_dim,
		int	k_dim,
		int	v_dim,
		int	num_heads,
		int	num_kv_heads
	)

Definition at line 25 of file recurrent_split_kernels.c.

                                                           {
    if (!packed_qkv || !q || !k || !v || rows <= 0 || q_dim <= 0 || k_dim <= 0 || v_dim <= 0 ||
        num_heads <= 0 || num_kv_heads <= 0) {
        return;
    }
 
    const int q_head_dim = q_dim / num_heads;
    const int k_head_dim = k_dim / num_kv_heads;
    const int v_head_dim = v_dim / num_kv_heads;
    if (q_head_dim <= 0 || k_head_dim <= 0 || v_head_dim <= 0) {
        return;
    }
    if (q_head_dim * num_heads != q_dim || k_head_dim * num_kv_heads != k_dim || v_head_dim * num_kv_heads != v_dim) {
        return;
    }
 
    const int packed_dim = q_dim + k_dim + v_dim;
    const size_t q_head_stride = (size_t) rows * (size_t) q_head_dim;
    const size_t k_head_stride = (size_t) rows * (size_t) k_head_dim;
    const size_t v_head_stride = (size_t) rows * (size_t) v_head_dim;
 
    for (int row = 0; row < rows; ++row) {
        const float *src = packed_qkv + (size_t) row * (size_t) packed_dim;
 
        for (int head = 0; head < num_heads; ++head) {
            const float *src_q = src + (size_t) head * (size_t) q_head_dim;
            float *dst_q = q + (size_t) head * q_head_stride + (size_t) row * (size_t) q_head_dim;
            memcpy(dst_q, src_q, (size_t) q_head_dim * sizeof(float));
        }
 
        const float *src_k_base = src + (size_t) q_dim;
        const float *src_v_base = src + (size_t) q_dim + (size_t) k_dim;
        for (int head = 0; head < num_kv_heads; ++head) {
            const float *src_k = src_k_base + (size_t) head * (size_t) k_head_dim;
            const float *src_v = src_v_base + (size_t) head * (size_t) v_head_dim;
            float *dst_k = k + (size_t) head * k_head_stride + (size_t) row * (size_t) k_head_dim;
            float *dst_v = v + (size_t) head * v_head_stride + (size_t) row * (size_t) v_head_dim;
            memcpy(dst_k, src_k, (size_t) k_head_dim * sizeof(float));
            memcpy(dst_v, src_v, (size_t) v_head_dim * sizeof(float));
        }
    }
}

◆ ssm_conv1d_backward()

void ssm_conv1d_backward	(	const float *	d_out,
		const float *	conv_x,
		const float *	kernel,
		float *	d_conv_x,
		float *	d_kernel,
		int	kernel_size,
		int	num_channels,
		int	num_tokens,
		int	num_seqs
	)

Definition at line 316 of file ssm_kernels.c.

{
    ssm_conv1d_backward_ref(d_out, conv_x, kernel, d_conv_x, d_kernel, kernel_size, num_channels, num_tokens, num_seqs);
}

References ssm_conv1d_backward_ref().

◆ ssm_conv1d_forward()

void ssm_conv1d_forward	(	const float *	conv_x,
		const float *	kernel,
		float *	out,
		int	kernel_size,
		int	num_channels,
		int	num_tokens,
		int	num_seqs
	)

Definition at line 128 of file ssm_kernels.c.

{
    ssm_conv1d_forward_ref(conv_x, kernel, out, kernel_size, num_channels, num_tokens, num_seqs);
}

References ssm_conv1d_forward_ref().

◆ ssm_conv1d_forward_llama_fma()

void ssm_conv1d_forward_llama_fma	(	const float *	conv_x,
		const float *	kernel,
		float *	out,
		int	kernel_size,
		int	num_channels,
		int	num_tokens,
		int	num_seqs
	)

Definition at line 272 of file ssm_kernels.c.

{
    if (!conv_x || !kernel || !out || kernel_size <= 0 || num_channels <= 0 ||
        num_tokens < 0 || num_seqs <= 0) {
        return;
    }
    ck_ssm_conv1d_llama_args_t args = {
        conv_x, kernel, out, kernel_size, num_channels, num_tokens, num_seqs,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    const int workers = pool ? ck_threadpool_n_threads(pool) : 1;
    const int active = workers < num_channels ? workers : num_channels;
    if (active > 1 && num_tokens > 1) {
        ck_threadpool_parallel_for_n(
            pool, active, 0, num_channels, 32,
            ck_ssm_conv1d_llama_fma_channel_range, &args);
    } else {
        ck_ssm_conv1d_llama_fma_channel_range(0, num_channels, &args);
    }
}

References ck_ssm_conv1d_llama_fma_channel_range(), ck_threadpool_global(), ck_threadpool_n_threads(), and ck_threadpool_parallel_for_n().

◆ ssm_conv1d_forward_llama_production()

void ssm_conv1d_forward_llama_production	(	const float *	conv_x,
		const float *	kernel,
		float *	out,
		int	kernel_size,
		int	num_channels,
		int	num_tokens,
		int	num_seqs
	)

Definition at line 207 of file ssm_kernels.c.

{
    if (!conv_x || !kernel || !out || kernel_size <= 0 || num_channels <= 0 ||
        num_tokens < 0 || num_seqs <= 0) {
        return;
    }
    ck_ssm_conv1d_llama_args_t args = {
        conv_x, kernel, out, kernel_size, num_channels, num_tokens, num_seqs,
    };
    ck_threadpool_t *pool = ck_threadpool_global();
    const int workers = pool ? ck_threadpool_n_threads(pool) : 1;
    const int active = workers < num_channels ? workers : num_channels;
    if (active > 1 && num_tokens > 1) {
        ck_threadpool_parallel_for_n(
            pool, active, 0, num_channels, 32,
            ck_ssm_conv1d_llama_channel_range, &args);
    } else {
        ck_ssm_conv1d_llama_channel_range(0, num_channels, &args);
    }
}

References ck_ssm_conv1d_llama_channel_range(), ck_threadpool_global(), ck_threadpool_n_threads(), and ck_threadpool_parallel_for_n().

◆ ssm_conv1d_forward_llama_production_serial()

void ssm_conv1d_forward_llama_production_serial	(	const float *	conv_x,
		const float *	kernel,
		float *	out,
		int	kernel_size,
		int	num_channels,
		int	num_tokens,
		int	num_seqs
	)

Definition at line 189 of file ssm_kernels.c.

{
    if (!conv_x || !kernel || !out || kernel_size <= 0 || num_channels <= 0 ||
        num_tokens < 0 || num_seqs <= 0) {
        return;
    }
    ck_ssm_conv1d_llama_args_t args = {
        conv_x, kernel, out, kernel_size, num_channels, num_tokens, num_seqs,
    };
    ck_ssm_conv1d_llama_channel_range(0, num_channels, &args);
}

References ck_ssm_conv1d_llama_channel_range().

◆ ssm_conv1d_forward_pytorch_bf16_storage()

void ssm_conv1d_forward_pytorch_bf16_storage	(	const float *	conv_x,
		const float *	kernel,
		float *	out,
		int	kernel_size,
		int	num_channels,
		int	num_tokens,
		int	num_seqs
	)

Definition at line 299 of file ssm_kernels.c.

{
    ssm_conv1d_forward_ref(
        conv_x, kernel, out, kernel_size, num_channels, num_tokens, num_seqs);
    const size_t count =
        (size_t)num_seqs * (size_t)num_tokens * (size_t)num_channels;
    for (size_t i = 0; i < count; ++i) {
        out[i] = bf16_to_float(float_to_bf16(out[i]));
    }
}

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

◆ swiglu_backward()

void swiglu_backward	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		int	tokens,
		int	dim
	)

SwiGLU backward pass

Test:

test_swiglu.py::TestSwiGLUBackward::test_backward_tokens

test_swiglu.py::TestSwiGLUBackward::test_backward_single

test_parity.py::test_swiglu_backward_parity

Computes dGate and dUp given dY. dGate = dy * b * silu'(a), dUp = dy * silu(a)

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

Definition at line 386 of file swiglu_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        swiglu_backward_exact(input, d_output, d_input, tokens, dim);
        return;
    }
 
    int T = tokens;
    int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const float *row = input + (size_t)t * (2 * D);
        const float *dy_row = d_output + (size_t)t * D;
        float *dx_row = d_input + (size_t)t * (2 * D);
        int d = 0;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        __m512 one = _mm512_set1_ps(1.0f);
        for (; d + 16 <= D; d += 16) {
            __m512 a = _mm512_loadu_ps(&row[d]);         // gate
            __m512 b = _mm512_loadu_ps(&row[D + d]);     // value
            __m512 dy = _mm512_loadu_ps(&dy_row[d]);
 
            __m512 s = sigmoid512_fast(a);              // sigmoid(a)
            __m512 silu = _mm512_mul_ps(a, s);          // silu(a) = a * s
            __m512 one_minus_s = _mm512_sub_ps(one, s);
            __m512 inner = _mm512_fmadd_ps(a, one_minus_s, one);      // 1 + a * (1 - s)
            __m512 silu_prime = _mm512_mul_ps(s, inner);              // s * (1 + a * (1 - s))
 
            // dA = dy * b * silu_prime
            __m512 dA = _mm512_mul_ps(dy, _mm512_mul_ps(b, silu_prime));
            // dB = dy * silu
            __m512 dB = _mm512_mul_ps(dy, silu);
 
            _mm512_storeu_ps(&dx_row[d], dA);
            _mm512_storeu_ps(&dx_row[D + d], dB);
        }
#elif defined(__AVX2__)
        // AVX2: Process 8 floats at a time
        __m256 one = _mm256_set1_ps(1.0f);
        for (; d + 8 <= D; d += 8) {
            __m256 a = _mm256_loadu_ps(&row[d]);         // gate
            __m256 b = _mm256_loadu_ps(&row[D + d]);     // value
            __m256 dy = _mm256_loadu_ps(&dy_row[d]);
 
            __m256 s = sigmoid256_fast(a);              // sigmoid(a)
            __m256 silu = _mm256_mul_ps(a, s);          // silu(a) = a * s
            __m256 one_minus_s = _mm256_sub_ps(one, s);
            __m256 inner = _mm256_fmadd_ps(a, one_minus_s, one);      // 1 + a * (1 - s)
            __m256 silu_prime = _mm256_mul_ps(s, inner);              // s * (1 + a * (1 - s))
 
            // dA = dy * b * silu_prime
            __m256 dA = _mm256_mul_ps(dy, _mm256_mul_ps(b, silu_prime));
            // dB = dy * silu
            __m256 dB = _mm256_mul_ps(dy, silu);
 
            _mm256_storeu_ps(&dx_row[d], dA);
            _mm256_storeu_ps(&dx_row[D + d], dB);
        }
#elif defined(__AVX__)
        // AVX1: Vectorize arithmetic, use scalar sigmoid
        __m256 one = _mm256_set1_ps(1.0f);
        float a_arr[8] __attribute__((aligned(32)));
        float s_arr[8] __attribute__((aligned(32)));
 
        for (; d + 8 <= D; d += 8) {
            __m256 a = _mm256_loadu_ps(&row[d]);         // gate
            __m256 b = _mm256_loadu_ps(&row[D + d]);     // value
            __m256 dy = _mm256_loadu_ps(&dy_row[d]);
 
            // Compute sigmoid scalarly
            _mm256_store_ps(a_arr, a);
            for (int j = 0; j < 8; ++j) {
                s_arr[j] = sigmoid_scalar(a_arr[j]);
            }
            __m256 s = _mm256_load_ps(s_arr);
 
            __m256 silu = _mm256_mul_ps(a, s);                        // silu(a) = a * s
            __m256 one_minus_s = _mm256_sub_ps(one, s);
            __m256 a_one_minus_s = _mm256_mul_ps(a, one_minus_s);
            __m256 inner = _mm256_add_ps(one, a_one_minus_s);         // 1 + a * (1 - s)
            __m256 silu_prime = _mm256_mul_ps(s, inner);              // s * (1 + a * (1 - s))
 
            // dA = dy * b * silu_prime
            __m256 dA = _mm256_mul_ps(dy, _mm256_mul_ps(b, silu_prime));
            // dB = dy * silu
            __m256 dB = _mm256_mul_ps(dy, silu);
 
            _mm256_storeu_ps(&dx_row[d], dA);
            _mm256_storeu_ps(&dx_row[D + d], dB);
        }
#endif
 
        // Scalar fallback for remaining elements
        for (; d < D; ++d) {
            float a = row[d];       // gate
            float b = row[D + d];   // value
            float dy = dy_row[d];
 
            float s = sigmoid_scalar(a);               // sigmoid(a)
            float silu = a * s;                        // silu(a)
            float silu_prime = s * (1.0f + a * (1.0f - s)); // silu'(a), PyTorch form
 
            float dA = dy * b * silu_prime;
            float dB = dy * silu;
 
            dx_row[d] = dA;
            dx_row[D + d] = dB;
        }
    }
}

References __attribute__(), ck_strict_parity_enabled(), sigmoid_scalar(), silu(), and swiglu_backward_exact().

Referenced by ck_layer_backward_rmsnorm_swiglu().

◆ swiglu_backward_bf16()

void swiglu_backward_bf16	(	const uint16_t *	input,
		const uint16_t *	d_output,
		uint16_t *	d_input,
		int	tokens,
		int	dim
	)

Definition at line 108 of file swiglu_kernels_bf16.c.

{
    if (!input || !d_output || !d_input || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const int T = tokens;
    const int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const uint16_t *row = input + (size_t)t * (size_t)(2 * D);
        const uint16_t *dy_row = d_output + (size_t)t * (size_t)D;
        uint16_t *dx_row = d_input + (size_t)t * (size_t)(2 * D);
        int d = 0;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        __m512 one = _mm512_set1_ps(1.0f);
        for (; d + 16 <= D; d += 16) {
            __m512 a = bf16_loadu_cvt_fp32(&row[d]);         // gate
            __m512 b = bf16_loadu_cvt_fp32(&row[D + d]);     // value
            __m512 dy = bf16_loadu_cvt_fp32(&dy_row[d]);
 
            __m512 s = sigmoid512_fast_bf16(a);             // sigmoid(a)
            __m512 silu = _mm512_mul_ps(a, s);              // silu(a) = a * s
            __m512 s_prime = _mm512_mul_ps(s, _mm512_sub_ps(one, s)); // s * (1 - s)
            __m512 silu_prime = _mm512_fmadd_ps(a, s_prime, s);       // s + a * s_prime
 
            // dA = dy * b * silu_prime
            __m512 dA = _mm512_mul_ps(dy, _mm512_mul_ps(b, silu_prime));
            // dB = dy * silu
            __m512 dB = _mm512_mul_ps(dy, silu);
 
            fp32_cvt_storeu_bf16(&dx_row[d], dA);
            fp32_cvt_storeu_bf16(&dx_row[D + d], dB);
        }
#endif
 
        // Scalar fallback for remaining elements
        for (; d < D; ++d) {
            float a = bf16_to_float(row[d]);
            float b = bf16_to_float(row[D + d]);
            float dy = bf16_to_float(dy_row[d]);
 
            float s = sigmoid_scalar(a);
            float silu = a * s;
            float s_prime = s * (1.0f - s);
            float silu_prime = s + a * s_prime;
 
            float dA = dy * b * silu_prime;
            float dB = dy * silu;
 
            dx_row[d] = float_to_bf16(dA);
            dx_row[D + d] = float_to_bf16(dB);
        }
    }
}

References bf16_to_float(), float_to_bf16(), sigmoid_scalar(), and silu().

◆ swiglu_backward_exact()

void swiglu_backward_exact	(	const float *	input,
		const float *	d_output,
		float *	d_input,
		int	tokens,
		int	dim
	)

SwiGLU backward pass (exact version using stdlib sigmoid)

Test:

test_swiglu.py::TestSwiGLUBackward::test_exact_vs_fast

test_swiglu.py::TestSwiGLUBackward::test_exact_single

Uses standard library expf for numerical accuracy reference.

After changes: make test

Definition at line 719 of file swiglu_kernels.c.

{
    int T = tokens;
    int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const float *row = input + (size_t)t * (2 * D);
        const float *dy_row = d_output + (size_t)t * D;
        float *dx_row = d_input + (size_t)t * (2 * D);
 
        for (int d = 0; d < D; ++d) {
            float a = row[d];       // gate
            float b = row[D + d];   // value
            float dy = dy_row[d];
 
            float s = sigmoid_scalar_parity(a); // sigmoid(a)
            float silu = a * s;                 // silu(a)
            float silu_prime = s * (1.0f + a * (1.0f - s)); // silu'(a), PyTorch form
 
            float dA = dy * b * silu_prime;
            float dB = dy * silu;
 
            dx_row[d] = dA;
            dx_row[D + d] = dB;
        }
    }
}

References sigmoid_scalar_parity(), and silu().

Referenced by swiglu_backward().

◆ swiglu_forward()

void swiglu_forward	(	const float *	input,
		float *	output,
		int	tokens,
		int	dim
	)

SwiGLU forward pass

Test:

test_swiglu.py::TestSwiGLUForward::test_forward_tokens

test_swiglu.py::TestSwiGLUForward::test_forward_single

test_mlp.py::TestMLPForward::test_swiglu_mlp

test_fused_swiglu_decode.py::TestFusedSwiGLUDecode::test_fused_swiglu_decode

test_parity.py::test_swiglu_parity

SwiGLU: y = silu(gate) * up where silu(x) = x * sigmoid(x)

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

Definition at line 237 of file swiglu_kernels.c.

{
    const char *fast_env = getenv("CK_SWIGLU_FAST");
    const char *exact_env = getenv("CK_SWIGLU_EXACT");
    if (ck_strict_parity_enabled() ||
        !(fast_env && atoi(fast_env) != 0) ||
        (exact_env && atoi(exact_env) != 0)) {
        swiglu_forward_exact(input, output, tokens, dim);
        return;
    }
 
    int T = tokens;
    int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const float *row = input + (size_t)t * (2 * D);
        float *out_row = output + (size_t)t * D;
        int d = 0;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        for (; d + 16 <= D; d += 16) {
            __m512 a = _mm512_loadu_ps(&row[d]);         // gate
            __m512 b = _mm512_loadu_ps(&row[D + d]);     // value
 
            __m512 s = sigmoid512_fast(a);              // sigmoid(a)
            __m512 silu = _mm512_mul_ps(a, s);          // silu(a) = a * sigmoid(a)
            __m512 y = _mm512_mul_ps(silu, b);          // y = silu(a) * b
 
            _mm512_storeu_ps(&out_row[d], y);
        }
#elif defined(__AVX2__)
        // AVX2: Process 8 floats at a time
        for (; d + 8 <= D; d += 8) {
            __m256 a = _mm256_loadu_ps(&row[d]);         // gate
            __m256 b = _mm256_loadu_ps(&row[D + d]);     // value
 
            __m256 s = sigmoid256_fast(a);              // sigmoid(a)
            __m256 silu = _mm256_mul_ps(a, s);          // silu(a) = a * sigmoid(a)
            __m256 y = _mm256_mul_ps(silu, b);          // y = silu(a) * b
 
            _mm256_storeu_ps(&out_row[d], y);
        }
#elif defined(__AVX__)
        // AVX1: Vectorize arithmetic, use scalar sigmoid
        float a_arr[8] __attribute__((aligned(32)));
        float s_arr[8] __attribute__((aligned(32)));
 
        for (; d + 8 <= D; d += 8) {
            __m256 a = _mm256_loadu_ps(&row[d]);         // gate
            __m256 b = _mm256_loadu_ps(&row[D + d]);     // value
 
            // Compute sigmoid scalarly
            _mm256_store_ps(a_arr, a);
            for (int j = 0; j < 8; ++j) {
                s_arr[j] = sigmoid_scalar(a_arr[j]);
            }
            __m256 s = _mm256_load_ps(s_arr);
 
            __m256 silu = _mm256_mul_ps(a, s);          // silu(a) = a * sigmoid(a)
            __m256 y = _mm256_mul_ps(silu, b);          // y = silu(a) * b
 
            _mm256_storeu_ps(&out_row[d], y);
        }
#endif
 
        // Scalar fallback for remaining elements
        for (; d < D; ++d) {
            float a = row[d];       // gate
            float b = row[D + d];   // value
 
            float s = sigmoid_scalar(a);         // sigmoid(a)
            float silu = a * s;                  // silu(a) = a * sigmoid(a)
 
            out_row[d] = silu * b;
        }
    }
}

References __attribute__(), ck_strict_parity_enabled(), sigmoid_scalar(), silu(), and swiglu_forward_exact().

◆ swiglu_forward_bf16()

void swiglu_forward_bf16	(	const uint16_t *	input,
		uint16_t *	output,
		int	tokens,
		int	dim
	)

Definition at line 66 of file swiglu_kernels_bf16.c.

{
    if (!input || !output || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const int T = tokens;
    const int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const uint16_t *row = input + (size_t)t * (size_t)(2 * D);
        uint16_t *out_row = output + (size_t)t * (size_t)D;
        int d = 0;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        for (; d + 16 <= D; d += 16) {
            __m512 a = bf16_loadu_cvt_fp32(&row[d]);         // gate
            __m512 b = bf16_loadu_cvt_fp32(&row[D + d]);     // value
 
            __m512 s = sigmoid512_fast_bf16(a);             // sigmoid(a)
            __m512 silu = _mm512_mul_ps(a, s);              // silu(a) = a * sigmoid(a)
            __m512 y = _mm512_mul_ps(silu, b);              // y = silu(a) * b
 
            fp32_cvt_storeu_bf16(&out_row[d], y);
        }
#endif
 
        // Scalar fallback for remaining elements
        for (; d < D; ++d) {
            float a = bf16_to_float(row[d]);
            float b = bf16_to_float(row[D + d]);
            float s = sigmoid_scalar(a);
            float silu = a * s;
            out_row[d] = float_to_bf16(silu * b);
        }
    }
}

References bf16_to_float(), float_to_bf16(), sigmoid_scalar(), and silu().

◆ swiglu_forward_exact()

void swiglu_forward_exact	(	const float *	input,
		float *	output,
		int	tokens,
		int	dim
	)

SwiGLU forward pass (exact version using stdlib sigmoid)

Test:

test_swiglu.py::TestSwiGLUForward::test_exact_vs_fast

test_swiglu.py::TestSwiGLUForward::test_exact_single

Uses standard library expf for numerical accuracy reference.

After changes: make test

Definition at line 515 of file swiglu_kernels.c.

{
    int T = tokens;
    int D = dim;
 
    for (int t = 0; t < T; ++t) {
        const float *row = input + (size_t)t * (2 * D);
        float *out_row = output + (size_t)t * D;
 
        for (int d = 0; d < D; ++d) {
            float a = row[d];       // gate
            float b = row[D + d];   // value
 
            float s = sigmoid_scalar_parity(a); // sigmoid(a)
            float silu = a * s;                 // silu(a)
            out_row[d] = silu * b;
        }
    }
}

References sigmoid_scalar_parity(), and silu().

Referenced by ck_mlp_swiglu_forward(), ck_mlp_swiglu_forward_ref(), and swiglu_forward().

◆ swiglu_forward_ggml()

void swiglu_forward_ggml	(	const float *	input,
		float *	output,
		int	tokens,
		int	dim
	)

Definition at line 538 of file swiglu_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *row = input + (size_t)t * (2 * dim);
        float *out_row = output + (size_t)t * dim;
        int d = 0;
 
#if defined(__AVX512F__) && defined(__AVX512DQ__)
        for (; d + 16 <= dim; d += 16) {
            const __m512 gate = _mm512_loadu_ps(row + d);
            const __m512 up = _mm512_loadu_ps(row + dim + d);
            const __m512 neg_gate = _mm512_sub_ps(_mm512_setzero_ps(), gate);
            const __m512 denom = _mm512_add_ps(
                _mm512_set1_ps(1.0f), ck_ggml_expf_avx512(neg_gate));
            const __m512 silu = _mm512_div_ps(gate, denom);
            _mm512_storeu_ps(out_row + d, _mm512_mul_ps(silu, up));
        }
#elif defined(__AVX2__) && defined(__FMA__)
        for (; d + 8 <= dim; d += 8) {
            const __m256 gate = _mm256_loadu_ps(row + d);
            const __m256 up = _mm256_loadu_ps(row + dim + d);
            const __m256 neg_gate = _mm256_sub_ps(_mm256_setzero_ps(), gate);
            const __m256 denom = _mm256_add_ps(
                _mm256_set1_ps(1.0f), ck_ggml_expf_avx2(neg_gate));
            const __m256 silu = _mm256_div_ps(gate, denom);
            _mm256_storeu_ps(out_row + d, _mm256_mul_ps(silu, up));
        }
#endif
        for (; d < dim; ++d) {
            const float gate = row[d];
            out_row[d] = (gate / (1.0f + expf(-gate))) * row[dim + d];
        }
    }
}

References silu().

Referenced by ck_moe_q4k_mixed_route_work(), ck_moe_q4k_q5k_bucket_work(), ck_moe_q4k_q5k_route_work(), moe_swiglu_expert_forward_q4k_q4k_workspace(), moe_swiglu_expert_forward_q4k_q5_0_workspace(), moe_swiglu_expert_forward_q4k_q5k_workspace(), moe_swiglu_expert_forward_q4k_q6k_workspace(), moe_swiglu_expert_forward_q4k_q8_0_workspace(), moe_swiglu_shared_forward_q4k_q4k_workspace(), moe_swiglu_shared_forward_q4k_q6k_workspace(), and moe_swiglu_shared_forward_q8_0_gated_workspace().

◆ swiglu_forward_ggml_split()

void swiglu_forward_ggml_split	(	const float *	gate,
		const float *	up,
		float *	output,
		int	tokens,
		int	dim
	)

Definition at line 576 of file swiglu_kernels.c.

{
    if (!gate || !up || !output || tokens <= 0 || dim <= 0) {
        return;
    }
    for (int t = 0; t < tokens; ++t) {
        const float *gate_row = gate + (size_t)t * (size_t)dim;
        const float *up_row = up + (size_t)t * (size_t)dim;
        float *out_row = output + (size_t)t * (size_t)dim;
        int d = 0;
 
#if defined(__AVX512F__) && defined(__AVX512DQ__)
        for (; d + 16 <= dim; d += 16) {
            const __m512 gate_v = _mm512_loadu_ps(gate_row + d);
            const __m512 up_v = _mm512_loadu_ps(up_row + d);
            const __m512 neg_gate = _mm512_sub_ps(_mm512_setzero_ps(), gate_v);
            const __m512 denom = _mm512_add_ps(
                _mm512_set1_ps(1.0f), ck_ggml_expf_avx512(neg_gate));
            const __m512 silu = _mm512_div_ps(gate_v, denom);
            _mm512_storeu_ps(out_row + d, _mm512_mul_ps(silu, up_v));
        }
#elif defined(__AVX2__) && defined(__FMA__)
        for (; d + 8 <= dim; d += 8) {
            const __m256 gate_v = _mm256_loadu_ps(gate_row + d);
            const __m256 up_v = _mm256_loadu_ps(up_row + d);
            const __m256 neg_gate = _mm256_sub_ps(_mm256_setzero_ps(), gate_v);
            const __m256 denom = _mm256_add_ps(
                _mm256_set1_ps(1.0f), ck_ggml_expf_avx2(neg_gate));
            const __m256 silu = _mm256_div_ps(gate_v, denom);
            _mm256_storeu_ps(out_row + d, _mm256_mul_ps(silu, up_v));
        }
#endif
        for (; d < dim; ++d) {
            const float gate_v = gate_row[d];
            out_row[d] = (gate_v / (1.0f + expf(-gate_v))) * up_row[d];
        }
    }
}

References silu().

Referenced by ck_moe_q4k_q5k_bucket_work(), and ck_moe_shared_q4k_gated_workspace().

◆ swiglu_forward_pytorch_bf16_storage()

void swiglu_forward_pytorch_bf16_storage	(	const float *	input,
		float *	output,
		int	tokens,
		int	dim
	)

Definition at line 658 of file swiglu_kernels.c.

{
    if (!input || !output || tokens < 0 || dim < 0) {
        fprintf(stderr, "HARD KERNEL CONTRACT FAULT: invalid PyTorch BF16 SwiGLU arguments\n");
        abort();
    }
 
#if defined(__AVX512F__)
    pthread_once(&ck_pytorch_swiglu_once, ck_bind_pytorch_swiglu_sleef);
    if (!ck_pytorch_swiglu_expf16) {
        fprintf(stderr,
                "HARD KERNEL CONTRACT FAULT: PyTorch BF16 SwiGLU requires "
                "SLEEF Sleef_expf16_u10; set CK_SLEEF_LIBRARY\n");
        abort();
    }
#endif
 
    for (int t = 0; t < tokens; ++t) {
        const float *row = input + (size_t)t * (size_t)(2 * dim);
        float *out_row = output + (size_t)t * (size_t)dim;
        int d = 0;
 
#if defined(__AVX512F__)
        for (; d + 16 <= dim; d += 16) {
            const __m512 gate = _mm512_loadu_ps(row + d);
            const __m512 denominator = _mm512_add_ps(
                _mm512_set1_ps(1.0f),
                ck_pytorch_swiglu_expf16(_mm512_sub_ps(_mm512_setzero_ps(), gate)));
            const __m512 silu = _mm512_div_ps(gate, denominator);
            float silu_lanes[16] __attribute__((aligned(64)));
            _mm512_store_ps(silu_lanes, silu);
            for (int lane = 0; lane < 16; ++lane) {
                const float silu_bf16 = bf16_to_float(float_to_bf16(silu_lanes[lane]));
                const float up_bf16 = bf16_to_float(float_to_bf16(row[dim + d + lane]));
                out_row[d + lane] = bf16_to_float(
                    float_to_bf16(silu_bf16 * up_bf16));
            }
        }
#endif
        for (; d < dim; ++d) {
            const float gate_bf16 = bf16_to_float(float_to_bf16(row[d]));
            const float up_bf16 = bf16_to_float(float_to_bf16(row[dim + d]));
            const float silu = gate_bf16 / (1.0f + expf(-gate_bf16));
            const float silu_bf16 = bf16_to_float(float_to_bf16(silu));
            out_row[d] = bf16_to_float(float_to_bf16(silu_bf16 * up_bf16));
        }
    }
}

References __attribute__(), bf16_to_float(), float_to_bf16(), and silu().

Referenced by moe_swiglu_packed_expert_forward_bf16().

◆ swiglu_forward_q8_k()

void swiglu_forward_q8_k	(	const float *	input,
		void *	output_q8,
		int	tokens,
		int	dim
	)

Definition at line 319 of file swiglu_kernels.c.

{
    if (!input || !output_q8 || tokens <= 0 || dim <= 0) {
        return;
    }
    if ((dim % QK_K) != 0) {
        return;
    }
 
    const char *fast_env = getenv("CK_SWIGLU_FAST");
    const char *exact_env = getenv("CK_SWIGLU_EXACT");
    const int use_fast = !ck_strict_parity_enabled() &&
                         (fast_env && atoi(fast_env) != 0) &&
                         !(exact_env && atoi(exact_env) != 0);
 
    const int blocks_per_row = dim / QK_K;
    block_q8_K *q8 = (block_q8_K *)output_q8;
    float tmp[QK_K];
 
    for (int t = 0; t < tokens; ++t) {
        const float *row = input + (size_t)t * (size_t)(2 * dim);
        block_q8_K *q8_row = q8 + (size_t)t * (size_t)blocks_per_row;
 
        for (int block = 0; block < blocks_per_row; ++block) {
            const int base = block * QK_K;
            int d = 0;
 
#if defined(__AVX2__)
            if (use_fast) {
                for (; d + 8 <= QK_K; d += 8) {
                    const __m256 a = _mm256_loadu_ps(row + base + d);
                    const __m256 b = _mm256_loadu_ps(row + dim + base + d);
                    const __m256 s = sigmoid256_fast(a);
                    const __m256 y = _mm256_mul_ps(_mm256_mul_ps(a, s), b);
                    _mm256_storeu_ps(tmp + d, y);
                }
            }
#else
            (void)use_fast;
#endif
 
            for (; d < QK_K; ++d) {
                const float a = row[base + d];
                const float b = row[dim + base + d];
                const float s = use_fast ? sigmoid_scalar(a) : sigmoid_scalar_parity(a);
                tmp[d] = (a * s) * b;
            }
 
            quantize_row_q8_k(tmp, (void *)&q8_row[block], QK_K);
        }
    }
}

References ck_strict_parity_enabled(), QK_K, quantize_row_q8_k(), sigmoid_scalar(), and sigmoid_scalar_parity().

◆ topk_batched_f32()

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

Batched top-K selection for multiple tokens.

Parameters

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

Definition at line 465 of file topk_kernels.c.

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

References topk_f32(), and topk_softmax_f32().

◆ topk_f32()

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

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

Parameters

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

Definition at line 51 of file topk_kernels.c.

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

Referenced by moe_softmax_topk_router_llama_f32_workspace(), topk_batched_f32(), and topk_softmax_f32().

◆ topk_softmax_backward_f32()

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

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

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

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

Definition at line 408 of file topk_kernels.c.

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

◆ topk_softmax_f32()

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

Find top-K indices with softmax-normalized weights.

Parameters

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

Definition at line 136 of file topk_kernels.c.

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

References topk_f32().

Referenced by topk_batched_f32().

◆ unfused_rmsnorm_qkv_prefill()

void unfused_rmsnorm_qkv_prefill	(	const float *	x,
		const float *	gamma,
		const float *	Wq,
		const float *	Wk,
		const float *	Wv,
		float *	x_norm,
		float *	Q,
		float *	K,
		float *	V,
		int	seq_len,
		int	hidden,
		int	q_dim,
		int	kv_dim,
		float	eps
	)

Unfused version for benchmarking comparison.

Definition at line 667 of file prefill_fused_gemm.c.

{
    /* Step 1: Full RMSNorm → writes x_norm to memory */
    rmsnorm_tile(x, gamma, x_norm, seq_len, hidden, hidden, eps);
 
    /* Step 2: Separate GEMMs with N-outer tiling for weight reuse */
    /* Q projection */
    for (int n_start = 0; n_start < q_dim; n_start += PREFILL_TILE_N) {
        int tile_n = (n_start + PREFILL_TILE_N <= q_dim)
                         ? PREFILL_TILE_N : (q_dim - n_start);
        const float *W_tile = Wq + (size_t)n_start * hidden;
 
        for (int m_start = 0; m_start < seq_len; m_start += PREFILL_TILE_M) {
            int tile_m = (m_start + PREFILL_TILE_M <= seq_len)
                             ? PREFILL_TILE_M : (seq_len - m_start);
            const float *x_tile = x_norm + (size_t)m_start * hidden;
            float *out_tile = Q + (size_t)m_start * q_dim + n_start;
            gemm_tile_nt_strided(x_tile, W_tile, out_tile,
                                  tile_m, tile_n, hidden, q_dim);
        }
    }
 
    /* K projection */
    for (int n_start = 0; n_start < kv_dim; n_start += PREFILL_TILE_N) {
        int tile_n = (n_start + PREFILL_TILE_N <= kv_dim)
                         ? PREFILL_TILE_N : (kv_dim - n_start);
        const float *W_tile = Wk + (size_t)n_start * hidden;
 
        for (int m_start = 0; m_start < seq_len; m_start += PREFILL_TILE_M) {
            int tile_m = (m_start + PREFILL_TILE_M <= seq_len)
                             ? PREFILL_TILE_M : (seq_len - m_start);
            const float *x_tile = x_norm + (size_t)m_start * hidden;
            float *out_tile = K + (size_t)m_start * kv_dim + n_start;
            gemm_tile_nt_strided(x_tile, W_tile, out_tile,
                                  tile_m, tile_n, hidden, kv_dim);
        }
    }
 
    /* V projection */
    for (int n_start = 0; n_start < kv_dim; n_start += PREFILL_TILE_N) {
        int tile_n = (n_start + PREFILL_TILE_N <= kv_dim)
                         ? PREFILL_TILE_N : (kv_dim - n_start);
        const float *W_tile = Wv + (size_t)n_start * hidden;
 
        for (int m_start = 0; m_start < seq_len; m_start += PREFILL_TILE_M) {
            int tile_m = (m_start + PREFILL_TILE_M <= seq_len)
                             ? PREFILL_TILE_M : (seq_len - m_start);
            const float *x_tile = x_norm + (size_t)m_start * hidden;
            float *out_tile = V + (size_t)m_start * kv_dim + n_start;
            gemm_tile_nt_strided(x_tile, W_tile, out_tile,
                                  tile_m, tile_n, hidden, kv_dim);
        }
    }
}

References gemm_tile_nt_strided(), PREFILL_TILE_M, PREFILL_TILE_N, and rmsnorm_tile().

◆ vec_dot_q6_k_q8_k()

void vec_dot_q6_k_q8_k	(	int	n,
		float *	s,
		const void *	vx,
		const void *	vy
	)

Q6_K x Q8_K dot product (single row)

Definition at line 1324 of file gemm_kernels_q6k_q8k.c.

{
    if (!s || !vx || !vy || n <= 0) {
        return;
    }
 
    const block_q6_K *x = (const block_q6_K *)vx;
    const block_q8_K *y = (const block_q8_K *)vy;
 
    /* This is the architecture-neutral scalar oracle. Production x86 dispatch
     * uses the separately parity-tested SIMD reduction tree. */
    *s = dot_q6_k_q8_k_ref(x, y, n);
}

References dot_q6_k_q8_k_ref().

◆ vision_position_ids_2d_merge()

void vision_position_ids_2d_merge	(	int32_t *	positions,
		int	grid_h,
		int	grid_w,
		int	merge_size
	)

Build merged 2D vision position IDs in the layout expected by vision M-RoPE.

Output layout: [4, grid_h * grid_w] flattened as [y_stream | x_stream | y_stream_dup | x_stream_dup]

Tokens are emitted in merged-tile traversal order so the position buffer matches the same 2x2 grouping used by Qwen-style vision encoders.

Definition at line 529 of file vision_kernels.c.

{
    if (!positions || grid_h <= 0 || grid_w <= 0 || merge_size <= 0) {
        return;
    }
 
    const int num_tokens = grid_h * grid_w;
    int ptr = 0;
 
    for (int y = 0; y < grid_h; y += merge_size) {
        for (int x = 0; x < grid_w; x += merge_size) {
            for (int dy = 0; dy < merge_size; ++dy) {
                for (int dx = 0; dx < merge_size; ++dx) {
                    const int yy = y + dy;
                    const int xx = x + dx;
                    if (yy >= grid_h || xx >= grid_w || ptr >= num_tokens) {
                        continue;
                    }
                    positions[ptr] = yy;
                    positions[num_tokens + ptr] = xx;
                    positions[2 * num_tokens + ptr] = yy;
                    positions[3 * num_tokens + ptr] = xx;
                    ++ptr;
                }
            }
        }
    }
 
    for (; ptr < num_tokens; ++ptr) {
        positions[ptr] = 0;
        positions[num_tokens + ptr] = 0;
        positions[2 * num_tokens + ptr] = 0;
        positions[3 * num_tokens + ptr] = 0;
    }
}

◆ weighted_sum_f32()

void weighted_sum_f32	(	float *	y,
		const float **	vectors,
		const float *	weights,
		int	k,
		int	n
	)

Weighted sum of k vectors: y = sum_i(weights[i] * vectors[i])

Parameters

y	Output vector [n]
vectors	Array of k input vector pointers, each [n]
weights	Array of k scalar weights
k	Number of vectors to combine
n	Vector length

Definition at line 206 of file axpy_kernels.c.

{
    if (!y || !vectors || !weights || k <= 0 || n <= 0) {
        return;
    }
 
    /* Initialize with first vector */
    scal_copy_f32(y, vectors[0], weights[0], n);
 
    /* Accumulate rest */
    for (int i = 1; i < k; i++) {
        axpy_f32(y, vectors[i], weights[i], n);
    }
}

References axpy_f32(), and scal_copy_f32().

◆ yarn_rope_cache_contiguous_positions_f32()

void yarn_rope_cache_contiguous_positions_f32	(	float *	cos_cache,
		float *	sin_cache,
		int	num_tokens,
		int	rotary_dim,
		float	freq_base,
		float	factor,
		int	original_context,
		float	beta_fast,
		float	beta_slow,
		float	mscale,
		float	mscale_all_dim
	)

Definition at line 517 of file rope_kernels.c.

{
    yarn_rope_cache_explicit_positions_impl(
        cos_cache, sin_cache, NULL, NULL, NULL, num_tokens, rotary_dim,
        freq_base, factor, original_context, beta_fast, beta_slow, mscale,
        mscale_all_dim);
}

References yarn_rope_cache_explicit_positions_impl().

◆ yarn_rope_cache_explicit_positions_bf16()

void yarn_rope_cache_explicit_positions_bf16	(	uint16_t *	cos_cache,
		uint16_t *	sin_cache,
		const int32_t *	positions,
		int	num_tokens,
		int	rotary_dim,
		float	freq_base,
		float	factor,
		int	original_context,
		float	beta_fast,
		float	beta_slow,
		float	mscale,
		float	mscale_all_dim
	)

Definition at line 535 of file rope_kernels.c.

{
    yarn_rope_cache_explicit_positions_impl(
        NULL, NULL, cos_cache, sin_cache, positions, num_tokens, rotary_dim,
        freq_base, factor, original_context, beta_fast, beta_slow, mscale,
        mscale_all_dim);
}

References yarn_rope_cache_explicit_positions_impl().

◆ yarn_rope_cache_explicit_positions_f32()

void yarn_rope_cache_explicit_positions_f32	(	float *	cos_cache,
		float *	sin_cache,
		const int32_t *	positions,
		int	num_tokens,
		int	rotary_dim,
		float	freq_base,
		float	factor,
		int	original_context,
		float	beta_fast,
		float	beta_slow,
		float	mscale,
		float	mscale_all_dim
	)

Definition at line 498 of file rope_kernels.c.

{
    yarn_rope_cache_explicit_positions_impl(
        cos_cache, sin_cache, NULL, NULL, positions, num_tokens, rotary_dim,
        freq_base, factor, original_context, beta_fast, beta_slow, mscale,
        mscale_all_dim);
}

References yarn_rope_cache_explicit_positions_impl().

Data Structures

Enumerations

Functions

Enumeration Type Documentation

◆ ck_attention_prefill_schedule_t

◆ ck_attention_reduction_t

◆ ck_attention_status_t

Function Documentation

◆ adamw_clip_update_multi_f32()

◆ adamw_update_f32()

◆ add_backward_bf16()

◆ add_forward_2d_bf16()

◆ add_forward_bf16()

◆ add_forward_f32()

◆ add_inplace_bf16()

◆ add_inplace_f32()

◆ add_scaled_forward_bf16()

◆ add_scaled_inplace_bf16()

◆ add_stream_inplace()

◆ add_stream_reorder_2d()

◆ argmax_f32()

◆ assistant_layer_scale_forward()

◆ attention_backward_causal_head_major()

◆ attention_backward_causal_head_major_gqa()

◆ attention_backward_causal_head_major_gqa_bf16()

◆ attention_flash_decode()

◆ 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_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_sliding()

◆ attention_forward_causal_head_major_gqa_flash_strided_sliding_gemma4()

◆ attention_forward_causal_head_major_gqa_flash_strided_sliding_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_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_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_full_bf16cache_pytorch_contract()

◆ attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_contract_workspace()

◆ attention_forward_causal_head_major_shared_kv_gemma4()

◆ attention_forward_causal_head_major_shared_kv_sliding_gemma4()

◆ attention_forward_chunk_head_major_gqa_flash_gemma4()

◆ attention_forward_decode_head_major_gqa_bf16cache_pytorch_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_f16kv()

◆ attention_forward_decode_head_major_gqa_flash_gemma4()

◆ attention_forward_decode_head_major_gqa_flash_sliding()

◆ attention_forward_decode_head_major_gqa_flash_sliding_gemma4()

◆ attention_forward_decode_head_major_gqa_llama_regular_sliding_workspace()

◆ attention_forward_decode_head_major_gqa_regular()

◆ attention_forward_decode_head_major_shared_kv_gemma4()

◆ attention_forward_decode_head_major_shared_kv_sliding_gemma4()

◆ attention_forward_full_head_major_gqa_exact_strided()

◆ attention_forward_full_head_major_gqa_flash()

◆ attention_forward_full_head_major_gqa_flash_strided()

◆ attention_forward_full_head_major_gqa_flash_strided_bf16_storage()

◆ attention_forward_full_head_major_gqa_flash_strided_gemma4()

◆ attention_forward_full_head_major_gqa_ggml_strided()

◆ attention_forward_full_head_major_gqa_ggml_strided_workspace()

◆ attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage()

◆ attention_forward_full_head_major_gqa_pytorch_cpu_flash_bf16_storage_token_output()

◆ attention_forward_full_head_major_gqa_sdpa_bf16_storage()

◆ attention_forward_full_head_major_gqa_tiled336_f16kv_fp32_strided()

◆ attention_forward_full_head_major_gqa_tiled64_f16kv_fp32_strided()