(size_per_head, is_fp16)
| 20 | USE_CACHE_BATCH_MAJOR_ATTENTION = True |
| 21 | |
| 22 | def get_op_cache_config(size_per_head, is_fp16): |
| 23 | x = 8 if is_fp16 else 4 |
| 24 | use_batch_major_op_cache = True if USE_CACHE_BATCH_MAJOR_ATTENTION == True and \ |
| 25 | size_per_head % x == 0 \ |
| 26 | else False |
| 27 | x = x if use_batch_major_op_cache else 1 |
| 28 | return use_batch_major_op_cache, x |
| 29 | |
| 30 | def init_op_cache(layer_num, batch_size, beam_width, max_seq_len, \ |
| 31 | decoding_max_seq_len, head_num, size_per_head, hidden_dim, is_fp16): |