| 192 | """ |
| 193 | |
| 194 | def fuse3(qq, qk, qv): |
| 195 | if not hasattr(qq, '_amax') or not hasattr(qk, '_amax') or not hasattr(qv, '_amax'): |
| 196 | logger.warn('missing amax buffer, unable to fuse') |
| 197 | return |
| 198 | q = qq._amax.detach().item() |
| 199 | k = qk._amax.detach().item() |
| 200 | v = qv._amax.detach().item() |
| 201 | |
| 202 | amax = max(q, k, v) |
| 203 | qq._amax.fill_(amax) |
| 204 | qk._amax.fill_(amax) |
| 205 | qv._amax.fill_(amax) |
| 206 | logger.info(f' q={q:7.4f} k={k:7.4f} v={v:7.4f} -> {amax:7.4f}') |
| 207 | |
| 208 | for name, mod in model.named_modules(): |
| 209 | if name.endswith('.attn'): |