Add the attention layer
(prefix, config, init_dict, network, input_tensor, mask_idx, cu_seqlens, max_seqlen)
| 96 | layer.get_output(out_idx).set_dynamic_range(-maxval, maxval) |
| 97 | |
| 98 | def attention_layer_opt(prefix, config, init_dict, network, input_tensor, mask_idx, cu_seqlens, max_seqlen): |
| 99 | """ |
| 100 | Add the attention layer |
| 101 | """ |
| 102 | hidden_size = config.hidden_size |
| 103 | num_heads = config.num_attention_heads |
| 104 | head_size = int(hidden_size / num_heads) |
| 105 | |
| 106 | Wall = init_dict[prefix + WQKV] |
| 107 | Ball = init_dict[prefix + BQKV] |
| 108 | |
| 109 | # FC_attention |
| 110 | if config.use_int8: |
| 111 | mult_all = network.add_convolution_nd(input_tensor, 3 * hidden_size, (1, 1), Wall, Ball) |
| 112 | else: |
| 113 | mult_all = network.add_fully_connected(input_tensor, 3 * hidden_size, Wall, Ball) |
| 114 | |
| 115 | if config.use_qat: |
| 116 | dr_qkv = max( |
| 117 | init_dict[prefix + 'self_qv_a_input_quantizer_amax'], |
| 118 | init_dict[prefix + 'self_qv_b_input_quantizer_amax'], |
| 119 | init_dict[prefix + 'self_av_b_input_quantizer_amax'], |
| 120 | ) |
| 121 | set_output_range(mult_all, dr_qkv) |
| 122 | set_output_name(mult_all, prefix, "qkv_mult") |
| 123 | |
| 124 | # QKV2CTX |
| 125 | dtype = config.get_trt_dtype() |
| 126 | |
| 127 | pf_type = trt.PluginField("type_id", np.array([int(dtype)], np.int32), trt.PluginFieldType.INT32) |
| 128 | pf_hidden_size = trt.PluginField("hidden_size", np.array([hidden_size], np.int32), trt.PluginFieldType.INT32) |
| 129 | pf_num_heads = trt.PluginField("num_heads", np.array([num_heads], np.int32), trt.PluginFieldType.INT32) |
| 130 | pf_has_mask = trt.PluginField("has_mask", np.array([1], np.int32), trt.PluginFieldType.INT32) |
| 131 | pf_var_seqlen = trt.PluginField("var_seqlen", np.array([int(1)], np.int32), trt.PluginFieldType.FLOAT32) |
| 132 | |
| 133 | if config.use_qat: |
| 134 | dr_probs = init_dict[prefix + 'self_av_a_input_quantizer_amax'] |
| 135 | dq_probs = dr_probs / 127.0 |
| 136 | pf_dq_probs = trt.PluginField("dq_probs", np.array([dq_probs], np.float32), trt.PluginFieldType.FLOAT32) |
| 137 | fields = [pf_hidden_size, pf_num_heads, pf_dq_probs] |
| 138 | else: |
| 139 | fields = [pf_hidden_size, pf_num_heads] |
| 140 | |
| 141 | if config.use_int8 and config.interleaved: |
| 142 | pfc = trt.PluginFieldCollection(fields) |
| 143 | qkv2ctx_plug = mha_plg_creator3.create_plugin("qkv2ctx", pfc) |
| 144 | qkv_in = [mult_all.get_output(0), cu_seqlens, max_seqlen] |
| 145 | else: |
| 146 | fields.append(pf_has_mask) |
| 147 | fields.append(pf_type) |
| 148 | fields.append(pf_var_seqlen) |
| 149 | pfc = trt.PluginFieldCollection(fields) |
| 150 | qkv2ctx_plug = mha_plg_creator2.create_plugin("qkv2ctx", pfc) |
| 151 | qkv_in = [mult_all.get_output(0), mask_idx, cu_seqlens, max_seqlen] |
| 152 | qkv2ctx = network.add_plugin_v2(qkv_in, qkv2ctx_plug) |
| 153 | qkv2ctx.name = prefix + 'qkv_to_ctx' |
| 154 | |
| 155 | if config.use_qat: |
no test coverage detected