softmax_kernels_bf16.c File Reference

Softmax kernels for BF16 tensors. More...

#include <stddef.h>
#include <stdint.h>
#include "bf16_utils.h"
#include "ckernel_engine.h"

Go to the source code of this file.

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

Detailed Description

Softmax kernels for BF16 tensors.

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_bf16.c.

Function Documentation

backward_causal_softmax_head_major_bf16()

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

Definition at line 53 of file softmax_kernels_bf16.c.

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

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

causal_softmax_head_major_bf16()

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

Definition at line 31 of file softmax_kernels_bf16.c.

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

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