AXPY kernels for FP32: y = y + alpha * x. More...

#include <stdint.h>
#include <stddef.h>
#include <stdatomic.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "ckernel_engine.h"
#include "ckernel_dtype.h"
#include "ck_threadpool.h"
#include "bf16_utils.h"

Typedefs
typedef int(*	ck_moe_expert_workspace_fn) (const float , const int , const float , const void , const void , const void , float , int, int, int, int, int, void , size_t)

typedef int(*	ck_moe_shared_gated_workspace_fn) (const float , const float , const void , const void , const void , const float , float , int, int, int, void , size_t)

typedef int(*	ck_moe_shared_workspace_fn) (const float , const float , const void , const void , const void , float , int, int, int, void *, size_t)

Enumerations
enum	{ CK_MOE_Q4K_Q5K_TASK_ROWS = 16 }

enum	ck_moe_down_kind_t { CK_MOE_DOWN_Q4_K , CK_MOE_DOWN_Q6_K , CK_MOE_DOWN_Q5_0 , CK_MOE_DOWN_Q8_0 }

Functions
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.

static size_t	ck_moe_align64 (size_t value)

static float	ck_moe_bf16_round (float x)

static int	ck_moe_bucket_expert_for_position (const int *offsets, int n_experts, int position)

static int	ck_moe_bucket_layout (int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k, ck_moe_q4k_q5k_bucket_layout_t *layout)

static int	ck_moe_debug_enabled (void)

static void	ck_moe_debug_finite (const char name, const float x, size_t n)

static size_t	ck_moe_down_idx (int e, int h, int i, int hidden_dim, int intermediate_dim)

static float	ck_moe_dsilu_f32 (float x)

static void	ck_moe_llama_weighted_accumulate (float output, const float expert_output, float route_weight, int n)

static void	ck_moe_q4k_llama_projection (float output, const void weights, const void input_q8, int output_dim, int input_dim, void scratch)

static size_t	ck_moe_q4k_llama_projection_scratch_bytes (int output_dim, int input_dim)

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

static int	ck_moe_q4k_mixed_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, ck_moe_expert_workspace_fn serial_fn, ck_moe_down_kind_t down_kind)

static int	ck_moe_q4k_mixed_route_parallel (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 hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes, size_t workspace_stride, ck_threadpool_t *pool, ck_moe_down_kind_t down_kind)

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

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

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

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

static int	ck_moe_q4k_q5k_route_parallel (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 hidden_dim, int intermediate_dim, int n_experts, int top_k, void workspace, size_t workspace_bytes, size_t workspace_stride, ck_threadpool_t *pool)

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

static int	ck_moe_shared_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, size_t stride, ck_moe_shared_gated_workspace_fn serial_fn)

static int	ck_moe_shared_q4k_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, void(down_projection)(float , const void , const void , int, int))

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

static int	ck_moe_shared_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, ck_moe_shared_workspace_fn serial_fn)

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

static float	ck_moe_sigmoid_f32 (float x)

static float	ck_moe_silu_f32 (float x)

static int	ck_moe_size_add (size_t a, size_t b, size_t *result)

static int	ck_moe_size_mul (size_t a, size_t b, size_t *result)

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

static size_t	ck_moe_up_idx (int e, int i, int h, int intermediate_dim, int hidden_dim)

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_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	gemm_nt_q4_k_packed_meta_x8_q8_k_gemv_order (const void input, const void packed_weight, const float bias, float output, int rows, int output_dim, int input_dim)

void	gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch (const void input, const void weight, const float bias, float output, int rows, int output_dim, int input_dim)

void	gemm_q4_k_q8_k_compact_rows4 (float output, int output_stride, const void weights, const void *const input_rows[4], int rows, int output_dim, int input_dim)

void	gemm_q4_k_q8_k_packed_vnni_x8_compact_order_rows4 (float output, const void weights_packed, const void *input_q8, int rows, int output_dim, int input_dim)

void	gemm_q5_k_q8_k_compact_rows4 (float output, int output_stride, const void weights, const void *const input_rows[4], int rows, int output_dim, int input_dim)

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)

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

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_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_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_bucketed_workspace_bytes (int rows, int hidden_dim, int intermediate_dim, int n_experts, int top_k)

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_backward_f32 (const float d_output, const float hidden, const float shared_gate, const float shared_up, const float shared_down, float d_hidden, float d_routed, float d_shared_gate, float d_shared_up, float d_shared_down, int rows, int hidden_dim, int intermediate_dim)

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_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_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	pack_q4_k_to_packed_meta_x8 (const void source, void destination, int output_dim, int input_dim)

size_t	q4_k_packed_meta_x8_block_size (void)

size_t	q4_k_packed_vnni_x8_block_size (void)

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

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

Detailed Description

AXPY kernels for FP32: y = y + alpha * x.

CK-ENGINE KERNEL RULES:

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

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

Classic BLAS Level-1 operation used in MoE expert output accumulation. When gathering expert outputs: output += weight[i] * expert_output[i]

Operations:

axpy_f32: y += alpha * x (in-place)
axpy_strided_f32: strided version for non-contiguous memory
weighted_sum_f32: sum multiple vectors with weights

Definition in file axpy_kernels.c.

Typedef Documentation

◆ ck_moe_expert_workspace_fn

typedef int(* ck_moe_expert_workspace_fn) (const float *, const int *, const float *, const void *, const void *, const void *, float *, int, int, int, int, int, void *, size_t)

Definition at line 1541 of file axpy_kernels.c.

◆ ck_moe_shared_gated_workspace_fn

typedef int(* ck_moe_shared_gated_workspace_fn) (const float *, const float *, const void *, const void *, const void *, const float *, float *, int, int, int, void *, size_t)

Definition at line 2848 of file axpy_kernels.c.

◆ ck_moe_shared_workspace_fn

typedef int(* ck_moe_shared_workspace_fn) (const float *, const float *, const void *, const void *, const void *, float *, int, int, int, void *, size_t)

Definition at line 1950 of file axpy_kernels.c.

Enumeration Type Documentation

◆ anonymous enum

anonymous enum

Enumerator
CK_MOE_Q4K_Q5K_TASK_ROWS

Definition at line 2265 of file axpy_kernels.c.

2265{ CK_MOE_Q4K_Q5K_TASK_ROWS = 16 };

CK_MOE_Q4K_Q5K_TASK_ROWS

@ CK_MOE_Q4K_Q5K_TASK_ROWS

Definition axpy_kernels.c:2265

◆ ck_moe_down_kind_t

enum ck_moe_down_kind_t

Enumerator
CK_MOE_DOWN_Q4_K
CK_MOE_DOWN_Q6_K
CK_MOE_DOWN_Q5_0
CK_MOE_DOWN_Q8_0

Definition at line 1545 of file axpy_kernels.c.

             {
    CK_MOE_DOWN_Q4_K,
    CK_MOE_DOWN_Q6_K,
    CK_MOE_DOWN_Q5_0,
    CK_MOE_DOWN_Q8_0,
} ck_moe_down_kind_t;

Function Documentation

◆ 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().

◆ ck_moe_align64()

static size_t ck_moe_align64 ( size_t value )

static

Definition at line 672 of file axpy_kernels.c.

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

◆ ck_moe_bf16_round()

static float ck_moe_bf16_round ( float x )

inlinestatic

Definition at line 439 of file axpy_kernels.c.

{
    return bf16_to_float(float_to_bf16(x));
}

References bf16_to_float(), and float_to_bf16().

Referenced by moe_swiglu_packed_expert_forward_bf16(), and moe_swiglu_shared_forward_bf16_gated_row_range().

◆ ck_moe_bucket_expert_for_position()

static int ck_moe_bucket_expert_for_position	(	const int *	offsets,
		int	n_experts,
		int	position
	)

static

Definition at line 2267 of file axpy_kernels.c.

{
    int lo = 0;
    int hi = n_experts;
    while (lo < hi) {
        const int mid = lo + (hi - lo) / 2;
        if (offsets[mid + 1] <= position) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    return lo;
}

References offsets.

Referenced by ck_moe_q4k_q5k_bucket_work().

◆ ck_moe_bucket_layout()

static int ck_moe_bucket_layout	(	int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		ck_moe_q4k_q5k_bucket_layout_t *	layout
	)

static

Definition at line 2116 of file axpy_kernels.c.

{
    if (!layout || rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0 ||
        n_experts <= 0 || top_k <= 0 || top_k > n_experts ||
        hidden_dim % 256 != 0 || intermediate_dim % 256 != 0) {
        return -1;
    }
 
    memset(layout, 0, sizeof(*layout));
    layout->hidden_q8_row_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
 
    size_t gate_up_bytes = 0;
    size_t worker_bytes = 0;
    if (ck_moe_size_mul(8u * sizeof(float), (size_t)intermediate_dim,
                        &gate_up_bytes) != 0) {
        return -1;
    }
    worker_bytes = ck_moe_align64(gate_up_bytes);
    if (ck_moe_size_add(
            worker_bytes,
            ck_moe_align64(4u * ck_dtype_row_bytes(
                CK_DT_Q8_K, (size_t)hidden_dim)),
            &worker_bytes) != 0 ||
        ck_moe_size_add(
            worker_bytes,
            ck_moe_align64(4u * ck_dtype_row_bytes(
                CK_DT_Q8_K, (size_t)intermediate_dim)),
            &worker_bytes) != 0 ||
        ck_moe_size_add(
            worker_bytes,
            ck_moe_align64(4u * (size_t)hidden_dim * sizeof(float)),
            &worker_bytes) != 0) {
        return -1;
    }
    layout->worker_stride = worker_bytes;
 
    size_t cursor = 0;
    size_t bytes = 0;
    layout->hidden_q8_offset = cursor;
    if (ck_moe_size_mul((size_t)rows, layout->hidden_q8_row_bytes, &bytes) != 0 ||
        ck_moe_size_add(cursor, ck_moe_align64(bytes), &cursor) != 0) {
        return -1;
    }
 
    layout->route_rows_offset = cursor;
    if (ck_moe_size_mul((size_t)rows, (size_t)top_k, &bytes) != 0 ||
        ck_moe_size_mul(bytes, sizeof(int), &bytes) != 0 ||
        ck_moe_size_add(cursor, ck_moe_align64(bytes), &cursor) != 0) {
        return -1;
    }
 
    layout->slot_offsets_offset = cursor;
    if (ck_moe_size_mul((size_t)top_k, (size_t)n_experts + 1u, &bytes) != 0 ||
        ck_moe_size_mul(bytes, sizeof(int), &bytes) != 0 ||
        ck_moe_size_add(cursor, ck_moe_align64(bytes), &cursor) != 0) {
        return -1;
    }
 
    layout->counts_offset = cursor;
    if (ck_moe_size_mul((size_t)n_experts, sizeof(int), &bytes) != 0 ||
        ck_moe_size_add(cursor, ck_moe_align64(bytes), &cursor) != 0) {
        return -1;
    }
 
    layout->cursors_offset = cursor;
    if (ck_moe_size_mul((size_t)n_experts, sizeof(int), &bytes) != 0 ||
        ck_moe_size_add(cursor, ck_moe_align64(bytes), &cursor) != 0) {
        return -1;
    }
 
    layout->workers_offset = cursor;
    if (ck_moe_size_mul((size_t)CK_THREADPOOL_MAX_THREADS,
                        layout->worker_stride, &bytes) != 0 ||
        ck_moe_size_add(cursor, ck_moe_align64(bytes), &cursor) != 0) {
        return -1;
    }
    layout->total_bytes = cursor;
    return 0;
}

References CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), ck_moe_size_add(), ck_moe_size_mul(), and CK_THREADPOOL_MAX_THREADS.

Referenced by ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl(), and moe_swiglu_expert_q4k_q5k_bucketed_workspace_bytes().

◆ ck_moe_debug_enabled()

static int ck_moe_debug_enabled ( void )

static

Definition at line 333 of file axpy_kernels.c.

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

Referenced by ck_moe_debug_finite(), moe_relu2_expert_forward_q5_0_q5_0(), moe_relu2_expert_forward_q5_0_q8_0(), and moe_relu2_shared_forward_q5_1_q8_0().

◆ ck_moe_debug_finite()

static void ck_moe_debug_finite	(	const char *	name,
		const float *	x,
		size_t	n
	)

static

Definition at line 339 of file axpy_kernels.c.

{
    if (!ck_moe_debug_enabled() || !x) {
        return;
    }
    size_t finite = 0;
    size_t nan = 0;
    size_t inf = 0;
    float min_v = 0.0f;
    float max_v = 0.0f;
    int have = 0;
    for (size_t i = 0; i < n; ++i) {
        const float v = x[i];
        if (isnan(v)) {
            ++nan;
        } else if (isinf(v)) {
            ++inf;
        } else {
            ++finite;
            if (!have || v < min_v) min_v = v;
            if (!have || v > max_v) max_v = v;
            have = 1;
        }
    }
    fprintf(stderr,
            "[CK_DEBUG_MOE] %s finite=%zu/%zu nan=%zu inf=%zu min=%g max=%g\n",
            name,
            finite,
            n,
            nan,
            inf,
            have ? min_v : 0.0f,
            have ? max_v : 0.0f);
}

References ck_moe_debug_enabled().

Referenced by moe_relu2_expert_forward_q5_0_q5_0(), moe_relu2_expert_forward_q5_0_q8_0(), and moe_relu2_shared_forward_q5_1_q8_0().

◆ ck_moe_down_idx()

static size_t ck_moe_down_idx	(	int	e,
		int	h,
		int	i,
		int	hidden_dim,
		int	intermediate_dim
	)

inlinestatic

Definition at line 328 of file axpy_kernels.c.

{
    return ((size_t)e * (size_t)hidden_dim + (size_t)h) * (size_t)intermediate_dim + (size_t)i;
}

Referenced by moe_relu2_expert_backward_f32(), moe_relu2_expert_forward_f32(), moe_swiglu_expert_backward_f32(), moe_swiglu_expert_forward_bf16_row_range(), moe_swiglu_expert_forward_f32(), and moe_swiglu_packed_expert_forward_bf16().

◆ ck_moe_dsilu_f32()

static float ck_moe_dsilu_f32 ( float x )

inlinestatic

Definition at line 444 of file axpy_kernels.c.

{
    const float sig = ck_moe_sigmoid_f32(x);
    return sig + x * sig * (1.0f - sig);
}

References ck_moe_sigmoid_f32().

Referenced by moe_swiglu_expert_backward_f32(), and moe_swiglu_shared_backward_f32().

◆ ck_moe_llama_weighted_accumulate()

static void ck_moe_llama_weighted_accumulate	(	float *	output,
		const float *	expert_output,
		float	route_weight,
		int	n
	)

static

Definition at line 718 of file axpy_kernels.c.

{
    for (int i = 0; i < n; ++i) {
        volatile float weighted = expert_output[i] * route_weight;
        output[i] += weighted;
    }
}

Referenced by ck_moe_q4k_mixed_route_parallel(), moe_swiglu_expert_forward_q4k_q5_0_workspace(), and moe_swiglu_expert_forward_q4k_q8_0_workspace().

◆ ck_moe_q4k_llama_projection()

static void ck_moe_q4k_llama_projection	(	float *	output,
		const void *	weights,
		const void *	input_q8,
		int	output_dim,
		int	input_dim,
		void *	scratch
	)

static

Definition at line 695 of file axpy_kernels.c.

{
#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__)
    if (scratch && output_dim > 0 && input_dim > 0 &&
        output_dim % 16 == 0 && input_dim % 256 == 0) {
        pack_q4_k_to_packed_meta_x8(weights, scratch, output_dim, input_dim);
        gemm_nt_q4_k_packed_meta_x8_q8_k_gemv_order(
            input_q8, scratch, NULL, output, 1, output_dim, input_dim);
        return;
    }
    gemv_q4_k_q8_k(output, weights, input_q8, output_dim, input_dim);
#else
    (void)scratch;
    gemv_q4_k_q8_k(output, weights, input_q8, output_dim, input_dim);
#endif
}

References gemm_nt_q4_k_packed_meta_x8_q8_k_gemv_order(), gemv_q4_k_q8_k(), and pack_q4_k_to_packed_meta_x8().

Referenced by ck_moe_q4k_mixed_route_work(), moe_swiglu_expert_forward_q4k_q5_0_workspace(), and moe_swiglu_expert_forward_q4k_q8_0_workspace().

◆ ck_moe_q4k_llama_projection_scratch_bytes()

static size_t ck_moe_q4k_llama_projection_scratch_bytes	(	int	output_dim,
		int	input_dim
	)

static

Definition at line 677 of file axpy_kernels.c.

{
#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__)
    if (output_dim <= 0 || input_dim <= 0 || output_dim % 16 != 0 ||
        input_dim % 256 != 0) {
        return 0;
    }
    return ck_moe_align64(
        (size_t)((output_dim + 7) / 8) * (size_t)(input_dim / 256) *
        q4_k_packed_meta_x8_block_size());
#else
    (void)output_dim;
    (void)input_dim;
    return 0;
#endif
}

References ck_moe_align64(), and q4_k_packed_meta_x8_block_size().

Referenced by moe_swiglu_expert_q4k_q8_0_workspace_bytes().

◆ ck_moe_q4k_mixed_parallel_work()

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

static

Definition at line 1571 of file axpy_kernels.c.

{
    ck_moe_q4k_mixed_parallel_args_t *args =
        (ck_moe_q4k_mixed_parallel_args_t *)opaque;
    const int begin = (args->rows * ith) / nth;
    const int end = (args->rows * (ith + 1)) / nth;
    if (begin >= end) {
        args->status[ith] = 0;
        return;
    }
    args->status[ith] = args->serial_fn(
        args->hidden + (size_t)begin * (size_t)args->hidden_dim,
        args->indices + (size_t)begin * (size_t)args->top_k,
        args->routing_weights + (size_t)begin * (size_t)args->top_k,
        args->expert_gate,
        args->expert_up,
        args->expert_down,
        args->output + (size_t)begin * (size_t)args->hidden_dim,
        end - begin,
        args->hidden_dim,
        args->intermediate_dim,
        args->n_experts,
        args->top_k,
        args->workspace + (size_t)ith * args->workspace_stride,
        args->workspace_stride);
}

References end.

Referenced by ck_moe_q4k_mixed_parallel_workspace().

◆ ck_moe_q4k_mixed_parallel_workspace()

static int ck_moe_q4k_mixed_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,
		ck_moe_expert_workspace_fn	serial_fn,
		ck_moe_down_kind_t	down_kind
	)

static

Definition at line 1781 of file axpy_kernels.c.

{
    const int q8_0_activation =
        down_kind == CK_MOE_DOWN_Q5_0 || down_kind == CK_MOE_DOWN_Q8_0;
    const size_t stride = q8_0_activation
        ? moe_swiglu_expert_q4k_q8_0_workspace_bytes(hidden_dim, intermediate_dim)
        : moe_swiglu_expert_q4k_q5k_workspace_bytes(hidden_dim, intermediate_dim);
    if (!hidden || !indices || !routing_weights || !expert_gate || !expert_up ||
        !expert_down || !output || !workspace || !serial_fn || 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_mixed_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, down_kind);
        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 serial_fn(
            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_mixed_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,
        .serial_fn = serial_fn,
        .status = {0},
    };
    ck_threadpool_dispatch_n(
        pool, active, ck_moe_q4k_mixed_parallel_work, &args);
    for (int ith = 0; ith < active; ++ith) {
        if (args.status[ith] != 0) return args.status[ith];
    }
    return 0;
}

References CK_MOE_DOWN_Q5_0, CK_MOE_DOWN_Q8_0, ck_moe_q4k_mixed_parallel_work(), ck_moe_q4k_mixed_route_parallel(), ck_threadpool_dispatch_n(), ck_threadpool_global(), CK_THREADPOOL_MAX_THREADS, ck_threadpool_n_threads(), moe_swiglu_expert_q4k_q5k_workspace_bytes(), and moe_swiglu_expert_q4k_q8_0_workspace_bytes().

Referenced by moe_swiglu_expert_forward_q4k_q4k_parallel_workspace(), moe_swiglu_expert_forward_q4k_q5_0_parallel_workspace(), moe_swiglu_expert_forward_q4k_q6k_parallel_workspace(), and moe_swiglu_expert_forward_q4k_q8_0_parallel_workspace().

◆ ck_moe_q4k_mixed_route_parallel()

static int ck_moe_q4k_mixed_route_parallel	(	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	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes,
		size_t	workspace_stride,
		ck_threadpool_t *	pool,
		ck_moe_down_kind_t	down_kind
	)

static

Definition at line 1713 of file axpy_kernels.c.

{
    if (!pool || top_k <= 1 || top_k > CK_THREADPOOL_MAX_THREADS ||
        ck_threadpool_n_threads(pool) < top_k ||
        workspace_bytes < workspace_stride * (size_t)top_k) {
        return 1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    void *hidden_q8 = workspace;
    quantize_row_q8_k(hidden, hidden_q8, hidden_dim);
 
    ck_moe_q4k_mixed_route_args_t args = {
        .indices = indices,
        .expert_gate = expert_gate,
        .expert_up = expert_up,
        .expert_down = expert_down,
        .hidden_q8 = hidden_q8,
        .workspace = (uint8_t *)workspace,
        .workspace_stride = workspace_stride,
        .hidden_q8_bytes = hidden_q8_bytes,
        .hidden_dim = hidden_dim,
        .intermediate_dim = intermediate_dim,
        .n_experts = n_experts,
        .top_k = top_k,
        .down_kind = down_kind,
        .expert_output = {0},
        .status = {0},
    };
    ck_threadpool_dispatch_n(
        pool, top_k, ck_moe_q4k_mixed_route_work, &args);
 
    memset(output, 0, (size_t)hidden_dim * sizeof(float));
    for (int slot = 0; slot < top_k; ++slot) {
        if (args.status[slot] != 0 || !args.expert_output[slot]) {
            return args.status[slot] != 0 ? args.status[slot] : -1;
        }
        if (down_kind == CK_MOE_DOWN_Q5_0 ||
            down_kind == CK_MOE_DOWN_Q8_0) {
            ck_moe_llama_weighted_accumulate(
                output, args.expert_output[slot], routing_weights[slot],
                hidden_dim);
        } else {
            axpy_f32(
                output, args.expert_output[slot], routing_weights[slot],
                hidden_dim);
        }
    }
    return 0;
}

References axpy_f32(), CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), CK_MOE_DOWN_Q5_0, CK_MOE_DOWN_Q8_0, ck_moe_llama_weighted_accumulate(), ck_moe_q4k_mixed_route_work(), ck_threadpool_dispatch_n(), CK_THREADPOOL_MAX_THREADS, ck_threadpool_n_threads(), and quantize_row_q8_k().

Referenced by ck_moe_q4k_mixed_parallel_workspace().

◆ ck_moe_q4k_mixed_route_work()

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

static

Definition at line 1616 of file axpy_kernels.c.

{
    ck_moe_q4k_mixed_route_args_t *args =
        (ck_moe_q4k_mixed_route_args_t *)opaque;
    if (ith >= nth || ith >= args->top_k) return;
 
    const int expert = args->indices[ith];
    if (expert < 0 || expert >= args->n_experts) {
        args->status[ith] = -2;
        return;
    }
 
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)args->intermediate_dim * sizeof(float));
    const int q8_0_activation =
        args->down_kind == CK_MOE_DOWN_Q5_0 ||
        args->down_kind == CK_MOE_DOWN_Q8_0;
    const size_t act_q8_bytes = ck_moe_align64(ck_dtype_row_bytes(
        q8_0_activation ? CK_DT_Q8_0 : CK_DT_Q8_K,
        (size_t)args->intermediate_dim));
    uint8_t *cursor = args->workspace + (size_t)ith * args->workspace_stride;
    cursor += args->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;
    args->expert_output[ith] = expert_output;
    cursor += ck_moe_align64((size_t)args->hidden_dim * sizeof(float));
    void *projection_scratch = cursor;
 
    const size_t gate_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q4_K, (size_t)args->hidden_dim);
    const size_t gate_offset =
        (size_t)expert * (size_t)args->intermediate_dim * gate_row_bytes;
    if (q8_0_activation) {
        ck_moe_q4k_llama_projection(
            gate_up, (const uint8_t *)args->expert_gate + gate_offset,
            args->hidden_q8, args->intermediate_dim, args->hidden_dim,
            projection_scratch);
        ck_moe_q4k_llama_projection(
            gate_up + args->intermediate_dim,
            (const uint8_t *)args->expert_up + gate_offset,
            args->hidden_q8, args->intermediate_dim, args->hidden_dim,
            projection_scratch);
    } else {
        gemv_q4_k_q8_k(
            gate_up, (const uint8_t *)args->expert_gate + gate_offset,
            args->hidden_q8, args->intermediate_dim, args->hidden_dim);
        gemv_q4_k_q8_k(
            gate_up + args->intermediate_dim,
            (const uint8_t *)args->expert_up + gate_offset,
            args->hidden_q8, args->intermediate_dim, args->hidden_dim);
    }
    swiglu_forward_ggml(gate_up, gate_up, 1, args->intermediate_dim);
    if (q8_0_activation) {
        quantize_row_q8_0(gate_up, act_q8, args->intermediate_dim);
    } else {
        quantize_row_q8_k(gate_up, act_q8, args->intermediate_dim);
    }
 
    if (args->down_kind == CK_MOE_DOWN_Q6_K) {
        const size_t row_bytes = ck_dtype_row_bytes(
            CK_DT_Q6_K, (size_t)args->intermediate_dim);
        const size_t offset =
            (size_t)expert * (size_t)args->hidden_dim * row_bytes;
        gemv_q6_k_q8_k(
            expert_output, (const uint8_t *)args->expert_down + offset,
            act_q8, args->hidden_dim, args->intermediate_dim);
    } else if (args->down_kind == CK_MOE_DOWN_Q5_0) {
        const size_t row_bytes = ck_dtype_row_bytes(
            CK_DT_Q5_0, (size_t)args->intermediate_dim);
        const size_t offset =
            (size_t)expert * (size_t)args->hidden_dim * row_bytes;
        gemv_q5_0_q8_0(
            expert_output, (const uint8_t *)args->expert_down + offset,
            act_q8, args->hidden_dim, args->intermediate_dim);
    } else if (args->down_kind == CK_MOE_DOWN_Q8_0) {
        const size_t row_bytes = ck_dtype_row_bytes(
            CK_DT_Q8_0, (size_t)args->intermediate_dim);
        const size_t offset =
            (size_t)expert * (size_t)args->hidden_dim * row_bytes;
        gemv_q8_0_q8_0(
            expert_output, (const uint8_t *)args->expert_down + offset,
            act_q8, args->hidden_dim, args->intermediate_dim);
    } else {
        const size_t row_bytes = ck_dtype_row_bytes(
            CK_DT_Q4_K, (size_t)args->intermediate_dim);
        const size_t offset =
            (size_t)expert * (size_t)args->hidden_dim * row_bytes;
        gemv_q4_k_q8_k(
            expert_output, (const uint8_t *)args->expert_down + offset,
            act_q8, args->hidden_dim, args->intermediate_dim);
    }
    args->status[ith] = 0;
}

References CK_DT_Q4_K, CK_DT_Q5_0, CK_DT_Q6_K, CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), CK_MOE_DOWN_Q5_0, CK_MOE_DOWN_Q6_K, CK_MOE_DOWN_Q8_0, ck_moe_q4k_llama_projection(), gemv_q4_k_q8_k(), gemv_q5_0_q8_0(), gemv_q6_k_q8_k(), gemv_q8_0_q8_0(), quantize_row_q8_0(), quantize_row_q8_k(), and swiglu_forward_ggml().

Referenced by ck_moe_q4k_mixed_route_parallel().

◆ ck_moe_q4k_q5k_bucket_work()

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

static

Definition at line 2284 of file axpy_kernels.c.

{
    (void)nth;
    ck_moe_q4k_q5k_bucket_work_t *args =
        (ck_moe_q4k_q5k_bucket_work_t *)opaque;
    uint8_t *cursor = args->workers + (size_t)ith * args->worker_stride;
    float *gate_up = (float *)cursor;
    cursor += ck_moe_align64(
        8u * (size_t)args->intermediate_dim * sizeof(float));
    uint8_t *hidden_q8_batch = cursor;
    const size_t hidden_q8_batch_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q8_K, (size_t)args->hidden_dim);
    cursor += ck_moe_align64(4u * hidden_q8_batch_row_bytes);
    void *act_q8 = cursor;
    cursor += ck_moe_align64(
        4u * ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)args->intermediate_dim));
    float *expert_output = (float *)cursor;
    const size_t act_q8_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q8_K, (size_t)args->intermediate_dim);
 
    for (;;) {
        const int task = atomic_fetch_add_explicit(
            &args->next_task, 1, memory_order_relaxed);
        if (task >= args->total_tasks) break;
 
        int position = task * CK_MOE_Q4K_Q5K_TASK_ROWS;
        const int task_end = position + CK_MOE_Q4K_Q5K_TASK_ROWS <
                args->bucket_offsets[args->n_experts]
            ? position + CK_MOE_Q4K_Q5K_TASK_ROWS
            : args->bucket_offsets[args->n_experts];
        int expert = ck_moe_bucket_expert_for_position(
            args->bucket_offsets, args->n_experts, position);
 
        while (position < task_end && expert < args->n_experts) {
            const int expert_end = args->bucket_offsets[expert + 1];
            const int segment_end = expert_end < task_end
                ? expert_end : task_end;
            const uint8_t *gate = args->gate_base +
                (size_t)expert * args->q4_expert_stride;
            const uint8_t *up = args->up_base +
                (size_t)expert * args->q4_expert_stride;
            const uint8_t *gate_packed = args->gate_packed_base
                ? args->gate_packed_base +
                    (size_t)expert * args->q4_packed_expert_stride
                : NULL;
            const uint8_t *up_packed = args->up_packed_base
                ? args->up_packed_base +
                    (size_t)expert * args->q4_packed_expert_stride
                : NULL;
            const uint8_t *down = args->down_base +
                (size_t)expert * args->q5_expert_stride;
 
            for (int i = position; i < segment_end; i += 4) {
                const int batch_rows = segment_end - i < 4
                    ? segment_end - i : 4;
                const void *hidden_rows[4] = {NULL, NULL, NULL, NULL};
                const void *activation_rows[4] = {NULL, NULL, NULL, NULL};
                int output_rows[4] = {0, 0, 0, 0};
                for (int batch_row = 0; batch_row < batch_rows; ++batch_row) {
                    const int row = args->bucket_rows[i + batch_row];
                    output_rows[batch_row] = row;
                    hidden_rows[batch_row] = args->hidden_q8 +
                        (size_t)row * args->hidden_q8_row_bytes;
                    if (gate_packed && up_packed) {
                        memcpy(
                            hidden_q8_batch +
                                (size_t)batch_row * hidden_q8_batch_row_bytes,
                            hidden_rows[batch_row],
                            args->hidden_q8_row_bytes);
                    }
                }
                for (int batch_row = batch_rows; batch_row < 4; ++batch_row) {
                    hidden_rows[batch_row] = hidden_rows[0];
                }
 
                if (gate_packed && up_packed) {
                    float *up_rows = gate_up +
                        4u * (size_t)args->intermediate_dim;
                    gemm_q4_k_q8_k_packed_vnni_x8_compact_order_rows4(
                        gate_up, gate_packed, hidden_q8_batch,
                        batch_rows, args->intermediate_dim,
                        args->hidden_dim);
                    gemm_q4_k_q8_k_packed_vnni_x8_compact_order_rows4(
                        up_rows, up_packed, hidden_q8_batch,
                        batch_rows, args->intermediate_dim,
                        args->hidden_dim);
                    swiglu_forward_ggml_split(
                        gate_up, up_rows, gate_up, batch_rows,
                        args->intermediate_dim);
                } else {
                    const int gate_up_stride = 2 * args->intermediate_dim;
                    gemm_q4_k_q8_k_compact_rows4(
                        gate_up, gate_up_stride, gate, hidden_rows, batch_rows,
                        args->intermediate_dim, args->hidden_dim);
                    gemm_q4_k_q8_k_compact_rows4(
                        gate_up + args->intermediate_dim, gate_up_stride, up,
                        hidden_rows, batch_rows,
                        args->intermediate_dim, args->hidden_dim);
                    swiglu_forward_ggml(
                        gate_up, gate_up, batch_rows, args->intermediate_dim);
                }
                for (int batch_row = 0; batch_row < batch_rows; ++batch_row) {
                    void *activation = (uint8_t *)act_q8 +
                        (size_t)batch_row * act_q8_row_bytes;
                    quantize_row_q8_k(
                        gate_up + (size_t)batch_row *
                            (size_t)args->intermediate_dim,
                        activation, args->intermediate_dim);
                    activation_rows[batch_row] = activation;
                }
                for (int batch_row = batch_rows; batch_row < 4; ++batch_row) {
                    activation_rows[batch_row] = activation_rows[0];
                }
                gemm_q5_k_q8_k_compact_rows4(
                    expert_output, args->hidden_dim, down, activation_rows,
                    batch_rows, args->hidden_dim, args->intermediate_dim);
 
                for (int batch_row = 0; batch_row < batch_rows; ++batch_row) {
                    const int row = output_rows[batch_row];
                    const size_t route_index = (size_t)row *
                        (size_t)args->top_k + (size_t)args->slot;
                    axpy_f32(
                        args->output + (size_t)row *
                            (size_t)args->hidden_dim,
                        expert_output + (size_t)batch_row *
                            (size_t)args->hidden_dim,
                        args->routing_weights[route_index], args->hidden_dim);
                }
            }
            position = segment_end;
            ++expert;
        }
    }
}

References axpy_f32(), CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), ck_moe_bucket_expert_for_position(), CK_MOE_Q4K_Q5K_TASK_ROWS, gemm_q4_k_q8_k_compact_rows4(), gemm_q4_k_q8_k_packed_vnni_x8_compact_order_rows4(), gemm_q5_k_q8_k_compact_rows4(), quantize_row_q8_k(), swiglu_forward_ggml(), and swiglu_forward_ggml_split().

Referenced by ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl().

◆ ck_moe_q4k_q5k_parallel_work()

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

static

Definition at line 1317 of file axpy_kernels.c.

{
    ck_moe_q4k_q5k_parallel_args_t *args =
        (ck_moe_q4k_q5k_parallel_args_t *)opaque;
    const int begin = (args->rows * ith) / nth;
    const int end = (args->rows * (ith + 1)) / nth;
    if (begin >= end) {
        args->status[ith] = 0;
        return;
    }
    args->status[ith] = moe_swiglu_expert_forward_q4k_q5k_workspace(
        args->hidden + (size_t)begin * (size_t)args->hidden_dim,
        args->indices + (size_t)begin * (size_t)args->top_k,
        args->routing_weights + (size_t)begin * (size_t)args->top_k,
        args->expert_gate,
        args->expert_up,
        args->expert_down,
        args->output + (size_t)begin * (size_t)args->hidden_dim,
        end - begin,
        args->hidden_dim,
        args->intermediate_dim,
        args->n_experts,
        args->top_k,
        args->workspace + (size_t)ith * args->workspace_stride,
        args->workspace_stride);
}

References end, and moe_swiglu_expert_forward_q4k_q5k_workspace().

Referenced by moe_swiglu_expert_forward_q4k_q5k_parallel_workspace().

◆ ck_moe_q4k_q5k_quantize_work()

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

static

Definition at line 2225 of file axpy_kernels.c.

{
    ck_moe_q4k_q5k_quantize_args_t *args =
        (ck_moe_q4k_q5k_quantize_args_t *)opaque;
    const int begin = (args->rows * ith) / nth;
    const int end = (args->rows * (ith + 1)) / nth;
    for (int row = begin; row < end; ++row) {
        quantize_row_q8_k(
            args->hidden + (size_t)row * (size_t)args->hidden_dim,
            args->hidden_q8 + (size_t)row * args->hidden_q8_row_bytes,
            args->hidden_dim);
    }
}

References end, and quantize_row_q8_k().

Referenced by ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl().

◆ ck_moe_q4k_q5k_route_parallel()

static int ck_moe_q4k_q5k_route_parallel	(	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	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k,
		void *	workspace,
		size_t	workspace_bytes,
		size_t	workspace_stride,
		ck_threadpool_t *	pool
	)

static

Definition at line 1412 of file axpy_kernels.c.

{
    if (!pool || top_k <= 1 || top_k > CK_THREADPOOL_MAX_THREADS ||
        ck_threadpool_n_threads(pool) < top_k ||
        workspace_bytes < workspace_stride * (size_t)top_k) {
        return 1;
    }
 
    const size_t hidden_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim));
    void *hidden_q8 = workspace;
    quantize_row_q8_k(hidden, hidden_q8, hidden_dim);
 
    ck_moe_q4k_q5k_route_args_t args = {
        .indices = indices,
        .expert_gate = expert_gate,
        .expert_up = expert_up,
        .expert_down = expert_down,
        .hidden_q8 = hidden_q8,
        .workspace = (uint8_t *)workspace,
        .workspace_stride = workspace_stride,
        .hidden_q8_bytes = hidden_q8_bytes,
        .hidden_dim = hidden_dim,
        .intermediate_dim = intermediate_dim,
        .n_experts = n_experts,
        .top_k = top_k,
        .expert_output = {0},
        .status = {0},
    };
    ck_threadpool_dispatch_n(
        pool, top_k, ck_moe_q4k_q5k_route_work, &args);
 
    memset(output, 0, (size_t)hidden_dim * sizeof(float));
    for (int slot = 0; slot < top_k; ++slot) {
        if (args.status[slot] != 0 || !args.expert_output[slot]) {
            return args.status[slot] != 0 ? args.status[slot] : -1;
        }
        axpy_f32(
            output, args.expert_output[slot], routing_weights[slot], hidden_dim);
    }
    return 0;
}

References axpy_f32(), CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), ck_moe_q4k_q5k_route_work(), ck_threadpool_dispatch_n(), CK_THREADPOOL_MAX_THREADS, ck_threadpool_n_threads(), and quantize_row_q8_k().

Referenced by moe_swiglu_expert_forward_q4k_q5k_parallel_workspace().

◆ ck_moe_q4k_q5k_route_work()

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

static

Definition at line 1361 of file axpy_kernels.c.

{
    ck_moe_q4k_q5k_route_args_t *args =
        (ck_moe_q4k_q5k_route_args_t *)opaque;
    if (ith >= nth || ith >= args->top_k) return;
 
    const int expert = args->indices[ith];
    if (expert < 0 || expert >= args->n_experts) {
        args->status[ith] = -2;
        return;
    }
 
    const size_t gate_up_bytes = ck_moe_align64(
        2u * (size_t)args->intermediate_dim * sizeof(float));
    const size_t act_q8_bytes = ck_moe_align64(
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)args->intermediate_dim));
    uint8_t *cursor = args->workspace + (size_t)ith * args->workspace_stride;
    cursor += args->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;
    args->expert_output[ith] = expert_output;
 
    const size_t q4_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q4_K, (size_t)args->hidden_dim);
    const size_t q5_row_bytes = ck_dtype_row_bytes(
        CK_DT_Q5_K, (size_t)args->intermediate_dim);
    const size_t q4_expert_offset =
        (size_t)expert * (size_t)args->intermediate_dim * q4_row_bytes;
    const size_t q5_expert_offset =
        (size_t)expert * (size_t)args->hidden_dim * q5_row_bytes;
 
    gemv_q4_k_q8_k(
        gate_up,
        (const uint8_t *)args->expert_gate + q4_expert_offset,
        args->hidden_q8, args->intermediate_dim, args->hidden_dim);
    gemv_q4_k_q8_k(
        gate_up + args->intermediate_dim,
        (const uint8_t *)args->expert_up + q4_expert_offset,
        args->hidden_q8, args->intermediate_dim, args->hidden_dim);
    swiglu_forward_ggml(gate_up, gate_up, 1, args->intermediate_dim);
    quantize_row_q8_k(gate_up, act_q8, args->intermediate_dim);
    gemv_q5_k_q8_k(
        expert_output,
        (const uint8_t *)args->expert_down + q5_expert_offset,
        act_q8, args->hidden_dim, args->intermediate_dim);
    args->status[ith] = 0;
}

References 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(), quantize_row_q8_k(), and swiglu_forward_ggml().

Referenced by ck_moe_q4k_q5k_route_parallel().

◆ ck_moe_shared_gated_parallel_workspace()

static int ck_moe_shared_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,
		size_t	stride,
		ck_moe_shared_gated_workspace_fn	serial_fn
	)

static

Definition at line 2897 of file axpy_kernels.c.

{
    if (!hidden || !shared_gate || !shared_up || !shared_down ||
        !shared_gate_input || !output || !workspace || !serial_fn ||
        stride == 0 || rows <= 0) {
        return -1;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    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 serial_fn(
            hidden, routed, shared_gate, shared_up, shared_down,
            shared_gate_input, output, rows, hidden_dim, intermediate_dim,
            workspace, stride);
    }
 
    ck_moe_shared_q8_0_parallel_args_t args = {
        .hidden = hidden,
        .routed = routed,
        .shared_gate = shared_gate,
        .shared_up = shared_up,
        .shared_down = shared_down,
        .shared_gate_input = shared_gate_input,
        .output = output,
        .rows = rows,
        .hidden_dim = hidden_dim,
        .intermediate_dim = intermediate_dim,
        .workspace = (uint8_t *)workspace,
        .workspace_stride = stride,
        .serial_fn = serial_fn,
        .status = {0},
    };
    ck_threadpool_dispatch_n(
        pool, active, ck_moe_shared_q8_0_parallel_work, &args);
    for (int ith = 0; ith < active; ++ith) {
        if (args.status[ith] != 0) return args.status[ith];
    }
    return 0;
}

References ck_moe_shared_q8_0_parallel_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), CK_THREADPOOL_MAX_THREADS, and ck_threadpool_n_threads().

Referenced by moe_swiglu_shared_forward_q8_0_gated_parallel_workspace().

◆ ck_moe_shared_q4k_gated_workspace()

static int ck_moe_shared_q4k_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,
		void()(float , const void , const void , int, int)	down_projection
	)

static

Definition at line 2709 of file axpy_kernels.c.

{
    const size_t required = moe_swiglu_shared_q4k_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;
    }
 
    enum { CK_SHARED_Q4K_BATCH_ROWS = 4 };
    const size_t hidden_q8_stride =
        ck_dtype_row_bytes(CK_DT_Q8_K, (size_t)hidden_dim);
    const size_t activation_q8_stride =
        ck_dtype_row_bytes(CK_DT_Q8_0, (size_t)intermediate_dim);
    const size_t hidden_q8_bytes = ck_moe_align64(
        CK_SHARED_Q4K_BATCH_ROWS * hidden_q8_stride);
    const size_t projection_bytes = ck_moe_align64(
        CK_SHARED_Q4K_BATCH_ROWS * (size_t)intermediate_dim * sizeof(float));
    const size_t activation_q8_bytes = ck_moe_align64(
        CK_SHARED_Q4K_BATCH_ROWS * activation_q8_stride);
    uint8_t *cursor = (uint8_t *)workspace;
    void *hidden_q8 = cursor;
    cursor += hidden_q8_bytes;
    float *gate_values = (float *)cursor;
    cursor += projection_bytes;
    float *up_values = (float *)cursor;
    cursor += projection_bytes;
    void *activation_q8 = cursor;
    cursor += activation_q8_bytes;
    float *shared_output = (float *)cursor;
 
    /* llama.cpp evaluates Q4_K prefill projections in complete four-row GEMM
     * groups, then uses the GEMV reduction order only for the residual tail.
     * Keeping that boundary here matters: tiny per-row reduction differences
     * can cross the following Q8_0 activation quantization threshold. */
    for (int row0 = 0; row0 < rows; row0 += CK_SHARED_Q4K_BATCH_ROWS) {
        int batch_rows = rows - row0;
        if (batch_rows > CK_SHARED_Q4K_BATCH_ROWS) {
            batch_rows = CK_SHARED_Q4K_BATCH_ROWS;
        }
        for (int local_row = 0; local_row < batch_rows; ++local_row) {
            const float *x = hidden +
                (size_t)(row0 + local_row) * (size_t)hidden_dim;
            quantize_row_q8_k(
                x,
                (uint8_t *)hidden_q8 +
                    (size_t)local_row * hidden_q8_stride,
                hidden_dim);
        }
        gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch(
            hidden_q8, shared_gate, NULL, gate_values,
            batch_rows, intermediate_dim, hidden_dim);
        gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch(
            hidden_q8, shared_up, NULL, up_values,
            batch_rows, intermediate_dim, hidden_dim);
        swiglu_forward_ggml_split(
            gate_values, up_values, gate_values, batch_rows, intermediate_dim);
 
        for (int local_row = 0; local_row < batch_rows; ++local_row) {
            const int row = row0 + local_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;
            void *activation_q8_row = (uint8_t *)activation_q8 +
                (size_t)local_row * activation_q8_stride;
            float *shared_output_row = shared_output +
                (size_t)local_row * (size_t)hidden_dim;
 
            quantize_row_q8_0(
                gate_values + (size_t)local_row * (size_t)intermediate_dim,
                activation_q8_row, intermediate_dim);
            down_projection(
                shared_output_row, shared_down, activation_q8_row,
                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_row[h] * gate_scale;
                output_row[h] = routed_value + gated_shared;
            }
        }
    }
    return 0;
}

References CK_DT_Q8_0, CK_DT_Q8_K, ck_dtype_row_bytes(), ck_moe_align64(), gemm_nt_f32_llama_production(), gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch(), moe_swiglu_shared_q4k_q8_0_gated_workspace_bytes(), quantize_row_q8_0(), quantize_row_q8_k(), and swiglu_forward_ggml_split().

Referenced by moe_swiglu_shared_forward_q4k_q5_0_gated_workspace(), and moe_swiglu_shared_forward_q4k_q8_0_gated_workspace().

◆ ck_moe_shared_q4k_parallel_work()

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

static

Definition at line 1970 of file axpy_kernels.c.

{
    ck_moe_shared_q4k_parallel_args_t *args =
        (ck_moe_shared_q4k_parallel_args_t *)opaque;
    const int begin = (args->rows * ith) / nth;
    const int end = (args->rows * (ith + 1)) / nth;
    if (begin >= end) {
        args->status[ith] = 0;
        return;
    }
    args->status[ith] = args->serial_fn(
        args->hidden + (size_t)begin * (size_t)args->hidden_dim,
        args->routed
            ? args->routed + (size_t)begin * (size_t)args->hidden_dim
            : NULL,
        args->shared_gate,
        args->shared_up,
        args->shared_down,
        args->output + (size_t)begin * (size_t)args->hidden_dim,
        end - begin,
        args->hidden_dim,
        args->intermediate_dim,
        args->workspace + (size_t)ith * args->workspace_stride,
        args->workspace_stride);
}

References end.

Referenced by ck_moe_shared_q4k_parallel_workspace().

◆ ck_moe_shared_q4k_parallel_workspace()

static int ck_moe_shared_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,
		ck_moe_shared_workspace_fn	serial_fn
	)

static

Definition at line 1996 of file axpy_kernels.c.

{
    const size_t stride = moe_swiglu_expert_q4k_q5k_workspace_bytes(
        hidden_dim, intermediate_dim);
    if (!hidden || !shared_gate || !shared_up || !shared_down || !output ||
        !workspace || !serial_fn || stride == 0 || rows <= 0) {
        return -1;
    }
    ck_threadpool_t *pool = ck_threadpool_global();
    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 serial_fn(
            hidden, routed, shared_gate, shared_up, shared_down, output, rows,
            hidden_dim, intermediate_dim, workspace, stride);
    }
 
    ck_moe_shared_q4k_parallel_args_t args = {
        .hidden = hidden,
        .routed = routed,
        .shared_gate = shared_gate,
        .shared_up = shared_up,
        .shared_down = shared_down,
        .output = output,
        .rows = rows,
        .hidden_dim = hidden_dim,
        .intermediate_dim = intermediate_dim,
        .workspace = (uint8_t *)workspace,
        .workspace_stride = stride,
        .serial_fn = serial_fn,
        .status = {0},
    };
    ck_threadpool_dispatch_n(
        pool, active, ck_moe_shared_q4k_parallel_work, &args);
    for (int ith = 0; ith < active; ++ith) {
        if (args.status[ith] != 0) return args.status[ith];
    }
    return 0;
}

References ck_moe_shared_q4k_parallel_work(), ck_threadpool_dispatch_n(), ck_threadpool_global(), CK_THREADPOOL_MAX_THREADS, ck_threadpool_n_threads(), and moe_swiglu_expert_q4k_q5k_workspace_bytes().

Referenced by moe_swiglu_shared_forward_q4k_q4k_parallel_workspace(), and moe_swiglu_shared_forward_q4k_q6k_parallel_workspace().

◆ ck_moe_shared_q8_0_parallel_work()

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

static

Definition at line 2869 of file axpy_kernels.c.

{
    ck_moe_shared_q8_0_parallel_args_t *args =
        (ck_moe_shared_q8_0_parallel_args_t *)opaque;
    const int begin = (args->rows * ith) / nth;
    const int end = (args->rows * (ith + 1)) / nth;
    if (begin >= end) {
        args->status[ith] = 0;
        return;
    }
 
    args->status[ith] = args->serial_fn(
        args->hidden + (size_t)begin * (size_t)args->hidden_dim,
        args->routed
            ? args->routed + (size_t)begin * (size_t)args->hidden_dim
            : NULL,
        args->shared_gate,
        args->shared_up,
        args->shared_down,
        args->shared_gate_input,
        args->output + (size_t)begin * (size_t)args->hidden_dim,
        end - begin,
        args->hidden_dim,
        args->intermediate_dim,
        args->workspace + (size_t)ith * args->workspace_stride,
        args->workspace_stride);
}

References end.

Referenced by ck_moe_shared_gated_parallel_workspace().

◆ ck_moe_sigmoid_f32()

static float ck_moe_sigmoid_f32 ( float x )

inlinestatic

Definition at line 429 of file axpy_kernels.c.

{
    return 1.0f / (1.0f + expf(-x));
}

Referenced by ck_moe_dsilu_f32(), ck_moe_silu_f32(), and moe_swiglu_shared_forward_bf16_gated_row_range().

◆ ck_moe_silu_f32()

static float ck_moe_silu_f32 ( float x )

inlinestatic

Definition at line 434 of file axpy_kernels.c.

{
    return x * ck_moe_sigmoid_f32(x);
}

References ck_moe_sigmoid_f32().

Referenced by farskip_swiglu_shared_combine_bf16_row_range(), moe_swiglu_expert_backward_f32(), moe_swiglu_expert_forward_bf16_row_range(), moe_swiglu_expert_forward_f32(), moe_swiglu_shared_backward_f32(), moe_swiglu_shared_forward_bf16_gated_row_range(), moe_swiglu_shared_forward_bf16_row_range(), and moe_swiglu_shared_forward_f32().

◆ ck_moe_size_add()

static int ck_moe_size_add	(	size_t	a,
		size_t	b,
		size_t *	result
	)

static

Definition at line 2102 of file axpy_kernels.c.

{
    if (!result || a > SIZE_MAX - b) return -1;
    *result = a + b;
    return 0;
}

Referenced by ck_moe_bucket_layout().

◆ ck_moe_size_mul()

static int ck_moe_size_mul	(	size_t	a,
		size_t	b,
		size_t *	result
	)

static

Definition at line 2109 of file axpy_kernels.c.

{
    if (!result || (a != 0 && b > SIZE_MAX / a)) return -1;
    *result = a * b;
    return 0;
}

Referenced by ck_moe_bucket_layout().

◆ ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl()

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

static

Definition at line 2419 of file axpy_kernels.c.

{
    ck_moe_q4k_q5k_bucket_layout_t layout;
    if (!hidden || !indices || !routing_weights || !expert_gate || !expert_up ||
        !expert_down || !output || !workspace ||
        ((expert_gate_packed == NULL) != (expert_up_packed == NULL)) ||
        ck_moe_bucket_layout(rows, hidden_dim, intermediate_dim, n_experts,
                             top_k, &layout) != 0 ||
        workspace_bytes < layout.total_bytes) {
        return -1;
    }
 
    uint8_t *base = (uint8_t *)workspace;
    uint8_t *hidden_q8 = base + layout.hidden_q8_offset;
    int *route_rows = (int *)(base + layout.route_rows_offset);
    int *slot_offsets = (int *)(base + layout.slot_offsets_offset);
    int *counts = (int *)(base + layout.counts_offset);
    int *cursors = (int *)(base + layout.cursors_offset);
    uint8_t *workers = base + layout.workers_offset;
 
    for (int slot = 0; slot < top_k; ++slot) {
        memset(counts, 0, (size_t)n_experts * sizeof(*counts));
        for (int row = 0; row < rows; ++row) {
            const int expert = indices[(size_t)row * (size_t)top_k +
                                       (size_t)slot];
            if (expert < 0 || expert >= n_experts) return -2;
            counts[expert] += 1;
        }
 
        int *offsets = slot_offsets + (size_t)slot * ((size_t)n_experts + 1u);
        offsets[0] = 0;
        for (int expert = 0; expert < n_experts; ++expert) {
            offsets[expert + 1] = offsets[expert] + counts[expert];
            cursors[expert] = offsets[expert];
        }
        int *rows_for_slot = route_rows + (size_t)slot * (size_t)rows;
        for (int row = 0; row < rows; ++row) {
            const int expert = indices[(size_t)row * (size_t)top_k +
                                       (size_t)slot];
            rows_for_slot[cursors[expert]++] = row;
        }
    }
 
    memset(output, 0, (size_t)rows * (size_t)hidden_dim * sizeof(float));
    ck_threadpool_t *pool = ck_threadpool_global();
    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;
    if (active < 1) active = 1;
 
    ck_moe_q4k_q5k_quantize_args_t quantize_args = {
        .hidden = hidden,
        .hidden_q8 = hidden_q8,
        .rows = rows,
        .hidden_dim = hidden_dim,
        .hidden_q8_row_bytes = layout.hidden_q8_row_bytes,
    };
    if (active > 1 && pool) {
        ck_threadpool_dispatch_n(
            pool, active, ck_moe_q4k_q5k_quantize_work, &quantize_args);
    } else {
        ck_moe_q4k_q5k_quantize_work(0, 1, &quantize_args);
    }
 
    const size_t q4_expert_stride = (size_t)intermediate_dim *
        ck_dtype_row_bytes(CK_DT_Q4_K, (size_t)hidden_dim);
    const size_t q4_packed_expert_stride =
        (size_t)((intermediate_dim + 7) / 8) *
        (size_t)(hidden_dim / 256) * q4_k_packed_vnni_x8_block_size();
    const size_t q5_expert_stride = (size_t)hidden_dim *
        ck_dtype_row_bytes(CK_DT_Q5_K, (size_t)intermediate_dim);
    for (int slot = 0; slot < top_k; ++slot) {
        ck_moe_q4k_q5k_bucket_work_t args = {
            .bucket_rows = route_rows + (size_t)slot * (size_t)rows,
            .bucket_offsets = slot_offsets +
                (size_t)slot * ((size_t)n_experts + 1u),
            .routing_weights = routing_weights,
            .hidden_q8 = hidden_q8,
            .gate_base = (const uint8_t *)expert_gate,
            .up_base = (const uint8_t *)expert_up,
            .gate_packed_base = (const uint8_t *)expert_gate_packed,
            .up_packed_base = (const uint8_t *)expert_up_packed,
            .down_base = (const uint8_t *)expert_down,
            .output = output,
            .workers = workers,
            .worker_stride = layout.worker_stride,
            .hidden_q8_row_bytes = layout.hidden_q8_row_bytes,
            .q4_expert_stride = q4_expert_stride,
            .q4_packed_expert_stride = q4_packed_expert_stride,
            .q5_expert_stride = q5_expert_stride,
            .hidden_dim = hidden_dim,
            .intermediate_dim = intermediate_dim,
            .n_experts = n_experts,
            .top_k = top_k,
            .slot = slot,
            .total_tasks = (rows + CK_MOE_Q4K_Q5K_TASK_ROWS - 1) /
                CK_MOE_Q4K_Q5K_TASK_ROWS,
        };
        atomic_init(&args.next_task, 0);
        int task_threads = active;
        if (task_threads > args.total_tasks) task_threads = args.total_tasks;
        if (active > 1 && pool) {
            ck_threadpool_dispatch_n(
                pool, task_threads, ck_moe_q4k_q5k_bucket_work, &args);
        } else {
            ck_moe_q4k_q5k_bucket_work(0, 1, &args);
        }
    }
    return 0;
}

References CK_DT_Q4_K, CK_DT_Q5_K, ck_dtype_row_bytes(), ck_moe_bucket_layout(), ck_moe_q4k_q5k_bucket_work(), ck_moe_q4k_q5k_quantize_work(), CK_MOE_Q4K_Q5K_TASK_ROWS, ck_threadpool_dispatch_n(), ck_threadpool_global(), CK_THREADPOOL_MAX_THREADS, ck_threadpool_n_threads(), offsets, and q4_k_packed_vnni_x8_block_size().

Referenced by moe_swiglu_expert_forward_q4k_q5k_bucketed_prepared_workspace(), and moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace().

◆ ck_moe_up_idx()

static size_t ck_moe_up_idx	(	int	e,
		int	i,
		int	h,
		int	intermediate_dim,
		int	hidden_dim
	)

inlinestatic

Definition at line 323 of file axpy_kernels.c.

{
    return ((size_t)e * (size_t)intermediate_dim + (size_t)i) * (size_t)hidden_dim + (size_t)h;
}

Referenced by moe_relu2_expert_backward_f32(), moe_relu2_expert_forward_f32(), moe_swiglu_expert_backward_f32(), moe_swiglu_expert_forward_bf16_row_range(), and moe_swiglu_expert_forward_f32().

◆ 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_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().

◆ gemm_nt_q4_k_packed_meta_x8_q8_k_gemv_order()

void gemm_nt_q4_k_packed_meta_x8_q8_k_gemv_order	(	const void *	input,
		const void *	packed_weight,
		const float *	bias,
		float *	output,
		int	rows,
		int	output_dim,
		int	input_dim
	)

Definition at line 2055 of file gemm_kernels_q4k_q8k_vnni.c.

{
    if (!A_q8 || !B_packed_x8 || !C || M <= 0 || N <= 0 || K <= 0 || (K % QK_K) != 0) {
        return;
    }
    const block_q8_K *A = (const block_q8_K *)A_q8;
    const block_q4_K_packed_meta_x8 *W = (const block_q4_K_packed_meta_x8 *)B_packed_x8;
    const int blocks_per_row = K / QK_K;
    const int groups = (N + 7) / 8;
    for (int m = 0; m < M; ++m) {
        const block_q8_K *a_row = A + (size_t)m * (size_t)blocks_per_row;
        float *c_row = C + (size_t)m * (size_t)N;
        for (int g = 0; g < groups; ++g) {
            const int n0 = g * 8;
            const int active = (n0 + 8 <= N) ? 8 : (N - n0);
            float acc[8] = {0};
            float acc_min[8] = {0};
            for (int b = 0; b < blocks_per_row; ++b) {
                const block_q4_K_packed_meta_x8 *w_group =
                        W + (size_t)g * (size_t)blocks_per_row + (size_t)b;
                accum_q4_k_packed_meta_x8_q8_k_gemv_block(
                        acc, acc_min, w_group, active, &a_row[b]);
            }
            float values[8];
#if defined(__AVX2__)
            _mm256_storeu_ps(values, _mm256_sub_ps(
                    _mm256_loadu_ps(acc), _mm256_loadu_ps(acc_min)));
#else
            for (int lane = 0; lane < active; ++lane) values[lane] = acc[lane] - acc_min[lane];
#endif
            for (int lane = 0; lane < active; ++lane) {
                c_row[n0 + lane] = values[lane] + (bias ? bias[n0 + lane] : 0.0f);
            }
        }
    }
}

References accum_q4_k_packed_meta_x8_q8_k_gemv_block(), C, and QK_K.

Referenced by ck_moe_q4k_llama_projection().

◆ gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch()

void gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch	(	const void *	input,
		const void *	weight,
		const float *	bias,
		float *	output,
		int	rows,
		int	output_dim,
		int	input_dim
	)

Referenced by ck_moe_shared_q4k_gated_workspace().

◆ gemm_q4_k_q8_k_compact_rows4()

void gemm_q4_k_q8_k_compact_rows4	(	float *	output,
		int	output_stride,
		const void *	weights,
		const void *const	input_rows[4],
		int	rows,
		int	output_dim,
		int	input_dim
	)

Definition at line 3615 of file gemm_kernels_q4k_q8k_vnni.c.

{
    if (!output || !weights || !input_rows || rows <= 0 || rows > 4 ||
        output_stride < output_dim || output_dim <= 0 || input_dim <= 0 ||
        (input_dim % QK_K) != 0) {
        return;
    }
    for (int row = 0; row < rows; ++row) {
        if (!input_rows[row]) return;
    }
 
#if defined(__AVX512VNNI__) && defined(__AVX512VL__)
    const block_q4_K *blocks = (const block_q4_K *)weights;
    const int blocks_per_row = input_dim / QK_K;
    const block_q8_K *inputs[4] = {
        (const block_q8_K *)input_rows[0],
        (const block_q8_K *)input_rows[rows > 1 ? 1 : 0],
        (const block_q8_K *)input_rows[rows > 2 ? 2 : 0],
        (const block_q8_K *)input_rows[rows > 3 ? 3 : 0],
    };
    for (int n = 0; n < output_dim; ++n) {
        const block_q4_K *weight_row =
            blocks + (size_t)n * (size_t)blocks_per_row;
        float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f};
        for (int block = 0; block < blocks_per_row; ++block) {
            const block_q8_K *block_rows[4] = {
                &inputs[0][block], &inputs[1][block],
                &inputs[2][block], &inputs[3][block],
            };
            float block_sums[4];
            dot_q4_k_q8_k_vnni_block_rows4(
                &weight_row[block], block_rows, rows, block_sums);
            for (int row = 0; row < rows; ++row) {
                sums[row] += block_sums[row];
            }
        }
        for (int row = 0; row < rows; ++row) {
            output[(size_t)row * (size_t)output_stride + (size_t)n] = sums[row];
        }
    }
#else
    for (int row = 0; row < rows; ++row) {
        gemv_q4_k_q8_k(
            output + (size_t)row * (size_t)output_stride,
            weights, input_rows[row], output_dim, input_dim);
    }
#endif
}

References gemv_q4_k_q8_k(), and QK_K.

Referenced by ck_moe_q4k_q5k_bucket_work().

◆ gemm_q4_k_q8_k_packed_vnni_x8_compact_order_rows4()

void gemm_q4_k_q8_k_packed_vnni_x8_compact_order_rows4	(	float *	output,
		const void *	weights_packed,
		const void *	input_q8,
		int	rows,
		int	output_dim,
		int	input_dim
	)

Definition at line 1123 of file gemm_kernels_q4k_q8k_vnni.c.

{
    if (!output || !weights_packed || !input_q8 || rows <= 0 || rows > 4 ||
        output_dim <= 0 || input_dim <= 0 || (input_dim % QK_K) != 0) {
        return;
    }
#if defined(CK_HAS_AVX_VNNI_256)
    const block_q8_K *input = (const block_q8_K *)input_q8;
    const block_q4_K_packed_vnni_x8 *weights =
        (const block_q4_K_packed_vnni_x8 *)weights_packed;
    const int blocks_per_row = input_dim / QK_K;
    const int groups = (output_dim + 7) / 8;
    for (int group = 0; group < groups; ++group) {
        const int n0 = group * 8;
        const int active = n0 + 8 <= output_dim ? 8 : output_dim - n0;
        float acc[4][8] = {{0}};
        for (int block = 0; block < blocks_per_row; ++block) {
            float block_sums[4][8] = {{0}};
            const block_q8_K *input_rows[4] = {NULL, NULL, NULL, NULL};
            for (int row = 0; row < rows; ++row) {
                input_rows[row] = input +
                    (size_t)row * (size_t)blocks_per_row + (size_t)block;
            }
            dot_q4_k_packed_vnni_x8_q8_k_compact_order(
                block_sums,
                weights + (size_t)group * (size_t)blocks_per_row +
                    (size_t)block,
                input_rows,
                rows);
            for (int row = 0; row < rows; ++row) {
                const __m256 prior = _mm256_loadu_ps(acc[row]);
                const __m256 current = _mm256_loadu_ps(block_sums[row]);
                _mm256_storeu_ps(acc[row], _mm256_add_ps(prior, current));
            }
        }
        for (int row = 0; row < rows; ++row) {
            for (int lane = 0; lane < active; ++lane) {
                output[(size_t)row * (size_t)output_dim +
                       (size_t)n0 + (size_t)lane] = acc[row][lane];
            }
        }
    }
#else
    (void)rows;
    (void)output_dim;
    (void)input_dim;
#endif
}

References dot_q4_k_packed_vnni_x8_q8_k_compact_order(), and QK_K.

Referenced by ck_moe_q4k_q5k_bucket_work().

◆ gemm_q5_k_q8_k_compact_rows4()

void gemm_q5_k_q8_k_compact_rows4	(	float *	output,
		int	output_stride,
		const void *	weights,
		const void *const	input_rows[4],
		int	rows,
		int	output_dim,
		int	input_dim
	)

Definition at line 937 of file gemm_kernels_q5_k.c.

{
    if (!output || !weights || !input_rows || rows <= 0 || rows > 4 ||
        output_stride < output_dim || output_dim <= 0 || input_dim <= 0 ||
        (input_dim % QK_K) != 0) {
        return;
    }
    for (int row = 0; row < rows; ++row) {
        if (!input_rows[row]) return;
    }
 
#if defined(__AVX2__)
    const block_q5_K *blocks = (const block_q5_K *)weights;
    const int blocks_per_row = input_dim / QK_K;
    const block_q8_K *inputs[4] = {
        (const block_q8_K *)input_rows[0],
        (const block_q8_K *)input_rows[rows > 1 ? 1 : 0],
        (const block_q8_K *)input_rows[rows > 2 ? 2 : 0],
        (const block_q8_K *)input_rows[rows > 3 ? 3 : 0],
    };
    for (int n = 0; n < output_dim; ++n) {
        float values[4];
        dot_q5_k_q8_k_rows4_avx2(
            blocks + (size_t)n * (size_t)blocks_per_row,
            inputs, rows, blocks_per_row, values);
        for (int row = 0; row < rows; ++row) {
            output[(size_t)row * (size_t)output_stride + (size_t)n] =
                values[row];
        }
    }
#else
    for (int row = 0; row < rows; ++row) {
        gemv_q5_k_q8_k(
            output + (size_t)row * (size_t)output_stride,
            weights, input_rows[row], output_dim, input_dim);
    }
#endif
}

References gemv_q5_k_q8_k(), and QK_K.

Referenced by ck_moe_q4k_q5k_bucket_work().

◆ 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_swiglu_expert_backward_f32()

void moe_swiglu_expert_backward_f32	(	const float *	d_output,
		const float *	hidden,
		const int *	indices,
		const float *	routing_weights,
		const float *	expert_gate,
		const float *	expert_up,
		const float *	expert_down,
		float *	d_hidden,
		float *	d_routing_weights,
		float *	d_expert_gate,
		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 3020 of file axpy_kernels.c.

{
    if (!d_output || !hidden || !indices || !routing_weights || !expert_gate || !expert_up || !expert_down ||
        !d_hidden || !d_routing_weights || !d_expert_gate || !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_gate[p] = 0.0f;
        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 gate[intermediate_dim];
    float up[intermediate_dim];
    float silu_gate[intermediate_dim];
    float act[intermediate_dim];
    float d_act[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 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;
                silu_gate[i] = ck_moe_silu_f32(gv);
                act[i] = silu_gate[i] * uv;
                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) {
                const float d_up = d_act[i] * silu_gate[i];
                const float d_gate = d_act[i] * up[i] * ck_moe_dsilu_f32(gate[i]);
                for (int h = 0; h < hidden_dim; ++h) {
                    d_expert_up[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)] += d_up * x[h];
                    d_expert_gate[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)] += d_gate * x[h];
                    dx[h] += d_up * expert_up[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)] +
                             d_gate * expert_gate[ck_moe_up_idx(e, i, h, intermediate_dim, hidden_dim)];
                }
            }
        }
    }
}

References ck_moe_down_idx(), ck_moe_dsilu_f32(), ck_moe_silu_f32(), and ck_moe_up_idx().

◆ 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_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_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_bucketed_workspace_bytes()

size_t moe_swiglu_expert_q4k_q5k_bucketed_workspace_bytes	(	int	rows,
		int	hidden_dim,
		int	intermediate_dim,
		int	n_experts,
		int	top_k
	)

Definition at line 2202 of file axpy_kernels.c.

{
    ck_moe_q4k_q5k_bucket_layout_t layout;
    if (ck_moe_bucket_layout(rows, hidden_dim, intermediate_dim, n_experts,
                             top_k, &layout) != 0) {
        return 0;
    }
    return layout.total_bytes;
}

References ck_moe_bucket_layout().

◆ 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_backward_f32()

void moe_swiglu_shared_backward_f32	(	const float *	d_output,
		const float *	hidden,
		const float *	shared_gate,
		const float *	shared_up,
		const float *	shared_down,
		float *	d_hidden,
		float *	d_routed,
		float *	d_shared_gate,
		float *	d_shared_up,
		float *	d_shared_down,
		int	rows,
		int	hidden_dim,
		int	intermediate_dim
	)

Definition at line 3400 of file axpy_kernels.c.

{
    if (!d_output || !hidden || !shared_gate || !shared_up || !shared_down ||
        !d_hidden || !d_routed || !d_shared_gate || !d_shared_up || !d_shared_down ||
        rows <= 0 || hidden_dim <= 0 || intermediate_dim <= 0) {
        return;
    }
 
    for (size_t p = 0; p < (size_t)rows * (size_t)hidden_dim; ++p) {
        d_hidden[p] = 0.0f;
        d_routed[p] = d_output[p];
    }
    for (size_t p = 0; p < (size_t)intermediate_dim * (size_t)hidden_dim; ++p) {
        d_shared_gate[p] = 0.0f;
        d_shared_up[p] = 0.0f;
    }
    for (size_t p = 0; p < (size_t)hidden_dim * (size_t)intermediate_dim; ++p) d_shared_down[p] = 0.0f;
 
    float gate[intermediate_dim];
    float up[intermediate_dim];
    float silu_gate[intermediate_dim];
    float act[intermediate_dim];
    float d_act[intermediate_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 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];
            }
            gate[i] = gv;
            up[i] = uv;
            silu_gate[i] = ck_moe_silu_f32(gv);
            act[i] = silu_gate[i] * uv;
            d_act[i] = 0.0f;
        }
 
        for (int h = 0; h < hidden_dim; ++h) {
            for (int i = 0; i < intermediate_dim; ++i) {
                d_shared_down[(size_t)h * (size_t)intermediate_dim + (size_t)i] += dy[h] * act[i];
                d_act[i] += dy[h] * shared_down[(size_t)h * (size_t)intermediate_dim + (size_t)i];
            }
        }
 
        for (int i = 0; i < intermediate_dim; ++i) {
            const float d_up = d_act[i] * silu_gate[i];
            const float d_gate = d_act[i] * up[i] * ck_moe_dsilu_f32(gate[i]);
            for (int h = 0; h < hidden_dim; ++h) {
                d_shared_up[(size_t)i * (size_t)hidden_dim + (size_t)h] += d_up * x[h];
                d_shared_gate[(size_t)i * (size_t)hidden_dim + (size_t)h] += d_gate * x[h];
                dx[h] += d_up * shared_up[(size_t)i * (size_t)hidden_dim + (size_t)h] +
                         d_gate * shared_gate[(size_t)i * (size_t)hidden_dim + (size_t)h];
            }
        }
    }
}

References ck_moe_dsilu_f32(), and ck_moe_silu_f32().

◆ 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_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_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().

◆ pack_q4_k_to_packed_meta_x8()

void pack_q4_k_to_packed_meta_x8	(	const void *	source,
		void *	destination,
		int	output_dim,
		int	input_dim
	)

Definition at line 525 of file gemm_kernels_q4k_q8k_vnni.c.

{
    if (!src || !dst || N <= 0 || K <= 0 || (K % QK_K) != 0) {
        return;
    }
    const block_q4_K *in = (const block_q4_K *)src;
    block_q4_K_packed_meta_x8 *out = (block_q4_K_packed_meta_x8 *)dst;
    const int blocks_per_row = K / QK_K;
    const int groups = (N + 7) / 8;
    memset(out, 0, (size_t)groups * (size_t)blocks_per_row * sizeof(*out));
 
    for (int g = 0; g < groups; ++g) {
        const int n0 = g * 8;
        const int active = (n0 + 8 <= N) ? 8 : (N - n0);
        for (int b = 0; b < blocks_per_row; ++b) {
            block_q4_K_packed_meta_x8 *pb = out + (size_t)g * (size_t)blocks_per_row + (size_t)b;
            pb->active = (uint8_t)active;
            for (int lane = 0; lane < active; ++lane) {
                const block_q4_K *sb = in + (size_t)(n0 + lane) * (size_t)blocks_per_row + (size_t)b;
                pb->d[lane] = sb->d;
                pb->dmin[lane] = sb->dmin;
                unpack_q4_k_scales(sb->scales, pb->sc[lane], pb->m[lane]);
                memcpy(pb->qs[lane], sb->qs, sizeof(pb->qs[lane]));
            }
        }
    }
}

References block_q4_K::d, block_q4_K::dmin, QK_K, block_q4_K::qs, block_q4_K::scales, and unpack_q4_k_scales().

Referenced by ck_moe_q4k_llama_projection().

◆ q4_k_packed_meta_x8_block_size()

size_t q4_k_packed_meta_x8_block_size ( void )

Definition at line 464 of file gemm_kernels_q4k_q8k_vnni.c.

{
    return sizeof(block_q4_K_packed_meta_x8);
}

Referenced by ck_moe_q4k_llama_projection_scratch_bytes().

◆ q4_k_packed_vnni_x8_block_size()

size_t q4_k_packed_vnni_x8_block_size ( void )

Definition at line 267 of file gemm_kernels_q4k_q8k_vnni.c.

{
    return sizeof(block_q4_K_packed_vnni_x8);
}

Referenced by ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl().

◆ 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().

◆ 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().

Typedefs

Enumerations

Functions

Detailed Description

CK-ENGINE KERNEL RULES:

Typedef Documentation

◆ ck_moe_expert_workspace_fn

◆ ck_moe_shared_gated_workspace_fn

◆ ck_moe_shared_workspace_fn

Enumeration Type Documentation

◆ anonymous enum

◆ ck_moe_down_kind_t

Function Documentation

◆ axpy_2d_f32()

◆ axpy_f32()

◆ axpy_zero_f32()

◆ ck_moe_align64()

◆ ck_moe_bf16_round()

◆ ck_moe_bucket_expert_for_position()

◆ ck_moe_bucket_layout()

◆ ck_moe_debug_enabled()

◆ ck_moe_debug_finite()

◆ ck_moe_down_idx()

◆ ck_moe_dsilu_f32()

◆ ck_moe_llama_weighted_accumulate()

◆ ck_moe_q4k_llama_projection()

◆ ck_moe_q4k_llama_projection_scratch_bytes()

◆ ck_moe_q4k_mixed_parallel_work()

◆ ck_moe_q4k_mixed_parallel_workspace()

◆ ck_moe_q4k_mixed_route_parallel()

◆ ck_moe_q4k_mixed_route_work()

◆ ck_moe_q4k_q5k_bucket_work()

◆ ck_moe_q4k_q5k_parallel_work()

◆ ck_moe_q4k_q5k_quantize_work()

◆ ck_moe_q4k_q5k_route_parallel()

◆ ck_moe_q4k_q5k_route_work()

◆ ck_moe_shared_gated_parallel_workspace()

◆ ck_moe_shared_q4k_gated_workspace()

◆ ck_moe_shared_q4k_parallel_work()

◆ ck_moe_shared_q4k_parallel_workspace()

◆ ck_moe_shared_q8_0_parallel_work()

◆ ck_moe_sigmoid_f32()

◆ ck_moe_silu_f32()

◆ ck_moe_size_add()

◆ ck_moe_size_mul()

◆ ck_moe_swiglu_expert_forward_q4k_q5k_bucketed_impl()

◆ ck_moe_up_idx()

◆ farskip_swiglu_shared_combine_bf16()

◆ farskip_swiglu_shared_combine_bf16_row_range()

◆ gemm_nt_q4_k_packed_meta_x8_q8_k_gemv_order()

◆ gemm_nt_q4_k_q8_k_pairwise_split_min_parallel_dispatch()

◆ gemm_q4_k_q8_k_compact_rows4()

◆ gemm_q4_k_q8_k_packed_vnni_x8_compact_order_rows4()

◆ gemm_q5_k_q8_k_compact_rows4()

◆ moe_accumulate_expert_f32()

◆ moe_relu2_expert_backward_f32()

◆ moe_relu2_expert_forward_f32()

◆ moe_relu2_expert_forward_q5_0_q5_0()

◆ moe_relu2_expert_forward_q5_0_q8_0()

◆ moe_relu2_shared_forward_q5_1_q8_0()

◆ moe_swiglu_expert_backward_f32()

◆ moe_swiglu_expert_forward_bf16()

◆ moe_swiglu_expert_forward_bf16_row_range()

◆ moe_swiglu_expert_forward_f32()

◆ moe_swiglu_expert_forward_q4k_q4k_parallel_workspace()

◆ moe_swiglu_expert_forward_q4k_q4k_workspace()

◆ moe_swiglu_expert_forward_q4k_q5_0_parallel_workspace()

◆ moe_swiglu_expert_forward_q4k_q5_0_workspace()

◆ moe_swiglu_expert_forward_q4k_q5k_auto_workspace()

◆ moe_swiglu_expert_forward_q4k_q5k_bucketed_prepared_workspace()

◆ moe_swiglu_expert_forward_q4k_q5k_bucketed_workspace()

◆ moe_swiglu_expert_forward_q4k_q5k_parallel_workspace()

◆ moe_swiglu_expert_forward_q4k_q5k_workspace()

◆ moe_swiglu_expert_forward_q4k_q6k_parallel_workspace()

◆ moe_swiglu_expert_forward_q4k_q6k_workspace()

◆ moe_swiglu_expert_forward_q4k_q8_0_parallel_workspace()

◆ moe_swiglu_expert_forward_q4k_q8_0_workspace()

◆ moe_swiglu_expert_q4k_q5k_bucketed_workspace_bytes()

◆ moe_swiglu_expert_q4k_q5k_workspace_bytes()

◆ moe_swiglu_expert_q4k_q8_0_workspace_bytes()