From b4d93e7b9901f673abb8a6da70416f8ca5a69795 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 3 May 2026 22:13:34 +0200 Subject: [PATCH 01/19] fused fwd --- src/kernels/kernels.cpp | 20 +++++ src/kernels/kernels.h | 26 ++++++ src/kernels/qk_norm.cu | 174 +++++++++++++++++++++++++++++++++++++++- 3 files changed, 217 insertions(+), 3 deletions(-) diff --git a/src/kernels/kernels.cpp b/src/kernels/kernels.cpp index cec30e9..d3e1547 100644 --- a/src/kernels/kernels.cpp +++ b/src/kernels/kernels.cpp @@ -121,6 +121,26 @@ void qk_norm_forward(Tensor& out, Tensor& r_rms, const Tensor& inp, const Tensor } } +void qk_norm_and_rope_forward(Tensor& out, Tensor& r_rms, float* abs_max_ptr, + const Tensor& inp, + const Tensor& q_wgt, const Tensor& k_wgt, + const Tensor& freqs_cis, + float epsilon, + int B, int T, int Nq, int Nkv, int HeadDim, + cudaStream_t stream) { + if (out.DType == ETensorDType::BF16) { + qk_norm_and_rope_forward(out.get(), r_rms.get(), abs_max_ptr, + inp.get(), q_wgt.get(), k_wgt.get(), + freqs_cis.get(), epsilon, B, T, Nq, Nkv, HeadDim, stream); + } else if (out.DType == ETensorDType::FP32) { + qk_norm_and_rope_forward(out.get(), r_rms.get(), abs_max_ptr, + inp.get(), q_wgt.get(), k_wgt.get(), + freqs_cis.get(), epsilon, B, T, Nq, Nkv, HeadDim, stream); + } else { + UNSUPPORTED_DTYPE(out, r_rms, inp, q_wgt, k_wgt, freqs_cis); + } +} + void qk_norm_backward(Tensor& dinp, Tensor& dq_wgt, Tensor& dk_wgt, Tensor& scratch, const Tensor& dout, const Tensor& inp, const Tensor& q_wgt, const Tensor& k_wgt, diff --git a/src/kernels/kernels.h b/src/kernels/kernels.h index 9c4f09a..214a9d3 100644 --- a/src/kernels/kernels.h +++ b/src/kernels/kernels.h @@ -129,6 +129,32 @@ void qk_norm_forward(Tensor& out, Tensor& r_rms, const Tensor& inp, float epsilon, int BT, int Nq, int Nkv, int HeadDim, cudaStream_t stream); + +void qk_norm_and_rope_forward(float* out, float* r_rms, float* abs_max_ptr, + const float* inp, + const float* q_wgt, const float* k_wgt, + const float* freqs_cis, + float epsilon, + int B, int T, int Nq, int Nkv, int HeadDim, + cudaStream_t stream); + +void qk_norm_and_rope_forward(nv_bfloat16* out, float* r_rms, float* abs_max_ptr, + const nv_bfloat16* inp, + const nv_bfloat16* q_wgt, const nv_bfloat16* k_wgt, + const half* freqs_cis, + float epsilon, + int B, int T, int Nq, int Nkv, int HeadDim, + cudaStream_t stream); + +void qk_norm_and_rope_forward(Tensor& out, Tensor& r_rms, float* abs_max_ptr, + const Tensor& inp, + const Tensor& q_wgt, const Tensor& k_wgt, + const Tensor& freqs_cis, + float epsilon, + int B, int T, int Nq, int Nkv, int HeadDim, + cudaStream_t stream); + + std::size_t qk_norm_backward_scratch_size(int Nq, int Nkv, int HeadDim, ETensorDType dtype, const cudaDeviceProp& dp); void qk_norm_backward(float* dinp, float* dq_wgt, float* dk_wgt, std::byte* scratch, diff --git a/src/kernels/qk_norm.cu b/src/kernels/qk_norm.cu index f67bf5c..5d73903 100644 --- a/src/kernels/qk_norm.cu +++ b/src/kernels/qk_norm.cu @@ -20,7 +20,7 @@ template struct QKHelpers { using x128 = GenericVector; - static __device__ __forceinline__ float norm_head(Float* out, x128* s_in, const x128* wgt_src, const Float* inp, int HeadDim, float epsilon) { + static __device__ __forceinline__ float norm_head(Float* out, x128* s_in, const x128* wgt_src, const Float* inp, int HeadDim, float epsilon, int active_mask) { using x128 = GenericVector; const int lane = threadIdx.x % GroupSize; @@ -35,7 +35,7 @@ struct QKHelpers { } } - acc = reduce_group_sum(acc, 0xffffffff) / HeadDim; + acc = reduce_group_sum(acc, active_mask) / HeadDim; float s = rsqrtf(acc + epsilon); for(int c = lane * x128::size; c < HeadDim; c += GroupSize * x128::size) { @@ -106,7 +106,7 @@ __global__ void qk_norm_forward_simple_kernel(Float* out, float* r_rms, const Fl return; } - float s = QKHelpers::norm_head(out, s_in, wgt_src, inp, HeadDim, epsilon); + float s = QKHelpers::norm_head(out, s_in, wgt_src, inp, HeadDim, epsilon, 0xffffffff); // store the rms, no need to cache it if(lane == 0 && r_rms != nullptr) { @@ -307,6 +307,119 @@ qk_norm_backward_reduce_kernel(Float* dq_wgt, Float* dk_wgt, const std::byte* sc } } +// fused kernels + +template +__global__ void qk_norm_and_rope_fwd_kernel(Float* out, float* r_rms, float* abs_max_ptr, + const Float* inp, + const Float* q_wgt, const Float* k_wgt, + const FloatFreq* freqs_cis, + float epsilon, int BT, int T, int Nq, int Nkv, int HeadDim) { + static_assert(sizeof(Float) == sizeof(FloatFreq), "Float and FloatFreq must have the same size"); + using x128 = GenericVector; + using freq128 = GenericVector; + + const int lane = threadIdx.x % GroupSize; + const int group = threadIdx.x / GroupSize; + __shared__ float block_abs_max; + if (abs_max_ptr && threadIdx.x == 0) { + block_abs_max = 0.f; + } + float thread_abs_max = 0.f; + + // load weights into shared memory + // do this before we allow any threads to exit! + extern __shared__ char* smem[]; + + // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so + // let's keep everything as x128 + x128* s_q_wgt = reinterpret_cast(smem); + x128* s_k_wgt = reinterpret_cast(smem) + (HeadDim / x128::size); + x128* s_in = reinterpret_cast(s_k_wgt) + (HeadDim / x128::size) + group * (HeadDim / x128::size); + + for(int i = threadIdx.x * x128::size; i < HeadDim; i += blockDim.x * x128::size) { + s_q_wgt[i/x128::size] = x128::load(q_wgt + i); + s_k_wgt[i/x128::size] = x128::load(k_wgt + i); + } + + __syncthreads(); + int idx = blockIdx.x * (blockDim.x / GroupSize) + group; + + int h = idx % (Nq + 2 * Nkv); + int bt = idx / (Nq + 2 * Nkv); + int t = bt % T; + + if (bt >= BT) return; + + // adjust pointers to current token + const Float* inp_h = inp + idx * HeadDim; + Float* out_h = out + idx * HeadDim; + + const x128* wgt_src = nullptr; + bool is_value_head = false; + if (h < Nq) { + wgt_src = s_q_wgt; + } else if (h < Nq + Nkv) { + wgt_src = s_k_wgt; + } else { + is_value_head = true; + } + + unsigned int active_mask = __ballot_sync(0xffffffffu, !is_value_head); + if (is_value_head) { + // ---- value head: pass-through, but still contribute to abs-max ---- + for (int c = lane * x128::size; c < HeadDim; c += GroupSize * x128::size) { + x128 v = x128::load_cs(inp_h + c); + for (int k = 0; k < x128::size; ++k) { + thread_abs_max = fmaxf(thread_abs_max, fabsf((float)v[k])); + } + if (inp_h != out_h) v.store(out_h + c); + } + } else { + // ---- Q or K head: norm + scale + RoPE ---- + + float s = QKHelpers::norm_head(&s_in[0][0], s_in, wgt_src, inp_h, HeadDim, epsilon, active_mask); + __syncwarp(active_mask); + + int head_dim_half = HeadDim / 2; + using x64 = GenericVector; + Float* s_in_f = reinterpret_cast(s_in); + for (int c = lane * x64::size; c < head_dim_half; c += GroupSize * x64::size) { + x64 v_real = x64::load(s_in_f + c); + x64 v_imag = x64::load(s_in_f + c + head_dim_half); + + freq128 freqs_vec = freq128::load_ldg(freqs_cis + t * HeadDim + 2 * c); + + x64 o_real, o_imag; + for (int k = 0; k < x64::size; ++k) { + float cos = (float)freqs_vec[2*k]; + float sin = (float)freqs_vec[2*k+1]; + float real = (float)v_real[k]; + float imag = (float)v_imag[k]; + float or_ = real * cos - imag * sin; + float oi_ = real * sin + imag * cos; + o_real[k] = (Float)or_; + o_imag[k] = (Float)oi_; + if (abs_max_ptr) { + thread_abs_max = fmaxf(thread_abs_max, fabsf(or_)); + thread_abs_max = fmaxf(thread_abs_max, fabsf(oi_)); + } + } + o_real.store(out_h + c); + o_imag.store(out_h + c + head_dim_half); + } + + // store the rms, no need to cache it + if (lane == 0 && r_rms != nullptr) { + __stcs(r_rms + idx, s); + } + } + + if (abs_max_ptr) { + handle_absmax_reduction(abs_max_ptr, &block_abs_max, thread_abs_max); + } +} + template void qk_norm_forward_imp(Float* out, float* r_rms, const Float* inp, @@ -339,6 +452,38 @@ void qk_norm_forward_imp(Float* out, float* r_rms, CUDA_CHECK(cudaGetLastError()); } +template +void qk_norm_and_rope_forward(Float* out, float* r_rms, float* abs_max_ptr, + const Float* inp, + const Float* q_wgt, const Float* k_wgt, + const FloatFreq* freqs_cis, + float epsilon, + int B, int T, int Nq, int Nkv, int HeadDim, + cudaStream_t stream) { + static_assert(sizeof(Float) == sizeof(FloatFreq), "Float and FloatFreq must have the same size"); + constexpr int block_size = 512; // larger blocks mean fewer redundant weight loads + static_assert(block_size % GroupSize == 0); + + const int BT = B * T; + const int groups_per_block = block_size / GroupSize; + const int total_heads = BT * (Nq + 2 * Nkv); + const int grid_size = div_ceil(total_heads, groups_per_block); + + // smem: q_wgt + k_wgt + one input buffer per group + size_t smem = (2 + groups_per_block) * HeadDim * sizeof(Float); + + CUDA_CHECK(cudaFuncSetAttribute( + qk_norm_and_rope_fwd_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem)); + + qk_norm_and_rope_fwd_kernel<<>>( + out, r_rms, abs_max_ptr, inp, q_wgt, k_wgt, freqs_cis, + epsilon, BT, T, Nq, Nkv, HeadDim); + + CUDA_CHECK(cudaGetLastError()); +} + // Get the amount of smem per block for backward template size_t qk_norm_backward_smem(int HeadDim) { @@ -427,6 +572,29 @@ void qk_norm_forward(nv_bfloat16* out, float* r_rms, const nv_bfloat16* inp, qk_norm_forward_imp(out, r_rms, inp, q_wgt, k_wgt, epsilon, BT, Nq, Nkv, HeadDim, stream); } + +void qk_norm_and_rope_forward(float* out, float* r_rms, float* abs_max_ptr, + const float* inp, + const float* q_wgt, const float* k_wgt, + const float* freqs_cis, + float epsilon, + int B, int T, int Nq, int Nkv, int HeadDim, + cudaStream_t stream) { + qk_norm_and_rope_forward(out, r_rms, abs_max_ptr, inp, q_wgt, k_wgt, + freqs_cis, epsilon, B, T, Nq, Nkv, HeadDim, stream); +} + +void qk_norm_and_rope_forward(nv_bfloat16* out, float* r_rms, float* abs_max_ptr, + const nv_bfloat16* inp, + const nv_bfloat16* q_wgt, const nv_bfloat16* k_wgt, + const half* freqs_cis, + float epsilon, + int B, int T, int Nq, int Nkv, int HeadDim, + cudaStream_t stream) { + qk_norm_and_rope_forward(out, r_rms, abs_max_ptr, inp, q_wgt, k_wgt, + freqs_cis, epsilon, B, T, Nq, Nkv, HeadDim, stream); +} + void qk_norm_backward(float* dinp, float* dq_wgt, float* dk_wgt, std::byte* scratch, const float* dout, const float* inp, const float* q_wgt, const float* k_wgt, From ace29001eda7ed818f4b8ad1d83d68f5bc61923e Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 3 May 2026 22:35:57 +0200 Subject: [PATCH 02/19] bindings and test --- scripts/test_kernels.py | 41 ++++++++++++++++++++++++++++++++++ src/binding/kernel_binding.cpp | 30 +++++++++++++++++++++++++ src/binding/python/kernels.py | 8 +++++++ 3 files changed, 79 insertions(+) diff --git a/scripts/test_kernels.py b/scripts/test_kernels.py index dead372..eece9ec 100644 --- a/scripts/test_kernels.py +++ b/scripts/test_kernels.py @@ -508,6 +508,47 @@ def test_qk_norm_forward_inplace(B, T, Nq, Nkv, HeadDim, dtype): rtol, atol = TOL[dtype] assert out.float().cpu().numpy() == pytest.approx(ref_out.float().cpu().numpy(), rel=rtol, abs=atol) + +@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ + (1, 4, 2, 1, 16), + (2, 8, 4, 2, 64), + (4, 16, 8, 4, 128), +]) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("with_abs_max", [False, True]) +def test_qk_norm_and_rope_forward(B, T, Nq, Nkv, HeadDim, dtype, with_abs_max): + torch.manual_seed(0) + device = "cuda" + N = Nq + 2 * Nkv + eps = 1e-6 + + inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + out = torch.empty_like(inp) + # NOTE: the kernel does not write rrms for value heads; init to 1 to match reference + r_rms = torch.ones((B, T, N), device=device, dtype=torch.float32) + + freq_dtype = torch.float32 if dtype == torch.float32 else torch.float16 + freqs = _make_rope_freqs(T, HeadDim, 1_000_000.0, freq_dtype, device) + + abs_max = torch.zeros((1,), device=device, dtype=torch.float32) if with_abs_max else None + + K.qk_norm_and_rope_forward(out, r_rms, inp, q_wgt, k_wgt, freqs, abs_max, eps, Nq, Nkv) + + # reference: qk-norm first, then rope on the result + ref_normed, ref_r_rms = _qk_norm_reference(inp, q_wgt, k_wgt, Nq, Nkv, eps) + ref_out = _rope_python(ref_normed.view(B, T, N, HeadDim), freqs, Nq, Nkv).view(B, T, N * HeadDim) + + rtol, atol = TOL[dtype] + assert r_rms.cpu().numpy() == pytest.approx(ref_r_rms.cpu().numpy(), rel=rtol, abs=atol) + assert out.float().cpu().numpy() == pytest.approx(ref_out.float().cpu().numpy(), rel=rtol, abs=atol) + + if with_abs_max: + # kernel takes abs-max over Q+K post-rope and V pass-through + ref_abs_max = ref_out.float().abs().max().item() + assert abs_max.item() == pytest.approx(ref_abs_max, rel=rtol, abs=atol) + @pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ (1, 4, 2, 1, 16), (2, 8, 4, 2, 64), diff --git a/src/binding/kernel_binding.cpp b/src/binding/kernel_binding.cpp index f976ab9..c3cfcf8 100644 --- a/src/binding/kernel_binding.cpp +++ b/src/binding/kernel_binding.cpp @@ -244,6 +244,33 @@ void bind_qk_norm_forward(const CudaArray& out, const CudaArray& r_std, const Cu epsilon, B * T, Nq, Nkv, HeadDim, as_stream(stream)); } +void bind_qk_norm_and_rope_forward(const CudaArray& out, const CudaArray& r_std, + const CudaArray& inp, + const CudaArray& q_wgt, const CudaArray& k_wgt, + const CudaArray& freqs_cis, + const std::optional& abs_max, + float epsilon, int Nq, int Nkv, + const std::uintptr_t stream) { + NB_CHECK_NDIMS(out, 3); + NB_CHECK_NDIMS(inp, 3); + NB_CHECK_NDIMS(r_std, 3); + NB_CHECK_NDIMS(q_wgt, 1); + NB_CHECK_NDIMS(k_wgt, 1); + NB_CHECK_NDIMS(freqs_cis, 2); + + const long B = get_dimension_checked({out.shape(0), inp.shape(0), r_std.shape(0)}, "B"); + const long T = get_dimension_checked({out.shape(1), inp.shape(1), r_std.shape(1), freqs_cis.shape(0)}, "T"); + const long HeadDim = get_dimension_checked({q_wgt.shape(0), k_wgt.shape(0), freqs_cis.shape(1)}, "HeadDim"); + (void)get_dimension_checked({out.shape(2), inp.shape(2), (std::size_t)((Nq + 2 * Nkv) * HeadDim)}, "NTotal*HeadDim"); + (void)get_dimension_checked({r_std.shape(2), (std::size_t)(Nq + 2 * Nkv)}, "Nq+2*Nkv"); + + Tensor out_t = to_tensor(out); + Tensor r_std_t = to_tensor(r_std); + qk_norm_and_rope_forward(out_t, r_std_t, get_abs_max_ptr(abs_max), + to_tensor(inp), to_tensor(q_wgt), to_tensor(k_wgt), to_tensor(freqs_cis), + epsilon, B, T, Nq, Nkv, HeadDim, as_stream(stream)); +} + void bind_qk_norm_backward(const CudaArray& dinp, const CudaArray& dq_wgt, const CudaArray& dk_wgt, const CudaArray& scratch, const CudaArray& dout, const CudaArray& inp, @@ -599,6 +626,9 @@ void register_kernels(nanobind::module_& m) { // QK Norm m.def("qk_norm_forward", &bind_qk_norm_forward, nb::arg("out"), nb::arg("r_rms"), nb::arg("inp"), nb::arg("q_wgt"), nb::arg("k_wgt"), nb::arg("epsilon"), nb::arg("Nq"), nb::arg("Nkv"), nb::arg("stream") = 0); + m.def("qk_norm_and_rope_forward", &bind_qk_norm_and_rope_forward, nb::arg("out"), nb::arg("r_rms"), nb::arg("inp"), + nb::arg("q_wgt"), nb::arg("k_wgt"), nb::arg("freqs_cis"), nb::arg("abs_max") = std::nullopt, + nb::arg("epsilon"), nb::arg("Nq"), nb::arg("Nkv"), nb::arg("stream") = 0); m.def("qk_norm_backward", &bind_qk_norm_backward, nb::arg("dinp"), nb::arg("dq_wgt"), nb::arg("dk_wgt"), nb::arg("scratch"), nb::arg("dout"), nb::arg("inp"), nb::arg("q_wgt"), nb::arg("k_wgt"), nb::arg("rstd"), nb::arg("abs_max") = std::nullopt, nb::arg("Nq"), nb::arg("Nkv"), nb::arg("stream") = 0); diff --git a/src/binding/python/kernels.py b/src/binding/python/kernels.py index b1f257f..9059905 100644 --- a/src/binding/python/kernels.py +++ b/src/binding/python/kernels.py @@ -65,6 +65,14 @@ def qk_norm_forward(out: torch.Tensor, r_rms: torch.Tensor, inp: torch.Tensor, epsilon: float, Nq: int, Nkv: int, stream: int = 0) -> None: _pyllmq.qk_norm_forward(out, r_rms, inp, q_wgt, k_wgt, epsilon, Nq, Nkv, stream) +@torch.library.custom_op("llmq::qk_norm_and_rope_forward", mutates_args=("out", "r_rms", "abs_max")) +def qk_norm_and_rope_forward(out: torch.Tensor, r_rms: torch.Tensor, inp: torch.Tensor, + q_wgt: torch.Tensor, k_wgt: torch.Tensor, freqs_cis: torch.Tensor, + abs_max: torch.Tensor | None, + epsilon: float, Nq: int, Nkv: int, stream: int = 0) -> None: + _pyllmq.qk_norm_and_rope_forward(out, r_rms, inp, q_wgt, k_wgt, freqs_cis, abs_max, + epsilon, Nq, Nkv, stream) + @torch.library.custom_op("llmq::qk_norm_backward", mutates_args=("dinp", "dq_wgt", "dk_wgt", "scratch", "abs_max")) def qk_norm_backward(dinp: torch.Tensor, dq_wgt: torch.Tensor, dk_wgt: torch.Tensor, scratch: torch.Tensor, dout: torch.Tensor, inp: torch.Tensor, From b76c31f274f433bd9e007d0cad3ad54a962fbb4c Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 26 May 2026 00:18:24 +0200 Subject: [PATCH 03/19] split tests file --- scripts/test_kernels.py | 1226 ---------------------------------- test/conftest.py | 10 + test/test_kernels_encoder.py | 114 ++++ test/test_kernels_matmul.py | 209 ++++++ test/test_kernels_other.py | 218 ++++++ test/test_kernels_qk_norm.py | 201 ++++++ test/test_kernels_quant.py | 91 +++ test/test_kernels_rmsnorm.py | 143 ++++ test/test_kernels_rope.py | 73 ++ test/test_kernels_sr.py | 128 ++++ test/test_kernels_swiglu.py | 74 ++ 11 files changed, 1261 insertions(+), 1226 deletions(-) delete mode 100644 scripts/test_kernels.py create mode 100644 test/conftest.py create mode 100644 test/test_kernels_encoder.py create mode 100644 test/test_kernels_matmul.py create mode 100644 test/test_kernels_other.py create mode 100644 test/test_kernels_qk_norm.py create mode 100644 test/test_kernels_quant.py create mode 100644 test/test_kernels_rmsnorm.py create mode 100644 test/test_kernels_rope.py create mode 100644 test/test_kernels_sr.py create mode 100644 test/test_kernels_swiglu.py diff --git a/scripts/test_kernels.py b/scripts/test_kernels.py deleted file mode 100644 index eece9ec..0000000 --- a/scripts/test_kernels.py +++ /dev/null @@ -1,1226 +0,0 @@ -import math - -import pytest -import torch -import torch.nn.functional as F -from pyllmq import kernels as K - - -DTYPES = [torch.float32, torch.bfloat16] - -# Tolerances by dtype for approximate kernels: (rtol, atol) -TOL = { - torch.float32: (5e-4, 5e-5), - torch.bfloat16: (1e-2, 5e-3), -} - - -# ============================================================ -# fill_constant -# Fills every element with a compile-time constant: result must -# be bit-exact, so rel=0, abs=0. -# ============================================================ - -@pytest.mark.parametrize("shape", [(8, 16), (1,), (128, 256), (4, 8, 32)]) -@pytest.mark.parametrize("value", [0.0, 1.0, 3.25, -2.5]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_fill_constant_basic(shape, value, dtype): - x = torch.empty(shape, device="cuda", dtype=dtype) - K.fill_constant(x, value) - ref = torch.full(shape, value, dtype=dtype) - # Exact: every element is independently written to the same constant. - assert x.float().cpu() == pytest.approx(ref.float().cpu(), rel=0, abs=0) - - -# ============================================================ -# transpose -# Byte-level shuffle — no arithmetic, so result must be exact. -# ============================================================ - -@pytest.mark.parametrize("rows,cols", [(7, 11), (1024, 2048), (64, 128), (3, 512)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_transpose_matches_torch(rows, cols, dtype): - src = torch.randn((rows, cols), device="cuda", dtype=dtype) - dst = torch.empty((cols, rows), device="cuda", dtype=dtype) - K.transpose(dst, src) - # Transposing rearranges values without touching bits — must be exact. - assert dst.float().cpu() == pytest.approx(src.t().float().cpu(), rel=0, abs=0) - - -# ============================================================ -# abs_max -# Reduction over exact fp values — output is one of the input -# values, so the scalar result must be exact. -# ============================================================ - -@pytest.mark.parametrize("shape", [(16, 64), (4, 128, 32)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_abs_max_writes_scalar(shape, dtype): - x = torch.randn(shape, device="cuda", dtype=dtype) - ref = x.abs().max() - result = torch.empty((), device="cuda", dtype=torch.float32) - K.abs_max(result, x) - assert torch.isfinite(result) - # abs_max just selects an existing value — no arithmetic error possible. - assert result.cpu() == pytest.approx(ref.float().cpu(), rel=0, abs=0) - - -# ============================================================ -# global_norm_squared -# Involves floating-point accumulation so tolerances are needed. -# ============================================================ - -@pytest.mark.parametrize("shape", [(128, 24524), (256, 1024), (1, 4096)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_global_norm_squared(shape, dtype): - x = torch.randn(*shape, device="cuda", dtype=dtype) - out = torch.zeros((256,), device="cuda", dtype=torch.float32) - K.global_norm_squared(out, x) - ref = (x.float() ** 2).sum() - rtol = 1e-5 if dtype == torch.float32 else 1e-2 - assert out.sum().cpu() == pytest.approx(ref.cpu(), rel=rtol) - - -# ============================================================ -# rmsnorm -# ============================================================ - -def _rmsnorm_reference(inp, weight, eps): - var = (inp.float() ** 2).mean(dim=-1, keepdim=True) - rms = torch.sqrt(var + eps) - r_rms = 1.0 / rms - out = inp.float() * r_rms * weight.float() - return out.to(inp.dtype), r_rms.squeeze(-1) - - -def _rmsnorm_backward_reference(dout, inp, weight, rstd): - # dout, inp: (B, T, C), weight: (C,), rstd: (B, T) - B, T, C = inp.shape - dout_f = dout.float() - inp_f = inp.float() - w_f = weight.float() - r = rstd.unsqueeze(-1) # (B, T, 1) - - # dweight: sum over B, T - dweight = (dout_f * inp_f * r).sum(dim=(0, 1)) - - # dinp - normed = inp_f * r # (B, T, C) - dy_w = dout_f * w_f # (B, T, C) - dot = (dy_w * normed).sum(dim=-1, keepdim=True) # (B, T, 1) - dinp = r * (dy_w - normed * dot / C) - - return dinp.to(inp.dtype), dweight.to(weight.dtype) - - -@pytest.mark.parametrize("B,T,C", [(2, 3, 16), (1, 2, 64), (8, 256, 1024)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_rmsnorm_forward_matches_reference(B, T, C, dtype): - torch.manual_seed(0) - inp = torch.randn((B, T, C), device="cuda", dtype=dtype) - weight = torch.randn((C,), device="cuda", dtype=dtype) - out = torch.empty_like(inp) - rms = torch.empty((B, T), device="cuda", dtype=torch.float32) - eps = 1e-6 - - K.rmsnorm_forward(out, rms, inp, weight, None, eps) - ref_out, ref_rms = _rmsnorm_reference(inp, weight, eps) - - rtol, atol = TOL[dtype] - assert rms.cpu() == pytest.approx(ref_rms.cpu(), rel=rtol, abs=atol) - assert out.float().cpu() == pytest.approx(ref_out.float().cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("B,T,C", [(2, 3, 16), (1, 4, 64), (4, 8, 256)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_rmsnorm_backward(B, T, C, dtype): - torch.manual_seed(0) - inp = torch.randn((B, T, C), device="cuda", dtype=dtype) - weight = torch.randn((C,), device="cuda", dtype=dtype) - dout = torch.randn((B, T, C), device="cuda", dtype=dtype) - eps = 1e-6 - - # Forward to get rstd - _, rstd = _rmsnorm_reference(inp, weight, eps) - rstd = rstd.to(torch.float32).cuda() - - dinp = torch.empty_like(inp) - dweight = torch.zeros((C,), device="cuda", dtype=dtype) - scratch = torch.zeros(K.get_rmsnorm_backward_scratch_size(C), device="cuda", dtype=torch.float32) - dresidual = torch.zeros_like(inp) - - K.rmsnorm_backward(dinp, dweight, scratch, dresidual, dout, inp, weight, rstd, None) - - ref_dinp, ref_dweight = _rmsnorm_backward_reference(dout, inp, weight, rstd) - rtol, atol = TOL[dtype] - assert dinp.float().cpu() == pytest.approx(ref_dinp.float().cpu(), rel=rtol, abs=atol) - assert dweight.float().cpu() == pytest.approx(ref_dweight.float().cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("B,T,C", [(2, 3, 16), (1, 4, 64), (4, 8, 256)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_rmsnorm_backward_dresidual_accumulate(B, T, C, dtype): - torch.manual_seed(1) - inp = torch.randn((B, T, C), device="cuda", dtype=dtype) - weight = torch.randn((C,), device="cuda", dtype=dtype) - dout = torch.randn((B, T, C), device="cuda", dtype=dtype) - eps = 1e-6 - - _, rstd = _rmsnorm_reference(inp, weight, eps) - rstd = rstd.to(torch.float32).cuda() - - dinp = torch.empty_like(inp) - dweight = torch.zeros((C,), device="cuda", dtype=dtype) - scratch = torch.zeros(K.get_rmsnorm_backward_scratch_size(C), device="cuda", dtype=torch.float32) - - # Pre-fill dresidual with a non-zero sentinel so accumulation is detectable. - dresidual = torch.randn_like(inp) - dresidual_initial = dresidual.clone() - ref_dinp, _ = _rmsnorm_backward_reference(dout, inp, weight, rstd) - expected_dinp = dresidual_initial.float() + ref_dinp.float() - - K.rmsnorm_backward(dinp, dweight, scratch, dresidual, dout, inp, weight, rstd, None) - - rtol, atol = TOL[dtype] - if dtype == torch.bfloat16: - # TODO check why we need so much tolerance - atol = 8e-3 - rtol = 0.015 - assert dinp.float().cpu() == pytest.approx( - expected_dinp.cpu(), rel=rtol, abs=atol - ) - -# ============================================================ -# swiglu_forward -# ============================================================ - -def _swiglu_reference(x: torch.Tensor) -> torch.Tensor: - a, b = torch.tensor_split(x.float(), 2, dim=-1) - return (torch.nn.functional.silu(b) * a).to(x.dtype) - - -def _swiglu_backward_reference(dout, inp): - inp_f = inp.detach().float().requires_grad_(True) - out = _swiglu_reference(inp_f) - out.backward(dout.float()) - return inp_f.grad.to(inp.dtype) - - -@pytest.mark.parametrize("B,T,C", [(1, 16, 128), (8, 8, 16), (5, 32, 32), (32, 32, 1024)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_swiglu_forward_matches_reference(B, T, C, dtype): - torch.manual_seed(123) - inp = torch.randn((B, T, 2 * C), device="cuda", dtype=dtype) - out = torch.empty((B, T, C), device="cuda", dtype=dtype) - - K.swiglu_forward(out, inp, None) - - ref = _swiglu_reference(inp) - rtol, atol = TOL[dtype] - assert out.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("B,T,C", [(2, 8, 64), (4, 16, 128)]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_swiglu_forward_quant_fp8(B, T, C, dtype): - """swiglu_forward_quant writes fp8 output, using the supplied abs-max to scale values""" - torch.manual_seed(42) - inp = torch.randn((B, T, 2 * C), device="cuda", dtype=dtype) - out = torch.empty((B, T, C), device="cuda", dtype=torch.float8_e4m3fn) - scale = torch.empty((), device="cuda", dtype=torch.float32) - - ref = _swiglu_reference(inp.float()) - abs_max = ref.abs().max().float() - K.swiglu_forward_quant(out, scale, inp, abs_max) - - dequant = out.float() * scale - # fp8 has limited precision — use loose tolerance - assert dequant.cpu() == pytest.approx(ref.cpu(), rel=0.125, abs=1e-2) - - expected_scale = abs_max / 448.0 - assert scale.cpu() == pytest.approx(expected_scale.cpu(), rel=1e-3) - -@pytest.mark.parametrize("B,T,C", [(1, 8, 256), (4, 16, 64)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_swiglu_backward(B, T, C, dtype): - torch.manual_seed(0) - inp = torch.randn((B, T, 2 * C), device="cuda", dtype=dtype) - dout = torch.randn((B, T, C), device="cuda", dtype=dtype) - dinp = torch.empty_like(inp) - - K.swiglu_backward(dinp, dout, inp, None) - - ref = _swiglu_backward_reference(dout, inp) - rtol, atol = TOL[dtype] - - assert dinp.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) - - -# ============================================================ -# grouped_loss_sum -# ============================================================ - -@pytest.mark.parametrize("B,T", [(1, 512), (2, 1024), (7, 2048), (4, 512)]) -def test_grouped_loss_sum_basic(B, T): - losses = torch.rand((B, T), device="cuda", dtype=torch.float32) - out = torch.empty((T // 512,), device="cuda", dtype=torch.float32) - K.grouped_loss_sum(out, losses) - ref = losses.reshape(B, -1, 512).sum(dim=2).sum(dim=0) - assert out.cpu() == pytest.approx(ref.cpu(), rel=1e-5, abs=1e-6) - - -# ============================================================ -# fill_normal -# ============================================================ - -@pytest.mark.parametrize("mean,std", [(0.0, 1.0), (1.5, 0.25), (-1.0, 2.0)]) -def test_fill_normal_stats(mean, std): - N = 256 * 1024 - seed = 0xABCDEF01 - x = torch.empty((N,), device="cuda", dtype=torch.float32) - K.fill_normal(x, float(mean), float(std), seed, 0) - - m = x.mean().item() - s = x.std(unbiased=True).item() - - assert abs(m - mean) < max(3e-3, 0.02 * abs(std)) - assert abs(s - std) / max(std, 1e-12) < 0.03 - - # Determinism: same seed + subsequence must give bit-identical output. - y = torch.empty_like(x) - K.fill_normal(y, float(mean), float(std), seed, 0) - assert x.cpu() == pytest.approx(y.cpu(), rel=0, abs=0) - -@pytest.mark.parametrize("mean,std", [(0.0, 1.0), (1.5, 0.25)]) -def test_fill_normal_subsequence_independence(mean, std): - """Different subsequence offsets with the same seed must produce - non-identical, uncorrelated draws.""" - N = 64 * 1024 - seed = 0xABCDEF01 - - tensors = [] - for subseq in range(4): - x = torch.empty((N,), device="cuda", dtype=torch.float32) - K.fill_normal(x, float(mean), float(std), seed, subseq) - tensors.append(x.cpu()) - - for i in range(len(tensors)): - for j in range(i + 1, len(tensors)): - # Bit-identity: the subsequence argument must actually be used. - assert not torch.equal(tensors[i], tensors[j]), ( - f"Subsequences {i} and {j} produced identical output" - ) - - # Pearson |r| < 0.05 is a very conservative bound at N=64k - # (3-sigma ≈ 0.012), so a genuine failure stands out clearly. - a = tensors[i] - tensors[i].mean() - b = tensors[j] - tensors[j].mean() - r = (a * b).mean() / (a.std() * b.std()) - assert abs(r.item()) < 0.05, ( - f"Subsequences {i} and {j} are correlated: r={r.item():.4f}" - ) - -@pytest.mark.parametrize("mean,std", [(0.0, 1.0), (1.5, 0.25), (-1.0, 2.0)]) -def test_fill_normal_stats_bf16(mean, std): - """BF16 should satisfy the same mean/std checks as float32 with - tolerances relaxed to reflect its ~0.8 % per-value precision.""" - N = 256 * 1024 - seed = 0xABCDEF01 - x = torch.empty((N,), device="cuda", dtype=torch.bfloat16) - K.fill_normal(x, float(mean), float(std), seed, 0) - - # Upcast before reducing — BF16 accumulation is itself lossy. - xf = x.float() - m = xf.mean().item() - s = xf.std(unbiased=True).item() - - assert abs(m - mean) < max(1e-2, 0.05 * abs(std)) - assert abs(s - std) / max(std, 1e-12) < 0.05 - - # Determinism check mirrors the float32 test. - y = torch.empty_like(x) - K.fill_normal(y, float(mean), float(std), seed, 0) - assert torch.equal(x.cpu(), y.cpu()) - -# ============================================================ -# rope forward + backward -# ============================================================ - -def _make_rope_freqs(T, head_dim, theta, dtype, device): - assert head_dim % 2 == 0 - half = head_dim // 2 - idx = torch.arange(half, device=device, dtype=torch.float32) - inv_freq = theta ** (-2 * idx / head_dim) - t = torch.arange(T, device=device, dtype=torch.float32).unsqueeze(1) - angles = t * inv_freq.unsqueeze(0) - cos = torch.cos(angles).to(dtype) - sin = torch.sin(angles).to(dtype) - return torch.stack([cos, sin], dim=-1).flatten(start_dim=1) # (T, head_dim) - - -def _rope_python(x, freqs_cis, Nq, Nkv, backward=False): - B, T, N, HD = x.shape - half = HD // 2 - cos = freqs_cis[:, 0::2].float() - sin = freqs_cis[:, 1::2].float() - if backward: - sin = -sin - cos = cos[None, :, None, :] - sin = sin[None, :, None, :] - - q = x[:, :, :Nq, :] - k = x[:, :, Nq:Nq + Nkv, :] - v = x[:, :, Nq + Nkv:, :] - - def rotate(h): - h = h.float() - return torch.cat([ - h[..., :half] * cos - h[..., half:] * sin, - h[..., :half] * sin + h[..., half:] * cos, - ], dim=-1).to(x.dtype) - - return torch.cat([rotate(q), rotate(k), v], dim=2) - - -@pytest.mark.parametrize("B,T,Nq,Nkv,HD", [ - (1, 8, 2, 1, 8), - (2, 4, 1, 2, 16), - (2, 16, 4, 2, 32), - (4, 32, 8, 4, 64), -]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_rope_forward_backward_matches_python(B, T, Nq, Nkv, HD, dtype): - device = "cuda" - C = (Nq + 2 * Nkv) * HD - x = torch.randn((B, T, C), device=device, dtype=dtype) - out = torch.empty_like(x) - - freq_dtype = torch.float32 if dtype == torch.float32 else torch.float16 - freqs = _make_rope_freqs(T, HD, 1_000_000.0, freq_dtype, device) - - K.rope_forward(out, x, freqs, None, Nq, Nkv) - ref = _rope_python(x.view(B, T, Nq + 2 * Nkv, HD), freqs, Nq, Nkv).view(B, T, C) - - rtol, atol = TOL[dtype] - assert out.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) - - dout = torch.randn_like(x) - dinp = torch.empty_like(x) - K.rope_backward(dinp, dout, freqs, None, Nq, Nkv) - ref_bw = _rope_python(dout.view(B, T, Nq + 2 * Nkv, HD), freqs, Nq, Nkv, backward=True).view(B, T, C) - assert dinp.float().cpu() == pytest.approx(ref_bw.float().cpu(), rel=rtol, abs=atol) - - -# ============================================================ -# qk_norm_forward -# ============================================================ - -def _qk_norm_reference(inp, q_wgt, k_wgt, Nq, Nkv, eps): - B, T, _ = inp.shape - N = Nq + 2 * Nkv - HeadDim = q_wgt.shape[0] - x = inp.float().view(B, T, N, HeadDim) - - # Only norm Q and K heads - Nqk = Nq + Nkv - x_qk = x[:, :, :Nqk, :] # (B, T, Nq+Nkv, HeadDim) - - var = (x_qk ** 2).mean(dim=-1, keepdim=True) - s = torch.rsqrt(var + eps) # (B, T, Nq+Nkv, 1) - - wgt = torch.cat([ - q_wgt.float().unsqueeze(0).expand(Nq, -1), - k_wgt.float().unsqueeze(0).expand(Nkv, -1), - ], dim=0).view(1, 1, Nqk, HeadDim) - - x_qk_normed = ((x_qk * s).to(inp.dtype) * wgt).to(inp.dtype) # (B, T, Nq+Nkv, HeadDim) - - # V heads are passed through unchanged - x_v = x[:, :, Nqk:, :].to(inp.dtype) # (B, T, Nkv, HeadDim) - - # r_rms: pad V slots with 1.0 (or zeros) since they aren't normed - r_rms_qk = s.squeeze(-1).float() # (B, T, Nq+Nkv) - r_rms_v = torch.ones(B, T, Nkv, dtype=torch.float32, device=inp.device) - r_rms = torch.cat([r_rms_qk, r_rms_v], dim=-1) # (B, T, N) - - out = torch.cat([x_qk_normed, x_v], dim=2).view(B, T, N * HeadDim) - return out, r_rms - - -@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ - (1, 4, 2, 1, 16), - (2, 8, 4, 2, 64), - (4, 16, 8, 4, 128), -]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_qk_norm_forward(B, T, Nq, Nkv, HeadDim, dtype): - torch.manual_seed(0) - device = "cuda" - N = Nq + 2 * Nkv - eps = 1e-6 - - inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) - q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) - k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) - out = torch.empty_like(inp) - # NOTE: the kernel does not write rrms for value heads at all; init to one to match reference - r_rms = torch.ones((B, T, N), device=device, dtype=torch.float32) - - K.qk_norm_forward(out, r_rms, inp, q_wgt, k_wgt, eps, Nq, Nkv) - ref_out, ref_r_std = _qk_norm_reference(inp, q_wgt, k_wgt, Nq, Nkv, eps) - - rtol, atol = TOL[dtype] - assert r_rms.cpu().numpy() == pytest.approx(ref_r_std.cpu().numpy(), rel=rtol, abs=atol) - assert out.float().cpu().numpy() == pytest.approx(ref_out.float().cpu().numpy(), rel=rtol, abs=atol) - - -# ---- In-place forward: out is inp ---- -# Exercises the early-return path for V-heads (inp == out) and the normalization path for QK-heads. -@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ - (1, 1, 4, 2, 64), - (2, 4, 2, 1, 128), -]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_qk_norm_forward_inplace(B, T, Nq, Nkv, HeadDim, dtype): - torch.manual_seed(0) - device = "cuda" - N = Nq + 2 * Nkv - eps = 1e-6 - - inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) - q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) - k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) - v_start = Nq + Nkv - v_end = N - - inp_copy = inp.clone() - - out = inp - r_rms = torch.ones((B, T, N), device=device, dtype=torch.float32) - - K.qk_norm_forward(out, r_rms, inp, q_wgt, k_wgt, eps, Nq, Nkv) - - actual_v = out[:, :, v_start * HeadDim : v_end * HeadDim] - orig_v = inp_copy[:, :, v_start * HeadDim : v_end * HeadDim] - assert torch.allclose(actual_v, orig_v, rtol=0, atol=0), "V-heads changed in in-place forward" - - ref_out, _ = _qk_norm_reference(inp_copy, q_wgt, k_wgt, Nq, Nkv, eps) - rtol, atol = TOL[dtype] - assert out.float().cpu().numpy() == pytest.approx(ref_out.float().cpu().numpy(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ - (1, 4, 2, 1, 16), - (2, 8, 4, 2, 64), - (4, 16, 8, 4, 128), -]) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("with_abs_max", [False, True]) -def test_qk_norm_and_rope_forward(B, T, Nq, Nkv, HeadDim, dtype, with_abs_max): - torch.manual_seed(0) - device = "cuda" - N = Nq + 2 * Nkv - eps = 1e-6 - - inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) - q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) - k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) - out = torch.empty_like(inp) - # NOTE: the kernel does not write rrms for value heads; init to 1 to match reference - r_rms = torch.ones((B, T, N), device=device, dtype=torch.float32) - - freq_dtype = torch.float32 if dtype == torch.float32 else torch.float16 - freqs = _make_rope_freqs(T, HeadDim, 1_000_000.0, freq_dtype, device) - - abs_max = torch.zeros((1,), device=device, dtype=torch.float32) if with_abs_max else None - - K.qk_norm_and_rope_forward(out, r_rms, inp, q_wgt, k_wgt, freqs, abs_max, eps, Nq, Nkv) - - # reference: qk-norm first, then rope on the result - ref_normed, ref_r_rms = _qk_norm_reference(inp, q_wgt, k_wgt, Nq, Nkv, eps) - ref_out = _rope_python(ref_normed.view(B, T, N, HeadDim), freqs, Nq, Nkv).view(B, T, N * HeadDim) - - rtol, atol = TOL[dtype] - assert r_rms.cpu().numpy() == pytest.approx(ref_r_rms.cpu().numpy(), rel=rtol, abs=atol) - assert out.float().cpu().numpy() == pytest.approx(ref_out.float().cpu().numpy(), rel=rtol, abs=atol) - - if with_abs_max: - # kernel takes abs-max over Q+K post-rope and V pass-through - ref_abs_max = ref_out.float().abs().max().item() - assert abs_max.item() == pytest.approx(ref_abs_max, rel=rtol, abs=atol) - -@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ - (1, 4, 2, 1, 16), - (2, 8, 4, 2, 64), - (4, 16, 8, 4, 64), - (1, 1, 4, 2, 128), -]) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("in_place", [True, False]) -def test_qk_norm_backward(B, T, Nq, Nkv, HeadDim, dtype, in_place: bool): - torch.manual_seed(0) - device = "cuda" - N = Nq + 2 * Nkv - eps = 1e-6 - - inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) - q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) - k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) - dout = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) - - # Reference in fp32 - inp_f = inp.float().detach().requires_grad_(True) - q_wgt_f = q_wgt.float().detach().requires_grad_(True) - k_wgt_f = k_wgt.float().detach().requires_grad_(True) - ref_out_f, _ = _qk_norm_reference(inp_f, q_wgt_f, k_wgt_f, Nq, Nkv, eps) - ref_out_f.backward(dout.float()) - ref_dinp = inp_f.grad - ref_dq_wgt = q_wgt_f.grad - ref_dk_wgt = k_wgt_f.grad - - # Run forward to get the kernel's rstd. - r_rms = torch.empty((B, T, N), device=device, dtype=torch.float32) - r_rms.fill_(1.0) # V-slot sentinel; kernel only writes QK slots - K.qk_norm_forward(torch.empty_like(inp), r_rms, inp, q_wgt, k_wgt, eps, Nq, Nkv) - - scratch_size = K.get_qknorm_backward_scratch_size(Nq, Nkv, HeadDim, inp.dtype) - scratch = torch.zeros((scratch_size,), device=device, dtype=torch.uint8) - dq_wgt_k = torch.zeros_like(q_wgt) - dk_wgt_k = torch.zeros_like(k_wgt) - - if in_place: - dinp_buf = dout.clone() - K.qk_norm_backward(dinp_buf, dq_wgt_k, dk_wgt_k, scratch, - dinp_buf, inp, q_wgt, k_wgt, r_rms, - None, Nq, Nkv) - dinp = dinp_buf - else: - dinp = torch.zeros_like(inp) - K.qk_norm_backward(dinp, dq_wgt_k, dk_wgt_k, scratch, - dout, inp, q_wgt, k_wgt, r_rms, - None, Nq, Nkv) - - rtol, atol = TOL[dtype] - - assert dinp.float().cpu().numpy() == pytest.approx(ref_dinp.cpu().numpy(), rel=rtol, abs=atol) - assert dq_wgt_k.float().cpu().numpy() == pytest.approx(ref_dq_wgt.cpu().numpy(), rel=rtol, abs=atol) - assert dk_wgt_k.float().cpu().numpy() == pytest.approx(ref_dk_wgt.cpu().numpy(), rel=rtol, abs=atol) - - -# ============================================================ -# fused_residual_rmsnorm_forward -# ============================================================ - -@pytest.mark.parametrize("B,T,C", [(2, 5, 64), (1, 3, 256), (4, 8, 128), (8, 16, 512)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_fused_residual_rmsnorm_forward_reference(B, T, C, dtype): - device = "cuda" - torch.manual_seed(0) - inp1 = torch.randn((B, T, C), device=device, dtype=dtype) - inp2 = torch.randn_like(inp1) - weight = torch.randn((C,), device=device, dtype=dtype) - residual = torch.empty_like(inp1) - normed = torch.empty_like(inp1) - rrms = torch.empty((B, T), device=device, dtype=torch.float32) - eps = 1e-6 - - K.fused_residual_rmsnorm_forward(residual, normed, rrms, inp1, inp2, weight, None, eps) - - res_ref = (inp1.float() + inp2.float()).to(dtype) - var = (res_ref.float() ** 2).mean(dim=-1, keepdim=True) - r_rms = 1.0 / torch.sqrt(var + eps) - norm_ref = (res_ref.float() * r_rms * weight.float()).to(dtype) - - rtol, atol = TOL[dtype] - assert residual.float().cpu() == pytest.approx(res_ref.float().cpu(), rel=rtol, abs=atol) - assert normed.float().cpu() == pytest.approx(norm_ref.float().cpu(), rel=rtol, abs=atol) - assert rrms.cpu() == pytest.approx(r_rms.squeeze(-1).float().cpu(), rel=rtol, abs=atol) - - -# ============================================================ -# vector_add_sr -# ============================================================ - -@pytest.mark.parametrize("nelem", [4096, 16384, 65536]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_vector_add_sr_determinism_and_accuracy(nelem, dtype): - device = "cuda" - a = torch.randn((nelem,), device=device, dtype=dtype) - b = torch.randn_like(a) - out1 = torch.empty_like(a) - out2 = torch.empty_like(a) - seed = 12345 - scale = torch.tensor(0.75, dtype=torch.float32) - K.vector_add_sr(out1, a, b, scale, seed) - K.vector_add_sr(out2, a, b, scale, seed) - - # Determinism: same seed must give bit-identical output. - assert out1.float().cpu() == pytest.approx(out2.float().cpu(), rel=0, abs=0) - - # Accuracy vs fp32 reference; stochastic rounding may introduce ~1 ulp at target dtype. - ref = scale.item() * (a.float() + b.float()) - assert out1.float().cpu() == pytest.approx(ref.cpu(), rel=1e-2, abs=5e-3) - - -# ============================================================ -# quantize_with_abs_max -# ============================================================ - -@pytest.mark.parametrize("N", [1024, 8192, 65536]) -@pytest.mark.parametrize("dtype", [torch.float32]) -def test_quantize_with_abs_max_bf16(N, dtype): - device = "cuda" - x = torch.randn((N,), device=device, dtype=dtype) - abs_max_val = torch.tensor([x.float().abs().max().item()], device=device, dtype=torch.float32) - out = torch.empty((N,), device=device, dtype=torch.bfloat16) - scale = torch.empty((), device=device, dtype=torch.float32) - K.quantize_with_abs_max(out, scale, x, abs_max_val) - assert scale.item() == pytest.approx(1.0, rel=0, abs=0) - assert out.float().cpu() == pytest.approx(x.bfloat16().float().cpu(), rel=0.01) - - -@pytest.mark.parametrize("N", [1024, 8192]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_quantize_with_abs_max_fp8(N, dtype): - device = "cuda" - x = torch.randn((N,), device=device, dtype=dtype) - abs_max_val = torch.tensor([x.float().abs().max().item()], device=device, dtype=torch.float32) - out = torch.empty((N,), device=device, dtype=torch.float8_e4m3fn) - scale = torch.empty((), device=device, dtype=torch.float32) - K.quantize_with_abs_max(out, scale, x, abs_max_val) - assert scale.item() == pytest.approx(abs_max_val.item() / 448.0) - dequant = out.float() * scale - assert dequant.cpu() == pytest.approx(x.float().cpu(), rel=0.1) - - -# ============================================================ -# fused_classifier -# ============================================================ - - -@pytest.mark.parametrize("write_dlogits", [False, True]) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("B,T,V", [(2, 4, 16), (1, 8, 32), (4, 2, 64)]) -def test_fused_classifier_losses(B, T, V, write_dlogits, dtype): - """Per-token losses, lse, and (optionally) dlogits must match reference with no z-regularisation.""" - torch.manual_seed(0) - logits = torch.randn((B, T, V), device="cuda", dtype=dtype) - targets = torch.randint(0, V, (B, T), device="cuda", dtype=torch.int32) - - losses = torch.zeros((B, T), device="cuda", dtype=torch.float32) - lse = torch.empty((B, T), device="cuda", dtype=torch.float32) - dloss = 1.0 - logits_copy = logits.clone() # kernel mutates logits in place - - K.fused_classifier(logits_copy, losses, lse, dloss, targets, 0.0, write_dlogits) - - ref_lse = torch.logsumexp(logits.float(), dim=-1) - assert lse.cpu() == pytest.approx(ref_lse.cpu(), rel=1e-4, abs=1e-5) - - ref_losses = F.cross_entropy(logits.reshape(B * T, V).float(), targets.reshape(B * T).long(), reduction="none").reshape(B, T) - assert losses.cpu() == pytest.approx(ref_losses.cpu(), rel=1e-4, abs=1e-5) - - if write_dlogits: - probs = torch.softmax(logits.float(), dim=-1) - onehot = torch.zeros_like(probs) - onehot.scatter_(-1, targets.long().unsqueeze(-1), 1.0) - ref_dlogits = probs - onehot # dloss=1 everywhere - assert logits_copy.float().cpu() == pytest.approx(ref_dlogits.cpu(), rel=1e-2, abs=1e-3) - - - -# ============================================================ -# adamw_update -# ============================================================ - -def _adamw_reference(params, grads, m, v, lr, beta1, beta2, t, eps, wd): - m_new = beta1 * m.float() + (1 - beta1) * grads.float() - v_new = beta2 * v.float() + (1 - beta2) * grads.float() ** 2 - m_hat = m_new / (1 - beta1 ** t) - v_hat = v_new / (1 - beta2 ** t) - params_new = params.float() * (1 - lr * wd) - lr * m_hat / (v_hat.sqrt() + eps) - return params_new, m_new, v_new - - -@pytest.mark.parametrize("N", [1024, 8192, 65536]) -@pytest.mark.parametrize("p_dtype, g_dtype, m_dtype, v_dtype", [ - (torch.float32, torch.float32, torch.float32, torch.float32), - (torch.bfloat16, torch.bfloat16, torch.float32, torch.float32), - (torch.bfloat16, torch.bfloat16, torch.bfloat16, torch.float32), - (torch.bfloat16, torch.bfloat16, torch.bfloat16, torch.bfloat16), -]) -def test_adamw_update_matches_reference(N, p_dtype, g_dtype, m_dtype, v_dtype): - torch.manual_seed(0) - params = torch.randn((N,), device="cuda", dtype=p_dtype) - grads = torch.randn((N,), device="cuda", dtype=g_dtype) - m = torch.randn((N,), device="cuda", dtype=m_dtype) * 0.1 - v = torch.rand((N,), device="cuda", dtype=v_dtype) * 0.01 + 1e-8 - g_scale = torch.rand((), device="cuda", dtype=torch.float32) * 0.5 + 0.5 - - lr, beta1, beta2, t, eps, wd = 1e-3, 0.9, 0.999, 1, 1e-8, 0.1 - ref_params, ref_m, ref_v = _adamw_reference( - params.clone().float(), grads.float() * g_scale, m.clone().float(), v.clone().float(), lr, beta1, beta2, t, eps, wd - ) - - K.adamw_update(params, grads, m, v, lr, beta1, beta2, t, eps, wd, g_scale) - - rtol, atol = TOL[p_dtype] - assert params.float().cpu() == pytest.approx(ref_params.cpu(), rel=rtol, abs=atol) - rtol, atol = TOL[m_dtype] - assert m.float().cpu() == pytest.approx(ref_m.cpu(), rel=rtol, abs=atol) - rtol, atol = TOL[v_dtype] - assert v.float().cpu() == pytest.approx(ref_v.cpu(), rel=rtol, abs=atol) - - -# ============================================================ -# vector_reduce_sr -# ============================================================ -@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192), (8, 16384)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_vector_reduce_sr_sum_matches_reference(n_shards, nelem, dtype): - """With scale=1, the reduction must equal the sum across shards.""" - torch.manual_seed(0) - src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) - dest = torch.empty((nelem,), device="cuda", dtype=dtype) - scale = torch.tensor(1.0, dtype=torch.float32) - - K.vector_reduce_sr(dest, src, scale, n_shards, seed=0) - - ref = src.view(n_shards, nelem).float().sum(dim=0) - rtol, atol = TOL[dtype] - assert dest.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) - - - -@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_vector_reduce_sr_determinism(n_shards, nelem, dtype): - src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) - dest1 = torch.empty((nelem,), device="cuda", dtype=dtype) - dest2 = torch.empty((nelem,), device="cuda", dtype=dtype) - scale = torch.tensor(0.5, dtype=torch.float32) - - K.vector_reduce_sr(dest1, src, scale, n_shards, seed=99) - K.vector_reduce_sr(dest2, src, scale, n_shards, seed=99) - - assert dest1.float().cpu() == pytest.approx(dest2.float().cpu(), rel=0, abs=0) - - -@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_vector_reduce_sr_scale_zero(n_shards, nelem, dtype): - src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) - dest = torch.empty((nelem,), device="cuda", dtype=dtype) - scale = torch.tensor(0.0, dtype=torch.float32) - - K.vector_reduce_sr(dest, src, scale, n_shards, seed=0) - - assert dest.float().cpu() == pytest.approx(torch.zeros(nelem).cpu(), rel=0, abs=0) - - -@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_vector_reduce_sr_skip(n_shards, nelem, dtype): - """Skipped shard should not contribute to the result.""" - torch.manual_seed(42) - src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) - dest = torch.empty((nelem,), device="cuda", dtype=dtype) - scale = torch.tensor(1.0, dtype=torch.float32) - - skip = 1 - K.vector_reduce_sr(dest, src, scale, n_shards, skip=skip, seed=0) - - shards = src.view(n_shards, nelem).float() - ref = sum(shards[k] for k in range(n_shards) if k != skip) - rtol, atol = TOL[dtype] - assert dest.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_vector_reduce_sr_accumulate(n_shards, nelem, dtype): - """With accumulate=True, dest's initial values should be included in the sum.""" - torch.manual_seed(7) - src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) - dest = torch.randn((nelem,), device="cuda", dtype=dtype) - dest_initial = dest.clone() - scale = torch.tensor(1.0, dtype=torch.float32) - - K.vector_reduce_sr(dest, src, scale, n_shards, accumulate=True, seed=0) - - ref = src.view(n_shards, nelem).float().sum(dim=0) + dest_initial.float() - rtol, atol = TOL[dtype] - assert dest.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("dtype", DTYPES) -def test_vector_reduce_sr_stochastic_rounding_unbiased(dtype: torch.dtype, nelem=4096): - """SR rounding should be unbiased: mean of many rounded values ≈ true mean.""" - n_shards = 2 - # Use values that are exactly halfway between representable values to maximise rounding effect - src = torch.ones((n_shards * nelem,), device="cuda", dtype=dtype) * 0.5 - results = [] - for seed in range(20): - dest = torch.empty((nelem,), device="cuda", dtype=dtype) - scale = torch.tensor(1.0, dtype=torch.float32) - K.vector_reduce_sr(dest, src, scale, n_shards, seed=seed) - results.append(dest.float().cpu()) - - mean_result = torch.stack(results).mean(dim=0) - # True answer is 1.0 (sum of two 0.5 shards); mean over seeds should be close - assert mean_result == pytest.approx(torch.ones(nelem).cpu(), rel=0.01, abs=0.01) - - -# ============================================================ -# backward_bias -# ============================================================ -@pytest.mark.parametrize("B,T,OC", [(2, 5, 64), (1, 3, 256), (4, 8, 128), (8, 16, 512)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_backward_bias(B, T, OC, dtype): - device = "cuda" - torch.manual_seed(0) - - dout = torch.randn((B, T, OC), device=device, dtype=dtype) - dbias = torch.zeros((OC,), device=device, dtype=dtype) - dbias_buffer = torch.empty(K.get_bias_backward_scratch_size(dtype, OC) // 4, device=device, dtype=torch.float32) - - K.backward_bias(dbias, dout, None, None, dbias_buffer) - - ref = dout.float().sum(dim=(0, 1)).to(dtype) - - rtol, atol = TOL[dtype] - assert dbias.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("B,T,OC", [(2, 5, 64), (1, 3, 256), (4, 8, 128), (8, 16, 512)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_backward_bias_with_scale(B, T, OC, dtype): - device = "cuda" - torch.manual_seed(0) - - dout = torch.randn((B, T, OC), device=device, dtype=dtype) - dbias = torch.zeros((OC,), device=device, dtype=dtype) - dbias_buffer = torch.zeros(K.get_bias_backward_scratch_size(dtype, OC) // 4, device=device, dtype=torch.float32) - scale_a = torch.tensor(0.25, device=device, dtype=torch.float32) - scale_b = torch.tensor(2.0, device=device, dtype=torch.float32) - - scale = scale_a.item() * scale_b.item() - ref = (dout.float().sum(dim=(0, 1)) * scale).to(dtype) - K.backward_bias(dbias, dout, scale_a, scale_b, dbias_buffer) - rtol, atol = TOL[dtype] - assert dbias.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("B,T,OC", [(2, 5, 64), (1, 3, 256), (4, 8, 128), (8, 16, 512)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_backward_bias_accumulates(B, T, OC, dtype): - """Verify that backward_bias adds into dbias rather than overwriting it.""" - device = "cuda" - torch.manual_seed(0) - - dout = torch.randn((B, T, OC), device=device, dtype=dtype) - initial = torch.randn((OC,), device=device, dtype=dtype) - dbias = initial.clone() - dbias_buffer = torch.zeros(K.get_bias_backward_scratch_size(dtype, OC) // 4, device=device, dtype=torch.float32) - - K.backward_bias(dbias, dout, None, None, dbias_buffer) - - ref = (initial.float() + dout.float().sum(dim=(0, 1))).to(dtype) - - rtol, atol = TOL[dtype] - assert dbias.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) - -# ============================================================ -# matmul -# ============================================================ - -@pytest.fixture(scope="module") -def cublas_handle(): - handle = K.create_cublas_handle() - try: - yield handle - finally: - K.destroy_cublas_handle(handle) - -@pytest.fixture(scope="module") -def workspace(): - # 32MB workspace, typical for cublasLt - return torch.empty(32 * 1024 * 1024, device="cuda", dtype=torch.uint8) - - -# EMMTranspose: TT=0, TN=1, NT=2, NN=3 -MODES = { - "NN": 3, # C = A @ B - "NT": 2, # C = A @ B.T - "TN": 1, # C = A.T @ B - "TT": 0, # C = A.T @ B.T -} - - -def ref_matmul(a: torch.Tensor, b: torch.Tensor, bias, mode_str: str, accumulate: bool, c_ref: torch.Tensor) -> torch.Tensor: - a_f = a.float() - b_f = b.float() - if mode_str == "NN": - out = a_f @ b_f - elif mode_str == "NT": - out = a_f @ b_f.transpose(-1, -2) - elif mode_str == "TN": - out = a_f.transpose(-1, -2) @ b_f - elif mode_str == "TT": - out = a_f.transpose(-1, -2) @ b_f.transpose(-1, -2) - if bias is not None: - out = out + bias.float() - if accumulate: - out = out + c_ref.float() - return out - - -@pytest.mark.parametrize("M,K_dim,N", [ - (1, 64, 64), - (10, 64, 128), - (3, 256, 256), - (32, 128, 64), -]) -@pytest.mark.parametrize("mode_str", list(MODES.keys())) -@pytest.mark.parametrize("dtype", DTYPES) -def test_matmul_basic(M, K_dim, N, mode_str, dtype, cublas_handle, workspace): - torch.manual_seed(0) - device = "cuda" - mode = MODES[mode_str] - rtol, atol = TOL[dtype] - - if mode_str in ("NN", "NT"): - a = torch.randn(M, K_dim, device=device, dtype=dtype) - else: # TN, TT: A is [K, M] - a = torch.randn(K_dim, M, device=device, dtype=dtype) - - if mode_str in ("NN", "TN"): - b = torch.randn(K_dim, N, device=device, dtype=dtype) - else: # NT, TT: B is [N, K] - b = torch.randn(N, K_dim, device=device, dtype=dtype) - - c = torch.empty(M, N, device=device, dtype=dtype) - K.matmul(c, a, b, None, None, None, cublas_handle, workspace, mode, False) - - ref = ref_matmul(a, b, None, mode_str, False, c) - assert c.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("M,K_dim,N", [(32, 64, 128), (64, 128, 64)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_matmul_bias(M, K_dim, N, dtype, cublas_handle, workspace): - torch.manual_seed(0) - device = "cuda" - rtol, atol = TOL[dtype] - - a = torch.randn(M, K_dim, device=device, dtype=dtype) - b = torch.randn(K_dim, N, device=device, dtype=dtype) - bias = torch.randn(N, device=device, dtype=dtype) - c = torch.zeros(M, N, device=device, dtype=dtype) - - K.matmul(c, a, b, bias, None, None, cublas_handle, workspace, MODES["NN"], False) - - ref = ref_matmul(a, b, bias, "NN", False, c) - assert c.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("M,K_dim,N", [(32, 64, 128), (64, 128, 64)]) -@pytest.mark.parametrize("mode_str", list(MODES.keys())) -@pytest.mark.parametrize("dtype", DTYPES) -def test_matmul_accumulate(M, K_dim, N, mode_str, dtype, cublas_handle, workspace): - torch.manual_seed(0) - device = "cuda" - mode = MODES[mode_str] - rtol, atol = TOL[dtype] - - if mode_str in ("NN", "NT"): - a = torch.randn(M, K_dim, device=device, dtype=dtype) - else: - a = torch.randn(K_dim, M, device=device, dtype=dtype) - - if mode_str in ("NN", "TN"): - b = torch.randn(K_dim, N, device=device, dtype=dtype) - else: - b = torch.randn(N, K_dim, device=device, dtype=dtype) - - c = torch.randn(M, N, device=device, dtype=dtype) - c_ref = c.clone() - - K.matmul(c, a, b, None, None, None, cublas_handle, workspace, mode, True) - - ref = ref_matmul(a, b, None, mode_str, True, c_ref) - assert c.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) - - -@pytest.mark.parametrize("M,K_dim,N", [(32, 64, 128)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_matmul_bias_and_accumulate(M, K_dim, N, dtype, cublas_handle, workspace): - torch.manual_seed(0) - device = "cuda" - rtol, atol = TOL[dtype] - if dtype == torch.bfloat16: - rtol = 1.5e-2 - atol = 6e-2 - - a = torch.randn(M, K_dim, device=device, dtype=dtype) - b = torch.randn(K_dim, N, device=device, dtype=dtype) - bias = torch.randn(N, device=device, dtype=dtype) - c = torch.randn(M, N, device=device, dtype=dtype) - c_ref = c.clone() - - K.matmul(c, a, b, bias, None, None, cublas_handle, workspace, MODES["NN"], True) - - ref = ref_matmul(a, b, bias, "NN", True, c_ref) - assert c.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) - -# ============================================================ -# encoder_forward / encoder_backward -# ============================================================ - -def _encoder_forward_reference(inp: torch.Tensor, wte: torch.Tensor, wpe: torch.Tensor | None) -> torch.Tensor: - """Token embedding + optional positional embedding lookup.""" - out = wte[inp] # (B, T, C) - if wpe is not None: - T = inp.shape[1] - out = out + wpe[:T] - return out - - -@pytest.mark.parametrize("B,T,V,C", [(2, 8, 64, 32), (1, 16, 128, 64), (4, 4, 32, 16)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_encoder_forward_with_wpe(B, T, V, C, dtype): - """Token + position embedding lookup must match reference, with and without wpe.""" - torch.manual_seed(0) - device = "cuda" - - inp = torch.randint(0, V, (B, T), device=device, dtype=torch.int32) - wte = torch.randn((V, C), device=device, dtype=dtype) - wpe = torch.randn((T, C), device=device, dtype=dtype) - out = torch.empty((B, T, C), device=device, dtype=dtype) - - K.encoder_forward(out, inp, wte, wpe) - - ref = _encoder_forward_reference(inp, wte, wpe) - # Embedding lookup is a pure gather — no arithmetic error possible. - assert out.float().cpu() == pytest.approx(ref.float().cpu(), rel=0, abs=0) - - -@pytest.mark.parametrize("B,T,V,C", [(2, 8, 64, 32), (3, 12, 50, 48)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_encoder_forward_no_wpe(B, T, V, C, dtype): - """With wpe=None the output must equal wte[inp] exactly — no position bias added.""" - torch.manual_seed(1) - device = "cuda" - - inp = torch.randint(0, V, (B, T), device=device, dtype=torch.int32) - wte = torch.randn((V, C), device=device, dtype=dtype) - out = torch.empty((B, T, C), device=device, dtype=dtype) - - K.encoder_forward(out, inp, wte, None) - - ref = wte[inp] # shape (B, T, C) - assert out.float().cpu() == pytest.approx(ref.float().cpu(), rel=0, abs=0) - - -def _make_encoder_backward_buffers(V: int, C: int, B: int, T: int, device: str, dtype: torch.dtype): - """Allocate the auxiliary buffers required by encoder_backward.""" - dwte = torch.zeros((V, C), device=device, dtype=dtype) - cg_max = int(math.ceil(C / 32)) - scratch = torch.zeros((B, T, 5*cg_max), device=device, dtype=torch.int32) - workload_indices = torch.zeros((B, T, cg_max), device="cpu", dtype=torch.int32) - bucket_info = torch.zeros((B, T, 4*cg_max), device="cpu", dtype=torch.int32) - return dwte, scratch, workload_indices, bucket_info - - -@pytest.mark.parametrize("B,T,V,C", [(2, 8, 16, 32), (1, 4, 8, 16), (3, 6, 32, 64)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_encoder_backward_gradient_accumulation(B, T, V, C, dtype): - """dwte[token] must accumulate dout contributions from every position that used that token.""" - torch.manual_seed(0) - device = "cuda" - event = torch.cuda.Event() - event.record() # pytorch events are lazy. - - # Repeat token index 0 at every position to stress accumulation. - inp = torch.zeros((B, T), device=device, dtype=torch.int32) - inp_cpu = inp.cpu() - dout = torch.randn((B, T, C), device=device, dtype=dtype) - dwte, scratch, workload_indices, bucket_info = _make_encoder_backward_buffers(V, C, B, T, device, dtype) - - K.encoder_backward(dwte, scratch, workload_indices, bucket_info, dout, inp, inp_cpu, sync_event=event.cuda_event, seed=0) - - # Reference: scatter-add over all (b, t) pairs. - ref_dwte = torch.zeros((V, C), dtype=torch.float32) - for b in range(B): - for t in range(T): - ref_dwte[inp_cpu[b, t]] += dout[b, t].float().cpu() - - rtol, atol = TOL[dtype] - assert dwte[0].float().cpu() == pytest.approx(ref_dwte[0].cpu(), rel=rtol, abs=atol) - # Tokens that were never used must stay zero. - assert dwte[1:].abs().max().item() == 0.0 - - -@pytest.mark.parametrize("B,T,V,C", [(2, 8, 16, 32), (4, 4, 32, 64)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_encoder_backward_seed_determinism(B, T, V, C, dtype): - """Same seed must produce bit-identical dwte results.""" - torch.manual_seed(42) - device = "cuda" - event = torch.cuda.Event() - event.record() # pytorch events are lazy. - - inp = torch.randint(0, V, (B, T), device=device, dtype=torch.int32) - inp_cpu = inp.cpu() - dout = torch.randn((B, T, C), device=device, dtype=dtype) - - seed = 0xDEADBEEF - - dwte1, scratch1, wi1, bi1 = _make_encoder_backward_buffers(V, C, B, T, device, dtype) - K.encoder_backward(dwte1, scratch1, wi1, bi1, dout, inp, inp_cpu, seed=seed, sync_event=event.cuda_event) - - dwte2, scratch2, wi2, bi2 = _make_encoder_backward_buffers(V, C, B, T, device, dtype) - K.encoder_backward(dwte2, scratch2, wi2, bi2, dout, inp, inp_cpu, seed=seed, sync_event=event.cuda_event) - - # Bit-exact: same seed must yield identical output. - assert dwte1.float().cpu() == pytest.approx(dwte2.float().cpu(), rel=0, abs=0) - - -# ============================================================ -# quantize_and_transpose_with_abs_max -# ============================================================ - -@pytest.mark.parametrize("rows,cols", [(32, 32), (128, 256), (512, 1024)]) -@pytest.mark.parametrize("dtype", [torch.float32]) -def test_quantize_and_transpose_with_abs_max_bf16(rows, cols, dtype): - device = "cuda" - x = torch.randn((rows, cols), device=device, dtype=dtype) - abs_max_val = torch.tensor([x.float().abs().max().item()], device=device, dtype=torch.float32) - out = torch.empty((cols, rows), device=device, dtype=torch.bfloat16) - scale = torch.empty((), device=device, dtype=torch.float32) - K.quantize_and_transpose_with_abs_max(out, scale, x, abs_max_val) - - # no assert for scale, as scale is unused for bf16 - - expected = x.bfloat16().T.contiguous() - assert out.float().cpu() == pytest.approx(expected.float().cpu(), rel=0.01) - - -@pytest.mark.parametrize("rows,cols", [(32, 32), (128, 256), (512, 1024)]) -@pytest.mark.parametrize("dtype", DTYPES) -def test_quantize_and_transpose_with_abs_max_fp8(rows, cols, dtype): - device = "cuda" - x = torch.randn((rows, cols), device=device, dtype=dtype) - abs_max_val = torch.tensor([x.float().abs().max().item()], device=device, dtype=torch.float32) - out = torch.empty((cols, rows), device=device, dtype=torch.float8_e4m3fn) - scale = torch.empty((), device=device, dtype=torch.float32) - K.quantize_and_transpose_with_abs_max(out, scale, x, abs_max_val) - - assert scale.item() == pytest.approx(abs_max_val.item() / 448.0) - - # Dequantize and verify values match input (with fp8 tolerance) - dequant = out.float() * scale - - # Verify transpose: reshape input as [rows, cols], transpose to [cols, rows] - expected = x.float().T.contiguous() - assert dequant.cpu() == pytest.approx(expected.cpu(), rel=0.1, abs=1e-4) diff --git a/test/conftest.py b/test/conftest.py new file mode 100644 index 0000000..d08cc6a --- /dev/null +++ b/test/conftest.py @@ -0,0 +1,10 @@ +import torch + + +DTYPES = [torch.float32, torch.bfloat16] + +# Tolerances by dtype for approximate kernels: (rtol, atol) +TOL = { + torch.float32: (5e-4, 5e-5), + torch.bfloat16: (1e-2, 5e-3), +} diff --git a/test/test_kernels_encoder.py b/test/test_kernels_encoder.py new file mode 100644 index 0000000..16dcf30 --- /dev/null +++ b/test/test_kernels_encoder.py @@ -0,0 +1,114 @@ +import math + +import pytest +import torch +from pyllmq import kernels as K +from conftest import DTYPES, TOL + +def _encoder_forward_reference(inp: torch.Tensor, wte: torch.Tensor, wpe: torch.Tensor | None) -> torch.Tensor: + """Token embedding + optional positional embedding lookup.""" + out = wte[inp] # (B, T, C) + if wpe is not None: + T = inp.shape[1] + out = out + wpe[:T] + return out + + +@pytest.mark.parametrize("B,T,V,C", [(2, 8, 64, 32), (1, 16, 128, 64), (4, 4, 32, 16)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_encoder_forward_with_wpe(B, T, V, C, dtype): + """Token + position embedding lookup must match reference, with and without wpe.""" + torch.manual_seed(0) + device = "cuda" + + inp = torch.randint(0, V, (B, T), device=device, dtype=torch.int32) + wte = torch.randn((V, C), device=device, dtype=dtype) + wpe = torch.randn((T, C), device=device, dtype=dtype) + out = torch.empty((B, T, C), device=device, dtype=dtype) + + K.encoder_forward(out, inp, wte, wpe) + + ref = _encoder_forward_reference(inp, wte, wpe) + # Embedding lookup is a pure gather — no arithmetic error possible. + assert out.float().cpu() == pytest.approx(ref.float().cpu(), rel=0, abs=0) + + +@pytest.mark.parametrize("B,T,V,C", [(2, 8, 64, 32), (3, 12, 50, 48)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_encoder_forward_no_wpe(B, T, V, C, dtype): + """With wpe=None the output must equal wte[inp] exactly — no position bias added.""" + torch.manual_seed(1) + device = "cuda" + + inp = torch.randint(0, V, (B, T), device=device, dtype=torch.int32) + wte = torch.randn((V, C), device=device, dtype=dtype) + out = torch.empty((B, T, C), device=device, dtype=dtype) + + K.encoder_forward(out, inp, wte, None) + + ref = wte[inp] # shape (B, T, C) + assert out.float().cpu() == pytest.approx(ref.float().cpu(), rel=0, abs=0) + + +def _make_encoder_backward_buffers(V: int, C: int, B: int, T: int, device: str, dtype: torch.dtype): + """Allocate the auxiliary buffers required by encoder_backward.""" + dwte = torch.zeros((V, C), device=device, dtype=dtype) + cg_max = int(math.ceil(C / 32)) + scratch = torch.zeros((B, T, 5*cg_max), device=device, dtype=torch.int32) + workload_indices = torch.zeros((B, T, cg_max), device="cpu", dtype=torch.int32) + bucket_info = torch.zeros((B, T, 4*cg_max), device="cpu", dtype=torch.int32) + return dwte, scratch, workload_indices, bucket_info + + +@pytest.mark.parametrize("B,T,V,C", [(2, 8, 16, 32), (1, 4, 8, 16), (3, 6, 32, 64)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_encoder_backward_gradient_accumulation(B, T, V, C, dtype): + """dwte[token] must accumulate dout contributions from every position that used that token.""" + torch.manual_seed(0) + device = "cuda" + event = torch.cuda.Event() + event.record() # pytorch events are lazy. + + # Repeat token index 0 at every position to stress accumulation. + inp = torch.zeros((B, T), device=device, dtype=torch.int32) + inp_cpu = inp.cpu() + dout = torch.randn((B, T, C), device=device, dtype=dtype) + dwte, scratch, workload_indices, bucket_info = _make_encoder_backward_buffers(V, C, B, T, device, dtype) + + K.encoder_backward(dwte, scratch, workload_indices, bucket_info, dout, inp, inp_cpu, sync_event=event.cuda_event, seed=0) + + # Reference: scatter-add over all (b, t) pairs. + ref_dwte = torch.zeros((V, C), dtype=torch.float32) + for b in range(B): + for t in range(T): + ref_dwte[inp_cpu[b, t]] += dout[b, t].float().cpu() + + rtol, atol = TOL[dtype] + assert dwte[0].float().cpu() == pytest.approx(ref_dwte[0].cpu(), rel=rtol, abs=atol) + # Tokens that were never used must stay zero. + assert dwte[1:].abs().max().item() == 0.0 + + +@pytest.mark.parametrize("B,T,V,C", [(2, 8, 16, 32), (4, 4, 32, 64)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_encoder_backward_seed_determinism(B, T, V, C, dtype): + """Same seed must produce bit-identical dwte results.""" + torch.manual_seed(42) + device = "cuda" + event = torch.cuda.Event() + event.record() # pytorch events are lazy. + + inp = torch.randint(0, V, (B, T), device=device, dtype=torch.int32) + inp_cpu = inp.cpu() + dout = torch.randn((B, T, C), device=device, dtype=dtype) + + seed = 0xDEADBEEF + + dwte1, scratch1, wi1, bi1 = _make_encoder_backward_buffers(V, C, B, T, device, dtype) + K.encoder_backward(dwte1, scratch1, wi1, bi1, dout, inp, inp_cpu, seed=seed, sync_event=event.cuda_event) + + dwte2, scratch2, wi2, bi2 = _make_encoder_backward_buffers(V, C, B, T, device, dtype) + K.encoder_backward(dwte2, scratch2, wi2, bi2, dout, inp, inp_cpu, seed=seed, sync_event=event.cuda_event) + + # Bit-exact: same seed must yield identical output. + assert dwte1.float().cpu() == pytest.approx(dwte2.float().cpu(), rel=0, abs=0) diff --git a/test/test_kernels_matmul.py b/test/test_kernels_matmul.py new file mode 100644 index 0000000..8e83f73 --- /dev/null +++ b/test/test_kernels_matmul.py @@ -0,0 +1,209 @@ +import pytest +import torch +from pyllmq import kernels as K +from conftest import DTYPES, TOL + + +@pytest.fixture(scope="module") +def cublas_handle(): + handle = K.create_cublas_handle() + try: + yield handle + finally: + K.destroy_cublas_handle(handle) + +@pytest.fixture(scope="module") +def workspace(): + # 32MB workspace, typical for cublasLt + return torch.empty(32 * 1024 * 1024, device="cuda", dtype=torch.uint8) + + +# EMMTranspose: TT=0, TN=1, NT=2, NN=3 +MODES = { + "NN": 3, # C = A @ B + "NT": 2, # C = A @ B.T + "TN": 1, # C = A.T @ B + "TT": 0, # C = A.T @ B.T +} + + +# ============================================================ +# matmul +# ============================================================ + +def ref_matmul(a: torch.Tensor, b: torch.Tensor, bias, mode_str: str, accumulate: bool, c_ref: torch.Tensor) -> torch.Tensor: + a_f = a.float() + b_f = b.float() + if mode_str == "NN": + out = a_f @ b_f + elif mode_str == "NT": + out = a_f @ b_f.transpose(-1, -2) + elif mode_str == "TN": + out = a_f.transpose(-1, -2) @ b_f + elif mode_str == "TT": + out = a_f.transpose(-1, -2) @ b_f.transpose(-1, -2) + if bias is not None: + out = out + bias.float() + if accumulate: + out = out + c_ref.float() + return out + + +@pytest.mark.parametrize("M,K_dim,N", [ + (1, 64, 64), + (10, 64, 128), + (3, 256, 256), + (32, 128, 64), +]) +@pytest.mark.parametrize("mode_str", list(MODES.keys())) +@pytest.mark.parametrize("dtype", DTYPES) +def test_matmul_basic(M, K_dim, N, mode_str, dtype, cublas_handle, workspace): + torch.manual_seed(0) + device = "cuda" + mode = MODES[mode_str] + rtol, atol = TOL[dtype] + + if mode_str in ("NN", "NT"): + a = torch.randn(M, K_dim, device=device, dtype=dtype) + else: # TN, TT: A is [K, M] + a = torch.randn(K_dim, M, device=device, dtype=dtype) + + if mode_str in ("NN", "TN"): + b = torch.randn(K_dim, N, device=device, dtype=dtype) + else: # NT, TT: B is [N, K] + b = torch.randn(N, K_dim, device=device, dtype=dtype) + + c = torch.empty(M, N, device=device, dtype=dtype) + K.matmul(c, a, b, None, None, None, cublas_handle, workspace, mode, False) + + ref = ref_matmul(a, b, None, mode_str, False, c) + assert c.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("M,K_dim,N", [(32, 64, 128), (64, 128, 64)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_matmul_bias(M, K_dim, N, dtype, cublas_handle, workspace): + torch.manual_seed(0) + device = "cuda" + rtol, atol = TOL[dtype] + + a = torch.randn(M, K_dim, device=device, dtype=dtype) + b = torch.randn(K_dim, N, device=device, dtype=dtype) + bias = torch.randn(N, device=device, dtype=dtype) + c = torch.zeros(M, N, device=device, dtype=dtype) + + K.matmul(c, a, b, bias, None, None, cublas_handle, workspace, MODES["NN"], False) + + ref = ref_matmul(a, b, bias, "NN", False, c) + assert c.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("M,K_dim,N", [(32, 64, 128), (64, 128, 64)]) +@pytest.mark.parametrize("mode_str", list(MODES.keys())) +@pytest.mark.parametrize("dtype", DTYPES) +def test_matmul_accumulate(M, K_dim, N, mode_str, dtype, cublas_handle, workspace): + torch.manual_seed(0) + device = "cuda" + mode = MODES[mode_str] + rtol, atol = TOL[dtype] + + if mode_str in ("NN", "NT"): + a = torch.randn(M, K_dim, device=device, dtype=dtype) + else: + a = torch.randn(K_dim, M, device=device, dtype=dtype) + + if mode_str in ("NN", "TN"): + b = torch.randn(K_dim, N, device=device, dtype=dtype) + else: + b = torch.randn(N, K_dim, device=device, dtype=dtype) + + c = torch.randn(M, N, device=device, dtype=dtype) + c_ref = c.clone() + + K.matmul(c, a, b, None, None, None, cublas_handle, workspace, mode, True) + + ref = ref_matmul(a, b, None, mode_str, True, c_ref) + assert c.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("M,K_dim,N", [(32, 64, 128)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_matmul_bias_and_accumulate(M, K_dim, N, dtype, cublas_handle, workspace): + torch.manual_seed(0) + device = "cuda" + rtol, atol = TOL[dtype] + if dtype == torch.bfloat16: + rtol = 1.5e-2 + atol = 6e-2 + + a = torch.randn(M, K_dim, device=device, dtype=dtype) + b = torch.randn(K_dim, N, device=device, dtype=dtype) + bias = torch.randn(N, device=device, dtype=dtype) + c = torch.randn(M, N, device=device, dtype=dtype) + c_ref = c.clone() + + K.matmul(c, a, b, bias, None, None, cublas_handle, workspace, MODES["NN"], True) + + ref = ref_matmul(a, b, bias, "NN", True, c_ref) + assert c.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) + + + +# ============================================================ +# backward_bias +# ============================================================ +@pytest.mark.parametrize("B,T,OC", [(2, 5, 64), (1, 3, 256), (4, 8, 128), (8, 16, 512)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_backward_bias(B, T, OC, dtype): + device = "cuda" + torch.manual_seed(0) + + dout = torch.randn((B, T, OC), device=device, dtype=dtype) + dbias = torch.zeros((OC,), device=device, dtype=dtype) + dbias_buffer = torch.empty(K.get_bias_backward_scratch_size(dtype, OC) // 4, device=device, dtype=torch.float32) + + K.backward_bias(dbias, dout, None, None, dbias_buffer) + + ref = dout.float().sum(dim=(0, 1)).to(dtype) + + rtol, atol = TOL[dtype] + assert dbias.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("B,T,OC", [(2, 5, 64), (1, 3, 256), (4, 8, 128), (8, 16, 512)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_backward_bias_with_scale(B, T, OC, dtype): + device = "cuda" + torch.manual_seed(0) + + dout = torch.randn((B, T, OC), device=device, dtype=dtype) + dbias = torch.zeros((OC,), device=device, dtype=dtype) + dbias_buffer = torch.zeros(K.get_bias_backward_scratch_size(dtype, OC) // 4, device=device, dtype=torch.float32) + scale_a = torch.tensor(0.25, device=device, dtype=torch.float32) + scale_b = torch.tensor(2.0, device=device, dtype=torch.float32) + + scale = scale_a.item() * scale_b.item() + ref = (dout.float().sum(dim=(0, 1)) * scale).to(dtype) + K.backward_bias(dbias, dout, scale_a, scale_b, dbias_buffer) + rtol, atol = TOL[dtype] + assert dbias.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("B,T,OC", [(2, 5, 64), (1, 3, 256), (4, 8, 128), (8, 16, 512)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_backward_bias_accumulates(B, T, OC, dtype): + """Verify that backward_bias adds into dbias rather than overwriting it.""" + device = "cuda" + torch.manual_seed(0) + + dout = torch.randn((B, T, OC), device=device, dtype=dtype) + initial = torch.randn((OC,), device=device, dtype=dtype) + dbias = initial.clone() + dbias_buffer = torch.zeros(K.get_bias_backward_scratch_size(dtype, OC) // 4, device=device, dtype=torch.float32) + + K.backward_bias(dbias, dout, None, None, dbias_buffer) + + ref = (initial.float() + dout.float().sum(dim=(0, 1))).to(dtype) + + rtol, atol = TOL[dtype] + assert dbias.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) diff --git a/test/test_kernels_other.py b/test/test_kernels_other.py new file mode 100644 index 0000000..736515a --- /dev/null +++ b/test/test_kernels_other.py @@ -0,0 +1,218 @@ +import pytest +import torch +import torch.nn.functional as F +from pyllmq import kernels as K +from conftest import DTYPES, TOL + + +# ============================================================ +# fill_constant +# Fills every element with a compile-time constant: result must +# be bit-exact, so rel=0, abs=0. +# ============================================================ + +@pytest.mark.parametrize("shape", [(8, 16), (1,), (128, 256), (4, 8, 32)]) +@pytest.mark.parametrize("value", [0.0, 1.0, 3.25, -2.5]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_fill_constant_basic(shape, value, dtype): + x = torch.empty(shape, device="cuda", dtype=dtype) + K.fill_constant(x, value) + ref = torch.full(shape, value, dtype=dtype) + # Exact: every element is independently written to the same constant. + assert x.float().cpu() == pytest.approx(ref.float().cpu(), rel=0, abs=0) + + +# ============================================================ +# transpose +# Byte-level shuffle — no arithmetic, so result must be exact. +# ============================================================ + +@pytest.mark.parametrize("rows,cols", [(7, 11), (1024, 2048), (64, 128), (3, 512)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_transpose_matches_torch(rows, cols, dtype): + src = torch.randn((rows, cols), device="cuda", dtype=dtype) + dst = torch.empty((cols, rows), device="cuda", dtype=dtype) + K.transpose(dst, src) + # Transposing rearranges values without touching bits — must be exact. + assert dst.float().cpu() == pytest.approx(src.t().float().cpu(), rel=0, abs=0) + + +# ============================================================ +# global_norm_squared +# Involves floating-point accumulation so tolerances are needed. +# ============================================================ + +@pytest.mark.parametrize("shape", [(128, 24524), (256, 1024), (1, 4096)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_global_norm_squared(shape, dtype): + x = torch.randn(*shape, device="cuda", dtype=dtype) + out = torch.zeros((256,), device="cuda", dtype=torch.float32) + K.global_norm_squared(out, x) + ref = (x.float() ** 2).sum() + rtol = 1e-5 if dtype == torch.float32 else 1e-2 + assert out.sum().cpu() == pytest.approx(ref.cpu(), rel=rtol) + + +# ============================================================ +# grouped_loss_sum +# ============================================================ + +@pytest.mark.parametrize("B,T", [(1, 512), (2, 1024), (7, 2048), (4, 512)]) +def test_grouped_loss_sum_basic(B, T): + losses = torch.rand((B, T), device="cuda", dtype=torch.float32) + out = torch.empty((T // 512,), device="cuda", dtype=torch.float32) + K.grouped_loss_sum(out, losses) + ref = losses.reshape(B, -1, 512).sum(dim=2).sum(dim=0) + assert out.cpu() == pytest.approx(ref.cpu(), rel=1e-5, abs=1e-6) + + +# ============================================================ +# fill_normal +# ============================================================ + +@pytest.mark.parametrize("mean,std", [(0.0, 1.0), (1.5, 0.25), (-1.0, 2.0)]) +def test_fill_normal_stats(mean, std): + N = 256 * 1024 + seed = 0xABCDEF01 + x = torch.empty((N,), device="cuda", dtype=torch.float32) + K.fill_normal(x, float(mean), float(std), seed, 0) + + m = x.mean().item() + s = x.std(unbiased=True).item() + + assert abs(m - mean) < max(3e-3, 0.02 * abs(std)) + assert abs(s - std) / max(std, 1e-12) < 0.03 + + # Determinism: same seed + subsequence must give bit-identical output. + y = torch.empty_like(x) + K.fill_normal(y, float(mean), float(std), seed, 0) + assert x.cpu() == pytest.approx(y.cpu(), rel=0, abs=0) + +@pytest.mark.parametrize("mean,std", [(0.0, 1.0), (1.5, 0.25)]) +def test_fill_normal_subsequence_independence(mean, std): + """Different subsequence offsets with the same seed must produce + non-identical, uncorrelated draws.""" + N = 64 * 1024 + seed = 0xABCDEF01 + + tensors = [] + for subseq in range(4): + x = torch.empty((N,), device="cuda", dtype=torch.float32) + K.fill_normal(x, float(mean), float(std), seed, subseq) + tensors.append(x.cpu()) + + for i in range(len(tensors)): + for j in range(i + 1, len(tensors)): + # Bit-identity: the subsequence argument must actually be used. + assert not torch.equal(tensors[i], tensors[j]), ( + f"Subsequences {i} and {j} produced identical output" + ) + + # Pearson |r| < 0.05 is a very conservative bound at N=64k + # (3-sigma ≈ 0.012), so a genuine failure stands out clearly. + a = tensors[i] - tensors[i].mean() + b = tensors[j] - tensors[j].mean() + r = (a * b).mean() / (a.std() * b.std()) + assert abs(r.item()) < 0.05, ( + f"Subsequences {i} and {j} are correlated: r={r.item():.4f}" + ) + +@pytest.mark.parametrize("mean,std", [(0.0, 1.0), (1.5, 0.25), (-1.0, 2.0)]) +def test_fill_normal_stats_bf16(mean, std): + """BF16 should satisfy the same mean/std checks as float32 with + tolerances relaxed to reflect its ~0.8 % per-value precision.""" + N = 256 * 1024 + seed = 0xABCDEF01 + x = torch.empty((N,), device="cuda", dtype=torch.bfloat16) + K.fill_normal(x, float(mean), float(std), seed, 0) + + # Upcast before reducing — BF16 accumulation is itself lossy. + xf = x.float() + m = xf.mean().item() + s = xf.std(unbiased=True).item() + + assert abs(m - mean) < max(1e-2, 0.05 * abs(std)) + assert abs(s - std) / max(std, 1e-12) < 0.05 + + # Determinism check mirrors the float32 test. + y = torch.empty_like(x) + K.fill_normal(y, float(mean), float(std), seed, 0) + assert torch.equal(x.cpu(), y.cpu()) + +# ============================================================ +# fused_classifier +# ============================================================ + + +@pytest.mark.parametrize("write_dlogits", [False, True]) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("B,T,V", [(2, 4, 16), (1, 8, 32), (4, 2, 64)]) +def test_fused_classifier_losses(B, T, V, write_dlogits, dtype): + """Per-token losses, lse, and (optionally) dlogits must match reference with no z-regularisation.""" + torch.manual_seed(0) + logits = torch.randn((B, T, V), device="cuda", dtype=dtype) + targets = torch.randint(0, V, (B, T), device="cuda", dtype=torch.int32) + + losses = torch.zeros((B, T), device="cuda", dtype=torch.float32) + lse = torch.empty((B, T), device="cuda", dtype=torch.float32) + dloss = 1.0 + logits_copy = logits.clone() # kernel mutates logits in place + + K.fused_classifier(logits_copy, losses, lse, dloss, targets, 0.0, write_dlogits) + + ref_lse = torch.logsumexp(logits.float(), dim=-1) + assert lse.cpu() == pytest.approx(ref_lse.cpu(), rel=1e-4, abs=1e-5) + + ref_losses = F.cross_entropy(logits.reshape(B * T, V).float(), targets.reshape(B * T).long(), reduction="none").reshape(B, T) + assert losses.cpu() == pytest.approx(ref_losses.cpu(), rel=1e-4, abs=1e-5) + + if write_dlogits: + probs = torch.softmax(logits.float(), dim=-1) + onehot = torch.zeros_like(probs) + onehot.scatter_(-1, targets.long().unsqueeze(-1), 1.0) + ref_dlogits = probs - onehot # dloss=1 everywhere + assert logits_copy.float().cpu() == pytest.approx(ref_dlogits.cpu(), rel=1e-2, abs=1e-3) + + + +# ============================================================ +# adamw_update +# ============================================================ + +def _adamw_reference(params, grads, m, v, lr, beta1, beta2, t, eps, wd): + m_new = beta1 * m.float() + (1 - beta1) * grads.float() + v_new = beta2 * v.float() + (1 - beta2) * grads.float() ** 2 + m_hat = m_new / (1 - beta1 ** t) + v_hat = v_new / (1 - beta2 ** t) + params_new = params.float() * (1 - lr * wd) - lr * m_hat / (v_hat.sqrt() + eps) + return params_new, m_new, v_new + + +@pytest.mark.parametrize("N", [1024, 8192, 65536]) +@pytest.mark.parametrize("p_dtype, g_dtype, m_dtype, v_dtype", [ + (torch.float32, torch.float32, torch.float32, torch.float32), + (torch.bfloat16, torch.bfloat16, torch.float32, torch.float32), + (torch.bfloat16, torch.bfloat16, torch.bfloat16, torch.float32), + (torch.bfloat16, torch.bfloat16, torch.bfloat16, torch.bfloat16), +]) +def test_adamw_update_matches_reference(N, p_dtype, g_dtype, m_dtype, v_dtype): + torch.manual_seed(0) + params = torch.randn((N,), device="cuda", dtype=p_dtype) + grads = torch.randn((N,), device="cuda", dtype=g_dtype) + m = torch.randn((N,), device="cuda", dtype=m_dtype) * 0.1 + v = torch.rand((N,), device="cuda", dtype=v_dtype) * 0.01 + 1e-8 + g_scale = torch.rand((), device="cuda", dtype=torch.float32) * 0.5 + 0.5 + + lr, beta1, beta2, t, eps, wd = 1e-3, 0.9, 0.999, 1, 1e-8, 0.1 + ref_params, ref_m, ref_v = _adamw_reference( + params.clone().float(), grads.float() * g_scale, m.clone().float(), v.clone().float(), lr, beta1, beta2, t, eps, wd + ) + + K.adamw_update(params, grads, m, v, lr, beta1, beta2, t, eps, wd, g_scale) + + rtol, atol = TOL[p_dtype] + assert params.float().cpu() == pytest.approx(ref_params.cpu(), rel=rtol, abs=atol) + rtol, atol = TOL[m_dtype] + assert m.float().cpu() == pytest.approx(ref_m.cpu(), rel=rtol, abs=atol) + rtol, atol = TOL[v_dtype] + assert v.float().cpu() == pytest.approx(ref_v.cpu(), rel=rtol, abs=atol) diff --git a/test/test_kernels_qk_norm.py b/test/test_kernels_qk_norm.py new file mode 100644 index 0000000..2283515 --- /dev/null +++ b/test/test_kernels_qk_norm.py @@ -0,0 +1,201 @@ +import pytest +import torch +from pyllmq import kernels as K +from conftest import DTYPES, TOL +from test_kernels_rope import _make_rope_freqs, _rope_python + + +# ============================================================ +# qk_norm_forward +# ============================================================ + +def _qk_norm_reference(inp, q_wgt, k_wgt, Nq, Nkv, eps): + B, T, _ = inp.shape + N = Nq + 2 * Nkv + HeadDim = q_wgt.shape[0] + x = inp.float().view(B, T, N, HeadDim) + + # Only norm Q and K heads + Nqk = Nq + Nkv + x_qk = x[:, :, :Nqk, :] # (B, T, Nq+Nkv, HeadDim) + + var = (x_qk ** 2).mean(dim=-1, keepdim=True) + s = torch.rsqrt(var + eps) # (B, T, Nq+Nkv, 1) + + wgt = torch.cat([ + q_wgt.float().unsqueeze(0).expand(Nq, -1), + k_wgt.float().unsqueeze(0).expand(Nkv, -1), + ], dim=0).view(1, 1, Nqk, HeadDim) + + x_qk_normed = ((x_qk * s).to(inp.dtype) * wgt).to(inp.dtype) # (B, T, Nq+Nkv, HeadDim) + + # V heads are passed through unchanged + x_v = x[:, :, Nqk:, :].to(inp.dtype) # (B, T, Nkv, HeadDim) + + # r_rms: pad V slots with 1.0 (or zeros) since they aren't normed + r_rms_qk = s.squeeze(-1).float() # (B, T, Nq+Nkv) + r_rms_v = torch.ones(B, T, Nkv, dtype=torch.float32, device=inp.device) + r_rms = torch.cat([r_rms_qk, r_rms_v], dim=-1) # (B, T, N) + + out = torch.cat([x_qk_normed, x_v], dim=2).view(B, T, N * HeadDim) + return out, r_rms + + +@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ + (1, 4, 2, 1, 16), + (2, 8, 4, 2, 64), + (4, 16, 8, 4, 128), +]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_qk_norm_forward(B, T, Nq, Nkv, HeadDim, dtype): + torch.manual_seed(0) + device = "cuda" + N = Nq + 2 * Nkv + eps = 1e-6 + + inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + out = torch.empty_like(inp) + # NOTE: the kernel does not write rrms for value heads at all; init to one to match reference + r_rms = torch.ones((B, T, N), device=device, dtype=torch.float32) + + K.qk_norm_forward(out, r_rms, inp, q_wgt, k_wgt, eps, Nq, Nkv) + ref_out, ref_r_std = _qk_norm_reference(inp, q_wgt, k_wgt, Nq, Nkv, eps) + + rtol, atol = TOL[dtype] + assert r_rms.cpu().numpy() == pytest.approx(ref_r_std.cpu().numpy(), rel=rtol, abs=atol) + assert out.float().cpu().numpy() == pytest.approx(ref_out.float().cpu().numpy(), rel=rtol, abs=atol) + + +# ---- In-place forward: out is inp ---- +# Exercises the early-return path for V-heads (inp == out) and the normalization path for QK-heads. +@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ + (1, 1, 4, 2, 64), + (2, 4, 2, 1, 128), +]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_qk_norm_forward_inplace(B, T, Nq, Nkv, HeadDim, dtype): + torch.manual_seed(0) + device = "cuda" + N = Nq + 2 * Nkv + eps = 1e-6 + + inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + v_start = Nq + Nkv + v_end = N + + inp_copy = inp.clone() + + out = inp + r_rms = torch.ones((B, T, N), device=device, dtype=torch.float32) + + K.qk_norm_forward(out, r_rms, inp, q_wgt, k_wgt, eps, Nq, Nkv) + + actual_v = out[:, :, v_start * HeadDim : v_end * HeadDim] + orig_v = inp_copy[:, :, v_start * HeadDim : v_end * HeadDim] + assert torch.allclose(actual_v, orig_v, rtol=0, atol=0), "V-heads changed in in-place forward" + + ref_out, _ = _qk_norm_reference(inp_copy, q_wgt, k_wgt, Nq, Nkv, eps) + rtol, atol = TOL[dtype] + assert out.float().cpu().numpy() == pytest.approx(ref_out.float().cpu().numpy(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ + (1, 4, 2, 1, 16), + (2, 8, 4, 2, 64), + (4, 16, 8, 4, 128), +]) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("with_abs_max", [False, True]) +def test_qk_norm_and_rope_forward(B, T, Nq, Nkv, HeadDim, dtype, with_abs_max): + torch.manual_seed(0) + device = "cuda" + N = Nq + 2 * Nkv + eps = 1e-6 + + inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + out = torch.empty_like(inp) + # NOTE: the kernel does not write rrms for value heads; init to 1 to match reference + r_rms = torch.ones((B, T, N), device=device, dtype=torch.float32) + + freq_dtype = torch.float32 if dtype == torch.float32 else torch.float16 + freqs = _make_rope_freqs(T, HeadDim, 1_000_000.0, freq_dtype, device) + + abs_max = torch.zeros((1,), device=device, dtype=torch.float32) if with_abs_max else None + + K.qk_norm_and_rope_forward(out, r_rms, inp, q_wgt, k_wgt, freqs, abs_max, eps, Nq, Nkv) + + # reference: qk-norm first, then rope on the result + ref_normed, ref_r_rms = _qk_norm_reference(inp, q_wgt, k_wgt, Nq, Nkv, eps) + ref_out = _rope_python(ref_normed.view(B, T, N, HeadDim), freqs, Nq, Nkv).view(B, T, N * HeadDim) + + rtol, atol = TOL[dtype] + assert r_rms.cpu().numpy() == pytest.approx(ref_r_rms.cpu().numpy(), rel=rtol, abs=atol) + assert out.float().cpu().numpy() == pytest.approx(ref_out.float().cpu().numpy(), rel=rtol, abs=atol) + + if with_abs_max: + # kernel takes abs-max over Q+K post-rope and V pass-through + ref_abs_max = ref_out.float().abs().max().item() + assert abs_max.item() == pytest.approx(ref_abs_max, rel=rtol, abs=atol) + +@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ + (1, 4, 2, 1, 16), + (2, 8, 4, 2, 64), + (4, 16, 8, 4, 64), + (1, 1, 4, 2, 128), +]) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("in_place", [True, False]) +def test_qk_norm_backward(B, T, Nq, Nkv, HeadDim, dtype, in_place: bool): + torch.manual_seed(0) + device = "cuda" + N = Nq + 2 * Nkv + eps = 1e-6 + + inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + dout = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + + # Reference in fp32 + inp_f = inp.float().detach().requires_grad_(True) + q_wgt_f = q_wgt.float().detach().requires_grad_(True) + k_wgt_f = k_wgt.float().detach().requires_grad_(True) + ref_out_f, _ = _qk_norm_reference(inp_f, q_wgt_f, k_wgt_f, Nq, Nkv, eps) + ref_out_f.backward(dout.float()) + ref_dinp = inp_f.grad + ref_dq_wgt = q_wgt_f.grad + ref_dk_wgt = k_wgt_f.grad + + # Run forward to get the kernel's rstd. + r_rms = torch.empty((B, T, N), device=device, dtype=torch.float32) + r_rms.fill_(1.0) # V-slot sentinel; kernel only writes QK slots + K.qk_norm_forward(torch.empty_like(inp), r_rms, inp, q_wgt, k_wgt, eps, Nq, Nkv) + + scratch_size = K.get_qknorm_backward_scratch_size(Nq, Nkv, HeadDim, inp.dtype) + scratch = torch.zeros((scratch_size,), device=device, dtype=torch.uint8) + dq_wgt_k = torch.zeros_like(q_wgt) + dk_wgt_k = torch.zeros_like(k_wgt) + + if in_place: + dinp_buf = dout.clone() + K.qk_norm_backward(dinp_buf, dq_wgt_k, dk_wgt_k, scratch, + dinp_buf, inp, q_wgt, k_wgt, r_rms, + None, Nq, Nkv) + dinp = dinp_buf + else: + dinp = torch.zeros_like(inp) + K.qk_norm_backward(dinp, dq_wgt_k, dk_wgt_k, scratch, + dout, inp, q_wgt, k_wgt, r_rms, + None, Nq, Nkv) + + rtol, atol = TOL[dtype] + + assert dinp.float().cpu().numpy() == pytest.approx(ref_dinp.cpu().numpy(), rel=rtol, abs=atol) + assert dq_wgt_k.float().cpu().numpy() == pytest.approx(ref_dq_wgt.cpu().numpy(), rel=rtol, abs=atol) + assert dk_wgt_k.float().cpu().numpy() == pytest.approx(ref_dk_wgt.cpu().numpy(), rel=rtol, abs=atol) diff --git a/test/test_kernels_quant.py b/test/test_kernels_quant.py new file mode 100644 index 0000000..2727ada --- /dev/null +++ b/test/test_kernels_quant.py @@ -0,0 +1,91 @@ +import pytest +import torch +from pyllmq import kernels as K +from conftest import DTYPES, TOL + + +# ============================================================ +# abs_max +# Reduction over exact fp values — output is one of the input +# values, so the scalar result must be exact. +# ============================================================ + +@pytest.mark.parametrize("shape", [(16, 64), (4, 128, 32)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_abs_max_writes_scalar(shape, dtype): + x = torch.randn(shape, device="cuda", dtype=dtype) + ref = x.abs().max() + result = torch.empty((), device="cuda", dtype=torch.float32) + K.abs_max(result, x) + assert torch.isfinite(result) + # abs_max just selects an existing value — no arithmetic error possible. + assert result.cpu() == pytest.approx(ref.float().cpu(), rel=0, abs=0) + +# ============================================================ +# quantize_with_abs_max +# ============================================================ + +@pytest.mark.parametrize("N", [1024, 8192, 65536]) +@pytest.mark.parametrize("dtype", [torch.float32]) +def test_quantize_with_abs_max_bf16(N, dtype): + device = "cuda" + x = torch.randn((N,), device=device, dtype=dtype) + abs_max_val = torch.tensor([x.float().abs().max().item()], device=device, dtype=torch.float32) + out = torch.empty((N,), device=device, dtype=torch.bfloat16) + scale = torch.empty((), device=device, dtype=torch.float32) + K.quantize_with_abs_max(out, scale, x, abs_max_val) + assert scale.item() == pytest.approx(1.0, rel=0, abs=0) + assert out.float().cpu() == pytest.approx(x.bfloat16().float().cpu(), rel=0.01) + + +@pytest.mark.parametrize("N", [1024, 8192]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_quantize_with_abs_max_fp8(N, dtype): + device = "cuda" + x = torch.randn((N,), device=device, dtype=dtype) + abs_max_val = torch.tensor([x.float().abs().max().item()], device=device, dtype=torch.float32) + out = torch.empty((N,), device=device, dtype=torch.float8_e4m3fn) + scale = torch.empty((), device=device, dtype=torch.float32) + K.quantize_with_abs_max(out, scale, x, abs_max_val) + assert scale.item() == pytest.approx(abs_max_val.item() / 448.0) + dequant = out.float() * scale + assert dequant.cpu() == pytest.approx(x.float().cpu(), rel=0.1, abs=1e-3) + +# ============================================================ +# quantize_and_transpose_with_abs_max +# ============================================================ + +@pytest.mark.parametrize("rows,cols", [(32, 32), (128, 256), (512, 1024)]) +@pytest.mark.parametrize("dtype", [torch.float32]) +def test_quantize_and_transpose_with_abs_max_bf16(rows, cols, dtype): + device = "cuda" + x = torch.randn((rows, cols), device=device, dtype=dtype) + abs_max_val = torch.tensor([x.float().abs().max().item()], device=device, dtype=torch.float32) + out = torch.empty((cols, rows), device=device, dtype=torch.bfloat16) + scale = torch.empty((), device=device, dtype=torch.float32) + K.quantize_and_transpose_with_abs_max(out, scale, x, abs_max_val) + + # no assert for scale, as scale is unused for bf16 + + expected = x.bfloat16().T.contiguous() + assert out.float().cpu() == pytest.approx(expected.float().cpu(), rel=0.01) + + +@pytest.mark.parametrize("rows,cols", [(32, 32), (128, 256), (512, 1024)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_quantize_and_transpose_with_abs_max_fp8(rows, cols, dtype): + device = "cuda" + x = torch.randn((rows, cols), device=device, dtype=dtype) + abs_max_val = torch.tensor([x.float().abs().max().item()], device=device, dtype=torch.float32) + out = torch.empty((cols, rows), device=device, dtype=torch.float8_e4m3fn) + scale = torch.empty((), device=device, dtype=torch.float32) + K.quantize_and_transpose_with_abs_max(out, scale, x, abs_max_val) + + assert scale.item() == pytest.approx(abs_max_val.item() / 448.0) + + # Dequantize and verify values match input (with fp8 tolerance) + dequant = out.float() * scale + + # Verify transpose: reshape input as [rows, cols], transpose to [cols, rows] + expected = x.float().T.contiguous() + assert dequant.cpu() == pytest.approx(expected.cpu(), rel=0.1, abs=1e-4) diff --git a/test/test_kernels_rmsnorm.py b/test/test_kernels_rmsnorm.py new file mode 100644 index 0000000..b7a7a05 --- /dev/null +++ b/test/test_kernels_rmsnorm.py @@ -0,0 +1,143 @@ +import pytest +import torch +from pyllmq import kernels as K +from conftest import DTYPES, TOL + +# ============================================================ +# rmsnorm +# ============================================================ + +def _rmsnorm_reference(inp, weight, eps): + var = (inp.float() ** 2).mean(dim=-1, keepdim=True) + rms = torch.sqrt(var + eps) + r_rms = 1.0 / rms + out = inp.float() * r_rms * weight.float() + return out.to(inp.dtype), r_rms.squeeze(-1) + + +def _rmsnorm_backward_reference(dout, inp, weight, rstd): + # dout, inp: (B, T, C), weight: (C,), rstd: (B, T) + B, T, C = inp.shape + dout_f = dout.float() + inp_f = inp.float() + w_f = weight.float() + r = rstd.unsqueeze(-1) # (B, T, 1) + + # dweight: sum over B, T + dweight = (dout_f * inp_f * r).sum(dim=(0, 1)) + + # dinp + normed = inp_f * r # (B, T, C) + dy_w = dout_f * w_f # (B, T, C) + dot = (dy_w * normed).sum(dim=-1, keepdim=True) # (B, T, 1) + dinp = r * (dy_w - normed * dot / C) + + return dinp.to(inp.dtype), dweight.to(weight.dtype) + + +@pytest.mark.parametrize("B,T,C", [(2, 3, 16), (1, 2, 64), (8, 256, 1024)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_rmsnorm_forward_matches_reference(B, T, C, dtype): + torch.manual_seed(0) + inp = torch.randn((B, T, C), device="cuda", dtype=dtype) + weight = torch.randn((C,), device="cuda", dtype=dtype) + out = torch.empty_like(inp) + rms = torch.empty((B, T), device="cuda", dtype=torch.float32) + eps = 1e-6 + + K.rmsnorm_forward(out, rms, inp, weight, None, eps) + ref_out, ref_rms = _rmsnorm_reference(inp, weight, eps) + + rtol, atol = TOL[dtype] + assert rms.cpu() == pytest.approx(ref_rms.cpu(), rel=rtol, abs=atol) + assert out.float().cpu() == pytest.approx(ref_out.float().cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("B,T,C", [(2, 3, 16), (1, 4, 64), (4, 8, 256)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_rmsnorm_backward(B, T, C, dtype): + torch.manual_seed(0) + inp = torch.randn((B, T, C), device="cuda", dtype=dtype) + weight = torch.randn((C,), device="cuda", dtype=dtype) + dout = torch.randn((B, T, C), device="cuda", dtype=dtype) + eps = 1e-6 + + # Forward to get rstd + _, rstd = _rmsnorm_reference(inp, weight, eps) + rstd = rstd.to(torch.float32).cuda() + + dinp = torch.empty_like(inp) + dweight = torch.zeros((C,), device="cuda", dtype=dtype) + scratch = torch.zeros(K.get_rmsnorm_backward_scratch_size(C), device="cuda", dtype=torch.float32) + dresidual = torch.zeros_like(inp) + + K.rmsnorm_backward(dinp, dweight, scratch, dresidual, dout, inp, weight, rstd, None) + + ref_dinp, ref_dweight = _rmsnorm_backward_reference(dout, inp, weight, rstd) + rtol, atol = TOL[dtype] + assert dinp.float().cpu() == pytest.approx(ref_dinp.float().cpu(), rel=rtol, abs=atol) + assert dweight.float().cpu() == pytest.approx(ref_dweight.float().cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("B,T,C", [(2, 3, 16), (1, 4, 64), (4, 8, 256)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_rmsnorm_backward_dresidual_accumulate(B, T, C, dtype): + torch.manual_seed(1) + inp = torch.randn((B, T, C), device="cuda", dtype=dtype) + weight = torch.randn((C,), device="cuda", dtype=dtype) + dout = torch.randn((B, T, C), device="cuda", dtype=dtype) + eps = 1e-6 + + _, rstd = _rmsnorm_reference(inp, weight, eps) + rstd = rstd.to(torch.float32).cuda() + + dinp = torch.empty_like(inp) + dweight = torch.zeros((C,), device="cuda", dtype=dtype) + scratch = torch.zeros(K.get_rmsnorm_backward_scratch_size(C), device="cuda", dtype=torch.float32) + + # Pre-fill dresidual with a non-zero sentinel so accumulation is detectable. + dresidual = torch.randn_like(inp) + dresidual_initial = dresidual.clone() + ref_dinp, _ = _rmsnorm_backward_reference(dout, inp, weight, rstd) + expected_dinp = dresidual_initial.float() + ref_dinp.float() + + K.rmsnorm_backward(dinp, dweight, scratch, dresidual, dout, inp, weight, rstd, None) + + rtol, atol = TOL[dtype] + if dtype == torch.bfloat16: + # TODO check why we need so much tolerance + atol = 8e-3 + rtol = 0.015 + assert dinp.float().cpu() == pytest.approx( + expected_dinp.cpu(), rel=rtol, abs=atol + ) + + +# ============================================================ +# fused_residual_rmsnorm_forward +# ============================================================ + +@pytest.mark.parametrize("B,T,C", [(2, 5, 64), (1, 3, 256), (4, 8, 128), (8, 16, 512)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_fused_residual_rmsnorm_forward_reference(B, T, C, dtype): + device = "cuda" + torch.manual_seed(0) + inp1 = torch.randn((B, T, C), device=device, dtype=dtype) + inp2 = torch.randn_like(inp1) + weight = torch.randn((C,), device=device, dtype=dtype) + residual = torch.empty_like(inp1) + normed = torch.empty_like(inp1) + rrms = torch.empty((B, T), device=device, dtype=torch.float32) + eps = 1e-6 + + K.fused_residual_rmsnorm_forward(residual, normed, rrms, inp1, inp2, weight, None, eps) + + res_ref = (inp1.float() + inp2.float()).to(dtype) + var = (res_ref.float() ** 2).mean(dim=-1, keepdim=True) + r_rms = 1.0 / torch.sqrt(var + eps) + norm_ref = (res_ref.float() * r_rms * weight.float()).to(dtype) + + rtol, atol = TOL[dtype] + assert residual.float().cpu() == pytest.approx(res_ref.float().cpu(), rel=rtol, abs=atol) + assert normed.float().cpu() == pytest.approx(norm_ref.float().cpu(), rel=rtol, abs=atol) + assert rrms.cpu() == pytest.approx(r_rms.squeeze(-1).float().cpu(), rel=rtol, abs=atol) diff --git a/test/test_kernels_rope.py b/test/test_kernels_rope.py new file mode 100644 index 0000000..ff3d3dc --- /dev/null +++ b/test/test_kernels_rope.py @@ -0,0 +1,73 @@ +import pytest +import torch +from pyllmq import kernels as K +from conftest import DTYPES, TOL + + +# ============================================================ +# rope forward + backward +# ============================================================ + +def _make_rope_freqs(T, head_dim, theta, dtype, device): + assert head_dim % 2 == 0 + half = head_dim // 2 + idx = torch.arange(half, device=device, dtype=torch.float32) + inv_freq = theta ** (-2 * idx / head_dim) + t = torch.arange(T, device=device, dtype=torch.float32).unsqueeze(1) + angles = t * inv_freq.unsqueeze(0) + cos = torch.cos(angles).to(dtype) + sin = torch.sin(angles).to(dtype) + return torch.stack([cos, sin], dim=-1).flatten(start_dim=1) # (T, head_dim) + + +def _rope_python(x, freqs_cis, Nq, Nkv, backward=False): + B, T, N, HD = x.shape + half = HD // 2 + cos = freqs_cis[:, 0::2].float() + sin = freqs_cis[:, 1::2].float() + if backward: + sin = -sin + cos = cos[None, :, None, :] + sin = sin[None, :, None, :] + + q = x[:, :, :Nq, :] + k = x[:, :, Nq:Nq + Nkv, :] + v = x[:, :, Nq + Nkv:, :] + + def rotate(h): + h = h.float() + return torch.cat([ + h[..., :half] * cos - h[..., half:] * sin, + h[..., :half] * sin + h[..., half:] * cos, + ], dim=-1).to(x.dtype) + + return torch.cat([rotate(q), rotate(k), v], dim=2) + + +@pytest.mark.parametrize("B,T,Nq,Nkv,HD", [ + (1, 8, 2, 1, 8), + (2, 4, 1, 2, 16), + (2, 16, 4, 2, 32), + (4, 32, 8, 4, 64), +]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_rope_forward_backward_matches_python(B, T, Nq, Nkv, HD, dtype): + device = "cuda" + C = (Nq + 2 * Nkv) * HD + x = torch.randn((B, T, C), device=device, dtype=dtype) + out = torch.empty_like(x) + + freq_dtype = torch.float32 if dtype == torch.float32 else torch.float16 + freqs = _make_rope_freqs(T, HD, 1_000_000.0, freq_dtype, device) + + K.rope_forward(out, x, freqs, None, Nq, Nkv) + ref = _rope_python(x.view(B, T, Nq + 2 * Nkv, HD), freqs, Nq, Nkv).view(B, T, C) + + rtol, atol = TOL[dtype] + assert out.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) + + dout = torch.randn_like(x) + dinp = torch.empty_like(x) + K.rope_backward(dinp, dout, freqs, None, Nq, Nkv) + ref_bw = _rope_python(dout.view(B, T, Nq + 2 * Nkv, HD), freqs, Nq, Nkv, backward=True).view(B, T, C) + assert dinp.float().cpu() == pytest.approx(ref_bw.float().cpu(), rel=rtol, abs=atol) diff --git a/test/test_kernels_sr.py b/test/test_kernels_sr.py new file mode 100644 index 0000000..aa20622 --- /dev/null +++ b/test/test_kernels_sr.py @@ -0,0 +1,128 @@ +import pytest +import torch +from pyllmq import kernels as K +from conftest import DTYPES, TOL + + +# ============================================================ +# vector_add_sr +# ============================================================ + +@pytest.mark.parametrize("nelem", [4096, 16384, 65536]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_vector_add_sr_determinism_and_accuracy(nelem, dtype): + device = "cuda" + a = torch.randn((nelem,), device=device, dtype=dtype) + b = torch.randn_like(a) + out1 = torch.empty_like(a) + out2 = torch.empty_like(a) + seed = 12345 + scale = torch.tensor(0.75, dtype=torch.float32) + K.vector_add_sr(out1, a, b, scale, seed) + K.vector_add_sr(out2, a, b, scale, seed) + + # Determinism: same seed must give bit-identical output. + assert out1.float().cpu() == pytest.approx(out2.float().cpu(), rel=0, abs=0) + + # Accuracy vs fp32 reference; stochastic rounding may introduce ~1 ulp at target dtype. + ref = scale.item() * (a.float() + b.float()) + assert out1.float().cpu() == pytest.approx(ref.cpu(), rel=1e-2, abs=5e-3) + + +# ============================================================ +# vector_reduce_sr +# ============================================================ +@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192), (8, 16384)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_vector_reduce_sr_sum_matches_reference(n_shards, nelem, dtype): + """With scale=1, the reduction must equal the sum across shards.""" + torch.manual_seed(0) + src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) + dest = torch.empty((nelem,), device="cuda", dtype=dtype) + scale = torch.tensor(1.0, dtype=torch.float32) + + K.vector_reduce_sr(dest, src, scale, n_shards, seed=0) + + ref = src.view(n_shards, nelem).float().sum(dim=0) + rtol, atol = TOL[dtype] + assert dest.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) + + + +@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_vector_reduce_sr_determinism(n_shards, nelem, dtype): + src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) + dest1 = torch.empty((nelem,), device="cuda", dtype=dtype) + dest2 = torch.empty((nelem,), device="cuda", dtype=dtype) + scale = torch.tensor(0.5, dtype=torch.float32) + + K.vector_reduce_sr(dest1, src, scale, n_shards, seed=99) + K.vector_reduce_sr(dest2, src, scale, n_shards, seed=99) + + assert dest1.float().cpu() == pytest.approx(dest2.float().cpu(), rel=0, abs=0) + + +@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_vector_reduce_sr_scale_zero(n_shards, nelem, dtype): + src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) + dest = torch.empty((nelem,), device="cuda", dtype=dtype) + scale = torch.tensor(0.0, dtype=torch.float32) + + K.vector_reduce_sr(dest, src, scale, n_shards, seed=0) + + assert dest.float().cpu() == pytest.approx(torch.zeros(nelem).cpu(), rel=0, abs=0) + + +@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_vector_reduce_sr_skip(n_shards, nelem, dtype): + """Skipped shard should not contribute to the result.""" + torch.manual_seed(42) + src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) + dest = torch.empty((nelem,), device="cuda", dtype=dtype) + scale = torch.tensor(1.0, dtype=torch.float32) + + skip = 1 + K.vector_reduce_sr(dest, src, scale, n_shards, skip=skip, seed=0) + + shards = src.view(n_shards, nelem).float() + ref = sum(shards[k] for k in range(n_shards) if k != skip) + rtol, atol = TOL[dtype] + assert dest.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("n_shards,nelem", [(2, 4096), (4, 8192)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_vector_reduce_sr_accumulate(n_shards, nelem, dtype): + """With accumulate=True, dest's initial values should be included in the sum.""" + torch.manual_seed(7) + src = torch.randn((n_shards * nelem,), device="cuda", dtype=dtype) + dest = torch.randn((nelem,), device="cuda", dtype=dtype) + dest_initial = dest.clone() + scale = torch.tensor(1.0, dtype=torch.float32) + + K.vector_reduce_sr(dest, src, scale, n_shards, accumulate=True, seed=0) + + ref = src.view(n_shards, nelem).float().sum(dim=0) + dest_initial.float() + rtol, atol = TOL[dtype] + assert dest.float().cpu() == pytest.approx(ref.cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("dtype", DTYPES) +def test_vector_reduce_sr_stochastic_rounding_unbiased(dtype: torch.dtype, nelem=4096): + """SR rounding should be unbiased: mean of many rounded values ≈ true mean.""" + n_shards = 2 + # Use values that are exactly halfway between representable values to maximise rounding effect + src = torch.ones((n_shards * nelem,), device="cuda", dtype=dtype) * 0.5 + results = [] + for seed in range(20): + dest = torch.empty((nelem,), device="cuda", dtype=dtype) + scale = torch.tensor(1.0, dtype=torch.float32) + K.vector_reduce_sr(dest, src, scale, n_shards, seed=seed) + results.append(dest.float().cpu()) + + mean_result = torch.stack(results).mean(dim=0) + # True answer is 1.0 (sum of two 0.5 shards); mean over seeds should be close + assert mean_result == pytest.approx(torch.ones(nelem).cpu(), rel=0.01, abs=0.01) diff --git a/test/test_kernels_swiglu.py b/test/test_kernels_swiglu.py new file mode 100644 index 0000000..63c111a --- /dev/null +++ b/test/test_kernels_swiglu.py @@ -0,0 +1,74 @@ +import pytest +import torch +from pyllmq import kernels as K +from conftest import DTYPES, TOL + + +# ============================================================ +# swiglu_forward +# ============================================================ + +def _swiglu_reference(x: torch.Tensor) -> torch.Tensor: + a, b = torch.tensor_split(x.float(), 2, dim=-1) + return (torch.nn.functional.silu(b) * a).to(x.dtype) + + +def _swiglu_backward_reference(dout, inp): + inp_f = inp.detach().float().requires_grad_(True) + out = _swiglu_reference(inp_f) + out.backward(dout.float()) + return inp_f.grad.to(inp.dtype) + + +@pytest.mark.parametrize("B,T,C", [(1, 16, 128), (8, 8, 16), (5, 32, 32), (32, 32, 1024)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_swiglu_forward_matches_reference(B, T, C, dtype): + torch.manual_seed(123) + inp = torch.randn((B, T, 2 * C), device="cuda", dtype=dtype) + out = torch.empty((B, T, C), device="cuda", dtype=dtype) + + K.swiglu_forward(out, inp, None) + + ref = _swiglu_reference(inp) + rtol, atol = TOL[dtype] + assert out.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) + + +@pytest.mark.parametrize("B,T,C", [(2, 8, 64), (4, 16, 128)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_swiglu_forward_quant_fp8(B, T, C, dtype): + """swiglu_forward_quant writes fp8 output, using the supplied abs-max to scale values""" + torch.manual_seed(42) + inp = torch.randn((B, T, 2 * C), device="cuda", dtype=dtype) + out = torch.empty((B, T, C), device="cuda", dtype=torch.float8_e4m3fn) + scale = torch.empty((), device="cuda", dtype=torch.float32) + + ref = _swiglu_reference(inp.float()) + abs_max = ref.abs().max().float() + K.swiglu_forward_quant(out, scale, inp, abs_max) + + dequant = out.float() * scale + # fp8 has limited precision — use loose tolerance + assert dequant.cpu() == pytest.approx(ref.cpu(), rel=0.125, abs=1e-2) + + expected_scale = abs_max / 448.0 + assert scale.cpu() == pytest.approx(expected_scale.cpu(), rel=1e-3) + +# ============================================================ +# swiglu_backward +# ============================================================ + +@pytest.mark.parametrize("B,T,C", [(1, 8, 256), (4, 16, 64)]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_swiglu_backward(B, T, C, dtype): + torch.manual_seed(0) + inp = torch.randn((B, T, 2 * C), device="cuda", dtype=dtype) + dout = torch.randn((B, T, C), device="cuda", dtype=dtype) + dinp = torch.empty_like(inp) + + K.swiglu_backward(dinp, dout, inp, None) + + ref = _swiglu_backward_reference(dout, inp) + rtol, atol = TOL[dtype] + + assert dinp.float().cpu() == pytest.approx(ref.float().cpu(), rel=rtol, abs=atol) From 40b3ba446c618103ec605a206356da7fbdf8ed0f Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 26 May 2026 01:10:02 +0200 Subject: [PATCH 04/19] fixup --- src/kernels/qk_norm.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/kernels/qk_norm.cu b/src/kernels/qk_norm.cu index 5d73903..26f30bb 100644 --- a/src/kernels/qk_norm.cu +++ b/src/kernels/qk_norm.cu @@ -66,7 +66,7 @@ __global__ void qk_norm_forward_simple_kernel(Float* out, float* r_rms, const Fl // load weights into shared memory // do this before we allow any threads to exit! - extern __shared__ char* smem[]; + extern __shared__ int4 smem[]; // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so // let's keep everything as x128 @@ -147,7 +147,7 @@ qk_norm_backward_kernel(Float* dinp, std::byte* scratch, // load weights into shared memory // do this before we allow any threads to exit! - extern __shared__ char* smem[]; + extern __shared__ int4 smem[]; __shared__ float block_abs_max; float thread_abs_max = 0.f; @@ -329,7 +329,7 @@ __global__ void qk_norm_and_rope_fwd_kernel(Float* out, float* r_rms, float* abs // load weights into shared memory // do this before we allow any threads to exit! - extern __shared__ char* smem[]; + extern __shared__ int4 smem[]; // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so // let's keep everything as x128 From 52de2571466ab071108834057510b074ba7313c5 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 26 May 2026 01:24:35 +0200 Subject: [PATCH 05/19] bench --- scripts/bench_qk.py | 68 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 scripts/bench_qk.py diff --git a/scripts/bench_qk.py b/scripts/bench_qk.py new file mode 100644 index 0000000..f5bbd7c --- /dev/null +++ b/scripts/bench_qk.py @@ -0,0 +1,68 @@ +import pytest +import torch +from pyllmq import kernels as K + + +DTYPES = [torch.float32, torch.bfloat16] + +# Tolerances by dtype for approximate kernels: (rtol, atol) +TOL = { + torch.float32: (5e-4, 5e-5), + torch.bfloat16: (1e-2, 5e-3), +} + + +@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ + (4, 1024, 8, 4, 128), +]) +@pytest.mark.parametrize("dtype", DTYPES) +def test_qk_norm_forward(B, T, Nq, Nkv, HeadDim, dtype): + torch.manual_seed(0) + device = "cuda" + N = Nq + 2 * Nkv + eps = 1e-6 + + inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + out = torch.empty_like(inp) + # NOTE: the kernel does not write rrms for value heads at all; init to one to match reference + r_rms = torch.ones((B, T, N), device=device, dtype=torch.float32) + + K.qk_norm_forward(out, r_rms, inp, q_wgt, k_wgt, eps, Nq, Nkv) + +@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ + (4, 1024, 8, 4, 128), +]) +@pytest.mark.parametrize("dtype", [torch.float32]) +def test_qk_norm_backward(B, T, Nq, Nkv, HeadDim, dtype): + torch.manual_seed(0) + device = "cuda" + N = Nq + 2 * Nkv + eps = 1e-6 + + inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype, requires_grad=True) + q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype, requires_grad=True) + k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype, requires_grad=True) + dout = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + + + scratch_size = K.get_qknorm_backward_scratch_size(Nq, Nkv, HeadDim, inp.dtype) + scratch = torch.zeros((scratch_size,), device=device, dtype=torch.uint8) + dinp = torch.zeros_like(inp) + dq_wgt_k = torch.zeros_like(q_wgt) + dk_wgt_k = torch.zeros_like(k_wgt) + r_rms = torch.ones((B, T, N), device=device, dtype=torch.float32) + K.qk_norm_forward(torch.empty_like(inp), r_rms, inp, q_wgt, k_wgt, eps, Nq, Nkv) + + print(f"BW: dout {dout.nbytes // 1024 // 1024} MB inp {inp.nbytes // 1024 // 1024} MB dinp {dinp.nbytes // 1024 // 1024} MB") + K.qk_norm_backward(dinp, dq_wgt_k, dk_wgt_k, scratch, + dout, inp, q_wgt, k_wgt, r_rms, + None, Nq, Nkv) + + +if __name__ == "__main__": + #test_qk_norm_forward(4, 1024, 16, 2, 128, torch.bfloat16) + # qk_norm_backward_kernel + test_qk_norm_backward(4, 1024, 16, 2, 128, torch.bfloat16) + print("QK norm backward kernel test passed") From 832656445a37a3d3f8458ece8d4a0fdedb483c1e Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 26 May 2026 02:11:40 +0200 Subject: [PATCH 06/19] enable caching of x in smem --- src/kernels/qk_norm.cu | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/src/kernels/qk_norm.cu b/src/kernels/qk_norm.cu index 26f30bb..817a678 100644 --- a/src/kernels/qk_norm.cu +++ b/src/kernels/qk_norm.cu @@ -155,6 +155,9 @@ qk_norm_backward_kernel(Float* dinp, std::byte* scratch, Float* s_wgt = reinterpret_cast(smem); float* s_d_wgt_base = reinterpret_cast(s_wgt + HeadDim); float* s_d_wgt = s_d_wgt_base + HeadDim * group; + // Per-group cache for input x, to avoid the second global load. + x128* s_in_base = reinterpret_cast(s_d_wgt_base + HeadDim * num_groups); + x128* s_in = s_in_base + group * (HeadDim / x128::size); for(int i = threadIdx.x * x128::size; i < HeadDim; i += blockDim.x * x128::size) { if(h < Nq) { @@ -220,6 +223,7 @@ qk_norm_backward_kernel(Float* dinp, std::byte* scratch, x128 o = x128::load(dout_i + i); x128 x = x128::load(inp_i + i); x128 w = x128::load(s_wgt + i); + s_in[i / x128::size] = x; for (int k = 0; k < x128::size; k++) { sum_xow += (float)w[k] * (float)o[k] * (float)x[k]; } @@ -230,7 +234,7 @@ qk_norm_backward_kernel(Float* dinp, std::byte* scratch, for (int i = lane * x128::size; i < HeadDim; i += GroupSize * x128::size) { x128 o = x128::load_cs(dout_i + i); - x128 x = x128::load_cs(inp_i + i); + x128 x = s_in[i / x128::size]; x128 w = x128::load(s_wgt + i); x128 dx = x128::zeros(); @@ -486,10 +490,12 @@ void qk_norm_and_rope_forward(Float* out, float* r_rms, float* abs_max_ptr, // Get the amount of smem per block for backward template - size_t qk_norm_backward_smem(int HeadDim) { +size_t qk_norm_backward_smem(int HeadDim) { constexpr int block_size = 512; constexpr int num_groups = block_size / GroupSize; - return HeadDim * sizeof(Float) + num_groups * HeadDim * sizeof(float); + return HeadDim * sizeof(Float) // s_wgt + + num_groups * HeadDim * sizeof(float) // s_d_wgt + + num_groups * HeadDim * sizeof(Float); // s_in } // Get the maximum number of concurrent blocks in x direction @@ -543,6 +549,12 @@ void qk_norm_backward_imp(Float* dinp, Float* dq_wgt, Float* dk_wgt, std::byte* constexpr int block_size = 512; const int Nh = Nq + 2 * Nkv; const size_t smem = qk_norm_backward_smem(HeadDim); + + CUDA_CHECK(cudaFuncSetAttribute( + qk_norm_backward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem)); + const int x_blocks = qk_norm_backward_x_blocks(Nq, Nkv, HeadDim, dp); dim3 grid(x_blocks, Nh); From f7b6c51405a7a94d569f0817140228361d4c07a8 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 26 May 2026 02:41:21 +0200 Subject: [PATCH 07/19] copy backward kernel for fused implementation --- src/kernels/qk_norm.cu | 164 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 164 insertions(+) diff --git a/src/kernels/qk_norm.cu b/src/kernels/qk_norm.cu index 817a678..91e5c08 100644 --- a/src/kernels/qk_norm.cu +++ b/src/kernels/qk_norm.cu @@ -424,6 +424,170 @@ __global__ void qk_norm_and_rope_fwd_kernel(Float* out, float* r_rms, float* abs } } +// ============================================================ +// Fused QK-norm + RoPE backward +// ============================================================ + +template +__global__ void __launch_bounds__(512, 2) +qk_norm_and_rope_backward_kernel(Float* dinp, std::byte* scratch, + const Float* dout, const Float* inp, + const Float* q_wgt, const Float* k_wgt, + const float* rstd, const FloatFreq* freqs_cis, + float* abs_max_ptr, + int BT, int T, int Nq, int Nkv, int HeadDim) { + // formulas: + // rms = sqrt(sum x_j^2 / C + eps) + // y_i = w_i x_i / rms + // + // o_i := dL/dy_i + // dw_i = sum (x_i o_i / rms) + // xow := sum x_i o_i w_i + // dy_j/drms = - w_j x_j / rms² + // drms/dx_j = x_j/(C rms) + // dx_i = dL/dy_j (dy_j/dx_i + dy_j/drms drms/dx_i) + // = o_i w_i / rms - sum_j o_j w_j x_j / rms² x_i/(C rms) + // = o_i w_i / rms - x_i (xow/C) / rms³ + // strategy: y-dimension of block indicates which head is operated on + // x-dimension goes over tokens in groups of eight threads + using x128 = GenericVector; + using f128 = GenericVector; + using fvec = GenericVector; + + const int lane = threadIdx.x % GroupSize; + const int group = threadIdx.x / GroupSize; + const int num_groups = blockDim.x / GroupSize; + const int h = blockIdx.y; + const int Nh = Nq + 2 * Nkv; + + // load weights into shared memory + // do this before we allow any threads to exit! + extern __shared__ int4 smem[]; + + __shared__ float block_abs_max; + float thread_abs_max = 0.f; + + Float* s_wgt = reinterpret_cast(smem); + float* s_d_wgt_base = reinterpret_cast(s_wgt + HeadDim); + float* s_d_wgt = s_d_wgt_base + HeadDim * group; + // Per-group cache for input x, to avoid the second global load. + x128* s_in_base = reinterpret_cast(s_d_wgt_base + HeadDim * num_groups); + x128* s_in = s_in_base + group * (HeadDim / x128::size); + + for(int i = threadIdx.x * x128::size; i < HeadDim; i += blockDim.x * x128::size) { + if(h < Nq) { + x128::load(q_wgt + i).store(s_wgt + i); + } else if (h < Nq + Nkv){ + x128::load(k_wgt + i).store(s_wgt + i); + } + // v-heads need no loading + } + + for(int i = threadIdx.x * f128::size; i < HeadDim * num_groups; i += blockDim.x * f128::size) { + f128::zeros().store(s_d_wgt_base + i); + } + + if (abs_max_ptr && threadIdx.x == 0) { + block_abs_max = 0.f; + } + + __syncthreads(); + + const int groups_in_grid = num_groups * gridDim.x; + const int start_idx = blockIdx.x * (blockDim.x / GroupSize) + group; + + for (int bt = start_idx; ; bt += groups_in_grid) { + bool valid = bt < BT; + unsigned int active_mask = __ballot_sync(0xffffffffu, valid); + bool all_finished = !__any_sync(0xffffffffu, valid); + if (all_finished) + break; + if (!valid) + continue; + + // adjusted pointers to current token + const Float* inp_i = inp + bt * Nh * HeadDim + h * HeadDim; + const Float* dout_i = dout + bt * Nh * HeadDim + h * HeadDim; + Float* dinp_i = dinp + bt * Nh * HeadDim + h * HeadDim; + const float rstd_i = rstd[bt * Nh + h]; + + // V heads + if (h >= Nq + Nkv) { + if (abs_max_ptr) { + for (int c = lane * x128::size; c < HeadDim; c += GroupSize * x128::size) { + x128 in_data = x128::load_cs(dout_i + c); + for (int k = 0; k < x128::size; k++) { + thread_abs_max = fmaxf(thread_abs_max, fabsf(in_data[k])); + } + if (dout_i != dinp_i) { + in_data.store(dinp_i + c); + } + } + } else if (dout_i != dinp_i) { + for (int c = lane * x128::size; c < HeadDim; c += GroupSize * x128::size) { + x128 in_data = x128::load_cs(dout_i + c); + in_data.store(dinp_i + c); + } + } + continue; + } + + // QK heads + float sum_xow = 0.0f; + for (int i = lane * x128::size; i < HeadDim; i += GroupSize * x128::size) { + x128 o = x128::load(dout_i + i); + x128 x = x128::load(inp_i + i); + x128 w = x128::load(s_wgt + i); + s_in[i / x128::size] = x; + for (int k = 0; k < x128::size; k++) { + sum_xow += (float)w[k] * (float)o[k] * (float)x[k]; + } + } + + sum_xow = reduce_group_sum(sum_xow, active_mask); + const float xow_norm = sum_xow / HeadDim * rstd_i; + + for (int i = lane * x128::size; i < HeadDim; i += GroupSize * x128::size) { + x128 o = x128::load_cs(dout_i + i); + x128 x = s_in[i / x128::size]; + x128 w = x128::load(s_wgt + i); + x128 dx = x128::zeros(); + + fvec dw = fvec::load(s_d_wgt + i); + for (int k = 0; k < x128::size; k++) { + float xn = (float)x[k] * rstd_i; + dw[k] += xn * (float)o[k]; + float dx_k = ((float)o[k] * (float)w[k] - xn * xow_norm) * rstd_i + (float)dx[k]; + thread_abs_max = fmaxf(thread_abs_max, fabsf(dx_k)); + dx[k] = static_cast(dx_k); + } + + dx.store(dinp_i + i); + + // Cache per-warp partial dweight in shared memory + dw.store(s_d_wgt_base + group * HeadDim + i); + } + } + + __syncthreads(); + // reduce across the block + if (threadIdx.x < 32) { + float* scratch_dweight = reinterpret_cast(scratch) + HeadDim * (blockIdx.x + blockIdx.y * gridDim.x); + for (int i = threadIdx.x * f128::size; i < HeadDim; i += 32 * f128::size) { + f128 accumulated = f128::zeros(); + for (int g = 0; g < num_groups; ++g) { + f128 dw_128 = f128::load(s_d_wgt_base + g * HeadDim + i); + for (int k = 0; k < f128::size; k++) { + accumulated[k] += dw_128[k]; + } + } + accumulated.store(scratch_dweight + i); + } + } + + handle_absmax_reduction(abs_max_ptr, &block_abs_max, thread_abs_max); +} + template void qk_norm_forward_imp(Float* out, float* r_rms, const Float* inp, From add1226f9d03354f51c5fb18335d06918bc80b54 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 26 May 2026 02:43:09 +0200 Subject: [PATCH 08/19] fused-qk+rope backward kernel V1 --- src/kernels/qk_norm.cu | 294 +++++++++++++++++++++++++++++++---------- 1 file changed, 222 insertions(+), 72 deletions(-) diff --git a/src/kernels/qk_norm.cu b/src/kernels/qk_norm.cu index 91e5c08..5b5c186 100644 --- a/src/kernels/qk_norm.cu +++ b/src/kernels/qk_norm.cu @@ -436,54 +436,49 @@ qk_norm_and_rope_backward_kernel(Float* dinp, std::byte* scratch, const float* rstd, const FloatFreq* freqs_cis, float* abs_max_ptr, int BT, int T, int Nq, int Nkv, int HeadDim) { - // formulas: - // rms = sqrt(sum x_j^2 / C + eps) - // y_i = w_i x_i / rms - // - // o_i := dL/dy_i - // dw_i = sum (x_i o_i / rms) - // xow := sum x_i o_i w_i - // dy_j/drms = - w_j x_j / rms² - // drms/dx_j = x_j/(C rms) - // dx_i = dL/dy_j (dy_j/dx_i + dy_j/drms drms/dx_i) - // = o_i w_i / rms - sum_j o_j w_j x_j / rms² x_i/(C rms) - // = o_i w_i / rms - x_i (xow/C) / rms³ - // strategy: y-dimension of block indicates which head is operated on - // x-dimension goes over tokens in groups of eight threads - using x128 = GenericVector; - using f128 = GenericVector; - using fvec = GenericVector; + static_assert(sizeof(Float) == sizeof(FloatFreq), + "Float and FloatFreq must have the same size"); - const int lane = threadIdx.x % GroupSize; - const int group = threadIdx.x / GroupSize; - const int num_groups = blockDim.x / GroupSize; - const int h = blockIdx.y; - const int Nh = Nq + 2 * Nkv; + using x64 = GenericVector; + using x128 = GenericVector; + using f128 = GenericVector; + using freq128 = GenericVector; + using fvec64 = GenericVector; - // load weights into shared memory - // do this before we allow any threads to exit! + const int lane = threadIdx.x % GroupSize; + const int group = threadIdx.x / GroupSize; + const int num_groups = blockDim.x / GroupSize; + const int h = blockIdx.y; + const int Nh = Nq + 2 * Nkv; + const int HDh = HeadDim / 2; + + // Shared memory layout (matches non-fused qk_norm_backward_kernel): + // s_wgt: HeadDim * sizeof(Float) + // s_d_wgt_base: HeadDim * num_groups * sizeof(float) + // s_in_base: HeadDim * num_groups * sizeof(Float) (caches x per group) extern __shared__ int4 smem[]; __shared__ float block_abs_max; float thread_abs_max = 0.f; - Float* s_wgt = reinterpret_cast(smem); + Float* s_wgt = reinterpret_cast(smem); float* s_d_wgt_base = reinterpret_cast(s_wgt + HeadDim); - float* s_d_wgt = s_d_wgt_base + HeadDim * group; - // Per-group cache for input x, to avoid the second global load. - x128* s_in_base = reinterpret_cast(s_d_wgt_base + HeadDim * num_groups); - x128* s_in = s_in_base + group * (HeadDim / x128::size); + float* s_d_wgt = s_d_wgt_base + HeadDim * group; + x64* s_in_base = reinterpret_cast(s_d_wgt_base + HeadDim * num_groups); + x64* s_in = s_in_base + group * (HeadDim / x64::size); - for(int i = threadIdx.x * x128::size; i < HeadDim; i += blockDim.x * x128::size) { - if(h < Nq) { + // Load weights for this head (QK only; V leaves s_wgt undefined and untouched). + for (int i = threadIdx.x * x128::size; i < HeadDim; i += blockDim.x * x128::size) { + if (h < Nq) { x128::load(q_wgt + i).store(s_wgt + i); - } else if (h < Nq + Nkv){ + } else if (h < Nq + Nkv) { x128::load(k_wgt + i).store(s_wgt + i); } // v-heads need no loading } - for(int i = threadIdx.x * f128::size; i < HeadDim * num_groups; i += blockDim.x * f128::size) { + // Zero per-group dweight accumulator. + for (int i = threadIdx.x * f128::size; i < HeadDim * num_groups; i += blockDim.x * f128::size) { f128::zeros().store(s_d_wgt_base + i); } @@ -494,30 +489,29 @@ qk_norm_and_rope_backward_kernel(Float* dinp, std::byte* scratch, __syncthreads(); const int groups_in_grid = num_groups * gridDim.x; - const int start_idx = blockIdx.x * (blockDim.x / GroupSize) + group; + const int start_idx = blockIdx.x * (blockDim.x / GroupSize) + group; for (int bt = start_idx; ; bt += groups_in_grid) { - bool valid = bt < BT; + bool valid = bt < BT; unsigned int active_mask = __ballot_sync(0xffffffffu, valid); - bool all_finished = !__any_sync(0xffffffffu, valid); - if (all_finished) - break; - if (!valid) - continue; + bool all_finished = !__any_sync(0xffffffffu, valid); + if (all_finished) break; + if (!valid) continue; - // adjusted pointers to current token - const Float* inp_i = inp + bt * Nh * HeadDim + h * HeadDim; + const int t = bt % T; + + const Float* inp_i = inp + bt * Nh * HeadDim + h * HeadDim; const Float* dout_i = dout + bt * Nh * HeadDim + h * HeadDim; - Float* dinp_i = dinp + bt * Nh * HeadDim + h * HeadDim; - const float rstd_i = rstd[bt * Nh + h]; + Float* dinp_i = dinp + bt * Nh * HeadDim + h * HeadDim; + const float rstd_i = rstd[bt * Nh + h]; - // V heads + // -------- V heads: pass-through (identical to non-fused) -------- if (h >= Nq + Nkv) { if (abs_max_ptr) { for (int c = lane * x128::size; c < HeadDim; c += GroupSize * x128::size) { x128 in_data = x128::load_cs(dout_i + c); for (int k = 0; k < x128::size; k++) { - thread_abs_max = fmaxf(thread_abs_max, fabsf(in_data[k])); + thread_abs_max = fmaxf(thread_abs_max, fabsf((float)in_data[k])); } if (dout_i != dinp_i) { in_data.store(dinp_i + c); @@ -532,51 +526,104 @@ qk_norm_and_rope_backward_kernel(Float* dinp, std::byte* scratch, continue; } - // QK heads + // -------- QK heads -------- + // + // load_and_derotate(c): + // Loads the two halves of dout at (c, c + HD/2), derotates by -theta_{t,c}, + // returns the two halves as fp32 vectors. Equivalent to rope-backward fused + // into the load. Called once per (token, channel-pair) per pass; the rotation + // itself is cheap (two muls per element + one freq128 load). + auto load_and_derotate = [=] (int c, fvec64& o_r, fvec64& o_i) { + x64 dy_r = x64::load(dout_i + c); + x64 dy_i = x64::load(dout_i + c + HDh); + freq128 f = freq128::load_ldg(freqs_cis + t * HeadDim + 2 * c); + #pragma unroll + for (int k = 0; k < x64::size; ++k) { + float cs = (float)f[2 * k]; + float sn = (float)f[2 * k + 1]; + float r = (float)dy_r[k]; + float im = (float)dy_i[k]; + // rope^T: rotate by -theta <=> sin negated relative to forward + o_r[k] = r * cs + im * sn; + o_i[k] = -r * sn + im * cs; + } + }; + + // Phase A: sum_xow reduction. Also primes s_in with x. float sum_xow = 0.0f; - for (int i = lane * x128::size; i < HeadDim; i += GroupSize * x128::size) { - x128 o = x128::load(dout_i + i); - x128 x = x128::load(inp_i + i); - x128 w = x128::load(s_wgt + i); - s_in[i / x128::size] = x; - for (int k = 0; k < x128::size; k++) { - sum_xow += (float)w[k] * (float)o[k] * (float)x[k]; + for (int c = lane * x64::size; c < HDh; c += GroupSize * x64::size) { + fvec64 o_r, o_i; + load_and_derotate(c, o_r, o_i); + + x64 x_r = x64::load(inp_i + c); + x64 x_i = x64::load(inp_i + c + HDh); + s_in[c / x64::size] = x_r; + s_in[c / x64::size + HDh / x64::size] = x_i; + + x64 w_r = x64::load(s_wgt + c); + x64 w_i = x64::load(s_wgt + c + HDh); + + #pragma unroll + for (int k = 0; k < x64::size; ++k) { + sum_xow += (float)w_r[k] * o_r[k] * (float)x_r[k]; + sum_xow += (float)w_i[k] * o_i[k] * (float)x_i[k]; } } sum_xow = reduce_group_sum(sum_xow, active_mask); const float xow_norm = sum_xow / HeadDim * rstd_i; - for (int i = lane * x128::size; i < HeadDim; i += GroupSize * x128::size) { - x128 o = x128::load_cs(dout_i + i); - x128 x = s_in[i / x128::size]; - x128 w = x128::load(s_wgt + i); - x128 dx = x128::zeros(); + // Phase B: per-element dx, dw. Rotates dout again (matches design choice; + // freq128 is small and cache-friendly, so the second freqs load is cheap). + for (int c = lane * x64::size; c < HDh; c += GroupSize * x64::size) { + fvec64 o_r, o_i; + load_and_derotate(c, o_r, o_i); - fvec dw = fvec::load(s_d_wgt + i); - for (int k = 0; k < x128::size; k++) { - float xn = (float)x[k] * rstd_i; - dw[k] += xn * (float)o[k]; - float dx_k = ((float)o[k] * (float)w[k] - xn * xow_norm) * rstd_i + (float)dx[k]; - thread_abs_max = fmaxf(thread_abs_max, fabsf(dx_k)); - dx[k] = static_cast(dx_k); - } + x64 x_r = s_in[c / x64::size]; + x64 x_i = s_in[c / x64::size + HDh / x64::size]; + x64 w_r = x64::load(s_wgt + c); + x64 w_i = x64::load(s_wgt + c + HDh); - dx.store(dinp_i + i); + fvec64 dw_r = fvec64::load(s_d_wgt + c); + fvec64 dw_i = fvec64::load(s_d_wgt + c + HDh); + x64 dx_r, dx_i; - // Cache per-warp partial dweight in shared memory - dw.store(s_d_wgt_base + group * HeadDim + i); + #pragma unroll + for (int k = 0; k < x64::size; ++k) { + float xn_r = (float)x_r[k] * rstd_i; + float xn_i = (float)x_i[k] * rstd_i; + dw_r[k] += xn_r * o_r[k]; + dw_i[k] += xn_i * o_i[k]; + + float dxk_r = (o_r[k] * (float)w_r[k] - xn_r * xow_norm) * rstd_i; + float dxk_i = (o_i[k] * (float)w_i[k] - xn_i * xow_norm) * rstd_i; + + thread_abs_max = fmaxf(thread_abs_max, fabsf(dxk_r)); + thread_abs_max = fmaxf(thread_abs_max, fabsf(dxk_i)); + + dx_r[k] = (Float)dxk_r; + dx_i[k] = (Float)dxk_i; + } + + dx_r.store(dinp_i + c); + dx_i.store(dinp_i + c + HDh); + dw_r.store(s_d_wgt + c); + dw_i.store(s_d_wgt + c + HDh); } } __syncthreads(); - // reduce across the block + + // Reduce per-group dweight to per-block scratch (identical to non-fused; + // f128-stride reads are independent of the per-group write stride). if (threadIdx.x < 32) { - float* scratch_dweight = reinterpret_cast(scratch) + HeadDim * (blockIdx.x + blockIdx.y * gridDim.x); + float* scratch_dweight = reinterpret_cast(scratch) + + HeadDim * (blockIdx.x + blockIdx.y * gridDim.x); for (int i = threadIdx.x * f128::size; i < HeadDim; i += 32 * f128::size) { f128 accumulated = f128::zeros(); for (int g = 0; g < num_groups; ++g) { f128 dw_128 = f128::load(s_d_wgt_base + g * HeadDim + i); + #pragma unroll for (int k = 0; k < f128::size; k++) { accumulated[k] += dw_128[k]; } @@ -733,6 +780,83 @@ void qk_norm_backward_imp(Float* dinp, Float* dq_wgt, Float* dk_wgt, std::byte* CUDA_CHECK(cudaGetLastError()); } + +// ------------------------------------------------------------ +// Host-side launch wrapper + explicit instantiations. +// Scratch sizing is identical to non-fused qk_norm_backward; we reuse +// qk_norm_backward_smem() and qk_norm_backward_x_blocks(). +// ------------------------------------------------------------ + +template +static int qk_norm_and_rope_backward_x_blocks(int Nq, int Nkv, int HeadDim, + const cudaDeviceProp& dp) { + const int Nh = Nq + 2 * Nkv; + const size_t smem = qk_norm_backward_smem(HeadDim); + int blocks_per_sm; + // FloatFreq parametrization doesn't affect occupancy: same kernel size, + // same smem, same launch bounds. Use the bf16-freq variant as representative. + using FloatFreq = std::conditional_t, float, half>; + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks_per_sm, + qk_norm_and_rope_backward_kernel, + 512, smem)); + int total_blocks = blocks_per_sm * dp.multiProcessorCount; + if (total_blocks < Nh) { + return 1; + } + return total_blocks / Nh; +} + +std::size_t qk_norm_and_rope_backward_scratch_size(int Nq, int Nkv, int HeadDim, + ETensorDType dtype, + const cudaDeviceProp& dp) { + const int Nh = Nq + 2 * Nkv; + int x_blocks; + switch (dtype) { + case ETensorDType::FP32: + x_blocks = qk_norm_and_rope_backward_x_blocks(Nq, Nkv, HeadDim, dp); + break; + case ETensorDType::BF16: + x_blocks = qk_norm_and_rope_backward_x_blocks(Nq, Nkv, HeadDim, dp); + break; + default: + throw std::invalid_argument("Unsupported dtype"); + } + return (size_t)HeadDim * x_blocks * Nh * sizeof(float); +} + +template +void qk_norm_and_rope_backward_imp(Float* dinp, Float* dq_wgt, Float* dk_wgt, + std::byte* scratch, + const Float* dout, const Float* inp, + const Float* q_wgt, const Float* k_wgt, + const float* rstd, const FloatFreq* freqs_cis, + float* abs_max_ptr, + int B, int T, int Nq, int Nkv, int HeadDim, + const cudaDeviceProp& dp, cudaStream_t stream) { + if (HeadDim % 16 != 0) { + throw std::invalid_argument("HeadDim must be a multiple of 16"); + } + + constexpr int block_size = 512; + const int BT = B * T; + const int Nh = Nq + 2 * Nkv; + const size_t smem = qk_norm_backward_smem(HeadDim); + const int x_blocks = qk_norm_and_rope_backward_x_blocks(Nq, Nkv, HeadDim, dp); + + dim3 grid(x_blocks, Nh); + qk_norm_and_rope_backward_kernel<<>>( + dinp, scratch, dout, inp, q_wgt, k_wgt, rstd, freqs_cis, abs_max_ptr, + BT, T, Nq, Nkv, HeadDim); + CUDA_CHECK(cudaGetLastError()); + + // Final reduction across blocks_x into dq_wgt / dk_wgt -- reuse non-fused reducer. + dim3 reduce_grid(1, 2); + qk_norm_backward_reduce_kernel<<>>( + dq_wgt, dk_wgt, scratch, x_blocks, Nq, Nkv, HeadDim); + CUDA_CHECK(cudaGetLastError()); +} + // Explicit instantiations void qk_norm_forward(float* out, float* r_rms, const float* inp, const float* q_wgt, const float* k_wgt, @@ -788,3 +912,29 @@ void qk_norm_backward(nv_bfloat16* dinp, nv_bfloat16* dq_wgt, nv_bfloat16* dk_wg const cudaDeviceProp& dp, cudaStream_t stream) { qk_norm_backward_imp(dinp, dq_wgt, dk_wgt, scratch, dout, inp, q_wgt, k_wgt, rstd, abs_max_ptr, BT, Nq, Nkv, HeadDim, dp, stream); } + +void qk_norm_and_rope_backward(float* dinp, float* dq_wgt, float* dk_wgt, + std::byte* scratch, + const float* dout, const float* inp, + const float* q_wgt, const float* k_wgt, + const float* rstd, const float* freqs_cis, + float* abs_max_ptr, + int B, int T, int Nq, int Nkv, int HeadDim, + const cudaDeviceProp& dp, cudaStream_t stream) { + qk_norm_and_rope_backward_imp( + dinp, dq_wgt, dk_wgt, scratch, dout, inp, q_wgt, k_wgt, rstd, freqs_cis, + abs_max_ptr, B, T, Nq, Nkv, HeadDim, dp, stream); +} + +void qk_norm_and_rope_backward(nv_bfloat16* dinp, nv_bfloat16* dq_wgt, nv_bfloat16* dk_wgt, + std::byte* scratch, + const nv_bfloat16* dout, const nv_bfloat16* inp, + const nv_bfloat16* q_wgt, const nv_bfloat16* k_wgt, + const float* rstd, const half* freqs_cis, + float* abs_max_ptr, + int B, int T, int Nq, int Nkv, int HeadDim, + const cudaDeviceProp& dp, cudaStream_t stream) { + qk_norm_and_rope_backward_imp( + dinp, dq_wgt, dk_wgt, scratch, dout, inp, q_wgt, k_wgt, rstd, freqs_cis, + abs_max_ptr, B, T, Nq, Nkv, HeadDim, dp, stream); +} From f74661e85f65c76bb13937b787b759d986277165 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 8 Jun 2026 03:22:32 +0200 Subject: [PATCH 09/19] fix launch config --- src/kernels/qk_norm.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/src/kernels/qk_norm.cu b/src/kernels/qk_norm.cu index 5b5c186..bd43dde 100644 --- a/src/kernels/qk_norm.cu +++ b/src/kernels/qk_norm.cu @@ -715,6 +715,7 @@ static int qk_norm_backward_x_blocks(int Nq, int Nkv, int HeadDim, const cudaDeviceProp& dp) { const int Nh = Nq + 2 * Nkv; const size_t smem = qk_norm_backward_smem(HeadDim); + CUDA_CHECK(cudaFuncSetAttribute(qk_norm_backward_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); int blocks_per_sm; CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( &blocks_per_sm, qk_norm_backward_kernel, 512, smem)); From c37cae53d9fd9b8d7643e36b5d01b7970c0ae680 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 9 Jun 2026 01:34:26 +0200 Subject: [PATCH 10/19] implement using cooperative groups --- src/kernels/qk_norm.cu | 74 +++++++++++++++++++----------------------- 1 file changed, 34 insertions(+), 40 deletions(-) diff --git a/src/kernels/qk_norm.cu b/src/kernels/qk_norm.cu index bd43dde..2e7106b 100644 --- a/src/kernels/qk_norm.cu +++ b/src/kernels/qk_norm.cu @@ -6,21 +6,18 @@ #include "utilities/utils.h" #include "utilities/dtype.h" #include "utilities/vec.cuh" +#include +#include -constexpr const int GroupSize = 8; +namespace cg = cooperative_groups; -__device__ float reduce_group_sum(float acc, unsigned int active_mask) { - for (int offset = GroupSize / 2; offset > 0; offset >>= 1) { - acc += __shfl_xor_sync(active_mask, acc, offset); - } - return acc; -} +constexpr const int GroupSize = 8; template struct QKHelpers { using x128 = GenericVector; - static __device__ __forceinline__ float norm_head(Float* out, x128* s_in, const x128* wgt_src, const Float* inp, int HeadDim, float epsilon, int active_mask) { + static __device__ __forceinline__ float norm_head(Float* out, x128* s_in, const x128* wgt_src, const Float* inp, int HeadDim, float epsilon, const cg::thread_block_tile<8, cg::thread_block>& tile) { using x128 = GenericVector; const int lane = threadIdx.x % GroupSize; @@ -35,7 +32,7 @@ struct QKHelpers { } } - acc = reduce_group_sum(acc, active_mask) / HeadDim; + acc = cg::reduce(tile, acc, cg::plus{}) / HeadDim; float s = rsqrtf(acc + epsilon); for(int c = lane * x128::size; c < HeadDim; c += GroupSize * x128::size) { @@ -59,10 +56,12 @@ struct QKHelpers { template __global__ void qk_norm_forward_simple_kernel(Float* out, float* r_rms, const Float* inp, const Float* q_wgt, const Float* k_wgt, float epsilon, int BT, int Nq, int Nkv, int HeadDim) { + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<8, cg::thread_block> tile = cg::tiled_partition(block); using x128 = GenericVector; - const int lane = threadIdx.x % GroupSize; - const int group = threadIdx.x / GroupSize; + const int lane = tile.thread_rank(); + const int group = tile.meta_group_rank(); // load weights into shared memory // do this before we allow any threads to exit! @@ -74,7 +73,7 @@ __global__ void qk_norm_forward_simple_kernel(Float* out, float* r_rms, const Fl x128* s_k_wgt = reinterpret_cast(smem) + (HeadDim / x128::size); x128* s_in = reinterpret_cast(s_k_wgt) + (HeadDim / x128::size) + group * (HeadDim / x128::size); - for(int i = threadIdx.x * x128::size; i < HeadDim; i += blockDim.x * x128::size) { + for(int i = threadIdx.x * x128::size; i < HeadDim; i += blockDim.x * x128::size) { s_q_wgt[i/x128::size] = x128::load(q_wgt + i); s_k_wgt[i/x128::size] = x128::load(k_wgt + i); } @@ -106,7 +105,7 @@ __global__ void qk_norm_forward_simple_kernel(Float* out, float* r_rms, const Fl return; } - float s = QKHelpers::norm_head(out, s_in, wgt_src, inp, HeadDim, epsilon, 0xffffffff); + float s = QKHelpers::norm_head(out, s_in, wgt_src, inp, HeadDim, epsilon, tile); // store the rms, no need to cache it if(lane == 0 && r_rms != nullptr) { @@ -139,8 +138,11 @@ qk_norm_backward_kernel(Float* dinp, std::byte* scratch, using f128 = GenericVector; using fvec = GenericVector; - const int lane = threadIdx.x % GroupSize; - const int group = threadIdx.x / GroupSize; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<8, cg::thread_block> tile = cg::tiled_partition(block); + + const int lane = tile.thread_rank(); + const int group = tile.meta_group_rank(); const int num_groups = blockDim.x / GroupSize; const int h = blockIdx.y; const int Nh = Nq + 2 * Nkv; @@ -181,15 +183,7 @@ qk_norm_backward_kernel(Float* dinp, std::byte* scratch, const int groups_in_grid = num_groups * gridDim.x; const int start_idx = blockIdx.x * (blockDim.x / GroupSize) + group; - for (int bt = start_idx; ; bt += groups_in_grid) { - bool valid = bt < BT; - unsigned int active_mask = __ballot_sync(0xffffffffu, valid); - bool all_finished = !__any_sync(0xffffffffu, valid); - if (all_finished) - break; - if (!valid) - continue; - + for (int bt = start_idx; bt < BT; bt += groups_in_grid) { // adjusted pointers to current token const Float* inp_i = inp + bt * Nh * HeadDim + h * HeadDim; const Float* dout_i = dout + bt * Nh * HeadDim + h * HeadDim; @@ -229,7 +223,7 @@ qk_norm_backward_kernel(Float* dinp, std::byte* scratch, } } - sum_xow = reduce_group_sum(sum_xow, active_mask); + sum_xow = cg::reduce(tile, sum_xow, cg::plus{}); const float xow_norm = sum_xow / HeadDim * rstd_i; for (int i = lane * x128::size; i < HeadDim; i += GroupSize * x128::size) { @@ -242,7 +236,7 @@ qk_norm_backward_kernel(Float* dinp, std::byte* scratch, for (int k = 0; k < x128::size; k++) { float xn = (float)x[k] * rstd_i; dw[k] += xn * (float)o[k]; - float dx_k = ((float)o[k] * (float)w[k] - xn * xow_norm) * rstd_i + (float)dx[k]; + float dx_k = ((float)o[k] * (float)w[k] - xn * xow_norm) * rstd_i;// + (float)dx[k]; thread_abs_max = fmaxf(thread_abs_max, fabsf(dx_k)); dx[k] = static_cast(dx_k); } @@ -323,8 +317,11 @@ __global__ void qk_norm_and_rope_fwd_kernel(Float* out, float* r_rms, float* abs using x128 = GenericVector; using freq128 = GenericVector; - const int lane = threadIdx.x % GroupSize; - const int group = threadIdx.x / GroupSize; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<8, cg::thread_block> tile = cg::tiled_partition(block); + + const int lane = tile.thread_rank(); + const int group = tile.meta_group_rank(); __shared__ float block_abs_max; if (abs_max_ptr && threadIdx.x == 0) { block_abs_max = 0.f; @@ -369,7 +366,6 @@ __global__ void qk_norm_and_rope_fwd_kernel(Float* out, float* r_rms, float* abs is_value_head = true; } - unsigned int active_mask = __ballot_sync(0xffffffffu, !is_value_head); if (is_value_head) { // ---- value head: pass-through, but still contribute to abs-max ---- for (int c = lane * x128::size; c < HeadDim; c += GroupSize * x128::size) { @@ -382,8 +378,8 @@ __global__ void qk_norm_and_rope_fwd_kernel(Float* out, float* r_rms, float* abs } else { // ---- Q or K head: norm + scale + RoPE ---- - float s = QKHelpers::norm_head(&s_in[0][0], s_in, wgt_src, inp_h, HeadDim, epsilon, active_mask); - __syncwarp(active_mask); + float s = QKHelpers::norm_head(&s_in[0][0], s_in, wgt_src, inp_h, HeadDim, epsilon, tile); + tile.sync(); int head_dim_half = HeadDim / 2; using x64 = GenericVector; @@ -445,8 +441,11 @@ qk_norm_and_rope_backward_kernel(Float* dinp, std::byte* scratch, using freq128 = GenericVector; using fvec64 = GenericVector; - const int lane = threadIdx.x % GroupSize; - const int group = threadIdx.x / GroupSize; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<8, cg::thread_block> tile = cg::tiled_partition(block); + + const int lane = tile.thread_rank(); + const int group = tile.meta_group_rank(); const int num_groups = blockDim.x / GroupSize; const int h = blockIdx.y; const int Nh = Nq + 2 * Nkv; @@ -491,12 +490,7 @@ qk_norm_and_rope_backward_kernel(Float* dinp, std::byte* scratch, const int groups_in_grid = num_groups * gridDim.x; const int start_idx = blockIdx.x * (blockDim.x / GroupSize) + group; - for (int bt = start_idx; ; bt += groups_in_grid) { - bool valid = bt < BT; - unsigned int active_mask = __ballot_sync(0xffffffffu, valid); - bool all_finished = !__any_sync(0xffffffffu, valid); - if (all_finished) break; - if (!valid) continue; + for (int bt = start_idx; bt < BT; bt += groups_in_grid) { const int t = bt % T; @@ -570,7 +564,7 @@ qk_norm_and_rope_backward_kernel(Float* dinp, std::byte* scratch, } } - sum_xow = reduce_group_sum(sum_xow, active_mask); + sum_xow = cg::reduce(tile, sum_xow, cg::plus{}); const float xow_norm = sum_xow / HeadDim * rstd_i; // Phase B: per-element dx, dw. Rotates dout again (matches design choice; From fb3e94a942348c6402fdc939d7f7b5df657c08f2 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 10 Jun 2026 08:34:45 +0200 Subject: [PATCH 11/19] backward --- src/binding/kernel_binding.cpp | 46 +++++++++++++++++++ src/binding/python/kernels.py | 14 ++++++ src/kernels/kernels.cpp | 24 ++++++++++ src/kernels/kernels.h | 28 ++++++++++++ test/test_kernels_qk_norm.py | 81 ++++++++++++++++++++++++++++++++++ 5 files changed, 193 insertions(+) diff --git a/src/binding/kernel_binding.cpp b/src/binding/kernel_binding.cpp index c3cfcf8..0d40461 100644 --- a/src/binding/kernel_binding.cpp +++ b/src/binding/kernel_binding.cpp @@ -312,6 +312,48 @@ long bind_get_qknorm_backward_scratch_size(int Nq, int Nkv, int HeadDim, std::st return qk_norm_backward_scratch_size(Nq, Nkv, HeadDim, dtype_from_str(dtype), get_device_prop(did)); } +void bind_qk_norm_and_rope_backward(const CudaArray& dinp, const CudaArray& dq_wgt, const CudaArray& dk_wgt, + const CudaArray& scratch, + const CudaArray& dout, const CudaArray& inp, + const CudaArray& q_wgt, const CudaArray& k_wgt, + const CudaArray& rstd, const CudaArray& freqs_cis, + const std::optional& abs_max, + int Nq, int Nkv, + const std::uintptr_t stream) { + NB_CHECK_NDIMS(dinp, 3); + NB_CHECK_NDIMS(dq_wgt, 1); + NB_CHECK_NDIMS(dk_wgt, 1); + NB_CHECK_NDIMS(scratch, 1); + NB_CHECK_NDIMS(dout, 3); + NB_CHECK_NDIMS(inp, 3); + NB_CHECK_NDIMS(q_wgt, 1); + NB_CHECK_NDIMS(k_wgt, 1); + NB_CHECK_NDIMS(rstd, 3); + NB_CHECK_NDIMS(freqs_cis, 2); + + auto dp = get_device_prop(inp.device_id()); + const long B = get_dimension_checked({dinp.shape(0), dout.shape(0), inp.shape(0), rstd.shape(0)}, "B"); + const long T = get_dimension_checked({dinp.shape(1), dout.shape(1), inp.shape(1), rstd.shape(1), freqs_cis.shape(0)}, "T"); + const long HeadDim = get_dimension_checked({q_wgt.shape(0), k_wgt.shape(0), dq_wgt.shape(0), dk_wgt.shape(0), freqs_cis.shape(1)}, "HeadDim"); + (void)get_dimension_checked({dinp.shape(2), dout.shape(2), inp.shape(2), (std::size_t)((Nq + 2 * Nkv) * HeadDim)}, "NTotal*HeadDim"); + (void)get_dimension_checked({rstd.shape(2), (std::size_t)(Nq + 2 * Nkv)}, "Nq+2*Nkv"); + + Tensor dinp_t = to_tensor(dinp); + Tensor dq_wgt_t = to_tensor(dq_wgt); + Tensor dk_wgt_t = to_tensor(dk_wgt); + Tensor scratch_t = to_tensor(scratch); + qk_norm_and_rope_backward(dinp_t, dq_wgt_t, dk_wgt_t, scratch_t, + to_tensor(dout), to_tensor(inp), to_tensor(q_wgt), to_tensor(k_wgt), to_tensor(rstd), to_tensor(freqs_cis), + get_abs_max_ptr(abs_max), B, T, Nq, Nkv, HeadDim, + dp, as_stream(stream)); +} + +long bind_get_qknorm_and_rope_backward_scratch_size(int Nq, int Nkv, int HeadDim, std::string dtype) { + int did; + CUDA_CHECK(cudaGetDevice(&did)); + return qk_norm_and_rope_backward_scratch_size(Nq, Nkv, HeadDim, dtype_from_str(dtype), get_device_prop(did)); +} + // ---- SwiGLU ---- void bind_swiglu_forward(const CudaArray& out, const CudaArray& inp, const std::optional& abs_max, const std::uintptr_t stream) { @@ -633,6 +675,10 @@ void register_kernels(nanobind::module_& m) { nb::arg("scratch"), nb::arg("dout"), nb::arg("inp"), nb::arg("q_wgt"), nb::arg("k_wgt"), nb::arg("rstd"), nb::arg("abs_max") = std::nullopt, nb::arg("Nq"), nb::arg("Nkv"), nb::arg("stream") = 0); m.def("get_qknorm_backward_scratch_size", &bind_get_qknorm_backward_scratch_size, nb::arg("Nq"), nb::arg("Nkv"), nb::arg("HeadDim"), nb::arg("dtype")); + m.def("qk_norm_and_rope_backward", &bind_qk_norm_and_rope_backward, nb::arg("dinp"), nb::arg("dq_wgt"), nb::arg("dk_wgt"), + nb::arg("scratch"), nb::arg("dout"), nb::arg("inp"), nb::arg("q_wgt"), nb::arg("k_wgt"), + nb::arg("rstd"), nb::arg("freqs_cis"), nb::arg("abs_max") = std::nullopt, nb::arg("Nq"), nb::arg("Nkv"), nb::arg("stream") = 0); + m.def("get_qknorm_and_rope_backward_scratch_size", &bind_get_qknorm_and_rope_backward_scratch_size, nb::arg("Nq"), nb::arg("Nkv"), nb::arg("HeadDim"), nb::arg("dtype")); // SwiGLU m.def("swiglu_forward", &bind_swiglu_forward, nb::arg("out"), nb::arg("inp"), nb::arg("absmax") = std::nullopt, nb::arg("stream") = 0); diff --git a/src/binding/python/kernels.py b/src/binding/python/kernels.py index 9059905..53212ba 100644 --- a/src/binding/python/kernels.py +++ b/src/binding/python/kernels.py @@ -84,6 +84,18 @@ def qk_norm_backward(dinp: torch.Tensor, dq_wgt: torch.Tensor, dk_wgt: torch.Ten def get_qknorm_backward_scratch_size(Nq: int, Nkv: int, HeadDim: int, dtype: torch.dtype): return _pyllmq.get_qknorm_backward_scratch_size(Nq, Nkv, HeadDim, _TORCH_TO_TYPE_NAME[dtype]) +@torch.library.custom_op("llmq::qk_norm_and_rope_backward", mutates_args=("dinp", "dq_wgt", "dk_wgt", "scratch", "abs_max")) +def qk_norm_and_rope_backward(dinp: torch.Tensor, dq_wgt: torch.Tensor, dk_wgt: torch.Tensor, + scratch: torch.Tensor, dout: torch.Tensor, inp: torch.Tensor, + q_wgt: torch.Tensor, k_wgt: torch.Tensor, rstd: torch.Tensor, + freqs_cis: torch.Tensor, abs_max: torch.Tensor | None, + Nq: int, Nkv: int, stream: int = 0) -> None: + _pyllmq.qk_norm_and_rope_backward(dinp, dq_wgt, dk_wgt, scratch, dout, inp, q_wgt, k_wgt, + rstd, freqs_cis, abs_max, Nq, Nkv, stream) + +def get_qknorm_and_rope_backward_scratch_size(Nq: int, Nkv: int, HeadDim: int, dtype: torch.dtype): + return _pyllmq.get_qknorm_and_rope_backward_scratch_size(Nq, Nkv, HeadDim, _TORCH_TO_TYPE_NAME[dtype]) + # SwiGLU @torch.library.custom_op("llmq::swiglu_forward", mutates_args=("out", "absmax")) def swiglu_forward(out: torch.Tensor, inp: torch.Tensor, absmax: torch.Tensor | None, stream: int = 0) -> None: @@ -236,6 +248,8 @@ def adamw_update(params: torch.Tensor, grads: torch.Tensor, m: torch.Tensor, v: 'encoder_forward', 'rmsnorm_forward', 'rmsnorm_backward', 'fused_residual_rmsnorm_forward', 'rope_forward', 'rope_backward', 'qk_norm_forward', 'qk_norm_backward', + 'qk_norm_and_rope_forward', 'qk_norm_and_rope_backward', + 'get_qknorm_and_rope_backward_scratch_size', 'swiglu_forward', 'swiglu_forward_quant', 'swiglu_backward', 'attention_forward', 'attention_backward', 'fused_classifier', 'grouped_loss_sum', diff --git a/src/kernels/kernels.cpp b/src/kernels/kernels.cpp index d3e1547..7b867f9 100644 --- a/src/kernels/kernels.cpp +++ b/src/kernels/kernels.cpp @@ -141,6 +141,30 @@ void qk_norm_and_rope_forward(Tensor& out, Tensor& r_rms, float* abs_max_ptr, } } +void qk_norm_and_rope_backward(Tensor& dinp, Tensor& dq_wgt, Tensor& dk_wgt, Tensor& scratch, + const Tensor& dout, const Tensor& inp, + const Tensor& q_wgt, const Tensor& k_wgt, + const Tensor& rstd, const Tensor& freqs_cis, + float* abs_max_ptr, + int B, int T, int Nq, int Nkv, int HeadDim, + const cudaDeviceProp& dp, cudaStream_t stream) { + if (dinp.DType == ETensorDType::FP32) { + qk_norm_and_rope_backward(dinp.get(), dq_wgt.get(), dk_wgt.get(), scratch.get(), + dout.get(), inp.get(), + q_wgt.get(), k_wgt.get(), + rstd.get(), freqs_cis.get(), abs_max_ptr, + B, T, Nq, Nkv, HeadDim, dp, stream); + } else if (dinp.DType == ETensorDType::BF16) { + qk_norm_and_rope_backward(dinp.get(), dq_wgt.get(), dk_wgt.get(), scratch.get(), + dout.get(), inp.get(), + q_wgt.get(), k_wgt.get(), + rstd.get(), freqs_cis.get(), abs_max_ptr, + B, T, Nq, Nkv, HeadDim, dp, stream); + } else { + UNSUPPORTED_DTYPE(dinp, dq_wgt, dk_wgt, dout, inp, q_wgt, k_wgt, rstd, freqs_cis); + } +} + void qk_norm_backward(Tensor& dinp, Tensor& dq_wgt, Tensor& dk_wgt, Tensor& scratch, const Tensor& dout, const Tensor& inp, const Tensor& q_wgt, const Tensor& k_wgt, diff --git a/src/kernels/kernels.h b/src/kernels/kernels.h index 214a9d3..c78e885 100644 --- a/src/kernels/kernels.h +++ b/src/kernels/kernels.h @@ -178,6 +178,34 @@ void qk_norm_backward(Tensor& dinp, Tensor& dq_wgt, Tensor& dk_wgt, Tensor& scra int BT, int Nq, int Nkv, int HeadDim, const cudaDeviceProp& dp, cudaStream_t stream); +std::size_t qk_norm_and_rope_backward_scratch_size(int Nq, int Nkv, int HeadDim, ETensorDType dtype, const cudaDeviceProp& dp); + +void qk_norm_and_rope_backward(float* dinp, float* dq_wgt, float* dk_wgt, + std::byte* scratch, + const float* dout, const float* inp, + const float* q_wgt, const float* k_wgt, + const float* rstd, const float* freqs_cis, + float* abs_max_ptr, + int B, int T, int Nq, int Nkv, int HeadDim, + const cudaDeviceProp& dp, cudaStream_t stream); + +void qk_norm_and_rope_backward(nv_bfloat16* dinp, nv_bfloat16* dq_wgt, nv_bfloat16* dk_wgt, + std::byte* scratch, + const nv_bfloat16* dout, const nv_bfloat16* inp, + const nv_bfloat16* q_wgt, const nv_bfloat16* k_wgt, + const float* rstd, const half* freqs_cis, + float* abs_max_ptr, + int B, int T, int Nq, int Nkv, int HeadDim, + const cudaDeviceProp& dp, cudaStream_t stream); + +void qk_norm_and_rope_backward(Tensor& dinp, Tensor& dq_wgt, Tensor& dk_wgt, Tensor& scratch, + const Tensor& dout, const Tensor& inp, + const Tensor& q_wgt, const Tensor& k_wgt, + const Tensor& rstd, const Tensor& freqs_cis, + float* abs_max_ptr, + int B, int T, int Nq, int Nkv, int HeadDim, + const cudaDeviceProp& dp, cudaStream_t stream); + // swiglu assumes that input is the concatenation of gate and up projection. void swiglu_forward(nv_bfloat16* out, const nv_bfloat16* inp, float* abs_max_ptr, int B, int T, int C, cudaStream_t stream); void swiglu_forward(float* out, const float* inp, float* abs_max_ptr, int B, int T, int C, cudaStream_t stream); diff --git a/test/test_kernels_qk_norm.py b/test/test_kernels_qk_norm.py index 2283515..d4afa8e 100644 --- a/test/test_kernels_qk_norm.py +++ b/test/test_kernels_qk_norm.py @@ -143,6 +143,11 @@ def test_qk_norm_and_rope_forward(B, T, Nq, Nkv, HeadDim, dtype, with_abs_max): ref_abs_max = ref_out.float().abs().max().item() assert abs_max.item() == pytest.approx(ref_abs_max, rel=rtol, abs=atol) + +# ============================================================ +# qk_norm_backward +# ============================================================ + @pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ (1, 4, 2, 1, 16), (2, 8, 4, 2, 64), @@ -199,3 +204,79 @@ def test_qk_norm_backward(B, T, Nq, Nkv, HeadDim, dtype, in_place: bool): assert dinp.float().cpu().numpy() == pytest.approx(ref_dinp.cpu().numpy(), rel=rtol, abs=atol) assert dq_wgt_k.float().cpu().numpy() == pytest.approx(ref_dq_wgt.cpu().numpy(), rel=rtol, abs=atol) assert dk_wgt_k.float().cpu().numpy() == pytest.approx(ref_dk_wgt.cpu().numpy(), rel=rtol, abs=atol) + +@pytest.mark.parametrize("B,T,Nq,Nkv,HeadDim", [ + (1, 4, 2, 1, 16), + (2, 8, 4, 2, 64), + (4, 16, 8, 4, 64), + (1, 1, 4, 2, 128), +]) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("in_place", [True, False]) +def test_qk_norm_and_rope_backward(B, T, Nq, Nkv, HeadDim, dtype, in_place: bool): + torch.manual_seed(0) + device = "cuda" + N = Nq + 2 * Nkv + eps = 1e-6 + + inp = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + q_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + k_wgt = torch.randn((HeadDim,), device=device, dtype=dtype) + dout = torch.randn((B, T, N * HeadDim), device=device, dtype=dtype) + + freq_dtype = torch.float32 if dtype == torch.float32 else torch.float16 + freqs = _make_rope_freqs(T, HeadDim, 1_000_000.0, freq_dtype, device) + + # Reference in fp32: qk-norm then rope, backprop through both. + inp_f = inp.float().detach().requires_grad_(True) + q_wgt_f = q_wgt.float().detach().requires_grad_(True) + k_wgt_f = k_wgt.float().detach().requires_grad_(True) + ref_normed_f, _ = _qk_norm_reference(inp_f, q_wgt_f, k_wgt_f, Nq, Nkv, eps) + ref_out_f = _rope_python( + ref_normed_f.view(B, T, N, HeadDim), freqs, Nq, Nkv + ).view(B, T, N * HeadDim) + ref_out_f.backward(dout.float()) + ref_dinp = inp_f.grad + ref_dq_wgt = q_wgt_f.grad + ref_dk_wgt = k_wgt_f.grad + + # Run the fused forward to populate r_rms (kernel only writes QK slots). + r_rms = torch.empty((B, T, N), device=device, dtype=torch.float32) + r_rms.fill_(1.0) + K.qk_norm_and_rope_forward( + torch.empty_like(inp), r_rms, inp, q_wgt, k_wgt, freqs, None, eps, Nq, Nkv, + ) + + scratch_size = K.get_qknorm_and_rope_backward_scratch_size(Nq, Nkv, HeadDim, inp.dtype) + scratch = torch.zeros((scratch_size,), device=device, dtype=torch.uint8) + dq_wgt_k = torch.zeros_like(q_wgt) + dk_wgt_k = torch.zeros_like(k_wgt) + + v_slice = slice((Nq + Nkv) * HeadDim, N * HeadDim) + + if in_place: + dinp_buf = dout.clone() + v_before = dinp_buf[:, :, v_slice].clone() + K.qk_norm_and_rope_backward( + dinp_buf, dq_wgt_k, dk_wgt_k, scratch, + dinp_buf, inp, q_wgt, k_wgt, r_rms, freqs, + None, Nq, Nkv, + ) + dinp = dinp_buf + # In-place V pass-through: kernel must not write V slots + # (mirrors qk_norm_backward_kernel's `if (dout_i != dinp_i)` guard). + assert torch.equal(dinp[:, :, v_slice], v_before) + else: + dinp = torch.zeros_like(inp) + K.qk_norm_and_rope_backward( + dinp, dq_wgt_k, dk_wgt_k, scratch, + dout, inp, q_wgt, k_wgt, r_rms, freqs, + None, Nq, Nkv, + ) + # Out-of-place V pass-through: dtype copy is bit-exact. + assert torch.equal(dinp[:, :, v_slice], dout[:, :, v_slice]) + + rtol, atol = TOL[dtype] + assert dinp.float().cpu().numpy() == pytest.approx(ref_dinp.cpu().numpy(), rel=rtol, abs=atol) + assert dq_wgt_k.float().cpu().numpy() == pytest.approx(ref_dq_wgt.cpu().numpy(), rel=rtol, abs=atol) + assert dk_wgt_k.float().cpu().numpy() == pytest.approx(ref_dk_wgt.cpu().numpy(), rel=rtol, abs=atol) From 51513bcf839717c411fa3c674919d7727282acd2 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 19 Jul 2026 01:32:24 +0200 Subject: [PATCH 12/19] opt in to large dynamic smem for the fused qk-norm+rope backward The fused kernel additionally caches x in shared memory, which exceeds the default 48 KiB limit for HeadDim=128 (the Qwen3 case); set cudaFuncAttributeMaxDynamicSharedMemorySize like the non-fused launcher already does. Co-Authored-By: Claude Fable 5 --- src/kernels/qk_norm.cu | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/kernels/qk_norm.cu b/src/kernels/qk_norm.cu index 2e7106b..19ff0c5 100644 --- a/src/kernels/qk_norm.cu +++ b/src/kernels/qk_norm.cu @@ -791,6 +791,8 @@ static int qk_norm_and_rope_backward_x_blocks(int Nq, int Nkv, int HeadDim, // FloatFreq parametrization doesn't affect occupancy: same kernel size, // same smem, same launch bounds. Use the bf16-freq variant as representative. using FloatFreq = std::conditional_t, float, half>; + CUDA_CHECK(cudaFuncSetAttribute(qk_norm_and_rope_backward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( &blocks_per_sm, qk_norm_and_rope_backward_kernel, @@ -839,6 +841,11 @@ void qk_norm_and_rope_backward_imp(Float* dinp, Float* dq_wgt, Float* dk_wgt, const size_t smem = qk_norm_backward_smem(HeadDim); const int x_blocks = qk_norm_and_rope_backward_x_blocks(Nq, Nkv, HeadDim, dp); + CUDA_CHECK(cudaFuncSetAttribute( + qk_norm_and_rope_backward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem)); + dim3 grid(x_blocks, Nh); qk_norm_and_rope_backward_kernel<<>>( dinp, scratch, dout, inp, q_wgt, k_wgt, rstd, freqs_cis, abs_max_ptr, From 5bfcfb456a762f4dc0c6951b61ef16494d4b4db3 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 4 Mar 2026 13:45:37 +0100 Subject: [PATCH 13/19] config preparations for Qwen3 --- src/models/llama_model.cpp | 3 ++ src/training/transformer_config.cpp | 76 ++++++++++++++++++++++------- src/training/transformer_config.h | 2 + 3 files changed, 63 insertions(+), 18 deletions(-) diff --git a/src/models/llama_model.cpp b/src/models/llama_model.cpp index 80b366c..3e131f8 100644 --- a/src/models/llama_model.cpp +++ b/src/models/llama_model.cpp @@ -17,6 +17,9 @@ LLamaModel::LLamaModel(TransformerConfig config, const LLamaOptions& options, int rank, int world, const std::shared_ptr& alloc) : Config(config), Options(options), Allocator(alloc ? alloc : std::make_shared()) { + if (Config.UseQKNorm) + throw std::runtime_error("UseQKNorm is not yet supported"); + Parameters = LLamaWeightsManager::create(Config, options, rank, world, *Allocator); } diff --git a/src/training/transformer_config.cpp b/src/training/transformer_config.cpp index a2966d8..f0709f4 100644 --- a/src/training/transformer_config.cpp +++ b/src/training/transformer_config.cpp @@ -27,6 +27,8 @@ TransformerConfig load_transformer_config(const char* file_name, ETensorDType dt arch_id = TransformerConfig::LLAMA; } else if(archs.front() == "Qwen2ForCausalLM") { arch_id = TransformerConfig::QWEN2; + } else if(archs.front() == "Qwen3ForCausalLM") { + arch_id = TransformerConfig::QWEN3; } else { throw std::runtime_error(fmt::format("unknown architecture {}", archs.front())); } @@ -55,6 +57,7 @@ TransformerConfig load_transformer_config(const char* file_name, ETensorDType dt result.RmsNormEps = result.Architecture == TransformerConfig::LLAMA ? 1e-5 : 1e-6; } + result.UseQKNorm = arch_id == TransformerConfig::QWEN3; result.UseQKVBias = arch_id == TransformerConfig::QWEN2; return result; @@ -62,6 +65,8 @@ TransformerConfig load_transformer_config(const char* file_name, ETensorDType dt [[nodiscard]] std::string_view TransformerConfig::model_name() const { switch(Architecture) { + case TransformerConfig::QWEN3: + return "Qwen3"; case TransformerConfig::QWEN2: return "Qwen2"; case TransformerConfig::LLAMA: @@ -78,8 +83,10 @@ void save_transformer_config(const TransformerConfig& config, const char* file_n } std::vector archs; - if(config.Architecture == TransformerConfig::QWEN2) { + if (config.Architecture == TransformerConfig::QWEN2) { archs = {"Qwen2ForCausalLM"}; + } else if(config.Architecture == TransformerConfig::QWEN3) { + archs = {"Qwen3ForCausalLM"}; } else if (config.Architecture == TransformerConfig::LLAMA) { archs = {"LlamaForCausalLM"}; } @@ -104,8 +111,8 @@ void save_transformer_config(const TransformerConfig& config, const char* file_n config_json["initializer_range"] = 0.02f; config_json["hidden_act"] = "silu"; config_json["use_cache"] = true; - if(config.Architecture == TransformerConfig::QWEN2) { - config_json["model_type"] = "qwen2"; + if (config.Architecture == TransformerConfig::QWEN2 || config.Architecture == TransformerConfig::QWEN3) { + config_json["model_type"] = config.Architecture == TransformerConfig::QWEN2 ? "qwen2" : "qwen3"; config_json["max_window_layers"] = config.NumLayers; config_json["sliding_window"] = config.MaxPositionEmbeddings; config_json["use_sliding_window"] = false; @@ -119,7 +126,7 @@ void save_transformer_config(const TransformerConfig& config, const char* file_n file << config_json.dump(4); } -static TransformerConfig create_qwen25_config(int hidden_size, int intermediate_size, int q_heads, int kv_heads, int depth, float rms, bool tied, ETensorDType dtype) { +static TransformerConfig create_qwen2_config(int hidden_size, int intermediate_size, int q_heads, int kv_heads, int depth, float rms, bool tied, ETensorDType dtype) { return { .Architecture = TransformerConfig::QWEN2, .BosTokenId = 151643, @@ -139,6 +146,27 @@ static TransformerConfig create_qwen25_config(int hidden_size, int intermediate_ }; } +static TransformerConfig create_qwen3_config(int hidden_size, int intermediate_size, int q_heads, int kv_heads, int depth, float rms, bool tied, ETensorDType dtype) { + return { + .Architecture = TransformerConfig::QWEN3, + .BosTokenId = 151643, + .EosTokenId = 151645, + .HiddenSize = hidden_size, + .IntermediateSize = intermediate_size, + .VocabSize = 151936, + .NumQueryHeads = q_heads, + .NumKeyValHeads = kv_heads, + .NumLayers = depth, + .MaxPositionEmbeddings = 40960, + .RopeTheta = 1'000'000.f, + .RmsNormEps = rms, + .TiedWordEmbeddings = tied, + .UseQKVBias = false, + .UseQKNorm = true, + .DType = dtype + }; +} + static TransformerConfig create_llama2_config(int hidden_size, int intermediate_size, int heads, int depth, ETensorDType dtype) { return { .Architecture = TransformerConfig::LLAMA, @@ -182,20 +210,32 @@ static TransformerConfig create_llama3_config(int hidden_size, int intermediate_ } TransformerConfig create_config_from_name(std::string_view name, ETensorDType dtype) { - if(iequals(name, "Qwen2.5-0.5B")) { - return create_qwen25_config(896, 4864, 14, 2, 24, 1e-06f, true, dtype); - } else if(iequals(name, "Qwen2.5-1.5B")) { - return create_qwen25_config(1536, 8960, 12, 2, 28, 1e-06f, true, dtype); - } else if(iequals(name, "Qwen2.5-3B")) { - return create_qwen25_config(2048, 11008, 16, 2, 36, 1e-06f, true, dtype); - } else if(iequals(name, "Qwen2.5-7B")) { - return create_qwen25_config(3584, 18944, 28, 4, 28, 1e-06f, false, dtype); - } else if(iequals(name, "Qwen2.5-14B")) { - return create_qwen25_config(5120, 13824, 40, 8, 48, 1e-05f, false, dtype); - } else if(iequals(name, "Qwen2.5-32B")) { - return create_qwen25_config(5120, 27648, 40, 8, 64, 1e-05f, false, dtype); - } else if(iequals(name, "Qwen2.5-72B")) { - return create_qwen25_config(8192, 29568, 64, 8, 80, 1e-05f, false, dtype); + if(iequals(name, "Qwen2.5-0.5B") || iequals(name, "Qwen2-0.5B")) { + return create_qwen2_config(896, 4864, 14, 2, 24, 1e-06f, true, dtype); + } else if(iequals(name, "Qwen2.5-1.5B") || iequals(name, "Qwen2-1.5B")) { + return create_qwen2_config(1536, 8960, 12, 2, 28, 1e-06f, true, dtype); + } else if(iequals(name, "Qwen2.5-3B") || iequals(name, "Qwen2-3B")) { + return create_qwen2_config(2048, 11008, 16, 2, 36, 1e-06f, true, dtype); + } else if(iequals(name, "Qwen2.5-7B") || iequals(name, "Qwen2-7B")) { + return create_qwen2_config(3584, 18944, 28, 4, 28, 1e-06f, false, dtype); + } else if(iequals(name, "Qwen2.5-14B") || iequals(name, "Qwen2-14B")) { + return create_qwen2_config(5120, 13824, 40, 8, 48, 1e-05f, false, dtype); + } else if(iequals(name, "Qwen2.5-32B") || iequals(name, "Qwen2-32B")) { + return create_qwen2_config(5120, 27648, 40, 8, 64, 1e-05f, false, dtype); + } else if(iequals(name, "Qwen2.5-72B") || iequals(name, "Qwen2-72B")) { + return create_qwen2_config(8192, 29568, 64, 8, 80, 1e-05f, false, dtype); + } else if (iequals(name, "Qwen3-0.6B")) { + return create_qwen3_config(1024, 3072, 16, 8, 28, 1e-6f, true, dtype); + } else if (iequals(name, "Qwen3-1.7B")) { + return create_qwen3_config(2048, 6144, 16, 8, 28, 1e-6f, true, dtype); + } else if (iequals(name, "Qwen3-4B")) { + return create_qwen3_config(2560, 9728, 32, 8, 36, 1e-6f, true, dtype); + } else if (iequals(name, "Qwen3-8B")) { + return create_qwen3_config(4096, 12288, 32, 8, 36, 1e-6f, false, dtype); + } else if (iequals(name, "Qwen3-14B")) { + return create_qwen3_config(5120, 17408, 40, 8, 40, 1e-6f, false, dtype); + } else if (iequals(name, "Qwen3-32B")) { + return create_qwen3_config(5120, 25600, 64, 8, 64, 1e-6f, false, dtype); } else if (iequals(name, "llama-2-7b")) { return create_llama2_config(4096, 11008, 32, 32, dtype); } else if (iequals(name, "llama-2-13b")) { diff --git a/src/training/transformer_config.h b/src/training/transformer_config.h index 21b526d..de3e3e6 100644 --- a/src/training/transformer_config.h +++ b/src/training/transformer_config.h @@ -16,6 +16,7 @@ struct TransformerConfig { enum EArchitecture { LLAMA, QWEN2, + QWEN3, } Architecture; int BosTokenId; int EosTokenId; @@ -33,6 +34,7 @@ struct TransformerConfig { float RmsNormEps; bool TiedWordEmbeddings; bool UseQKVBias; + bool UseQKNorm = false; ETensorDType DType = ETensorDType::BF16; From 0c6041804c122654bbb7d9db3a91b72fc198ace2 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 4 Mar 2026 15:31:18 +0100 Subject: [PATCH 14/19] cleanup --- src/models/llama_weights.cpp | 34 ---------------------------------- 1 file changed, 34 deletions(-) diff --git a/src/models/llama_weights.cpp b/src/models/llama_weights.cpp index e6977a1..0f8e79b 100644 --- a/src/models/llama_weights.cpp +++ b/src/models/llama_weights.cpp @@ -84,40 +84,6 @@ void fill_non_matrix_shapes(sLLamaBlockWeights& target, const Trans create_vector_shard(target.Attn_QKV_b, config.UseQKVBias ? attn_intermediate_size : 0); } -std::size_t aligned_size(std::size_t raw, int num_shards) { - return div_ceil(div_exact(raw, static_cast(num_shards)), static_cast(4096)) * 4096; -} - -std::size_t bytes_for_block_matrices(const TransformerConfig& config, ETensorDType dtype, int num_shards) { - std::size_t C = config.HiddenSize; - std::size_t HS = config.head_size(); - - std::size_t total = 2 * aligned_size(C * get_dtype_size(dtype), num_shards); // norms - long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * HS; - if(config.UseQKVBias) { - total += aligned_size(attn_intermediate_size * get_dtype_size(dtype), num_shards); // QKV bias - } - return total; -} - -std::size_t bytes_for_block_non_matrix(const TransformerConfig& config, ETensorDType dtype, int num_shards) { - std::size_t C = config.HiddenSize; - long H = config.IntermediateSize; - long HS = C / config.NumQueryHeads; - long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * HS; - - std::size_t total = 0; - total += aligned_size(attn_intermediate_size * C * get_dtype_size(dtype), num_shards); // QKV - total += aligned_size(C * C * get_dtype_size(dtype), num_shards); // out - total += aligned_size(2 * C * H * get_dtype_size(dtype), num_shards); // up - total += aligned_size(H * C * get_dtype_size(dtype), num_shards); // down - return total; -} - -std::size_t bytes_for_block(const TransformerConfig& config, ETensorDType matrix_dtype, ETensorDType other_dtype, int num_shards) { - return bytes_for_block_non_matrix(config, other_dtype, num_shards) + bytes_for_block_matrices(config, matrix_dtype, num_shards); -} - sLLamaBlockWeights allocate_block_full(const TransformerConfig& config, ETensorDType matrix_dtype, ETensorDType other_dtype, EAllocationType kind, TensorAllocator& alloc) { sLLamaBlockWeights layer; allocate_matrix_params(layer, config, matrix_dtype, kind, 0, 1, alloc); From dfd34aa8f573aa9d46afe3dbb7a5266cb5f41b3e Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 4 Mar 2026 15:31:36 +0100 Subject: [PATCH 15/19] add QKNorm weights --- src/models/llama_gradients.cpp | 4 ++++ src/models/llama_model.cpp | 4 +++- src/models/llama_optimizer.cpp | 2 ++ src/models/llama_weights.cpp | 36 +++++++++++++++++++++++++++++++--- src/models/llama_weights.h | 12 +++++++----- 5 files changed, 49 insertions(+), 9 deletions(-) diff --git a/src/models/llama_gradients.cpp b/src/models/llama_gradients.cpp index d73f362..788c3fb 100644 --- a/src/models/llama_gradients.cpp +++ b/src/models/llama_gradients.cpp @@ -21,6 +21,8 @@ void LLamaGradientsUnsharded::on_first_micro_step(cudaStream_t stream) { fill_zero(layer.get_tensor(LN1_W), stream); fill_zero(layer.get_tensor(LN2_W), stream); fill_zero(layer.get_tensor(QKV_B), stream); + fill_zero(layer.get_tensor(QNORM_W), stream); + fill_zero(layer.get_tensor(KNORM_W), stream); // no need to zero out the matrix weights, we'll just overwrite them on the first // grad accumulation step } @@ -42,6 +44,8 @@ void LLamaGradientsBlockShardedBase::on_get_block(SimpleTensorContainer& block, fill_zero(block.get_tensor(LLamaWeightID::LN1_W), stream); fill_zero(block.get_tensor(LLamaWeightID::LN2_W), stream); fill_zero(block.get_tensor(LLamaWeightID::QKV_B), stream); + fill_zero(block.get_tensor(LLamaWeightID::QNORM_W), stream); + fill_zero(block.get_tensor(LLamaWeightID::KNORM_W), stream); } diff --git a/src/models/llama_model.cpp b/src/models/llama_model.cpp index 3e131f8..85a6404 100644 --- a/src/models/llama_model.cpp +++ b/src/models/llama_model.cpp @@ -720,7 +720,7 @@ IRunState& LLamaModel::get_run_state() const { } std::size_t LLamaModel::num_block_tensors() const { - return 7; + return 9; } void LLamaModel::fill_block_shapes(GenericTensorContainer& target, const TransformerConfig& config, @@ -746,6 +746,8 @@ void LLamaModel::fill_block_shapes(GenericTensorContainer& target, const Transfo create(target.get_tensor(LLamaWeightID::LN1_W), C, 0, other_dtype); create(target.get_tensor(LLamaWeightID::LN2_W), C, 0, other_dtype); create(target.get_tensor(LLamaWeightID::QKV_B), config.UseQKVBias ? attn_intermediate_size : 0, 0, other_dtype); + create(target.get_tensor(LLamaWeightID::QNORM_W), config.UseQKNorm ? HS : 0, 0, other_dtype); + create(target.get_tensor(LLamaWeightID::KNORM_W), config.UseQKNorm ? HS : 0, 0, other_dtype); } std::size_t LLamaModel::num_non_block_tensors() const { diff --git a/src/models/llama_optimizer.cpp b/src/models/llama_optimizer.cpp index 52a1f8e..e66b245 100644 --- a/src/models/llama_optimizer.cpp +++ b/src/models/llama_optimizer.cpp @@ -39,6 +39,8 @@ void OptStateWrapper::iterate_tensors(const std::function& target, const Transformer target.LN1_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "ln1_w", {C}, kind); target.LN2_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "ln2_w", {C}, kind); + if (config.UseQKNorm) { + target.QNorm_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "qnorm_w", {HS}, kind); + target.KNorm_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "knorm_w", {HS}, kind); + } else { + target.QNorm_w = Tensor{}; + target.KNorm_w = Tensor{}; + } + long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * HS; if(config.UseQKVBias) { target.Attn_QKV_b = alloc.allocate_shard(dtype, shard_idx, num_shards, "att_qkv_b", {attn_intermediate_size}, kind); @@ -80,6 +88,8 @@ void fill_non_matrix_shapes(sLLamaBlockWeights& target, const Trans create_vector_shard(target.LN1_w, C); create_vector_shard(target.LN2_w, C); + create_vector_shard(target.QNorm_w, config.UseQKNorm ? HS : 0); + create_vector_shard(target.KNorm_w, config.UseQKNorm ? HS : 0); long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * HS; create_vector_shard(target.Attn_QKV_b, config.UseQKVBias ? attn_intermediate_size : 0); } @@ -107,6 +117,8 @@ sLLamaBlockWeights shard_block(const sLLamaBlockWeights& bl result.Attn_QKV_b = shard_view(block.Attn_QKV_b, shard_idx, num_shards); result.LN1_w = shard_view(block.LN1_w, shard_idx, num_shards); result.LN2_w = shard_view(block.LN2_w, shard_idx, num_shards); + result.QNorm_w = shard_view(block.QNorm_w, shard_idx, num_shards); + result.KNorm_w = shard_view(block.KNorm_w, shard_idx, num_shards); return result; } @@ -221,7 +233,7 @@ LLamaWeightsManager::~LLamaWeightsManager() { void LLamaWeightsManager::setup_scales(TensorAllocator& alloc) { int layers = mMaster.Blocks.size(); - mAbsMaxes = alloc.allocate(ETensorDType::FP32, "abs_maxes", EAllocationType::ON_DEVICE, {6 + layers * 14}); + mAbsMaxes = alloc.allocate(ETensorDType::FP32, "abs_maxes", EAllocationType::ON_DEVICE, {6 + layers * 18}); float* abs_maxes = mAbsMaxes.get(); mMaster.NonBlocks.Embeddings.Stats = abs_maxes + 0; mMaster.NonBlocks.LNF_w.Stats = abs_maxes + 2; @@ -235,14 +247,16 @@ void LLamaWeightsManager::setup_scales(TensorAllocator& alloc) { mMaster.Blocks[i].Attn_QKV_b.Stats = a + 8; mMaster.Blocks[i].LN1_w.Stats = a + 10; mMaster.Blocks[i].LN2_w.Stats = a + 12; + mMaster.Blocks[i].QNorm_w.Stats = a + 14; + mMaster.Blocks[i].KNorm_w.Stats = a + 16; } } std::pair LLamaWeightsManager::get_scales_for_block(int layer_idx) { float* abs_maxes = mAbsMaxes.get(); - float* begin = abs_maxes + 6 + layer_idx * 14; - return {begin + 0, begin + 14}; + float* begin = abs_maxes + 6 + layer_idx * 18; + return {begin + 0, begin + 18}; } @@ -305,6 +319,8 @@ void LLamaWeightsManager::begin_optimizer(DeviceMemoryStack& memory, cudaStream_ buf.LN1_w = mMaster.Blocks[0].LN1_w; buf.LN2_w = mMaster.Blocks[0].LN2_w; buf.Attn_QKV_b = mMaster.Blocks[0].Attn_QKV_b; + buf.QNorm_w = mMaster.Blocks[0].QNorm_w; + buf.KNorm_w = mMaster.Blocks[0].KNorm_w; } mMasterDeviceDoubleBufferStorage[i] = alloc.commit(memory, "master"); @@ -555,6 +571,12 @@ void sLLamaWeights::iterate_tensors(const std::function shard_non_block(const sLLamaNonBlockWeights& target, const TransformerConfig& config, ETensorDType dtype, int shard_idx, int num_shards); void fill_non_matrix_shapes(sLLamaBlockWeights& target, const TransformerConfig& config, ETensorDType dtype, int shard_idx, int num_shards); -std::size_t bytes_for_block(const TransformerConfig& config, ETensorDType matrix_dtype, ETensorDType other_dtype, int num_shards); -std::size_t bytes_for_block_matrices(const TransformerConfig& config, ETensorDType dtype, int num_shards); -std::size_t bytes_for_block_non_matrix(const TransformerConfig& config, ETensorDType dtype, int num_shards); - #endif //LLMQ_LLAMA_WEIGHTS_H From 865e0e286f8da93822839ea27cd2ac5979e08a8f Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 19 Jul 2026 00:22:55 +0200 Subject: [PATCH 16/19] decouple head_dim from hidden size Qwen3 models specify head_dim explicitly (128), which does not equal HiddenSize / NumQueryHeads for the 0.6B/4B/32B variants. Add an optional HeadDim to TransformerConfig (0 = derived as before, so existing models are unchanged) and route all users through head_size() / the new attn_channels(): - o_proj weight becomes {C, attn_channels} - the o_proj matmuls take Hq*Hs input channels - att activations and d_att_y are sized (B, T, attn_channels); d_att_y can no longer alias d_lnf when the sizes differ Co-Authored-By: Claude Fable 5 --- src/models/llama_model.cpp | 8 ++++---- src/models/llama_run_state.cpp | 10 ++++++---- src/models/llama_run_state.h | 4 ++-- src/models/llama_weights.cpp | 10 ++++------ src/training/transformer_config.cpp | 7 +++++++ src/training/transformer_config.h | 6 +++++- 6 files changed, 28 insertions(+), 17 deletions(-) diff --git a/src/models/llama_model.cpp b/src/models/llama_model.cpp index 85a6404..2205b82 100644 --- a/src/models/llama_model.cpp +++ b/src/models/llama_model.cpp @@ -206,7 +206,7 @@ void LLamaModel::_forward_block(sLLamaBlockWeights& weights, sLLamaLayer forward_qmm(acts.AttO, acts.Att, weights.Attn_Out_w, Tensor{}, rs->CublasLtHandle, rs->CuBlasWorkspace, - B, T, C, C, + B, T, Hq * Hs, C, rs->DeviceProp, false, main_stream, rs->MatmulBackend); fused_residual_rmsnorm_forward(acts.ResidualAtt, acts.LN2.Value, acts.LN2_Rstd, residual, acts.AttO, weights.LN2_w, @@ -590,7 +590,7 @@ void LLamaModel::_recompute_block(sLLamaBlockWeights& weights, sLLamaLay if (opt.RecomputeBlock) { forward_qmm(acts.AttO, acts.Att, weights.Attn_Out_w, Tensor{}, rs->CublasLtHandle, rs->CuBlasWorkspace, - B, T, C, C, + B, T, Hq * Hs, C, rs->DeviceProp, false, main_stream, rs->MatmulBackend); } } @@ -660,7 +660,7 @@ void LLamaModel::_backward_block(bool accumulate, sLLamaBlockWeights& we bool recompute_ln1 = rs->Options.RecomputeRMSNorm || rs->Options.RecomputeAtt; backward_qmm(d_acts.DAttY, d_weights.get_tensor(ATTO_W), Tensor{}, d_acts.DResAtt, acts.Att, weights.Attn_Out_w, Tensor{}, - accumulate, *rs, B, T, C, C, false, main_stream); + accumulate, *rs, B, T, Hq * Hs, C, false, main_stream); rs->temp_acquire(d_acts.DQKV.Value); rs->temp_acquire(rs->CuDNNWorkspace); @@ -739,7 +739,7 @@ void LLamaModel::fill_block_shapes(GenericTensorContainer& target, const Transfo long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * HS; create(target.get_tensor(LLamaWeightID::QKV_W), attn_intermediate_size, C, matrix_dtype); - create(target.get_tensor(LLamaWeightID::ATTO_W), C, C, matrix_dtype); + create(target.get_tensor(LLamaWeightID::ATTO_W), C, config.attn_channels(), matrix_dtype); create(target.get_tensor(LLamaWeightID::UP_W), 2 * H, C, matrix_dtype); create(target.get_tensor(LLamaWeightID::DOWN_W), C, H, matrix_dtype); diff --git a/src/models/llama_run_state.cpp b/src/models/llama_run_state.cpp index bd7f208..a261582 100644 --- a/src/models/llama_run_state.cpp +++ b/src/models/llama_run_state.cpp @@ -17,7 +17,7 @@ constexpr const int QWEN2_NUM_LINEAR_OPS = 4; class RunStateBuilder { public: RunStateBuilder(TransformerConfig config, LLamaOptions options, int B, int T, std::shared_ptr alloc) - : Config(config), Options(options), B(B), T(T), C(config.HiddenSize), H(config.IntermediateSize), Alloc(alloc) + : Config(config), Options(options), B(B), T(T), C(config.HiddenSize), H(config.IntermediateSize), AC(config.attn_channels()), Alloc(alloc) { } @@ -56,6 +56,7 @@ class RunStateBuilder { long T; long C; // Config.HiddenSize; long H; // Config.IntermediateSize; + long AC; // Config.attn_channels(); std::shared_ptr Alloc; Tensor tSwiGluBuffer; @@ -96,7 +97,7 @@ LLamaRunState::LayerActivations RunStateBuilder::allocate_basic_fwd_tensors(Tens Tensor qkv = allocate_or_reuse(Options.RecomputeQKV, tQKVBuffer, Config.DType, "qkv", B, T, Config.qkv_channels()); Tensor res_att = allocate_or_reuse(Options.RecomputeBlock, tResAttBuffer, Config.DType, "res_att", B, T, C); Tensor lse = allocate(ETensorDType::FP32, "lse", B, T, Config.NumQueryHeads); - Tensor att_v = allocate_or_reuse(Options.RecomputeAtt, tAttBuffer, Config.DType, "att_v", B, T, C); + Tensor att_v = allocate_or_reuse(Options.RecomputeAtt, tAttBuffer, Config.DType, "att_v", B, T, AC); // not needed for backward, so can reuse an existing buffer // we can use the same buffer as for the rms norms, because those support // inplace transforms. @@ -123,7 +124,7 @@ void RunStateBuilder::allocate_fwd_quant_tensors(LLamaRunState::LayerActivations // allocate a new buffer for every forward quantization act.LN1.Quant = allocate(matmul_dtype, "ln1.q", B, T, C); act.LN2.Quant = allocate(matmul_dtype, "ln2.q", B, T, C); - act.Att.Quant = allocate(matmul_dtype, "att.q", B, T, C); + act.Att.Quant = allocate(matmul_dtype, "att.q", B, T, AC); act.SwiGLu.Quant = allocate(matmul_dtype, "swiglu.q", B, T, H); } @@ -163,7 +164,8 @@ LLamaRunState::LayerGradients RunStateBuilder::allocate_basic_bwd_tensors(Tensor Tensor d_swiglu = Tensor{Config.DType, {B, T, H}, nullptr, nullptr, 3, d_lnf.Device}; QuantizableTensor d_mlp_up{}; // this will be handled in-place Tensor d_ln2 = Options.KeepAllActivations ? allocate(Config.DType, "d_ln2", B, T, C) : d_lnf; - Tensor d_att_y = Options.KeepAllActivations ? allocate(Config.DType, "d_att_y", B, T, C) : d_lnf; + // d_lnf has shape (B, T, C), so it can only be reused if the attention channels match the hidden size + Tensor d_att_y = (Options.KeepAllActivations || AC != C) ? allocate(Config.DType, "d_att_y", B, T, AC) : d_lnf; QuantizableTensor d_qkv{Tensor{Config.DType, {B, T, Config.qkv_channels()}, nullptr, nullptr, 3, d_lnf.Device}}; Tensor d_ln1 = Options.KeepAllActivations ? allocate(Config.DType, "d_ln1", B, T, C) : d_lnf; QuantizableTensor d_res_att = Options.KeepAllActivations ? QuantizableTensor{allocate(Config.DType, "d_res_att", B, T, C)} : d_res_ffn; diff --git a/src/models/llama_run_state.h b/src/models/llama_run_state.h index b1857d1..307beb6 100644 --- a/src/models/llama_run_state.h +++ b/src/models/llama_run_state.h @@ -32,7 +32,7 @@ struct sLLamaLayerActivations { QTensor LN2; // (B, T, C) Tensor QKV; // (B, T, QKV_C) Tensor LSE; // (B, T) - QTensor Att; // (B, T, C) + QTensor Att; // (B, T, AC) Tensor AttO; // (B, T, C) Tensor ResidualAtt; // (B, T, C) Tensor MlpUp; // (B, T, 2*Ch) @@ -48,7 +48,7 @@ struct sLLamaLayerGradients { QTensor DMlpUp; // (B, T, 2*Ch) Tensor DLN2; // (B, T, C) QTensor DResAtt; // (B, T, C) - Tensor DAttY; // (B, T, C) + Tensor DAttY; // (B, T, AC) QTensor DQKV; // (B, T, QKV_C) Tensor DLN1; // (B, T, C) }; diff --git a/src/models/llama_weights.cpp b/src/models/llama_weights.cpp index d782553..58ea83f 100644 --- a/src/models/llama_weights.cpp +++ b/src/models/llama_weights.cpp @@ -40,10 +40,9 @@ void allocate_matrix_params(sLLamaBlockWeights& target, const TransformerConf long C = config.HiddenSize; long H = config.IntermediateSize; - long head_size = C / config.NumQueryHeads; - long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * head_size; + long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * config.head_size(); target.Attn_QKV_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "att_qkv_w", {attn_intermediate_size, C}, kind); - target.Attn_Out_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "attproj_w", {C, C}, kind); + target.Attn_Out_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "attproj_w", {C, (long)config.attn_channels()}, kind); target.MLP_Up_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "mlp_up_w", {2 * H, C}, kind); target.MLP_Down_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "mlp_down_w", {C, H}, kind); } @@ -64,10 +63,9 @@ void fill_matrix_shapes(sLLamaBlockWeights& target, const Transform tgt.GlobalShape[1] = cols; }; - long head_size = C / config.NumQueryHeads; - long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * head_size; + long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * config.head_size(); create_matrix_shard(target.Attn_QKV_w, attn_intermediate_size, C); - create_matrix_shard(target.Attn_Out_w, C, C); + create_matrix_shard(target.Attn_Out_w, C, config.attn_channels()); create_matrix_shard(target.MLP_Up_w, 2 * H, C); create_matrix_shard(target.MLP_Down_w, C, H); } diff --git a/src/training/transformer_config.cpp b/src/training/transformer_config.cpp index f0709f4..f20f4ec 100644 --- a/src/training/transformer_config.cpp +++ b/src/training/transformer_config.cpp @@ -48,6 +48,9 @@ TransformerConfig load_transformer_config(const char* file_name, ETensorDType dt result.NumQueryHeads = config_json["num_attention_heads"].get(); result.NumKeyValHeads = config_json["num_key_value_heads"].get(); result.NumLayers = config_json["num_hidden_layers"].get(); + if (config_json.contains("head_dim")) { + result.HeadDim = config_json["head_dim"].get(); + } result.MaxPositionEmbeddings = config_json["max_position_embeddings"].get(); result.RopeTheta = config_json["rope_theta"].get(); result.TiedWordEmbeddings = config_json["tie_word_embeddings"].get(); @@ -101,6 +104,9 @@ void save_transformer_config(const TransformerConfig& config, const char* file_n config_json["num_attention_heads"] = config.NumQueryHeads; config_json["num_key_value_heads"] = config.NumKeyValHeads; config_json["num_hidden_layers"] = config.NumLayers; + if (config.HeadDim != 0) { + config_json["head_dim"] = config.HeadDim; + } config_json["max_position_embeddings"] = config.MaxPositionEmbeddings; config_json["rope_theta"] = config.RopeTheta; config_json["rms_norm_eps"] = config.RmsNormEps; @@ -157,6 +163,7 @@ static TransformerConfig create_qwen3_config(int hidden_size, int intermediate_s .NumQueryHeads = q_heads, .NumKeyValHeads = kv_heads, .NumLayers = depth, + .HeadDim = 128, .MaxPositionEmbeddings = 40960, .RopeTheta = 1'000'000.f, .RmsNormEps = rms, diff --git a/src/training/transformer_config.h b/src/training/transformer_config.h index de3e3e6..9c93c13 100644 --- a/src/training/transformer_config.h +++ b/src/training/transformer_config.h @@ -28,6 +28,7 @@ struct TransformerConfig { int NumQueryHeads; int NumKeyValHeads; int NumLayers; + int HeadDim = 0; // 0: derived as HiddenSize / NumQueryHeads int MaxPositionEmbeddings; float RopeTheta; @@ -38,8 +39,11 @@ struct TransformerConfig { ETensorDType DType = ETensorDType::BF16; - [[nodiscard]] int head_size() const { return HiddenSize / NumQueryHeads; } + [[nodiscard]] int head_size() const { return HeadDim != 0 ? HeadDim : HiddenSize / NumQueryHeads; } + //! number of channels consumed by the attention operation [[nodiscard]] int qkv_channels() const { return head_size() * (NumQueryHeads + 2 * NumKeyValHeads); } + //! number of channels produced by the attention operation + [[nodiscard]] int attn_channels() const { return head_size() * NumQueryHeads; } [[nodiscard]] std::string_view model_name() const; }; From 06bc00b9348d03b55f70c0e664c458683a2fe18c Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 19 Jul 2026 02:58:43 +0200 Subject: [PATCH 17/19] enable qk-norm --- src/binding/py_train.cpp | 4 ++ src/models/llama_model.cpp | 67 +++++++++++++++++++++++++++++----- src/models/llama_model.h | 5 +++ src/models/llama_run_state.cpp | 29 ++++++++++++++- src/models/llama_run_state.h | 4 ++ 5 files changed, 97 insertions(+), 12 deletions(-) diff --git a/src/binding/py_train.cpp b/src/binding/py_train.cpp index 5f875d3..94f6378 100644 --- a/src/binding/py_train.cpp +++ b/src/binding/py_train.cpp @@ -297,6 +297,10 @@ std::vector> MultiGPUPyTrainer::get_gradients(int if (block.get_tensor(QKV_B)) result.emplace_back(prefix + ".self_attn.qkv.bias", block.get_tensor(QKV_B)); result.emplace_back(prefix + ".self_attn.o_proj.weight", block.get_tensor(ATTO_W)); + if (block.get_tensor(QNORM_W)) + result.emplace_back(prefix + ".self_attn.q_norm.weight", block.get_tensor(QNORM_W)); + if (block.get_tensor(KNORM_W)) + result.emplace_back(prefix + ".self_attn.k_norm.weight", block.get_tensor(KNORM_W)); result.emplace_back(prefix + ".mlp.up.weight", block.get_tensor(UP_W)); result.emplace_back(prefix + ".mlp.down_proj.weight", block.get_tensor(DOWN_W)); result.emplace_back(prefix + ".input_layernorm.weight", block.get_tensor(LN1_W)); diff --git a/src/models/llama_model.cpp b/src/models/llama_model.cpp index 2205b82..5303717 100644 --- a/src/models/llama_model.cpp +++ b/src/models/llama_model.cpp @@ -17,9 +17,6 @@ LLamaModel::LLamaModel(TransformerConfig config, const LLamaOptions& options, int rank, int world, const std::shared_ptr& alloc) : Config(config), Options(options), Allocator(alloc ? alloc : std::make_shared()) { - if (Config.UseQKNorm) - throw std::runtime_error("UseQKNorm is not yet supported"); - Parameters = LLamaWeightsManager::create(Config, options, rank, world, *Allocator); } @@ -195,10 +192,22 @@ void LLamaModel::_forward_block(sLLamaBlockWeights& weights, sLLamaLayer rs->CublasLtHandle, rs->CuBlasWorkspace, B, T, C, Config.qkv_channels(), rs->DeviceProp, false, main_stream, rs->MatmulBackend); - // 2) apply RoPE to q,k (potentially in place) - rope_forward(acts.QKV, acts.QKV, rs->FreqCis, nullptr, B, T, Hq, Hkv, Hs, main_stream); + // 2) apply qk-norm (if enabled) and RoPE to q,k. With qk-norm, acts.QKV keeps the + // pre-norm projection for the backward pass; only the rotated copy is norm-scaled. // 3) attention: att <- softmax(qk^T)v - attention_forward_cudnn(acts.Att.Value, acts.LSE, acts.QKV, rs->CuBlasWorkspace, rs->CudnnHandle, B, T, Hq, Hkv, Hs, main_stream); + if (Config.UseQKNorm) { + bool transient = Options.recompute_qk_rope(); + if (transient) rs->temp_acquire(rs->QKVRope); + Tensor& post_rope = transient ? rs->QKVRope : acts.PostRopeQKV; + qk_norm_and_rope_forward(post_rope, acts.QK_Rstd, nullptr, acts.QKV, + weights.QNorm_w, weights.KNorm_w, rs->FreqCis, + Config.RmsNormEps, B, T, Hq, Hkv, Hs, main_stream); + attention_forward_cudnn(acts.Att.Value, acts.LSE, post_rope, rs->CuBlasWorkspace, rs->CudnnHandle, B, T, Hq, Hkv, Hs, main_stream); + if (transient) rs->temp_free(rs->QKVRope); + } else { + rope_forward(acts.QKV, acts.QKV, rs->FreqCis, nullptr, B, T, Hq, Hkv, Hs, main_stream); + attention_forward_cudnn(acts.Att.Value, acts.LSE, acts.QKV, rs->CuBlasWorkspace, rs->CudnnHandle, B, T, Hq, Hkv, Hs, main_stream); + } // quantize attention if necessary if(acts.Att.Quant) { abs_max(acts.Att.Quant.abs_max(), acts.Att.Value, acts.Att.Value.nelem(), rs->DeviceProp, main_stream); @@ -580,11 +589,24 @@ void LLamaModel::_recompute_block(sLLamaBlockWeights& weights, sLLamaLay rs->CublasLtHandle, rs->CuBlasWorkspace, B, T, C, Config.qkv_channels(), rs->DeviceProp, !recompute_ln1, main_stream, rs->MatmulBackend); - rope_forward(acts.QKV, acts.QKV, rs->FreqCis, nullptr, B, T, Hq, Hkv, Hs, main_stream); + // with qk-norm, acts.QKV keeps the pre-norm projection + if (!Config.UseQKNorm) { + rope_forward(acts.QKV, acts.QKV, rs->FreqCis, nullptr, B, T, Hq, Hkv, Hs, main_stream); + } } if (recompute_att) { - attention_forward_cudnn(acts.Att.Value, acts.LSE, acts.QKV, rs->CuBlasWorkspace, rs->CudnnHandle, B, T, Hq, Hkv, Hs, main_stream); + // recompute-att implies recompute_qk_rope(), so PostRopeQKV is never populated here + if (Config.UseQKNorm) { + rs->temp_acquire(rs->QKVRope); + qk_norm_and_rope_forward(rs->QKVRope, acts.QK_Rstd, nullptr, acts.QKV, + weights.QNorm_w, weights.KNorm_w, rs->FreqCis, + Config.RmsNormEps, B, T, Hq, Hkv, Hs, main_stream); + attention_forward_cudnn(acts.Att.Value, acts.LSE, rs->QKVRope, rs->CuBlasWorkspace, rs->CudnnHandle, B, T, Hq, Hkv, Hs, main_stream); + rs->temp_free(rs->QKVRope); + } else { + attention_forward_cudnn(acts.Att.Value, acts.LSE, acts.QKV, rs->CuBlasWorkspace, rs->CudnnHandle, B, T, Hq, Hkv, Hs, main_stream); + } // AttO not needed in backward pass; but if we want to recompute the entire transformer block, we need its output // to recompute the FFN part if (opt.RecomputeBlock) { @@ -663,6 +685,16 @@ void LLamaModel::_backward_block(bool accumulate, sLLamaBlockWeights& we accumulate, *rs, B, T, Hq * Hs, C, false, main_stream); rs->temp_acquire(d_acts.DQKV.Value); + bool rematerialize = Config.UseQKNorm && Options.recompute_qk_rope(); + if (rematerialize) { + // rematerialize the post-norm+rope qkv for attention backward; also rewrites QK_Rstd + rs->temp_acquire(rs->QKVRope); + qk_norm_and_rope_forward(rs->QKVRope, acts.QK_Rstd, nullptr, acts.QKV, + weights.QNorm_w, weights.KNorm_w, rs->FreqCis, + Config.RmsNormEps, B, T, Hq, Hkv, Hs, main_stream); + } + const Tensor& att_qkv = !Config.UseQKNorm ? acts.QKV + : (rematerialize ? rs->QKVRope : acts.PostRopeQKV); rs->temp_acquire(rs->CuDNNWorkspace); for (int i=0; i < Options.AttBwdChunks; ++i) { long chunk_batch_size = div_exact(B, (long)Options.AttBwdChunks); @@ -670,12 +702,23 @@ void LLamaModel::_backward_block(bool accumulate, sLLamaBlockWeights& we Tensor lse = shard_view(acts.LSE, i, Options.AttBwdChunks); Tensor att = shard_view(acts.Att.Value, i, Options.AttBwdChunks); Tensor d_atty = shard_view(d_acts.DAttY, i, Options.AttBwdChunks); - Tensor qkv = shard_view(acts.QKV, i, Options.AttBwdChunks); + Tensor qkv = shard_view(att_qkv, i, Options.AttBwdChunks); attention_backward_cudnn(d_qkv, lse, att, d_atty, qkv, rs->CuDNNWorkspace, rs->CudnnHandle, chunk_batch_size, T, Hq, Hkv, Hs, main_stream); } rs->temp_free(rs->CuDNNWorkspace); - rope_backward(d_acts.DQKV.Value, d_acts.DQKV.Value, rs->FreqCis, d_acts.DQKV.Quant.abs_max(), B, T, Hq, Hkv, Hs, main_stream); + if (Config.UseQKNorm) { + if (rematerialize) rs->temp_free(rs->QKVRope); + // dq_wgt/dk_wgt accumulate, like the other norm-weight backwards + qk_norm_and_rope_backward(d_acts.DQKV.Value, + d_weights.get_tensor(QNORM_W), d_weights.get_tensor(KNORM_W), + rs->QKNormScratch, d_acts.DQKV.Value, acts.QKV, + weights.QNorm_w, weights.KNorm_w, acts.QK_Rstd, rs->FreqCis, + d_acts.DQKV.Quant.abs_max(), + B, T, Hq, Hkv, Hs, rs->DeviceProp, main_stream); + } else { + rope_backward(d_acts.DQKV.Value, d_acts.DQKV.Value, rs->FreqCis, d_acts.DQKV.Quant.abs_max(), B, T, Hq, Hkv, Hs, main_stream); + } backward_qmm(d_acts.DLN1, d_weights.get_tensor(QKV_W), d_weights.get_tensor(QKV_B), d_acts.DQKV, acts.LN1, weights.Attn_QKV_w, rs->MatmulBiasScratch, accumulate, *rs, B, T, C, Config.qkv_channels(), !recompute_ln1, main_stream); @@ -860,6 +903,10 @@ void LLamaModel::update(NCCLCommunicator& comm, float learning_rate, float beta_ auto& sm = OptimizerState->get_block_scales_m(i); run_update(bw.get_tensor(LN1_W), bg.get_tensor(LN1_W), bm.get_tensor(LN1_W), bv.get_tensor(LN1_W), sm.get_tensor(LN1_W), 0.f); run_update(bw.get_tensor(LN2_W), bg.get_tensor(LN2_W), bm.get_tensor(LN2_W), bv.get_tensor(LN2_W), sm.get_tensor(LN2_W), 0.f); + if(Config.UseQKNorm) { + run_update(bw.get_tensor(QNORM_W), bg.get_tensor(QNORM_W), bm.get_tensor(QNORM_W), bv.get_tensor(QNORM_W), sm.get_tensor(QNORM_W), 0.f); + run_update(bw.get_tensor(KNORM_W), bg.get_tensor(KNORM_W), bm.get_tensor(KNORM_W), bv.get_tensor(KNORM_W), sm.get_tensor(KNORM_W), 0.f); + } run_update(bw.get_tensor(QKV_W), bg.get_tensor(QKV_W), bm.get_tensor(QKV_W), bv.get_tensor(QKV_W), sm.get_tensor(QKV_W), weight_decay); diff --git a/src/models/llama_model.h b/src/models/llama_model.h index d590c25..3696f4f 100644 --- a/src/models/llama_model.h +++ b/src/models/llama_model.h @@ -61,6 +61,11 @@ struct LLamaOptions { return MatmulType.value_or(ModelType.value()); } + //! the post-norm+rope qkv is dropped and rematerialized along with the buffers it derives from + bool recompute_qk_rope() const { + return RecomputeQKV || RecomputeAtt || RecomputeBlock; + } + ETensorDType grad_dtype() const { return GradientType.value_or(matmul_dtype()); } diff --git a/src/models/llama_run_state.cpp b/src/models/llama_run_state.cpp index a261582..efee959 100644 --- a/src/models/llama_run_state.cpp +++ b/src/models/llama_run_state.cpp @@ -62,6 +62,7 @@ class RunStateBuilder { Tensor tSwiGluBuffer; Tensor tMlpUpBuffer; Tensor tQKVBuffer; + Tensor tQKRstdBuffer; Tensor tAttBuffer; Tensor tLN1Buffer; Tensor tResAttBuffer; @@ -95,6 +96,17 @@ LLamaRunState::LayerActivations RunStateBuilder::allocate_basic_fwd_tensors(Tens Tensor ln2_v = allocate_or_reuse(reuse_ln_buffer || Options.RecomputeFFN, lnf, Config.DType, "ln2", B, T, C); Tensor qkv = allocate_or_reuse(Options.RecomputeQKV, tQKVBuffer, Config.DType, "qkv", B, T, Config.qkv_channels()); + // qk-norm rstd, saved for backward; shareable across layers iff rematerialization rewrites it + Tensor qk_rstd; + Tensor post_rope_qkv; + if (Config.UseQKNorm) { + qk_rstd = allocate_or_reuse(Options.recompute_qk_rope(), tQKRstdBuffer, + ETensorDType::FP32, "qk_rstd", + B, T, (long)Config.NumQueryHeads + 2 * Config.NumKeyValHeads); + if (!Options.recompute_qk_rope()) { + post_rope_qkv = allocate(Config.DType, "post_rope_qkv", B, T, Config.qkv_channels()); + } + } Tensor res_att = allocate_or_reuse(Options.RecomputeBlock, tResAttBuffer, Config.DType, "res_att", B, T, C); Tensor lse = allocate(ETensorDType::FP32, "lse", B, T, Config.NumQueryHeads); Tensor att_v = allocate_or_reuse(Options.RecomputeAtt, tAttBuffer, Config.DType, "att_v", B, T, AC); @@ -116,7 +128,7 @@ LLamaRunState::LayerActivations RunStateBuilder::allocate_basic_fwd_tensors(Tens Tensor mlp_down = allocate_or_reuse(true, lnf, Config.DType, "mlp_down", B, T, C); return LLamaRunState::LayerActivations{ln1_rstd, ln1, ln2_rstd, ln2, qkv, lse, att, atto, - res_att, mlp_up, mlp_down, swiglu}; + res_att, mlp_up, mlp_down, swiglu, qk_rstd, post_rope_qkv}; } void RunStateBuilder::allocate_fwd_quant_tensors(LLamaRunState::LayerActivations& act) { @@ -278,6 +290,12 @@ LLamaRunState::LLamaRunState(TransformerConfig config, LLamaOptions options, lon LNF = alloc->allocate(Config.DType, "lnf", {B, T, C}); LNF_Rstd = alloc->allocate(ETensorDType::FP32, "lnf_rstd", {B, T}); DLNF = alloc->allocate(Config.DType, "d_lnf", {B, T, C}); + if (Config.UseQKNorm) { + long qknorm_scratch_size = qk_norm_and_rope_backward_scratch_size( + Config.NumQueryHeads, Config.NumKeyValHeads, Config.head_size(), Config.DType, DeviceProp); + QKNormScratch = alloc->allocate(ETensorDType::BYTE, "qknorm_scratch", {qknorm_scratch_size}); + QKVRope = Tensor{Config.DType, {B, T, (long)Config.qkv_channels()}, nullptr, nullptr, 3, Inputs.Device}; + } long rms_scratch_size = get_rmsnorm_backward_scratch_size(C, DeviceProp); long bias_scratch_size = get_bias_backward_scratch_size(Config.DType, Config.qkv_channels(), DeviceProp); RMSNormScratch = alloc->allocate(ETensorDType::BYTE, "rms_scratch", {rms_scratch_size}); @@ -376,7 +394,14 @@ LLamaRunState::LLamaRunState(TransformerConfig config, LLamaOptions options, lon // simulate to determine required stack size auto mlp_up = stack.allocate(Config.DType, {B, T, 2 * Config.IntermediateSize}, "mlp_up"); auto ws = stack.allocate(CuDNNWorkspace.bytes(), "workspace"); - stack.free(stack.allocate(DActs[0].DQKV.Value.bytes(), "dqkv")); // attention + if (Config.UseQKNorm && Options.recompute_qk_rope()) { + // backward holds dqkv, the rematerialized post-rope qkv, and the workspace at once + auto dqkv = stack.allocate(DActs[0].DQKV.Value.bytes(), "dqkv"); + stack.free(stack.allocate(QKVRope.bytes(), "qkv_rope")); + stack.free(dqkv); + } else { + stack.free(stack.allocate(DActs[0].DQKV.Value.bytes(), "dqkv")); // attention + } stack.free(ws); // attention auto dswi = stack.allocate(DActs[0].DSwiGLU.bytes(), "dswiglu"); diff --git a/src/models/llama_run_state.h b/src/models/llama_run_state.h index 307beb6..dd3007f 100644 --- a/src/models/llama_run_state.h +++ b/src/models/llama_run_state.h @@ -38,6 +38,8 @@ struct sLLamaLayerActivations { Tensor MlpUp; // (B, T, 2*Ch) Tensor MlpDown; // (B, T, C) QTensor SwiGLu; // (B, T, Ch) + Tensor QK_Rstd; // (B, T, Hq+2*Hkv); only if qk-norm is enabled + Tensor PostRopeQKV; // (B, T, QKV_C); only with qk-norm, unless recompute_qk_rope() }; struct sLLamaLayerGradients { @@ -94,6 +96,8 @@ struct LLamaRunState : public IRunState { // scratch buffers Tensor RMSNormScratch; // (#Blocks*C+128) + Tensor QKNormScratch; // per-block dweight partials; only if qk-norm is enabled + Tensor QKVRope; // (B, T, QKV_C) stack temp for the post-norm+rope qkv; only if qk-norm is enabled Tensor MatmulBiasScratch; // TODO Tensor CuDNNWorkspace; Tensor EncoderBwdScratch; // (B, T, 5 * C / (x128::size * 32)) From a2ea2b33b529df973f424734e9ed76217d598d70 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 19 Jul 2026 13:40:48 +0200 Subject: [PATCH 18/19] ignore redundant lm_head.weight when importing tied checkpoints Qwen3-0.6B ships an explicit lm_head.weight despite tie_word_embeddings=true; skip it instead of rejecting the file. Co-Authored-By: Claude Fable 5 --- src/models/llama_weights.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/models/llama_weights.cpp b/src/models/llama_weights.cpp index 58ea83f..6841dda 100644 --- a/src/models/llama_weights.cpp +++ b/src/models/llama_weights.cpp @@ -930,6 +930,8 @@ void LLamaWeightsManager::import_from_file(const std::string& file_name, bool al } else { throw std::runtime_error("Unexpected tensor name: " + entry.name()); } + } else if (entry.name() == "lm_head.weight" && mConfig.TiedWordEmbeddings) { + // some checkpoints of tied models store a redundant copy of the embeddings } else { throw std::runtime_error("Unexpected tensor name: " + entry.name()); } From e9708ab5dd8b0a2906c3f1381dccf6dfee2ad433 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 19 Jul 2026 14:13:37 +0200 Subject: [PATCH 19/19] use qkv_channels() instead of recomputing it by hand Also makes the distinction to attn_channels() (Nq*Hd, the o_proj input) obvious at the allocation sites. Co-Authored-By: Claude Fable 5 --- src/models/llama_model.cpp | 2 +- src/models/llama_weights.cpp | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/models/llama_model.cpp b/src/models/llama_model.cpp index 5303717..94b33c0 100644 --- a/src/models/llama_model.cpp +++ b/src/models/llama_model.cpp @@ -780,7 +780,7 @@ void LLamaModel::fill_block_shapes(GenericTensorContainer& target, const Transfo tgt.Sizes[1] = cols; }; - long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * HS; + long attn_intermediate_size = config.qkv_channels(); create(target.get_tensor(LLamaWeightID::QKV_W), attn_intermediate_size, C, matrix_dtype); create(target.get_tensor(LLamaWeightID::ATTO_W), C, config.attn_channels(), matrix_dtype); create(target.get_tensor(LLamaWeightID::UP_W), 2 * H, C, matrix_dtype); diff --git a/src/models/llama_weights.cpp b/src/models/llama_weights.cpp index 6841dda..756f0ec 100644 --- a/src/models/llama_weights.cpp +++ b/src/models/llama_weights.cpp @@ -27,7 +27,7 @@ void allocate_non_matrix_params(sLLamaBlockWeights& target, const Transformer target.KNorm_w = Tensor{}; } - long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * HS; + long attn_intermediate_size = config.qkv_channels(); if(config.UseQKVBias) { target.Attn_QKV_b = alloc.allocate_shard(dtype, shard_idx, num_shards, "att_qkv_b", {attn_intermediate_size}, kind); } else { @@ -40,7 +40,7 @@ void allocate_matrix_params(sLLamaBlockWeights& target, const TransformerConf long C = config.HiddenSize; long H = config.IntermediateSize; - long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * config.head_size(); + long attn_intermediate_size = config.qkv_channels(); target.Attn_QKV_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "att_qkv_w", {attn_intermediate_size, C}, kind); target.Attn_Out_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "attproj_w", {C, (long)config.attn_channels()}, kind); target.MLP_Up_w = alloc.allocate_shard(dtype, shard_idx, num_shards, "mlp_up_w", {2 * H, C}, kind); @@ -63,7 +63,7 @@ void fill_matrix_shapes(sLLamaBlockWeights& target, const Transform tgt.GlobalShape[1] = cols; }; - long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * config.head_size(); + long attn_intermediate_size = config.qkv_channels(); create_matrix_shard(target.Attn_QKV_w, attn_intermediate_size, C); create_matrix_shard(target.Attn_Out_w, C, config.attn_channels()); create_matrix_shard(target.MLP_Up_w, 2 * H, C); @@ -88,7 +88,7 @@ void fill_non_matrix_shapes(sLLamaBlockWeights& target, const Trans create_vector_shard(target.LN2_w, C); create_vector_shard(target.QNorm_w, config.UseQKNorm ? HS : 0); create_vector_shard(target.KNorm_w, config.UseQKNorm ? HS : 0); - long attn_intermediate_size = (config.NumQueryHeads + 2 * config.NumKeyValHeads) * HS; + long attn_intermediate_size = config.qkv_channels(); create_vector_shard(target.Attn_QKV_b, config.UseQKVBias ? attn_intermediate_size : 0); }