issue/1118: qyblas_gptq_w4a16_gemm

Author: qinyiqun

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);
 

src/infinicore/nn/linear.cc

@@ -4,6 +4,7 @@
 #include "infinicore/ops/distributed/allreduce.hpp"
 #include "infinicore/ops/linear.hpp"
 #include "infinicore/ops/linear_w4a16_awq.hpp"
+#include "infinicore/ops/linear_w4a16_gptq_qy.hpp"
 #include "infinicore/ops/linear_w8a8i8.hpp"
 #include <optional>
 #include <spdlog/spdlog.h>
@@ -53,6 +54,19 @@ Tensor BaseLinear::compute_linear(Tensor &input) const {
         auto output = infinicore::op::linear_w4a16_awq(input_contiguous->contiguous(), qweight, scales, qzeros, bias_opt);
         return output;
     }
+    case infinicore::quantization::QuantScheme::GPTQ_W4A16_QY: {
+        Tensor input_contiguous = input->is_contiguous() ? input : input->contiguous();
+        Tensor qweight = static_cast<const Tensor &>(weight_);
+        Tensor qzeros = static_cast<const Tensor &>(weight_zeros_);
+        Tensor scales = static_cast<const Tensor &>(weight_scale_);
+        Tensor g_idx = static_cast<const Tensor &>(gidx_);
+        std::optional<Tensor> bias_opt = has_bias_ ? std::make_optional<Tensor>(static_cast<const Tensor &>(bias_)) : std::nullopt;
+        auto output = infinicore::op::linear_w4a16_gptq_qy(input_contiguous->contiguous(), qweight, qzeros, scales, 0, 4);
+        if (bias_opt.has_value()) {
+            infinicore::op::add_(output, output, bias_opt.value()->as_strided(output->shape(), {0, 0, 1}));
+        }
+        return output;
+    }
     default: {
         // Ensure input is contiguous before creating views (required for matmul)
         // This prevents hanging when input tensor has non-contiguous memory layout
@@ -140,6 +154,23 @@ Linear::Linear(size_t in_features, size_t out_features,
         }
         break;
     }
+    case infinicore::quantization::QuantScheme::GPTQ_W4A16_QY: {
+        weight_ = infinicore::nn::Parameter({in_features / 2, out_features}, infinicore::DataType::U8, device);
+        this->register_parameter("qweight", weight_);
+        weight_zeros_ = infinicore::nn::Parameter({in_features / 128, out_features}, dtype_, device);
+        this->register_parameter("qzeros", weight_zeros_);
+        weight_scale_ = infinicore::nn::Parameter({in_features / 128, out_features}, dtype_, device);
+        this->register_parameter("scales", weight_scale_);
+
+        gidx_ = infinicore::nn::Parameter({in_features}, infinicore::DataType::I32, device);
+        this->register_parameter("g_idx", gidx_);
+        if (bias) {
+            INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device));
+        } else {
+            bias_ = Parameter();
+        }
+        break;
+    }
     default: {
         // Initialize parameters using macro
         INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, dtype_, device));
@@ -247,6 +278,27 @@ ColumnParallelLinear::ColumnParallelLinear(size_t in_features, size_t out_featur
         }
         break;
     }
+    case infinicore::quantization::QuantScheme::GPTQ_W4A16_QY: {
+        auto gptq_ptr = std::static_pointer_cast<infinicore::quantization::GPTQ_QY>(this->quantization_);
+        int group_size = gptq_ptr->get_group_size();
+        int packing_num = gptq_ptr->get_packing_num();
+        weight_ = infinicore::nn::Parameter({in_features / 2, out_features}, infinicore::DataType::U8, device, 1, tp_rank_, tp_size_);
+        this->register_parameter("qweight", weight_);
+        weight_zeros_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 1, tp_rank_, tp_size_);
+        this->register_parameter("qzeros", weight_zeros_);
+        weight_scale_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 1, tp_rank_, tp_size_);
+        this->register_parameter("scales", weight_scale_);
+        gidx_ = infinicore::nn::Parameter({in_features},
+                                          infinicore::DataType::I32,
+                                          device, 0, tp_rank_, tp_size_);
+        this->register_parameter("g_idx", gidx_);
+        if (bias) {
+            INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, tp_rank_, tp_size_));
+        } else {
+            bias_ = Parameter();
+        }
+        break;
+    }
     default: {
         // Initialize parameters using macro
         INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, dtype_, device,
@@ -356,6 +408,34 @@ RowParallelLinear::RowParallelLinear(size_t in_features, size_t out_features, st
         }
         break;
     }
+    case infinicore::quantization::QuantScheme::GPTQ_W4A16_QY: {
+        // GPTQ W4A16 QY for RowParallelLinear:切分维度为 in_features（权重矩阵的第1维）
+        // - Weight: packed int4 in U8 containers (8 int4 per U8)
+        // - Group-wise quantization with group_size=128
+        // - Scale and zero points stored per group along in_features dimension
+
+        auto gptq_ptr = std::static_pointer_cast<infinicore::quantization::GPTQ_QY>(this->quantization_);
+        int group_size = gptq_ptr->get_group_size();
+        int packing_num = gptq_ptr->get_packing_num();
+
+        weight_ = infinicore::nn::Parameter({in_features / 2, out_features}, infinicore::DataType::U8, device, 0, tp_rank_, tp_size_);
+        this->register_parameter("qweight", weight_);
+        weight_zeros_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 0, tp_rank_, tp_size_);
+        this->register_parameter("qzeros", weight_zeros_);
+        weight_scale_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 0, tp_rank_, tp_size_);
+        this->register_parameter("scales", weight_scale_);
+
+        gidx_ = infinicore::nn::Parameter({in_features},
+                                          infinicore::DataType::I32,
+                                          device, 0, tp_rank_, tp_size_);
+        this->register_parameter("g_idx", gidx_);
+        if (bias && (0 == tp_rank_)) {
+            INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, 0, 1));
+        } else {
+            bias_ = Parameter();
+        }
+        break;
+    }
     default: {
         // Initialize parameters using macro
         INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, dtype_, device,