(self, x)
| 42 | return self.down_proj(self.our_llama_mlp(x)) |
| 43 | |
| 44 | def our_llama_mlp(self, x): |
| 45 | out_shape = x.shape[:-1] + (self.intermediate_size,) |
| 46 | x = x.reshape(-1, x.shape[-1]) |
| 47 | |
| 48 | if x.shape[0] <= 8: |
| 49 | gate_output = awq_inference_engine.gemv_forward_cuda( |
| 50 | x, |
| 51 | self.gate_proj_qweight, |
| 52 | self.gate_proj_scales, |
| 53 | self.gate_proj_qzeros, |
| 54 | self.w_bit, |
| 55 | self.down_proj.group_size, |
| 56 | ) |
| 57 | gate_output = F.silu(gate_output) |
| 58 | up_output = awq_inference_engine.gemv_forward_cuda( |
| 59 | x, |
| 60 | self.up_proj_qweight, |
| 61 | self.up_proj_scales, |
| 62 | self.up_proj_qzeros, |
| 63 | self.w_bit, |
| 64 | self.down_proj.group_size, |
| 65 | ) |
| 66 | else: |
| 67 | if USE_TRITON: |
| 68 | gate_output = quant_matmul_v2(x, |
| 69 | self.gate_proj_qweight.T.contiguous(), |
| 70 | self.gate_proj_qzeros.T.contiguous(), |
| 71 | self.gate_proj_scales.T.contiguous(), |
| 72 | M=out_shape[-2], |
| 73 | N=out_shape[-1], |
| 74 | K=self.in_features, |
| 75 | pack_num=self.split_k_iters, |
| 76 | group_size=self.down_proj.group_size, |
| 77 | w_bit=self.w_bit, |
| 78 | offset=self.offset |
| 79 | ) |
| 80 | gate_output = F.silu(gate_output) |
| 81 | up_output = quant_matmul_v2(x, |
| 82 | self.up_proj_qweight.T.contiguous(), |
| 83 | self.up_proj_qzeros.T.contiguous(), |
| 84 | self.up_proj_scales.T.contiguous(), |
| 85 | M=out_shape[-2], |
| 86 | N=out_shape[-1], |
| 87 | K=self.in_features, |
| 88 | pack_num=self.split_k_iters, |
| 89 | group_size=self.down_proj.group_size, |
| 90 | w_bit=self.w_bit, |
| 91 | offset=self.offset |
| 92 | ) |
| 93 | |
| 94 | # gate_output = awq_inference_engine.gemm_forward_cuda( |
| 95 | # x, |
| 96 | # self.gate_proj_qweight, |
| 97 | # self.gate_proj_scales, |
| 98 | # self.gate_proj_qzeros, |
| 99 | # self.down_proj.group_size, |
| 100 | # self.split_k_iters, |
| 101 | # ) |
no test coverage detected