Skip to main content

ruda_tensor_device/dispatch/
frozen_awq.rs

1use crate::{DeviceBackend, DeviceRuntime, FloatElement, IntElement, element::BoolElement};
2use ruda_tensor::{frozen_awq::FrozenAwqOps, tensor::{FloatTensor, IntTensor}};
3use rublas::tensor_int4::{AwqGemm, Int4Error};
4
5impl<R, F, I, BT> FrozenAwqOps for DeviceBackend<R, F, I, BT>
6where R: DeviceRuntime, F: FloatElement, I: IntElement, BT: BoolElement {
7    type AwqError = Int4Error;
8
9    fn frozen_awq_forward(
10        input: FloatTensor<Self>, qweight: IntTensor<Self>, qzeros: IntTensor<Self>,
11        scales: FloatTensor<Self>, bias: Option<FloatTensor<Self>>, group_size: usize,
12    ) -> Result<FloatTensor<Self>, Self::AwqError> {
13        AwqGemm::new(qweight, qzeros, scales, bias, group_size)?.forward_with_input_dtype(input)
14    }
15
16    fn frozen_awq_input_backward(
17        gradient: FloatTensor<Self>, qweight: IntTensor<Self>, qzeros: IntTensor<Self>,
18        scales: FloatTensor<Self>, group_size: usize,
19    ) -> Result<FloatTensor<Self>, Self::AwqError> {
20        AwqGemm::new(qweight, qzeros, scales, None, group_size)?.input_backward(gradient)
21    }
22}