recurrent_qk_norm_kernels.c

Go to the documentation of this file.

#include "bf16_utils.h"
#include "ckernel_engine.h"
#include "ck_threadpool.h"
 
#include <limits.h>
#include <math.h>
#include <stdint.h>
 
#if defined(__AVX2__)
#include <immintrin.h>
#endif
 
static void recurrent_l2_norm_rows_forward_one(float *x,
                                               int rows,
                                               int dim,
                                               int head_dim,
                                               float eps) {
    const int num_heads = dim / head_dim;
    for (int row = 0; row < rows; ++row) {
        float *row_ptr = x + (size_t)row * (size_t)dim;
        for (int head = 0; head < num_heads; ++head) {
            float *head_ptr = row_ptr + (size_t)head * (size_t)head_dim;
            double sum_sq = 0.0;
            for (int col = 0; col < head_dim; ++col) {
                sum_sq += (double)(head_ptr[col] * head_ptr[col]);
            }
            const float norm = sqrtf((float)sum_sq);
            const float inv_norm = 1.0f / fmaxf(norm, eps);
            for (int col = 0; col < head_dim; ++col) {
                head_ptr[col] *= inv_norm;
            }
        }
    }
}
 
typedef struct {
    float *q;
    float *k;
    int q_heads;
    int k_heads;
    int head_dim;
    float eps;
} ck_recurrent_qk_l2_args_t;
 
static void recurrent_qk_l2_norm_head_range(int begin,
                                            int end,
                                            void *opaque) {
    const ck_recurrent_qk_l2_args_t *args =
        (const ck_recurrent_qk_l2_args_t *)opaque;
    const int heads_per_row = args->q_heads + args->k_heads;
    for (int job = begin; job < end; ++job) {
        const int row = job / heads_per_row;
        const int local_head = job - row * heads_per_row;
        float *head_ptr;
        if (local_head < args->q_heads) {
            head_ptr = args->q +
                ((size_t)row * (size_t)args->q_heads +
                 (size_t)local_head) * (size_t)args->head_dim;
        } else {
            head_ptr = args->k +
                ((size_t)row * (size_t)args->k_heads +
                 (size_t)(local_head - args->q_heads)) *
                    (size_t)args->head_dim;
        }
 
        double sum_sq = 0.0;
        for (int col = 0; col < args->head_dim; ++col) {
            sum_sq += (double)(head_ptr[col] * head_ptr[col]);
        }
        const float norm = sqrtf((float)sum_sq);
        const float inv_norm = 1.0f / fmaxf(norm, args->eps);
        for (int col = 0; col < args->head_dim; ++col) {
            head_ptr[col] *= inv_norm;
        }
    }
}
 
static void recurrent_l2_norm_rows_backward_one(const float *d_out,
                                                const float *x,
                                                float *d_x,
                                                int rows,
                                                int dim,
                                                int head_dim,
                                                float eps) {
    if (!d_out || !x || !d_x || rows <= 0 || dim <= 0 || head_dim <= 0) {
        return;
    }
    const int num_heads = dim / head_dim;
    if (num_heads <= 0 || num_heads * head_dim != dim) {
        return;
    }
 
    for (int row = 0; row < rows; ++row) {
        const float *d_row = d_out + (size_t) row * (size_t) dim;
        const float *x_row = x + (size_t) row * (size_t) dim;
        float *dx_row = d_x + (size_t) row * (size_t) dim;
        for (int head = 0; head < num_heads; ++head) {
            const float *d_head = d_row + (size_t) head * (size_t) head_dim;
            const float *x_head = x_row + (size_t) head * (size_t) head_dim;
            float *dx_head = dx_row + (size_t) head * (size_t) head_dim;
 
            double sum_sq = 0.0;
            double dot = 0.0;
            for (int col = 0; col < head_dim; ++col) {
                sum_sq += (double) (x_head[col] * x_head[col]);
                dot += (double) (d_head[col] * x_head[col]);
            }
            const float norm = sqrtf((float) sum_sq);
            const float inv_norm = 1.0f / fmaxf(norm, eps);
            if (norm <= eps) {
                for (int col = 0; col < head_dim; ++col) {
                    dx_head[col] = d_head[col] / eps;
                }
                continue;
            }
            const float proj_scale = inv_norm * inv_norm * inv_norm * (float) dot;
            for (int col = 0; col < head_dim; ++col) {
                dx_head[col] = inv_norm * d_head[col] - proj_scale * x_head[col];
            }
        }
    }
}
 
void recurrent_qk_l2_norm_forward(float *q,
                                  float *k,
                                  int rows,
                                  int q_dim,
                                  int k_dim,
                                  int head_dim,
                                  float eps) {
    if (!q || !k || rows <= 0 || q_dim <= 0 || k_dim <= 0 ||
        head_dim <= 0 || q_dim % head_dim != 0 || k_dim % head_dim != 0) {
        return;
    }
    const int q_heads = q_dim / head_dim;
    const int k_heads = k_dim / head_dim;
    if (q_heads > INT_MAX - k_heads) {
        return;
    }
    const int heads_per_row = q_heads + k_heads;
    if (rows > INT_MAX / heads_per_row) {
        return;
    }
 
    ck_recurrent_qk_l2_args_t args = {
        q, k, q_heads, k_heads, head_dim, eps,
    };
    const int jobs = rows * heads_per_row;
    const size_t elements = (size_t)jobs * (size_t)head_dim;
    const size_t minimum_elements_per_thread = 32768u;
    const uintptr_t q_begin = (uintptr_t)q;
    const uintptr_t q_end = q_begin +
        (size_t)rows * (size_t)q_dim * sizeof(float);
    const uintptr_t k_begin = (uintptr_t)k;
    const uintptr_t k_end = k_begin +
        (size_t)rows * (size_t)k_dim * sizeof(float);
    const int ranges_overlap = q_begin < k_end && k_begin < q_end;
    if (ranges_overlap || elements < 2u * minimum_elements_per_thread) {
        recurrent_l2_norm_rows_forward_one(q, rows, q_dim, head_dim, eps);
        recurrent_l2_norm_rows_forward_one(k, rows, k_dim, head_dim, eps);
        return;
    }
 
    ck_threadpool_t *pool = ck_threadpool_global();
    int active = pool ? ck_threadpool_n_threads(pool) : 1;
    size_t useful_threads =
        (elements + minimum_elements_per_thread - 1u) /
        minimum_elements_per_thread;
    if ((size_t)active > useful_threads) {
        active = (int)useful_threads;
    }
    if (active > jobs) {
        active = jobs;
    }
 
    if (active > 1 && ck_threadpool_thread_id(pool) <= 0) {
        ck_threadpool_parallel_for_n(
            pool, active, 0, jobs, 8,
            recurrent_qk_l2_norm_head_range, &args);
    } else {
        recurrent_l2_norm_rows_forward_one(q, rows, q_dim, head_dim, eps);
        recurrent_l2_norm_rows_forward_one(k, rows, k_dim, head_dim, eps);
    }
}
 
static int recurrent_ceil_log2(int value)
{
    int result = 0;
    int remaining = value - 1;
    while (remaining > 0) {
        ++result;
        remaining >>= 1;
    }
    return result;
}
 
static float recurrent_pytorch_fp32_square_sum(const float *x, int dim)
{
#if defined(__AVX2__)
    if (dim >= 8) {
        enum { ilp_factor = 4, num_levels = 4, vector_width = 8 };
        const int vector_count = dim / vector_width;
        const int cascade_count = vector_count / ilp_factor;
        const int level_power =
            recurrent_ceil_log2(cascade_count) / num_levels > 4
                ? recurrent_ceil_log2(cascade_count) / num_levels
                : 4;
        const int level_step = 1 << level_power;
        const int level_mask = level_step - 1;
        __m256 acc[num_levels][ilp_factor];
        for (int level = 0; level < num_levels; ++level) {
            for (int lane = 0; lane < ilp_factor; ++lane) {
                acc[level][lane] = _mm256_setzero_ps();
            }
        }
 
        int group = 0;
        while (group + level_step <= cascade_count) {
            for (int offset = 0; offset < level_step; ++offset, ++group) {
                for (int lane = 0; lane < ilp_factor; ++lane) {
                    const int vector_index = group * ilp_factor + lane;
                    const __m256 value = _mm256_loadu_ps(
                        x + vector_index * vector_width);
                    const __m256 square = _mm256_mul_ps(value, value);
                    acc[0][lane] = _mm256_add_ps(acc[0][lane], square);
                }
            }
            for (int level = 1; level < num_levels; ++level) {
                for (int lane = 0; lane < ilp_factor; ++lane) {
                    acc[level][lane] =
                        _mm256_add_ps(acc[level][lane], acc[level - 1][lane]);
                    acc[level - 1][lane] = _mm256_setzero_ps();
                }
                const int mask = level_mask << (level * level_power);
                if ((group & mask) != 0) {
                    break;
                }
            }
        }
        for (; group < cascade_count; ++group) {
            for (int lane = 0; lane < ilp_factor; ++lane) {
                const int vector_index = group * ilp_factor + lane;
                const __m256 value = _mm256_loadu_ps(
                    x + vector_index * vector_width);
                const __m256 square = _mm256_mul_ps(value, value);
                acc[0][lane] = _mm256_add_ps(acc[0][lane], square);
            }
        }
        for (int level = 1; level < num_levels; ++level) {
            for (int lane = 0; lane < ilp_factor; ++lane) {
                acc[0][lane] =
                    _mm256_add_ps(acc[0][lane], acc[level][lane]);
            }
        }
 
        int vector_index = cascade_count * ilp_factor;
        for (; vector_index < vector_count; ++vector_index) {
            const __m256 value = _mm256_loadu_ps(
                x + vector_index * vector_width);
            const __m256 square = _mm256_mul_ps(value, value);
            acc[0][0] = _mm256_add_ps(acc[0][0], square);
        }
        for (int lane = 1; lane < ilp_factor; ++lane) {
            acc[0][0] = _mm256_add_ps(acc[0][0], acc[0][lane]);
        }
 
        _Alignas(32) float lanes[8];
        _mm256_store_ps(lanes, acc[0][0]);
        volatile float sum = 0.0f;
        for (int lane = 0; lane < 8; ++lane) {
            sum = sum + lanes[lane];
        }
        for (int d = vector_count * vector_width; d < dim; ++d) {
            sum = sum + x[d] * x[d];
        }
        return sum;
    }
#endif
    volatile float sum = 0.0f;
    for (int d = 0; d < dim; ++d) {
        sum = sum + x[d] * x[d];
    }
    return sum;
}
 
static void recurrent_pytorch_fp32_l2_rows(float *x,
                                            int rows,
                                            int dim,
                                            int head_dim,
                                            float eps)
{
    if (!x || rows <= 0 || dim <= 0 || head_dim <= 0 ||
        dim % head_dim != 0) {
        return;
    }
    const int num_heads = dim / head_dim;
    for (int row = 0; row < rows; ++row) {
        float *row_ptr = x + (size_t)row * (size_t)dim;
        for (int head = 0; head < num_heads; ++head) {
            float *head_ptr = row_ptr + (size_t)head * (size_t)head_dim;
            const float sum = recurrent_pytorch_fp32_square_sum(
                head_ptr, head_dim);
            const float inverse = 1.0f / sqrtf(sum + eps);
#if defined(__AVX512F__)
            const __m512 inv = _mm512_set1_ps(inverse);
            int d = 0;
            for (; d + 16 <= head_dim; d += 16) {
                _mm512_storeu_ps(
                    head_ptr + d,
                    _mm512_mul_ps(_mm512_loadu_ps(head_ptr + d), inv));
            }
            for (; d < head_dim; ++d) {
                head_ptr[d] = head_ptr[d] * inverse;
            }
#elif defined(__AVX2__)
            const __m256 inv = _mm256_set1_ps(inverse);
            int d = 0;
            for (; d + 8 <= head_dim; d += 8) {
                _mm256_storeu_ps(
                    head_ptr + d,
                    _mm256_mul_ps(_mm256_loadu_ps(head_ptr + d), inv));
            }
            for (; d < head_dim; ++d) {
                head_ptr[d] = head_ptr[d] * inverse;
            }
#else
            for (int d = 0; d < head_dim; ++d) {
                head_ptr[d] = head_ptr[d] * inverse;
            }
#endif
        }
    }
}
 
void recurrent_qk_l2_norm_pytorch_fp32_output(float *q,
                                               float *k,
                                               int rows,
                                               int q_dim,
                                               int k_dim,
                                               int head_dim,
                                               float eps)
{
    recurrent_pytorch_fp32_l2_rows(q, rows, q_dim, head_dim, eps);
    recurrent_pytorch_fp32_l2_rows(k, rows, k_dim, head_dim, eps);
}
 
static float recurrent_pytorch_bf16_square_sum(const float *x, int dim)
{
#if defined(__AVX2__)
    if (dim >= 32) {
        __m256 streams[4] = {
            _mm256_setzero_ps(), _mm256_setzero_ps(),
            _mm256_setzero_ps(), _mm256_setzero_ps()
        };
        int d = 0;
        for (; d + 32 <= dim; d += 32) {
            for (int stream = 0; stream < 4; ++stream) {
                _Alignas(32) float lanes[8];
                const __m256 values = _mm256_loadu_ps(x + d + stream * 8);
                _mm256_store_ps(lanes, _mm256_mul_ps(values, values));
                for (int lane = 0; lane < 8; ++lane) {
                    lanes[lane] =
                        bf16_to_float(float_to_bf16(lanes[lane]));
                }
                streams[stream] =
                    _mm256_add_ps(streams[stream], _mm256_load_ps(lanes));
            }
        }
        __m256 reduced = _mm256_add_ps(streams[0], streams[1]);
        reduced = _mm256_add_ps(reduced, streams[2]);
        reduced = _mm256_add_ps(reduced, streams[3]);
        _Alignas(32) float lanes[8];
        _mm256_store_ps(lanes, reduced);
        volatile float sum = 0.0f;
        for (int lane = 0; lane < 8; ++lane) {
            sum = sum + lanes[lane];
        }
        for (; d < dim; ++d) {
            const float value = bf16_to_float(float_to_bf16(x[d]));
            const float square =
                bf16_to_float(float_to_bf16(value * value));
            sum = sum + square;
        }
        return sum;
    }
#endif
    volatile float sum = 0.0f;
    for (int d = 0; d < dim; ++d) {
        const float value = bf16_to_float(float_to_bf16(x[d]));
        const float square = bf16_to_float(float_to_bf16(value * value));
        sum = sum + square;
    }
    return sum;
}
 
static void recurrent_pytorch_bf16_l2_rows(float *x,
                                            int rows,
                                            int dim,
                                            int expanded_heads,
                                            int head_dim,
                                            float eps)
{
    if (!x || rows <= 0 || dim <= 0 || head_dim <= 0 ||
        dim % head_dim != 0) {
        return;
    }
    const int num_heads = dim / head_dim;
    if (expanded_heads <= 0 || expanded_heads % num_heads != 0) {
        return;
    }
    const int heads_per_group = expanded_heads / num_heads;
    const int total_heads = rows * expanded_heads;
    const int vector_limit = (total_heads / 32) * 32;
    for (int row = 0; row < rows; ++row) {
        float *row_ptr = x + (size_t)row * (size_t)dim;
        for (int head = 0; head < num_heads; ++head) {
            float *head_ptr =
                row_ptr + (size_t)head * (size_t)head_dim;
            float sum = recurrent_pytorch_bf16_square_sum(
                head_ptr, head_dim);
            sum = bf16_to_float(float_to_bf16(sum));
            const float denominator =
                bf16_to_float(float_to_bf16(sum + eps));
            /*
             * Qwen3Next repeat_interleave() expands compact Q/K groups to
             * value heads before l2norm.  TensorIterator's rsqrt tail is
             * therefore determined by the expanded tensor, not this compact
             * storage.  Select the arithmetic used by the first repeated
             * lane; all lanes are vector lanes for production prefill.
             */
            const int flat_head =
                (row * expanded_heads) + head * heads_per_group;
            float inverse;
            if (flat_head < vector_limit) {
                /* ATen Vec<BF16> evaluates 32 scalar outputs per vector
                 * iteration using FP32 rsqrt followed by BF16 storage. */
                inverse = bf16_to_float(
                    float_to_bf16(1.0f / sqrtf(denominator)));
            } else {
                /* TensorIterator's BF16 scalar tail materializes sqrt to
                 * BF16 before applying the BF16 reciprocal. */
                const float root =
                    bf16_to_float(float_to_bf16(sqrtf(denominator)));
                inverse =
                    bf16_to_float(float_to_bf16(1.0f / root));
            }
            for (int col = 0; col < head_dim; ++col) {
                const float value =
                    bf16_to_float(float_to_bf16(head_ptr[col]));
                head_ptr[col] = bf16_to_float(
                    float_to_bf16(value * inverse));
            }
        }
    }
}
 
void recurrent_qk_l2_norm_pytorch_bf16_storage(float *q,
                                                float *k,
                                                int rows,
                                                int q_dim,
                                                int k_dim,
                                                int expanded_heads,
                                                int head_dim,
                                                float eps)
{
    recurrent_pytorch_bf16_l2_rows(
        q, rows, q_dim, expanded_heads, head_dim, eps);
    recurrent_pytorch_bf16_l2_rows(
        k, rows, k_dim, expanded_heads, head_dim, eps);
}
 
void recurrent_qk_l2_norm_backward(const float *d_q_out,
                                   const float *d_k_out,
                                   const float *q,
                                   const float *k,
                                   float *d_q,
                                   float *d_k,
                                   int rows,
                                   int q_dim,
                                   int k_dim,
                                   int head_dim,
                                   float eps) {
    recurrent_l2_norm_rows_backward_one(d_q_out, q, d_q, rows, q_dim, head_dim, eps);
    recurrent_l2_norm_rows_backward_one(d_k_out, k, d_k, rows, k_dim, head_dim, eps);
}