215 lines
8.2 KiB
Plaintext
215 lines
8.2 KiB
Plaintext
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include <float.h>
|
|
|
|
#include "cpu_kernels/local_attention.cuh"
|
|
|
|
//! TODO: Include your local attention kernel here!
|
|
#include "kernels/local_attention.cuh"
|
|
|
|
#define cudaCheck(call) do { \
|
|
cudaError_t err = call; \
|
|
if (err != cudaSuccess) { \
|
|
fprintf(stderr, "CUDA error in %s at line %d: %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
|
|
exit(EXIT_FAILURE); \
|
|
} \
|
|
} while(0)
|
|
|
|
void read_input_from_file(const char* filename, int* batch_size, int* seq_length, int* num_heads, int* head_dim, float** input) {
|
|
FILE* file = fopen(filename, "r");
|
|
if (!file) {
|
|
fprintf(stderr, "Error opening input file: %s\n", filename);
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
|
|
fscanf(file, "%d\n%d\n%d\n%d", batch_size, seq_length, num_heads, head_dim);
|
|
|
|
int B = *batch_size; int T = *seq_length;
|
|
int NH = *num_heads; int HS = *head_dim;
|
|
*input = (float*) malloc(B * T * 3 * NH * HS * sizeof(float));
|
|
|
|
// read Q, K, V
|
|
for (int b = 0; b < B; ++b) {
|
|
for (int t = 0; t < T; ++t) {
|
|
for (int nh = 0; nh < NH; ++nh) {
|
|
for (int hs = 0; hs < HS; ++hs) {
|
|
int idx = b * T * 3 * NH * HS + t * 3 * NH * HS + 0 * NH * HS + nh * HS + hs;
|
|
fscanf(file, "%f", &((*input)[idx]));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
fgetc(file);
|
|
for (int b = 0; b < B; ++b) {
|
|
for (int t = 0; t < T; ++t) {
|
|
for (int nh = 0; nh < NH; ++nh) {
|
|
for (int hs = 0; hs < HS; ++hs) {
|
|
int idx = b * T * 3 * NH * HS + t * 3 * NH * HS + 1 * NH * HS + nh * HS + hs;
|
|
fscanf(file, "%f", &((*input)[idx]));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
fgetc(file);
|
|
for (int b = 0; b < B; ++b) {
|
|
for (int t = 0; t < T; ++t) {
|
|
for (int nh = 0; nh < NH; ++nh) {
|
|
for (int hs = 0; hs < HS; ++hs) {
|
|
int idx = b * T * 3 * NH * HS + t * 3 * NH * HS + 2 * NH * HS + nh * HS + hs;
|
|
fscanf(file, "%f", &((*input)[idx]));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
fclose(file);
|
|
}
|
|
|
|
void read_output_from_file(const char* filename, int* batch_size, int* seq_length, int* num_heads, int* head_dim, float** outputs) {
|
|
FILE* file = fopen(filename, "r");
|
|
if (!file) {
|
|
fprintf(stderr, "Error opening input file: %s\n", filename);
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
|
|
fscanf(file, "%d\n%d\n%d\n%d", batch_size, seq_length, num_heads, head_dim);
|
|
|
|
int tensor_size = (*batch_size) * (*seq_length) * (*num_heads) * (*head_dim);
|
|
*outputs = (float*) malloc(tensor_size * sizeof(float));
|
|
|
|
for (int i = 0; i < tensor_size; ++i) fscanf(file, "%f", &((*outputs)[i]));
|
|
fgetc(file);
|
|
fclose(file);
|
|
}
|
|
|
|
void compare_arrays(float* computed_outputs, float* expected_outputs, int size, const char* test_name) {
|
|
int diff_count = 0;
|
|
double max_relative_error = 0.0;
|
|
double rmse = 0.0;
|
|
int first_mismatch_idx = -1;
|
|
const double ERROR_THRESHOLD = 0.1; // 10% relative error threshold
|
|
|
|
for (int i = 0; i < size; ++i) {
|
|
double abs_error = fabs(computed_outputs[i] - expected_outputs[i]);
|
|
double rel_error = abs_error / (fabs(expected_outputs[i]) + 1e-6);
|
|
rmse += abs_error * abs_error;
|
|
|
|
if (rel_error > max_relative_error) {
|
|
max_relative_error = rel_error;
|
|
}
|
|
|
|
if (rel_error > ERROR_THRESHOLD) {
|
|
if (first_mismatch_idx == -1) {
|
|
first_mismatch_idx = i;
|
|
}
|
|
++diff_count;
|
|
}
|
|
}
|
|
|
|
rmse = sqrt(rmse / size);
|
|
|
|
if (diff_count > 0) {
|
|
printf("+++++++++++++++ %s TEST FAILED with %d differing values (%.2f%%) +++++++++++++++\n",
|
|
test_name, diff_count, (100.0 * diff_count) / size);
|
|
printf("Max Relative Error: %.6f, RMSE: %.6f\n", max_relative_error, rmse);
|
|
|
|
if (first_mismatch_idx >= 0) {
|
|
printf("First mismatch at index %d: Expected %.6f, Got %.6f\n",
|
|
first_mismatch_idx, expected_outputs[first_mismatch_idx], computed_outputs[first_mismatch_idx]);
|
|
}
|
|
} else {
|
|
printf("-------------------- Test Passed! (Max Relative Error: %.6f) --------------------\n",
|
|
max_relative_error);
|
|
}
|
|
}
|
|
|
|
int main() {
|
|
printf("============= Local Attention Skeleton Test Code (Inputs and Golden Solution Provided) =============\n");
|
|
|
|
cudaCheck(cudaSetDevice(0));
|
|
|
|
int num_examples = 8;
|
|
|
|
for (int example_idx = 0; example_idx < num_examples; example_idx++) {
|
|
int batch_size = 0;
|
|
int seq_length = 0;
|
|
int num_heads = 0;
|
|
int head_dim = 0;
|
|
|
|
float *input_data = NULL;
|
|
float *expected_outputs = NULL;
|
|
|
|
char input_filepath[512];
|
|
char output_filepath[512];
|
|
snprintf(input_filepath, sizeof(input_filepath), "/work/hdd/bche/Project_GPT/local_attn_examples/example_%i.tcin", example_idx);
|
|
snprintf(output_filepath, sizeof(output_filepath), "/work/hdd/bche/Project_GPT/local_attn_examples/example_%i.tcout", example_idx);
|
|
|
|
printf("\n Testing Example %d: \n", example_idx);
|
|
|
|
// read provided input
|
|
read_input_from_file(input_filepath, &batch_size, &seq_length, &num_heads, &head_dim, &input_data);
|
|
// read golden solution values
|
|
read_output_from_file(output_filepath, &batch_size, &seq_length, &num_heads, &head_dim, &expected_outputs);
|
|
|
|
// --------------- CPU reference implementation test example ---------------
|
|
|
|
int size = batch_size * seq_length * num_heads * head_dim;
|
|
|
|
// separate copy of input for CPU to avoid input memory being overwritten
|
|
float* att_inp_cpu = (float*) malloc(size * 3 * sizeof(float));
|
|
float* qkvr_local = (float*) malloc(size * 3 * sizeof(float));
|
|
float* att_out = (float*) malloc(size * sizeof(float));
|
|
|
|
memcpy(att_inp_cpu, input_data, size * 3 * sizeof(float));
|
|
|
|
// calls the CPU implementation to compute windowed attention outputs
|
|
local_attention_forward_cpu(att_out, qkvr_local, att_inp_cpu, batch_size, seq_length, num_heads, head_dim);
|
|
|
|
//! TODO: uncomment the below line to see how to compare outputs of the provided CPU code with golden solution
|
|
// compare_arrays(att_out, expected_outputs, size, input_filepath);
|
|
|
|
free(att_inp_cpu);
|
|
free(qkvr_local);
|
|
free(att_out);
|
|
|
|
//! TODO: Modify the below code to test your GPU implementation!
|
|
|
|
// --------------- Your GPU implementation test ---------------
|
|
|
|
int tensor_size = batch_size * seq_length * num_heads * head_dim;
|
|
|
|
float *device_input = NULL;
|
|
float *device_qkvr = NULL;
|
|
float *device_output = NULL;
|
|
|
|
cudaCheck(cudaMalloc(&device_input, tensor_size * 3 * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&device_qkvr, tensor_size * 3 * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&device_output, tensor_size * sizeof(float)));
|
|
cudaCheck(cudaMemcpy(device_input, input_data, tensor_size * 3 * sizeof(float),
|
|
cudaMemcpyHostToDevice));
|
|
// //! TODO: Call your GPU implementation here (including permute and unpermute)
|
|
local_attention_forward_gpu(device_output, device_qkvr, NULL, device_input, batch_size, seq_length, num_heads, head_dim);
|
|
|
|
cudaCheck(cudaGetLastError());
|
|
cudaCheck(cudaDeviceSynchronize());
|
|
|
|
// // Copy results back to host for comparison
|
|
float* host_outputs = (float*)calloc(tensor_size, sizeof(float));
|
|
cudaCheck(cudaMemcpy(host_outputs, device_output, tensor_size * sizeof(float),
|
|
cudaMemcpyDeviceToHost));
|
|
|
|
compare_arrays(host_outputs, expected_outputs, tensor_size, input_filepath);
|
|
|
|
// //! TODO: Free allocated memory in host and device
|
|
free(input_data);
|
|
free(expected_outputs);
|
|
free(host_outputs);
|
|
cudaCheck(cudaFree(device_input));
|
|
cudaCheck(cudaFree(device_qkvr));
|
|
cudaCheck(cudaFree(device_output));
|
|
}
|
|
|
|
printf("\n================== End of Tests Reached ==================\n");
|
|
return 0;
|
|
} |