tiny.accy.kernel.library.attention

1const std = @import("std");2const gpu = @import("gpu");34const entry = @import("entry.zig");5const kernel = @import("../root.zig");67pub const ScaledDotProduct = struct {8    batch: u64,9    queries: u64,10    keys: u64,11    qk: u64,12    value: u64,13    threads: entry.Threads3D = .{},14    batch_axis: []const u8 = "b",15    query_axis: []const u8 = "q",16    key_axis: []const u8 = "k",17    qk_axis: []const u8 = "h",18    value_axis: []const u8 = "v",19};2021fn indexUpper(comptime extent: u64) i64 {22    if (extent > @as(u64, @intCast(std.math.maxInt(i64)))) {23        @compileError("kernel library attention extent overflows index range");24    }25    return @intCast(extent);26}2728fn checkedProduct(comptime lhs: u64, comptime rhs: u64) u64 {29    return std.math.mul(u64, lhs, rhs) catch @compileError("kernel library attention shape product overflow");30}3132fn floatExtent(comptime extent: u64) f64 {33    return @floatFromInt(extent);34}3536fn queryShape(comptime spec: ScaledDotProduct) entry.Shape {37    return entry.shape3D(spec.batch_axis, spec.batch, spec.query_axis, spec.queries, spec.qk_axis, spec.qk);38}3940fn keyShape(comptime spec: ScaledDotProduct) entry.Shape {41    return entry.shape3D(spec.batch_axis, spec.batch, spec.key_axis, spec.keys, spec.qk_axis, spec.qk);42}4344fn valueShape(comptime spec: ScaledDotProduct) entry.Shape {45    return entry.shape3D(spec.batch_axis, spec.batch, spec.key_axis, spec.keys, spec.value_axis, spec.value);46}4748fn outputShape(comptime spec: ScaledDotProduct) entry.Shape {49    return entry.shape3D(spec.batch_axis, spec.batch, spec.query_axis, spec.queries, spec.value_axis, spec.value);50}5152fn queryRowShape(comptime spec: ScaledDotProduct) entry.Shape {53    return entry.shape2D(spec.batch_axis, spec.batch, spec.query_axis, spec.queries);54}5556fn scoreShape(comptime spec: ScaledDotProduct) entry.Shape {57    return entry.shape3D(spec.batch_axis, spec.batch, spec.query_axis, spec.queries, spec.key_axis, spec.keys);58}5960fn scaledDotProductSpecialization(comptime spec: ScaledDotProduct) entry.Specialization {61    return .{62        .dtype = .f32,63        .operation = .{ .attention = .scaled_dot_product },64        .equation = "bqh,bkh,bkv->bqv",65        .inputs = &.{66            queryShape(spec),67            keyShape(spec),68            valueShape(spec),69        },70        .outputs = &.{outputShape(spec)},71        .reductions = &.{72            entry.reduction("score_dot", .dot_product, entry.shape1D(spec.qk_axis, spec.qk)),73            entry.dependentReduction("score_max", .maximum, entry.shape1D(spec.key_axis, spec.keys), &.{"score_dot"}),74            entry.dependentReduction("score_exp_sum", .sum_exp_shifted, entry.shape1D(spec.key_axis, spec.keys), &.{"score_max"}),75            entry.dependentReduction("value_weighted_sum", .weighted_sum, entry.shape1D(spec.key_axis, spec.keys), &.{"score_exp_sum"}),76        },77        .reduction_reuse = &.{78            entry.reductionReuse("score_dot", scoreShape(spec)),79            entry.reductionReuse("score_max", queryRowShape(spec)),80            entry.reductionReuse("score_exp_sum", queryRowShape(spec)),81        },82        .launch = entry.launch3D(spec.value, spec.queries, spec.batch, spec.threads.x, spec.threads.y, spec.threads.z),83        .schedule = entry.threadBlocks3D(spec.value_axis, spec.value, spec.query_axis, spec.queries, spec.batch_axis, spec.batch, spec.threads.x, spec.threads.y, spec.threads.z),84    };85}8687fn queryIndex(88    inner: anytype,89    comptime spec: ScaledDotProduct,90    batch: kernel.Value,91    query: kernel.Value,92    feature: kernel.Value,93) !kernel.Value {94    const batch_stride = try inner.constantIndex(indexUpper(checkedProduct(spec.queries, spec.qk)));95    const query_stride = try inner.constantIndex(indexUpper(spec.qk));96    const batch_offset = try inner.mul(batch, batch_stride);97    const query_offset = try inner.mul(query, query_stride);98    const batch_query_offset = try inner.add(batch_offset, query_offset);99    return inner.add(batch_query_offset, feature);100}101102fn keyIndex(103    inner: anytype,104    comptime spec: ScaledDotProduct,105    batch: kernel.Value,106    key: kernel.Value,107    feature: kernel.Value,108) !kernel.Value {109    const batch_stride = try inner.constantIndex(indexUpper(checkedProduct(spec.keys, spec.qk)));110    const key_stride = try inner.constantIndex(indexUpper(spec.qk));111    const batch_offset = try inner.mul(batch, batch_stride);112    const key_offset = try inner.mul(key, key_stride);113    const batch_key_offset = try inner.add(batch_offset, key_offset);114    return inner.add(batch_key_offset, feature);115}116117fn valueIndex(118    inner: anytype,119    comptime spec: ScaledDotProduct,120    batch: kernel.Value,121    key: kernel.Value,122    value: kernel.Value,123) !kernel.Value {124    const batch_stride = try inner.constantIndex(indexUpper(checkedProduct(spec.keys, spec.value)));125    const key_stride = try inner.constantIndex(indexUpper(spec.value));126    const batch_offset = try inner.mul(batch, batch_stride);127    const key_offset = try inner.mul(key, key_stride);128    const batch_key_offset = try inner.add(batch_offset, key_offset);129    return inner.add(batch_key_offset, value);130}131132fn outputIndex(133    inner: anytype,134    comptime spec: ScaledDotProduct,135    batch: kernel.Value,136    query: kernel.Value,137    value: kernel.Value,138) !kernel.Value {139    const batch_stride = try inner.constantIndex(indexUpper(checkedProduct(spec.queries, spec.value)));140    const query_stride = try inner.constantIndex(indexUpper(spec.value));141    const batch_offset = try inner.mul(batch, batch_stride);142    const query_offset = try inner.mul(query, query_stride);143    const batch_query_offset = try inner.add(batch_offset, query_offset);144    return inner.add(batch_query_offset, value);145}146147fn scaled_score_dot(fold_inner: anytype, feature: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {148    const query_index = try queryIndex(fold_inner, ctx.spec, ctx.batch, ctx.query, feature);149    const key_index = try keyIndex(fold_inner, ctx.spec, ctx.batch, ctx.key, feature);150    const query_value = try ctx.query_buffer.load(fold_inner, query_index);151    const key_value = try ctx.key_buffer.load(fold_inner, key_index);152    const product = try query_value.mul(fold_inner, key_value);153    return fold_inner.add(acc, product.raw());154}155156fn scaledScore(157    inner: anytype,158    comptime spec: ScaledDotProduct,159    query_buffer: anytype,160    key_buffer: anytype,161    batch: kernel.Value,162    query: kernel.Value,163    key: kernel.Value,164) !kernel.Value {165    const zero = try inner.constantFloat(.f32, 0.0);166    const dot = try inner.foldRange(0, indexUpper(spec.qk), 1, zero, .{167        .spec = spec,168        .query_buffer = query_buffer,169        .key_buffer = key_buffer,170        .batch = batch,171        .query = query,172        .key = key,173    }, scaled_score_dot);174    const scale = try inner.constantFloat(.f32, 1.0 / @sqrt(floatExtent(spec.qk)));175    return inner.mul(dot, scale);176}177178fn score_max_reduce(fold_inner: anytype, key: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {179    const score = try scaledScore(fold_inner, ctx.spec, ctx.query_buffer, ctx.key_buffer, ctx.batch, ctx.query, key);180    return fold_inner.max(acc, score);181}182183fn scoreMax(184    inner: anytype,185    comptime spec: ScaledDotProduct,186    query_buffer: anytype,187    key_buffer: anytype,188    batch: kernel.Value,189    query: kernel.Value,190) !kernel.Value {191    const first_key = try inner.constantIndex(0);192    const first = try scaledScore(inner, spec, query_buffer, key_buffer, batch, query, first_key);193    return inner.foldRange(1, indexUpper(spec.keys), 1, first, .{194        .spec = spec,195        .query_buffer = query_buffer,196        .key_buffer = key_buffer,197        .batch = batch,198        .query = query,199    }, score_max_reduce);200}201202fn score_denominator_reduce(fold_inner: anytype, key: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {203    const score = try scaledScore(fold_inner, ctx.spec, ctx.query_buffer, ctx.key_buffer, ctx.batch, ctx.query, key);204    const shifted = try fold_inner.sub(score, ctx.row_max);205    const exp_score = try fold_inner.exp(shifted);206    return fold_inner.add(acc, exp_score);207}208209fn scoreDenominator(210    inner: anytype,211    comptime spec: ScaledDotProduct,212    query_buffer: anytype,213    key_buffer: anytype,214    batch: kernel.Value,215    query: kernel.Value,216    row_max: kernel.Value,217) !kernel.Value {218    const zero = try inner.constantFloat(.f32, 0.0);219    return inner.foldRange(0, indexUpper(spec.keys), 1, zero, .{220        .spec = spec,221        .query_buffer = query_buffer,222        .key_buffer = key_buffer,223        .batch = batch,224        .query = query,225        .row_max = row_max,226    }, score_denominator_reduce);227}228229fn weighted_value_sum_reduce(fold_inner: anytype, key: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {230    const score = try scaledScore(fold_inner, ctx.spec, ctx.query_buffer, ctx.key_buffer, ctx.batch, ctx.query, key);231    const shifted = try fold_inner.sub(score, ctx.row_max);232    const exp_score = try fold_inner.exp(shifted);233    const weight = try fold_inner.div(exp_score, ctx.denominator);234    const value_index = try valueIndex(fold_inner, ctx.spec, ctx.batch, key, ctx.value);235    const value_item = try ctx.value_buffer.load(fold_inner, value_index);236    const weighted = try fold_inner.mul(weight, value_item.raw());237    return fold_inner.add(acc, weighted);238}239240fn weightedValueSum(241    inner: anytype,242    comptime spec: ScaledDotProduct,243    query_buffer: anytype,244    key_buffer: anytype,245    value_buffer: anytype,246    batch: kernel.Value,247    query: kernel.Value,248    value: kernel.Value,249) !kernel.Value {250    const row_max = try scoreMax(inner, spec, query_buffer, key_buffer, batch, query);251    const denominator = try scoreDenominator(inner, spec, query_buffer, key_buffer, batch, query, row_max);252    const zero = try inner.constantFloat(.f32, 0.0);253    return inner.foldRange(0, indexUpper(spec.keys), 1, zero, .{254        .spec = spec,255        .query_buffer = query_buffer,256        .key_buffer = key_buffer,257        .value_buffer = value_buffer,258        .batch = batch,259        .query = query,260        .value = value,261        .row_max = row_max,262        .denominator = denominator,263    }, weighted_value_sum_reduce);264}265266fn scaled_dot_product_each(inner: anytype, index: kernel.Index3D, ctx: anytype) !void {267    const output_value = try weightedValueSum(inner, ctx.spec, ctx.args.param(.query), ctx.args.param(.key), ctx.args.param(.value), index.z.index, index.y.index, index.x.index);268    const out_index = try outputIndex(inner, ctx.spec, index.z.index, index.y.index, index.x.index);269    try ctx.args.param(.dst).store(inner, output_value, out_index);270}271272fn scaledDotProductProgram(comptime spec: ScaledDotProduct) type {273    const Body = struct {274        fn run(k: anytype, args: anytype) !void {275            _ = try k.forEach3D(.{276                .x = kernel.logical.axis(spec.value_axis, spec.value),277                .y = kernel.logical.axis(spec.query_axis, spec.queries),278                .z = kernel.logical.axis(spec.batch_axis, spec.batch),279            }, .{ .spec = spec, .args = args }, scaled_dot_product_each);280        }281    };282283    return kernel.logical.Program(.{284        .name = std.fmt.comptimePrint(285            "accy_kernel_attention_sdpa{}x{}x{}x{}x{}_{}x{}x{}_f32",286            .{ spec.batch, spec.queries, spec.keys, spec.qk, spec.value, spec.threads.x, spec.threads.y, spec.threads.z },287        ),288        .parameters = .{289            .dst = kernel.dynamicBuffer(.f32),290            .query = kernel.dynamicBuffer(.f32),291            .key = kernel.dynamicBuffer(.f32),292            .value = kernel.dynamicBuffer(.f32),293        },294        .body = Body.run,295    }).withSchedule(kernel.logical.schedule.threadBlocks(.{296        .x = spec.threads.x,297        .y = spec.threads.y,298        .z = spec.threads.z,299    }));300}301302pub fn scaledDotProductF32(comptime spec: ScaledDotProduct) type {303    return entry.Entry(scaledDotProductProgram(spec), .{304        .target = std.fmt.comptimePrint(305            "accy.kernel.attention.sdpa{}x{}x{}x{}x{}_{}x{}x{}_f32",306            .{ spec.batch, spec.queries, spec.keys, spec.qk, spec.value, spec.threads.x, spec.threads.y, spec.threads.z },307        ),308        .layer = .logical,309        .category = .attention,310        .specialization = scaledDotProductSpecialization(spec),311    });312}313314pub const ScaledDotProductAttention2x2x3x2x2F32 = scaledDotProductF32(.{315    .batch = 2,316    .queries = 2,317    .keys = 3,318    .qk = 2,319    .value = 2,320    .threads = .{ .x = 2, .y = 2, .z = 1 },321});322323fn expectedOutput(comptime spec: ScaledDotProduct, query_values: []const f32, key_values: []const f32, value_values: []const f32, batch: usize, query: usize, value: usize) f32 {324    var scores: [spec.keys]f32 = undefined;325    const scale: f32 = @floatCast(1.0 / @sqrt(floatExtent(spec.qk)));326    var max_score = -std.math.inf(f32);327    for (0..spec.keys) |key| {328        var dot: f32 = 0.0;329        for (0..spec.qk) |feature| {330            const query_offset = batch * spec.queries * spec.qk + query * spec.qk + feature;331            const key_offset = batch * spec.keys * spec.qk + key * spec.qk + feature;332            dot += query_values[query_offset] * key_values[key_offset];333        }334        const score = dot * scale;335        scores[key] = score;336        if (score > max_score) max_score = score;337    }338339    var denominator: f32 = 0.0;340    for (scores) |score| {341        denominator += @exp(score - max_score);342    }343344    var sum: f32 = 0.0;345    for (0..spec.keys) |key| {346        const weight = @exp(scores[key] - max_score) / denominator;347        const value_offset = batch * spec.keys * spec.value + key * spec.value + value;348        sum += weight * value_values[value_offset];349    }350    return sum;351}352353test "attention scaled dot product entry runs on CPU and records schedule" {354    var query_values = [_]f32{355        1.0,  0.0,356        0.0,  1.0,357        1.0,  1.0,358        -1.0, 0.5,359    };360    var key_values = [_]f32{361        1.0,  0.0,362        0.0,  1.0,363        1.0,  1.0,364        0.5,  -1.0,365        1.0,  2.0,366        -1.0, 1.0,367    };368    var value_values = [_]f32{369        1.0,  0.0,370        0.0,  2.0,371        3.0,  1.0,372        -1.0, 1.0,373        2.0,  0.0,374        0.0,  4.0,375    };376    var dst = @as([8]f32, @splat(0.0));377378    try ScaledDotProductAttention2x2x3x2x2F32.runCpu(std.testing.allocator, ScaledDotProductAttention2x2x3x2x2F32.Limits.testing, &.{379        kernel.argumentBuffer(f32, dst[0..]),380        kernel.argumentBuffer(f32, query_values[0..]),381        kernel.argumentBuffer(f32, key_values[0..]),382        kernel.argumentBuffer(f32, value_values[0..]),383    });384385    for (0..2) |batch| {386        for (0..2) |query| {387            for (0..2) |value| {388                const output_offset = batch * 4 + query * 2 + value;389                const expected = expectedOutput(.{ .batch = 2, .queries = 2, .keys = 3, .qk = 2, .value = 2 }, query_values[0..], key_values[0..], value_values[0..], batch, query, value);390                try std.testing.expectApproxEqAbs(expected, dst[output_offset], 0.0001);391            }392        }393    }394395    const launch_value = try ScaledDotProductAttention2x2x3x2x2F32.launch(std.testing.allocator, ScaledDotProductAttention2x2x3x2x2F32.Limits.testing);396    try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);397    try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);398    try std.testing.expectEqual(@as(u32, 2), launch_value.grid[2]);399    try std.testing.expectEqual(@as(u32, 2), launch_value.block[0]);400    try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);401    try std.testing.expectEqual(@as(u32, 1), launch_value.block[2]);402}403404test "attention scaled dot product entry carries specialization metadata" {405    const ScaledDotProductAttention3x2x4x5x3F32 = scaledDotProductF32(.{406        .batch = 3,407        .queries = 2,408        .keys = 4,409        .qk = 5,410        .value = 3,411        .threads = .{ .x = 3, .y = 2, .z = 1 },412    });413414    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.operationIs(.{ .attention = .scaled_dot_product }));415    try std.testing.expectEqualStrings("bqh,bkh,bkv->bqv", ScaledDotProductAttention3x2x4x5x3F32.specialization.equation.?);416    try std.testing.expectEqualStrings("accy.kernel.attention.sdpa3x2x4x5x3_3x2x1_f32", ScaledDotProductAttention3x2x4x5x3F32.target);417    try std.testing.expectEqual(@as(usize, 3), ScaledDotProductAttention3x2x4x5x3F32.specialization.inputs.len);418    try std.testing.expectEqual(@as(u64, 30), ScaledDotProductAttention3x2x4x5x3F32.specialization.inputs[0].elementCount().?);419    try std.testing.expectEqual(@as(u64, 60), ScaledDotProductAttention3x2x4x5x3F32.specialization.inputs[1].elementCount().?);420    try std.testing.expectEqual(@as(u64, 36), ScaledDotProductAttention3x2x4x5x3F32.specialization.inputs[2].elementCount().?);421    try std.testing.expectEqual(@as(u64, 18), ScaledDotProductAttention3x2x4x5x3F32.specialization.outputs[0].elementCount().?);422    try std.testing.expectEqual(@as(usize, 4), ScaledDotProductAttention3x2x4x5x3F32.specialization.reductions.len);423    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.reductionMatches(0, .{ .name = "score_dot", .operator = .dot_product, .extents = &.{5} }));424    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.reductionMatches(1, .{ .name = "score_max", .operator = .maximum, .extents = &.{4}, .dependencies = &.{"score_dot"} }));425    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.reductionMatches(2, .{ .name = "score_exp_sum", .operator = .sum_exp_shifted, .extents = &.{4}, .dependencies = &.{"score_max"} }));426    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.reductionMatches(3, .{ .name = "value_weighted_sum", .operator = .weighted_sum, .extents = &.{4}, .dependencies = &.{"score_exp_sum"} }));427    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.reductionReuseScopesAreValid());428    try std.testing.expectEqual(@as(usize, 3), ScaledDotProductAttention3x2x4x5x3F32.specialization.reduction_reuse.len);429    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.reductionReuseMatches(0, .{ .reduction = "score_dot", .extents = &.{ 3, 2, 4 } }));430    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.reductionReuseMatches(1, .{ .reduction = "score_max", .extents = &.{ 3, 2 } }));431    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.reductionReuseMatches(2, .{ .reduction = "score_exp_sum", .extents = &.{ 3, 2 } }));432    try std.testing.expectEqual(@as(u64, 432), ScaledDotProductAttention3x2x4x5x3F32.specialization.estimatedElementOps().?);433    try std.testing.expectEqual(@as(u32, 1), ScaledDotProductAttention3x2x4x5x3F32.specialization.launch.?.grid[0]);434    try std.testing.expectEqual(@as(u32, 1), ScaledDotProductAttention3x2x4x5x3F32.specialization.launch.?.grid[1]);435    try std.testing.expectEqual(@as(u32, 3), ScaledDotProductAttention3x2x4x5x3F32.specialization.launch.?.grid[2]);436    try std.testing.expectEqualDeep(ScaledDotProductAttention3x2x4x5x3F32.specialization.launch.?, ScaledDotProductAttention3x2x4x5x3F32.specialization.schedule.?.launch());437438    var snapshot = try ScaledDotProductAttention3x2x4x5x3F32.scheduleSnapshot(std.testing.allocator, ScaledDotProductAttention3x2x4x5x3F32.Limits.testing);439    defer snapshot.deinit(std.testing.allocator);440    try std.testing.expect(ScaledDotProductAttention3x2x4x5x3F32.specialization.schedule.?.matchesSnapshot(&snapshot));441}442443test "attention scaled dot product entry creates registry-ready artifact" {444    const allocator = std.testing.allocator;445    var state = gpu.recording.BackendState{446        .allocator = allocator,447        .kind = .cuda,448        .format = .cuda_ptx,449    };450451    var call_artifact = try ScaledDotProductAttention2x2x3x2x2F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = ScaledDotProductAttention2x2x3x2x2F32.Limits.testing });452    defer call_artifact.deinit();453454    const artifact = call_artifact.registry().find(ScaledDotProductAttention2x2x3x2x2F32.target, ScaledDotProductAttention2x2x3x2x2F32.version, .cuda_ptx) orelse {455        return error.TestExpectedKernelCallArtifact;456    };457    try std.testing.expectEqualStrings(ScaledDotProductAttention2x2x3x2x2F32.name, artifact.entry_name);458    try std.testing.expectEqual(@as(u32, 4), artifact.argument_count);459    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);460    switch (artifact.launch) {461        .fixed => |geometry| {462            try std.testing.expectEqual(ScaledDotProductAttention2x2x3x2x2F32.specialization.launch.?.grid[0], geometry.grid[0]);463            try std.testing.expectEqual(ScaledDotProductAttention2x2x3x2x2F32.specialization.launch.?.grid[2], geometry.grid[2]);464            try std.testing.expectEqual(ScaledDotProductAttention2x2x3x2x2F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);465            try std.testing.expectEqual(ScaledDotProductAttention2x2x3x2x2F32.specialization.launch.?.threadgroup[2], geometry.threadgroup[2]);466        },467        else => return error.TestExpectedFixedLaunch,468    }469}

1pub const attention = @import("attention.zig");