| 270 | } |
| 271 | |
| 272 | void attention_backward(Context& ctx, float* dinp, float* dpreatt, float* datt, |
| 273 | float* dout, float* inp, float* att, |
| 274 | int B, int T, int C, int NH){ |
| 275 | struct AttentionParams { |
| 276 | uint32_t B; |
| 277 | uint32_t T; |
| 278 | uint32_t C; |
| 279 | uint32_t NH; |
| 280 | }; |
| 281 | unsigned long b = static_cast<unsigned long>(B); |
| 282 | unsigned long t = static_cast<unsigned long>(T); |
| 283 | unsigned long c = static_cast<unsigned long>(C); |
| 284 | unsigned long nh = static_cast<unsigned long>(NH); |
| 285 | setLogLevel(kError); |
| 286 | Tensor dinp_t = createTensor(ctx, Shape{b * t * c * 3}, kf32, dinp); |
| 287 | Tensor dpreatt_t = createTensor(ctx, Shape{b * nh * t * t}, kf32, dpreatt); |
| 288 | Tensor datt_t = createTensor(ctx, Shape{b * nh * t * t}, kf32, datt); |
| 289 | Tensor dout_t = createTensor(ctx, Shape{b * t * c}, kf32, dout); |
| 290 | Tensor inp_t = createTensor(ctx, Shape{b * t * c * 3}, kf32, inp); |
| 291 | Tensor att_t = createTensor(ctx, Shape{b * nh * t * t}, kf32, att); |
| 292 | std::promise<void> promise; |
| 293 | std::future<void> future = promise.get_future(); |
| 294 | Kernel op = createKernel(ctx, {kShaderAttentionBackward, 256, kf32}, |
| 295 | Bindings{dinp_t, dpreatt_t, datt_t, dout_t, inp_t, att_t}, |
| 296 | /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, |
| 297 | /* params */ |
| 298 | AttentionParams{ |
| 299 | static_cast<uint32_t>(b), |
| 300 | static_cast<uint32_t>(t), |
| 301 | static_cast<uint32_t>(c), |
| 302 | static_cast<uint32_t>(nh) |
| 303 | }); |
| 304 | dispatchKernel(ctx, op, promise); |
| 305 | wait(ctx, future); |
| 306 | toCPU(ctx, dinp_t, dinp, b * t * c * 3 * sizeof(float)); |
| 307 | toCPU(ctx, dpreatt_t, dpreatt, b * nh * t * t * sizeof(float)); |
| 308 | toCPU(ctx, datt_t, datt, b * nh * t * t * sizeof(float)); |
| 309 | } |
| 310 | |
| 311 | void gelu_forward(Context& ctx, float* out, float* inp, int n) { |
| 312 | unsigned long N = static_cast<unsigned long>(n); |
no test coverage detected