NOLINTNEXTLINE(readability-function-size)
| 30 | |
| 31 | // NOLINTNEXTLINE(readability-function-size) |
| 32 | static migraphx::program create_gqa_program(const size_t batch_size, |
| 33 | const size_t num_heads, |
| 34 | const size_t kv_num_heads, |
| 35 | const size_t sequence_length, |
| 36 | const size_t head_size, |
| 37 | const size_t past_sequence_length, |
| 38 | const size_t max_sequence_length, |
| 39 | const bool do_rotary, |
| 40 | const float scale, |
| 41 | const bool test_rotary = false, |
| 42 | const bool test_concat = false, |
| 43 | const int local_window_size = -1) |
| 44 | { |
| 45 | migraphx::program p; |
| 46 | auto* mm = p.get_main_module(); |
| 47 | std::vector<size_t> query_lens{ |
| 48 | batch_size, sequence_length, head_size * (num_heads + 2 * kv_num_heads)}; |
| 49 | std::vector<size_t> kv_lens{batch_size, kv_num_heads, max_sequence_length, head_size}; |
| 50 | std::vector<size_t> slk_lens{batch_size, 1}; |
| 51 | std::vector<size_t> cs_cache_lens{max_sequence_length, head_size / 2}; |
| 52 | auto dtype = migraphx::shape::half_type; |
| 53 | migraphx::shape query_s{dtype, query_lens}; |
| 54 | migraphx::shape kv_s{dtype, kv_lens}; |
| 55 | migraphx::shape slk_s{migraphx::shape::int32_type, slk_lens}; |
| 56 | migraphx::shape cs_cache_s{dtype, cs_cache_lens}; |
| 57 | auto query = mm->add_parameter("query", query_s); |
| 58 | std::vector<int> slk_vec(slk_s.elements(), past_sequence_length); |
| 59 | std::vector<float> cs_min_vec(cs_cache_s.elements(), -1.0); |
| 60 | std::vector<float> cs_max_vec(cs_cache_s.elements(), 1.0); |
| 61 | auto k = mm->add_parameter("k", kv_s); |
| 62 | auto v = mm->add_parameter("v", kv_s); |
| 63 | auto slk = mm->add_parameter("slk", slk_s); |
| 64 | auto slk_lit = mm->add_literal(slk_s, slk_vec); |
| 65 | slk = mm->add_instruction(migraphx::make_op("clip"), slk, slk_lit, slk_lit); |
| 66 | auto cs_min = mm->add_literal(cs_cache_s, cs_min_vec); |
| 67 | auto cs_max = mm->add_literal(cs_cache_s, cs_max_vec); |
| 68 | auto cos_cache = mm->add_parameter("cos_cache", cs_cache_s); |
| 69 | auto sin_cache = mm->add_parameter("sin_cache", cs_cache_s); |
| 70 | cos_cache = mm->add_instruction(migraphx::make_op("clip"), cos_cache, cs_min, cs_max); |
| 71 | sin_cache = mm->add_instruction(migraphx::make_op("clip"), sin_cache, cs_min, cs_max); |
| 72 | |
| 73 | std::vector<std::size_t> bsnh{ |
| 74 | batch_size, sequence_length, num_heads + 2 * kv_num_heads, head_size}; |
| 75 | |
| 76 | auto transposed_qkv = |
| 77 | mm->add_instruction(migraphx::make_op("reshape", {{"dims", bsnh}}), query); |
| 78 | |
| 79 | transposed_qkv = mm->add_instruction( |
| 80 | migraphx::make_op("transpose", {{"permutation", {0, 2, 1, 3}}}), transposed_qkv); |
| 81 | |
| 82 | auto rotary_qkv = transposed_qkv; |
| 83 | if(do_rotary) |
| 84 | { |
| 85 | std::vector<migraphx::instruction_ref> rotary_inputs{ |
| 86 | transposed_qkv, slk, cos_cache, sin_cache}; |
| 87 | rotary_qkv = mm->add_instruction( |
| 88 | migraphx::make_op( |
| 89 | "gqa_rotary_embedding", |
no test coverage detected