Skip to content

issue/1083: NVIDIA机器添加gptq_marlin_gemm算子 #1110

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xgqdut2016 wants to merge 10 commits into
base: main
Choose a base branch
from
Open

issue/1083: NVIDIA机器添加gptq_marlin_gemm算子 #1110

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
0dff925
issue/1083: gptq_marlin_gemm
xgqdut2016 Mar 31, 2026
5d28903
issue/1083: modified format
xgqdut2016 Mar 31, 2026
d6a778d
issue/1083: modified global
xgqdut2016 Apr 1, 2026
c3e1fba
issue/1083: success gptq_marlin
xgqdut2016 May 20, 2026
71e4b4c
issue/1083: modified format
xgqdut2016 May 20, 2026
598cec9
issue/1083: delete tvm
xgqdut2016 May 20, 2026
c4b99a8
issue/1083: add README.md
xgqdut2016 May 20, 2026
d8a23e0
issue/1083: add TVM_ROOT
xgqdut2016 May 20, 2026
6d82fe4
issue/1083: modified nvidia.lua
xgqdut2016 May 20, 2026
50cc3da
issue/1083: add commit id tvm-ffi
xgqdut2016 May 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,19 @@ python scripts/install.py [XMAKE_CONFIG_FLAGS]

```

##### 试验功能 -- 编译marlin相关算子

```shell

# 需要从github上克隆tvm_ffi仓库，克隆命令参考
## tvm-ffi commit: 35c99d0ac4cb784862115d0089f60c603acec8f9
git clone https://github.com/apache/tvm-ffi.git --recursive

# 设置TVM_ROOT
export TVM_ROOT=<path-to>/tvm-ffi #用来搜索tvm相关头文件

Copy link
Copy Markdown
Collaborator

@qinyiqun qinyiqun May 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

需要指定commit id，尽量使用https的git地址

```

2. 编译安装

默认安装路径为 `$HOME/.infini`。
Expand Down
42 changes: 42 additions & 0 deletions include/infiniop/ops/gptq_marlin_gemm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
#ifndef __INFINIOP_GPTQ_MARLIN_GEMM_API_H__
#define __INFINIOP_GPTQ_MARLIN_GEMM_API_H__

#include "../operator_descriptor.h"
#include <cstdint>

typedef struct InfiniopDescriptor *infiniopGptqMarlinGemmDescriptor_t;

__INFINI_C __export infiniStatus_t infiniopCreateGptqMarlinGemmDescriptor(infiniopHandle_t handle,
infiniopGptqMarlinGemmDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t out_desc,
infiniopTensorDescriptor_t a_desc,
infiniopTensorDescriptor_t b_desc,
infiniopTensorDescriptor_t b_scales_desc,
infiniopTensorDescriptor_t global_scale_desc,
infiniopTensorDescriptor_t b_zeros_desc,
infiniopTensorDescriptor_t g_idx_desc,
infiniopTensorDescriptor_t perm_desc);

__INFINI_C __export infiniStatus_t infiniopGetGptqMarlinGemmWorkspaceSize(infiniopGptqMarlinGemmDescriptor_t desc, size_t *size);

__INFINI_C __export infiniStatus_t infiniopGptqMarlinGemm(infiniopGptqMarlinGemmDescriptor_t desc,
void *workspace,
size_t workspace_size,
void *out,
const void *a,
const void *b,
void *b_scales,
void *global_scale,
void *b_zeros,
void *g_idx,
void *perm,
int64_t b_q_type_id,
bool is_k_full,
bool use_atomic_add,
bool use_fp32_reduce,
bool is_zp_float,
void *stream);

__INFINI_C __export infiniStatus_t infiniopDestroyGptqMarlinGemmDescriptor(infiniopGptqMarlinGemmDescriptor_t desc);

#endif
66 changes: 66 additions & 0 deletions src/infiniop/ops/gptq_marlin_gemm/gptq_marlin_gemm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
#ifndef __GPTQ_MARLIN_GEMM_H__
#define __GPTQ_MARLIN_GEMM_H__

#include "../../../utils.h"
#include "../../operator.h"
#include "../../tensor.h"
#include "info.h"

#define DESCRIPTOR(NAMESPACE) \
\
namespace op::gptq_marlin_gemm::NAMESPACE { \
class Descriptor final : public InfiniopDescriptor { \
struct Opaque; \
Opaque *_opaque; \
GptqMarlinGemmInfo _info; \
size_t _workspace_size; \
\
Descriptor( \
size_t workspace_size_, \
Opaque *opaque, \
GptqMarlinGemmInfo info, \
infiniDevice_t device_type, \
int device_id) \
: InfiniopDescriptor{device_type, device_id}, \
_opaque(opaque), \
_info(info), \
_workspace_size(workspace_size_) {} \
\
public: \
~Descriptor(); \
\
size_t workspaceSize() const { return _workspace_size; } \
\
static infiniStatus_t create( \
infiniopHandle_t handle, \
Descriptor **desc_ptr, \
infiniopTensorDescriptor_t out_desc, \
infiniopTensorDescriptor_t a_desc, \
infiniopTensorDescriptor_t b_desc, \
infiniopTensorDescriptor_t b_scales_desc, \
infiniopTensorDescriptor_t global_scale_desc, \
infiniopTensorDescriptor_t b_zeros_desc, \
infiniopTensorDescriptor_t g_idx_desc, \
infiniopTensorDescriptor_t perm_desc); \
\
infiniStatus_t calculate( \
void *workspace, \
size_t workspace_size, \
void *out, \
const void *a, \
const void *b, \
void *b_scales, \
void *global_scale, \
void *b_zeros, \
void *g_idx, \
void *perm, \
int64_t b_q_type_id, \
bool is_k_full, \
bool use_atomic_add, \
bool use_fp32_reduce, \
bool is_zp_float, \
void *stream) const; \
}; \
}

#endif //__GPTQ_MARLIN_GEMM_H__
59 changes: 59 additions & 0 deletions src/infiniop/ops/gptq_marlin_gemm/info.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
#ifndef __GPTQ_MARLIN_GEMM_INFO_H__
#define __GPTQ_MARLIN_GEMM_INFO_H__

#include "../../../utils.h"
#include "../../tensor.h"
#include <vector>

#include <cassert>

namespace op::gptq_marlin_gemm {

class GptqMarlinGemmInfo {
GptqMarlinGemmInfo() = default;

public:
infiniDtype_t dtype;
size_t M, K, N, b_q_size_1;
int num_groups;
ptrdiff_t a_stride_0;

static utils::Result<GptqMarlinGemmInfo> create(
infiniopTensorDescriptor_t out_desc,
infiniopTensorDescriptor_t a_desc,
infiniopTensorDescriptor_t b_desc,
infiniopTensorDescriptor_t b_scales_desc,
infiniopTensorDescriptor_t global_scale_desc,
infiniopTensorDescriptor_t b_zeros_desc,
infiniopTensorDescriptor_t g_idx_desc,
infiniopTensorDescriptor_t perm_desc) {
CHECK_OR_RETURN(
out_desc != nullptr && a_desc != nullptr && b_desc != nullptr && b_scales_desc != nullptr,
INFINI_STATUS_NULL_POINTER);
const infiniDtype_t dtype = a_desc->dtype();
size_t M = out_desc->dim(0);
size_t N = out_desc->dim(1);
size_t K = a_desc->dim(1);
size_t b_q_size_1 = b_desc->dim(1);
int num_groups = static_cast<int>(b_scales_desc->dim(0));
ptrdiff_t a_stride_0 = a_desc->strides()[0];

auto ndim = out_desc->ndim();
CHECK_OR_RETURN(ndim == 2
&& a_desc->ndim() == ndim
&& b_desc->ndim() == ndim
&& b_scales_desc->ndim() == ndim,
INFINI_STATUS_BAD_TENSOR_SHAPE);

CHECK_OR_RETURN(b_scales_desc->shape()[1] == N
&& a_stride_0 % 8 == 0,
INFINI_STATUS_BAD_TENSOR_SHAPE);

return utils::Result<GptqMarlinGemmInfo>(
GptqMarlinGemmInfo{dtype, M, K, N, b_q_size_1, num_groups, a_stride_0});
}
};

} // namespace op::gptq_marlin_gemm

#endif // __GPTQ_MARLIN_GEMM_INFO_H__
Loading
Loading