issue/1118: QY机器添加gptq_qyblas_gemm算子 (#1123)

* issue/1118: qyblas error

* issue/1118: success qy int8 test

* issue/1118: bit=4 error

* issue/1118: debug w4

* issue/1118: success w4 kernel

* issue/1118: success quant_type and bit test

* issue/1118: qyblas_gptq_w4a16_gemm

---------

Co-authored-by: qinyiqun <qinyiqun@outlook.com>

Author: xgqdut2016

include/infinicore/nn/linear.hpp

@@ -34,6 +34,7 @@ class BaseLinear : public Module {
     Tensor bias() const { return bias_; }
     Tensor weight_scale() const { return weight_scale_; }
     Tensor weight_zeros() const { return weight_zeros_; }
+    Tensor gidx() const { return gidx_; }
 
     std::shared_ptr<infinicore::quantization::BaseQuantization> get_quantization() const { return quantization_; }
 
@@ -45,6 +46,8 @@ class BaseLinear : public Module {
     INFINICORE_NN_PARAMETER(weight_scale);
     INFINICORE_NN_PARAMETER(weight_zeros);
 
+    INFINICORE_NN_PARAMETER(gidx);
+
 protected:
     // Helper method for common forward computation
     Tensor compute_linear(Tensor &input) const;

include/infinicore/ops/linear_w4a16_gptq_qy.hpp

@@ -0,0 +1,12 @@
+#pragma once
+
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+Tensor linear_w4a16_gptq_qy(Tensor in, Tensor qweight, Tensor qzeros, Tensor scales, int64_t quant_type, int64_t bit);
+
+void linear_w4a16_gptq_qy_(Tensor out, Tensor in, Tensor qweights, Tensor scales, Tensor qzeros, int64_t quant_type, int64_t bit);
+
+} // namespace infinicore::op

include/infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp

@@ -0,0 +1,13 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(GptqQyblasGemm, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, int64_t, int64_t);
+
+void scaled_mm_w4a16_gptq_qy_(Tensor out, const Tensor &in, const Tensor &qweight, const Tensor &scales, const Tensor &qzeros, int64_t quant_type, int64_t bit);
+} // namespace infinicore::op

include/infinicore/quantization.hpp

@@ -3,5 +3,6 @@
 #include "quantization/awq.hpp"
 #include "quantization/base_quantization.hpp"
 #include "quantization/compressed_tensors.hpp"
+#include "quantization/gptq_qy.hpp"
 #include "quantization/none_quantizaiton.hpp"
 #include "quantization/quantization_scheme.hpp"

include/infinicore/quantization/base_quantization.hpp

@@ -6,7 +6,7 @@ namespace infinicore::quantization {
 class BaseQuantization {
     // Base class for quantization schemes. Intended to be extended to support various quantization methods.
 public:
-    explicit BaseQuantization(const nlohmann::json &quant_config) : quant_config_(quant_config){};
+    explicit BaseQuantization(const nlohmann::json &quant_config) : quant_config_(quant_config) {};
     virtual ~BaseQuantization() = default;
 
     virtual infinicore::quantization::QuantScheme get_quant_scheme() const = 0;

include/infinicore/quantization/gptq_qy.hpp

@@ -0,0 +1,30 @@
+#pragma once
+#include "base_quantization.hpp"
+namespace infinicore::quantization {
+
+class GPTQ_QY : public BaseQuantization {
+    // This is a temporary class that currently only returns GPTQ W4A16.
+    // Future enhancements should parse quant_config to extract detailed quantization
+    // information and support multiple quantization schemes.
+public:
+    explicit GPTQ_QY(const nlohmann::json &quant_config)
+        : BaseQuantization(quant_config) {};
+
+    infinicore::quantization::QuantScheme
+    get_quant_scheme() const override {
+        return infinicore::quantization::QuantScheme::GPTQ_W4A16_QY;
+    };
+
+    int get_packing_num() const {
+        // For GPTQ, we pack 8 int4 weights into a single int32 value.
+        return 32 / this->get_or<int>("bits", 4); // Default to 8 if not specified in config
+    }
+
+    int get_group_size() const {
+        // For simplicity, we return a fixed group size here. In a more complete implementation,
+        // this could be extracted from quant_config_ to support different group sizes.
+        return this->get_or<int>("group_size", 128); // Standard GPTQ group size
+    }
+};
+
+} // namespace infinicore::quantization

include/infinicore/quantization/quantization_scheme.hpp

@@ -7,6 +7,7 @@ enum class QuantScheme {
     NONE,
     COMPRESSED_TENSOR_W8A8I8,
     AWQ_W4A16,
+    GPTQ_W4A16_QY,
 };
 
 enum class KVQuantAlgo {

include/infiniop.h

@@ -48,6 +48,7 @@
 #include "infiniop/ops/fmod.h"
 #include "infiniop/ops/gelu.h"
 #include "infiniop/ops/gemm.h"
+#include "infiniop/ops/gptq_qyblas_gemm.h"
 #include "infiniop/ops/hardswish.h"
 #include "infiniop/ops/hardtanh.h"
 #include "infiniop/ops/hinge_embedding_loss.h"

include/infiniop/ops/gemm.h

@@ -6,22 +6,22 @@
 typedef struct InfiniopDescriptor *infiniopGemmDescriptor_t;
 
 __INFINI_C __export infiniStatus_t infiniopCreateGemmDescriptor(infiniopHandle_t handle,
-                                                         infiniopGemmDescriptor_t *desc_ptr,
-                                                         infiniopTensorDescriptor_t c_desc,
-                                                         infiniopTensorDescriptor_t a_desc,
-                                                         infiniopTensorDescriptor_t b_desc);
+                                                                infiniopGemmDescriptor_t *desc_ptr,
+                                                                infiniopTensorDescriptor_t c_desc,
+                                                                infiniopTensorDescriptor_t a_desc,
+                                                                infiniopTensorDescriptor_t b_desc);
 
 __INFINI_C __export infiniStatus_t infiniopGetGemmWorkspaceSize(infiniopGemmDescriptor_t desc, size_t *size);
 
 __INFINI_C __export infiniStatus_t infiniopGemm(infiniopGemmDescriptor_t desc,
-                                         void *workspace,
-                                         size_t workspace_size,
-                                         void *c,
-                                         void const *a,
-                                         void const *b,
-                                         float alpha,
-                                         float beta,
-                                         void *stream);
+                                                void *workspace,
+                                                size_t workspace_size,
+                                                void *c,
+                                                void const *a,
+                                                void const *b,
+                                                float alpha,
+                                                float beta,
+                                                void *stream);
 
 __INFINI_C __export infiniStatus_t infiniopDestroyGemmDescriptor(infiniopGemmDescriptor_t desc);
 

include/infiniop/ops/gptq_qyblas_gemm.h

@@ -0,0 +1,37 @@
+#ifndef __INFINIOP_GPTQ_QYBLAS_GEMM_API_H__
+#define __INFINIOP_GPTQ_QYBLAS_GEMM_API_H__
+
+#include "../operator_descriptor.h"
+#include <cstdint>
+
+typedef struct InfiniopDescriptor *infiniopGptqQyblasGemmDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateGptqQyblasGemmDescriptor(
+    infiniopHandle_t handle,
+    infiniopGptqQyblasGemmDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t a_desc,
+    infiniopTensorDescriptor_t b_desc,
+    infiniopTensorDescriptor_t b_scales_desc,
+    infiniopTensorDescriptor_t b_zeros_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetGptqQyblasGemmWorkspaceSize(
+    infiniopGptqQyblasGemmDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopGptqQyblasGemm(
+    infiniopGptqQyblasGemmDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *out,
+    const void *a,
+    const void *b,
+    void *b_scale,
+    void *b_zero,
+    int64_t quant_type,
+    int64_t bit,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyGptqQyblasGemmDescriptor(
+    infiniopGptqQyblasGemmDescriptor_t desc);
+#endif