logit_kernels.c File Reference

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

Go to the source code of this file.

Functions
void	final_logit_scale_f32 (float *logits, int tokens, int vocab_size, float scale)
void	final_logit_scale_muse_pytorch_bf16_storage (float *logits, int tokens, int vocab_size, float scale)
void	final_logit_softcap_muse_pytorch_bf16_storage (float *logits, int tokens, int vocab_size, float cap)

Function Documentation

final_logit_scale_f32()

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

Definition at line 9 of file logit_kernels.c.

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

References vocab_size.

final_logit_scale_muse_pytorch_bf16_storage()

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

Definition at line 24 of file logit_kernels.c.

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

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

final_logit_softcap_muse_pytorch_bf16_storage()

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

Definition at line 40 of file logit_kernels.c.

{
    if (!logits || tokens <= 0 || vocab_size <= 0 || cap <= 0.0f) {
        return;
    }
 
    const size_t total = (size_t)tokens * (size_t)vocab_size;
    for (size_t i = 0; i < total; ++i) {
        const float value = bf16_to_float(float_to_bf16(logits[i]));
        const float divided = bf16_to_float(float_to_bf16(value / cap));
        const float activated = bf16_to_float(float_to_bf16(tanhf(divided)));
        logits[i] = bf16_to_float(float_to_bf16(activated * cap));
    }
}

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