issue/1395 - six more operators for metax

src/infiniop/ops/dequantize_gptq/metax/dequantize_w42f16_metax.h

@@ -0,0 +1,8 @@
+#ifndef __DEQUANTIZE_GPTQ_METAX_H__
+#define __DEQUANTIZE_GPTQ_METAX_H__
+
+#include "../dequantize_gptq.h"
+
+DESCRIPTOR(metax)
+
+#endif // __DEQUANTIZE_GPTQ_METAX_H__

src/infiniop/ops/dequantize_gptq/metax/dequantize_w42f16_metax.maca

@@ -0,0 +1,136 @@
+#ifdef ENABLE_METAX_MC_API
+#include <mc_runtime.h>
+#else
+#include <hc_runtime.h>
+#endif
+
+#if defined(ENABLE_METAX_API) || defined(ENABLE_QY_API)
+
+#include "../../../devices/metax/metax_common.h"
+#include "../../../devices/metax/metax_handle.h"
+#include "../../../devices/metax/metax_kernel_common.h"
+#include "../dequantize_gptq.h"
+#include "dequantize_w42f16_metax.h"
+
+namespace op::dequantize_gptq::metax {
+
+struct Descriptor::Opaque {
+    std::shared_ptr<device::metax::Handle::Internal> internal;
+};
+
+Descriptor::~Descriptor() { delete _opaque; }
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle_,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t qweight_desc,
+    infiniopTensorDescriptor_t scales_desc,
+    infiniopTensorDescriptor_t zeros_desc,
+    infiniopTensorDescriptor_t g_idx_desc) {
+
+    auto handle = reinterpret_cast<device::metax::Handle *>(handle_);
+    auto result = DequantizeGPTQInfo::create(out_desc, qweight_desc, scales_desc, zeros_desc, g_idx_desc);
+
+    *desc_ptr = new Descriptor(
+        0,
+        new Opaque{handle->internal()},
+        result.take(),
+        handle->device, handle->device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+// 仅保留这一版内核（支持 g_idx）
+// qweight: [in_packed, out_features] packing 8 input channels per word
+// zeros: [num_groups, out_packed] packing 8 output channels per word
+// scales: [num_groups, out_features], g_idx: [in_features]
+__global__ void __launch_bounds__(128)
+    dequantize_weights_gptq(const uint32_t *__restrict__ qweight,
+                            const half *__restrict__ scales,
+                            const uint32_t *__restrict__ zeros,
+                            const int *__restrict__ g_idx,
+                            half *__restrict__ out,
+                            int in_features,
+                            int out_features,
+                            int out_packed, // ceil(out_features / 8)
+                            int num_groups) {
+    // Each thread handles one packed output column (8 real output cols).
+    const int col_pack = blockIdx.x * blockDim.x + threadIdx.x; // packed output column
+    const int row = blockIdx.y * blockDim.y + threadIdx.y;      // real input row
+    if (col_pack >= out_packed || row >= in_features) {
+        return;
+    }
+
+    // Clamp gid to valid range
+    const int gid_raw = g_idx ? g_idx[row] : 0;
+    const int gid = ((gid_raw % num_groups) + num_groups) % num_groups;
+
+    const int pack_row = row >> 3; // packed input row
+
+    const int zero_idx = gid * out_packed + col_pack; // zeros layout: [num_groups, out_packed]
+    const uint32_t zeros_loaded = zeros[zero_idx];
+
+    const int q_shift = (row & 7) * 4;  // qweight packs 8 input rows
+    const int col_base = col_pack << 3; // 8 real cols per pack
+    const int scale_base = gid * out_features + col_base;
+
+#pragma unroll
+    for (int j = 0; j < 8; ++j) {
+        const int col = col_base + j;
+        if (col >= out_features) {
+            break;
+        }
+
+        const uint32_t q_loaded = qweight[pack_row * out_features + col];
+        const int q_nib = (q_loaded >> q_shift) & 0xF;
+
+        const int z_nib = (zeros_loaded >> (j * 4)) & 0xF;
+        const half scale = scales[scale_base + j];
+
+        // GPTQ quirk: The stored zero point is usually offset by 1.
+        // Standard formula: (q - (z + 1)) * s
+        const float v = float(q_nib - (z_nib + 1)) * __half2float(scale);
+        out[row * out_features + col] = __float2half(v);
+    }
+}
+
+infiniStatus_t
+Descriptor::calculate(
+    void *workspace, size_t workspace_size,
+    void *out,
+    const void *qweight,
+    const void *scales,
+    const void *zeros,
+    const void *g_idx,
+    void *stream) const {
+
+    const int in_features = _info.in_features();
+    const int out_features = _info.out_features();
+    const int out_packed = _info.out_packed();
+    const int in_packed = _info.in_packed();
+    const int num_groups = _info.num_groups();
+
+    if (num_groups <= 0 || in_features <= 0 || out_features <= 0 || out_packed <= 0 || in_packed <= 0) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+    constexpr int BLOCK_X = 16; // packed columns
+    constexpr int BLOCK_Y = 4;  // rows
+    dim3 threads(BLOCK_X, BLOCK_Y);
+    dim3 blocks((out_packed + BLOCK_X - 1) / BLOCK_X,
+                (in_features + BLOCK_Y - 1) / BLOCK_Y);
+
+    dequantize_weights_gptq<<<blocks, threads, 0,
+                              reinterpret_cast<hcStream_t>(stream)>>>(
+        reinterpret_cast<const uint32_t *>(qweight),
+        reinterpret_cast<const half *>(scales),
+        reinterpret_cast<const uint32_t *>(zeros),
+        reinterpret_cast<const int *>(g_idx),
+        reinterpret_cast<half *>(out),
+        in_features, out_features, out_packed, num_groups);
+    return INFINI_STATUS_SUCCESS;
+}
+
+} // namespace op::dequantize_gptq::metax
+
+#endif

src/infiniop/ops/dequantize_gptq/operator.cc

@@ -5,6 +5,10 @@
 #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API)
 #include "nvidia/dequantize_w42f16_nvidia.cuh"
 #endif
+
+#ifdef ENABLE_METAX_API
+#include "metax/dequantize_w42f16_metax.h"
+#endif
 #ifdef ENABLE_MOORE_API
 #include "moore/dequantize_w42f16_moore.h"
 #endif
@@ -34,6 +38,9 @@ __INFINI_C infiniStatus_t infiniopCreateDequantizeGPTQDescriptor(
 #ifdef ENABLE_NVIDIA_API
         CREATE(INFINI_DEVICE_NVIDIA, nvidia);
 #endif
+#ifdef ENABLE_METAX_API
+        CREATE(INFINI_DEVICE_METAX, metax);
+#endif
 #ifdef ENABLE_MOORE_API
         CREATE(INFINI_DEVICE_MOORE, moore);
 #endif
@@ -61,6 +68,9 @@ __INFINI_C infiniStatus_t infiniopGetDequantizeGPTQWorkspaceSize(infiniopDequant
 #ifdef ENABLE_NVIDIA_API
         GET(INFINI_DEVICE_NVIDIA, nvidia);
 #endif
+#ifdef ENABLE_METAX_API
+        GET(INFINI_DEVICE_METAX, metax);
+#endif
 #ifdef ENABLE_MOORE_API
         GET(INFINI_DEVICE_MOORE, moore);
 #endif
@@ -96,6 +106,9 @@ __INFINI_C infiniStatus_t infiniopDequantizeGPTQ(
 #ifdef ENABLE_NVIDIA_API
         CALCULATE(INFINI_DEVICE_NVIDIA, nvidia);
 #endif
+#ifdef ENABLE_METAX_API
+        CALCULATE(INFINI_DEVICE_METAX, metax);
+#endif
 #ifdef ENABLE_MOORE_API
         CALCULATE(INFINI_DEVICE_MOORE, moore);
 #endif
@@ -124,6 +137,9 @@ infiniopDestroyDequantizeGPTQDescriptor(infiniopDequantizeGPTQDescriptor_t desc)
 #ifdef ENABLE_NVIDIA_API
         DELETE(INFINI_DEVICE_NVIDIA, nvidia);
 #endif
+#ifdef ENABLE_METAX_API
+        DELETE(INFINI_DEVICE_METAX, metax);
+#endif
 #ifdef ENABLE_MOORE_API
         DELETE(INFINI_DEVICE_MOORE, moore);
 #endif

src/infiniop/ops/mamba_selective_scan/metax/mamba_selective_scan_metax.h

@@ -0,0 +1,5 @@
+#ifndef __MAMBA_SELECTIVE_SCAN_METAX_H__
+#define __MAMBA_SELECTIVE_SCAN_METAX_H__
+#include "../mamba_selective_scan.h"
+DESCRIPTOR(metax)
+#endif

src/infiniop/ops/mamba_selective_scan/metax/mamba_selective_scan_metax.maca

@@ -0,0 +1,118 @@
+#include "../../../devices/metax/metax_common.h"
+#include "../../../devices/metax/metax_kernel_common.h"
+#include "mamba_selective_scan_metax.h"
+#ifdef ENABLE_METAX_MC_API
+#include <mc_runtime.h>
+#else
+#include <hc_runtime.h>
+#endif
+
+namespace op::mamba_selective_scan::metax {
+
+struct Descriptor::Opaque {
+    std::shared_ptr<device::metax::Handle::Internal> internal;
+};
+Descriptor::~Descriptor() { delete _opaque; }
+
+namespace {
+template <typename T>
+__device__ float to_float(T v) { return static_cast<float>(v); }
+template <>
+__device__ float to_float<half>(half v) { return __half2float(v); }
+template <>
+__device__ float to_float<__nv_bfloat16>(__nv_bfloat16 v) { return __bfloat162float(v); }
+template <typename T>
+__device__ T from_float(float v) { return static_cast<T>(v); }
+template <>
+__device__ half from_float<half>(float v) { return __float2half_rn(v); }
+template <>
+__device__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { return __float2bfloat16_rn(v); }
+__device__ float softplusf_stable(float x) { return x > 20.0f ? x : log1pf(expf(x)); }
+__device__ float siluf(float x) { return x / (1.0f + expf(-x)); }
+
+template <typename T>
+__global__ void mamba_selective_scan_kernel(
+    T *out, const T *x, const T *dt, const T *b, const T *c, const T *a_log,
+    const T *d, const T *gate, const T *dt_bias, float *state,
+    size_t batch, size_t seq_len, size_t intermediate, size_t state_size) {
+    size_t linear = blockIdx.x * blockDim.x + threadIdx.x;
+    size_t total = batch * intermediate;
+    if (linear >= total) {
+        return;
+    }
+    size_t ch = linear % intermediate;
+    size_t batch_idx = linear / intermediate;
+    size_t state_base = (batch_idx * intermediate + ch) * state_size;
+    for (size_t t = 0; t < seq_len; ++t) {
+        size_t x_idx = (batch_idx * seq_len + t) * intermediate + ch;
+        float xt = to_float<T>(x[x_idx]);
+        float dtv = softplusf_stable(to_float<T>(dt[x_idx]) + to_float<T>(dt_bias[ch]));
+        float y = 0.0f;
+        for (size_t n = 0; n < state_size; ++n) {
+            float a = -expf(to_float<T>(a_log[ch * state_size + n]));
+            float discrete_a = expf(a * dtv);
+            float bn = to_float<T>(b[(batch_idx * seq_len + t) * state_size + n]);
+            float cn = to_float<T>(c[(batch_idx * seq_len + t) * state_size + n]);
+            float s = discrete_a * state[state_base + n] + dtv * bn * xt;
+            state[state_base + n] = s;
+            y += s * cn;
+        }
+        y += xt * to_float<T>(d[ch]);
+        y *= siluf(to_float<T>(gate[x_idx]));
+        out[x_idx] = from_float<T>(y);
+    }
+}
+
+template <typename T>
+infiniStatus_t launch_typed(const MambaSelectiveScanInfo &info, void *out, const void *x,
+                            const void *dt, const void *b, const void *c, const void *a_log,
+                            const void *d, const void *gate, const void *dt_bias, void *state,
+                            hcStream_t stream) {
+    constexpr int threads = 256;
+    size_t total = info.batch * info.intermediate;
+    dim3 blocks((total + threads - 1) / threads);
+    mamba_selective_scan_kernel<T><<<blocks, threads, 0, stream>>>(
+        reinterpret_cast<T *>(out), reinterpret_cast<const T *>(x), reinterpret_cast<const T *>(dt),
+        reinterpret_cast<const T *>(b), reinterpret_cast<const T *>(c), reinterpret_cast<const T *>(a_log),
+        reinterpret_cast<const T *>(d), reinterpret_cast<const T *>(gate), reinterpret_cast<const T *>(dt_bias),
+        reinterpret_cast<float *>(state), info.batch, info.seq_len, info.intermediate, info.state_size);
+    CHECK_METAX(hcGetLastError());
+    return INFINI_STATUS_SUCCESS;
+}
+} // namespace
+
+infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr,
+                                  infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t x_desc,
+                                  infiniopTensorDescriptor_t dt_desc, infiniopTensorDescriptor_t b_desc,
+                                  infiniopTensorDescriptor_t c_desc, infiniopTensorDescriptor_t a_log_desc,
+                                  infiniopTensorDescriptor_t d_desc, infiniopTensorDescriptor_t gate_desc,
+                                  infiniopTensorDescriptor_t dt_bias_desc, infiniopTensorDescriptor_t state_desc) {
+    auto result = MambaSelectiveScanInfo::create(out_desc, x_desc, dt_desc, b_desc, c_desc, a_log_desc, d_desc, gate_desc, dt_bias_desc, state_desc);
+    CHECK_RESULT(result);
+    auto info = result.take();
+    *desc_ptr = new Descriptor(new Opaque{reinterpret_cast<device::metax::Handle *>(handle)->internal()}, info, 0, handle->device, handle->device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, void *out,
+                                     const void *x, const void *dt, const void *b, const void *c,
+                                     const void *a_log, const void *d, const void *gate,
+                                     const void *dt_bias, void *state, void *stream_) const {
+    (void)workspace;
+    if (workspace_size < _workspace_size) {
+        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
+    }
+    auto stream = reinterpret_cast<hcStream_t>(stream_);
+    if (_info.dtype == INFINI_DTYPE_F16) {
+        return launch_typed<half>(_info, out, x, dt, b, c, a_log, d, gate, dt_bias, state, stream);
+    }
+    if (_info.dtype == INFINI_DTYPE_BF16) {
+        return launch_typed<__nv_bfloat16>(_info, out, x, dt, b, c, a_log, d, gate, dt_bias, state, stream);
+    }
+    if (_info.dtype == INFINI_DTYPE_F32) {
+        return launch_typed<float>(_info, out, x, dt, b, c, a_log, d, gate, dt_bias, state, stream);
+    }
+    return INFINI_STATUS_BAD_TENSOR_DTYPE;
+}
+
+} // namespace op::mamba_selective_scan::metax

src/infiniop/ops/mamba_selective_scan/operator.cc

@@ -5,6 +5,10 @@
 #include "nvidia/mamba_selective_scan_nvidia.cuh"
 #endif
 
+#ifdef ENABLE_METAX_API
+#include "metax/mamba_selective_scan_metax.h"
+#endif
+
 __INFINI_C infiniStatus_t infiniopCreateMambaSelectiveScanDescriptor(
     infiniopHandle_t handle, infiniopMambaSelectiveScanDescriptor_t *desc_ptr,
     infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t x_desc,
@@ -30,6 +34,9 @@ __INFINI_C infiniStatus_t infiniopCreateMambaSelectiveScanDescriptor(
 #endif
 #ifdef ENABLE_HYGON_API
         CREATE(INFINI_DEVICE_HYGON, nvidia);
+#endif
+#ifdef ENABLE_METAX_API
+        CREATE(INFINI_DEVICE_METAX, metax);
 #endif
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -57,6 +64,9 @@ __INFINI_C infiniStatus_t infiniopGetMambaSelectiveScanWorkspaceSize(infiniopMam
 #endif
 #ifdef ENABLE_HYGON_API
         GET(INFINI_DEVICE_HYGON, nvidia);
+#endif
+#ifdef ENABLE_METAX_API
+        GET(INFINI_DEVICE_METAX, metax);
 #endif
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -87,6 +97,9 @@ __INFINI_C infiniStatus_t infiniopMambaSelectiveScan(
 #endif
 #ifdef ENABLE_HYGON_API
         CALC(INFINI_DEVICE_HYGON, nvidia);
+#endif
+#ifdef ENABLE_METAX_API
+        CALC(INFINI_DEVICE_METAX, metax);
 #endif
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
@@ -114,6 +127,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMambaSelectiveScanDescriptor(infiniopMa
 #endif
 #ifdef ENABLE_HYGON_API
         DESTROY(INFINI_DEVICE_HYGON, nvidia);
+#endif
+#ifdef ENABLE_METAX_API
+        DESTROY(INFINI_DEVICE_METAX, metax);
 #endif
     default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;

src/infiniop/ops/moe_sum/metax/moe_sum_metax.h

@@ -0,0 +1,8 @@
+#ifndef __MOE_SUM_METAX_H__
+#define __MOE_SUM_METAX_H__
+
+#include "../moe_sum.h"
+
+DESCRIPTOR(metax)
+
+#endif