feat: add per-tensor FP8 e4m3 quant/dequant ops for KV cache quantiza…

include/infinicore/ops.hpp

@@ -64,6 +64,8 @@
 #include "ops/paged_attention_prefill.hpp"
 #include "ops/paged_caching.hpp"
 #include "ops/per_tensor_dequant_i8.hpp"
+#include "ops/per_tensor_dequant_fp8.hpp"
+#include "ops/per_tensor_quant_fp8.hpp"
 #include "ops/per_tensor_quant_i8.hpp"
 #include "ops/prepare_moe_input.hpp"
 #include "ops/quickgelu.hpp"

include/infinicore/ops/per_tensor_dequant_fp8.hpp

@@ -0,0 +1,11 @@
+#pragma once
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(PerTensorDequantFp8, const Tensor &, const Tensor &, const Tensor &);
+
+void per_tensor_dequant_fp8_(Tensor x, const Tensor &x_packed, const Tensor &x_scale);
+} // namespace infinicore::op

include/infinicore/ops/per_tensor_quant_fp8.hpp

@@ -0,0 +1,13 @@
+#pragma once
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(PerTensorQuantFp8, const Tensor &, Tensor, Tensor, bool);
+
+void per_tensor_quant_fp8_(const Tensor &x, Tensor x_packed, Tensor x_scale, bool is_static);
+
+Tensor per_tensor_quant_fp8(const Tensor &x, Tensor x_scale, bool is_static);
+} // namespace infinicore::op

include/infiniop.h

@@ -120,6 +120,8 @@
 #include "infiniop/ops/pixel_shuffle.h"
 #include "infiniop/ops/prepare_moe_input.h"
 #include "infiniop/ops/quant/per_channel_quant_int8.h"
+#include "infiniop/ops/dequant/per_tensor_dequant_fp8.h"
+#include "infiniop/ops/quant/per_tensor_quant_fp8.h"
 #include "infiniop/ops/quant/per_tensor_quant_int8.h"
 #include "infiniop/ops/quickgelu.h"
 #include "infiniop/ops/random_sample.h"

include/infiniop/ops/dequant/per_tensor_dequant_fp8.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_PER_TENSOR_DEQUANT_FP8_API_H__
+#define __INFINIOP_PER_TENSOR_DEQUANT_FP8_API_H__
+
+#include "../../operator_descriptor.h"
+
+typedef InfiniopDescriptor *infiniopPerTensorDequantFp8Descriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreatePerTensorDequantFp8Descriptor(infiniopHandle_t handle,
+                                                                               infiniopPerTensorDequantFp8Descriptor_t *desc_ptr,
+                                                                               infiniopTensorDescriptor_t x_desc,
+                                                                               infiniopTensorDescriptor_t x_packed_desc,
+                                                                               infiniopTensorDescriptor_t x_scale_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetPerTensorDequantFp8WorkspaceSize(infiniopPerTensorDequantFp8Descriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopPerTensorDequantFp8(infiniopPerTensorDequantFp8Descriptor_t desc,
+                                                               void *workspace,
+                                                               size_t workspace_size,
+                                                               void *x,
+                                                               const void *x_packed,
+                                                               const void *x_scale,
+                                                               void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyPerTensorDequantFp8Descriptor(infiniopPerTensorDequantFp8Descriptor_t desc);
+
+#endif

include/infiniop/ops/quant/per_tensor_quant_fp8.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_PER_TENSOR_QUANT_FP8_API_H__
+#define __INFINIOP_PER_TENSOR_QUANT_FP8_API_H__
+
+#include "../../operator_descriptor.h"
+
+typedef InfiniopDescriptor *infiniopPerTensorQuantFp8Descriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreatePerTensorQuantFp8Descriptor(infiniopHandle_t handle,
+                                                                             infiniopPerTensorQuantFp8Descriptor_t *desc_ptr,
+                                                                             infiniopTensorDescriptor_t x_packed_desc,
+                                                                             infiniopTensorDescriptor_t x_scale_desc,
+                                                                             infiniopTensorDescriptor_t x_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetPerTensorQuantFp8WorkspaceSize(infiniopPerTensorQuantFp8Descriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopPerTensorQuantFp8(infiniopPerTensorQuantFp8Descriptor_t desc,
+                                                             void *workspace,
+                                                             size_t workspace_size,
+                                                             void *x_packed,
+                                                             void *x_scale,
+                                                             const void *x,
+                                                             const bool is_static,
+                                                             void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyPerTensorQuantFp8Descriptor(infiniopPerTensorQuantFp8Descriptor_t desc);
+
+#endif

src/infinicore/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8.cc

@@ -0,0 +1,20 @@
+#include "../../../utils.hpp"
+#include "infinicore/ops/per_tensor_dequant_fp8.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PerTensorDequantFp8);
+
+PerTensorDequantFp8::PerTensorDequantFp8(const Tensor &x, const Tensor &x_packed, const Tensor &x_scale) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, x_packed, x_scale);
+    INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, x_packed, x_scale);
+}
+
+void PerTensorDequantFp8::execute(const Tensor &x, const Tensor &x_packed, const Tensor &x_scale) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(PerTensorDequantFp8, x, x_packed, x_scale);
+}
+
+void per_tensor_dequant_fp8_(Tensor x, const Tensor &x_packed, const Tensor &x_scale) {
+    PerTensorDequantFp8::execute(x, x_packed, x_scale);
+}
+} // namespace infinicore::op

src/infinicore/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8_infiniop.cc

@@ -0,0 +1,50 @@
+#include "../../infiniop_impl.hpp"
+#include "infinicore/ops/per_tensor_dequant_fp8.hpp"
+
+namespace infinicore::op::per_tensor_dequant_fp8_impl::infiniop {
+
+INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PerTensorDequantFp8, 100);
+
+struct PlannedMeta {
+    std::shared_ptr<Descriptor> descriptor;
+    graph::GraphTensor workspace, x, x_packed, x_scale;
+};
+
+void *plan(const Tensor &x, const Tensor &x_packed, const Tensor &x_scale) {
+    size_t seed = hash_combine(x, x_packed, x_scale);
+
+    INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(
+        Descriptor, descriptor, PerTensorDequantFp8,
+        seed,
+        x->desc(), x_packed->desc(), x_scale->desc());
+
+    INFINIOP_WORKSPACE_TENSOR(workspace, PerTensorDequantFp8, descriptor);
+
+    return new PlannedMeta{
+        descriptor,
+        graph::GraphTensor(workspace),
+        graph::GraphTensor(x),
+        graph::GraphTensor(x_packed),
+        graph::GraphTensor(x_scale)};
+}
+
+void run(void *planned_meta) {
+    auto planned = reinterpret_cast<PlannedMeta *>(planned_meta);
+    INFINICORE_CHECK_ERROR(infiniopPerTensorDequantFp8(
+        planned->descriptor->desc,
+        planned->workspace->data(),
+        planned->workspace->numel(),
+        planned->x->data(),
+        planned->x_packed->data(),
+        planned->x_scale->data(),
+        context::getStream()));
+}
+
+void cleanup(void **planned_meta_ptr) {
+    delete *reinterpret_cast<PlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(PerTensorDequantFp8, &plan, &run, &cleanup);
+
+} // namespace infinicore::op::per_tensor_dequant_fp8_impl::infiniop

src/infinicore/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8.cc

@@ -0,0 +1,26 @@
+#include "../../../utils.hpp"
+#include "infinicore/ops/per_tensor_quant_fp8.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PerTensorQuantFp8);
+
+PerTensorQuantFp8::PerTensorQuantFp8(const Tensor &x, Tensor x_packed, Tensor x_scale, bool is_static) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, x_packed, x_scale);
+    INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, x_packed, x_scale, is_static);
+}
+
+void PerTensorQuantFp8::execute(const Tensor &x, Tensor x_packed, Tensor x_scale, bool is_static) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(PerTensorQuantFp8, x, x_packed, x_scale, is_static);
+}
+
+void per_tensor_quant_fp8_(const Tensor &x, Tensor x_packed, Tensor x_scale, bool is_static) {
+    PerTensorQuantFp8::execute(x, x_packed, x_scale, is_static);
+}
+
+Tensor per_tensor_quant_fp8(const Tensor &x, Tensor x_scale, bool is_static) {
+    auto x_packed = Tensor::strided_empty(x->shape(), x->strides(), infinicore::DataType::F8, x->device());
+    PerTensorQuantFp8::execute(x, x_packed, x_scale, is_static);
+    return x_packed;
+}
+} // namespace infinicore::op

src/infinicore/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8_infiniop.cc

@@ -0,0 +1,54 @@
+#include "../../infiniop_impl.hpp"
+#include "infinicore/ops/per_tensor_quant_fp8.hpp"
+
+namespace infinicore::op::per_tensor_quant_fp8_impl::infiniop {
+
+INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PerTensorQuantFp8, 100);
+
+struct PlannedMeta {
+    std::shared_ptr<Descriptor> descriptor;
+    graph::GraphTensor workspace, x, x_packed, x_scale;
+    const bool is_static;
+};
+
+void *plan(const Tensor &x, Tensor x_packed, Tensor x_scale, bool is_static) {
+    size_t seed = hash_combine(x, x_packed, x_scale);
+
+    INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(
+        Descriptor, descriptor, PerTensorQuantFp8,
+        seed,
+        x_packed->desc(), x_scale->desc(), x->desc());
+
+    INFINIOP_WORKSPACE_TENSOR(workspace, PerTensorQuantFp8, descriptor);
+
+    return new PlannedMeta{
+        descriptor,
+        graph::GraphTensor(workspace),
+        graph::GraphTensor(x),
+        graph::GraphTensor(x_packed),
+        graph::GraphTensor(x_scale),
+        is_static};
+}
+
+void run(void *planned_meta) {
+    auto planned = reinterpret_cast<PlannedMeta *>(planned_meta);
+    const bool is_static = planned->is_static;
+    INFINICORE_CHECK_ERROR(infiniopPerTensorQuantFp8(
+        planned->descriptor->desc,
+        planned->workspace->data(),
+        planned->workspace->numel(),
+        planned->x_packed->data(),
+        planned->x_scale->data(),
+        planned->x->data(),
+        is_static,
+        context::getStream()));
+}
+
+void cleanup(void **planned_meta_ptr) {
+    delete *reinterpret_cast<PlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(PerTensorQuantFp8, &plan, &run, &cleanup);
+
+} // namespace infinicore::op::per_tensor_quant_fp8_impl::infiniop

src/infiniop/ops/dequant/per_tensor_dequant_fp8/cuda/kernel.cuh

@@ -0,0 +1,52 @@
+#ifndef __PER_TENSOR_DEQUANT_FP8_KERNEL_CUH__
+#define __PER_TENSOR_DEQUANT_FP8_KERNEL_CUH__
+
+// Convert an FP8 e4m3fn byte to float.
+// e4m3fn: 1 sign + 4 exponent (bias 7) + 3 mantissa; no inf/nan.
+__device__ __forceinline__ float e4m3_to_float(unsigned char b) {
+    unsigned int exp = (b >> 3) & 0x0fu;
+    unsigned int mant = b & 0x07u;
+    float val;
+    if (exp == 0) {
+        val = (float)mant * 0x1p-9f; // subnormal: m * 2^-9
+    } else {
+        // exp is unsigned: subtract 7 in signed arithmetic to avoid wrapping
+        // for exp < 7 (values below 1.0).
+        val = (1.0f + (float)mant / 8.0f) * exp2f((float)((int)exp - 7));
+    }
+    return (b & 0x80u) ? -val : val;
+}
+
+template <typename Tout>
+__device__ void perTensorDequantFp8SymKernel(
+    Tout *x, const unsigned char *x_packed, const float *x_scale,
+    size_t batch_size, size_t channel, size_t hidden_dim, size_t width,
+    ptrdiff_t strides_0, ptrdiff_t strides_1, ptrdiff_t strides_2, ptrdiff_t strides_3,
+    ptrdiff_t p_strides_0, ptrdiff_t p_strides_1, ptrdiff_t p_strides_2, ptrdiff_t p_strides_3,
+    int num_elements) {
+
+    unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    const int grid_size = blockDim.x * gridDim.x;
+    float x_scale_val = x_scale[0];
+    for (int ind = gid; ind < num_elements; ind += grid_size) {
+        int tid = ind;
+        int w = tid % (int)width;
+        tid = tid / (int)width;
+
+        int h = tid % (int)hidden_dim;
+        tid = tid / (int)hidden_dim;
+
+        int c = tid % (int)channel;
+        tid = tid / (int)channel;
+
+        int b = tid % (int)batch_size;
+
+        int index = w * (int)strides_3 + h * (int)strides_2 + c * (int)strides_1 + b * (int)strides_0;
+        int p_index = w * (int)p_strides_3 + h * (int)p_strides_2 + c * (int)p_strides_1 + b * (int)p_strides_0;
+
+        float val = e4m3_to_float(x_packed[p_index]) * x_scale_val;
+        x[index] = static_cast<Tout>(val);
+    }
+}
+
+#endif // __PER_TENSOR_DEQUANT_FP8_KERNEL_CUH__

src/infiniop/ops/dequant/per_tensor_dequant_fp8/info.h

@@ -0,0 +1,75 @@
+#ifndef __PER_TENSOR_DEQUANT_FP8_INFO_H__
+#define __PER_TENSOR_DEQUANT_FP8_INFO_H__
+
+#include "../../../../utils.h"
+#include "../../../operator.h"
+#include "../../../tensor.h"
+
+namespace op::per_tensor_dequant_fp8 {
+
+class PerTensorDequantFp8Info {
+private:
+    PerTensorDequantFp8Info() = default;
+
+public:
+    infiniDtype_t dtype, packed_type;
+    size_t batch_size, channel, hidden_dim, width;
+    ptrdiff_t strides_0, strides_1, strides_2, strides_3;
+    ptrdiff_t p_strides_0, p_strides_1, p_strides_2, p_strides_3;
+    int num_elements;
+
+    static utils::Result<PerTensorDequantFp8Info> createPerTensorDequantFp8Info(
+        infiniopTensorDescriptor_t x_desc,
+        infiniopTensorDescriptor_t x_packed_desc,
+        infiniopTensorDescriptor_t x_scale_desc) {
+
+        CHECK_OR_RETURN(
+            x_desc != nullptr && x_packed_desc != nullptr && x_scale_desc != nullptr,
+            INFINI_STATUS_NULL_POINTER);
+
+        const infiniDtype_t dtype = x_desc->dtype();
+        const infiniDtype_t packed_type = x_packed_desc->dtype();
+
+        CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32);
+        CHECK_DTYPE(packed_type, INFINI_DTYPE_F8);
+
+        auto shape = x_desc->shape();
+        CHECK_SAME_SHAPE(shape, x_packed_desc->shape());
+
+        auto ndim = x_desc->ndim();
+        CHECK_OR_RETURN(ndim <= 4,
+                        INFINI_STATUS_BAD_TENSOR_SHAPE);
+
+        size_t width = shape[ndim - 1];
+        size_t hidden_dim = (ndim > 1 ? shape[ndim - 2] : 1);
+        size_t channel = (ndim > 2 ? shape[ndim - 3] : 1);
+        size_t batch_size = (ndim > 3 ? shape[ndim - 4] : 1);
+
+        ptrdiff_t strides_3 = x_desc->strides()[ndim - 1];
+        ptrdiff_t strides_2 = (ndim > 1 ? x_desc->strides()[ndim - 2] : 0);
+        ptrdiff_t strides_1 = (ndim > 2 ? x_desc->strides()[ndim - 3] : 0);
+        ptrdiff_t strides_0 = (ndim > 3 ? x_desc->strides()[ndim - 4] : 0);
+
+        ptrdiff_t p_strides_3 = x_packed_desc->strides()[ndim - 1];
+        ptrdiff_t p_strides_2 = (ndim > 1 ? x_packed_desc->strides()[ndim - 2] : 0);
+        ptrdiff_t p_strides_1 = (ndim > 2 ? x_packed_desc->strides()[ndim - 3] : 0);
+        ptrdiff_t p_strides_0 = (ndim > 3 ? x_packed_desc->strides()[ndim - 4] : 0);
+
+        int num_elements = 1;
+        for (int i = 0; i < (int)ndim; i++) {
+            num_elements *= static_cast<int>(shape[i]);
+        }
+
+        return utils::Result<PerTensorDequantFp8Info>(PerTensorDequantFp8Info{
+            dtype,
+            packed_type,
+            batch_size, channel, hidden_dim, width,
+            strides_0, strides_1, strides_2, strides_3,
+            p_strides_0, p_strides_1, p_strides_2, p_strides_3,
+            num_elements});
+    }
+};
+
+} // namespace op::per_tensor_dequant_fp8
+
+#endif //  __PER_TENSOR_DEQUANT_FP8_INFO_H__