| 16 | USE_TRITON = False |
| 17 | |
| 18 | class QuantLlamaMLP(nn.Module): |
| 19 | def __init__( |
| 20 | self, |
| 21 | gate_proj, |
| 22 | down_proj, |
| 23 | up_proj, |
| 24 | ): |
| 25 | super().__init__() |
| 26 | self.register_buffer("gate_proj_qweight", gate_proj.qweight) |
| 27 | self.register_buffer("gate_proj_scales", gate_proj.scales) |
| 28 | self.register_buffer("gate_proj_qzeros", gate_proj.qzeros) |
| 29 | self.register_buffer("up_proj_qweight", up_proj.qweight) |
| 30 | self.register_buffer("up_proj_scales", up_proj.scales) |
| 31 | self.register_buffer("up_proj_qzeros", up_proj.qzeros) |
| 32 | |
| 33 | self.in_features = gate_proj.in_features |
| 34 | self.intermediate_size = gate_proj.out_features |
| 35 | self.out_features = down_proj.out_features |
| 36 | self.w_bit = gate_proj.w_bit |
| 37 | self.down_proj = down_proj |
| 38 | self.split_k_iters = down_proj.split_k_iters |
| 39 | self.offset = 0x0F if self.w_bit == 4 else 0x03 |
| 40 | |
| 41 | def forward(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, |