From e4c34c72167eb09f76aa405d85f06ed221317ddd Mon Sep 17 00:00:00 2001 From: bobodai Date: Sun, 23 Aug 2026 20:56:46 +0800 Subject: [PATCH 1/2] feat: add per-tensor FP8 e4m3 quant/dequant ops for KV cache quantization --- include/infinicore/ops.hpp | 2 + .../infinicore/ops/per_tensor_dequant_fp8.hpp | 11 + .../infinicore/ops/per_tensor_quant_fp8.hpp | 13 + include/infiniop.h | 2 + .../ops/dequant/per_tensor_dequant_fp8.h | 26 ++ .../infiniop/ops/quant/per_tensor_quant_fp8.h | 27 ++ .../per_tensor_dequant_fp8.cc | 20 + .../per_tensor_dequant_fp8_infiniop.cc | 50 +++ .../per_tensor_quant_fp8.cc | 26 ++ .../per_tensor_quant_fp8_infiniop.cc | 54 +++ .../per_tensor_dequant_fp8/cuda/kernel.cuh | 52 +++ .../ops/dequant/per_tensor_dequant_fp8/info.h | 75 ++++ .../nvidia/per_tensor_dequant_fp8_nvidia.cu | 105 ++++++ .../nvidia/per_tensor_dequant_fp8_nvidia.cuh | 7 + .../per_tensor_dequant_fp8/operator.cc | 99 +++++ .../per_tensor_dequant_fp8.h | 39 ++ .../per_tensor_quant_fp8/cuda/kernel.cuh | 152 ++++++++ .../ops/quant/per_tensor_quant_fp8/info.h | 76 ++++ .../nvidia/per_tensor_quant_fp8_nvidia.cu | 124 +++++++ .../nvidia/per_tensor_quant_fp8_nvidia.cuh | 7 + .../quant/per_tensor_quant_fp8/operator.cc | 100 +++++ .../per_tensor_quant_fp8.h | 39 ++ test/infiniop/libinfiniop/op_register.py | 70 ++++ test/infiniop/per_tensor_dequant_fp8.py | 212 +++++++++++ test/infiniop/per_tensor_quant_fp8.py | 286 +++++++++++++++ test/infiniop/w8a8fp8.py | 343 ++++++++++++++++++ 26 files changed, 2017 insertions(+) create mode 100644 include/infinicore/ops/per_tensor_dequant_fp8.hpp create mode 100644 include/infinicore/ops/per_tensor_quant_fp8.hpp create mode 100644 include/infiniop/ops/dequant/per_tensor_dequant_fp8.h create mode 100644 include/infiniop/ops/quant/per_tensor_quant_fp8.h create mode 100644 src/infinicore/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8.cc create mode 100644 src/infinicore/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8_infiniop.cc create mode 100644 src/infinicore/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8.cc create mode 100644 src/infinicore/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8_infiniop.cc create mode 100644 src/infiniop/ops/dequant/per_tensor_dequant_fp8/cuda/kernel.cuh create mode 100644 src/infiniop/ops/dequant/per_tensor_dequant_fp8/info.h create mode 100644 src/infiniop/ops/dequant/per_tensor_dequant_fp8/nvidia/per_tensor_dequant_fp8_nvidia.cu create mode 100644 src/infiniop/ops/dequant/per_tensor_dequant_fp8/nvidia/per_tensor_dequant_fp8_nvidia.cuh create mode 100644 src/infiniop/ops/dequant/per_tensor_dequant_fp8/operator.cc create mode 100644 src/infiniop/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8.h create mode 100644 src/infiniop/ops/quant/per_tensor_quant_fp8/cuda/kernel.cuh create mode 100644 src/infiniop/ops/quant/per_tensor_quant_fp8/info.h create mode 100644 src/infiniop/ops/quant/per_tensor_quant_fp8/nvidia/per_tensor_quant_fp8_nvidia.cu create mode 100644 src/infiniop/ops/quant/per_tensor_quant_fp8/nvidia/per_tensor_quant_fp8_nvidia.cuh create mode 100644 src/infiniop/ops/quant/per_tensor_quant_fp8/operator.cc create mode 100644 src/infiniop/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8.h create mode 100644 test/infiniop/per_tensor_dequant_fp8.py create mode 100644 test/infiniop/per_tensor_quant_fp8.py create mode 100644 test/infiniop/w8a8fp8.py diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index 5e93e1457..884c79d6b 100644 --- a/include/infinicore/ops.hpp +++ b/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" diff --git a/include/infinicore/ops/per_tensor_dequant_fp8.hpp b/include/infinicore/ops/per_tensor_dequant_fp8.hpp new file mode 100644 index 000000000..d8fbbf7c8 --- /dev/null +++ b/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 diff --git a/include/infinicore/ops/per_tensor_quant_fp8.hpp b/include/infinicore/ops/per_tensor_quant_fp8.hpp new file mode 100644 index 000000000..2f8cd4abd --- /dev/null +++ b/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 diff --git a/include/infiniop.h b/include/infiniop.h index 9f632e27f..27b6573c2 100644 --- a/include/infiniop.h +++ b/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" diff --git a/include/infiniop/ops/dequant/per_tensor_dequant_fp8.h b/include/infiniop/ops/dequant/per_tensor_dequant_fp8.h new file mode 100644 index 000000000..f68499885 --- /dev/null +++ b/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 diff --git a/include/infiniop/ops/quant/per_tensor_quant_fp8.h b/include/infiniop/ops/quant/per_tensor_quant_fp8.h new file mode 100644 index 000000000..faf9e9859 --- /dev/null +++ b/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 diff --git a/src/infinicore/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8.cc b/src/infinicore/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8.cc new file mode 100644 index 000000000..58035602d --- /dev/null +++ b/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 diff --git a/src/infinicore/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8_infiniop.cc b/src/infinicore/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8_infiniop.cc new file mode 100644 index 000000000..3c6949029 --- /dev/null +++ b/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; + 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(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(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 diff --git a/src/infinicore/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8.cc b/src/infinicore/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8.cc new file mode 100644 index 000000000..e2ad19baf --- /dev/null +++ b/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 diff --git a/src/infinicore/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8_infiniop.cc b/src/infinicore/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8_infiniop.cc new file mode 100644 index 000000000..139976778 --- /dev/null +++ b/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; + 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(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(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 diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_fp8/cuda/kernel.cuh b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/cuda/kernel.cuh new file mode 100644 index 000000000..f97977f1e --- /dev/null +++ b/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 +__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(val); + } +} + +#endif // __PER_TENSOR_DEQUANT_FP8_KERNEL_CUH__ diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_fp8/info.h b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/info.h new file mode 100644 index 000000000..b868a5e19 --- /dev/null +++ b/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 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(shape[i]); + } + + return utils::Result(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__ diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_fp8/nvidia/per_tensor_dequant_fp8_nvidia.cu b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/nvidia/per_tensor_dequant_fp8_nvidia.cu new file mode 100644 index 000000000..f3e9bd222 --- /dev/null +++ b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/nvidia/per_tensor_dequant_fp8_nvidia.cu @@ -0,0 +1,105 @@ +#include "../../../../devices/nvidia/nvidia_common.cuh" +#include "per_tensor_dequant_fp8_nvidia.cuh" + +#include "../../../../devices/nvidia/nvidia_kernel_common.cuh" +#include + +#include "../cuda/kernel.cuh" + +template +INFINIOP_CUDA_KERNEL perTensorDequantFp8Sym( + 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) { + perTensorDequantFp8SymKernel(x, x_packed, x_scale, + 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::nvidia { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, Descriptor **desc_ptr, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t x_packed_desc, + infiniopTensorDescriptor_t x_scale_desc) { + auto info = PerTensorDequantFp8Info::createPerTensorDequantFp8Info(x_desc, x_packed_desc, x_scale_desc); + CHECK_RESULT(info); + + *desc_ptr = new Descriptor( + new Opaque{reinterpret_cast(handle)->internal()}, + info.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +template +infiniStatus_t per_tensor_dequant_fp8Kernel(const PerTensorDequantFp8Info &info, Tdata *x, const unsigned char *x_packed, const float *x_scale, cudaStream_t stream) { + int num_elements = (int)info.num_elements; + int num_blocks = (num_elements + BLOCK_SIZE - 1) / BLOCK_SIZE; + + size_t batch_size = info.batch_size; + size_t channel = info.channel; + size_t hidden_dim = info.hidden_dim; + size_t width = info.width; + + ptrdiff_t strides_0 = info.strides_0; + ptrdiff_t strides_1 = info.strides_1; + ptrdiff_t strides_2 = info.strides_2; + ptrdiff_t strides_3 = info.strides_3; + + ptrdiff_t p_strides_0 = info.p_strides_0; + ptrdiff_t p_strides_1 = info.p_strides_1; + ptrdiff_t p_strides_2 = info.p_strides_2; + ptrdiff_t p_strides_3 = info.p_strides_3; + + perTensorDequantFp8Sym + <<>>(x, x_packed, x_scale, + 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); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, + void *x, const void *x_packed, const void *x_scale, + void *stream_) const { + cudaStream_t stream = (cudaStream_t)stream_; +#define DEQUANT(BLOCK_SIZE, TDATA) \ + per_tensor_dequant_fp8Kernel(_info, (TDATA *)x, (const unsigned char *)x_packed, (const float *)x_scale, stream) +#define DEQUANT_WITH_BLOCK_SIZE(BLOCK_SIZE) \ + { \ + if (_info.dtype == INFINI_DTYPE_F16) \ + return DEQUANT(BLOCK_SIZE, half); \ + else if (_info.dtype == INFINI_DTYPE_F32) \ + return DEQUANT(BLOCK_SIZE, float); \ + else if (_info.dtype == INFINI_DTYPE_BF16) \ + return DEQUANT(BLOCK_SIZE, __nv_bfloat16); \ + else \ + return INFINI_STATUS_BAD_TENSOR_DTYPE; \ + } + if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { + DEQUANT_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_1024) + } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { + DEQUANT_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_512) + } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { + DEQUANT_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_4096) + } else { + return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; + } + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::per_tensor_dequant_fp8::nvidia diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_fp8/nvidia/per_tensor_dequant_fp8_nvidia.cuh b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/nvidia/per_tensor_dequant_fp8_nvidia.cuh new file mode 100644 index 000000000..f876116ef --- /dev/null +++ b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/nvidia/per_tensor_dequant_fp8_nvidia.cuh @@ -0,0 +1,7 @@ +#ifndef __PER_TENSOR_DEQUANT_FP8_NVIDIA_API_H__ +#define __PER_TENSOR_DEQUANT_FP8_NVIDIA_API_H__ +#include "../per_tensor_dequant_fp8.h" + +DESCRIPTOR(nvidia) + +#endif // __PER_TENSOR_DEQUANT_FP8_NVIDIA_API_H__ diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_fp8/operator.cc b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/operator.cc new file mode 100644 index 000000000..06df37002 --- /dev/null +++ b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/operator.cc @@ -0,0 +1,99 @@ +#include "../../../operator.h" +#include "../../../handle.h" +#include "infiniop/ops/dequant/per_tensor_dequant_fp8.h" + +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) +#include "nvidia/per_tensor_dequant_fp8_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreatePerTensorDequantFp8Descriptor(infiniopHandle_t handle, + infiniopPerTensorDequantFp8Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t x_packed_desc, + infiniopTensorDescriptor_t x_scale_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::per_tensor_dequant_fp8::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + x_desc, \ + x_packed_desc, \ + x_scale_desc); + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia) +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia) +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetPerTensorDequantFp8WorkspaceSize(infiniopPerTensorDequantFp8Descriptor_t desc, size_t *size) { + switch (desc->device_type) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->minWorkspaceSize(); \ + return INFINI_STATUS_SUCCESS; +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia) +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia) +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C infiniStatus_t infiniopPerTensorDequantFp8(infiniopPerTensorDequantFp8Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *x, + const void *x_packed, + const void *x_scale, + void *stream) { +#define DEQUANT(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate( \ + workspace, workspace_size, x, x_packed, x_scale, stream); + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DEQUANT(INFINI_DEVICE_NVIDIA, nvidia) +#endif +#ifdef ENABLE_QY_API + DEQUANT(INFINI_DEVICE_QY, nvidia) +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DEQUANT +} + +__INFINI_C infiniStatus_t infiniopDestroyPerTensorDequantFp8Descriptor(infiniopPerTensorDequantFp8Descriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS; + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia) +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia) +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8.h b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8.h new file mode 100644 index 000000000..13d5b6255 --- /dev/null +++ b/src/infiniop/ops/dequant/per_tensor_dequant_fp8/per_tensor_dequant_fp8.h @@ -0,0 +1,39 @@ +#ifndef __PER_TENSOR_DEQUANT_FP8_H__ +#define __PER_TENSOR_DEQUANT_FP8_H__ + +#include "../../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + \ + namespace op::per_tensor_dequant_fp8::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + PerTensorDequantFp8Info _info; \ + size_t _workspace_size; \ + \ + Descriptor(Opaque *opaque, PerTensorDequantFp8Info info, \ + size_t workspace_size, \ + infiniDevice_t device_type, int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + ~Descriptor(); \ + \ + size_t minWorkspaceSize() const { return _workspace_size; } \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t x_desc, \ + infiniopTensorDescriptor_t x_packed_desc, \ + infiniopTensorDescriptor_t x_scale_desc); \ + \ + infiniStatus_t calculate( \ + void *workspace, size_t workspace_size, \ + void *x, const void *x_packed, const void *x_scale, void *stream) const; \ + }; \ + } + +#endif // __PER_TENSOR_DEQUANT_FP8_H__ diff --git a/src/infiniop/ops/quant/per_tensor_quant_fp8/cuda/kernel.cuh b/src/infiniop/ops/quant/per_tensor_quant_fp8/cuda/kernel.cuh new file mode 100644 index 000000000..fd2fbd0fe --- /dev/null +++ b/src/infiniop/ops/quant/per_tensor_quant_fp8/cuda/kernel.cuh @@ -0,0 +1,152 @@ +#ifndef __PER_TENSOR_QUANT_FP8_KERNEL_CUH__ +#define __PER_TENSOR_QUANT_FP8_KERNEL_CUH__ + +#include + +#ifndef WARP_SIZE +#define WARP_SIZE 32 +#endif + +#define FULL_MASK 0xffffffff + +// FP8 e4m3fn format: 1 sign + 4 exponent (bias 7) + 3 mantissa. +// Finite range: [-448, 448], min normal 2^-6, min subnormal 2^-9. +#define FP8_E4M3_MAX 448.0f + +// warp reduce max +__device__ __forceinline__ float warpReduceMax(float val) { + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + val = fmaxf(val, __shfl_xor_sync(FULL_MASK, val, offset)); + } + return val; +} + +// float atomic max (safe version) +__device__ __forceinline__ void atomicMaxFloat(float *addr, float val) { + int *addr_i = (int *)addr; + int old = *addr_i; + int assumed; + + do { + assumed = old; + float old_f = __int_as_float(assumed); + float new_f = fmaxf(val, old_f); + + old = atomicCAS(addr_i, assumed, __float_as_int(new_f)); + + } while (assumed != old); +} + +__device__ inline int round_half_away_from_zero(float x) { + float ax = fabsf(x); + float r = floorf(ax + 0.5f); + return (x >= 0.0f) ? (int)r : -(int)r; +} + +// Convert a float to an FP8 e4m3fn byte (round half away from zero, saturate). +__device__ __forceinline__ unsigned char float_to_e4m3(float x) { + unsigned int sign = 0; + if (x < 0.0f) { + sign = 0x80u; + x = -x; + } + if (x >= FP8_E4M3_MAX) { + return (unsigned char)(sign | 0x7eu); // 448 -> exp 15, mantissa 6 + } + if (x < 0x1p-9f) { + return (unsigned char)sign; // rounds to zero (0x1p-9 = 2^-9) + } + float e = floorf(log2f(x)); + int stored = (int)e + 7; + if (stored <= 0) { + // subnormal: value = m * 2^-9 + int m = (int)roundf(x * 512.0f); // x / 2^-9 + if (m > 7) { + m = 7; + } + return (unsigned char)(sign | (unsigned int)m); + } + float frac = x * exp2f(-e) - 1.0f; + int m = (int)roundf(frac * 8.0f); + if (m == 8) { + m = 0; + stored += 1; + if (stored > 15) { + return (unsigned char)(sign | 0x7eu); + } + } + return (unsigned char)(sign | ((unsigned int)stored << 3) | (unsigned int)m); +} + +template +__device__ void perTensorAbsmaxFp8Kernel(float *x_scale, const Tdata *x, + 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, + int num_elements) { + int idx = threadIdx.x; + int gid = blockIdx.x * blockDim.x + idx; + int grid_size = blockDim.x * gridDim.x; + + float local_max = 0.f; + + 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; + + float v = fabsf((float)x[index]); + + local_max = fmaxf(local_max, v); + } + + local_max = warpReduceMax(local_max); + if ((idx & (WARP_SIZE - 1)) == 0) { + atomicMaxFloat(x_scale, local_max / FP8_E4M3_MAX); + } +} + +template +__device__ void perTensorQuantFp8SymKernel( + unsigned char *x_packed, float *x_scale, const Tdata *x, + 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 scale_val = 1.0f / 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 qf = (float)x[index] * scale_val; + x_packed[p_index] = float_to_e4m3(qf); + } +} + +#endif // __PER_TENSOR_QUANT_FP8_KERNEL_CUH__ diff --git a/src/infiniop/ops/quant/per_tensor_quant_fp8/info.h b/src/infiniop/ops/quant/per_tensor_quant_fp8/info.h new file mode 100644 index 000000000..5c5486b05 --- /dev/null +++ b/src/infiniop/ops/quant/per_tensor_quant_fp8/info.h @@ -0,0 +1,76 @@ +#ifndef __PER_TENSOR_QUANT_FP8_INFO_H__ +#define __PER_TENSOR_QUANT_FP8_INFO_H__ + +#include "../../../../utils.h" +#include "../../../operator.h" +#include "../../../tensor.h" + +namespace op::per_tensor_quant_fp8 { + +class PerTensorQuantFp8Info { +private: + PerTensorQuantFp8Info() = 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; + bool is_static; + + static utils::Result createPerTensorQuantFp8Info( + infiniopTensorDescriptor_t x_packed_desc, + infiniopTensorDescriptor_t x_scale_desc, + infiniopTensorDescriptor_t x_desc) { + + CHECK_OR_RETURN( + x_packed_desc != nullptr && x_scale_desc != nullptr && x_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(shape[i]); + } + + return utils::Result(PerTensorQuantFp8Info{ + 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_quant_fp8 + +#endif // __PER_TENSOR_QUANT_FP8_INFO_H__ diff --git a/src/infiniop/ops/quant/per_tensor_quant_fp8/nvidia/per_tensor_quant_fp8_nvidia.cu b/src/infiniop/ops/quant/per_tensor_quant_fp8/nvidia/per_tensor_quant_fp8_nvidia.cu new file mode 100644 index 000000000..a56b130d9 --- /dev/null +++ b/src/infiniop/ops/quant/per_tensor_quant_fp8/nvidia/per_tensor_quant_fp8_nvidia.cu @@ -0,0 +1,124 @@ +#include "../../../../devices/nvidia/nvidia_common.cuh" +#include "per_tensor_quant_fp8_nvidia.cuh" + +#include "../../../../devices/nvidia/nvidia_kernel_common.cuh" +#include + +#include "../cuda/kernel.cuh" + +template +INFINIOP_CUDA_KERNEL perTensorAbsmaxFp8( + float *x_scale, const Tdata *x, + 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, + int num_elements) { + perTensorAbsmaxFp8Kernel(x_scale, x, + batch_size, channel, hidden_dim, width, + strides_0, strides_1, strides_2, strides_3, + num_elements); +} + +template +INFINIOP_CUDA_KERNEL perTensorQuantFp8Sym( + unsigned char *x_packed, float *x_scale, const Tdata *x, + 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) { + perTensorQuantFp8SymKernel(x_packed, x_scale, x, + 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_quant_fp8::nvidia { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, Descriptor **desc_ptr, + infiniopTensorDescriptor_t x_packed_desc, + infiniopTensorDescriptor_t x_scale_desc, + infiniopTensorDescriptor_t x_desc) { + auto info = PerTensorQuantFp8Info::createPerTensorQuantFp8Info(x_packed_desc, x_scale_desc, x_desc); + CHECK_RESULT(info); + + *desc_ptr = new Descriptor( + new Opaque{reinterpret_cast(handle)->internal()}, + info.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +template +infiniStatus_t per_tensor_quant_fp8Kernel(const PerTensorQuantFp8Info &info, unsigned char *x_packed, float *x_scale, const Tdata *x, const bool is_static, cudaStream_t stream) { + int num_elements = (int)info.num_elements; + int num_blocks = (num_elements + BLOCK_SIZE - 1) / BLOCK_SIZE; + + size_t batch_size = info.batch_size; + size_t channel = info.channel; + size_t hidden_dim = info.hidden_dim; + size_t width = info.width; + + ptrdiff_t strides_0 = info.strides_0; + ptrdiff_t strides_1 = info.strides_1; + ptrdiff_t strides_2 = info.strides_2; + ptrdiff_t strides_3 = info.strides_3; + + ptrdiff_t p_strides_0 = info.p_strides_0; + ptrdiff_t p_strides_1 = info.p_strides_1; + ptrdiff_t p_strides_2 = info.p_strides_2; + ptrdiff_t p_strides_3 = info.p_strides_3; + + if (!is_static) { + perTensorAbsmaxFp8 + <<>>(x_scale, x, + batch_size, channel, hidden_dim, width, + strides_0, strides_1, strides_2, strides_3, + num_elements); + } + perTensorQuantFp8Sym + <<>>(x_packed, x_scale, x, + 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); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, + void *x_packed, void *x_scale, const void *x, const bool is_static, + void *stream_) const { + cudaStream_t stream = (cudaStream_t)stream_; +#define QUANT(BLOCK_SIZE, TDATA) \ + per_tensor_quant_fp8Kernel(_info, (unsigned char *)x_packed, (float *)x_scale, (const TDATA *)x, is_static, stream) +#define QUANT_WITH_BLOCK_SIZE(BLOCK_SIZE) \ + { \ + if (_info.dtype == INFINI_DTYPE_F16) \ + return QUANT(BLOCK_SIZE, half); \ + else if (_info.dtype == INFINI_DTYPE_F32) \ + return QUANT(BLOCK_SIZE, float); \ + else if (_info.dtype == INFINI_DTYPE_BF16) \ + return QUANT(BLOCK_SIZE, __nv_bfloat16); \ + else \ + return INFINI_STATUS_BAD_TENSOR_DTYPE; \ + } + if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { + QUANT_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_1024) + } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { + QUANT_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_512) + } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { + QUANT_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_4096) + } else { + return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; + } + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::per_tensor_quant_fp8::nvidia diff --git a/src/infiniop/ops/quant/per_tensor_quant_fp8/nvidia/per_tensor_quant_fp8_nvidia.cuh b/src/infiniop/ops/quant/per_tensor_quant_fp8/nvidia/per_tensor_quant_fp8_nvidia.cuh new file mode 100644 index 000000000..329ec9e98 --- /dev/null +++ b/src/infiniop/ops/quant/per_tensor_quant_fp8/nvidia/per_tensor_quant_fp8_nvidia.cuh @@ -0,0 +1,7 @@ +#ifndef __PER_TENSOR_QUANT_FP8_NVIDIA_API_H__ +#define __PER_TENSOR_QUANT_FP8_NVIDIA_API_H__ +#include "../per_tensor_quant_fp8.h" + +DESCRIPTOR(nvidia) + +#endif // __PER_TENSOR_QUANT_FP8_NVIDIA_API_H__ diff --git a/src/infiniop/ops/quant/per_tensor_quant_fp8/operator.cc b/src/infiniop/ops/quant/per_tensor_quant_fp8/operator.cc new file mode 100644 index 000000000..dab4c8f4c --- /dev/null +++ b/src/infiniop/ops/quant/per_tensor_quant_fp8/operator.cc @@ -0,0 +1,100 @@ +#include "../../../operator.h" +#include "../../../handle.h" +#include "infiniop/ops/quant/per_tensor_quant_fp8.h" + +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) +#include "nvidia/per_tensor_quant_fp8_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreatePerTensorQuantFp8Descriptor(infiniopHandle_t handle, + infiniopPerTensorQuantFp8Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t x_packed_desc, + infiniopTensorDescriptor_t x_scale_desc, + infiniopTensorDescriptor_t x_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::per_tensor_quant_fp8::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + x_packed_desc, \ + x_scale_desc, \ + x_desc); + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia) +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia) +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetPerTensorQuantFp8WorkspaceSize(infiniopPerTensorQuantFp8Descriptor_t desc, size_t *size) { + switch (desc->device_type) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->minWorkspaceSize(); \ + return INFINI_STATUS_SUCCESS; +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia) +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia) +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C 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) { +#define QUANT(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate( \ + workspace, workspace_size, x_packed, x_scale, x, is_static, stream); + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + QUANT(INFINI_DEVICE_NVIDIA, nvidia) +#endif +#ifdef ENABLE_QY_API + QUANT(INFINI_DEVICE_QY, nvidia) +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef QUANT +} + +__INFINI_C infiniStatus_t infiniopDestroyPerTensorQuantFp8Descriptor(infiniopPerTensorQuantFp8Descriptor_t desc) { +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS; + + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia) +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia) +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8.h b/src/infiniop/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8.h new file mode 100644 index 000000000..b8d53919b --- /dev/null +++ b/src/infiniop/ops/quant/per_tensor_quant_fp8/per_tensor_quant_fp8.h @@ -0,0 +1,39 @@ +#ifndef __PER_TENSOR_QUANT_FP8_H__ +#define __PER_TENSOR_QUANT_FP8_H__ + +#include "../../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + \ + namespace op::per_tensor_quant_fp8::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + PerTensorQuantFp8Info _info; \ + size_t _workspace_size; \ + \ + Descriptor(Opaque *opaque, PerTensorQuantFp8Info info, \ + size_t workspace_size, \ + infiniDevice_t device_type, int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + ~Descriptor(); \ + \ + size_t minWorkspaceSize() const { return _workspace_size; } \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t x_packed_desc, \ + infiniopTensorDescriptor_t x_scale_desc, \ + infiniopTensorDescriptor_t x_desc); \ + \ + infiniStatus_t calculate( \ + void *workspace, size_t workspace_size, \ + void *x_packed, void *x_scale, const void *x, const bool is_static, void *stream) const; \ + }; \ + } + +#endif // __PER_TENSOR_QUANT_FP8_H__ diff --git a/test/infiniop/libinfiniop/op_register.py b/test/infiniop/libinfiniop/op_register.py index ec9add01a..1a4dfd6d5 100644 --- a/test/infiniop/libinfiniop/op_register.py +++ b/test/infiniop/libinfiniop/op_register.py @@ -1502,6 +1502,76 @@ def per_tensor_dequant_int8_(lib): ] +@OpRegister.operator +def per_tensor_quant_fp8_(lib): + lib.infiniopCreatePerTensorQuantFp8Descriptor.restype = c_int32 + lib.infiniopCreatePerTensorQuantFp8Descriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + ] + + lib.infiniopGetPerTensorQuantFp8WorkspaceSize.restype = c_int32 + lib.infiniopGetPerTensorQuantFp8WorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + + lib.infiniopPerTensorQuantFp8.restype = c_int32 + lib.infiniopPerTensorQuantFp8.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_bool, + c_void_p, + ] + + lib.infiniopDestroyPerTensorQuantFp8Descriptor.restype = c_int32 + lib.infiniopDestroyPerTensorQuantFp8Descriptor.argtypes = [ + infiniopOperatorDescriptor_t, + ] + + +@OpRegister.operator +def per_tensor_dequant_fp8_(lib): + lib.infiniopCreatePerTensorDequantFp8Descriptor.restype = c_int32 + lib.infiniopCreatePerTensorDequantFp8Descriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + ] + + lib.infiniopGetPerTensorDequantFp8WorkspaceSize.restype = c_int32 + lib.infiniopGetPerTensorDequantFp8WorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + + lib.infiniopPerTensorDequantFp8.restype = c_int32 + lib.infiniopPerTensorDequantFp8.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + + lib.infiniopDestroyPerTensorDequantFp8Descriptor.restype = c_int32 + lib.infiniopDestroyPerTensorDequantFp8Descriptor.argtypes = [ + infiniopOperatorDescriptor_t, + ] + + + @OpRegister.operator def gptq_marlin_gemm_(lib): lib.infiniopCreateGptqMarlinGemmDescriptor.restype = c_int32 diff --git a/test/infiniop/per_tensor_dequant_fp8.py b/test/infiniop/per_tensor_dequant_fp8.py new file mode 100644 index 000000000..123545fcb --- /dev/null +++ b/test/infiniop/per_tensor_dequant_fp8.py @@ -0,0 +1,212 @@ +import torch +import ctypes +from ctypes import c_uint64 +from libinfiniop import ( + LIBINFINIOP, + TestTensor, + get_test_devices, + check_error, + test_operator, + get_args, + debug, + get_tolerance, + profile_operation, + TestWorkspace, + InfiniDtype, + InfiniDtypeNames, + InfiniDeviceNames, + infiniopOperatorDescriptor_t, +) + +# ============================================================================== +# Configuration (Internal Use Only) +# ============================================================================== +# These are not meant to be imported from other modules +_TEST_CASES = [ + # x_shape, x_stride, x_packed_stride, extreme_bytes + ((16, 5632), None, None, False), + ((13, 4), (10, 1), None, False), + ((13, 4), (10, 1), (10, 1), False), + ((16, 5632), (13312, 1), (13312, 1), False), + ((4, 4, 5632), None, None, False), + ((4, 4, 5632), (45056, 5632, 1), (45056, 5632, 1), False), + ((1, 4, 132, 128), (67584, 16896, 128, 1), (67584, 16896, 128, 1), False), + ((1, 4, 132, 128), None, None, False), + # Deterministic byte coverage: 0, min subnormal, max normal, exact max 448, + # NaN/inf-reserved patterns (0x7F/0xFF), and both signs of each. + ((1, 16), None, None, True), +] + +_TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16, InfiniDtype.F32] + +_TOLERANCE_MAP = { + InfiniDtype.F16: {"atol": 1e-3, "rtol": 5e-2}, + InfiniDtype.BF16: {"atol": 1e-3, "rtol": 5e-2}, + InfiniDtype.F32: {"atol": 3e-5, "rtol": 5e-3}, +} + +DEBUG = False +PROFILE = False +NUM_PRERUN = 10 +NUM_ITERATIONS = 1000 + + +def e4m3_to_float(x: torch.Tensor) -> torch.Tensor: + """e4m3fn byte tensor -> float tensor (mirrors the CUDA kernel).""" + b = x.view(torch.uint8) + sign = torch.where((b >> 7) & 1 == 1, -1.0, 1.0) + exp = (b >> 3) & 0x0F + mant = b & 0x07 + val = torch.where( + exp == 0, + mant.float() * (2.0**-9), + # exp is uint8: subtract 7 in float to avoid wrapping for exp < 7 + (1.0 + mant.float() / 8.0) * torch.exp2(exp.float() - 7.0), + ) + return sign * val + + +# 16 deterministic bytes covering every e4m3 decode path (both signs): +# 0 (0x00/0x80), min/max subnormal (0x01/0x07), min normal (0x08), +# all-mantissa normal (0x77 = 240), exp-15 patterns (0x78 = 256), +# exact max 448 (0x7E/0xFE), NaN/inf-reserved patterns (0x7F/0xFF) which the +# kernel decodes as plain values (480/-480). +_EXTREME_BYTES = torch.tensor( + [ + 0x00, + 0x01, + 0x07, + 0x08, + 0x0F, + 0x10, + 0x3F, + 0x40, + 0x77, + 0x78, + 0x7E, + 0x7F, + 0x80, + 0x81, + 0xFF, + 0xFE, + ], + dtype=torch.uint8, +) + + +def per_tensor_dequant_fp8_torch(x_packed, x_scale, dtype): + dq = e4m3_to_float(x_packed) * x_scale.float() + return dq.to(dtype) + + +def test( + handle, + device, + x_shape, + x_stride, + x_packed_stride, + extreme=False, + dtype=InfiniDtype.F16, + sync=None, +): + print( + f"Testing Per Tensor Dequant Fp8 on {InfiniDeviceNames[device]} with x_shape:{x_shape}, x_stride:{x_stride}, x_packed_stride:{x_packed_stride}, extreme:{extreme}, dtype:{InfiniDtypeNames[dtype]}" + ) + + x = TestTensor(x_shape, x_stride, dtype, device, mode="zeros") + + if extreme: + x_packed = TestTensor( + x_shape, + x_packed_stride, + InfiniDtype.F8, + device, + mode="manual", + set_tensor=_EXTREME_BYTES.view(torch.float8_e4m3fn).view(x_shape), + ) + else: + x_packed = TestTensor( + x_shape, + x_packed_stride, + InfiniDtype.F8, + device, + mode="float8_e4m3fn", + ) + x_scale = TestTensor((1,), None, InfiniDtype.F32, device) + if sync is not None: + sync() + + ans = per_tensor_dequant_fp8_torch( + x_packed.torch_tensor(), x_scale.torch_tensor(), x.torch_tensor().dtype + ) + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreatePerTensorDequantFp8Descriptor( + handle, + ctypes.byref(descriptor), + x.descriptor, + x_packed.descriptor, + x_scale.descriptor, + ) + ) + + # Invalidate the shape and strides in the descriptor to prevent them from being directly used by the kernel + x_packed.destroy_desc() + x_scale.destroy_desc() + + workspace_size = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetPerTensorDequantFp8WorkspaceSize( + descriptor, ctypes.byref(workspace_size) + ) + ) + workspace = TestWorkspace(workspace_size.value, x.device) + + def lib_per_tensor_dequant_fp8(): + check_error( + LIBINFINIOP.infiniopPerTensorDequantFp8( + descriptor, + workspace.data(), + workspace_size.value, + x.data(), + x_packed.data(), + x_scale.data(), + None, + ) + ) + + lib_per_tensor_dequant_fp8() + + if sync is not None: + sync() + + atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype) + if DEBUG: + debug(x.actual_tensor().float(), ans.float(), atol=atol, rtol=rtol) + + assert torch.allclose(x.actual_tensor().float(), ans.float(), atol=atol, rtol=rtol) + + # Profiling workflow + if PROFILE: + # fmt: off + profile_operation("PyTorch", lambda: per_tensor_dequant_fp8_torch(x_packed.torch_tensor(), x_scale.torch_tensor(), x.torch_tensor().dtype), device, NUM_PRERUN, NUM_ITERATIONS) + profile_operation(" lib", lambda: lib_per_tensor_dequant_fp8(), device, NUM_PRERUN, NUM_ITERATIONS) + # fmt: on + + check_error(LIBINFINIOP.infiniopDestroyPerTensorDequantFp8Descriptor(descriptor)) + + +if __name__ == "__main__": + args = get_args() + + # Configure testing options + DEBUG = args.debug + PROFILE = args.profile + NUM_PRERUN = args.num_prerun + NUM_ITERATIONS = args.num_iterations + + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + + print("\033[92mTest passed!\033[0m") diff --git a/test/infiniop/per_tensor_quant_fp8.py b/test/infiniop/per_tensor_quant_fp8.py new file mode 100644 index 000000000..dee499bcf --- /dev/null +++ b/test/infiniop/per_tensor_quant_fp8.py @@ -0,0 +1,286 @@ +import torch +import ctypes +from ctypes import c_uint64 +from libinfiniop import ( + LIBINFINIOP, + TestTensor, + get_test_devices, + check_error, + test_operator, + get_args, + debug, + get_tolerance, + profile_operation, + TestWorkspace, + InfiniDtype, + InfiniDtypeNames, + InfiniDeviceNames, + infiniopOperatorDescriptor_t, +) + +# ============================================================================== +# Configuration (Internal Use Only) +# ============================================================================== +# These are not meant to be imported from other modules +# FP8 e4m3fn: symmetric per-tensor quant, scale = max / 448 (dynamic or static) +_TEST_CASES = [ + # x_shape, x_stride, x_packed_stride, is_static, extreme_values + ((16, 5632), None, None, False, False), + ((13, 4), (10, 1), None, True, False), + ((13, 4), (10, 1), (10, 1), False, False), + ((16, 5632), (13312, 1), (13312, 1), True, False), + ((4, 4, 5632), None, None, False, False), + ((4, 4, 5632), (45056, 5632, 1), (45056, 5632, 1), True, False), + ((1, 32, 4, 128), (147456, 4608, 128, 1), (147456, 4608, 128, 1), False, False), + ((1, 32, 4, 128), (16384, 512, 128, 1), (16384, 512, 128, 1), True, False), + # Deterministic boundary coverage: zero, subnormal (< 2^-9), normal, + # carry rounding, exact max 448, and saturation (|v| > 448 -> 0x7E/0xFE). + ((1, 16), None, None, True, True), + ((1, 16), None, None, False, True), +] + +_TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16, InfiniDtype.F32] + +_TOLERANCE_MAP = { + InfiniDtype.F16: {"atol": 1e-3, "rtol": 5e-2}, + InfiniDtype.BF16: {"atol": 1e-3, "rtol": 5e-2}, + InfiniDtype.F32: {"atol": 3e-5, "rtol": 5e-3}, +} + +DEBUG = False +PROFILE = False +NUM_PRERUN = 10 +NUM_ITERATIONS = 1000 + +FP8_E4M3_MAX = 448.0 + +# 16 deterministic values covering every e4m3 rounding path: +# 0, subnormal (0.0025 in [2^-9, 2^-8) -> code 0x01), normal, carry rounding, +# exact max 448, and saturation (> 448). +_EXTREME_VALUES = torch.tensor( + [ + 0.0, + 0.0025, + -0.0025, + 0.5, + -0.5, + 0.25, + -0.25, + 2.0, + 3.5, + 100.0, + -100.0, + 448.0, + -448.0, + 1000.0, + -1000.0, + 7.0, + ], + dtype=torch.float32, +) + + +def e4m3_to_float(x: torch.Tensor) -> torch.Tensor: + """e4m3fn byte tensor -> float tensor (mirrors the CUDA kernel).""" + b = x.view(torch.uint8) + sign = torch.where((b >> 7) & 1 == 1, -1.0, 1.0) + exp = (b >> 3) & 0x0F + mant = b & 0x07 + val = torch.where( + exp == 0, + mant.float() * (2.0**-9), + # exp is uint8: subtract 7 in float to avoid wrapping for exp < 7 + (1.0 + mant.float() / 8.0) * torch.exp2(exp.float() - 7.0), + ) + return sign * val + + +def round_half_away(x: torch.Tensor) -> torch.Tensor: + return torch.floor(x + 0.5) + + +def float_to_e4m3(x: torch.Tensor) -> torch.Tensor: + """float tensor -> e4m3fn byte tensor (mirrors the CUDA kernel).""" + sign = (x < 0).to(torch.uint8) * 0x80 + ax = x.abs() + + sat = ax >= FP8_E4M3_MAX + zero = ax < 2.0**-9 + + # subnormal: value = m * 2^-9 + sub_m = torch.clamp(round_half_away(ax * 512.0), 0, 7).to(torch.uint8) + + # normal + e = torch.floor(torch.log2(ax)) + stored = e + 7.0 + frac = ax * torch.exp2(-e) - 1.0 + m = round_half_away(frac * 8.0) + carry = m == 8.0 + stored2 = torch.where(carry, stored + 1.0, stored) + m2 = torch.where(carry, torch.zeros_like(m), m) + sat2 = stored2 > 15.0 + stored3 = torch.clamp(stored2, 0, 15) + m3 = torch.clamp(m2, 0, 7) + normal_bits = (stored3.to(torch.uint8) << 3) | m3.to(torch.uint8) + + is_sub = (stored <= 0) & ~zero & ~sat + bits = torch.where(is_sub, sub_m, normal_bits) + bits = torch.where(zero, torch.zeros_like(bits), bits) + bits = torch.where(sat | sat2, torch.full_like(bits, 0x7E), bits) + return (sign | bits).to(torch.uint8) + + +def per_tensor_quant_fp8_torch(x, x_scale, is_static): + x = x.float() + if is_static: + scale = x_scale.float() + x_packed = float_to_e4m3(x / scale).view(torch.float8_e4m3fn) + return x_packed, scale + else: + absmax = x.flatten().abs().max() + if absmax == 0: + q = torch.zeros_like(x, dtype=torch.float8_e4m3fn) + return q, torch.tensor(1.0, device=x.device, dtype=torch.float32) + scale = absmax / FP8_E4M3_MAX + x_packed = float_to_e4m3(x / scale).view(torch.float8_e4m3fn) + return x_packed, scale + + +def test( + handle, + device, + x_shape, + x_stride, + x_packed_stride, + is_static, + extreme=False, + dtype=InfiniDtype.F16, + sync=None, +): + + print( + f"Testing Per Tensor Quant Fp8 on {InfiniDeviceNames[device]} with x_shape:{x_shape}, x_stride:{x_stride}, x_packed_stride:{x_packed_stride}, is_static:{is_static}, extreme:{extreme}, dtype:{InfiniDtypeNames[dtype]}" + ) + + if extreme: + torch_dtype = { + InfiniDtype.F16: torch.float16, + InfiniDtype.BF16: torch.bfloat16, + InfiniDtype.F32: torch.float32, + }[dtype] + x = TestTensor( + x_shape, + x_stride, + dtype, + device, + mode="manual", + set_tensor=_EXTREME_VALUES.view(x_shape).to(torch_dtype), + ) + else: + x = TestTensor(x_shape, x_stride, dtype, device) + x_packed = TestTensor( + x_shape, x_packed_stride, InfiniDtype.F8, device, mode="zeros" + ) + if extreme and is_static: + # static scale = 1.0 keeps the raw pattern: |v| / 1.0 = |v|, so values + # beyond 448 exercise the saturation path. + x_scale = TestTensor( + (1,), + None, + InfiniDtype.F32, + device, + mode="manual", + set_tensor=torch.tensor([1.0], dtype=torch.float32), + ) + elif is_static: + x_scale = TestTensor((1,), None, InfiniDtype.F32, device) + else: + x_scale = TestTensor((1,), None, InfiniDtype.F32, device, mode="zeros") + if sync is not None: + sync() + + x_p, x_s = per_tensor_quant_fp8_torch( + x.torch_tensor(), x_scale.torch_tensor(), is_static + ) + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreatePerTensorQuantFp8Descriptor( + handle, + ctypes.byref(descriptor), + x_packed.descriptor, + x_scale.descriptor, + x.descriptor, + ) + ) + + # Invalidate the shape and strides in the descriptor to prevent them from being directly used by the kernel + x_packed.destroy_desc() + x_scale.destroy_desc() + + workspace_size = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetPerTensorQuantFp8WorkspaceSize( + descriptor, ctypes.byref(workspace_size) + ) + ) + workspace = TestWorkspace(workspace_size.value, x.device) + + def lib_per_tensor_quant_fp8(): + check_error( + LIBINFINIOP.infiniopPerTensorQuantFp8( + descriptor, + workspace.data(), + workspace_size.value, + x_packed.data(), + x_scale.data(), + x.data(), + is_static, + None, + ) + ) + + lib_per_tensor_quant_fp8() + + if sync is not None: + sync() + + atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype) + + # Compare in dequantized space: fp8 rounding-mode differences between the + # reference and the kernel are within fp8 granularity. + dq_actual = e4m3_to_float(x_packed.actual_tensor().view(torch.uint8)) * x_scale.actual_tensor() + dq_ref = e4m3_to_float(x_p.view(torch.uint8)) * x_s + if DEBUG: + debug(dq_actual, dq_ref, atol=atol, rtol=rtol) + debug(x_scale.actual_tensor(), x_s, atol=atol, rtol=rtol) + + # Both sides are fp8 round-trips of the same x; differences are at most one + # fp8 LSB (rounding-mode), so use an fp8-appropriate tolerance. + assert torch.allclose(dq_actual, dq_ref, atol=0.02, rtol=0.1) and torch.allclose( + x_scale.actual_tensor(), x_s, atol=atol, rtol=rtol + ) + + # Profiling workflow + if PROFILE: + # fmt: off + profile_operation("PyTorch", lambda: per_tensor_quant_fp8_torch(x.torch_tensor(), x_scale.torch_tensor(), is_static), device, NUM_PRERUN, NUM_ITERATIONS) + profile_operation(" lib", lambda: lib_per_tensor_quant_fp8(), device, NUM_PRERUN, NUM_ITERATIONS) + # fmt: on + + check_error(LIBINFINIOP.infiniopDestroyPerTensorQuantFp8Descriptor(descriptor)) + + +if __name__ == "__main__": + args = get_args() + + # Configure testing options + DEBUG = args.debug + PROFILE = args.profile + NUM_PRERUN = args.num_prerun + NUM_ITERATIONS = args.num_iterations + + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + + print("\033[92mTest passed!\033[0m") diff --git a/test/infiniop/w8a8fp8.py b/test/infiniop/w8a8fp8.py new file mode 100644 index 000000000..b515eef0f --- /dev/null +++ b/test/infiniop/w8a8fp8.py @@ -0,0 +1,343 @@ +import torch +import ctypes +from ctypes import c_uint64 +from enum import Enum, auto +from libinfiniop import ( + LIBINFINIOP, + TestTensor, + get_test_devices, + check_error, + test_operator, + get_args, + debug, + get_tolerance, + TestWorkspace, + InfiniDtype, + InfiniDtypeNames, + InfiniDeviceNames, + infiniopOperatorDescriptor_t, +) + +# ============================================================================== +# Configuration (Internal Use Only) +# ============================================================================== +# These are not meant to be imported from other modules +# W8A8-FP8: activations are quantized to FP8 e4m3 (per-tensor, dynamic +# scale = max/448). Weights are either kept in fp16 (activation-only FP8, +# the KV-cache-style use case) or quantized to FP8 as well (full symmetric +# W8A8-FP8); everything is dequantized and multiplied with an fp16 matmul. + + +class WMode(Enum): + FP16 = auto() # only activations are fp8-quantized; weights stay fp16 + FP8 = auto() # both activations and weights are fp8-quantized + + +_TEST_CASES_ = [ + # x_shape = [M,K], w_shape = [N, K], y_shape = [M, N] + ((100, 3584), (10752, 3584), (100, 10752)), + ((1000, 3584), (10752, 3584), (1000, 10752)), + ((1, 3584), (10752, 3584), (1, 10752)), + ((2000, 3584), (10752, 3584), (2000, 10752)), +] + +_WMODES = [ + WMode.FP16, + WMode.FP8, +] + +_TEST_CASES = [ + test_case + (wmode,) + for test_case in _TEST_CASES_ + for wmode in _WMODES +] + +_TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16] + +# fp8 e4m3 has ~6% relative precision; allow round-trip + rounding-mode slack. +_TOLERANCE_MAP = { + InfiniDtype.F16: {"atol": 3e-1, "rtol": 1.5e-1}, + InfiniDtype.BF16: {"atol": 3e-1, "rtol": 1.5e-1}, +} + +DEBUG = False +PROFILE = False +NUM_PRERUN = 10 +NUM_ITERATIONS = 1000 + +FP8_E4M3_MAX = 448.0 + + +def e4m3_to_float(x: torch.Tensor) -> torch.Tensor: + b = x.view(torch.uint8) + sign = torch.where((b >> 7) & 1 == 1, -1.0, 1.0) + exp = (b >> 3) & 0x0F + mant = b & 0x07 + val = torch.where( + exp == 0, + mant.float() * (2.0**-9), + # exp is uint8: subtract 7 in float to avoid wrapping for exp < 7 + (1.0 + mant.float() / 8.0) * torch.exp2(exp.float() - 7.0), + ) + return sign * val + + +def per_tensor_quant_fp8_torch(x): + """Dynamic per-tensor FP8 quant: returns (packed fp8, scale = max/448).""" + x = x.float() + absmax = x.flatten().abs().max() + if absmax == 0: + scale = torch.tensor(1.0, device=x.device, dtype=torch.float32) + q = torch.zeros_like(x, dtype=torch.float8_e4m3fn) + return q, scale + scale = absmax / FP8_E4M3_MAX + # round half away from zero, then convert via fp16 (values fit fp8 range) + x_q = torch.floor((x / scale).abs() + 0.5) * torch.sign(x) + q = torch.clamp(x_q, -FP8_E4M3_MAX, FP8_E4M3_MAX).to(torch.float8_e4m3fn) + return q, scale + + +def test( + handle, + device, + x_shape, + w_shape, + y_shape, + wmode=WMode.FP8, + dtype=InfiniDtype.BF16, + sync=None, +): + print( + f"Testing W8A8-Fp8 ({wmode.name}) on {InfiniDeviceNames[device]} with x_shape:{x_shape}, w_shape:{w_shape}, dtype:{InfiniDtypeNames[dtype]}" + ) + M, K = x_shape + N = w_shape[0] + out_dtype = torch.float16 if dtype == InfiniDtype.F16 else torch.bfloat16 + + x = TestTensor(x_shape, None, dtype, device) + w = TestTensor(w_shape, None, dtype, device) + + x_packed = TestTensor(x_shape, None, InfiniDtype.F8, device, mode="zeros") + x_scale = TestTensor((1,), None, InfiniDtype.F32, device, mode="zeros") + y = TestTensor(y_shape, None, dtype, device, mode="zeros") + + if wmode == WMode.FP8: + w_packed = TestTensor(w_shape, None, InfiniDtype.F8, device, mode="zeros") + w_scale = TestTensor((1,), None, InfiniDtype.F32, device, mode="zeros") + + # Reference: x is always quantized (dynamic per-tensor fp8); w is quantized + # only in full W8A8 mode, otherwise kept in fp16. + x_p, x_s = per_tensor_quant_fp8_torch(x.torch_tensor()) + if wmode == WMode.FP8: + w_p, w_s = per_tensor_quant_fp8_torch(w.torch_tensor()) + w_ref = e4m3_to_float(w_p) * w_s + else: + w_ref = w.torch_tensor().float() + ref = torch.matmul(e4m3_to_float(x_p) * x_s, w_ref.t()).to(out_dtype) + + # --- per_tensor_quant_fp8 on x --- + quant_x_desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreatePerTensorQuantFp8Descriptor( + handle, + ctypes.byref(quant_x_desc), + x_packed.descriptor, + x_scale.descriptor, + x.descriptor, + ) + ) + + # --- per_tensor_dequant_fp8 for x --- + dequant_x_desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreatePerTensorDequantFp8Descriptor( + handle, + ctypes.byref(dequant_x_desc), + x.descriptor, # output buffer (fp16/bf16) + x_packed.descriptor, + x_scale.descriptor, + ) + ) + + # Invalidate the tensor descriptors now that both ops are created + x_packed.destroy_desc() + x_scale.destroy_desc() + + quant_x_ws = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetPerTensorQuantFp8WorkspaceSize( + quant_x_desc, ctypes.byref(quant_x_ws) + ) + ) + quant_x_ws_t = TestWorkspace(quant_x_ws.value, x.device) + + def lib_quant_x(): + check_error( + LIBINFINIOP.infiniopPerTensorQuantFp8( + quant_x_desc, + quant_x_ws_t.data(), + quant_x_ws.value, + x_packed.data(), + x_scale.data(), + x.data(), + False, # is_static=False -> dynamic + None, + ) + ) + + dequant_x_ws = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetPerTensorDequantFp8WorkspaceSize( + dequant_x_desc, ctypes.byref(dequant_x_ws) + ) + ) + dequant_x_ws_t = TestWorkspace(dequant_x_ws.value, x.device) + + def lib_dequant_x(): + check_error( + LIBINFINIOP.infiniopPerTensorDequantFp8( + dequant_x_desc, + dequant_x_ws_t.data(), + dequant_x_ws.value, + x.data(), # NOTE: x is both input (pre-quant) and dequant output buffer + x_packed.data(), + x_scale.data(), + None, + ) + ) + + if wmode == WMode.FP8: + # --- per_tensor_quant_fp8 on w --- + quant_w_desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreatePerTensorQuantFp8Descriptor( + handle, + ctypes.byref(quant_w_desc), + w_packed.descriptor, + w_scale.descriptor, + w.descriptor, + ) + ) + + # --- per_tensor_dequant_fp8 for w --- + dequant_w_desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreatePerTensorDequantFp8Descriptor( + handle, + ctypes.byref(dequant_w_desc), + w.descriptor, + w_packed.descriptor, + w_scale.descriptor, + ) + ) + + # Invalidate the tensor descriptors now that both ops are created + w_packed.destroy_desc() + w_scale.destroy_desc() + + quant_w_ws = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetPerTensorQuantFp8WorkspaceSize( + quant_w_desc, ctypes.byref(quant_w_ws) + ) + ) + quant_w_ws_t = TestWorkspace(quant_w_ws.value, w.device) + + def lib_quant_w(): + check_error( + LIBINFINIOP.infiniopPerTensorQuantFp8( + quant_w_desc, + quant_w_ws_t.data(), + quant_w_ws.value, + w_packed.data(), + w_scale.data(), + w.data(), + False, + None, + ) + ) + + dequant_w_ws = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetPerTensorDequantFp8WorkspaceSize( + dequant_w_desc, ctypes.byref(dequant_w_ws) + ) + ) + dequant_w_ws_t = TestWorkspace(dequant_w_ws.value, w.device) + + def lib_dequant_w(): + check_error( + LIBINFINIOP.infiniopPerTensorDequantFp8( + dequant_w_desc, + dequant_w_ws_t.data(), + dequant_w_ws.value, + w.data(), + w_packed.data(), + w_scale.data(), + None, + ) + ) + + def lib_w8a8fp8(): + # quant x (and, in full W8A8 mode, w) to dynamic per-tensor fp8, + # dequant them back into the original buffers, then matmul on the + # dequantized values. + lib_quant_x() + lib_dequant_x() + if wmode == WMode.FP8: + lib_quant_w() + lib_dequant_w() + + lib_w8a8fp8() + + if sync is not None: + sync() + + # Final matmul on the dequantized values (fp16/bf16) + y_actual = torch.matmul( + x.torch_tensor().float(), w.torch_tensor().float().t() + ).to(out_dtype) + + atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype) + if DEBUG: + debug(y_actual, ref, atol=atol, rtol=rtol) + + assert torch.allclose(y_actual.float(), ref.float(), atol=atol, rtol=rtol) + + # Profiling workflow + if PROFILE: + def profile_operation(name, func, device, num_prerun, num_iterations): + for _ in range(num_prerun): + func() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(num_iterations): + func() + end.record() + torch.cuda.synchronize() + elapsed = start.elapsed_time(end) + print(f"{name} took {elapsed / num_iterations:.6f} ms over {num_iterations} iterations") + + profile_operation("lib w8a8fp8", lambda: lib_w8a8fp8(), device, NUM_PRERUN, NUM_ITERATIONS) + + check_error(LIBINFINIOP.infiniopDestroyPerTensorQuantFp8Descriptor(quant_x_desc)) + check_error(LIBINFINIOP.infiniopDestroyPerTensorDequantFp8Descriptor(dequant_x_desc)) + if wmode == WMode.FP8: + check_error(LIBINFINIOP.infiniopDestroyPerTensorQuantFp8Descriptor(quant_w_desc)) + check_error(LIBINFINIOP.infiniopDestroyPerTensorDequantFp8Descriptor(dequant_w_desc)) + + +if __name__ == "__main__": + args = get_args() + + DEBUG = args.debug + PROFILE = args.profile + NUM_PRERUN = args.num_prerun + NUM_ITERATIONS = args.num_iterations + + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + + print("\033[92mTest passed!\033[0m") From 036fb7061ee38a2734482b4eef1aeab279ac4ea2 Mon Sep 17 00:00:00 2001 From: bobodai Date: Sun, 23 Aug 2026 22:56:44 +0800 Subject: [PATCH 2/2] fix: guard per-tensor FP8 quant against zero scale (all-zero input) --- .../per_tensor_quant_fp8/cuda/kernel.cuh | 17 ++++++++-- test/infiniop/per_tensor_quant_fp8.py | 32 +++++++++++-------- 2 files changed, 34 insertions(+), 15 deletions(-) diff --git a/src/infiniop/ops/quant/per_tensor_quant_fp8/cuda/kernel.cuh b/src/infiniop/ops/quant/per_tensor_quant_fp8/cuda/kernel.cuh index fd2fbd0fe..bf93ca865 100644 --- a/src/infiniop/ops/quant/per_tensor_quant_fp8/cuda/kernel.cuh +++ b/src/infiniop/ops/quant/per_tensor_quant_fp8/cuda/kernel.cuh @@ -126,7 +126,18 @@ __device__ void perTensorQuantFp8SymKernel( unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x; const int grid_size = blockDim.x * gridDim.x; - float scale_val = 1.0f / x_scale[0]; + // Guard against a zero scale: all-zero input (dynamic mode, max|x| = 0) or + // a zero static scale. 1.0f / 0.0f = inf would turn 0 * inf into NaN and + // produce garbage bytes. When scale == 0 we skip the division entirely and + // write all-zero fp8 bytes (equivalent to memsetting the output), and + // normalize the scale to 1.0, matching the reference semantics + // (absmax == 0 -> scale = 1). This also covers inference engines feeding + // empty/all-zero sequences, avoiding a NaN storm downstream. + const bool zero_scale = (x_scale[0] == 0.0f); + if (zero_scale && threadIdx.x == 0 && blockIdx.x == 0) { + x_scale[0] = 1.0f; + } + const float scale_val = zero_scale ? 0.0f : 1.0f / x_scale[0]; for (int ind = gid; ind < num_elements; ind += grid_size) { int tid = ind; @@ -145,7 +156,9 @@ __device__ void perTensorQuantFp8SymKernel( int p_index = w * (int)p_strides_3 + h * (int)p_strides_2 + c * (int)p_strides_1 + b * (int)p_strides_0; float qf = (float)x[index] * scale_val; - x_packed[p_index] = float_to_e4m3(qf); + // zero_scale: write 0x00 directly, skipping float_to_e4m3 entirely + // (immune to NaN inputs; equivalent to memsetting the output). + x_packed[p_index] = zero_scale ? 0 : float_to_e4m3(qf); } } diff --git a/test/infiniop/per_tensor_quant_fp8.py b/test/infiniop/per_tensor_quant_fp8.py index dee499bcf..7e0352554 100644 --- a/test/infiniop/per_tensor_quant_fp8.py +++ b/test/infiniop/per_tensor_quant_fp8.py @@ -24,19 +24,22 @@ # These are not meant to be imported from other modules # FP8 e4m3fn: symmetric per-tensor quant, scale = max / 448 (dynamic or static) _TEST_CASES = [ - # x_shape, x_stride, x_packed_stride, is_static, extreme_values - ((16, 5632), None, None, False, False), - ((13, 4), (10, 1), None, True, False), - ((13, 4), (10, 1), (10, 1), False, False), - ((16, 5632), (13312, 1), (13312, 1), True, False), - ((4, 4, 5632), None, None, False, False), - ((4, 4, 5632), (45056, 5632, 1), (45056, 5632, 1), True, False), - ((1, 32, 4, 128), (147456, 4608, 128, 1), (147456, 4608, 128, 1), False, False), - ((1, 32, 4, 128), (16384, 512, 128, 1), (16384, 512, 128, 1), True, False), + # x_shape, x_stride, x_packed_stride, is_static, extreme_values, all_zero + ((16, 5632), None, None, False, False, False), + ((13, 4), (10, 1), None, True, False, False), + ((13, 4), (10, 1), (10, 1), False, False, False), + ((16, 5632), (13312, 1), (13312, 1), True, False, False), + ((4, 4, 5632), None, None, False, False, False), + ((4, 4, 5632), (45056, 5632, 1), (45056, 5632, 1), True, False, False), + ((1, 32, 4, 128), (147456, 4608, 128, 1), (147456, 4608, 128, 1), False, False, False), + ((1, 32, 4, 128), (16384, 512, 128, 1), (16384, 512, 128, 1), True, False, False), # Deterministic boundary coverage: zero, subnormal (< 2^-9), normal, # carry rounding, exact max 448, and saturation (|v| > 448 -> 0x7E/0xFE). - ((1, 16), None, None, True, True), - ((1, 16), None, None, False, True), + ((1, 16), None, None, True, True, False), + ((1, 16), None, None, False, True, False), + # All-zero input (dynamic): scale = max|x| / 448 = 0 must not divide by + # zero; expect scale normalized to 1.0 and all-zero fp8 output. + ((1, 16), None, None, False, False, True), ] _TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16, InfiniDtype.F32] @@ -154,15 +157,18 @@ def test( x_packed_stride, is_static, extreme=False, + all_zero=False, dtype=InfiniDtype.F16, sync=None, ): print( - f"Testing Per Tensor Quant Fp8 on {InfiniDeviceNames[device]} with x_shape:{x_shape}, x_stride:{x_stride}, x_packed_stride:{x_packed_stride}, is_static:{is_static}, extreme:{extreme}, dtype:{InfiniDtypeNames[dtype]}" + f"Testing Per Tensor Quant Fp8 on {InfiniDeviceNames[device]} with x_shape:{x_shape}, x_stride:{x_stride}, x_packed_stride:{x_packed_stride}, is_static:{is_static}, extreme:{extreme}, all_zero:{all_zero}, dtype:{InfiniDtypeNames[dtype]}" ) - if extreme: + if all_zero: + x = TestTensor(x_shape, x_stride, dtype, device, mode="zeros") + elif extreme: torch_dtype = { InfiniDtype.F16: torch.float16, InfiniDtype.BF16: torch.bfloat16,