perf(hip): the decode softmax stats and partial sums load sixteen ahead

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.

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<long long>(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<long long>(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<long long>(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<long long>(tile + u) * d];
+#pragma unroll
+        for (int u = 0; u < 16; ++u) acc += ahead[u];
+    }
+    for (; tile < tiles; ++tile) acc += col[static_cast<long long>(tile) * d];
     out[(static_cast<long long>(batch) * gridDim.x + head) * d + lane] = acc;
 }
 

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<long long>(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<long long>(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<long long>(tile + u) * d];
+#pragma unroll
+        for (int u = 0; u < 16; ++u) acc += ahead[u];
+    }
+    for (; tile < tiles; ++tile) acc += col[static_cast<long long>(tile) * d];
     store(out, kind, (static_cast<long long>(batch) * gridDim.x + head) * d + lane, acc);
 }
 __global__ void __launch_bounds__(kAttendThreads)

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 },
         },
     },
     .{

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");

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});
+}