diff --git a/CHANGELOG.md b/CHANGELOG.md index fb26ff2ed1..5b6b9a6a3c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- On HIP, decode's softmax stats and partial sums load sixteen values ahead and fold them in the same order, with the same bits (#1961). Thanks to @AbrahamCete. + - GLM's CUDA `TF_KV_PAGES` is read through the profile registry, and startup reports the selection it resolved; the value and where it is read are unchanged (#1959, part of #1660). Thanks to @garvrao80. - GLM on CUDA can keep the KV pages' latent rows in pinned host memory with `TF_GLM_KV_HOST`, for a 1M-token window on two 96 GB cards; unset, the engine is as before, with identical tokens at TP=2 on two GB10s (#1795, part of #1794). Thanks to @russellmm. diff --git a/zig/kernels/hip/attention/attention.hip b/zig/kernels/hip/attention/attention.hip index e4c8afa2c8..9eddd16eea 100644 --- a/zig/kernels/hip/attention/attention.hip +++ b/zig/kernels/hip/attention/attention.hip @@ -121,8 +121,18 @@ extern "C" __global__ void tf_softmax_stats(const float* __restrict__ scores, fl int head = blockIdx.x; int batch = blockIdx.y; const float* row = scores + (static_cast(batch) * gridDim.x + head) * span; + // a thread's keys i, i + blockDim.x, ...: sixteen loads go out together, then fold in that order + const int ahead_span = 15 * blockDim.x; float local = -INFINITY; - for (int i = threadIdx.x; i < visible; i += blockDim.x) local = fmaxf(local, row[i]); + int i = threadIdx.x; + for (; i + ahead_span < visible; i += 16 * blockDim.x) { + float ahead[16]; +#pragma unroll + for (int u = 0; u < 16; ++u) ahead[u] = row[i + u * blockDim.x]; +#pragma unroll + for (int u = 0; u < 16; ++u) local = fmaxf(local, ahead[u]); + } + for (; i < visible; i += blockDim.x) local = fmaxf(local, row[i]); __shared__ float buf[256]; buf[threadIdx.x] = local; __syncthreads(); @@ -133,7 +143,15 @@ extern "C" __global__ void tf_softmax_stats(const float* __restrict__ scores, fl float max_score = buf[0]; __syncthreads(); local = 0.f; - for (int i = threadIdx.x; i < visible; i += blockDim.x) local += expf(row[i] - max_score); + i = threadIdx.x; + for (; i + ahead_span < visible; i += 16 * blockDim.x) { + float ahead[16]; +#pragma unroll + for (int u = 0; u < 16; ++u) ahead[u] = row[i + u * blockDim.x]; +#pragma unroll + for (int u = 0; u < 16; ++u) local += expf(ahead[u] - max_score); + } + for (; i < visible; i += blockDim.x) local += expf(row[i] - max_score); buf[threadIdx.x] = local; __syncthreads(); for (int step = blockDim.x >> 1; step > 0; step >>= 1) { @@ -186,8 +204,17 @@ extern "C" __global__ void tf_sum_partials(const float* __restrict__ partials, f int lane = threadIdx.x; if (lane >= d) return; const float* col = partials + ((static_cast(batch) * gridDim.x + head) * tiles) * d + lane; + // sixteen tiles' loads go out together, then add in tile order float acc = 0.f; - for (int tile = 0; tile < tiles; ++tile) acc += col[static_cast(tile) * d]; + int tile = 0; + for (; tile + 16 <= tiles; tile += 16) { + float ahead[16]; +#pragma unroll + for (int u = 0; u < 16; ++u) ahead[u] = col[static_cast(tile + u) * d]; +#pragma unroll + for (int u = 0; u < 16; ++u) acc += ahead[u]; + } + for (; tile < tiles; ++tile) acc += col[static_cast(tile) * d]; out[(static_cast(batch) * gridDim.x + head) * d + lane] = acc; } diff --git a/zig/kernels/hip/decode/plan.hip b/zig/kernels/hip/decode/plan.hip index a59241c5e2..656d6be027 100644 --- a/zig/kernels/hip/decode/plan.hip +++ b/zig/kernels/hip/decode/plan.hip @@ -379,8 +379,17 @@ __global__ void tf_sum_partials_act(const float* __restrict__ partials, void* __ int lane = threadIdx.x; if (lane >= d) return; const float* col = partials + ((static_cast(batch) * gridDim.x + head) * tiles) * d + lane; + // sixteen tiles' loads go out together, then add in tile order float acc = 0.f; - for (int tile = 0; tile < tiles; ++tile) acc += col[static_cast(tile) * d]; + int tile = 0; + for (; tile + 16 <= tiles; tile += 16) { + float ahead[16]; +#pragma unroll + for (int u = 0; u < 16; ++u) ahead[u] = col[static_cast(tile + u) * d]; +#pragma unroll + for (int u = 0; u < 16; ++u) acc += ahead[u]; + } + for (; tile < tiles; ++tile) acc += col[static_cast(tile) * d]; store(out, kind, (static_cast(batch) * gridDim.x + head) * d + lane, acc); } __global__ void __launch_bounds__(kAttendThreads) diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig index e9af10f759..0b339e28ca 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/zig/src/core/contract/hip_records.zig @@ -191,10 +191,10 @@ const own_records = [_]Record{ .threads = 256, }, .usage = .{ - .{ .vgprs = 9, .sgprs = 16, .lds = 1024, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 9, .sgprs = 18, .lds = 1024, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 9, .sgprs = 16, .lds = 1024, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 9, .sgprs = 18, .lds = 1024, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 61, .sgprs = 32, .lds = 1024, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 70, .sgprs = 32, .lds = 1024, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 70, .sgprs = 32, .lds = 1024, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 70, .sgprs = 32, .lds = 1024, .private = 0, .bound = 1024, .wave = 32 }, }, }, .{ @@ -233,10 +233,10 @@ const own_records = [_]Record{ .threads = 256, }, .usage = .{ - .{ .vgprs = 5, .sgprs = 14, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 5, .sgprs = 18, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 5, .sgprs = 14, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 5, .sgprs = 16, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 23, .sgprs = 20, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 23, .sgprs = 21, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 23, .sgprs = 20, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 23, .sgprs = 21, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, }, }, .{ @@ -441,10 +441,10 @@ const own_records = [_]Record{ .threads = 256, }, .usage = .{ - .{ .vgprs = 7, .sgprs = 16, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 6, .sgprs = 18, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 6, .sgprs = 16, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 6, .sgprs = 18, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 24, .sgprs = 21, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 24, .sgprs = 21, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 24, .sgprs = 21, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 24, .sgprs = 22, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, }, }, .{ diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index d37295b4d2..ff98d00f33 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -24,6 +24,7 @@ test { _ = @import("kernel_tests/chain.zig"); _ = @import("kernel_tests/attention.zig"); _ = @import("kernel_tests/attention_pin.zig"); + _ = @import("kernel_tests/stats.zig"); _ = @import("kernel_tests/ops.zig"); _ = @import("kernel_tests/argmax.zig"); _ = @import("kernel_tests/tp.zig"); diff --git a/zig/src/hip/kernel_tests/stats.zig b/zig/src/hip/kernel_tests/stats.zig new file mode 100644 index 0000000000..32b98bea2c --- /dev/null +++ b/zig/src/hip/kernel_tests/stats.zig @@ -0,0 +1,130 @@ +//! The decode walk's softmax stats and partial sums past the one-block walk: the max, and each lane's tiles in order. + +const std = @import("std"); +const rig = @import("rig.zig"); +const util = @import("../launches/util.zig"); +const DeviceBuffer = @import("../memory.zig").DeviceBuffer; + +const Rig = rig.Rig; +const gpa = rig.gpa; +const at = rig.at; + +fn ptr(b: DeviceBuffer) ?*anyopaque { + return @ptrFromInt(at(b)); +} + +/// bf16 rounded to nearest even, as the kernels' float_to_bf16 rounds a finite value. +fn bf16Even(v: f32) u16 { + const u: u32 = @bitCast(v); + return @intCast((u + ((u >> 16) & 1) + 0x7fff) >> 16); +} + +test "partial sums add each lane's tiles in tile order, into fp32 and the activation kinds" { + const t = try Rig.open(1 << 20); + defer t.close(); + var rng: rig.Rng = .{ .state = 0x5A5 }; + const heads = 4; + var cases: usize = 0; + for ([_]usize{ 1, 33, 512, 1024 }) |tiles| for ([_]usize{ 64, 128, 256 }) |d| for ([_]usize{ 1, 3 }) |rows| { + const host = try gpa.alloc(f32, rows * heads * tiles * d); + defer gpa.free(host); + for (host) |*v| v.* = @floatCast(rng.normal()); + var partials = try t.upload(host); + defer partials.free(); + var out = try t.alloc(rows * heads * d * 4); + defer out.free(); + var half = try t.alloc(rows * heads * d * 2); + defer half.free(); + const want = try gpa.alloc(f32, rows * heads * d); + defer gpa.free(want); + // a lane adds its column's tiles one after another from tile 0 + for (0..rows * heads) |row| for (0..d) |lane| { + var acc: f32 = 0; + for (0..tiles) |tile| acc += host[(row * tiles + tile) * d + lane]; + want[row * d + lane] = acc; + }; + const s = t.stream.handle; + var e: util.Args = .{}; + try e.add(ptr(partials)); + try e.add(ptr(out)); + try e.add(@as(c_int, @intCast(tiles))); + try e.add(@as(c_int, @intCast(d))); + try t.on.go(t.on.sum_partials, util.dim(heads, rows, 1), util.dim(256, 1, 1), 0, s, &e); + try t.stream.synchronize(); + const got = try rig.download(f32, out); + defer gpa.free(got); + for (got, want, 0..) |g, w, i| try rig.expect( + @as(u32, @bitCast(g)) == @as(u32, @bitCast(w)), + "sum of {d} tiles, d {d}, {d} rows: element {d} is {e}, the sum in tile order {e}", + .{ tiles, d, rows, i, g, w }, + ); + for ([_]c_int{ 1, 2 }) |kind| { + var a: util.Args = .{}; + try a.add(ptr(partials)); + try a.add(ptr(half)); + try a.add(@as(c_int, @intCast(tiles))); + try a.add(@as(c_int, @intCast(d))); + try a.add(kind); + try t.on.go(t.on.plan.sum_act, util.dim(heads, rows, 1), util.dim(256, 1, 1), 0, s, &a); + try t.stream.synchronize(); + const bits = try rig.download(u16, half); + defer gpa.free(bits); + for (bits, want, 0..) |b, w, i| { + const r: u16 = if (kind == 1) @bitCast(@as(f16, @floatCast(w))) else bf16Even(w); + try rig.expect(b == r, "sum into kind {d} of {d} tiles, d {d}: element {d} is {x}, the rounded sum {x}", .{ kind, tiles, d, i, b, r }); + } + } + cases += 1; + }; + std.debug.print(" partial sums: {d} cases, each lane's tiles in tile order\n", .{cases}); +} + +test "softmax stats take the row's max exactly and a sum that bounds the row, past the one-block walk" { + const t = try Rig.open(1 << 20); + defer t.close(); + var rng: rig.Rng = .{ .state = 0x57A7 }; + const heads = 4; + var cases: usize = 0; + for ([_][2]usize{ .{ 4352, 4097 }, .{ 16384, 16384 }, .{ 16384, 9001 }, .{ 65536, 65536 }, .{ 65536, 65535 } }) |sv| for ([_]usize{ 1, 2 }) |rows| { + const span = sv[0]; + const visible = sv[1]; + const host = try gpa.alloc(f32, rows * heads * span); + defer gpa.free(host); + // keys past the visible ones hold -inf, as the score kernel writes them + for (host, 0..) |*v, i| v.* = if (i % span < visible) @floatCast(8 * rng.normal()) else -std.math.inf(f32); + // each row's max, one above the rest, sits at a key of its own: last visible, first, past 4096, mid-row + const tops = [_]usize{ visible - 1, 0, 4096, visible / 2 + 255 }; + for (0..rows * heads) |row| { + const keys = host[row * span ..][0..visible]; + var top: f32 = -std.math.inf(f32); + for (keys) |k| top = @max(top, k); + keys[tops[row % tops.len]] = top + 1; + } + var scores = try t.upload(host); + defer scores.free(); + var stats = try t.alloc(rows * heads * 2 * 4); + defer stats.free(); + var b: util.Args = .{}; + try b.add(ptr(scores)); + try b.add(ptr(stats)); + try b.add(@as(c_int, @intCast(span))); + try b.add(@as(c_int, @intCast(visible))); + try b.add(@as(?*const c_int, null)); + try t.on.go(t.on.softmax_stats, util.dim(heads, rows, 1), util.dim(256, 1, 1), 0, t.stream.handle, &b); + try t.stream.synchronize(); + const got = try rig.download(f32, stats); + defer gpa.free(got); + for (0..rows * heads) |row| { + const keys = host[row * span ..][0..visible]; + var top: f32 = -std.math.inf(f32); + for (keys) |k| top = @max(top, k); + var sum: f64 = 0; + for (keys) |k| sum += @exp(@as(f64, k) - top); + const g = got[row * 2 ..][0..2]; + try rig.expect(@as(u32, @bitCast(g[0])) == @as(u32, @bitCast(top)), "stats over {d} of {d} keys: max {e}, the row's {e}", .{ visible, span, g[0], top }); + try rig.expect(@abs(g[1] - sum) <= 1e-5 * sum, "stats over {d} of {d} keys: sum {e}, float64 {e}", .{ visible, span, g[1], sum }); + } + cases += 1; + }; + std.debug.print(" softmax stats: {d} cases, the max exactly and the sum within 1e-5\n", .{cases}); +}