Loading...
L2 normalization converts vectors to unit length, essential for cosine similarity, contrastive learning, and normalized embeddings. Efficient implementation fuses norm computation with division and handles the numerical edge case of zero-magnitude vectors.
Compute norm and normalize in same kernel.
__global__ void l2_normalize_fused(float* x, float* y, int N, int D) {
int row = blockIdx.x;
float* in = x + row * D;
float* out = y + row * D;
// Compute squared norm
float sq_sum = 0;
for (int i = threadIdx.x; i < D; i += blockDim.x) {
sq_sum += in[i] * in[i];
}
sq_sum = blockReduceSum(sq_sum);
__shared__ float inv_norm;
if (threadIdx.x == 0) {
inv_norm = (sq_sum > 1e-12f) ? rsqrtf(sq_sum) : 0.0f;
}
__syncthreads();
// Normalize
for (int i = threadIdx.x; i < D; i += blockDim.x) {
out[i] = in[i] * inv_norm;
}
}Two kernels, intermediate storage for norms.
void l2_normalize_naive(float* x, float* y, int N, int D) {
// Pass 1: compute norms
compute_row_norms<<<N, 256>>>(x, norms, N, D);
// Pass 2: divide by norms
divide_by_norms<<<blocks, threads>>>(x, y, norms, N, D);
}Float4 vectorization for 4x fewer memory transactions.
__global__ void l2_normalize_opt(float4* x, float4* y, int N, int D4) {
int row = blockIdx.x;
// Compute squared norm with float4
float sq_sum = 0;
for (int i = threadIdx.x; i < D4; i += blockDim.x) {
float4 v = x[row * D4 + i];
sq_sum += v.x*v.x + v.y*v.y + v.z*v.z + v.w*v.w;
}
sq_sum = blockReduceSum(sq_sum);
__shared__ float inv_norm;
if (threadIdx.x == 0) {
inv_norm = (sq_sum > 1e-12f) ? rsqrtf(sq_sum) : 0.0f;
}
__syncthreads();
// Normalize with float4
for (int i = threadIdx.x; i < D4; i += blockDim.x) {
float4 v = x[row * D4 + i];
y[row * D4 + i] = make_float4(
v.x * inv_norm, v.y * inv_norm,
v.z * inv_norm, v.w * inv_norm
);
}
}| Metric | Naive | Optimized | Improvement |
|---|---|---|---|
| Throughput (1M x 768) | 85 GB/s | 380 GB/s | 4.5x faster |
| Kernel launches | 2 | 1 | 2x fewer |
Return zero vector (not NaN). Check norm > epsilon before division. Some prefer returning small random unit vector.
L2 norm preserves relative magnitudes. Max norm (divide by max element) preserves sparsity patterns. L2 is standard for embeddings.
Ready to optimize your CUDA code? Download RightNow AI and get real-time performance analysis for your kernels.