diff --git a/README.md b/README.md index 3e53f3f..f1bc14f 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ Paper: [SyclKittens: A Tile Programming Model for Programmers and Coding Agents - Intel oneAPI DPC++ compiler with `icpx` and SYCL 2020 support - Intel Data Center GPU Max (PVC) for the supplied AOT target - C++20 -- XPU-enabled PyTorch and pybind11 for Python extension kernels +- XPU-enabled PyTorch and pybind11 for Python extension kernels (`CONFIG=pytorch`) - oneMKL for the decode GEMV fallback - Level Zero development headers for the direct-copy all-gather and distributed GEMM targets; set @@ -78,10 +78,36 @@ make -C kernels/collective all_reduce make -C kernels/context_attention ring_overlap ``` -Python extension targets use `PYTHON` to select an XPU-enabled environment: +### Build modes + +Each kernel in `attention`, `gemm`, `norm`, `rotary`, `activation`, and `decode` +builds in one of two modes, selected with `CONFIG`: + +| `CONFIG` | Output | Needs | +|---|---|---| +| `standalone` (default) | executable; `main()` comes from `_harness.impl` | SYCL only | +| `pytorch` | Python extension (`sk_*.so`) | XPU PyTorch + pybind11 | + +A standalone run times the kernel and checks it against a CPU fp32 reference; +arguments are `[iters] [warmup]`, and the exit code is non-zero on failure: + +```bash +make -C kernels/attention forward +./kernels/attention/build/gqa_forward 100 10 +``` + +To build and run every standalone harness at once (summary at the end, non-zero +exit on any failure): + +```bash +./kernels/run_standalone.sh [iters] [warmup] +``` + +Python extension targets use `CONFIG=pytorch` and `PYTHON` to select an +XPU-enabled environment: ```bash -make -C kernels/norm all \ +make -C kernels/norm all CONFIG=pytorch \ PYTHON=/path/to/xpu/python3 RMSNORM_D=4096 LN_D=4096 ``` diff --git a/kernels/activation/Makefile b/kernels/activation/Makefile index e9f9cfc..7597aee 100644 --- a/kernels/activation/Makefile +++ b/kernels/activation/Makefile @@ -4,7 +4,13 @@ include ../common.mk all: silu_mul +ifeq ($(CONFIG),standalone) +silu_mul: check-sycl + @mkdir -p $(BUILD_DIR) + $(CXX) $(COMMON_SYCL_FLAGS) silu_mul.dp.cpp -o $(BUILD_DIR)/silu_mul +else silu_mul: check-python @mkdir -p $(BUILD_DIR) $(CXX) $(EXTENSION_FLAGS) -DTORCH_EXTENSION_NAME=sk_silu_mul silu_mul.dp.cpp \ -o $(BUILD_DIR)/sk_silu_mul$(EXT_SUFFIX) $(TORCH_LIBS) +endif diff --git a/kernels/activation/silu_mul.dp.cpp b/kernels/activation/silu_mul.dp.cpp index 1cf44cb..93da6da 100644 --- a/kernels/activation/silu_mul.dp.cpp +++ b/kernels/activation/silu_mul.dp.cpp @@ -28,10 +28,6 @@ #include -#include "pyutils/torch_helpers.dp.hpp" -#include -#include - #ifndef SILU_VEC #define SILU_VEC 8 #endif @@ -111,10 +107,15 @@ void launch_silu_mul( // NO queue.wait() — caller synchronizes via torch.xpu.synchronize() } +#ifdef TORCH_COMPILE // ============================================================ // PyTorch dispatch // ============================================================ +#include "pyutils/torch_helpers.dp.hpp" +#include +#include + torch::Tensor dispatch_silu_mul( torch::Tensor gate, // (B, N, inter), bf16 torch::Tensor up // (B, N, inter), bf16 @@ -153,3 +154,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Fused SiLU*mul: gate,up (bf16 [B,N,inter]) -> out (bf16 [B,N,inter])", py::arg("gate"), py::arg("up")); } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "silu_mul_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/activation/silu_mul_harness.impl b/kernels/activation/silu_mul_harness.impl new file mode 100644 index 0000000..a0108b9 --- /dev/null +++ b/kernels/activation/silu_mul_harness.impl @@ -0,0 +1,78 @@ +// Standalone harness for silu_mul.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of silu_mul.dp.cpp when TORCH_COMPILE is not defined, +// so bf16_t and launch_silu_mul() are already in scope. +// +// usage: silu_mul [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. + +#include +#include +#include +#include +#include +#include + +// Deterministic values in [-4, 4) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return (static_cast(state >> 8) / static_cast(1u << 23) - 1.0f) * 4.0f; +} + +static uint16_t to_bits(float f) { return sycl::bit_cast(bf16_t(f)); } +static float from_bits(uint16_t b) { return static_cast(sycl::bit_cast(b)); } + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + constexpr int INTER = 14336; + + sycl::queue queue{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s\n", + queue.get_device().get_info().c_str()); + + bool all_pass = true; + for (int T : {1, 2048, 8192}) { + const size_t total = size_t(T) * INTER; + std::vector h_gate(total), h_up(total), h_out(total); + uint32_t seed = 12345u + T; + for (auto &x : h_gate) x = to_bits(lcg_uniform(seed)); + for (auto &x : h_up) x = to_bits(lcg_uniform(seed)); + + uint16_t *d_gate = sycl::malloc_device(total, queue); + uint16_t *d_up = sycl::malloc_device(total, queue); + uint16_t *d_out = sycl::malloc_device(total, queue); + queue.memcpy(d_gate, h_gate.data(), total * sizeof(uint16_t)); + queue.memcpy(d_up, h_up.data(), total * sizeof(uint16_t)); + queue.wait(); + + for (int i = 0; i < warmup; i++) launch_silu_mul(queue, d_gate, d_up, d_out, total); + queue.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) launch_silu_mul(queue, d_gate, d_up, d_out, total); + queue.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + const double gbps = 3.0 * total * sizeof(uint16_t) / (ms * 1e-3) / 1e9; // r gate, r up, w out + + // Full CPU reference (fp32 math on the bf16 inputs). + queue.memcpy(h_out.data(), d_out, total * sizeof(uint16_t)).wait(); + double err2 = 0.0, ref2 = 0.0; + for (size_t i = 0; i < total; i++) { + const float g = from_bits(h_gate[i]), u = from_bits(h_up[i]); + const double ref = double(g / (1.0f + std::exp(-g)) * u); + const double d = double(from_bits(h_out[i])) - ref; + err2 += d * d; + ref2 += ref * ref; + } + const double rel = std::sqrt(err2 / (ref2 + 1e-30)); + const bool pass = rel < 1e-2; + all_pass &= pass; + std::printf("silu_mul T=%d inter=%d: %.4f ms %.1f GB/s (iters=%d warmup=%d) rel_l2=%.2e %s\n", + T, INTER, ms, gbps, iters, warmup, rel, pass ? "PASS" : "FAIL"); + + sycl::free(d_gate, queue); + sycl::free(d_up, queue); + sycl::free(d_out, queue); + } + return all_pass ? 0 : 1; +} diff --git a/kernels/attention/Makefile b/kernels/attention/Makefile index 0c71715..fd30185 100644 --- a/kernels/attention/Makefile +++ b/kernels/attention/Makefile @@ -7,22 +7,33 @@ BWD_H ?= 16 BWD_H_KV ?= 16 BWD_D ?= 128 +FWD_DEFS := -DATTN_H_VALUE=$(FWD_H) -DATTN_H_KV_VALUE=$(FWD_H_KV) -DATTN_D_VALUE=$(FWD_D) +BWD_DEFS := -DATTN_H_VALUE=$(BWD_H) -DATTN_H_KV_VALUE=$(BWD_H_KV) -DATTN_D_VALUE=$(BWD_D) \ + -ftarget-register-alloc-mode=pvc:large + .PHONY: all forward backward all: forward backward +ifeq ($(CONFIG),standalone) +forward: check-sycl + @mkdir -p $(BUILD_DIR) + $(CXX) $(COMMON_SYCL_FLAGS) $(FWD_DEFS) gqa_forward.dp.cpp -o $(BUILD_DIR)/gqa_forward + +backward: check-sycl + @mkdir -p $(BUILD_DIR) + $(CXX) $(COMMON_SYCL_FLAGS) $(BWD_DEFS) gqa_backward.dp.cpp -o $(BUILD_DIR)/gqa_backward +else forward: check-python @mkdir -p $(BUILD_DIR) - $(CXX) $(EXTENSION_FLAGS) \ - -DATTN_H_VALUE=$(FWD_H) -DATTN_H_KV_VALUE=$(FWD_H_KV) -DATTN_D_VALUE=$(FWD_D) \ + $(CXX) $(EXTENSION_FLAGS) $(FWD_DEFS) \ -DTORCH_EXTENSION_NAME=sk_gqa_forward gqa_forward.dp.cpp \ -o $(BUILD_DIR)/sk_gqa_forward$(EXT_SUFFIX) $(TORCH_LIBS) backward: check-python @mkdir -p $(BUILD_DIR) - $(CXX) $(EXTENSION_FLAGS) \ - -DATTN_H_VALUE=$(BWD_H) -DATTN_H_KV_VALUE=$(BWD_H_KV) -DATTN_D_VALUE=$(BWD_D) \ - -ftarget-register-alloc-mode=pvc:large \ + $(CXX) $(EXTENSION_FLAGS) $(BWD_DEFS) \ -DTORCH_EXTENSION_NAME=sk_gqa_backward gqa_backward.dp.cpp \ -o $(BUILD_DIR)/sk_gqa_backward$(EXT_SUFFIX) \ $(TORCH_LIBS) +endif diff --git a/kernels/attention/gqa_backward.dp.cpp b/kernels/attention/gqa_backward.dp.cpp index 2d3d9fa..012226d 100644 --- a/kernels/attention/gqa_backward.dp.cpp +++ b/kernels/attention/gqa_backward.dp.cpp @@ -11,10 +11,6 @@ #include "kittens.dp.hpp" #include "ops/warp/register/tile/conversions.dp.hpp" -#include "pyutils/torch_helpers.dp.hpp" - -#include -#include using namespace kittens; // ============================================================ @@ -370,21 +366,13 @@ SYCL_EXTERNAL void attend_bwd_fused_ker( } // ============================================================ -// Dispatch +// Shared raw-pointer entries (both builds). All tensors BNHD except +// L/delta [B,H,N] f32. Q/K/V/O/dO bf16; dQ/dK/dV f32. // ============================================================ -void dispatch_prep(torch::Tensor O, torch::Tensor dO, torch::Tensor delta) { - CHECK_INPUT(O); CHECK_INPUT(dO); CHECK_INPUT(delta); - const int B = O.size(0), N = O.size(1), H = O.size(2); - - auto stream = c10::xpu::getCurrentXPUStream(O.device().index()); - auto &queue = stream.queue(); - queue.wait(); - - bf16 *d_O_bf = reinterpret_cast(O.data_ptr()); - bf16 *d_dO_bf = reinterpret_cast(dO.data_ptr()); - float *d_delta = delta.data_ptr(); - +// delta[b,h,q] = rowsum(dO * O). Blocks until done. +static void run_prep(sycl::queue &queue, bf16 *d_O_bf, bf16 *d_dO_bf, + float *d_delta, int B, int N, int H) { bwd_gl Og(d_O_bf, B, N, H, nullptr); bwd_gl dOg(d_dO_bf, B, N, H, nullptr); @@ -407,43 +395,13 @@ void dispatch_prep(torch::Tensor O, torch::Tensor dO, torch::Tensor delta) { queue.wait(); } -// Allocate the P^T / dS^T reorder scratch: [B, kv_blocks*128, H_KV, 128] bf16. -static torch::Tensor make_reorder_scratch(int B, int N, int H_KV, torch::Tensor like) { - auto opts = torch::TensorOptions().dtype(torch::kBFloat16).device(like.device()); - return torch::empty({(long)B, (long)N, (long)H_KV, (long)ATTN_D}, opts); -} - -void dispatch_bwd(torch::Tensor Q, torch::Tensor K, torch::Tensor V, - torch::Tensor dO, torch::Tensor L, torch::Tensor delta, - torch::Tensor dQ, torch::Tensor dK, torch::Tensor dV) { - CHECK_INPUT(Q); CHECK_INPUT(K); CHECK_INPUT(V); CHECK_INPUT(dO); - CHECK_INPUT(L); CHECK_INPUT(delta); - CHECK_INPUT(dQ); CHECK_INPUT(dK); CHECK_INPUT(dV); - - const int B = Q.size(0), N = Q.size(1), H = Q.size(2), H_KV = K.size(2); - - auto stream = c10::xpu::getCurrentXPUStream(Q.device().index()); - auto &queue = stream.queue(); - queue.wait(); - - bf16 *d_Q = reinterpret_cast(Q.data_ptr()); - bf16 *d_K = reinterpret_cast(K.data_ptr()); - bf16 *d_V = reinterpret_cast(V.data_ptr()); - bf16 *d_dO = reinterpret_cast(dO.data_ptr()); - float *d_L = L.data_ptr(); - float *d_delta = delta.data_ptr(); - float *d_dQ = dQ.data_ptr(); - float *d_dK = dK.data_ptr(); - float *d_dV = dV.data_ptr(); - - // Reorder scratch (bf16 P^T and dS^T), shaped [B, N(kv), H_KV, D(q=128)]. - torch::Tensor Pt = make_reorder_scratch(B, N, H_KV, Q); - torch::Tensor dSt = make_reorder_scratch(B, N, H_KV, Q); - torch::Tensor dSq = make_reorder_scratch(B, N, H_KV, Q); // dQ q-major scratch - bf16 *d_Pt = reinterpret_cast(Pt.data_ptr()); - bf16 *d_dSt = reinterpret_cast(dSt.data_ptr()); - bf16 *d_dSq = reinterpret_cast(dSq.data_ptr()); - +// Fused causal dV/dK/dQ. d_Pt/d_dSt/d_dSq are caller-owned bf16 reorder scratch +// buffers of B*N*H_KV*ATTN_D elements each. Blocks until done. +static void run_bwd(sycl::queue &queue, bf16 *d_Q, bf16 *d_K, bf16 *d_V, + bf16 *d_dO, float *d_L, float *d_delta, + float *d_dQ, float *d_dK, float *d_dV, + bf16 *d_Pt, bf16 *d_dSt, bf16 *d_dSq, + int B, int N, int H, int H_KV) { bwd_gl Qg(d_Q, B, N, H, nullptr); bwd_gl Kg(d_K, B, N, H_KV, nullptr); bwd_gl Vg(d_V, B, N, H_KV, nullptr); @@ -476,7 +434,70 @@ void dispatch_bwd(torch::Tensor Q, torch::Tensor K, torch::Tensor V, } } +#ifdef TORCH_COMPILE +// ============================================================ +// PyTorch dispatch +// ============================================================ +#include "pyutils/torch_helpers.dp.hpp" +#include +#include + +void dispatch_prep(torch::Tensor O, torch::Tensor dO, torch::Tensor delta) { + CHECK_INPUT(O); CHECK_INPUT(dO); CHECK_INPUT(delta); + const int B = O.size(0), N = O.size(1), H = O.size(2); + + auto stream = c10::xpu::getCurrentXPUStream(O.device().index()); + auto &queue = stream.queue(); + queue.wait(); + + run_prep(queue, reinterpret_cast(O.data_ptr()), + reinterpret_cast(dO.data_ptr()), + delta.data_ptr(), B, N, H); +} + +// Allocate the P^T / dS^T reorder scratch: [B, kv_blocks*128, H_KV, 128] bf16. +static torch::Tensor make_reorder_scratch(int B, int N, int H_KV, torch::Tensor like) { + auto opts = torch::TensorOptions().dtype(torch::kBFloat16).device(like.device()); + return torch::empty({(long)B, (long)N, (long)H_KV, (long)ATTN_D}, opts); +} + +void dispatch_bwd(torch::Tensor Q, torch::Tensor K, torch::Tensor V, + torch::Tensor dO, torch::Tensor L, torch::Tensor delta, + torch::Tensor dQ, torch::Tensor dK, torch::Tensor dV) { + CHECK_INPUT(Q); CHECK_INPUT(K); CHECK_INPUT(V); CHECK_INPUT(dO); + CHECK_INPUT(L); CHECK_INPUT(delta); + CHECK_INPUT(dQ); CHECK_INPUT(dK); CHECK_INPUT(dV); + + const int B = Q.size(0), N = Q.size(1), H = Q.size(2), H_KV = K.size(2); + + auto stream = c10::xpu::getCurrentXPUStream(Q.device().index()); + auto &queue = stream.queue(); + queue.wait(); + + // Reorder scratch (bf16 P^T and dS^T), shaped [B, N(kv), H_KV, D(q=128)]. + torch::Tensor Pt = make_reorder_scratch(B, N, H_KV, Q); + torch::Tensor dSt = make_reorder_scratch(B, N, H_KV, Q); + torch::Tensor dSq = make_reorder_scratch(B, N, H_KV, Q); // dQ q-major scratch + + run_bwd(queue, + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K.data_ptr()), + reinterpret_cast(V.data_ptr()), + reinterpret_cast(dO.data_ptr()), + L.data_ptr(), delta.data_ptr(), + dQ.data_ptr(), dK.data_ptr(), dV.data_ptr(), + reinterpret_cast(Pt.data_ptr()), + reinterpret_cast(dSt.data_ptr()), + reinterpret_cast(dSq.data_ptr()), + B, N, H, H_KV); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("dispatch_prep", &dispatch_prep); m.def("dispatch_bwd", &dispatch_bwd); } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "gqa_backward_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/attention/gqa_backward_harness.impl b/kernels/attention/gqa_backward_harness.impl new file mode 100644 index 0000000..640742a --- /dev/null +++ b/kernels/attention/gqa_backward_harness.impl @@ -0,0 +1,167 @@ +// Standalone harness for gqa_backward.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of gqa_backward.dp.cpp when TORCH_COMPILE is not defined, +// so ATTN_H/ATTN_H_KV/ATTN_D, BLOCK_N and run_prep()/run_bwd() are in scope. +// +// usage: gqa_backward [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. +// +// The kernel is causal-only. Runs B=1, N=256 (multiple of BLOCK_N) so a full +// fp32 CPU reference for delta, dQ, dK, dV is cheap; reports rel-L2 for each. + +#include +#include +#include +#include +#include +#include +#include + +// Deterministic values in [-1, 1) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} + +static double rel_l2(const std::vector &got, const std::vector &ref) { + double e = 0.0, r = 0.0; + for (size_t i = 0; i < ref.size(); i++) { + e += (got[i] - ref[i]) * (got[i] - ref[i]); + r += ref[i] * ref[i]; + } + return std::sqrt(e / (r + 1e-30)); +} + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + constexpr int B = 1, N = 256, H = ATTN_H, HKV = ATTN_H_KV, D = ATTN_D; + constexpr int GROUP = H / HKV; + static_assert(N % BLOCK_N == 0, "N must be a multiple of BLOCK_N"); + + sycl::queue queue{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s\n", + queue.get_device().get_info().c_str()); + + const size_t size_q = size_t(N) * H * D, size_kv = size_t(N) * HKV * D; + std::vector h_q(size_q), h_k(size_kv), h_v(size_kv), h_do(size_q), h_o(size_q); + uint32_t seed = 12345u; + for (auto &x : h_q) x = bf16(lcg_uniform(seed)); + for (auto &x : h_k) x = bf16(lcg_uniform(seed)); + for (auto &x : h_v) x = bf16(lcg_uniform(seed)); + for (auto &x : h_do) x = bf16(lcg_uniform(seed)); + + auto q_at = [&](int n, int h, int d) { return double(float(h_q[(size_t(n) * H + h) * D + d])); }; + auto k_at = [&](int n, int h, int d) { return double(float(h_k[(size_t(n) * HKV + h) * D + d])); }; + auto v_at = [&](int n, int h, int d) { return double(float(h_v[(size_t(n) * HKV + h) * D + d])); }; + auto do_at = [&](int n, int h, int d) { return double(float(h_do[(size_t(n) * H + h) * D + d])); }; + + // ---- CPU forward: P (causal softmax of S/sqrt(D)), O (bf16, fed to the kernel), L. + const double scale = 1.0 / std::sqrt(double(D)); + std::vector P(size_t(H) * N * N, 0.0); + std::vector h_l(size_t(H) * N); + for (int h = 0; h < H; h++) { + const int kvh = h / GROUP; + for (int i = 0; i < N; i++) { + double *p = &P[(size_t(h) * N + i) * N]; + double m = -1e300; + for (int j = 0; j <= i; j++) { + double dot = 0.0; + for (int d = 0; d < D; d++) dot += q_at(i, h, d) * k_at(j, kvh, d); + p[j] = dot * scale; + m = std::max(m, p[j]); + } + double sum = 0.0; + for (int j = 0; j <= i; j++) { p[j] = std::exp(p[j] - m); sum += p[j]; } + for (int j = 0; j <= i; j++) p[j] /= sum; + h_l[size_t(h) * N + i] = float(m + std::log(sum)); + for (int d = 0; d < D; d++) { + double o = 0.0; + for (int j = 0; j <= i; j++) o += p[j] * v_at(j, kvh, d); + h_o[(size_t(i) * H + h) * D + d] = bf16(float(o)); + } + } + } + auto o_at = [&](int n, int h, int d) { return double(float(h_o[(size_t(n) * H + h) * D + d])); }; + + // ---- CPU backward reference. + std::vector delta_ref(size_t(H) * N), dq_ref(size_q, 0.0), dk_ref(size_kv, 0.0), dv_ref(size_kv, 0.0); + std::vector dS(N); + for (int h = 0; h < H; h++) { + const int kvh = h / GROUP; + for (int i = 0; i < N; i++) { + double dl = 0.0; + for (int d = 0; d < D; d++) dl += do_at(i, h, d) * o_at(i, h, d); + delta_ref[size_t(h) * N + i] = dl; + const double *p = &P[(size_t(h) * N + i) * N]; + for (int j = 0; j <= i; j++) { + double dp = 0.0; + for (int d = 0; d < D; d++) dp += do_at(i, h, d) * v_at(j, kvh, d); + dS[j] = p[j] * (dp - dl); + } + for (int j = 0; j <= i; j++) + for (int d = 0; d < D; d++) { + dq_ref[(size_t(i) * H + h) * D + d] += dS[j] * k_at(j, kvh, d) * scale; + dk_ref[(size_t(j) * HKV + kvh) * D + d] += dS[j] * q_at(i, h, d) * scale; + dv_ref[(size_t(j) * HKV + kvh) * D + d] += p[j] * do_at(i, h, d); + } + } + } + + // ---- GPU. + bf16 *d_q = sycl::malloc_device(size_q, queue); + bf16 *d_k = sycl::malloc_device(size_kv, queue); + bf16 *d_v = sycl::malloc_device(size_kv, queue); + bf16 *d_do = sycl::malloc_device(size_q, queue); + bf16 *d_o = sycl::malloc_device(size_q, queue); + float *d_l = sycl::malloc_device(size_t(H) * N, queue); + float *d_delta = sycl::malloc_device(size_t(H) * N, queue); + float *d_dq = sycl::malloc_device(size_q, queue); + float *d_dk = sycl::malloc_device(size_kv, queue); + float *d_dv = sycl::malloc_device(size_kv, queue); + bf16 *d_pt = sycl::malloc_device(size_kv, queue); // reorder scratch, B*N*HKV*D each + bf16 *d_dst = sycl::malloc_device(size_kv, queue); + bf16 *d_dsq = sycl::malloc_device(size_kv, queue); + queue.memcpy(d_q, h_q.data(), size_q * sizeof(bf16)); + queue.memcpy(d_k, h_k.data(), size_kv * sizeof(bf16)); + queue.memcpy(d_v, h_v.data(), size_kv * sizeof(bf16)); + queue.memcpy(d_do, h_do.data(), size_q * sizeof(bf16)); + queue.memcpy(d_o, h_o.data(), size_q * sizeof(bf16)); + queue.memcpy(d_l, h_l.data(), h_l.size() * sizeof(float)); + queue.wait(); + + auto step = [&] { + run_prep(queue, d_o, d_do, d_delta, B, N, H); + run_bwd(queue, d_q, d_k, d_v, d_do, d_l, d_delta, d_dq, d_dk, d_dv, + d_pt, d_dst, d_dsq, B, N, H, HKV); + }; + for (int i = 0; i < warmup; i++) step(); + queue.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) step(); + queue.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + // Backward ~ 2.5x the causal forward FLOPs (same convention as ThunderKittens). + const double flops = 2.5 * 4.0 * B * H * double(N) * N * D * 0.5; + + std::vector g_delta(size_t(H) * N), g_dq(size_q), g_dk(size_kv), g_dv(size_kv); + queue.memcpy(g_delta.data(), d_delta, g_delta.size() * sizeof(float)); + queue.memcpy(g_dq.data(), d_dq, size_q * sizeof(float)); + queue.memcpy(g_dk.data(), d_dk, size_kv * sizeof(float)); + queue.memcpy(g_dv.data(), d_dv, size_kv * sizeof(float)); + queue.wait(); + + const double e_delta = rel_l2(g_delta, delta_ref), e_dq = rel_l2(g_dq, dq_ref); + const double e_dk = rel_l2(g_dk, dk_ref), e_dv = rel_l2(g_dv, dv_ref); + const bool pass = e_delta < 1e-2 && e_dq < 3e-2 && e_dk < 3e-2 && e_dv < 3e-2; + std::printf("GQA bwd B%d N%d H%d/%d D%d causal (prep+bwd): %.3f ms %.1f TFLOP/s " + "(iters=%d warmup=%d) rel_l2 delta=%.2e dQ=%.2e dK=%.2e dV=%.2e %s\n", + B, N, H, HKV, D, ms, flops / (ms * 1e-3) / 1e12, iters, warmup, + e_delta, e_dq, e_dk, e_dv, pass ? "PASS" : "FAIL"); + + for (void *p : {(void *)d_q, (void *)d_k, (void *)d_v, (void *)d_do, (void *)d_o, (void *)d_l, + (void *)d_delta, (void *)d_dq, (void *)d_dk, (void *)d_dv, + (void *)d_pt, (void *)d_dst, (void *)d_dsq}) + sycl::free(p, queue); + return pass ? 0 : 1; +} diff --git a/kernels/attention/gqa_forward.dp.cpp b/kernels/attention/gqa_forward.dp.cpp index 0841dd4..9592613 100644 --- a/kernels/attention/gqa_forward.dp.cpp +++ b/kernels/attention/gqa_forward.dp.cpp @@ -36,10 +36,6 @@ #include "kittens.dp.hpp" #include "ops/warp/register/tile/conversions.dp.hpp" -#include "pyutils/torch_helpers.dp.hpp" - -#include -#include using namespace kittens; @@ -467,6 +463,35 @@ static inline void route_fwd(sycl::queue &queue, #undef FWD_ROUTE } +// ============================================================ +// Shared raw-pointer entry (both builds). Q,K,V: bf16 BNHD. O: bf16 (o_bf16) +// or f32 BNHD. L: f32 [B,H,N] logsumexp. Enqueues only; the caller syncs. +// ============================================================ +static void run_fwd(sycl::queue &queue, bf16 *d_Q, bf16 *d_K, bf16 *d_V, + void *d_O, bool o_bf16, float *d_L, int B, int N, + bool is_causal) { + global_layout Qg(d_Q, B, N, ATTN_H, nullptr); + global_layout Kg(d_K, B, N, ATTN_H_KV, nullptr); + global_layout Vg(d_V, B, N, ATTN_H_KV, nullptr); + + // bf16 output (E2E inference): the store epilogue converts the float softmax + // accumulators -> bf16 in-register, so the composed path no longer needs a + // separate f32->bf16 elementwise cast (+ HBM round-trip) per decoder layer. + // f32 output (training / bwd recompute) keeps the exact legacy behaviour. + if (o_bf16) { + global_layout_outh Og(static_cast(d_O), B, N, ATTN_H, nullptr); + route_fwd(queue, Qg, Kg, Vg, Og, d_L, B, ATTN_H, N, is_causal); + } else { + global_layout_out Og(static_cast(d_O), B, N, ATTN_H, nullptr); + route_fwd(queue, Qg, Kg, Vg, Og, d_L, B, ATTN_H, N, is_causal); + } +} + +#ifdef TORCH_COMPILE +#include "pyutils/torch_helpers.dp.hpp" +#include +#include + void dispatch_fwd(torch::Tensor Q, torch::Tensor K, torch::Tensor V, torch::Tensor O, torch::Tensor L, bool is_causal) { CHECK_INPUT(Q); CHECK_INPUT(K); CHECK_INPUT(V); @@ -490,27 +515,12 @@ void dispatch_fwd(torch::Tensor Q, torch::Tensor K, torch::Tensor V, auto &queue = stream.queue(); queue.wait(); - bf16 *d_Q = reinterpret_cast(Q.data_ptr()); - bf16 *d_K = reinterpret_cast(K.data_ptr()); - bf16 *d_V = reinterpret_cast(V.data_ptr()); - float *d_L = L.data_ptr(); - - global_layout Qg(d_Q, B, N, H, nullptr); - global_layout Kg(d_K, B, N, H_KV, nullptr); - global_layout Vg(d_V, B, N, H_KV, nullptr); - - // bf16 output (E2E inference): the store epilogue converts the float softmax - // accumulators -> bf16 in-register, so the composed path no longer needs a - // separate f32->bf16 elementwise cast (+ HBM round-trip) per decoder layer. - // f32 output (training / bwd recompute) keeps the exact legacy behaviour. - if (O.dtype() == torch::kBFloat16) { - global_layout_outh Og( - reinterpret_cast(O.data_ptr()), B, N, H, nullptr); - route_fwd(queue, Qg, Kg, Vg, Og, d_L, B, H, N, is_causal); - } else { - global_layout_out Og(O.data_ptr(), B, N, H, nullptr); - route_fwd(queue, Qg, Kg, Vg, Og, d_L, B, H, N, is_causal); - } + run_fwd(queue, + reinterpret_cast(Q.data_ptr()), + reinterpret_cast(K.data_ptr()), + reinterpret_cast(V.data_ptr()), + O.data_ptr(), O.dtype() == torch::kBFloat16, L.data_ptr(), + B, N, is_causal); } // ============================================================ @@ -524,3 +534,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("Q"), py::arg("K"), py::arg("V"), py::arg("O"), py::arg("L"), py::arg("is_causal") = true); } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "gqa_forward_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/attention/gqa_forward_harness.impl b/kernels/attention/gqa_forward_harness.impl new file mode 100644 index 0000000..1de1784 --- /dev/null +++ b/kernels/attention/gqa_forward_harness.impl @@ -0,0 +1,116 @@ +// Standalone harness for gqa_forward.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of gqa_forward.dp.cpp when TORCH_COMPILE is not defined, +// so ATTN_H/ATTN_H_KV/ATTN_D and run_fwd() are already in scope. +// +// usage: gqa_forward [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. +// +// Runs B=1, N=4096 causal and non-causal (bf16 out) and checks sampled output +// rows and their logsumexp L against an fp32 CPU reference. + +#include +#include +#include +#include +#include +#include +#include + +// Deterministic values in [-1, 1) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + constexpr int B = 1, N = 4096, H = ATTN_H, HKV = ATTN_H_KV, D = ATTN_D; + constexpr int GROUP = H / HKV; + + sycl::queue queue{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s\n", + queue.get_device().get_info().c_str()); + + const size_t size_q = size_t(B) * N * H * D, size_kv = size_t(B) * N * HKV * D; + std::vector h_q(size_q), h_k(size_kv), h_v(size_kv); + uint32_t seed = 12345u; + for (auto &x : h_q) x = bf16(lcg_uniform(seed)); + for (auto &x : h_k) x = bf16(lcg_uniform(seed)); + for (auto &x : h_v) x = bf16(lcg_uniform(seed)); + + bf16 *d_q = sycl::malloc_device(size_q, queue); + bf16 *d_k = sycl::malloc_device(size_kv, queue); + bf16 *d_v = sycl::malloc_device(size_kv, queue); + bf16 *d_o = sycl::malloc_device(size_q, queue); + float *d_l = sycl::malloc_device(size_t(B) * H * N, queue); + queue.memcpy(d_q, h_q.data(), size_q * sizeof(bf16)); + queue.memcpy(d_k, h_k.data(), size_kv * sizeof(bf16)); + queue.memcpy(d_v, h_v.data(), size_kv * sizeof(bf16)); + queue.wait(); + + auto q_at = [&](int n, int h, int d) { return float(h_q[(size_t(n) * H + h) * D + d]); }; + auto k_at = [&](int n, int h, int d) { return float(h_k[(size_t(n) * HKV + h) * D + d]); }; + auto v_at = [&](int n, int h, int d) { return float(h_v[(size_t(n) * HKV + h) * D + d]); }; + + bool all_pass = true; + for (bool causal : {true, false}) { + for (int i = 0; i < warmup; i++) + run_fwd(queue, d_q, d_k, d_v, d_o, true, d_l, B, N, causal); + queue.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) + run_fwd(queue, d_q, d_k, d_v, d_o, true, d_l, B, N, causal); + queue.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + const double flops = 4.0 * B * H * double(N) * N * D * (causal ? 0.5 : 1.0); + + std::vector h_o(size_q); + std::vector h_l(size_t(B) * H * N); + queue.memcpy(h_o.data(), d_o, size_q * sizeof(bf16)); + queue.memcpy(h_l.data(), d_l, h_l.size() * sizeof(float)); + queue.wait(); + + // Sampled (head, query) rows, incl. the first and last query. + double max_err = 0.0, err2 = 0.0, ref2 = 0.0, max_l_err = 0.0; + const int rows[][2] = {{0, 0}, {0, N - 1}, {H - 1, 17}, {H / 2, N / 2}, + {1 % H, 1023}, {H - 1, N - 1}, {GROUP % H, 2049}, {H / 3, 3001}}; + std::vector s(N), o(D); + for (const auto &rq : rows) { + const int h = rq[0], qi = rq[1], kvh = h / GROUP; + const int n_kv = causal ? qi + 1 : N; + double m = -1e300; + for (int j = 0; j < n_kv; j++) { + double dot = 0.0; + for (int d = 0; d < D; d++) dot += double(q_at(qi, h, d)) * k_at(j, kvh, d); + s[j] = dot / std::sqrt(double(D)); + m = std::max(m, s[j]); + } + double sum = 0.0; + for (int j = 0; j < n_kv; j++) { s[j] = std::exp(s[j] - m); sum += s[j]; } + std::fill(o.begin(), o.end(), 0.0); + for (int j = 0; j < n_kv; j++) + for (int d = 0; d < D; d++) o[d] += s[j] / sum * v_at(j, kvh, d); + for (int d = 0; d < D; d++) { + const double got = float(h_o[(size_t(qi) * H + h) * D + d]); + max_err = std::max(max_err, std::fabs(got - o[d])); + err2 += (got - o[d]) * (got - o[d]); + ref2 += o[d] * o[d]; + } + const double l_ref = m + std::log(sum); + max_l_err = std::max(max_l_err, std::fabs(double(h_l[size_t(h) * N + qi]) - l_ref)); + } + const double rel = std::sqrt(err2 / (ref2 + 1e-30)); + const bool pass = max_err < 0.05 && rel < 2e-2 && max_l_err < 2e-2; + all_pass &= pass; + std::printf("GQA fwd B%d N%d H%d/%d D%d %s: %.3f ms %.1f TFLOP/s " + "(iters=%d warmup=%d) max_err=%.4f rel_l2=%.2e L_err=%.2e %s\n", + B, N, H, HKV, D, causal ? "causal" : "full", ms, flops / (ms * 1e-3) / 1e12, + iters, warmup, max_err, rel, max_l_err, pass ? "PASS" : "FAIL"); + } + + sycl::free(d_q, queue); sycl::free(d_k, queue); sycl::free(d_v, queue); + sycl::free(d_o, queue); sycl::free(d_l, queue); + return all_pass ? 0 : 1; +} diff --git a/kernels/common.mk b/kernels/common.mk index cb8ffa8..1b27f76 100644 --- a/kernels/common.mk +++ b/kernels/common.mk @@ -10,6 +10,14 @@ GPU_TARGET ?= intel_gpu_pvc BUILD_DIR ?= build SYCL_TARGET_FLAGS ?= -fsycl-targets=$(GPU_TARGET) +# Build mode for kernels that ship a harness (_harness.impl): +# standalone : executable whose main() comes from the harness; SYCL only, no torch (default) +# pytorch : Python extension built with -DTORCH_COMPILE +CONFIG ?= standalone +ifeq ($(filter $(CONFIG),pytorch standalone),) +$(error CONFIG must be pytorch or standalone, got '$(CONFIG)') +endif + SPIRV_FLAGS := -Xspirv-translator -spirv-ext=+SPV_INTEL_split_barrier,+SPV_INTEL_2d_block_io,+SPV_INTEL_subgroup_matrix_multiply_accumulate COMMON_SYCL_FLAGS := -fsycl $(SYCL_TARGET_FLAGS) -std=c++20 -O3 -fsycl-unnamed-lambda -Wno-deprecated-declarations $(SPIRV_FLAGS) -DKITTENS_INTEL -DKITTENS_XE -DKITTENS_FAST_BF16_CONVERT -I$(SK_ROOT)/include diff --git a/kernels/decode/Makefile b/kernels/decode/Makefile index 15db978..7c33632 100644 --- a/kernels/decode/Makefile +++ b/kernels/decode/Makefile @@ -4,7 +4,13 @@ include ../common.mk all: decode +ifeq ($(CONFIG),standalone) +decode: check-sycl + @mkdir -p $(BUILD_DIR) + $(CXX) $(COMMON_SYCL_FLAGS) decode.dp.cpp -o $(BUILD_DIR)/decode -qmkl +else decode: check-python @mkdir -p $(BUILD_DIR) $(CXX) $(EXTENSION_FLAGS) -DTORCH_EXTENSION_NAME=sk_decode decode.dp.cpp \ - -o $(BUILD_DIR)/sk_decode$(EXT_SUFFIX) $(TORCH_LIBS) -qmkl \ No newline at end of file + -o $(BUILD_DIR)/sk_decode$(EXT_SUFFIX) $(TORCH_LIBS) -qmkl +endif diff --git a/kernels/decode/decode.dp.cpp b/kernels/decode/decode.dp.cpp index 0b20de4..fe9d540 100644 --- a/kernels/decode/decode.dp.cpp +++ b/kernels/decode/decode.dp.cpp @@ -19,8 +19,10 @@ * at the end. Kernels are simple (decode is memory-bound; correctness > peak). */ +#ifdef TORCH_COMPILE #include #include +#endif #include #include @@ -29,9 +31,15 @@ using sbf16 = sycl::ext::oneapi::bfloat16; +#ifdef TORCH_COMPILE static inline sycl::queue &tq(int d) { return c10::xpu::getCurrentXPUStream(d).queue(); } +#else +// Standalone build: the harness owns one in-order queue per tile. +static std::vector g_queues; +static inline sycl::queue &tq(int d) { return g_queues[d]; } +#endif struct PtrPack { sbf16 *p[12]; }; @@ -369,28 +377,10 @@ struct LayerW { }; static std::vector g_layers; -static void register_layer(int li, int nl, - std::vector norm_in, std::vector Wqkv, - std::vector Wo, std::vector norm_post, - std::vector Wgate, std::vector Wup, - std::vector Wdown, std::vector kcache, - std::vector vcache) -{ +// Store one layer's per-tile weight / KV pointers (raw-pointer core of register_layer). +static void set_layer(int li, int nl, const LayerW &lw) { if ((int)g_layers.size() != nl) g_layers.assign(nl, LayerW{}); - const int N = g_ctx.n_gpus; - LayerW &lw = g_layers[li]; - for (int d = 0; d < N; d++) { - lw.norm_in[d] = reinterpret_cast(norm_in[d].data_ptr()); - lw.Wqkv[d] = reinterpret_cast(Wqkv[d].data_ptr()); - lw.Wo[d] = reinterpret_cast(Wo[d].data_ptr()); - lw.norm_post[d] = reinterpret_cast(norm_post[d].data_ptr()); - lw.Wgate[d] = reinterpret_cast(Wgate[d].data_ptr()); - lw.Wup[d] = reinterpret_cast(Wup[d].data_ptr()); - lw.Wdown[d] = reinterpret_cast(Wdown[d].data_ptr()); - lw.kc[d] = reinterpret_cast(kcache[d].data_ptr()); - lw.vc[d] = reinterpret_cast(vcache[d].data_ptr()); - } - lw.inter = Wgate[0].numel() / g_ctx.hidden; + g_layers[li] = lw; } static inline void do_layer(const LayerW &lw, sbf16 **xptr, sbf16 **cptr, @@ -463,23 +453,18 @@ static inline void do_layer(const LayerW &lw, sbf16 **xptr, sbf16 **cptr, static sbf16 *g_xbuf[12], *g_cosbuf[12], *g_sinbuf[12]; static bool g_io_registered = false; -static void register_io(std::vector xbuf, - std::vector cosbuf, - std::vector sinbuf) { - const int N = g_ctx.n_gpus; - for (int d = 0; d < N; d++) { - g_xbuf[d] = reinterpret_cast(xbuf[d].data_ptr()); - g_cosbuf[d] = reinterpret_cast(cosbuf[d].data_ptr()); - g_sinbuf[d] = reinterpret_cast(sinbuf[d].data_ptr()); +// Store the per-tile bf16 x / cos / sin buffers (raw-pointer core of register_io). +static void set_io(sbf16 *const *x, sbf16 *const *cos, sbf16 *const *sin) { + for (int d = 0; d < g_ctx.n_gpus; d++) { + g_xbuf[d] = x[d]; g_cosbuf[d] = cos[d]; g_sinbuf[d] = sin[d]; } g_io_registered = true; } -static void decode_step(int pos, int nq, int nkv, double eps) { - TORCH_CHECK(g_ctx.initialized && !g_layers.empty() && g_io_registered, - "sk_decode: init / register_layer / register_io first"); +// One decode step through all registered layers (x updated in place). Enqueues only. +static void run_layers(int pos, int nq, int nkv, float eps) { for (size_t li = 0; li < g_layers.size(); li++) - do_layer(g_layers[li], g_xbuf, g_cosbuf, g_sinbuf, pos, nq, nkv, (float)eps); + do_layer(g_layers[li], g_xbuf, g_cosbuf, g_sinbuf, pos, nq, nkv, eps); } // ============================================================ @@ -581,15 +566,13 @@ static inline void do_layer_2d(const LayerW &lw, sbf16 **xptr, sbf16 **cptr, } } -static void decode_2d_step(int pos, int nq, int nkv, double eps) { - TORCH_CHECK(g_ctx.initialized && !g_layers.empty() && g_io_registered && g_tp > 0, - "sk_decode: init / register_layer / register_io / set_2d first"); +static void run_layers_2d(int pos, int nq, int nkv, float eps) { for (size_t li = 0; li < g_layers.size(); li++) - do_layer_2d(g_layers[li], g_xbuf, g_cosbuf, g_sinbuf, pos, nq, nkv, (float)eps); + do_layer_2d(g_layers[li], g_xbuf, g_cosbuf, g_sinbuf, pos, nq, nkv, eps); } static void py_set_2d(int tp, int cp) { g_tp = tp; g_cp = cp; } -// ---- python interface ---- +// ---- setup / teardown (shared by both builds) ---- static void py_init(int n_gpus, int hidden, int inter_max, int qkv_max, int nq_max, int hd, int max_ctx) { g_ctx.init(n_gpus, hidden, inter_max, qkv_max, nq_max, hd, max_ctx); @@ -598,6 +581,55 @@ static void py_init(int n_gpus, int hidden, int inter_max, int qkv_max, static void py_cleanup() { g_ctx.cleanup(); g_layers.clear(); g_io_registered = false; } static bool py_is_initialized() { return g_ctx.initialized; } +#ifdef TORCH_COMPILE +// ---- python interface: tensors -> raw pointers, then the shared core ---- +static void register_layer(int li, int nl, + std::vector norm_in, std::vector Wqkv, + std::vector Wo, std::vector norm_post, + std::vector Wgate, std::vector Wup, + std::vector Wdown, std::vector kcache, + std::vector vcache) +{ + LayerW lw{}; + for (int d = 0; d < g_ctx.n_gpus; d++) { + lw.norm_in[d] = reinterpret_cast(norm_in[d].data_ptr()); + lw.Wqkv[d] = reinterpret_cast(Wqkv[d].data_ptr()); + lw.Wo[d] = reinterpret_cast(Wo[d].data_ptr()); + lw.norm_post[d] = reinterpret_cast(norm_post[d].data_ptr()); + lw.Wgate[d] = reinterpret_cast(Wgate[d].data_ptr()); + lw.Wup[d] = reinterpret_cast(Wup[d].data_ptr()); + lw.Wdown[d] = reinterpret_cast(Wdown[d].data_ptr()); + lw.kc[d] = reinterpret_cast(kcache[d].data_ptr()); + lw.vc[d] = reinterpret_cast(vcache[d].data_ptr()); + } + lw.inter = Wgate[0].numel() / g_ctx.hidden; + set_layer(li, nl, lw); +} + +static void register_io(std::vector xbuf, + std::vector cosbuf, + std::vector sinbuf) { + sbf16 *x[12], *c[12], *s[12]; + for (int d = 0; d < g_ctx.n_gpus; d++) { + x[d] = reinterpret_cast(xbuf[d].data_ptr()); + c[d] = reinterpret_cast(cosbuf[d].data_ptr()); + s[d] = reinterpret_cast(sinbuf[d].data_ptr()); + } + set_io(x, c, s); +} + +static void decode_step(int pos, int nq, int nkv, double eps) { + TORCH_CHECK(g_ctx.initialized && !g_layers.empty() && g_io_registered, + "sk_decode: init / register_layer / register_io first"); + run_layers(pos, nq, nkv, (float)eps); +} + +static void decode_2d_step(int pos, int nq, int nkv, double eps) { + TORCH_CHECK(g_ctx.initialized && !g_layers.empty() && g_io_registered && g_tp > 0, + "sk_decode: init / register_layer / register_io / set_2d first"); + run_layers_2d(pos, nq, nkv, (float)eps); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "SyclKittens registered all-layer decode orchestrator"; m.def("init", &py_init, pybind11::arg("n_gpus"), pybind11::arg("hidden"), @@ -627,3 +659,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { pybind11::arg("pos"), pybind11::arg("nq"), pybind11::arg("nkv"), pybind11::arg("eps")); } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "decode_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/decode/decode_harness.impl b/kernels/decode/decode_harness.impl new file mode 100644 index 0000000..e0d4e16 --- /dev/null +++ b/kernels/decode/decode_harness.impl @@ -0,0 +1,213 @@ +// Standalone harness for decode.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of decode.dp.cpp when TORCH_COMPILE is not defined, so the +// orchestrator core (g_queues, py_init, set_layer, set_io, run_layers) is in scope. +// +// Single tile (TP=1, CP=1), Llama-3-8B layer dims, NL layers of random weights. +// For each position: one decode step checked against a CPU fp32 reference of the +// same layer math (bf16 rounding emulated where the GPU stores bf16), then timing. +// +// usage: decode [iters=50] [warmup=5] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. + +#include +#include +#include +#include +#include + +namespace harness { + +constexpr int H = 4096, NQ = 32, NKV = 8, HD = 128, INTER = 14336, NL = 2; +constexpr int MAX_CTX = 4096, QKV_OUT = (NQ + 2 * NKV) * HD; +constexpr float EPS = 1e-5f; + +struct HostLayer { + std::vector norm_in, wqkv, wo, norm_post, wgate, wup, wdown, kc, vc; +}; + +static float lcg_uniform(uint32_t &state) { // deterministic, in [-1, 1) + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} +static std::vector rand_vec(size_t n, float scale, float offset, uint32_t &seed) { + std::vector v(n); + for (auto &x : v) x = sbf16(offset + scale * lcg_uniform(seed)); + return v; +} +static float bf(float v) { return float(sbf16(v)); } // round like a bf16 store + +static void matvec(const std::vector &W, const std::vector &x, + std::vector &y, int n_out, int k) { // W is [n_out, k] + y.assign(n_out, 0.f); + for (int n = 0; n < n_out; n++) { + const sbf16 *w = &W[size_t(n) * k]; + double acc = 0.0; + for (int i = 0; i < k; i++) acc += double(float(w[i])) * x[i]; + y[n] = bf(float(acc)); + } +} +static std::vector rmsnorm(const std::vector &x, const std::vector &w) { + double ss = 0.0; + for (float v : x) ss += double(v) * v; + const float inv = 1.0f / std::sqrt(float(ss / x.size()) + EPS); + std::vector y(x.size()); + for (size_t i = 0; i < x.size(); i++) y[i] = bf(x[i] * inv * float(w[i])); + return y; +} +static void rope(float *row, const std::vector &c, const std::vector &s) { + for (int i = 0; i < HD / 2; i++) { + const float a = row[i], b = row[i + HD / 2], cv = float(c[i]), sv = float(s[i]); + row[i] = bf(a * cv - b * sv); + row[i + HD / 2] = bf(b * cv + a * sv); + } +} + +// CPU reference of one decode step (all NL layers) at position `pos`. +static std::vector ref_step(const std::vector &layers, std::vector x, + int pos, const std::vector &c, const std::vector &s) { + const int group = NQ / NKV, L = pos + 1; + const float scale = 1.0f / std::sqrt(float(HD)); + std::vector h, qkv, o, gate, up, d; + for (const HostLayer &lw : layers) { + h = rmsnorm(x, lw.norm_in); + matvec(lw.wqkv, h, qkv, QKV_OUT, H); + for (int r = 0; r < NQ + NKV; r++) rope(&qkv[size_t(r) * HD], c, s); // q heads, then k heads + const float *knew = &qkv[size_t(NQ) * HD], *vnew = &qkv[size_t(NQ + NKV) * HD]; + std::vector attn(size_t(NQ) * HD), sc(L); + for (int qh = 0; qh < NQ; qh++) { + const int kvh = qh / group; + const float *q = &qkv[size_t(qh) * HD]; + auto key = [&](int t, int i) { + return t == pos ? knew[size_t(kvh) * HD + i] + : float(lw.kc[(size_t(kvh) * MAX_CTX + t) * HD + i]); }; + auto val = [&](int t, int i) { + return t == pos ? vnew[size_t(kvh) * HD + i] + : float(lw.vc[(size_t(kvh) * MAX_CTX + t) * HD + i]); }; + float m = -1e30f; + for (int t = 0; t < L; t++) { + float a = 0.f; + for (int i = 0; i < HD; i++) a += q[i] * key(t, i); + sc[t] = a * scale; m = std::max(m, sc[t]); + } + double sum = 0.0; + for (int t = 0; t < L; t++) { sc[t] = std::exp(sc[t] - m); sum += sc[t]; } + for (int i = 0; i < HD; i++) { + double acc = 0.0; + for (int t = 0; t < L; t++) acc += double(sc[t]) * val(t, i); + attn[size_t(qh) * HD + i] = bf(float(acc / sum)); + } + } + matvec(lw.wo, attn, o, H, NQ * HD); + std::vector x1(H); + for (int i = 0; i < H; i++) x1[i] = bf(x[i] + o[i]); + h = rmsnorm(x1, lw.norm_post); + matvec(lw.wgate, h, gate, INTER, H); + matvec(lw.wup, h, up, INTER, H); + std::vector hm(INTER); + for (int i = 0; i < INTER; i++) hm[i] = bf(gate[i] / (1.0f + std::exp(-gate[i])) * up[i]); + matvec(lw.wdown, hm, d, H, INTER); + for (int i = 0; i < H; i++) x[i] = bf(x1[i] + d[i]); + } + return x; +} + +template static T *to_device(const std::vector &h, sycl::queue &q) { + T *d = sycl::malloc_device(h.size(), q); + q.memcpy(d, h.data(), h.size() * sizeof(T)); + return d; +} + +} // namespace harness + +int main(int argc, char **argv) { + using namespace harness; + const int iters = argc > 1 ? std::atoi(argv[1]) : 50; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 5; + + g_queues.emplace_back(sycl::gpu_selector_v, sycl::property::queue::in_order{}); + sycl::queue &q = g_queues[0]; + std::printf("Device: %s\n", q.get_device().get_info().c_str()); + py_init(1, H, INTER, QKV_OUT, NQ, HD, MAX_CTX); + + // Random layers (weights scaled by 1/sqrt(in_dim)); random KV history. + uint32_t seed = 2024u; + std::vector layers(NL); + std::vector dev_bufs; + for (int li = 0; li < NL; li++) { + HostLayer &hl = layers[li]; + hl.norm_in = rand_vec(H, 0.1f, 1.0f, seed); + hl.wqkv = rand_vec(size_t(QKV_OUT) * H, 1.0f / std::sqrt(float(H)), 0.f, seed); + hl.wo = rand_vec(size_t(H) * NQ * HD, 1.0f / std::sqrt(float(NQ * HD)), 0.f, seed); + hl.norm_post = rand_vec(H, 0.1f, 1.0f, seed); + hl.wgate = rand_vec(size_t(INTER) * H, 1.0f / std::sqrt(float(H)), 0.f, seed); + hl.wup = rand_vec(size_t(INTER) * H, 1.0f / std::sqrt(float(H)), 0.f, seed); + hl.wdown = rand_vec(size_t(H) * INTER, 1.0f / std::sqrt(float(INTER)), 0.f, seed); + hl.kc = rand_vec(size_t(NKV) * MAX_CTX * HD, 1.0f, 0.f, seed); + hl.vc = rand_vec(size_t(NKV) * MAX_CTX * HD, 1.0f, 0.f, seed); + LayerW lw; + lw.norm_in[0] = to_device(hl.norm_in, q); lw.Wqkv[0] = to_device(hl.wqkv, q); + lw.Wo[0] = to_device(hl.wo, q); lw.norm_post[0] = to_device(hl.norm_post, q); + lw.Wgate[0] = to_device(hl.wgate, q); lw.Wup[0] = to_device(hl.wup, q); + lw.Wdown[0] = to_device(hl.wdown, q); + lw.kc[0] = to_device(hl.kc, q); lw.vc[0] = to_device(hl.vc, q); + lw.inter = INTER; + set_layer(li, NL, lw); + for (sbf16 *p : {lw.norm_in[0], lw.Wqkv[0], lw.Wo[0], lw.norm_post[0], lw.Wgate[0], + lw.Wup[0], lw.Wdown[0], lw.kc[0], lw.vc[0]}) dev_bufs.push_back(p); + } + const std::vector x0 = rand_vec(H, 1.0f, 0.f, seed); + sbf16 *d_x = sycl::malloc_device(H, q); + sbf16 *d_cos = sycl::malloc_device(HD / 2, q); + sbf16 *d_sin = sycl::malloc_device(HD / 2, q); + set_io(&d_x, &d_cos, &d_sin); + q.wait(); + + bool all_pass = true; + for (int pos : {127, 4095}) { + std::vector c(HD / 2), s(HD / 2); + for (int i = 0; i < HD / 2; i++) { // Llama-3 rope theta + const double ang = pos * std::pow(500000.0, -2.0 * i / HD); + c[i] = sbf16(float(std::cos(ang))); s[i] = sbf16(float(std::sin(ang))); + } + q.memcpy(d_cos, c.data(), c.size() * sizeof(sbf16)); + q.memcpy(d_sin, s.data(), s.size() * sizeof(sbf16)); + for (int li = 0; li < NL; li++) { // pristine KV history for the check + q.memcpy(g_layers[li].kc[0], layers[li].kc.data(), layers[li].kc.size() * sizeof(sbf16)); + q.memcpy(g_layers[li].vc[0], layers[li].vc.data(), layers[li].vc.size() * sizeof(sbf16)); + } + q.memcpy(d_x, x0.data(), H * sizeof(sbf16)); + run_layers(pos, NQ, NKV, EPS); + std::vector x_gpu(H); + q.memcpy(x_gpu.data(), d_x, H * sizeof(sbf16)).wait(); + + std::vector xin(H); + for (int i = 0; i < H; i++) xin[i] = float(x0[i]); + const std::vector x_ref = ref_step(layers, xin, pos, c, s); + double e2 = 0.0, r2 = 0.0; + for (int i = 0; i < H; i++) { + const double dlt = double(float(x_gpu[i])) - x_ref[i]; + e2 += dlt * dlt; r2 += double(x_ref[i]) * x_ref[i]; + } + const double rel = std::sqrt(e2 / (r2 + 1e-30)); + const bool pass = rel < 3e-2; + all_pass = all_pass && pass; + + for (int i = 0; i < warmup; i++) run_layers(pos, NQ, NKV, EPS); + q.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) run_layers(pos, NQ, NKV, EPS); + q.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + + std::printf("decode TP=1 layers=%d H=%d nq=%d nkv=%d hd=%d inter=%d pos=%d: " + "%.3f ms/token (%.3f ms/layer, iters=%d warmup=%d) rel_l2=%.2e %s\n", + NL, H, NQ, NKV, HD, INTER, pos, ms, ms / NL, iters, warmup, rel, + pass ? "PASS" : "FAIL"); + } + + for (sbf16 *p : dev_bufs) sycl::free(p, q); + sycl::free(d_x, q); sycl::free(d_cos, q); sycl::free(d_sin, q); + py_cleanup(); + return all_pass ? 0 : 1; +} diff --git a/kernels/gemm/Makefile b/kernels/gemm/Makefile index 4b32d2b..2794acc 100644 --- a/kernels/gemm/Makefile +++ b/kernels/gemm/Makefile @@ -4,17 +4,32 @@ M ?= 4096 N ?= 4096 K ?= 4096 +GEMM_DEFS := -DGEMM_M=$(M) -DGEMM_N=$(N) -DGEMM_K=$(K) + .PHONY: all bf16 qkv all: bf16 qkv +ifeq ($(CONFIG),standalone) +bf16: check-sycl + @mkdir -p $(BUILD_DIR) + $(CXX) $(COMMON_SYCL_FLAGS) $(GEMM_DEFS) bf16_gemm.dp.cpp \ + -o $(BUILD_DIR)/bf16_gemm_$(M)x$(N)x$(K) +else bf16: check-python @mkdir -p $(BUILD_DIR) - $(CXX) $(EXTENSION_FLAGS) -DGEMM_M=$(M) -DGEMM_N=$(N) -DGEMM_K=$(K) \ + $(CXX) $(EXTENSION_FLAGS) $(GEMM_DEFS) \ -DTORCH_EXTENSION_NAME=sk_bf16_gemm bf16_gemm.dp.cpp \ -o $(BUILD_DIR)/sk_bf16_gemm$(EXT_SUFFIX) $(TORCH_LIBS) +endif +ifeq ($(CONFIG),standalone) +qkv: check-sycl + @mkdir -p $(BUILD_DIR) + $(CXX) $(COMMON_SYCL_FLAGS) qkv_gemm.dp.cpp -o $(BUILD_DIR)/qkv_gemm +else qkv: check-python @mkdir -p $(BUILD_DIR) $(CXX) $(EXTENSION_FLAGS) -DTORCH_EXTENSION_NAME=sk_qkv_gemm qkv_gemm.dp.cpp \ -o $(BUILD_DIR)/sk_qkv_gemm$(EXT_SUFFIX) $(TORCH_LIBS) +endif diff --git a/kernels/gemm/bf16_gemm.dp.cpp b/kernels/gemm/bf16_gemm.dp.cpp index 0b1d983..c961942 100644 --- a/kernels/gemm/bf16_gemm.dp.cpp +++ b/kernels/gemm/bf16_gemm.dp.cpp @@ -499,6 +499,21 @@ static void launch_splitk(kittens::bf16 *d_A, kittens::bf16 *d_B, } #endif // GEMM_NO_SPLITK_CODE +// ============================================================ +// Shared entry: C = A @ B on raw device pointers (used by both builds). +// Enqueues only; the caller synchronizes. +// ============================================================ + +static void run_gemm(kittens::bf16 *d_A, kittens::bf16 *d_B, + typename _gl_C::dtype *d_C, sycl::queue &queue) { +#ifdef GEMM_NO_SPLITK_CODE + dispatch_micro(d_A, d_B, d_C, queue); +#else + launch_splitk(d_A, d_B, d_C, queue); +#endif +} + +#ifdef TORCH_COMPILE // ============================================================ // PyTorch dispatch — wall-clock timing done by caller // ============================================================ @@ -539,11 +554,7 @@ void dispatch_gemm(torch::Tensor A, torch::Tensor B, torch::Tensor C) { bf16 *d_C = reinterpret_cast(C.data_ptr()); #endif -#ifdef GEMM_NO_SPLITK_CODE - dispatch_micro(d_A, d_B, d_C, queue); -#else - launch_splitk(d_A, d_B, d_C, queue); -#endif + run_gemm(d_A, d_B, d_C, queue); // No queue.wait() — caller synchronizes via torch.xpu.synchronize() } @@ -561,3 +572,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.attr("N") = GN; m.attr("K") = GK; } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "bf16_gemm_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/gemm/bf16_gemm_harness.impl b/kernels/gemm/bf16_gemm_harness.impl new file mode 100644 index 0000000..e3c178a --- /dev/null +++ b/kernels/gemm/bf16_gemm_harness.impl @@ -0,0 +1,79 @@ +// Standalone harness for bf16_gemm.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of bf16_gemm.dp.cpp when TORCH_COMPILE is not defined, +// so GM/GN/GK, _gl_C and run_gemm() are already in scope. +// +// usage: bf16_gemm [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. + +#include +#include +#include +#include +#include +#include + +using c_type = typename _gl_C::dtype; + +// Deterministic values in [-1, 1) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + + sycl::queue queue{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s\n", + queue.get_device().get_info().c_str()); + + const size_t size_a = size_t(GM) * GK, size_b = size_t(GK) * GN, size_c = size_t(GM) * GN; + std::vector h_a(size_a), h_b(size_b); + uint32_t seed = 12345u; + for (auto &x : h_a) x = bf16(lcg_uniform(seed)); + for (auto &x : h_b) x = bf16(lcg_uniform(seed)); + + bf16 *d_a = sycl::malloc_device(size_a, queue); + bf16 *d_b = sycl::malloc_device(size_b, queue); + c_type *d_c = sycl::malloc_device(size_c, queue); + queue.memcpy(d_a, h_a.data(), size_a * sizeof(bf16)); + queue.memcpy(d_b, h_b.data(), size_b * sizeof(bf16)); + queue.wait(); + + for (int i = 0; i < warmup; i++) run_gemm(d_a, d_b, d_c, queue); + queue.wait(); + + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) run_gemm(d_a, d_b, d_c, queue); + queue.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + const double tflops = 2.0 * GM * GN * GK / (ms * 1e-3) / 1e12; + + // Correctness: rel-L2 over 256 sampled outputs against an fp32 host reference. + std::vector h_c(size_c); + queue.memcpy(h_c.data(), d_c, size_c * sizeof(c_type)).wait(); + double err2 = 0.0, ref2 = 0.0; + for (int s = 0; s < 256; s++) { + const int r = int((size_t(s) * 7919) % GM), c = int((size_t(s) * 104729) % GN); + double ref = 0.0; + for (int k = 0; k < GK; k++) + ref += double(float(h_a[size_t(r) * GK + k])) * double(float(h_b[size_t(k) * GN + c])); + const double got = double(float(h_c[size_t(r) * GN + c])); + err2 += (got - ref) * (got - ref); + ref2 += ref * ref; + } + const double rel = std::sqrt(err2 / (ref2 + 1e-30)); + const bool pass = rel < 1e-2; + + std::printf("GEMM %dx%dx%d (bf16 in, %s out): %.3f ms %.1f TFLOP/s " + "(iters=%d warmup=%d) rel_l2=%.2e %s\n", + GM, GN, GK, sizeof(c_type) == 4 ? "fp32" : "bf16", ms, tflops, + iters, warmup, rel, pass ? "PASS" : "FAIL"); + + sycl::free(d_a, queue); + sycl::free(d_b, queue); + sycl::free(d_c, queue); + return pass ? 0 : 1; +} diff --git a/kernels/gemm/qkv_gemm.dp.cpp b/kernels/gemm/qkv_gemm.dp.cpp index 659c236..1d55d09 100644 --- a/kernels/gemm/qkv_gemm.dp.cpp +++ b/kernels/gemm/qkv_gemm.dp.cpp @@ -213,6 +213,7 @@ static void launch_qkv(bf16 *d_A, bf16 *d_B, }); } +#ifdef TORCH_COMPILE // ============================================================ // PyTorch dispatch // ============================================================ @@ -292,3 +293,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.attr("BM") = BM; m.attr("BN") = BN; } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "qkv_gemm_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/gemm/qkv_gemm_harness.impl b/kernels/gemm/qkv_gemm_harness.impl new file mode 100644 index 0000000..4b1f0fd --- /dev/null +++ b/kernels/gemm/qkv_gemm_harness.impl @@ -0,0 +1,106 @@ +// Standalone harness for qkv_gemm.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of qkv_gemm.dp.cpp when TORCH_COMPILE is not defined, +// so GK/GN/NQ/NK/NV, BM and launch_qkv<>() are already in scope. +// +// usage: qkv_gemm [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. + +#include +#include +#include +#include +#include + +// Deterministic values in [-1, 1) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} + +// rel-L2 over 256 sampled outputs vs an fp32 host reference. out(r, c) reads the +// result for fused column c (0..GN), so the same check covers both output forms. +template +static double sampled_rel_l2(const std::vector &a, const std::vector &b, int M, Out out) { + double err2 = 0.0, ref2 = 0.0; + for (int s = 0; s < 256; s++) { + const int r = int((size_t(s) * 7919) % M), c = int((size_t(s) * 104729) % GN); + double ref = 0.0; + for (int k = 0; k < GK; k++) + ref += double(float(a[size_t(r) * GK + k])) * double(float(b[size_t(k) * GN + c])); + const double d = double(out(r, c)) - ref; + err2 += d * d; + ref2 += ref * ref; + } + return std::sqrt(err2 / (ref2 + 1e-30)); +} + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + + sycl::queue queue{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s\n", + queue.get_device().get_info().c_str()); + + std::vector h_b(size_t(GK) * GN); + uint32_t seed = 12345u; + for (auto &x : h_b) x = bf16(lcg_uniform(seed) * 0.05f); + bf16 *d_b = sycl::malloc_device(h_b.size(), queue); + queue.memcpy(d_b, h_b.data(), h_b.size() * sizeof(bf16)).wait(); + + bool all_pass = true; + for (int M : {512, 2048}) { + if (M % BM != 0) { std::printf("skip M=%d (not a multiple of BM=%d)\n", M, BM); continue; } + std::vector h_a(size_t(M) * GK); + for (auto &x : h_a) x = bf16(lcg_uniform(seed)); + bf16 *d_a = sycl::malloc_device(h_a.size(), queue); + bf16 *d_c = sycl::malloc_device(size_t(M) * GN, queue); + bf16 *d_q = sycl::malloc_device(size_t(M) * NQ, queue); + bf16 *d_k = sycl::malloc_device(size_t(M) * NK, queue); + bf16 *d_v = sycl::malloc_device(size_t(M) * NV, queue); + queue.memcpy(d_a, h_a.data(), h_a.size() * sizeof(bf16)).wait(); + + for (int split = 0; split < 2; split++) { + auto run = [&] { + if (split) launch_qkv(d_a, d_b, nullptr, d_q, d_k, d_v, M, queue); + else launch_qkv(d_a, d_b, d_c, d_c, d_c, d_c, M, queue); + }; + for (int i = 0; i < warmup; i++) run(); + queue.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) run(); + queue.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + const double tflops = 2.0 * M * GN * GK / (ms * 1e-3) / 1e12; + + double rel; + if (split) { + std::vector hq(size_t(M) * NQ), hk(size_t(M) * NK), hv(size_t(M) * NV); + queue.memcpy(hq.data(), d_q, hq.size() * sizeof(bf16)); + queue.memcpy(hk.data(), d_k, hk.size() * sizeof(bf16)); + queue.memcpy(hv.data(), d_v, hv.size() * sizeof(bf16)); + queue.wait(); + rel = sampled_rel_l2(h_a, h_b, M, [&](int r, int c) -> float { + if (c < NQ) return float(hq[size_t(r) * NQ + c]); + if (c < NQ + NK) return float(hk[size_t(r) * NK + (c - NQ)]); + return float(hv[size_t(r) * NV + (c - NQ - NK)]); + }); + } else { + std::vector hc(size_t(M) * GN); + queue.memcpy(hc.data(), d_c, hc.size() * sizeof(bf16)).wait(); + rel = sampled_rel_l2(h_a, h_b, M, [&](int r, int c) { return float(hc[size_t(r) * GN + c]); }); + } + const bool pass = rel < 1e-2; + all_pass &= pass; + std::printf("qkv_gemm %s M=%d K=%d N=%d (q %d k %d v %d): %.4f ms %.1f TFLOP/s " + "(iters=%d warmup=%d) rel_l2=%.2e %s\n", + split ? "split" : "fused", M, GK, GN, NQ, NK, NV, ms, tflops, + iters, warmup, rel, pass ? "PASS" : "FAIL"); + } + sycl::free(d_a, queue); sycl::free(d_c, queue); + sycl::free(d_q, queue); sycl::free(d_k, queue); sycl::free(d_v, queue); + } + sycl::free(d_b, queue); + return all_pass ? 0 : 1; +} diff --git a/kernels/norm/Makefile b/kernels/norm/Makefile index f1e3aa4..dff7784 100644 --- a/kernels/norm/Makefile +++ b/kernels/norm/Makefile @@ -4,18 +4,31 @@ RMSNORM_D ?= 4096 LN_D ?= 4096 LN_WARPS_PER_ROW ?= 8 +RMSNORM_DEFS := -DRMSNORM_D=$(RMSNORM_D) +LN_DEFS := -DLN_D=$(LN_D) -DLN_WARPS_PER_ROW=$(LN_WARPS_PER_ROW) + .PHONY: all rmsnorm layernorm all: rmsnorm layernorm +ifeq ($(CONFIG),standalone) +rmsnorm: check-sycl + @mkdir -p $(BUILD_DIR) + $(CXX) $(COMMON_SYCL_FLAGS) $(RMSNORM_DEFS) rmsnorm.dp.cpp -o $(BUILD_DIR)/rmsnorm + +layernorm: check-sycl + @mkdir -p $(BUILD_DIR) + $(CXX) $(COMMON_SYCL_FLAGS) $(LN_DEFS) layernorm.dp.cpp -o $(BUILD_DIR)/layernorm +else rmsnorm: check-python @mkdir -p $(BUILD_DIR) - $(CXX) $(EXTENSION_FLAGS) -DRMSNORM_D=$(RMSNORM_D) \ + $(CXX) $(EXTENSION_FLAGS) $(RMSNORM_DEFS) \ -DTORCH_EXTENSION_NAME=sk_rmsnorm rmsnorm.dp.cpp \ -o $(BUILD_DIR)/sk_rmsnorm$(EXT_SUFFIX) $(TORCH_LIBS) layernorm: check-python @mkdir -p $(BUILD_DIR) - $(CXX) $(EXTENSION_FLAGS) -DLN_D=$(LN_D) -DLN_WARPS_PER_ROW=$(LN_WARPS_PER_ROW) \ + $(CXX) $(EXTENSION_FLAGS) $(LN_DEFS) \ -DTORCH_EXTENSION_NAME=sk_layernorm layernorm.dp.cpp \ -o $(BUILD_DIR)/sk_layernorm$(EXT_SUFFIX) $(TORCH_LIBS) +endif diff --git a/kernels/norm/layernorm.dp.cpp b/kernels/norm/layernorm.dp.cpp index ba56524..720bc29 100644 --- a/kernels/norm/layernorm.dp.cpp +++ b/kernels/norm/layernorm.dp.cpp @@ -27,10 +27,6 @@ #include "kittens.dp.hpp" #pragma clang diagnostic pop -#include "pyutils/torch_helpers.dp.hpp" -#include -#include - #ifndef LN_D #define LN_D 4096 #endif @@ -177,10 +173,15 @@ void launch_layernorm_v6( }); } +#ifdef TORCH_COMPILE // ============================================================ // PyTorch dispatch functions // ============================================================ +#include "pyutils/torch_helpers.dp.hpp" +#include +#include + torch::Tensor dispatch_layernorm( torch::Tensor x, torch::Tensor weight, torch::Tensor bias) { CHECK_INPUT(x); @@ -251,3 +252,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Fused Add+LayerNorm: (x + residual) -> LN, residual_out", py::arg("x"), py::arg("residual"), py::arg("weight"), py::arg("bias")); } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "layernorm_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/norm/layernorm_harness.impl b/kernels/norm/layernorm_harness.impl new file mode 100644 index 0000000..1bd0b4b --- /dev/null +++ b/kernels/norm/layernorm_harness.impl @@ -0,0 +1,118 @@ +// Standalone harness for layernorm.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of layernorm.dp.cpp when TORCH_COMPILE is not defined, so +// LN_D, LN_EPS and launch_layernorm_v6() are already in scope. +// +// usage: layernorm [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. + +#include +#include +#include +#include +#include +#include + +// Deterministic values in [-1, 1) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} + +static double rel_l2(const std::vector &got, const std::vector &ref) { + double err2 = 0.0, ref2 = 0.0; + for (size_t i = 0; i < ref.size(); i++) { + const double d = double(float(got[i])) - ref[i]; + err2 += d * d; + ref2 += double(ref[i]) * ref[i]; + } + return std::sqrt(err2 / (ref2 + 1e-30)); +} + +// One case: T tokens (B=1, N=T). FUSED adds the residual and writes it back. +template +static bool run_case(sycl::queue &q, int T, int iters, int warmup) { + constexpr int D = LN_D; + const size_t n = size_t(T) * D; + + std::vector h_x(n), h_r(n), h_g(D), h_b(D); + uint32_t seed = 54321u + T; + for (auto &v : h_x) v = bf16(lcg_uniform(seed)); + for (auto &v : h_r) v = bf16(lcg_uniform(seed)); + for (auto &v : h_g) v = bf16(1.0f + 0.1f * lcg_uniform(seed)); + for (auto &v : h_b) v = bf16(0.1f * lcg_uniform(seed)); + + bf16 *d_x = sycl::malloc_device(n, q), *d_r = sycl::malloc_device(n, q); + bf16 *d_g = sycl::malloc_device(D, q), *d_b = sycl::malloc_device(D, q); + bf16 *d_o = sycl::malloc_device(n, q), *d_or = sycl::malloc_device(n, q); + q.memcpy(d_x, h_x.data(), n * sizeof(bf16)); + q.memcpy(d_r, h_r.data(), n * sizeof(bf16)); + q.memcpy(d_g, h_g.data(), D * sizeof(bf16)); + q.memcpy(d_b, h_b.data(), D * sizeof(bf16)); + q.wait(); + + auto launch = [&] { + launch_layernorm_v6( + q, d_x, FUSED ? d_r : nullptr, d_g, d_b, d_o, FUSED ? d_or : nullptr, + 1, T, LN_EPS); + }; + for (int i = 0; i < warmup; i++) launch(); + q.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) launch(); + q.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + // Compulsory traffic: read x (+residual), write out (+residual_out). + const double gbps = (FUSED ? 4.0 : 2.0) * n * sizeof(bf16) / (ms * 1e-3) / 1e9; + + std::vector h_o(n), h_or(n); + q.memcpy(h_o.data(), d_o, n * sizeof(bf16)); + if (FUSED) q.memcpy(h_or.data(), d_or, n * sizeof(bf16)); + q.wait(); + + // fp32 reference: h = x (+ r); out = (h - mean) * rsqrt(var + eps) * gamma + beta. + std::vector ref_o(n), ref_h(n); + for (int t = 0; t < T; t++) { + const size_t base = size_t(t) * D; + double s = 0.0, ss = 0.0; + for (int d = 0; d < D; d++) { + float h = float(h_x[base + d]) + (FUSED ? float(h_r[base + d]) : 0.0f); + ref_h[base + d] = h; + s += h; + ss += double(h) * h; + } + const double mean = s / D, var = ss / D - mean * mean; + const float inv = float(1.0 / std::sqrt(var + LN_EPS)); + for (int d = 0; d < D; d++) + ref_o[base + d] = (ref_h[base + d] - float(mean)) * inv * float(h_g[d]) + float(h_b[d]); + } + const double err = rel_l2(h_o, ref_o); + const double err_r = FUSED ? rel_l2(h_or, ref_h) : 0.0; + const bool pass = err < 1e-2 && err_r < 1e-2; + + char resid[48] = ""; + if (FUSED) std::snprintf(resid, sizeof(resid), " resid_rel_l2=%.2e", err_r); + std::printf("layernorm%s T=%-5d D=%d: %8.4f ms %7.1f GB/s rel_l2=%.2e%s %s\n", + FUSED ? "+residual" : " ", T, D, ms, gbps, err, resid, + pass ? "PASS" : "FAIL"); + + for (void *p : {(void *)d_x, (void *)d_r, (void *)d_g, (void *)d_b, (void *)d_o, (void *)d_or}) + sycl::free(p, q); + return pass; +} + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + + sycl::queue q{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s (iters=%d warmup=%d)\n", + q.get_device().get_info().c_str(), iters, warmup); + + bool ok = true; + for (int T : {1, 2048, 8192}) { + ok &= run_case(q, T, iters, warmup); + ok &= run_case(q, T, iters, warmup); + } + return ok ? 0 : 1; +} diff --git a/kernels/norm/rmsnorm.dp.cpp b/kernels/norm/rmsnorm.dp.cpp index 0670496..7508dd4 100644 --- a/kernels/norm/rmsnorm.dp.cpp +++ b/kernels/norm/rmsnorm.dp.cpp @@ -23,10 +23,6 @@ #include "kittens.dp.hpp" #pragma clang diagnostic pop -#include "pyutils/torch_helpers.dp.hpp" -#include -#include - #ifndef RMSNORM_D #define RMSNORM_D 2048 #endif @@ -159,10 +155,15 @@ void launch_rmsnorm_v6( }); } +#ifdef TORCH_COMPILE // ============================================================ // PyTorch dispatch functions // ============================================================ +#include "pyutils/torch_helpers.dp.hpp" +#include +#include + std::vector dispatch_rmsnorm( torch::Tensor x, torch::Tensor weight, double eps) { CHECK_INPUT(x); @@ -234,3 +235,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Fused add + RMSNorm: x,residual (bf16 [B,N,D]) -> (out, residual_out, inv_rms)", py::arg("x"), py::arg("residual"), py::arg("weight"), py::arg("eps")); } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "rmsnorm_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/norm/rmsnorm_harness.impl b/kernels/norm/rmsnorm_harness.impl new file mode 100644 index 0000000..329b221 --- /dev/null +++ b/kernels/norm/rmsnorm_harness.impl @@ -0,0 +1,115 @@ +// Standalone harness for rmsnorm.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of rmsnorm.dp.cpp when TORCH_COMPILE is not defined, so +// RMSNORM_D and launch_rmsnorm_v6() are already in scope. +// +// usage: rmsnorm [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. + +#include +#include +#include +#include +#include +#include + +// Deterministic values in [-1, 1) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} + +static double rel_l2(const std::vector &got, const std::vector &ref) { + double err2 = 0.0, ref2 = 0.0; + for (size_t i = 0; i < ref.size(); i++) { + const double d = double(float(got[i])) - ref[i]; + err2 += d * d; + ref2 += double(ref[i]) * ref[i]; + } + return std::sqrt(err2 / (ref2 + 1e-30)); +} + +// One case: T tokens (B=1, N=T). Returns true on PASS. +template +static bool run_case(sycl::queue &q, int T, int iters, int warmup) { + constexpr int D = RMSNORM_D; + constexpr float eps = 1e-6f; + const size_t n = size_t(T) * D; + + std::vector h_x(n), h_r(n), h_w(D); + uint32_t seed = 12345u + T; + for (auto &v : h_x) v = bf16(lcg_uniform(seed)); + for (auto &v : h_r) v = bf16(lcg_uniform(seed)); + for (auto &v : h_w) v = bf16(1.0f + 0.1f * lcg_uniform(seed)); + + bf16 *d_x = sycl::malloc_device(n, q), *d_r = sycl::malloc_device(n, q); + bf16 *d_w = sycl::malloc_device(D, q); + bf16 *d_o = sycl::malloc_device(n, q), *d_or = sycl::malloc_device(n, q); + float *d_inv = sycl::malloc_device(T, q); + q.memcpy(d_x, h_x.data(), n * sizeof(bf16)); + q.memcpy(d_r, h_r.data(), n * sizeof(bf16)); + q.memcpy(d_w, h_w.data(), D * sizeof(bf16)); + q.wait(); + + auto launch = [&] { + launch_rmsnorm_v6( + q, d_x, RESIDUAL ? d_r : nullptr, d_w, d_o, RESIDUAL ? d_or : nullptr, + d_inv, 1, T, eps); + }; + for (int i = 0; i < warmup; i++) launch(); + q.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) launch(); + q.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + // Compulsory traffic: read x (+residual), write out (+residual_out). + const double gbps = (RESIDUAL ? 4.0 : 2.0) * n * sizeof(bf16) / (ms * 1e-3) / 1e9; + + std::vector h_o(n), h_or(n); + q.memcpy(h_o.data(), d_o, n * sizeof(bf16)); + if (RESIDUAL) q.memcpy(h_or.data(), d_or, n * sizeof(bf16)); + q.wait(); + + // fp32 reference: h = x (+ r); out = h * rsqrt(mean(h^2) + eps) * w. + std::vector ref_o(n), ref_h(n); + for (int t = 0; t < T; t++) { + const size_t base = size_t(t) * D; + double ss = 0.0; + for (int d = 0; d < D; d++) { + float h = float(h_x[base + d]) + (RESIDUAL ? float(h_r[base + d]) : 0.0f); + ref_h[base + d] = h; + ss += double(h) * h; + } + const float inv = 1.0f / std::sqrt(float(ss / D) + eps); + for (int d = 0; d < D; d++) ref_o[base + d] = ref_h[base + d] * inv * float(h_w[d]); + } + const double err = rel_l2(h_o, ref_o); + const double err_r = RESIDUAL ? rel_l2(h_or, ref_h) : 0.0; + const bool pass = err < 1e-2 && err_r < 1e-2; + + char resid[48] = ""; + if (RESIDUAL) std::snprintf(resid, sizeof(resid), " resid_rel_l2=%.2e", err_r); + std::printf("rmsnorm%s T=%-5d D=%d: %8.4f ms %7.1f GB/s rel_l2=%.2e%s %s\n", + RESIDUAL ? "+residual" : " ", T, D, ms, gbps, err, resid, + pass ? "PASS" : "FAIL"); + + for (void *p : {(void *)d_x, (void *)d_r, (void *)d_w, (void *)d_o, (void *)d_or, (void *)d_inv}) + sycl::free(p, q); + return pass; +} + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + + sycl::queue q{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s (iters=%d warmup=%d)\n", + q.get_device().get_info().c_str(), iters, warmup); + + bool ok = true; + for (int T : {1, 2048, 8192}) { + ok &= run_case(q, T, iters, warmup); + ok &= run_case(q, T, iters, warmup); + } + return ok ? 0 : 1; +} diff --git a/kernels/rotary/Makefile b/kernels/rotary/Makefile index 8caa07c..5777ad1 100644 --- a/kernels/rotary/Makefile +++ b/kernels/rotary/Makefile @@ -3,24 +3,32 @@ include ../common.mk ROTARY_D ?= 128 ROTARY_NUM_WORKERS ?= 4 +ROTARY_DEFS := -DROTARY_D=$(ROTARY_D) -DROTARY_NUM_WORKERS=$(ROTARY_NUM_WORKERS) + .PHONY: all rotary rotary_qk rope_kvwrite all: rotary rotary_qk rope_kvwrite +ifeq ($(CONFIG),standalone) +rotary rotary_qk rope_kvwrite: check-sycl + @mkdir -p $(BUILD_DIR) + $(CXX) $(COMMON_SYCL_FLAGS) $(ROTARY_DEFS) $@.dp.cpp -o $(BUILD_DIR)/$@ +else rotary: check-python @mkdir -p $(BUILD_DIR) - $(CXX) $(EXTENSION_FLAGS) -DROTARY_D=$(ROTARY_D) -DROTARY_NUM_WORKERS=$(ROTARY_NUM_WORKERS) \ + $(CXX) $(EXTENSION_FLAGS) $(ROTARY_DEFS) \ -DTORCH_EXTENSION_NAME=sk_rotary rotary.dp.cpp \ -o $(BUILD_DIR)/sk_rotary$(EXT_SUFFIX) $(TORCH_LIBS) rotary_qk: check-python @mkdir -p $(BUILD_DIR) - $(CXX) $(EXTENSION_FLAGS) -DROTARY_D=$(ROTARY_D) -DROTARY_NUM_WORKERS=$(ROTARY_NUM_WORKERS) \ + $(CXX) $(EXTENSION_FLAGS) $(ROTARY_DEFS) \ -DTORCH_EXTENSION_NAME=sk_rotary_qk rotary_qk.dp.cpp \ -o $(BUILD_DIR)/sk_rotary_qk$(EXT_SUFFIX) $(TORCH_LIBS) rope_kvwrite: check-python @mkdir -p $(BUILD_DIR) - $(CXX) $(EXTENSION_FLAGS) -DROTARY_D=$(ROTARY_D) -DROTARY_NUM_WORKERS=$(ROTARY_NUM_WORKERS) \ + $(CXX) $(EXTENSION_FLAGS) $(ROTARY_DEFS) \ -DTORCH_EXTENSION_NAME=sk_rope_kvwrite rope_kvwrite.dp.cpp \ -o $(BUILD_DIR)/sk_rope_kvwrite$(EXT_SUFFIX) $(TORCH_LIBS) +endif diff --git a/kernels/rotary/rope_kvwrite.dp.cpp b/kernels/rotary/rope_kvwrite.dp.cpp index 6fcc64c..e38c7be 100644 --- a/kernels/rotary/rope_kvwrite.dp.cpp +++ b/kernels/rotary/rope_kvwrite.dp.cpp @@ -33,9 +33,6 @@ #include "kittens.dp.hpp" #pragma clang diagnostic pop -#include "pyutils/torch_helpers.dp.hpp" -#include -#include #ifndef ROTARY_D #define ROTARY_D 128 @@ -210,6 +207,11 @@ void launch_rope_kvwrite( queue.wait(); } +#ifdef TORCH_COMPILE +#include "pyutils/torch_helpers.dp.hpp" +#include +#include + // ============================================================ // PyTorch dispatch // ============================================================ @@ -281,3 +283,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("k"), py::arg("v"), py::arg("cos"), py::arg("sin"), py::arg("k_cache"), py::arg("v_cache"), py::arg("cache_pos")); } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "rope_kvwrite_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/rotary/rope_kvwrite_harness.impl b/kernels/rotary/rope_kvwrite_harness.impl new file mode 100644 index 0000000..d9ba6bb --- /dev/null +++ b/kernels/rotary/rope_kvwrite_harness.impl @@ -0,0 +1,120 @@ +// Standalone harness for rope_kvwrite.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of rope_kvwrite.dp.cpp when TORCH_COMPILE is not defined, +// so ROTARY_D and launch_rope_kvwrite() are already in scope. +// +// usage: rope_kvwrite [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. + +#include +#include +#include +#include +#include +#include + +// Deterministic values in [-1, 1) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} + +// k, v: (B, N, H, D); caches: (B, H, S, D) with S = cache_pos + N + 16 (slack rows +// must stay untouched). Checks rotated K, copied V, and the untouched cache rows. +static bool run_case(sycl::queue &queue, int B, int N, int H, long cache_pos, + int iters, int warmup) { + constexpr int D = ROTARY_D, HALF = D / 2; + const long S = cache_pos + N + 16; + const size_t nx = size_t(B) * N * H * D, ncs = size_t(N) * HALF, nc = size_t(B) * H * S * D; + const bf16 sentinel = bf16(7.0f); // outside the [-2, 2] range any written value can take + + uint32_t seed = 777u + uint32_t(N) * 7u; + std::vector h_k(nx), h_v(nx), h_cos(ncs), h_sin(ncs), h_kc(nc), h_vc(nc); + for (auto &x : h_k) x = bf16(lcg_uniform(seed)); + for (auto &x : h_v) x = bf16(lcg_uniform(seed)); + for (size_t i = 0; i < ncs; i++) { + const float ang = 3.14159265f * lcg_uniform(seed); + h_cos[i] = bf16(std::cos(ang)); + h_sin[i] = bf16(std::sin(ang)); + } + + bf16 *d_k = sycl::malloc_device(nx, queue); + bf16 *d_v = sycl::malloc_device(nx, queue); + bf16 *d_cos = sycl::malloc_device(ncs, queue); + bf16 *d_sin = sycl::malloc_device(ncs, queue); + bf16 *d_kc = sycl::malloc_device(nc, queue); + bf16 *d_vc = sycl::malloc_device(nc, queue); + queue.memcpy(d_k, h_k.data(), nx * sizeof(bf16)); + queue.memcpy(d_v, h_v.data(), nx * sizeof(bf16)); + queue.memcpy(d_cos, h_cos.data(), ncs * sizeof(bf16)); + queue.memcpy(d_sin, h_sin.data(), ncs * sizeof(bf16)); + queue.fill(d_kc, sentinel, nc); + queue.fill(d_vc, sentinel, nc); + queue.wait(); + + // The kernel is idempotent for fixed inputs, so repeated launches keep the same cache. + auto launch = [&] { + launch_rope_kvwrite(queue, d_k, d_v, d_cos, d_sin, d_kc, d_vc, B, N, H, S, S, cache_pos); + }; + for (int i = 0; i < warmup; i++) launch(); + queue.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) launch(); + queue.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + const double bytes = 4.0 * nx * sizeof(bf16) + 2.0 * ncs * sizeof(bf16); // read k,v + write both caches + + queue.memcpy(h_kc.data(), d_kc, nc * sizeof(bf16)); + queue.memcpy(h_vc.data(), d_vc, nc * sizeof(bf16)); + queue.wait(); + double max_err = 0.0; + long untouched_bad = 0; + for (int b = 0; b < B; b++) + for (int h = 0; h < H; h++) + for (long s = 0; s < S; s++) { + const size_t c_off = ((size_t(b) * H + h) * S + s) * D; + const long n = s - cache_pos; + if (n < 0 || n >= N) { // rows outside [cache_pos, cache_pos + N) must be untouched + for (int d = 0; d < D; d++) + untouched_bad += (float(h_kc[c_off + d]) != float(sentinel)) + + (float(h_vc[c_off + d]) != float(sentinel)); + continue; + } + const size_t x_off = ((size_t(b) * N + n) * H + h) * D; + for (int d = 0; d < HALF; d++) { + const float a = float(h_k[x_off + d]), bb = float(h_k[x_off + HALF + d]); + const float c = float(h_cos[size_t(n) * HALF + d]), sn = float(h_sin[size_t(n) * HALF + d]); + const double r1 = a * c - bb * sn, r2 = bb * c + a * sn; + max_err = std::fmax(max_err, std::fabs(double(float(h_kc[c_off + d])) - r1)); + max_err = std::fmax(max_err, std::fabs(double(float(h_kc[c_off + HALF + d])) - r2)); + } + for (int d = 0; d < D; d++) // V is a raw copy: must be bit-exact + max_err = std::fmax(max_err, std::fabs(double(float(h_vc[c_off + d])) - + double(float(h_v[x_off + d])))); + } + const bool pass = max_err <= 1e-2 && untouched_bad == 0; + std::printf("rope_kvwrite B=%d T=%d H=%d D=%d S=%ld pos=%ld: %.4f ms %.1f GB/s " + "max_err=%.2e untouched_bad=%ld %s\n", + B, N, H, D, S, cache_pos, ms, bytes / (ms * 1e-3) / 1e9, max_err, untouched_bad, + pass ? "PASS" : "FAIL"); + + sycl::free(d_k, queue); sycl::free(d_v, queue); + sycl::free(d_cos, queue); sycl::free(d_sin, queue); + sycl::free(d_kc, queue); sycl::free(d_vc, queue); + return pass; +} + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + + sycl::queue queue{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s (iters=%d warmup=%d)\n", + queue.get_device().get_info().c_str(), iters, warmup); + + bool ok = true; + // T=1 at a nonzero cache position (decode step), T=2048 from position 0 (prefill). + ok &= run_case(queue, /*B=*/1, /*T=*/1, /*H=*/8, /*cache_pos=*/100, iters, warmup); + ok &= run_case(queue, /*B=*/1, /*T=*/2048, /*H=*/8, /*cache_pos=*/0, iters, warmup); + return ok ? 0 : 1; +} diff --git a/kernels/rotary/rotary.dp.cpp b/kernels/rotary/rotary.dp.cpp index 706d704..d562434 100644 --- a/kernels/rotary/rotary.dp.cpp +++ b/kernels/rotary/rotary.dp.cpp @@ -20,10 +20,6 @@ #include "kittens.dp.hpp" #pragma clang diagnostic pop -#include "pyutils/torch_helpers.dp.hpp" -#include -#include - #ifndef ROTARY_D #define ROTARY_D 128 #endif @@ -172,6 +168,11 @@ void launch_rotary( queue.wait(); } +#ifdef TORCH_COMPILE +#include "pyutils/torch_helpers.dp.hpp" +#include +#include + // ============================================================ // PyTorch dispatch // ============================================================ @@ -253,3 +254,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "RoPE backward: dy (bf16 [B,N,H,D]), cos/sin (bf16 [N,D/2]) -> dx", py::arg("dy"), py::arg("cos"), py::arg("sin")); } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "rotary_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/rotary/rotary_harness.impl b/kernels/rotary/rotary_harness.impl new file mode 100644 index 0000000..7bc56a7 --- /dev/null +++ b/kernels/rotary/rotary_harness.impl @@ -0,0 +1,96 @@ +// Standalone harness for rotary.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of rotary.dp.cpp when TORCH_COMPILE is not defined, so +// ROTARY_D and launch_rotary() are already in scope. +// +// usage: rotary [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. + +#include +#include +#include +#include +#include +#include + +// Deterministic values in [-1, 1) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} + +// One forward (negate_sin=false) or backward (negate_sin=true) case; returns pass. +static bool run_case(sycl::queue &queue, int B, int N, int H, bool negate_sin, + int iters, int warmup) { + constexpr int D = ROTARY_D, HALF = D / 2; + const size_t nx = size_t(B) * N * H * D, ncs = size_t(N) * HALF; + + uint32_t seed = 12345u + uint32_t(N) * 7u + (negate_sin ? 1u : 0u); + std::vector h_x(nx), h_cos(ncs), h_sin(ncs), h_out(nx); + for (auto &v : h_x) v = bf16(lcg_uniform(seed)); + for (size_t i = 0; i < ncs; i++) { + const float ang = 3.14159265f * lcg_uniform(seed); + h_cos[i] = bf16(std::cos(ang)); + h_sin[i] = bf16(std::sin(ang)); + } + + bf16 *d_x = sycl::malloc_device(nx, queue); + bf16 *d_out = sycl::malloc_device(nx, queue); + bf16 *d_cos = sycl::malloc_device(ncs, queue); + bf16 *d_sin = sycl::malloc_device(ncs, queue); + queue.memcpy(d_x, h_x.data(), nx * sizeof(bf16)); + queue.memcpy(d_cos, h_cos.data(), ncs * sizeof(bf16)); + queue.memcpy(d_sin, h_sin.data(), ncs * sizeof(bf16)); + queue.wait(); + + auto launch = [&] { + if (negate_sin) launch_rotary(queue, d_x, d_out, d_cos, d_sin, B, N, H); + else launch_rotary(queue, d_x, d_out, d_cos, d_sin, B, N, H); + }; + for (int i = 0; i < warmup; i++) launch(); + queue.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) launch(); + queue.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + const double bytes = 2.0 * nx * sizeof(bf16) + 2.0 * ncs * sizeof(bf16); + + // Full fp32 reference: neox half-split rotation with shared (N, D/2) cos/sin. + queue.memcpy(h_out.data(), d_out, nx * sizeof(bf16)).wait(); + double max_err = 0.0; + for (size_t row = 0; row < size_t(B) * N * H; row++) { + const int n = int((row / H) % N); + for (int d = 0; d < HALF; d++) { + const float a = float(h_x[row * D + d]), b = float(h_x[row * D + HALF + d]); + const float c = float(h_cos[size_t(n) * HALF + d]); + const float s = negate_sin ? -float(h_sin[size_t(n) * HALF + d]) + : float(h_sin[size_t(n) * HALF + d]); + const double r1 = a * c - b * s, r2 = b * c + a * s; + max_err = std::fmax(max_err, std::fabs(double(float(h_out[row * D + d])) - r1)); + max_err = std::fmax(max_err, std::fabs(double(float(h_out[row * D + HALF + d])) - r2)); + } + } + const bool pass = max_err <= 1e-2; + std::printf("rotary %s B=%d T=%d H=%d D=%d: %.4f ms %.1f GB/s max_err=%.2e %s\n", + negate_sin ? "bwd" : "fwd", B, N, H, D, ms, bytes / (ms * 1e-3) / 1e9, + max_err, pass ? "PASS" : "FAIL"); + + sycl::free(d_x, queue); sycl::free(d_out, queue); + sycl::free(d_cos, queue); sycl::free(d_sin, queue); + return pass; +} + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + + sycl::queue queue{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s (iters=%d warmup=%d)\n", + queue.get_device().get_info().c_str(), iters, warmup); + + bool ok = true; + for (int T : {1, 2048}) + for (bool bwd : {false, true}) + ok &= run_case(queue, /*B=*/1, T, /*H=*/32, bwd, iters, warmup); + return ok ? 0 : 1; +} diff --git a/kernels/rotary/rotary_qk.dp.cpp b/kernels/rotary/rotary_qk.dp.cpp index ef156de..1927395 100644 --- a/kernels/rotary/rotary_qk.dp.cpp +++ b/kernels/rotary/rotary_qk.dp.cpp @@ -27,9 +27,6 @@ #include "kittens.dp.hpp" #pragma clang diagnostic pop -#include "pyutils/torch_helpers.dp.hpp" -#include -#include #include #ifndef ROTARY_D @@ -202,6 +199,11 @@ void launch_rotary_qk( queue.wait(); } +#ifdef TORCH_COMPILE +#include "pyutils/torch_helpers.dp.hpp" +#include +#include + // ============================================================ // PyTorch dispatch // ============================================================ @@ -307,3 +309,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "cos/sin (bf16 [N,D/2]) -> (dq_out, dk_out)", py::arg("dq"), py::arg("dk"), py::arg("cos"), py::arg("sin")); } + +#else +// Standalone executable (CONFIG=standalone): main() lives in the harness. +#include "rotary_qk_harness.impl" +#endif // TORCH_COMPILE diff --git a/kernels/rotary/rotary_qk_harness.impl b/kernels/rotary/rotary_qk_harness.impl new file mode 100644 index 0000000..fb6109e --- /dev/null +++ b/kernels/rotary/rotary_qk_harness.impl @@ -0,0 +1,113 @@ +// Standalone harness for rotary_qk.dp.cpp (CONFIG=standalone, no PyTorch). +// Included at the end of rotary_qk.dp.cpp when TORCH_COMPILE is not defined, +// so ROTARY_D and launch_rotary_qk() are already in scope. +// +// usage: rotary_qk [iters=100] [warmup=10] +// pick the GPU with ONEAPI_DEVICE_SELECTOR / ZE_AFFINITY_MASK. + +#include +#include +#include +#include +#include +#include + +// Deterministic values in [-1, 1) so the run is reproducible. +static float lcg_uniform(uint32_t &state) { + state = state * 1664525u + 1013904223u; + return static_cast(state >> 8) / static_cast(1u << 23) - 1.0f; +} + +// Max |out - ref| over one (B, N, heads, D) tensor rotated with (N, D/2) cos/sin. +static double rope_max_err(const std::vector &x, const std::vector &out, + const std::vector &cs, const std::vector &sn, + int B, int N, int heads, bool negate_sin) { + constexpr int D = ROTARY_D, HALF = D / 2; + double max_err = 0.0; + for (size_t row = 0; row < size_t(B) * N * heads; row++) { + const int n = int((row / heads) % N); + for (int d = 0; d < HALF; d++) { + const float a = float(x[row * D + d]), b = float(x[row * D + HALF + d]); + const float c = float(cs[size_t(n) * HALF + d]); + const float s = negate_sin ? -float(sn[size_t(n) * HALF + d]) : float(sn[size_t(n) * HALF + d]); + const double r1 = a * c - b * s, r2 = b * c + a * s; + max_err = std::fmax(max_err, std::fabs(double(float(out[row * D + d])) - r1)); + max_err = std::fmax(max_err, std::fabs(double(float(out[row * D + HALF + d])) - r2)); + } + } + return max_err; +} + +// One forward (negate_sin=false) or backward (negate_sin=true) case; returns pass. +static bool run_case(sycl::queue &queue, int B, int N, int H, int HKV, bool negate_sin, + int iters, int warmup) { + constexpr int D = ROTARY_D, HALF = D / 2; + const size_t nq = size_t(B) * N * H * D, nk = size_t(B) * N * HKV * D, ncs = size_t(N) * HALF; + + uint32_t seed = 54321u + uint32_t(N) * 7u + (negate_sin ? 1u : 0u); + std::vector h_q(nq), h_k(nk), h_cos(ncs), h_sin(ncs), h_qo(nq), h_ko(nk); + for (auto &v : h_q) v = bf16(lcg_uniform(seed)); + for (auto &v : h_k) v = bf16(lcg_uniform(seed)); + for (size_t i = 0; i < ncs; i++) { + const float ang = 3.14159265f * lcg_uniform(seed); + h_cos[i] = bf16(std::cos(ang)); + h_sin[i] = bf16(std::sin(ang)); + } + + bf16 *d_q = sycl::malloc_device(nq, queue); + bf16 *d_qo = sycl::malloc_device(nq, queue); + bf16 *d_k = sycl::malloc_device(nk, queue); + bf16 *d_ko = sycl::malloc_device(nk, queue); + bf16 *d_cos = sycl::malloc_device(ncs, queue); + bf16 *d_sin = sycl::malloc_device(ncs, queue); + queue.memcpy(d_q, h_q.data(), nq * sizeof(bf16)); + queue.memcpy(d_k, h_k.data(), nk * sizeof(bf16)); + queue.memcpy(d_cos, h_cos.data(), ncs * sizeof(bf16)); + queue.memcpy(d_sin, h_sin.data(), ncs * sizeof(bf16)); + queue.wait(); + + auto launch = [&] { + if (negate_sin) + launch_rotary_qk(queue, d_q, d_qo, d_k, d_ko, d_cos, d_sin, B, N, H, HKV); + else + launch_rotary_qk(queue, d_q, d_qo, d_k, d_ko, d_cos, d_sin, B, N, H, HKV); + }; + for (int i = 0; i < warmup; i++) launch(); + queue.wait(); + const auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < iters; i++) launch(); + queue.wait(); + const auto t1 = std::chrono::high_resolution_clock::now(); + const double ms = std::chrono::duration(t1 - t0).count() / iters; + const double bytes = 2.0 * (nq + nk) * sizeof(bf16) + 2.0 * ncs * sizeof(bf16); + + queue.memcpy(h_qo.data(), d_qo, nq * sizeof(bf16)); + queue.memcpy(h_ko.data(), d_ko, nk * sizeof(bf16)); + queue.wait(); + const double max_err = std::fmax(rope_max_err(h_q, h_qo, h_cos, h_sin, B, N, H, negate_sin), + rope_max_err(h_k, h_ko, h_cos, h_sin, B, N, HKV, negate_sin)); + const bool pass = max_err <= 1e-2; + std::printf("rotary_qk %s B=%d T=%d H=%d HKV=%d D=%d: %.4f ms %.1f GB/s max_err=%.2e %s\n", + negate_sin ? "bwd" : "fwd", B, N, H, HKV, D, ms, bytes / (ms * 1e-3) / 1e9, + max_err, pass ? "PASS" : "FAIL"); + + sycl::free(d_q, queue); sycl::free(d_qo, queue); + sycl::free(d_k, queue); sycl::free(d_ko, queue); + sycl::free(d_cos, queue); sycl::free(d_sin, queue); + return pass; +} + +int main(int argc, char **argv) { + const int iters = argc > 1 ? std::atoi(argv[1]) : 100; + const int warmup = argc > 2 ? std::atoi(argv[2]) : 10; + + sycl::queue queue{sycl::gpu_selector_v, sycl::property::queue::in_order{}}; + std::printf("Device: %s (iters=%d warmup=%d)\n", + queue.get_device().get_info().c_str(), iters, warmup); + + bool ok = true; + for (int T : {1, 2048}) + for (bool bwd : {false, true}) + ok &= run_case(queue, /*B=*/1, T, /*H=*/32, /*HKV=*/8, bwd, iters, warmup); + return ok ? 0 : 1; +} diff --git a/kernels/run_standalone.sh b/kernels/run_standalone.sh new file mode 100755 index 0000000..27e4a65 --- /dev/null +++ b/kernels/run_standalone.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# Build every kernel in standalone mode and run its harness; prints a PASS/FAIL summary. +# +# source ~/oneapi.sh # or env.src: icpx (and oneMKL for decode) on PATH +# ./kernels/run_standalone.sh [iters=50] [warmup=5] +# +# Pick the GPU with ZE_AFFINITY_MASK. Extra make variables go in MAKE_ARGS, +# e.g. MAKE_ARGS="GPU_TARGET=intel_gpu_pvc". Exit code is non-zero if anything fails. +set -u +cd "$(dirname "$0")" +ITERS=${1:-50} +WARMUP=${2:-5} + +command -v icpx > /dev/null || { echo "icpx not found: source your oneAPI environment first" >&2; exit 1; } +# decode includes ; oneAPI setvars does not put it on CPATH. +[ -n "${MKLROOT:-}" ] && export CPATH="$MKLROOT/include${CPATH:+:$CPATH}" + +DIRS="gemm attention norm rotary activation decode" +# dir/binary for every harness, built by `make -C all CONFIG=standalone` +BINS="gemm/bf16_gemm_4096x4096x4096 gemm/qkv_gemm attention/gqa_forward attention/gqa_backward + norm/rmsnorm norm/layernorm rotary/rotary rotary/rotary_qk rotary/rope_kvwrite + activation/silu_mul decode/decode" + +echo "== build (CONFIG=standalone), logs in /build/standalone.log" +declare -A pid +for d in $DIRS; do + mkdir -p "$d/build" + make -C "$d" all CONFIG=standalone ${MAKE_ARGS:-} > "$d/build/standalone.log" 2>&1 & + pid[$d]=$! +done +build_fail=0; broken=" " +for d in $DIRS; do + if ! wait "${pid[$d]}"; then + echo " $d: BUILD FAILED (see $d/build/standalone.log)" + build_fail=1; broken="$broken$d " + fi +done + +echo "== run (iters=$ITERS warmup=$WARMUP)" +pass=0; fail=0; failed="" +for b in $BINS; do + d="${b%%/*}"; exe="$d/build/${b#*/}" + # Never run an executable left over from an earlier build of a failed directory. + if [[ "$broken" == *" $d "* ]] || [ ! -x "$exe" ]; then + echo "-- $b: not built"; fail=$((fail + 1)); failed="$failed $b"; continue + fi + out=$("./$exe" "$ITERS" "$WARMUP" 2>&1); rc=$? + echo "-- $b" + echo "$out" | grep -E "PASS|FAIL" | sed 's/^/ /' + if [ $rc -eq 0 ] && ! echo "$out" | grep -q FAIL; then + pass=$((pass + 1)) + else + fail=$((fail + 1)); failed="$failed $b" + [ $rc -ne 0 ] && echo "$out" | tail -n 3 | sed 's/^/ ! /' + fi +done + +echo "== summary: $pass passed, $fail failed${failed:+ ($failed )}" +[ $fail -eq 0 ] && [ $build_fail -eq 0 ]