diff --git a/CHANGELOG.md b/CHANGELOG.md index d73cefe89b..20f8e37c63 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 +- 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. diff --git a/zig/kernels/hip/quant/mlx.hpp b/zig/kernels/hip/quant/mlx.hpp index 9d376a277d..abcbf908b4 100644 --- a/zig/kernels/hip/quant/mlx.hpp +++ b/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. diff --git a/zig/kernels/hip/quant/mlx_tiles.hpp b/zig/kernels/hip/quant/mlx_tiles.hpp index ec4f67f515..79f353d849 100644 --- a/zig/kernels/hip/quant/mlx_tiles.hpp +++ b/zig/kernels/hip/quant/mlx_tiles.hpp @@ -20,21 +20,21 @@ namespace rocm { template __global__ void __launch_bounds__(256) affine_gemm_block(Affine a) { - gemm_tile, IdentityAct, T, F32Out, RT>(a); + gemm_tile, IdentityAct, T, GemmOut, RT>(a); } // ---- the K-parallel tiles of a short prompt ---- template __global__ void __launch_bounds__(32 * WAVES) affine_gemm_kp(Affine a) { - gemm_kp_tile, IdentityAct, T, F32Out, CB, R, WAVES, RB, LOOP>(a); + gemm_kp_tile, IdentityAct, T, GemmOut, CB, R, WAVES, RB, LOOP>(a); } // ---- the prefill GEMM tile on the matrix cores of gfx11 ---- template __global__ void __launch_bounds__(256) affine_wmma_gemm(Affine a) { - wmma_gemm_tile, BF16Act, WmmaBF16, F32Out>(a); + wmma_gemm_tile, BF16Act, WmmaBF16, GemmOut>(a); } // ---- the decode stream tile ---- diff --git a/zig/kernels/hip/tiles/dot2_tiles.hip b/zig/kernels/hip/tiles/dot2_tiles.hip index 8bde31cd18..5dd79585ee 100644 --- a/zig/kernels/hip/tiles/dot2_tiles.hip +++ b/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::store(a, row, col, acc[i][j]); } } } diff --git a/zig/kernels/hip/tiles/epilogue.hpp b/zig/kernels/hip/tiles/epilogue.hpp index 97047ee727..0483313c44 100644 --- a/zig/kernels/hip/tiles/epilogue.hpp +++ b/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 - __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 __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 +struct GemmOut { + template + __device__ static void store(const Args& a, int row, int col, float v) { + if (a.out16 != nullptr) { + static_cast(static_cast(a.out16))[out_row(a, row) * a.n + col] = narrow(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 __device__ inline typename T::elem narrow_product(float a, float b) { diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index 934db3f04a..c1f3e1b506 100644 --- a/zig/src/hip/kernel_tests.zig +++ 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"); diff --git a/zig/src/hip/kernel_tests/affine/round.zig b/zig/src/hip/kernel_tests/affine/round.zig new file mode 100644 index 0000000000..780944f05b --- /dev/null +++ b/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}); +} diff --git a/zig/src/hip/launches/mlx_entries.zig b/zig/src/hip/launches/mlx_entries.zig index 9d71533476..55c3cf042b 100644 --- a/zig/src/hip/launches/mlx_entries.zig +++ b/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 }), }; } diff --git a/zig/src/hip/ops/project.zig b/zig/src/hip/ops/project.zig index 1333f090a4..a4d7e38914 100644 --- a/zig/src/hip/ops/project.zig +++ b/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 = .{