perf(hip): prompt GEMMs write the activation type, not fp32 and a cast

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: prompt projections round to the activation type in their tile instead of writing fp32 and launching a cast, with the same bytes; cold prompts are 2% to 3% faster on Qwen3.8-27B on gfx1151 (#1685). Thanks to @AbrahamCete.
+
 - On Macs, DeepSeek V4.1's MXFP4 experts have their own Metal kernels under the kernel contract, with Kimi's left as they are (#1707, part of #1676). Thanks to @gilby.
 
 - On Macs, DeepSeek V4.1's FP8 block-scaled rows and tiles have their own Metal kernels under the kernel contract, ready for the release checkpoint's serving parts (#1706, part of #1676). Thanks to @gilby.

zig/kernels/hip/quant/mlx.hpp

@@ -35,7 +35,7 @@ struct Affine {
     int m, n, k, bits, group;
     int fp16;  // 0 is bf16, 1 is fp16
     Routing route = {};
-    __half* out16 = nullptr;  // set: the decode tile rounds each output to fp16 here instead of writing out
+    __half* out16 = nullptr;  // set: the tile rounds each output to the activation type here instead of writing out
 };
 
 // A routed block takes item ``z``: its expert's weights and its rows. False when the item has no rows.

zig/kernels/hip/quant/mlx_tiles.hpp

@@ -20,21 +20,21 @@ namespace rocm {
 
 template <typename T, int BITS, int RT>
 __global__ void __launch_bounds__(256) affine_gemm_block(Affine a) {
-    gemm_tile<MlxDecoder<BITS>, IdentityAct<T>, T, F32Out, RT>(a);
+    gemm_tile<MlxDecoder<BITS>, IdentityAct<T>, T, GemmOut<T>, RT>(a);
 }
 
 // ---- the K-parallel tiles of a short prompt ----
 
 template <typename T, int BITS, int CB, int R, int WAVES, int RB, bool LOOP>
 __global__ void __launch_bounds__(32 * WAVES) affine_gemm_kp(Affine a) {
-    gemm_kp_tile<MlxDecoder<BITS>, IdentityAct<T>, T, F32Out, CB, R, WAVES, RB, LOOP>(a);
+    gemm_kp_tile<MlxDecoder<BITS>, IdentityAct<T>, T, GemmOut<T>, CB, R, WAVES, RB, LOOP>(a);
 }
 
 // ---- the prefill GEMM tile on the matrix cores of gfx11 ----
 
 template <int BITS>
 __global__ void __launch_bounds__(256) affine_wmma_gemm(Affine a) {
-    wmma_gemm_tile<MlxDecoder<BITS>, BF16Act, WmmaBF16, F32Out>(a);
+    wmma_gemm_tile<MlxDecoder<BITS>, BF16Act, WmmaBF16, GemmOut<WmmaBF16>>(a);
 }
 
 // ---- the decode stream tile ----

zig/kernels/hip/tiles/dot2_tiles.hip

@@ -326,7 +326,7 @@ __global__ void __launch_bounds__(256, 1) affine_dot2_block(Affine a) {
 #pragma unroll
         for (int j = 0; j < 8; ++j) {
             const int col = n0 + tx + 16 * j;
-            if (col < a.n) a.out[out_row(a, row) * a.n + col] = acc[i][j];
+            if (col < a.n) GemmOut<T>::store(a, row, col, acc[i][j]);
         }
     }
 }

zig/kernels/hip/tiles/epilogue.hpp

@@ -11,14 +11,6 @@
 namespace tf {
 namespace rocm {
 
-// The product in fp32 at its (row, col) of out, the plan's pair row for a routed item.
-struct F32Out {
-    template <class Args>
-    __device__ static void store(const Args& a, int row, int col, float v) {
-        a.out[out_row(a, row) * a.n + col] = v;
-    }
-};
-
 // A float as the activation type, rounded to nearest even.
 template <typename T>
 __device__ inline typename T::elem narrow(float v) {
@@ -31,6 +23,19 @@ __device__ inline typename T::elem narrow(float v) {
     }
 }
 
+// The product at its (row, col), the plan's pair row for a routed item: rounded to T into out16 when set, else fp32.
+template <typename T>
+struct GemmOut {
+    template <class Args>
+    __device__ static void store(const Args& a, int row, int col, float v) {
+        if (a.out16 != nullptr) {
+            static_cast<typename T::elem*>(static_cast<void*>(a.out16))[out_row(a, row) * a.n + col] = narrow<T>(v);
+        } else {
+            a.out[out_row(a, row) * a.n + col] = v;
+        }
+    }
+};
+
 // A float product rounded to T: FP16 is one rounding of the exact product (fma, zero addend), never fp32 then FP16.
 template <typename T>
 __device__ inline typename T::elem narrow_product(float a, float b) {

zig/src/hip/kernel_tests.zig

@@ -30,6 +30,7 @@ test {
     _ = @import("kernel_tests/affine/affine.zig");
     _ = @import("kernel_tests/affine/rows.zig");
     _ = @import("kernel_tests/affine/order.zig");
+    _ = @import("kernel_tests/affine/round.zig");
     _ = @import("kernel_tests/graph_split.zig");
     _ = @import("kernel_tests/flashnext.zig");
     _ = @import("kernel_tests/flashnext_attn.zig");

zig/src/hip/kernel_tests/affine/round.zig

@@ -0,0 +1,63 @@
+//! A prompt tile given out16 writes its sum there rounded as tf_cast rounds the fp32 one, and leaves out untouched.
+
+const std = @import("std");
+const base = @import("../rig.zig");
+const Buffer = base.Buffer;
+const cases = @import("cases.zig");
+const prod = @import("product.zig");
+const Product = prod.Product;
+const Rig = prod.Rig;
+
+fn dense(name: []const u8, rows: usize, group: u16, bits: u8) cases.Case {
+    return .{ .name = name, .path = .prefill, .rows = rows, .n = 1001, .k = 2112, .group = group, .bits = bits };
+}
+
+/// Row counts on every prompt tile (the K-parallel ones, then the 128-row tiles and their edges), groups and widths.
+const shapes = [_]cases.Case{
+    dense("1 row", 1, 64, 4),      dense("2 rows", 2, 64, 4),       dense("3 rows", 3, 64, 4),
+    dense("4 rows", 4, 64, 4),     dense("17 rows", 17, 64, 4),     dense("32 rows", 32, 64, 4),
+    dense("33 rows", 33, 64, 4),   dense("100 rows", 100, 64, 4),   dense("300 rows", 300, 64, 4),
+    dense("group 32", 100, 32, 4), dense("group 128", 100, 128, 4), dense("8-bit", 100, 64, 8),
+};
+
+fn digest(buf: Buffer, len: usize) !u64 {
+    return prod.digestOf(buf, len);
+}
+
+test "a prompt entry with out16 set writes tf_cast's bytes of its fp32 sum and leaves out untouched" {
+    const b = try base.Rig.open(1 << 20);
+    defer b.close();
+    var rig = Rig.open(b);
+    var checked: usize = 0;
+    for (shapes) |c| {
+        var p = try Product.init(&rig, c);
+        defer p.deinit();
+        const n = p.out_rows * c.n;
+        var wide = try Buffer.alloc(rig.base, n * 4);
+        defer wide.free();
+        var want = try Buffer.alloc(rig.base, n * 2);
+        defer want.free();
+        var got = try Buffer.alloc(rig.base, n * 2);
+        defer got.free();
+        var spare = try Buffer.alloc(rig.base, n * 4);
+        defer spare.free();
+        try spare.fill8(0xA5, null);
+        const untouched = try digest(spare, n * 4);
+        for (rig.kernels.reg.entries) |*e| {
+            if (e.op != .project or e.path != .prefill or !prod.takes(&p, e)) continue;
+            try wide.fill8(0xA5, null);
+            try prod.launchEntry(&p, e, wide, false, 0, 0);
+            try rig.base.on.tf_cast(@ptrFromInt(wide.ptr), 0, @ptrFromInt(want.ptr), rig.base.kind(), @intCast(n), rig.base.stream.handle);
+            try got.fill8(0x5A, null);
+            var a = p.arg(spare.ptr, 0);
+            a.out16 = got.ptr;
+            try e.launch(rig.kernels, .{ .r = &rig.base.r, .s = rig.base.stream.handle, .arg = a, .items = @intCast(p.items), .partial = 0, .parts = 1 });
+            try rig.base.stream.synchronize();
+            try base.expect(try digest(got, n * 2) == try digest(want, n * 2), "round {s}: {s} with out16 writes other bytes than tf_cast of its fp32 sum", .{ c.name, e.id });
+            try base.expect(try digest(spare, n * 4) == untouched, "round {s}: {s} with out16 also writes out", .{ c.name, e.id });
+            checked += 1;
+        }
+    }
+    std.debug.print("ROUND {d} prompt launches write tf_cast's bytes into out16\n", .{checked});
+    try base.expect(checked >= shapes.len, "round: only {d} prompt launches ran", .{checked});
+}

zig/src/hip/launches/mlx_entries.zig

@@ -304,14 +304,16 @@ const other_decode = [_]Entry{
 };
 
 fn prefill(comptime op: Op) [7]Entry {
+    // a plain prompt product rounds into out16 when asked; a routed one writes fp32 for the expert combine
+    const r: registry.Rounds = if (op == .project) .always else .never;
     return .{
-        entry(op, .prefill, "kp0", .{ .family = .gemm, .fits = kpFits(0), .launch = launchKp(0), .policy = &gemmOn }),
-        entry(op, .prefill, "kp1", .{ .family = .gemm, .fits = kpFits(1), .launch = launchKp(1), .policy = &gemmOn }),
-        entry(op, .prefill, "kp2", .{ .family = .gemm, .fits = kpFits(2), .launch = launchKp(2), .policy = &gemmOn }),
-        entry(op, .prefill, "kp3", .{ .family = .gemm, .fits = kpFits(3), .launch = launchKp(3), .policy = &gemmOn }),
-        entry(op, .prefill, "matrix", .{ .family = .gemm, .fits = &fitsPrefillMatrix, .launch = launchBlock(.matrix), .caps = &matrixCores, .policy = &matrixOn }),
-        entry(op, .prefill, "gemm", .{ .family = .gemm, .fits = &fitsPrefillGemm, .launch = launchBlock(.gemm), .policy = &gemmOn }),
-        entry(op, .prefill, "block", .{ .family = .gemm, .fits = &fitsPrefillBlock, .launch = launchBlock(.block) }),
+        entry(op, .prefill, "kp0", .{ .family = .gemm, .fits = kpFits(0), .launch = launchKp(0), .policy = &gemmOn, .rounds = r }),
+        entry(op, .prefill, "kp1", .{ .family = .gemm, .fits = kpFits(1), .launch = launchKp(1), .policy = &gemmOn, .rounds = r }),
+        entry(op, .prefill, "kp2", .{ .family = .gemm, .fits = kpFits(2), .launch = launchKp(2), .policy = &gemmOn, .rounds = r }),
+        entry(op, .prefill, "kp3", .{ .family = .gemm, .fits = kpFits(3), .launch = launchKp(3), .policy = &gemmOn, .rounds = r }),
+        entry(op, .prefill, "matrix", .{ .family = .gemm, .fits = &fitsPrefillMatrix, .launch = launchBlock(.matrix), .caps = &matrixCores, .policy = &matrixOn, .rounds = r }),
+        entry(op, .prefill, "gemm", .{ .family = .gemm, .fits = &fitsPrefillGemm, .launch = launchBlock(.gemm), .policy = &gemmOn, .rounds = r }),
+        entry(op, .prefill, "block", .{ .family = .gemm, .fits = &fitsPrefillBlock, .launch = launchBlock(.block), .rounds = r }),
     };
 }
 

zig/src/hip/ops/project.zig

@@ -236,7 +236,7 @@ pub fn affine(o: Ops, x: Tensor, w: Affine, m: usize, f32_out: bool) Error!Tenso
     const z = o.l;
     const shape = shapeOf(o, x, w, m);
     const tile = z.affine.choose(.project, if (o.prefill) .prefill else .decode, shape) orelse return error.BadShape;
-    half = !f32_out and !o.prefill and tile.roundsAct(shape);
+    half = !f32_out and tile.roundsAct(shape);
     const out = try o.arena.take(m * n * @as(usize, if (half) 2 else 4));
     {
         const arg: affine_launch.Arg = .{