lib/accy/src/kernel/library/attention.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3
4 const entry = @import("entry.zig");
5 const kernel = @import("../root.zig");
6
7 pub 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 };
20
21 fn 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 }
27
28 fn 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 }
31
32 fn floatExtent(comptime extent: u64) f64 {
33 return @floatFromInt(extent);
34 }
35
36 fn 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 }
39
40 fn 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 }
43
44 fn 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 }
47
48 fn 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 }
51
52 fn queryRowShape(comptime spec: ScaledDotProduct) entry.Shape {
53 return entry.shape2D(spec.batch_axis, spec.batch, spec.query_axis, spec.queries);
54 }
55
56 fn 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 }
59
60 fn 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 }
86
87 fn 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 }
101
102 fn 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 }
116
117 fn 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 }
131
132 fn 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 }
146
147 fn 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 }
155
156 fn 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 }
177
178 fn 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 }
182
183 fn 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 }
201
202 fn 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 }
208
209 fn 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 }
228
229 fn 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 }
239
240 fn 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 }
265
266 fn 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 }
271
272 fn 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 };
282
283 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 }
301
302 pub 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 }
313
314 pub 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 });
322
323 fn 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 }
338
339 var denominator: f32 = 0.0;
340 for (scores) |score| {
341 denominator += @exp(score - max_score);
342 }
343
344 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 }
352
353 test "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));
377
378 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 });
384
385 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 }
394
395 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 }
403
404 test "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 });
413
414 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());
437
438 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 }
442
443 test "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 };
450
451 var call_artifact = try ScaledDotProductAttention2x2x3x2x2F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = ScaledDotProductAttention2x2x3x2x2F32.Limits.testing });
452 defer call_artifact.deinit();
453
454 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 }