| import torch | |
| def install(model): | |
| """Install the selected BF16 fusion stack before prefill with BF16 weights, KV and recurrent storage.""" | |
| assert model.config.gdn_recurrent_state_storage == 'bf16_non_anchor' | |
| assert all(p.is_cuda and p.dtype == torch.bfloat16 for p in model.parameters()) | |
| from . import b1_bf16_fusions | |
| b1_bf16_fusions.install(model) | |
| from . import b1_runtime_joint_projection | |
| b1_runtime_joint_projection.install(model) | |
| from . import b1_bf16_crosslayer | |
| b1_bf16_crosslayer.install(model) | |
| from . import b1_bf16_swa_prepare | |
| b1_bf16_swa_prepare.install(model) | |
| from . import b1_bf16_output_fusion | |
| b1_bf16_output_fusion.install(model) | |
| return b1_bf16_output_fusion.generation_function() | |