7f59912916
The model code was spread across eight crates that had grown into each other:
ggml and cuda and mlx each owned part of a tensor runtime, llama and tts and
voice2 each owned part of a model, and libs/diffusion owned everything else.
They are now one tree with an explicit shape:
libs/ai/cuda — kernels and launch surface
libs/ai/metal — Metal shaders and the shim
libs/ai/llm — the language-model runtime (sessions, lanes, contexts,
the CUDA and Metal executors, the compiled Metal path)
libs/ai/models/ — common, flux, h3, music, paint, speech, stems, vision
libs/diffusion is not deleted but demoted: what remains is the VALIDATOR
crate — several dozen `*_validate.rs` oracles that check a native
implementation against a reference, which is where they belong now that the
implementations live next door.
The functional work inside the move is mostly in the LLM runtime: N lanes that
draft while one verify batch serves all of them, per-slot prefill over a shared
folded attention arena, speculation that survives batching, and a scheduler
that reports rather than publishes. And in the CUDA build: a machine without
usable CUDA must still LINK (and say so), the default kernel arch is the
building machine's GPU, `NO_CUDA` forces the stub even where the toolkit
exists, and kernels compile in parallel with progress.
libs/video_flow is new here: classical optical flow estimation and the `mkfl`
motion-field payload — a flow field measured from a clip without a model,
which is what drives free-rate bounce-looping playback and the uprez/tween
enhance pipe.
188 lines
8.4 KiB
Text
188 lines
8.4 KiB
Text
// SPDX-License-Identifier: MIT
|
|
// Copyright (c) 2023-2026 The ggml authors
|
|
//
|
|
// Substantial portions derived from ggml / llama.cpp
|
|
// (https://github.com/ggml-org/llama.cpp), MIT licensed.
|
|
// The original copyright notice and permission notice are retained.
|
|
// See libs/ai/NOTICE and, where present, LICENSE in this directory.
|
|
//
|
|
#pragma once
|
|
// Official llama.cpp ggml-cuda/norm.cu rms_norm_f32 fused with MUL (+ ADD).
|
|
// Source: /Users/admin/llama.cpp/ggml/src/ggml-cuda/norm.cu:74-384
|
|
// Dispatch: ggml-cuda.cu:3994-4004 ggml_cuda_op_rms_norm_fused[_add].
|
|
|
|
#include "common.cuh"
|
|
|
|
template <int block_size, bool do_multiply, bool do_add>
|
|
static __global__ void rms_norm_f32(
|
|
const float * x,
|
|
float * dst,
|
|
const int ncols,
|
|
const int64_t stride_row,
|
|
const int64_t stride_channel,
|
|
const int64_t stride_sample,
|
|
const float eps,
|
|
const int64_t dst_stride_row,
|
|
const int64_t dst_stride_channel,
|
|
const int64_t dst_stride_sample,
|
|
const float * mul,
|
|
const int64_t mul_stride_row,
|
|
const int64_t mul_stride_channel,
|
|
const int64_t mul_stride_sample,
|
|
const uint3 mul_ncols_packed,
|
|
const uint3 mul_nrows_packed,
|
|
const uint3 mul_nchannels_packed,
|
|
const uint3 mul_nsamples_packed,
|
|
const float * add,
|
|
const int64_t add_stride_row,
|
|
const int64_t add_stride_channel,
|
|
const int64_t add_stride_sample,
|
|
const uint3 add_ncols_packed,
|
|
const uint3 add_nrows_packed,
|
|
const uint3 add_nchannels_packed,
|
|
const uint3 add_nsamples_packed) {
|
|
const int nrows = gridDim.x;
|
|
const int nchannels = gridDim.y;
|
|
|
|
const int row = blockIdx.x;
|
|
const int channel = blockIdx.y;
|
|
const int sample = blockIdx.z;
|
|
const int tid = threadIdx.x;
|
|
|
|
static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying");
|
|
|
|
x += sample*stride_sample + channel*stride_channel + row*stride_row;
|
|
dst += sample*dst_stride_sample + channel*dst_stride_channel + row*dst_stride_row;
|
|
(void) nrows;
|
|
(void) nchannels;
|
|
|
|
if constexpr (do_multiply) {
|
|
const uint32_t mul_row = fastmodulo((uint32_t) row, mul_nrows_packed);
|
|
const uint32_t mul_channel = fastmodulo((uint32_t) channel, mul_nchannels_packed);
|
|
const uint32_t mul_sample = fastmodulo((uint32_t) sample, mul_nsamples_packed);
|
|
mul += (int64_t) mul_sample * mul_stride_sample
|
|
+ (int64_t) mul_channel * mul_stride_channel
|
|
+ (int64_t) mul_row * mul_stride_row;
|
|
}
|
|
|
|
if constexpr (do_add) {
|
|
const int add_row = (int) fastmodulo((uint32_t) row, add_nrows_packed);
|
|
const int add_channel = (int) fastmodulo((uint32_t) channel, add_nchannels_packed);
|
|
const int add_sample = (int) fastmodulo((uint32_t) sample, add_nsamples_packed);
|
|
add += (int64_t) add_sample * add_stride_sample
|
|
+ (int64_t) add_channel * add_stride_channel
|
|
+ (int64_t) add_row * add_stride_row;
|
|
}
|
|
|
|
float tmp = 0.0f;
|
|
|
|
for (int col = tid; col < ncols; col += block_size) {
|
|
const float xi = x[col];
|
|
tmp += xi * xi;
|
|
}
|
|
|
|
extern __shared__ float s_sum[];
|
|
tmp = block_reduce_sum<block_size>(tmp, s_sum);
|
|
|
|
const float mean = tmp / ncols;
|
|
const float scale = rsqrtf(mean + eps);
|
|
|
|
for (int col = tid; col < ncols; col += block_size) {
|
|
if constexpr (do_multiply && do_add) {
|
|
const int mul_col = (int) fastmodulo((uint32_t) col, mul_ncols_packed);
|
|
const int add_col = (int) fastmodulo((uint32_t) col, add_ncols_packed);
|
|
dst[col] = scale * x[col] * mul[mul_col] + add[add_col];
|
|
} else if constexpr (do_multiply) {
|
|
const int mul_col = (int) fastmodulo((uint32_t) col, mul_ncols_packed);
|
|
dst[col] = scale * x[col] * mul[mul_col];
|
|
} else {
|
|
dst[col] = scale * x[col];
|
|
}
|
|
}
|
|
}
|
|
|
|
static void rms_norm_mul_f32_cuda(
|
|
const float * x,
|
|
const float * mul,
|
|
const float * add,
|
|
float * dst,
|
|
const int ncols,
|
|
const int nrows,
|
|
const int nchannels,
|
|
const int nsamples,
|
|
const int64_t stride_row,
|
|
const int64_t stride_channel,
|
|
const int64_t stride_sample,
|
|
const int64_t dst_stride_row,
|
|
const int64_t dst_stride_channel,
|
|
const int64_t dst_stride_sample,
|
|
const int64_t mul_stride_row,
|
|
const int64_t mul_stride_channel,
|
|
const int64_t mul_stride_sample,
|
|
const uint32_t mul_ncols,
|
|
const uint32_t mul_nrows,
|
|
const uint32_t mul_nchannels,
|
|
const uint32_t mul_nsamples,
|
|
const int64_t add_stride_row,
|
|
const int64_t add_stride_channel,
|
|
const int64_t add_stride_sample,
|
|
const uint32_t add_ncols,
|
|
const uint32_t add_nrows,
|
|
const uint32_t add_nchannels,
|
|
const uint32_t add_nsamples,
|
|
const float eps,
|
|
cudaStream_t stream) {
|
|
const dim3 blocks_num((unsigned) nrows, (unsigned) nchannels, (unsigned) nsamples);
|
|
const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols);
|
|
const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows);
|
|
const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
|
|
const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples);
|
|
|
|
// llama.cpp norm.cu:342-367: do not pack add fastdiv when add is null
|
|
// (init_fastdiv_values(0) is integer divide-by-zero).
|
|
if (add == nullptr) {
|
|
if (ncols < 1024) {
|
|
const dim3 block_dims(256, 1, 1);
|
|
rms_norm_f32<256, true, false><<<blocks_num, block_dims, 32 * sizeof(float), stream>>>(
|
|
x, dst, ncols, stride_row, stride_channel, stride_sample, eps,
|
|
dst_stride_row, dst_stride_channel, dst_stride_sample,
|
|
mul, mul_stride_row, mul_stride_channel, mul_stride_sample,
|
|
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
|
|
nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0),
|
|
make_uint3(0, 0, 0), make_uint3(0, 0, 0));
|
|
} else {
|
|
const dim3 block_dims(1024, 1, 1);
|
|
rms_norm_f32<1024, true, false><<<blocks_num, block_dims, 32 * sizeof(float), stream>>>(
|
|
x, dst, ncols, stride_row, stride_channel, stride_sample, eps,
|
|
dst_stride_row, dst_stride_channel, dst_stride_sample,
|
|
mul, mul_stride_row, mul_stride_channel, mul_stride_sample,
|
|
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
|
|
nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0),
|
|
make_uint3(0, 0, 0), make_uint3(0, 0, 0));
|
|
}
|
|
} else {
|
|
const uint3 add_ncols_packed = init_fastdiv_values(add_ncols);
|
|
const uint3 add_nrows_packed = init_fastdiv_values(add_nrows);
|
|
const uint3 add_nchannels_packed = init_fastdiv_values(add_nchannels);
|
|
const uint3 add_nsamples_packed = init_fastdiv_values(add_nsamples);
|
|
if (ncols < 1024) {
|
|
const dim3 block_dims(256, 1, 1);
|
|
rms_norm_f32<256, true, true><<<blocks_num, block_dims, 32 * sizeof(float), stream>>>(
|
|
x, dst, ncols, stride_row, stride_channel, stride_sample, eps,
|
|
dst_stride_row, dst_stride_channel, dst_stride_sample,
|
|
mul, mul_stride_row, mul_stride_channel, mul_stride_sample,
|
|
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
|
|
add, add_stride_row, add_stride_channel, add_stride_sample,
|
|
add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed);
|
|
} else {
|
|
const dim3 block_dims(1024, 1, 1);
|
|
rms_norm_f32<1024, true, true><<<blocks_num, block_dims, 32 * sizeof(float), stream>>>(
|
|
x, dst, ncols, stride_row, stride_channel, stride_sample, eps,
|
|
dst_stride_row, dst_stride_channel, dst_stride_sample,
|
|
mul, mul_stride_row, mul_stride_channel, mul_stride_sample,
|
|
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
|
|
add, add_stride_row, add_stride_channel, add_stride_sample,
|
|
add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed);
|
|
}
|
|
}
|
|
}
|