| 233 | } |
| 234 | |
| 235 | void attention_forward(Context& ctx, float* out, float* preatt, float* att, |
| 236 | float* inp, |
| 237 | int B, int T, int C, int NH){ |
| 238 | struct AttentionParams { |
| 239 | uint32_t B; |
| 240 | uint32_t T; |
| 241 | uint32_t C; |
| 242 | uint32_t NH; |
| 243 | }; |
| 244 | unsigned long b = static_cast<unsigned long>(B); |
| 245 | unsigned long t = static_cast<unsigned long>(T); |
| 246 | unsigned long c = static_cast<unsigned long>(C); |
| 247 | unsigned long nh = static_cast<unsigned long>(NH); |
| 248 | setLogLevel(kError); |
| 249 | Tensor inp_t = createTensor(ctx, Shape{b * t * c * 3}, kf32, inp); |
| 250 | Tensor preatt_t = createTensor(ctx, Shape{b * nh * t * t}, kf32, preatt); |
| 251 | Tensor att_t = createTensor(ctx, Shape{b * nh * t * t}, kf32, att); |
| 252 | Tensor out_t = createTensor(ctx, Shape{b * t * c}, kf32); |
| 253 | std::promise<void> promise; |
| 254 | std::future<void> future = promise.get_future(); |
| 255 | Kernel op = createKernel(ctx, {kShaderAttention, 256, kf32}, |
| 256 | Bindings{inp_t, preatt_t, att_t, out_t}, |
| 257 | /* nWorkgroups */ {cdiv(b * t, 256), 1, 1}, |
| 258 | /* params */ |
| 259 | AttentionParams{ |
| 260 | static_cast<uint32_t>(b), |
| 261 | static_cast<uint32_t>(t), |
| 262 | static_cast<uint32_t>(c), |
| 263 | static_cast<uint32_t>(nh) |
| 264 | }); |
| 265 | dispatchKernel(ctx, op, promise); |
| 266 | wait(ctx, future); |
| 267 | toCPU(ctx, preatt_t, preatt, b * nh * t * t * sizeof(float)); |
| 268 | toCPU(ctx, att_t, att, b * nh * t * t * sizeof(float)); |
| 269 | toCPU(ctx, out_t, out, b * t * c * sizeof(float)); |
| 270 | } |
| 271 | |
| 272 | void attention_backward(Context& ctx, float* dinp, float* dpreatt, float* datt, |
| 273 | float* dout, float* inp, float* att, |
no test coverage detected