202 lines
8.6 KiB
C
202 lines
8.6 KiB
C
#include "gem/ops.h"
|
|
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <string.h>
|
|
|
|
static int close_enough(float actual, float expected, float tolerance)
|
|
{
|
|
if (fabsf(actual - expected) <= tolerance)
|
|
return 1;
|
|
fprintf(stderr, "expected %.8f, got %.8f\n", expected, actual);
|
|
return 0;
|
|
}
|
|
|
|
static int test_online_softmax_partitions(void)
|
|
{
|
|
enum { TOKENS = 13, DIMENSIONS = 17, MAX_PARTITIONS = 16 };
|
|
float scores[TOKENS];
|
|
float values[TOKENS * DIMENSIONS];
|
|
float reference_numerator[DIMENSIONS];
|
|
float reference[DIMENSIONS];
|
|
float partial_numerator[MAX_PARTITIONS][DIMENSIONS];
|
|
float partial_maximum[MAX_PARTITIONS];
|
|
float partial_denominator[MAX_PARTITIONS];
|
|
float merged_numerator[DIMENSIONS];
|
|
float scratch_numerator[DIMENSIONS];
|
|
float actual[DIMENSIONS];
|
|
for (size_t token = 0; token < TOKENS; ++token) {
|
|
scores[token] = ((int)(token * 19 % 23) - 11) * 0.75f;
|
|
for (size_t dimension = 0; dimension < DIMENSIONS; ++dimension)
|
|
values[token * DIMENSIONS + dimension] =
|
|
((int)((token * 31 + dimension * 7) % 41) - 20) * 0.125f;
|
|
}
|
|
float reference_maximum;
|
|
float reference_denominator;
|
|
if (gem_attention_online_partial_f32(
|
|
scores, values, TOKENS, DIMENSIONS, &reference_maximum,
|
|
&reference_denominator, reference_numerator) != 0 ||
|
|
gem_attention_online_finalize_f32(
|
|
reference_denominator, reference_numerator, DIMENSIONS, reference) != 0)
|
|
return 0;
|
|
|
|
static const size_t partition_counts[] = {1, 2, 3, 4, 7, 16};
|
|
for (size_t test = 0;
|
|
test < sizeof(partition_counts) / sizeof(partition_counts[0]); ++test) {
|
|
size_t partitions = partition_counts[test];
|
|
for (size_t partition = 0; partition < partitions; ++partition) {
|
|
size_t begin = TOKENS * partition / partitions;
|
|
size_t end = TOKENS * (partition + 1) / partitions;
|
|
if (gem_attention_online_partial_f32(
|
|
scores + begin, values + begin * DIMENSIONS, end - begin, DIMENSIONS,
|
|
&partial_maximum[partition], &partial_denominator[partition],
|
|
partial_numerator[partition]) != 0)
|
|
return 0;
|
|
}
|
|
float merged_maximum = -INFINITY;
|
|
float merged_denominator = 0.0f;
|
|
memset(merged_numerator, 0, sizeof(merged_numerator));
|
|
for (size_t partition = 0; partition < partitions; ++partition) {
|
|
float next_maximum;
|
|
float next_denominator;
|
|
if (gem_attention_online_merge_f32(
|
|
merged_maximum, merged_denominator, merged_numerator,
|
|
partial_maximum[partition], partial_denominator[partition],
|
|
partial_numerator[partition], DIMENSIONS,
|
|
&next_maximum, &next_denominator, scratch_numerator) != 0)
|
|
return 0;
|
|
merged_maximum = next_maximum;
|
|
merged_denominator = next_denominator;
|
|
memcpy(merged_numerator, scratch_numerator, sizeof(merged_numerator));
|
|
}
|
|
if (gem_attention_online_finalize_f32(
|
|
merged_denominator, merged_numerator, DIMENSIONS, actual) != 0)
|
|
return 0;
|
|
for (size_t dimension = 0; dimension < DIMENSIONS; ++dimension)
|
|
if (!close_enough(actual[dimension], reference[dimension], 2e-6f))
|
|
return 0;
|
|
}
|
|
|
|
float invalid_maximum;
|
|
float invalid_denominator;
|
|
float invalid_numerator[DIMENSIONS];
|
|
float invalid_scores[] = {NAN};
|
|
size_t maximum_bytes;
|
|
size_t denominator_bytes;
|
|
size_t numerator_bytes;
|
|
size_t total_bytes;
|
|
return gem_attention_online_partial_f32(
|
|
NULL, NULL, 0, DIMENSIONS, &invalid_maximum,
|
|
&invalid_denominator, invalid_numerator) == 0 &&
|
|
isinf(invalid_maximum) && invalid_maximum < 0.0f &&
|
|
invalid_denominator == 0.0f &&
|
|
gem_attention_online_finalize_f32(
|
|
invalid_denominator, invalid_numerator, DIMENSIONS, actual) != 0 &&
|
|
gem_attention_online_partial_f32(
|
|
invalid_scores, values, 1, DIMENSIONS, &invalid_maximum,
|
|
&invalid_denominator, invalid_numerator) != 0 &&
|
|
gem_attention_online_scratch_bytes(
|
|
32, 512, 32, &maximum_bytes, &denominator_bytes,
|
|
&numerator_bytes, &total_bytes) == 0 &&
|
|
maximum_bytes == 4096 && denominator_bytes == 4096 &&
|
|
numerator_bytes == 2097152 && total_bytes == 2105344 &&
|
|
gem_attention_online_scratch_bytes(
|
|
SIZE_MAX, 512, 32, &maximum_bytes, &denominator_bytes,
|
|
&numerator_bytes, &total_bytes) != 0;
|
|
}
|
|
|
|
int main(void)
|
|
{
|
|
if (!test_online_softmax_partitions())
|
|
return 1;
|
|
if (!close_enough(gem_f16_to_f32(UINT16_C(0x3c00)), 1.0f, 0.0f) ||
|
|
!close_enough(gem_f16_to_f32(UINT16_C(0xc000)), -2.0f, 0.0f) ||
|
|
!close_enough(gem_f16_to_f32(UINT16_C(0x0001)), 0x1p-24f, 0.0f))
|
|
return 1;
|
|
|
|
uint8_t q4[18] = {0x00, 0x3c};
|
|
memset(q4 + 2, 0xf0, 16);
|
|
float values[32];
|
|
if (gem_dequantize_q4_0(q4, 32, values) != 0)
|
|
return 1;
|
|
for (size_t i = 0; i < 16; ++i)
|
|
if (!close_enough(values[i], -8.0f, 0.0f) ||
|
|
!close_enough(values[i + 16], 7.0f, 0.0f))
|
|
return 1;
|
|
|
|
float input[32];
|
|
for (size_t i = 0; i < 32; ++i)
|
|
input[i] = 1.0f;
|
|
float product;
|
|
if (gem_matvec_q4_0(q4, 32, 1, input, &product) != 0 ||
|
|
!close_enough(product, -16.0f, 0.0f))
|
|
return 1;
|
|
|
|
uint8_t q8[34] = {0x00, 0x38}; /* scale 0.5 */
|
|
for (size_t i = 0; i < 32; ++i)
|
|
q8[2 + i] = (uint8_t)(int8_t)((int)i - 16);
|
|
if (gem_dequantize_q8_0(q8, 32, values) != 0)
|
|
return 1;
|
|
for (size_t i = 0; i < 32; ++i)
|
|
if (!close_enough(values[i], ((float)i - 16.0f) * 0.5f, 0.0f))
|
|
return 1;
|
|
if (gem_matvec_q8_0(q8, 32, 1, input, &product) != 0 ||
|
|
!close_enough(product, -8.0f, 0.0f))
|
|
return 1;
|
|
|
|
const float norm_input[] = {3.0f, 4.0f};
|
|
const float weights[] = {1.0f, 2.0f};
|
|
float norm_output[2];
|
|
gem_rms_norm_f32(norm_input, weights, 0.0f, 2, norm_output);
|
|
float inverse_rms = 1.0f / sqrtf(12.5f);
|
|
if (!close_enough(norm_output[0], 3.0f * inverse_rms, 1e-6f) ||
|
|
!close_enough(norm_output[1], 8.0f * inverse_rms, 1e-6f))
|
|
return 1;
|
|
|
|
float softcap[] = {-60.0f, 0.0f, 60.0f};
|
|
gem_softcap_f32(softcap, 3, 30.0f);
|
|
if (!close_enough(softcap[0], -tanhf(2.0f) * 30.0f, 1e-6f) ||
|
|
!close_enough(softcap[1], 0.0f, 0.0f) ||
|
|
!close_enough(softcap[2], tanhf(2.0f) * 30.0f, 1e-6f))
|
|
return 1;
|
|
|
|
const float rope_input[] = {1.0f, 2.0f, 3.0f, 4.0f};
|
|
float rope_output[4];
|
|
if (gem_rope_neox_f32(rope_input, 4, 4, 1, 10000.0f, 1.0f, NULL, rope_output) != 0 ||
|
|
!close_enough(rope_output[0], cosf(1.0f) - 3.0f * sinf(1.0f), 1e-6f) ||
|
|
!close_enough(rope_output[2], sinf(1.0f) + 3.0f * cosf(1.0f), 1e-6f) ||
|
|
!close_enough(rope_output[1], 2.0f * cosf(0.01f) - 4.0f * sinf(0.01f), 1e-6f) ||
|
|
!close_enough(rope_output[3], 2.0f * sinf(0.01f) + 4.0f * cosf(0.01f), 1e-6f))
|
|
return 1;
|
|
|
|
const float query[] = {1.0f, 0.0f, 0.0f, 1.0f};
|
|
const float keys[] = {1.0f, 0.0f, 0.0f, 1.0f};
|
|
const float values_cache[] = {10.0f, 20.0f, 30.0f, 40.0f};
|
|
float scores[2];
|
|
float attention[4];
|
|
float p = expf(1.0f) / (expf(1.0f) + 1.0f);
|
|
if (gem_attention_decode_f32(query, keys, values_cache, 2, 1, 2, 2, 0, 1.0f,
|
|
scores, attention) != 0 ||
|
|
!close_enough(attention[0], p * 10.0f + (1.0f - p) * 30.0f, 1e-5f) ||
|
|
!close_enough(attention[1], p * 20.0f + (1.0f - p) * 40.0f, 1e-5f) ||
|
|
!close_enough(attention[2], (1.0f - p) * 10.0f + p * 30.0f, 1e-5f) ||
|
|
!close_enough(attention[3], (1.0f - p) * 20.0f + p * 40.0f, 1e-5f))
|
|
return 1;
|
|
if (gem_attention_decode_f32(query, keys, values_cache, 2, 1, 2, 2, 1, 1.0f,
|
|
scores, attention) != 0)
|
|
return 1;
|
|
for (size_t i = 0; i < 4; i += 2)
|
|
if (!close_enough(attention[i], 30.0f, 0.0f) ||
|
|
!close_enough(attention[i + 1], 40.0f, 0.0f))
|
|
return 1;
|
|
|
|
const float ring_keys[] = {0.0f, 1.0f, 1.0f, 0.0f}; /* logical t2 in slot 0, t1 in slot 1 */
|
|
const float ring_values[] = {30.0f, 40.0f, 10.0f, 20.0f};
|
|
if (gem_attention_decode_ring_f32(query, ring_keys, ring_values, 2, 1, 2, 3, 2, 2,
|
|
1.0f, scores, attention) != 0 ||
|
|
!close_enough(attention[0], p * 10.0f + (1.0f - p) * 30.0f, 1e-5f) ||
|
|
!close_enough(attention[3], p * 40.0f + (1.0f - p) * 20.0f, 1e-5f))
|
|
return 1;
|
|
return 0;
|
|
}
|