softmax_kernels.c File Reference

Softmax forward/backward kernels with SIMD (SSE/AVX/AVX512) More...

#include <math.h>

Go to the source code of this file.

Functions
void	backward_causal_softmax_head_major (float *d_scores, const float *weights, int num_heads, int num_tokens, int aligned_context_window)
void	causal_softmax_head_major (float *scores, int num_heads, int num_tokens, int aligned_context_window)
void	causal_softmax_head_major_exact (float *scores, int num_heads, int num_tokens, int aligned_context_window)

Detailed Description

Softmax forward/backward kernels with SIMD (SSE/AVX/AVX512)

CK-ENGINE KERNEL RULES:

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

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

Softmax: y[i] = exp(x[i] - max(x)) / sum(exp(x - max(x)))

Definition in file softmax_kernels.c.

Function Documentation

backward_causal_softmax_head_major()

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

Definition at line 382 of file softmax_kernels.c.

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

Referenced by backward_causal_softmax_head_major_bf16().

causal_softmax_head_major()

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

Causal softmax (in-place, row-wise)

Test:

test_softmax.py::TestSoftmaxForward::test_causal_softmax

test_softmax.py::TestSoftmaxForward::test_causal_vs_softmax

test_attention.py::TestAttentionForward::test_softmax_correctness

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

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

Definition at line 144 of file softmax_kernels.c.

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

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

causal_softmax_head_major_exact()

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

Causal softmax (exact version using stdlib expf)

Test:

test_softmax.py::TestSoftmaxForward::test_causal_softmax_exact

test_softmax.py::TestSoftmaxForward::test_exact_vs_fast

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

After changes: make test

Definition at line 339 of file softmax_kernels.c.

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

Referenced by attention_forward_causal_head_major_exact(), and attention_forward_causal_head_major_gqa_exact().