112 lines
4.3 KiB
C
112 lines
4.3 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_kernel_norm.h
|
|
* @brief Normalization family: row softmax, RMS normalization over one
|
|
* vector or a batch of rows, and the per-head QK-RMSNorm the attention
|
|
* block needs. The reductions are cooperative (simd_sum + a 32-slot
|
|
* threadgroup staging array), except QK-RMSNorm which runs one SIMD group
|
|
* per head and needs no threadgroup memory at all.
|
|
*
|
|
* @status Complete: softmax, rms_norm, rms_norm_nx, qk_rms_norm. All three
|
|
* RMS entry points honor weight == NULL.
|
|
* @issues None
|
|
* @todo None
|
|
*
|
|
* OCCUPANCY OF THE ROW MAPPING: rms_norm_nx runs one threadgroup of
|
|
* min(n, 256) threads per row with consecutive lanes on consecutive
|
|
* elements, so at n_embd = 2048 every lane of all 8 SIMD groups carries
|
|
* 8 elements — no idle lane, and each load coalesces. Measured on M4:
|
|
* 512 rows x 2048 in 0.053 ms, 238 GB/s counting x read twice plus the
|
|
* write, i.e. above the ~120 GB/s DRAM roof because a 4 MB working set
|
|
* stays in cache. Lanes only idle for rows narrower than 32, which no
|
|
* shape in this model has. qk_rms_norm runs ONE SIMD group per head and
|
|
* head_dim = 256 gives each lane 8 elements, likewise fully fed; at 16
|
|
* heads it is dispatch-bound (0.009 ms) rather than bandwidth-bound.
|
|
*
|
|
* FUNCTION-CONSTANT INDICES OWNED BY THIS FAMILY: 110-119 (110 used).
|
|
* Function constants share ONE index space across the whole assembled
|
|
* shader library, so each family reserves a block and two families
|
|
* colliding on an index fails the compile of every kernel in the backend,
|
|
* not just theirs. Allocated so far: 0-2 probe · 30 Q3_K · 50 Q5_K ·
|
|
* 60 Q6_K · 70/80 Q8_0+Q4_K · 100-109 rope · 110-119 norm ·
|
|
* 120-129 elementwise (claimed, none used).
|
|
*/
|
|
|
|
#ifndef FOUNDRY_METAL_KERNEL_NORM_H
|
|
#define FOUNDRY_METAL_KERNEL_NORM_H
|
|
|
|
#include "metal/fm_metal.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/* Softmax over last dimension.
|
|
* x: [rows, cols], out: [rows, cols] */
|
|
void fm_softmax(fm_buffer *x, fm_buffer *out, int rows, int cols);
|
|
|
|
/**
|
|
* RMS normalization over one vector: out = x * rsqrt(mean(x^2) + eps) * weight.
|
|
*
|
|
* @param x Input, n floats
|
|
* @param weight Per-element gain, n floats, or NULL for no gain
|
|
* @param out Output, n floats (may alias x)
|
|
* @param n Vector length
|
|
* @param eps Added to mean(x^2) before the rsqrt
|
|
*
|
|
* Thread-safe: no (encodes into the global command buffer)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_rms_norm(fm_buffer *x, fm_buffer *weight, fm_buffer *out,
|
|
int n, float eps);
|
|
|
|
/**
|
|
* RMS normalization over a batch of contiguous rows — the prefill form of
|
|
* fm_rms_norm. Each row is normalized independently; the weight is shared
|
|
* by every row, exactly as the single-vector form applies it.
|
|
*
|
|
* @param x Input, [rows, n] row-major, no padding between rows
|
|
* @param weight Per-element gain, n floats, or NULL for no gain
|
|
* @param out Output, [rows, n] (may alias x)
|
|
* @param rows Number of rows
|
|
* @param n Elements per row, which is also the row stride
|
|
* @param eps Added to mean(x^2) before the rsqrt
|
|
*
|
|
* Thread-safe: no
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_rms_norm_nx(fm_buffer *x, fm_buffer *weight, fm_buffer *out,
|
|
int rows, int n, float eps);
|
|
|
|
/**
|
|
* Per-head QK-RMSNorm: RMS-normalizes each head of a [n_heads, head_dim]
|
|
* activation independently.
|
|
*
|
|
* The weight is SHARED across heads — it is head_dim long, not
|
|
* n_heads * head_dim, and head h reads weight[0 .. head_dim), not
|
|
* weight[h * head_dim ...]. That mirrors the GGUF layout, which stores one
|
|
* pre-baked norm weight per tensor and applies it to every head.
|
|
*
|
|
* @param x Input, [n_heads, head_dim] row-major
|
|
* @param weight Shared per-head gain, head_dim floats, or NULL for none
|
|
* @param out Output, [n_heads, head_dim] (may alias x)
|
|
* @param n_heads Number of heads
|
|
* @param head_dim Elements per head
|
|
* @param eps Added to mean(x^2) before the rsqrt
|
|
*
|
|
* Thread-safe: no
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void fm_qk_rms_norm(fm_buffer *x, fm_buffer *weight, fm_buffer *out,
|
|
int n_heads, int head_dim, float eps);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_METAL_KERNEL_NORM_H */
|