perf(hip): the decode tail issues a row's loads before its branches

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- HIP: the decode round's residual-add and norm tail issues each row's loads before its branches, and a MoE tail loads its slots' rows together, with output bytes unchanged. On a gfx1151, one-stream decode is 2.7% faster on Qwen3.6-35B-A3B and 1.1% on the 27B. `tf_tail` is now under the kernel contract (#2247). Thanks to @AbrahamCete.
+
 - Docs: `docs/optimization-ledger.md` lists which speed and memory optimizations apply to which native family, and where each stands, with links to the PRs that measured them (#2232). Thanks to @jschmied.
 
 - Docs: the pull request pre-flight asks a PR that adds, finishes or drops a family's optimization to update that family's cell in the optimization ledger (#2243). Thanks to @jschmied.

docs/kernel-contract.md

@@ -298,6 +298,9 @@ form no registration runs. A fixed plan keeps the linear. Gate 4 proves an fp16
 stress input (all-ones states: one weight at every width). Six controls: another K split (3, 4 and 5 bits; `linear_wc`),
 swapped scales, a swapped suh.
 HIP: the RMS norm rows, a wave a row (fp32, fp16 and bf16 out) and a block a row (fp16 and bf16 out).
+HIP also checks the decode tail (`tf_tail`): the residual add, of a row of the activation kind or of nine slots' fp32
+rows times their weights summed in slot order, and the RMS norm after it, against float64 at Qwen3.6-35B-A3B's and
+Qwen3.8-27B's widths, decode rows and windows, and the widest row the launch takes.
 HIP also checks the 64-row prompt attention tile (`prefill.hip`): `tf_fa_paged_bf16` and `tf_fa_paged_f16` over a
 permuted page table at the 27B's, the 35B-A3B's and Nemotron's head counts, each output within attention_ref's float64
 bound, and `tf_fa_wide_bf16` and `tf_fa_wide_f16`, which no model calls, bit-equal to them on flat keys. A build runs

zig/kernels/hip/decode/decode.hip

@@ -152,32 +152,67 @@ __device__ __forceinline__ void store_kind(void* p, int kind, long long i, float
 
 constexpr int kTailThreads = 512;
 constexpr int kTailWidth = 8192;
+constexpr int kTailAhead = 4;  // a thread's elements loaded before any is used, so their latencies overlap
+constexpr int kTailSlots = 8;  // the slots whose rows are loaded before any is summed
 
-// One row a block, one launch: x = round(x + t), t = y or the slot-ordered weighted sum; normed = rms(x) * weight.
-__device__ void tail_body(void* x, const void* y, const float* wts, const float* weight, void* normed, int kind,
-                          int slots, int width, float eps) {
-    __shared__ float row_v[kTailWidth];
-    __shared__ float part[kTailThreads / 32];
+// A thread's elements of the row: x = round(x + t), t = y or (SLOTS) the slots' sum in slot order; its sum of v * v.
+template <bool SLOTS>
+__device__ __forceinline__ float tail_add(void* x, const void* y, const float* wts, float* row_v, int kind, int slots,
+                                          int width) {
     const int row = blockIdx.x;
     const long long at = static_cast<long long>(row) * width;
     float sum = 0.f;
-    for (int i = threadIdx.x; i < width; i += kTailThreads) {
-        float t;
-        if (slots > 0) {
-            float acc = 0.f;
-            for (int s = 0; s < slots; ++s) {
-                acc += static_cast<const float*>(y)[(static_cast<long long>(row) * slots + s) * width + i] *
-                       wts[row * slots + s];
+    for (int base = threadIdx.x; base < width; base += kTailAhead * kTailThreads) {
+        // an element past the row loads the row's last, so no load waits behind a branch
+        int col[kTailAhead];
+        float xv[kTailAhead], t[kTailAhead];
+#pragma unroll
+        for (int k = 0; k < kTailAhead; ++k) {
+            col[k] = min(base + k * kTailThreads, width - 1);
+            xv[k] = load_kind(x, kind, at + col[k]);
+            t[k] = SLOTS ? 0.f : load_kind(y, kind, at + col[k]);
+        }
+        // each element's slots in slot order, a group of slots' loads issued before their sums
+        if constexpr (SLOTS) {
+            for (int s0 = 0; s0 < slots; s0 += kTailSlots) {
+                float yv[kTailSlots][kTailAhead];
+#pragma unroll
+                for (int j = 0; j < kTailSlots; ++j) {
+                    if (s0 + j >= slots) break;
+                    const long long first = (static_cast<long long>(row) * slots + s0 + j) * width;
+                    const float* ys = static_cast<const float*>(y) + first;
+#pragma unroll
+                    for (int k = 0; k < kTailAhead; ++k) yv[j][k] = ys[col[k]];
+                }
+#pragma unroll
+                for (int j = 0; j < kTailSlots; ++j) {
+                    if (s0 + j >= slots) break;
+                    const float w = wts[row * slots + s0 + j];
+#pragma unroll
+                    for (int k = 0; k < kTailAhead; ++k) t[k] += yv[j][k] * w;
+                }
             }
-            t = round_kind(acc, kind);
-        } else {
-            t = load_kind(y, kind, at + i);
         }
-        const float v = round_kind(load_kind(x, kind, at + i) + t, kind);
-        store_kind(x, kind, at + i, v);
-        row_v[i] = v;
-        sum += v * v;
+#pragma unroll
+        for (int k = 0; k < kTailAhead; ++k) {
+            const int i = base + k * kTailThreads;
+            if (i >= width) break;
+            const float v = round_kind(xv[k] + (SLOTS ? round_kind(t[k], kind) : t[k]), kind);
+            store_kind(x, kind, at + i, v);
+            row_v[i] = v;
+            sum += v * v;
+        }
     }
+    return sum;
+}
+
+// tail_body at one activation kind, so a thread's loads and stores take no branch on it.
+template <int KIND>
+__device__ __forceinline__ void tail_kind(void* x, const void* y, const float* wts, const float* weight, void* normed,
+                                          float* row_v, float* part, int slots, int width, float eps) {
+    const long long at = static_cast<long long>(blockIdx.x) * width;
+    float sum = slots > 0 ? tail_add<true>(x, y, wts, row_v, KIND, slots, width)
+                          : tail_add<false>(x, y, wts, row_v, KIND, slots, width);
 #pragma unroll
     for (int mask = 16; mask > 0; mask >>= 1) sum += __shfl_xor(sum, mask, 32);
     if ((threadIdx.x & 31) == 0) part[threadIdx.x >> 5] = sum;
@@ -186,10 +221,32 @@ __device__ void tail_body(void* x, const void* y, const float* wts, const float*
 #pragma unroll
     for (int w = 0; w < kTailThreads / 32; ++w) total += part[w];
     const float inv = rsqrtf(total / static_cast<float>(width) + eps);
-    for (int i = threadIdx.x; i < width; i += kTailThreads) {
-        float value = row_v[i] * inv;
-        if (weight != nullptr) value *= weight[i];
-        store_kind(normed, kind, at + i, value);
+    for (int base = threadIdx.x; base < width; base += kTailAhead * kTailThreads) {
+        float w[kTailAhead];
+#pragma unroll
+        for (int k = 0; k < kTailAhead; ++k) {
+            w[k] = weight != nullptr ? weight[min(base + k * kTailThreads, width - 1)] : 1.f;
+        }
+#pragma unroll
+        for (int k = 0; k < kTailAhead; ++k) {
+            const int i = base + k * kTailThreads;
+            if (i >= width) break;
+            float value = row_v[i] * inv;
+            if (weight != nullptr) value *= w[k];
+            store_kind(normed, KIND, at + i, value);
+        }
+    }
+}
+
+// One row a block, one launch: x = round(x + t), t = y or the slot-ordered weighted sum; normed = rms(x) * weight.
+__device__ void tail_body(void* x, const void* y, const float* wts, const float* weight, void* normed, int kind,
+                          int slots, int width, float eps) {
+    __shared__ float row_v[kTailWidth];
+    __shared__ float part[kTailThreads / 32];
+    if (kind == 1) {
+        tail_kind<1>(x, y, wts, weight, normed, row_v, part, slots, width, eps);
+    } else {
+        tail_kind<2>(x, y, wts, weight, normed, row_v, part, slots, width, eps);
     }
 }
 

zig/src/core/contract/hip_records.zig

@@ -494,7 +494,7 @@ const own_records = [_]Record{
         },
     },
 } ++ @import("hip_records_pair.zig").records ++ @import("hip_records_prompt.zig").records ++
-    @import("hip_records_score.zig").records ++ @import("hip_records_embed.zig").records;
+    @import("hip_records_score.zig").records ++ @import("hip_records_embed.zig").records ++ @import("hip_records_tail.zig").records;
 
 pub const records = own_records ++ @import("hip_records_stream.zig").records ++ @import("hip_records_router.zig").records ++
     @import("hip_records_item.zig").records ++ @import("hip_records_qwen3.zig").records ++ @import("hip_records_columns.zig").records ++

zig/src/core/contract/hip_records_families.zig

@@ -16,6 +16,6 @@ pub const all = join: {
 
 test "the HIP records stay main's; the family lists only regroup GLM-OCR's and the joint schema's" {
     const every = @import("hip_records.zig").records;
-    try std.testing.expectEqual(@as(usize, 192), every.len); // main's records, GLM-OCR's 12 and the joint schema's 10 among them
-    try std.testing.expectEqual(@as(usize, 170), every.len - all.len);
+    try std.testing.expectEqual(@as(usize, 193), every.len); // main's records, GLM-OCR's 12 and the joint schema's 10 among them
+    try std.testing.expectEqual(@as(usize, 171), every.len - all.len);
 }

zig/src/core/contract/hip_records_tail.zig

@@ -0,0 +1,12 @@
+//! Gate 5's record of the decode tail (decode.hip: tf_tail): a row a block of 512 threads, the row held in LDS.
+const Record = @import("hip_records.zig").Record;
+
+pub const records = [_]Record{.{
+    .launch = .{ .name = "decode tail tf_tail", .symbol = "tf_tail", .image = "decode", .threads = 512 },
+    .usage = .{
+        .{ .vgprs = 125, .sgprs = 58, .lds = 32832, .private = 0, .bound = 512, .wave = 32 },
+        .{ .vgprs = 120, .sgprs = 58, .lds = 32832, .private = 0, .bound = 512, .wave = 32 },
+        .{ .vgprs = 120, .sgprs = 58, .lds = 32832, .private = 0, .bound = 512, .wave = 32 },
+        .{ .vgprs = 120, .sgprs = 58, .lds = 32832, .private = 0, .bound = 512, .wave = 32 },
+    },
+}};

zig/src/hip/kernel_tests.zig

@@ -21,6 +21,7 @@ test {
     _ = @import("kernel_tests/gdn.zig");
     _ = @import("kernel_tests/mamba.zig");
     _ = @import("kernel_tests/decode.zig");
+    _ = @import("kernel_tests/tail_pin.zig");
     _ = @import("kernel_tests/chain.zig");
     _ = @import("kernel_tests/attention.zig");
     _ = @import("kernel_tests/attention_pin.zig");

zig/src/hip/kernel_tests/tail_pin.zig

@@ -0,0 +1,101 @@
+//! The decode tail's exact output bits (the new residual and its norm), pinned by a digest per architecture.
+const std = @import("std");
+const rig = @import("rig.zig");
+const Rig = rig.Rig;
+const gpa = rig.gpa;
+const at = rig.at;
+const kernels = @import("../kernels.zig");
+
+/// The digest of every case's residual and norm, in both activation kinds, on an architecture whose bits were recorded.
+const pins = [_]struct { arch: []const u8, digest: u64 }{
+    .{ .arch = "gfx1151", .digest = 0x325db42cb918dd99 },
+};
+
+/// `rows` rows of `width` plus a plain row (`slots` 0) or the slots' weighted sum; `cancel` makes the slot order show.
+const Case = struct { rows: usize, width: usize, slots: usize, weight: bool = true, cancel: bool = false };
+
+/// Cancelling slots (past a slot group too), one row and a window at the 35B's and 27B's widths, ragged, unweighted.
+const cases = [_]Case{
+    .{ .rows = 2, .width = 2048, .slots = 3, .cancel = true },
+    .{ .rows = 2, .width = 2048, .slots = 9, .cancel = true },
+    .{ .rows = 1, .width = 2048, .slots = 0 },
+    .{ .rows = 1, .width = 2048, .slots = 9 },
+    .{ .rows = 4, .width = 2048, .slots = 9 },
+    .{ .rows = 1, .width = 5120, .slots = 0 },
+    .{ .rows = 3, .width = 5120, .slots = 0 },
+    .{ .rows = 2, .width = 1000, .slots = 3 },
+    .{ .rows = 2, .width = 2048, .slots = 0, .weight = false },
+    .{ .rows = 1, .width = 8192, .slots = 2 },
+};
+
+/// FNV-1a over `bytes`, continuing from `h`.
+fn fnv(h: u64, bytes: []const u8) u64 {
+    var x = h;
+    for (bytes) |b| {
+        x ^= b;
+        x *%= 0x100000001b3;
+    }
+    return x;
+}
+
+/// One case's residual and norm through tf_tail at activation kind `kind` (1 fp16, 2 bf16), folded into `h`.
+fn run(t: *Rig, rng: *rig.Rng, c: Case, kind: c_int, h: u64) !u64 {
+    const bits = if (kind == 1) &rig.f16Bits else &rig.bf16Bits;
+    const n = c.rows * c.width;
+    const hx = try gpa.alloc(u16, n);
+    defer gpa.free(hx);
+    for (hx) |*v| v.* = bits(rng.unit());
+    const hy16 = try gpa.alloc(u16, if (c.slots > 0) 0 else n);
+    defer gpa.free(hy16);
+    for (hy16) |*v| v.* = bits(rng.unit() / 8.0);
+    const hy32 = try gpa.alloc(f32, c.rows * c.slots * c.width);
+    defer gpa.free(hy32);
+    for (hy32, 0..) |*v, i| {
+        const slot = (i / c.width) % @max(c.slots, 1);
+        const big: f32 = if (!c.cancel or slot > 1) 0 else if (slot == 0) 4096 else -4096;
+        v.* = big + rng.unit() / 8.0;
+    }
+    const hw = try gpa.alloc(f32, c.rows * @max(c.slots, 1));
+    defer gpa.free(hw);
+    for (hw) |*v| v.* = if (c.cancel) 0.5 else (rng.unit() + 1.0) / @as(f32, @floatFromInt(@max(c.slots, 1)));
+    const hn = try gpa.alloc(f32, c.width);
+    defer gpa.free(hn);
+    for (hn) |*v| v.* = 1.0 + rng.unit() / 4.0;
+    var x = try t.upload(hx);
+    defer x.free();
+    var y = if (c.slots > 0) try t.upload(hy32) else try t.upload(hy16);
+    defer y.free();
+    var wts = try t.upload(hw);
+    defer wts.free();
+    var weight = try t.upload(hn);
+    defer weight.free();
+    var normed = try t.alloc(n * 2);
+    defer normed.free();
+    const wp: ?[*]const f32 = if (c.slots > 0) @ptrFromInt(at(wts)) else null;
+    const np: ?[*]const f32 = if (c.weight) @ptrFromInt(at(weight)) else null;
+    try t.on.tf_tail(@ptrFromInt(at(x)), @ptrFromInt(at(y)), wp, np, @ptrFromInt(at(normed)), kind, @intCast(c.rows), @intCast(c.slots), @intCast(c.width), 1e-6, t.stream.handle);
+    try t.stream.synchronize();
+    const gx = try rig.download(u16, x);
+    defer gpa.free(gx);
+    const gn = try rig.download(u16, normed);
+    defer gpa.free(gn);
+    return fnv(fnv(h, std.mem.sliceAsBytes(gx)), std.mem.sliceAsBytes(gn));
+}
+
+test "the decode tail writes its pinned residual and norm bits in fp16 and bf16, plain, slot-summed, ragged and wide" {
+    const t = try Rig.open(1 << 20);
+    defer t.close();
+    var rng: rig.Rng = .{ .state = 0x9E3779B97F4A7C15 };
+    var h: u64 = 0xcbf29ce484222325;
+    for ([_]c_int{ 2, 1 }) |kind| {
+        for (cases) |c| h = try run(t, &rng, c, kind, h);
+    }
+    const want = for (pins) |p| {
+        if (std.mem.eql(u8, p.arch, kernels.arch)) break p.digest;
+    } else {
+        std.debug.print("decode tail: no digest pinned for {s}; this run wrote 0x{x:0>16}\n", .{ kernels.arch, h });
+        return error.SkipZigTest;
+    };
+    if (h != want) std.debug.print("decode tail digest on {s}: 0x{x:0>16}\n", .{ kernels.arch, h });
+    try std.testing.expectEqual(want, h);
+}

zig/tests/hip/contract_main.zig

@@ -15,6 +15,7 @@ const moe_k = @import("contract_nemo_moe.zig");
 const walk_k = @import("contract_walk.zig");
 const prefill_k = @import("contract_prefill.zig");
 const add_rms_k = @import("contract_add_rms.zig");
+const tail_k = @import("contract_tail.zig");
 const window_k = @import("contract_window.zig");
 const relu2_k = @import("contract_relu2.zig");
 const kv_k = @import("contract_kv_write.zig");
@@ -191,6 +192,10 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize {
     for (&add_rms, std.enums.values(add_rms_k.Mode)) |*m, w| m.* = .init(a, &r, s, &l, w);
     defer for (&add_rms) |*m| m.deinit();
     for (&add_rms) |*m| try kernels.append(a, m.kernel());
+    var tails: [2]tail_k.Tail = undefined;
+    for (&tails, std.enums.values(tail_k.Mode)) |*m, w| m.* = .init(a, &r, s, &l, w);
+    defer for (&tails) |*m| m.deinit();
+    for (&tails) |*m| try kernels.append(a, m.kernel());
     var convs: [2]window_k.ConvWindow = undefined;
     for (&convs, std.enums.values(window_k.Part)) |*m, p| m.init(a, &r, s, &l, p);
     defer for (&convs) |*m| m.deinit();