From 6558c50dcfea6511f39d724890cef563d9c0cc20 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Wed, 8 Jul 2026 09:21:31 +0800 Subject: [PATCH] feat: add DSV4 Hygon ops Added operators: - dsv4_per_token_quant_int8 - dsv4_rmsnorm_self - dsv4_silu_and_mul - dsv4_fused_rope - dsv4_topk_transform - dsv4_mask_topk_ids - dsv4_linear_bf16_fp32 - dsv4_silu_mul_quant - dsv4_add_rmsnorm_quant - dsv4_act_quant_fp8 - dsv4_swa_prefill_indices --- include/infinicore/ops.hpp | 11 + include/infinicore/ops/dsv4_act_quant_fp8.hpp | 9 + .../infinicore/ops/dsv4_add_rmsnorm_quant.hpp | 9 + include/infinicore/ops/dsv4_fused_rope.hpp | 15 + .../infinicore/ops/dsv4_linear_bf16_fp32.hpp | 15 + include/infinicore/ops/dsv4_mask_topk_ids.hpp | 14 + .../ops/dsv4_per_token_quant_int8.hpp | 17 + include/infinicore/ops/dsv4_rmsnorm_self.hpp | 15 + include/infinicore/ops/dsv4_silu_and_mul.hpp | 15 + .../infinicore/ops/dsv4_silu_mul_quant.hpp | 9 + .../ops/dsv4_swa_prefill_indices.hpp | 9 + .../infinicore/ops/dsv4_topk_transform.hpp | 15 + include/infiniop.h | 11 + include/infiniop/ops/dsv4_act_quant_fp8.h | 9 + include/infiniop/ops/dsv4_add_rmsnorm_quant.h | 9 + include/infiniop/ops/dsv4_fused_rope.h | 13 + include/infiniop/ops/dsv4_linear_bf16_fp32.h | 31 ++ include/infiniop/ops/dsv4_mask_topk_ids.h | 29 ++ .../infiniop/ops/dsv4_per_token_quant_int8.h | 31 ++ include/infiniop/ops/dsv4_rmsnorm_self.h | 18 + include/infiniop/ops/dsv4_silu_and_mul.h | 13 + include/infiniop/ops/dsv4_silu_mul_quant.h | 9 + .../infiniop/ops/dsv4_swa_prefill_indices.h | 9 + include/infiniop/ops/dsv4_topk_transform.h | 13 + .../dsv4_act_quant_fp8/dsv4_act_quant_fp8.cc | 13 + .../dsv4_act_quant_fp8_infiniop.cc | 25 ++ .../dsv4_add_rmsnorm_quant.cc | 13 + .../dsv4_add_rmsnorm_quant_infiniop.cc | 25 ++ .../ops/dsv4_fused_rope/dsv4_fused_rope.cc | 32 ++ .../dsv4_fused_rope_infiniop.cc | 59 ++++ .../dsv4_linear_bf16_fp32.cc | 30 ++ .../dsv4_linear_bf16_fp32_infiniop.cc | 42 +++ .../dsv4_mask_topk_ids/dsv4_mask_topk_ids.cc | 21 ++ .../dsv4_mask_topk_ids_infiniop.cc | 41 +++ .../dsv4_per_token_quant_int8.cc | 39 +++ .../dsv4_per_token_quant_int8_infiniop.cc | 42 +++ .../dsv4_rmsnorm_self/dsv4_rmsnorm_self.cc | 27 ++ .../dsv4_rmsnorm_self_infiniop.cc | 33 ++ .../dsv4_silu_and_mul/dsv4_silu_and_mul.cc | 27 ++ .../dsv4_silu_and_mul_infiniop.cc | 33 ++ .../dsv4_silu_mul_quant.cc | 13 + .../dsv4_silu_mul_quant_infiniop.cc | 25 ++ .../dsv4_swa_prefill_indices.cc | 12 + .../dsv4_swa_prefill_indices_infiniop.cc | 25 ++ .../dsv4_topk_transform.cc | 29 ++ .../dsv4_topk_transform_infiniop.cc | 33 ++ .../dsv4_act_quant_fp8/dsv4_act_quant_fp8.h | 19 ++ src/infiniop/ops/dsv4_act_quant_fp8/info.h | 24 ++ .../nvidia/dsv4_act_quant_fp8_nvidia.cu | 107 ++++++ .../nvidia/dsv4_act_quant_fp8_nvidia.cuh | 5 + .../ops/dsv4_act_quant_fp8/operator.cc | 108 ++++++ .../dsv4_add_rmsnorm_quant.h | 21 ++ .../ops/dsv4_add_rmsnorm_quant/info.h | 24 ++ .../nvidia/dsv4_add_rmsnorm_quant_nvidia.cu | 134 ++++++++ .../nvidia/dsv4_add_rmsnorm_quant_nvidia.cuh | 5 + .../ops/dsv4_add_rmsnorm_quant/operator.cc | 108 ++++++ .../ops/dsv4_fused_rope/dsv4_fused_rope.h | 19 ++ src/infiniop/ops/dsv4_fused_rope/info.h | 33 ++ .../nvidia/dsv4_fused_rope_nvidia.cu | 68 ++++ .../nvidia/dsv4_fused_rope_nvidia.cuh | 5 + src/infiniop/ops/dsv4_fused_rope/operator.cc | 112 ++++++ .../dsv4_linear_bf16_fp32.h | 26 ++ src/infiniop/ops/dsv4_linear_bf16_fp32/info.h | 36 ++ .../nvidia/dsv4_linear_bf16_fp32_nvidia.cu | 112 ++++++ .../nvidia/dsv4_linear_bf16_fp32_nvidia.cuh | 8 + .../ops/dsv4_linear_bf16_fp32/operator.cc | 134 ++++++++ .../dsv4_mask_topk_ids/dsv4_mask_topk_ids.h | 23 ++ src/infiniop/ops/dsv4_mask_topk_ids/info.h | 30 ++ .../nvidia/dsv4_mask_topk_ids_nvidia.cu | 55 +++ .../nvidia/dsv4_mask_topk_ids_nvidia.cuh | 8 + .../ops/dsv4_mask_topk_ids/operator.cc | 131 +++++++ .../dsv4_per_token_quant_int8.h | 22 ++ .../ops/dsv4_per_token_quant_int8/info.h | 40 +++ .../dsv4_per_token_quant_int8_nvidia.cu | 153 +++++++++ .../dsv4_per_token_quant_int8_nvidia.cuh | 5 + .../ops/dsv4_per_token_quant_int8/operator.cc | 117 +++++++ .../ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self.h | 21 ++ src/infiniop/ops/dsv4_rmsnorm_self/info.h | 31 ++ .../nvidia/dsv4_rmsnorm_self_nvidia.cu | 141 ++++++++ .../nvidia/dsv4_rmsnorm_self_nvidia.cuh | 5 + .../ops/dsv4_rmsnorm_self/operator.cc | 112 ++++++ .../ops/dsv4_silu_and_mul/dsv4_silu_and_mul.h | 19 ++ src/infiniop/ops/dsv4_silu_and_mul/info.h | 31 ++ .../nvidia/dsv4_silu_and_mul_nvidia.cu | 88 +++++ .../nvidia/dsv4_silu_and_mul_nvidia.cuh | 5 + .../ops/dsv4_silu_and_mul/operator.cc | 112 ++++++ .../dsv4_silu_mul_quant/dsv4_silu_mul_quant.h | 19 ++ src/infiniop/ops/dsv4_silu_mul_quant/info.h | 24 ++ .../nvidia/dsv4_silu_mul_quant_nvidia.cu | 97 ++++++ .../nvidia/dsv4_silu_mul_quant_nvidia.cuh | 5 + .../ops/dsv4_silu_mul_quant/operator.cc | 108 ++++++ .../dsv4_swa_prefill_indices.h | 19 ++ .../ops/dsv4_swa_prefill_indices/info.h | 19 ++ .../nvidia/dsv4_swa_prefill_indices_nvidia.cu | 29 ++ .../dsv4_swa_prefill_indices_nvidia.cuh | 5 + .../ops/dsv4_swa_prefill_indices/operator.cc | 108 ++++++ .../dsv4_topk_transform/dsv4_topk_transform.h | 19 ++ src/infiniop/ops/dsv4_topk_transform/info.h | 23 ++ .../nvidia/dsv4_topk_transform_nvidia.cu | 32 ++ .../nvidia/dsv4_topk_transform_nvidia.cuh | 5 + .../ops/dsv4_topk_transform/operator.cc | 112 ++++++ test/infiniop/_dsv4_common.py | 247 ++++++++++++++ test/infiniop/dsv4_act_quant_fp8.py | 72 ++++ test/infiniop/dsv4_add_rmsnorm_quant.py | 94 +++++ test/infiniop/dsv4_fused_rope.py | 14 + test/infiniop/dsv4_linear_bf16_fp32.py | 64 ++++ test/infiniop/dsv4_mask_topk_ids.py | 66 ++++ test/infiniop/dsv4_per_token_quant_int8.py | 14 + test/infiniop/dsv4_rmsnorm_self.py | 14 + test/infiniop/dsv4_silu_and_mul.py | 14 + test/infiniop/dsv4_silu_mul_quant.py | 77 +++++ test/infiniop/dsv4_swa_prefill_indices.py | 62 ++++ test/infiniop/dsv4_topk_transform.py | 14 + test/infiniop/libinfiniop/op_register.py | 323 ++++++++++++++++++ 114 files changed, 4943 insertions(+) create mode 100644 include/infinicore/ops/dsv4_act_quant_fp8.hpp create mode 100644 include/infinicore/ops/dsv4_add_rmsnorm_quant.hpp create mode 100644 include/infinicore/ops/dsv4_fused_rope.hpp create mode 100644 include/infinicore/ops/dsv4_linear_bf16_fp32.hpp create mode 100644 include/infinicore/ops/dsv4_mask_topk_ids.hpp create mode 100644 include/infinicore/ops/dsv4_per_token_quant_int8.hpp create mode 100644 include/infinicore/ops/dsv4_rmsnorm_self.hpp create mode 100644 include/infinicore/ops/dsv4_silu_and_mul.hpp create mode 100644 include/infinicore/ops/dsv4_silu_mul_quant.hpp create mode 100644 include/infinicore/ops/dsv4_swa_prefill_indices.hpp create mode 100644 include/infinicore/ops/dsv4_topk_transform.hpp create mode 100644 include/infiniop/ops/dsv4_act_quant_fp8.h create mode 100644 include/infiniop/ops/dsv4_add_rmsnorm_quant.h create mode 100644 include/infiniop/ops/dsv4_fused_rope.h create mode 100644 include/infiniop/ops/dsv4_linear_bf16_fp32.h create mode 100644 include/infiniop/ops/dsv4_mask_topk_ids.h create mode 100644 include/infiniop/ops/dsv4_per_token_quant_int8.h create mode 100644 include/infiniop/ops/dsv4_rmsnorm_self.h create mode 100644 include/infiniop/ops/dsv4_silu_and_mul.h create mode 100644 include/infiniop/ops/dsv4_silu_mul_quant.h create mode 100644 include/infiniop/ops/dsv4_swa_prefill_indices.h create mode 100644 include/infiniop/ops/dsv4_topk_transform.h create mode 100644 src/infinicore/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8.cc create mode 100644 src/infinicore/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant.cc create mode 100644 src/infinicore/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_fused_rope/dsv4_fused_rope.cc create mode 100644 src/infinicore/ops/dsv4_fused_rope/dsv4_fused_rope_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32.cc create mode 100644 src/infinicore/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids.cc create mode 100644 src/infinicore/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8.cc create mode 100644 src/infinicore/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self.cc create mode 100644 src/infinicore/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_silu_and_mul/dsv4_silu_and_mul.cc create mode 100644 src/infinicore/ops/dsv4_silu_and_mul/dsv4_silu_and_mul_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant.cc create mode 100644 src/infinicore/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices.cc create mode 100644 src/infinicore/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices_infiniop.cc create mode 100644 src/infinicore/ops/dsv4_topk_transform/dsv4_topk_transform.cc create mode 100644 src/infinicore/ops/dsv4_topk_transform/dsv4_topk_transform_infiniop.cc create mode 100644 src/infiniop/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8.h create mode 100644 src/infiniop/ops/dsv4_act_quant_fp8/info.h create mode 100644 src/infiniop/ops/dsv4_act_quant_fp8/nvidia/dsv4_act_quant_fp8_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_act_quant_fp8/nvidia/dsv4_act_quant_fp8_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_act_quant_fp8/operator.cc create mode 100644 src/infiniop/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant.h create mode 100644 src/infiniop/ops/dsv4_add_rmsnorm_quant/info.h create mode 100644 src/infiniop/ops/dsv4_add_rmsnorm_quant/nvidia/dsv4_add_rmsnorm_quant_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_add_rmsnorm_quant/nvidia/dsv4_add_rmsnorm_quant_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_add_rmsnorm_quant/operator.cc create mode 100644 src/infiniop/ops/dsv4_fused_rope/dsv4_fused_rope.h create mode 100644 src/infiniop/ops/dsv4_fused_rope/info.h create mode 100644 src/infiniop/ops/dsv4_fused_rope/nvidia/dsv4_fused_rope_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_fused_rope/nvidia/dsv4_fused_rope_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_fused_rope/operator.cc create mode 100644 src/infiniop/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32.h create mode 100644 src/infiniop/ops/dsv4_linear_bf16_fp32/info.h create mode 100644 src/infiniop/ops/dsv4_linear_bf16_fp32/nvidia/dsv4_linear_bf16_fp32_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_linear_bf16_fp32/nvidia/dsv4_linear_bf16_fp32_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_linear_bf16_fp32/operator.cc create mode 100644 src/infiniop/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids.h create mode 100644 src/infiniop/ops/dsv4_mask_topk_ids/info.h create mode 100644 src/infiniop/ops/dsv4_mask_topk_ids/nvidia/dsv4_mask_topk_ids_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_mask_topk_ids/nvidia/dsv4_mask_topk_ids_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_mask_topk_ids/operator.cc create mode 100644 src/infiniop/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8.h create mode 100644 src/infiniop/ops/dsv4_per_token_quant_int8/info.h create mode 100644 src/infiniop/ops/dsv4_per_token_quant_int8/nvidia/dsv4_per_token_quant_int8_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_per_token_quant_int8/nvidia/dsv4_per_token_quant_int8_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_per_token_quant_int8/operator.cc create mode 100644 src/infiniop/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self.h create mode 100644 src/infiniop/ops/dsv4_rmsnorm_self/info.h create mode 100644 src/infiniop/ops/dsv4_rmsnorm_self/nvidia/dsv4_rmsnorm_self_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_rmsnorm_self/nvidia/dsv4_rmsnorm_self_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_rmsnorm_self/operator.cc create mode 100644 src/infiniop/ops/dsv4_silu_and_mul/dsv4_silu_and_mul.h create mode 100644 src/infiniop/ops/dsv4_silu_and_mul/info.h create mode 100644 src/infiniop/ops/dsv4_silu_and_mul/nvidia/dsv4_silu_and_mul_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_silu_and_mul/nvidia/dsv4_silu_and_mul_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_silu_and_mul/operator.cc create mode 100644 src/infiniop/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant.h create mode 100644 src/infiniop/ops/dsv4_silu_mul_quant/info.h create mode 100644 src/infiniop/ops/dsv4_silu_mul_quant/nvidia/dsv4_silu_mul_quant_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_silu_mul_quant/nvidia/dsv4_silu_mul_quant_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_silu_mul_quant/operator.cc create mode 100644 src/infiniop/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices.h create mode 100644 src/infiniop/ops/dsv4_swa_prefill_indices/info.h create mode 100644 src/infiniop/ops/dsv4_swa_prefill_indices/nvidia/dsv4_swa_prefill_indices_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_swa_prefill_indices/nvidia/dsv4_swa_prefill_indices_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_swa_prefill_indices/operator.cc create mode 100644 src/infiniop/ops/dsv4_topk_transform/dsv4_topk_transform.h create mode 100644 src/infiniop/ops/dsv4_topk_transform/info.h create mode 100644 src/infiniop/ops/dsv4_topk_transform/nvidia/dsv4_topk_transform_nvidia.cu create mode 100644 src/infiniop/ops/dsv4_topk_transform/nvidia/dsv4_topk_transform_nvidia.cuh create mode 100644 src/infiniop/ops/dsv4_topk_transform/operator.cc create mode 100644 test/infiniop/_dsv4_common.py create mode 100644 test/infiniop/dsv4_act_quant_fp8.py create mode 100644 test/infiniop/dsv4_add_rmsnorm_quant.py create mode 100644 test/infiniop/dsv4_fused_rope.py create mode 100644 test/infiniop/dsv4_linear_bf16_fp32.py create mode 100644 test/infiniop/dsv4_mask_topk_ids.py create mode 100644 test/infiniop/dsv4_per_token_quant_int8.py create mode 100644 test/infiniop/dsv4_rmsnorm_self.py create mode 100644 test/infiniop/dsv4_silu_and_mul.py create mode 100644 test/infiniop/dsv4_silu_mul_quant.py create mode 100644 test/infiniop/dsv4_swa_prefill_indices.py create mode 100644 test/infiniop/dsv4_topk_transform.py diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index b5c4ff18f..b36c6213b 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -25,6 +25,17 @@ #include "ops/conv2d.hpp" #include "ops/cross_entropy.hpp" #include "ops/deepseek_moe.hpp" +#include "ops/dsv4_act_quant_fp8.hpp" +#include "ops/dsv4_add_rmsnorm_quant.hpp" +#include "ops/dsv4_fused_rope.hpp" +#include "ops/dsv4_linear_bf16_fp32.hpp" +#include "ops/dsv4_mask_topk_ids.hpp" +#include "ops/dsv4_per_token_quant_int8.hpp" +#include "ops/dsv4_rmsnorm_self.hpp" +#include "ops/dsv4_silu_and_mul.hpp" +#include "ops/dsv4_silu_mul_quant.hpp" +#include "ops/dsv4_swa_prefill_indices.hpp" +#include "ops/dsv4_topk_transform.hpp" #include "ops/embedding.hpp" #include "ops/flash_attention.hpp" #include "ops/fmin.hpp" diff --git a/include/infinicore/ops/dsv4_act_quant_fp8.hpp b/include/infinicore/ops/dsv4_act_quant_fp8.hpp new file mode 100644 index 000000000..1b4f7ca1d --- /dev/null +++ b/include/infinicore/ops/dsv4_act_quant_fp8.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" +namespace infinicore::op { +INFINICORE_GRAPH_OP_CLASS(Dsv4ActQuantFp8, Tensor, Tensor, const Tensor &, float); +void dsv4_act_quant_fp8_(Tensor xq, Tensor scale, const Tensor &x, float fp8_max); +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_add_rmsnorm_quant.hpp b/include/infinicore/ops/dsv4_add_rmsnorm_quant.hpp new file mode 100644 index 000000000..4a9bf38a1 --- /dev/null +++ b/include/infinicore/ops/dsv4_add_rmsnorm_quant.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" +namespace infinicore::op { +INFINICORE_GRAPH_OP_CLASS(Dsv4AddRMSNormQuant, Tensor, Tensor, Tensor, const Tensor &, const Tensor &, float); +void dsv4_add_rmsnorm_quant_(Tensor res, Tensor q, Tensor scale, const Tensor &x, const Tensor &weight, float epsilon); +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_fused_rope.hpp b/include/infinicore/ops/dsv4_fused_rope.hpp new file mode 100644 index 000000000..ad856c7ed --- /dev/null +++ b/include/infinicore/ops/dsv4_fused_rope.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Dsv4FusedRope, Tensor, Tensor, const Tensor &, const Tensor &, bool); + +void dsv4_fused_rope_(Tensor q, Tensor k, const Tensor &freq_real, const Tensor &freq_imag, bool has_k = false); +Tensor dsv4_fused_rope(const Tensor &q, const Tensor &freq_real, const Tensor &freq_imag); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_linear_bf16_fp32.hpp b/include/infinicore/ops/dsv4_linear_bf16_fp32.hpp new file mode 100644 index 000000000..e04c28357 --- /dev/null +++ b/include/infinicore/ops/dsv4_linear_bf16_fp32.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Dsv4LinearBf16Fp32, Tensor, const Tensor &, const Tensor &); + +Tensor dsv4_linear_bf16_fp32(const Tensor &x, const Tensor &w); +void dsv4_linear_bf16_fp32_(Tensor y, const Tensor &x, const Tensor &w); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_mask_topk_ids.hpp b/include/infinicore/ops/dsv4_mask_topk_ids.hpp new file mode 100644 index 000000000..1eb12d9aa --- /dev/null +++ b/include/infinicore/ops/dsv4_mask_topk_ids.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Dsv4MaskTopkIds, Tensor, const Tensor &); + +void dsv4_mask_topk_ids_(Tensor topk_ids, const Tensor &num_token_non_padded); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_per_token_quant_int8.hpp b/include/infinicore/ops/dsv4_per_token_quant_int8.hpp new file mode 100644 index 000000000..0fa9ec04a --- /dev/null +++ b/include/infinicore/ops/dsv4_per_token_quant_int8.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Dsv4PerTokenQuantInt8, Tensor, Tensor, const Tensor &); + +std::pair dsv4_per_token_quant_int8(const Tensor &x); +void dsv4_per_token_quant_int8_(Tensor q, Tensor scale, const Tensor &x); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_rmsnorm_self.hpp b/include/infinicore/ops/dsv4_rmsnorm_self.hpp new file mode 100644 index 000000000..177cf69b6 --- /dev/null +++ b/include/infinicore/ops/dsv4_rmsnorm_self.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Dsv4RMSNormSelf, Tensor, const Tensor &, float); + +Tensor dsv4_rmsnorm_self(const Tensor &x, float epsilon = 1e-6f); +void dsv4_rmsnorm_self_(Tensor y, const Tensor &x, float epsilon = 1e-6f); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_silu_and_mul.hpp b/include/infinicore/ops/dsv4_silu_and_mul.hpp new file mode 100644 index 000000000..b2ce429e1 --- /dev/null +++ b/include/infinicore/ops/dsv4_silu_and_mul.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Dsv4SiluAndMul, Tensor, const Tensor &, const Tensor &); + +Tensor dsv4_silu_and_mul(const Tensor &gate, const Tensor &up); +void dsv4_silu_and_mul_(Tensor y, const Tensor &gate, const Tensor &up); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_silu_mul_quant.hpp b/include/infinicore/ops/dsv4_silu_mul_quant.hpp new file mode 100644 index 000000000..238dcbb13 --- /dev/null +++ b/include/infinicore/ops/dsv4_silu_mul_quant.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" +namespace infinicore::op { +INFINICORE_GRAPH_OP_CLASS(Dsv4SiluMulQuant, Tensor, Tensor, const Tensor &, const Tensor &); +void dsv4_silu_mul_quant_(Tensor q, Tensor scale, const Tensor &gate, const Tensor &up); +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_swa_prefill_indices.hpp b/include/infinicore/ops/dsv4_swa_prefill_indices.hpp new file mode 100644 index 000000000..8693e6331 --- /dev/null +++ b/include/infinicore/ops/dsv4_swa_prefill_indices.hpp @@ -0,0 +1,9 @@ +#pragma once +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" +namespace infinicore::op { +INFINICORE_GRAPH_OP_CLASS(Dsv4SwaPrefillIndices, Tensor, int); +void dsv4_swa_prefill_indices_(Tensor indices, int window_size); +} // namespace infinicore::op diff --git a/include/infinicore/ops/dsv4_topk_transform.hpp b/include/infinicore/ops/dsv4_topk_transform.hpp new file mode 100644 index 000000000..85b453064 --- /dev/null +++ b/include/infinicore/ops/dsv4_topk_transform.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(Dsv4TopkTransform, Tensor, const Tensor &, const Tensor &, const Tensor &, int); + +Tensor dsv4_topk_transform(const Tensor &scores, const Tensor &seq_lens, const Tensor &page_tables, int page_size = 64); +void dsv4_topk_transform_(Tensor out, const Tensor &scores, const Tensor &seq_lens, const Tensor &page_tables, int page_size = 64); + +} // namespace infinicore::op diff --git a/include/infiniop.h b/include/infiniop.h index 7794dfc5c..a7bc25656 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -44,6 +44,17 @@ #include "infiniop/ops/digamma.h" #include "infiniop/ops/dist.h" #include "infiniop/ops/dot.h" +#include "infiniop/ops/dsv4_act_quant_fp8.h" +#include "infiniop/ops/dsv4_add_rmsnorm_quant.h" +#include "infiniop/ops/dsv4_fused_rope.h" +#include "infiniop/ops/dsv4_linear_bf16_fp32.h" +#include "infiniop/ops/dsv4_mask_topk_ids.h" +#include "infiniop/ops/dsv4_per_token_quant_int8.h" +#include "infiniop/ops/dsv4_rmsnorm_self.h" +#include "infiniop/ops/dsv4_silu_and_mul.h" +#include "infiniop/ops/dsv4_silu_mul_quant.h" +#include "infiniop/ops/dsv4_swa_prefill_indices.h" +#include "infiniop/ops/dsv4_topk_transform.h" #include "infiniop/ops/embedding.h" #include "infiniop/ops/equal.h" #include "infiniop/ops/erf.h" diff --git a/include/infiniop/ops/dsv4_act_quant_fp8.h b/include/infiniop/ops/dsv4_act_quant_fp8.h new file mode 100644 index 000000000..991d6623b --- /dev/null +++ b/include/infiniop/ops/dsv4_act_quant_fp8.h @@ -0,0 +1,9 @@ +#ifndef __INFINIOP_DSV4_ACT_QUANT_FP8_API_H__ +#define __INFINIOP_DSV4_ACT_QUANT_FP8_API_H__ +#include "../operator_descriptor.h" +typedef struct InfiniopDescriptor *infiniopDsv4ActQuantFp8Descriptor_t; +__INFINI_C __export infiniStatus_t infiniopCreateDsv4ActQuantFp8Descriptor(infiniopHandle_t, infiniopDsv4ActQuantFp8Descriptor_t *, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, float); +__INFINI_C __export infiniStatus_t infiniopGetDsv4ActQuantFp8WorkspaceSize(infiniopDsv4ActQuantFp8Descriptor_t, size_t *); +__INFINI_C __export infiniStatus_t infiniopDsv4ActQuantFp8(infiniopDsv4ActQuantFp8Descriptor_t, void *, size_t, void *, void *, const void *, void *); +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4ActQuantFp8Descriptor(infiniopDsv4ActQuantFp8Descriptor_t); +#endif diff --git a/include/infiniop/ops/dsv4_add_rmsnorm_quant.h b/include/infiniop/ops/dsv4_add_rmsnorm_quant.h new file mode 100644 index 000000000..307618664 --- /dev/null +++ b/include/infiniop/ops/dsv4_add_rmsnorm_quant.h @@ -0,0 +1,9 @@ +#ifndef __INFINIOP_DSV4_ADD_RMSNORM_QUANT_API_H__ +#define __INFINIOP_DSV4_ADD_RMSNORM_QUANT_API_H__ +#include "../operator_descriptor.h" +typedef struct InfiniopDescriptor *infiniopDsv4AddRMSNormQuantDescriptor_t; +__INFINI_C __export infiniStatus_t infiniopCreateDsv4AddRMSNormQuantDescriptor(infiniopHandle_t, infiniopDsv4AddRMSNormQuantDescriptor_t *, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, float); +__INFINI_C __export infiniStatus_t infiniopGetDsv4AddRMSNormQuantWorkspaceSize(infiniopDsv4AddRMSNormQuantDescriptor_t, size_t *); +__INFINI_C __export infiniStatus_t infiniopDsv4AddRMSNormQuant(infiniopDsv4AddRMSNormQuantDescriptor_t, void *, size_t, void *, void *, void *, const void *, const void *, void *); +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4AddRMSNormQuantDescriptor(infiniopDsv4AddRMSNormQuantDescriptor_t); +#endif diff --git a/include/infiniop/ops/dsv4_fused_rope.h b/include/infiniop/ops/dsv4_fused_rope.h new file mode 100644 index 000000000..70c425ce9 --- /dev/null +++ b/include/infiniop/ops/dsv4_fused_rope.h @@ -0,0 +1,13 @@ +#ifndef __INFINIOP_DSV4_FUSED_ROPE_API_H__ +#define __INFINIOP_DSV4_FUSED_ROPE_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDsv4FusedRopeDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDsv4FusedRopeDescriptor(infiniopHandle_t handle, infiniopDsv4FusedRopeDescriptor_t *desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t k_desc, infiniopTensorDescriptor_t freq_real_desc, infiniopTensorDescriptor_t freq_imag_desc, int has_k); +__INFINI_C __export infiniStatus_t infiniopGetDsv4FusedRopeWorkspaceSize(infiniopDsv4FusedRopeDescriptor_t desc, size_t *size); +__INFINI_C __export infiniStatus_t infiniopDsv4FusedRope(infiniopDsv4FusedRopeDescriptor_t desc, void *workspace, size_t workspace_size, void *q, void *k, const void *freq_real, const void *freq_imag, void *stream); +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4FusedRopeDescriptor(infiniopDsv4FusedRopeDescriptor_t desc); + +#endif // __INFINIOP_DSV4_FUSED_ROPE_API_H__ diff --git a/include/infiniop/ops/dsv4_linear_bf16_fp32.h b/include/infiniop/ops/dsv4_linear_bf16_fp32.h new file mode 100644 index 000000000..c55313e56 --- /dev/null +++ b/include/infiniop/ops/dsv4_linear_bf16_fp32.h @@ -0,0 +1,31 @@ +#ifndef __INFINIOP_DSV4_LINEAR_BF16_FP32_API_H__ +#define __INFINIOP_DSV4_LINEAR_BF16_FP32_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDsv4LinearBf16Fp32Descriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDsv4LinearBf16Fp32Descriptor( + infiniopHandle_t handle, + infiniopDsv4LinearBf16Fp32Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t w_desc); + +__INFINI_C __export infiniStatus_t infiniopGetDsv4LinearBf16Fp32WorkspaceSize( + infiniopDsv4LinearBf16Fp32Descriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDsv4LinearBf16Fp32( + infiniopDsv4LinearBf16Fp32Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + const void *w, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4LinearBf16Fp32Descriptor( + infiniopDsv4LinearBf16Fp32Descriptor_t desc); + +#endif diff --git a/include/infiniop/ops/dsv4_mask_topk_ids.h b/include/infiniop/ops/dsv4_mask_topk_ids.h new file mode 100644 index 000000000..ca9741b31 --- /dev/null +++ b/include/infiniop/ops/dsv4_mask_topk_ids.h @@ -0,0 +1,29 @@ +#ifndef __INFINIOP_DSV4_MASK_TOPK_IDS_API_H__ +#define __INFINIOP_DSV4_MASK_TOPK_IDS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDsv4MaskTopkIdsDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDsv4MaskTopkIdsDescriptor( + infiniopHandle_t handle, + infiniopDsv4MaskTopkIdsDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t topk_ids_desc, + infiniopTensorDescriptor_t num_token_non_padded_desc); + +__INFINI_C __export infiniStatus_t infiniopGetDsv4MaskTopkIdsWorkspaceSize( + infiniopDsv4MaskTopkIdsDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDsv4MaskTopkIds( + infiniopDsv4MaskTopkIdsDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *topk_ids, + const void *num_token_non_padded, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4MaskTopkIdsDescriptor( + infiniopDsv4MaskTopkIdsDescriptor_t desc); + +#endif diff --git a/include/infiniop/ops/dsv4_per_token_quant_int8.h b/include/infiniop/ops/dsv4_per_token_quant_int8.h new file mode 100644 index 000000000..5fcdaba00 --- /dev/null +++ b/include/infiniop/ops/dsv4_per_token_quant_int8.h @@ -0,0 +1,31 @@ +#ifndef __INFINIOP_DSV4_PER_TOKEN_QUANT_INT8_API_H__ +#define __INFINIOP_DSV4_PER_TOKEN_QUANT_INT8_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDsv4PerTokenQuantInt8Descriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDsv4PerTokenQuantInt8Descriptor( + infiniopHandle_t handle, + infiniopDsv4PerTokenQuantInt8Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t q_desc, + infiniopTensorDescriptor_t scale_desc, + infiniopTensorDescriptor_t x_desc); + +__INFINI_C __export infiniStatus_t infiniopGetDsv4PerTokenQuantInt8WorkspaceSize( + infiniopDsv4PerTokenQuantInt8Descriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopDsv4PerTokenQuantInt8( + infiniopDsv4PerTokenQuantInt8Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *q, + void *scale, + const void *x, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4PerTokenQuantInt8Descriptor( + infiniopDsv4PerTokenQuantInt8Descriptor_t desc); + +#endif // __INFINIOP_DSV4_PER_TOKEN_QUANT_INT8_API_H__ diff --git a/include/infiniop/ops/dsv4_rmsnorm_self.h b/include/infiniop/ops/dsv4_rmsnorm_self.h new file mode 100644 index 000000000..62d7acf0c --- /dev/null +++ b/include/infiniop/ops/dsv4_rmsnorm_self.h @@ -0,0 +1,18 @@ +#ifndef __INFINIOP_DSV4_RMSNORM_SELF_API_H__ +#define __INFINIOP_DSV4_RMSNORM_SELF_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDsv4RMSNormSelfDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDsv4RMSNormSelfDescriptor( + infiniopHandle_t handle, + infiniopDsv4RMSNormSelfDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + float epsilon); +__INFINI_C __export infiniStatus_t infiniopGetDsv4RMSNormSelfWorkspaceSize(infiniopDsv4RMSNormSelfDescriptor_t desc, size_t *size); +__INFINI_C __export infiniStatus_t infiniopDsv4RMSNormSelf(infiniopDsv4RMSNormSelfDescriptor_t desc, void *workspace, size_t workspace_size, void *y, const void *x, void *stream); +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4RMSNormSelfDescriptor(infiniopDsv4RMSNormSelfDescriptor_t desc); + +#endif // __INFINIOP_DSV4_RMSNORM_SELF_API_H__ diff --git a/include/infiniop/ops/dsv4_silu_and_mul.h b/include/infiniop/ops/dsv4_silu_and_mul.h new file mode 100644 index 000000000..397f91c18 --- /dev/null +++ b/include/infiniop/ops/dsv4_silu_and_mul.h @@ -0,0 +1,13 @@ +#ifndef __INFINIOP_DSV4_SILU_AND_MUL_API_H__ +#define __INFINIOP_DSV4_SILU_AND_MUL_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDsv4SiluAndMulDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDsv4SiluAndMulDescriptor(infiniopHandle_t handle, infiniopDsv4SiluAndMulDescriptor_t *desc_ptr, infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t gate_desc, infiniopTensorDescriptor_t up_desc); +__INFINI_C __export infiniStatus_t infiniopGetDsv4SiluAndMulWorkspaceSize(infiniopDsv4SiluAndMulDescriptor_t desc, size_t *size); +__INFINI_C __export infiniStatus_t infiniopDsv4SiluAndMul(infiniopDsv4SiluAndMulDescriptor_t desc, void *workspace, size_t workspace_size, void *y, const void *gate, const void *up, void *stream); +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4SiluAndMulDescriptor(infiniopDsv4SiluAndMulDescriptor_t desc); + +#endif // __INFINIOP_DSV4_SILU_AND_MUL_API_H__ diff --git a/include/infiniop/ops/dsv4_silu_mul_quant.h b/include/infiniop/ops/dsv4_silu_mul_quant.h new file mode 100644 index 000000000..c603f2dc5 --- /dev/null +++ b/include/infiniop/ops/dsv4_silu_mul_quant.h @@ -0,0 +1,9 @@ +#ifndef __INFINIOP_DSV4_SILU_MUL_QUANT_API_H__ +#define __INFINIOP_DSV4_SILU_MUL_QUANT_API_H__ +#include "../operator_descriptor.h" +typedef struct InfiniopDescriptor *infiniopDsv4SiluMulQuantDescriptor_t; +__INFINI_C __export infiniStatus_t infiniopCreateDsv4SiluMulQuantDescriptor(infiniopHandle_t, infiniopDsv4SiluMulQuantDescriptor_t *, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t, infiniopTensorDescriptor_t); +__INFINI_C __export infiniStatus_t infiniopGetDsv4SiluMulQuantWorkspaceSize(infiniopDsv4SiluMulQuantDescriptor_t, size_t *); +__INFINI_C __export infiniStatus_t infiniopDsv4SiluMulQuant(infiniopDsv4SiluMulQuantDescriptor_t, void *, size_t, void *, void *, const void *, const void *, void *); +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4SiluMulQuantDescriptor(infiniopDsv4SiluMulQuantDescriptor_t); +#endif diff --git a/include/infiniop/ops/dsv4_swa_prefill_indices.h b/include/infiniop/ops/dsv4_swa_prefill_indices.h new file mode 100644 index 000000000..5f8368934 --- /dev/null +++ b/include/infiniop/ops/dsv4_swa_prefill_indices.h @@ -0,0 +1,9 @@ +#ifndef __INFINIOP_DSV4_SWA_PREFILL_INDICES_API_H__ +#define __INFINIOP_DSV4_SWA_PREFILL_INDICES_API_H__ +#include "../operator_descriptor.h" +typedef struct InfiniopDescriptor *infiniopDsv4SwaPrefillIndicesDescriptor_t; +__INFINI_C __export infiniStatus_t infiniopCreateDsv4SwaPrefillIndicesDescriptor(infiniopHandle_t, infiniopDsv4SwaPrefillIndicesDescriptor_t *, infiniopTensorDescriptor_t, int); +__INFINI_C __export infiniStatus_t infiniopGetDsv4SwaPrefillIndicesWorkspaceSize(infiniopDsv4SwaPrefillIndicesDescriptor_t, size_t *); +__INFINI_C __export infiniStatus_t infiniopDsv4SwaPrefillIndices(infiniopDsv4SwaPrefillIndicesDescriptor_t, void *, size_t, void *, void *); +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4SwaPrefillIndicesDescriptor(infiniopDsv4SwaPrefillIndicesDescriptor_t); +#endif diff --git a/include/infiniop/ops/dsv4_topk_transform.h b/include/infiniop/ops/dsv4_topk_transform.h new file mode 100644 index 000000000..21d16cd6d --- /dev/null +++ b/include/infiniop/ops/dsv4_topk_transform.h @@ -0,0 +1,13 @@ +#ifndef __INFINIOP_DSV4_TOPK_TRANSFORM_API_H__ +#define __INFINIOP_DSV4_TOPK_TRANSFORM_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopDsv4TopkTransformDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateDsv4TopkTransformDescriptor(infiniopHandle_t handle, infiniopDsv4TopkTransformDescriptor_t *desc_ptr, infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t scores_desc, infiniopTensorDescriptor_t seq_lens_desc, infiniopTensorDescriptor_t page_tables_desc, int page_size); +__INFINI_C __export infiniStatus_t infiniopGetDsv4TopkTransformWorkspaceSize(infiniopDsv4TopkTransformDescriptor_t desc, size_t *size); +__INFINI_C __export infiniStatus_t infiniopDsv4TopkTransform(infiniopDsv4TopkTransformDescriptor_t desc, void *workspace, size_t workspace_size, void *out, const void *scores, const void *seq_lens, const void *page_tables, void *stream); +__INFINI_C __export infiniStatus_t infiniopDestroyDsv4TopkTransformDescriptor(infiniopDsv4TopkTransformDescriptor_t desc); + +#endif // __INFINIOP_DSV4_TOPK_TRANSFORM_API_H__ diff --git a/src/infinicore/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8.cc b/src/infinicore/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8.cc new file mode 100644 index 000000000..d08220213 --- /dev/null +++ b/src/infinicore/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8.cc @@ -0,0 +1,13 @@ +#include "infinicore/ops/dsv4_act_quant_fp8.hpp" +#include "../../utils.hpp" +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4ActQuantFp8); +Dsv4ActQuantFp8::Dsv4ActQuantFp8(Tensor xq, Tensor scale, const Tensor &x, float fp8_max) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(xq, scale, x); + INFINICORE_GRAPH_OP_DISPATCH(xq->device().getType(), xq, scale, x, fp8_max); +} +void Dsv4ActQuantFp8::execute(Tensor xq, Tensor scale, const Tensor &x, float fp8_max) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4ActQuantFp8, xq, scale, x, fp8_max); +} +void dsv4_act_quant_fp8_(Tensor xq, Tensor scale, const Tensor &x, float fp8_max) { Dsv4ActQuantFp8::execute(xq, scale, x, fp8_max); } +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8_infiniop.cc b/src/infinicore/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8_infiniop.cc new file mode 100644 index 000000000..4bcc932c1 --- /dev/null +++ b/src/infinicore/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8_infiniop.cc @@ -0,0 +1,25 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_act_quant_fp8.hpp" +#include "infiniop/ops/dsv4_act_quant_fp8.h" +namespace infinicore::op::dsv4_act_quant_fp8_impl::infiniop { +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4ActQuantFp8, 100); +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, xq, scale, x; +}; +void *plan(Tensor xq, Tensor scale, const Tensor &x, float fp8_max) { + size_t seed = hash_combine(xq, scale, x, fp8_max); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(Descriptor, descriptor, Dsv4ActQuantFp8, seed, xq->desc(), scale->desc(), x->desc(), fp8_max); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4ActQuantFp8, descriptor); + return new PlannedMeta{descriptor, graph::GraphTensor(workspace), graph::GraphTensor(xq), graph::GraphTensor(scale), graph::GraphTensor(x)}; +} +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4ActQuantFp8(p->descriptor->desc, p->workspace->data(), p->workspace->numel(), p->xq->data(), p->scale->data(), p->x->data(), context::getStream())); +} +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4ActQuantFp8, &plan, &run, &cleanup); +} // namespace infinicore::op::dsv4_act_quant_fp8_impl::infiniop diff --git a/src/infinicore/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant.cc b/src/infinicore/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant.cc new file mode 100644 index 000000000..80a8ac104 --- /dev/null +++ b/src/infinicore/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant.cc @@ -0,0 +1,13 @@ +#include "infinicore/ops/dsv4_add_rmsnorm_quant.hpp" +#include "../../utils.hpp" +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4AddRMSNormQuant); +Dsv4AddRMSNormQuant::Dsv4AddRMSNormQuant(Tensor res, Tensor q, Tensor scale, const Tensor &x, const Tensor &weight, float epsilon) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(res, q, scale, x, weight); + INFINICORE_GRAPH_OP_DISPATCH(res->device().getType(), res, q, scale, x, weight, epsilon); +} +void Dsv4AddRMSNormQuant::execute(Tensor res, Tensor q, Tensor scale, const Tensor &x, const Tensor &weight, float epsilon) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4AddRMSNormQuant, res, q, scale, x, weight, epsilon); +} +void dsv4_add_rmsnorm_quant_(Tensor res, Tensor q, Tensor scale, const Tensor &x, const Tensor &weight, float epsilon) { Dsv4AddRMSNormQuant::execute(res, q, scale, x, weight, epsilon); } +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant_infiniop.cc b/src/infinicore/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant_infiniop.cc new file mode 100644 index 000000000..fffeec177 --- /dev/null +++ b/src/infinicore/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant_infiniop.cc @@ -0,0 +1,25 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_add_rmsnorm_quant.hpp" +#include "infiniop/ops/dsv4_add_rmsnorm_quant.h" +namespace infinicore::op::dsv4_add_rmsnorm_quant_impl::infiniop { +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4AddRMSNormQuant, 100); +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, res, q, scale, x, weight; +}; +void *plan(Tensor res, Tensor q, Tensor scale, const Tensor &x, const Tensor &weight, float epsilon) { + size_t seed = hash_combine(res, q, scale, x, weight, epsilon); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(Descriptor, descriptor, Dsv4AddRMSNormQuant, seed, res->desc(), q->desc(), scale->desc(), x->desc(), weight->desc(), epsilon); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4AddRMSNormQuant, descriptor); + return new PlannedMeta{descriptor, graph::GraphTensor(workspace), graph::GraphTensor(res), graph::GraphTensor(q), graph::GraphTensor(scale), graph::GraphTensor(x), graph::GraphTensor(weight)}; +} +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4AddRMSNormQuant(p->descriptor->desc, p->workspace->data(), p->workspace->numel(), p->res->data(), p->q->data(), p->scale->data(), p->x->data(), p->weight->data(), context::getStream())); +} +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4AddRMSNormQuant, &plan, &run, &cleanup); +} // namespace infinicore::op::dsv4_add_rmsnorm_quant_impl::infiniop diff --git a/src/infinicore/ops/dsv4_fused_rope/dsv4_fused_rope.cc b/src/infinicore/ops/dsv4_fused_rope/dsv4_fused_rope.cc new file mode 100644 index 000000000..b04172dcc --- /dev/null +++ b/src/infinicore/ops/dsv4_fused_rope/dsv4_fused_rope.cc @@ -0,0 +1,32 @@ +#include "infinicore/ops/dsv4_fused_rope.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4FusedRope); + +Dsv4FusedRope::Dsv4FusedRope(Tensor q, Tensor k, const Tensor &freq_real, const Tensor &freq_imag, bool has_k) { + if (has_k) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q, k, freq_real, freq_imag); + } else { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q, freq_real, freq_imag); + } + INFINICORE_GRAPH_OP_DISPATCH(q->device().getType(), q, k, freq_real, freq_imag, has_k); +} + +void Dsv4FusedRope::execute(Tensor q, Tensor k, const Tensor &freq_real, const Tensor &freq_imag, bool has_k) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4FusedRope, q, k, freq_real, freq_imag, has_k); +} + +void dsv4_fused_rope_(Tensor q, Tensor k, const Tensor &freq_real, const Tensor &freq_imag, bool has_k) { + Dsv4FusedRope::execute(q, k, freq_real, freq_imag, has_k); +} + +Tensor dsv4_fused_rope(const Tensor &q, const Tensor &freq_real, const Tensor &freq_imag) { + auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); + out->copy_from(q); + dsv4_fused_rope_(out, Tensor(), freq_real, freq_imag, false); + return out; +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_fused_rope/dsv4_fused_rope_infiniop.cc b/src/infinicore/ops/dsv4_fused_rope/dsv4_fused_rope_infiniop.cc new file mode 100644 index 000000000..9568055f4 --- /dev/null +++ b/src/infinicore/ops/dsv4_fused_rope/dsv4_fused_rope_infiniop.cc @@ -0,0 +1,59 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_fused_rope.hpp" +#include "infiniop/ops/dsv4_fused_rope.h" + +#include + +namespace infinicore::op::dsv4_fused_rope_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4FusedRope, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, q, freq_real, freq_imag; + std::optional k; + bool has_k; +}; + +void *plan(Tensor q, Tensor k, const Tensor &freq_real, const Tensor &freq_imag, bool has_k) { + size_t seed = has_k ? hash_combine(q, k, freq_real, freq_imag, 1) : hash_combine(q, freq_real, freq_imag, 0); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, Dsv4FusedRope, + seed, + q->desc(), + has_k ? k->desc() : nullptr, + freq_real->desc(), + freq_imag->desc(), + has_k ? 1 : 0); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4FusedRope, descriptor); + return new PlannedMeta{ + descriptor, + graph::GraphTensor(workspace), + graph::GraphTensor(q), + graph::GraphTensor(freq_real), + graph::GraphTensor(freq_imag), + has_k ? std::optional(graph::GraphTensor(k)) : std::nullopt, + has_k}; +} + +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4FusedRope( + p->descriptor->desc, + p->workspace->data(), + p->workspace->numel(), + p->q->data(), + p->has_k ? p->k.value()->data() : nullptr, + p->freq_real->data(), + p->freq_imag->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4FusedRope, &plan, &run, &cleanup); + +} // namespace infinicore::op::dsv4_fused_rope_impl::infiniop diff --git a/src/infinicore/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32.cc b/src/infinicore/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32.cc new file mode 100644 index 000000000..5feaaa93d --- /dev/null +++ b/src/infinicore/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32.cc @@ -0,0 +1,30 @@ +#include "infinicore/ops/dsv4_linear_bf16_fp32.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4LinearBf16Fp32); + +Dsv4LinearBf16Fp32::Dsv4LinearBf16Fp32(Tensor y, const Tensor &x, const Tensor &w) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x, w); + INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x, w); +} + +void Dsv4LinearBf16Fp32::execute(Tensor y, const Tensor &x, const Tensor &w) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4LinearBf16Fp32, y, x, w); +} + +Tensor dsv4_linear_bf16_fp32(const Tensor &x, const Tensor &w) { + auto x_shape = x->shape(); + auto w_shape = w->shape(); + INFINICORE_ASSERT(x_shape.size() == 2 && w_shape.size() == 2 && x_shape[1] == w_shape[1]); + auto y = Tensor::empty({x_shape[0], w_shape[0]}, DataType::F32, x->device()); + dsv4_linear_bf16_fp32_(y, x, w); + return y; +} + +void dsv4_linear_bf16_fp32_(Tensor y, const Tensor &x, const Tensor &w) { + Dsv4LinearBf16Fp32::execute(y, x, w); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32_infiniop.cc b/src/infinicore/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32_infiniop.cc new file mode 100644 index 000000000..c38aa2774 --- /dev/null +++ b/src/infinicore/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32_infiniop.cc @@ -0,0 +1,42 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_linear_bf16_fp32.hpp" +#include "infiniop/ops/dsv4_linear_bf16_fp32.h" + +namespace infinicore::op::dsv4_linear_bf16_fp32_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4LinearBf16Fp32, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, y, x, w; +}; + +void *plan(Tensor y, const Tensor &x, const Tensor &w) { + size_t seed = hash_combine(y, x, w); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, Dsv4LinearBf16Fp32, seed, y->desc(), x->desc(), w->desc()); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4LinearBf16Fp32, descriptor); + return new PlannedMeta{ + descriptor, graph::GraphTensor(workspace), graph::GraphTensor(y), graph::GraphTensor(x), graph::GraphTensor(w)}; +} + +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4LinearBf16Fp32( + p->descriptor->desc, + p->workspace->data(), + p->workspace->numel(), + p->y->data(), + p->x->data(), + p->w->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4LinearBf16Fp32, &plan, &run, &cleanup); + +} // namespace infinicore::op::dsv4_linear_bf16_fp32_impl::infiniop diff --git a/src/infinicore/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids.cc b/src/infinicore/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids.cc new file mode 100644 index 000000000..7bb24ce2e --- /dev/null +++ b/src/infinicore/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids.cc @@ -0,0 +1,21 @@ +#include "infinicore/ops/dsv4_mask_topk_ids.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4MaskTopkIds); + +Dsv4MaskTopkIds::Dsv4MaskTopkIds(Tensor topk_ids, const Tensor &num_token_non_padded) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_ids, num_token_non_padded); + INFINICORE_GRAPH_OP_DISPATCH(topk_ids->device().getType(), topk_ids, num_token_non_padded); +} + +void Dsv4MaskTopkIds::execute(Tensor topk_ids, const Tensor &num_token_non_padded) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4MaskTopkIds, topk_ids, num_token_non_padded); +} + +void dsv4_mask_topk_ids_(Tensor topk_ids, const Tensor &num_token_non_padded) { + Dsv4MaskTopkIds::execute(topk_ids, num_token_non_padded); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids_infiniop.cc b/src/infinicore/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids_infiniop.cc new file mode 100644 index 000000000..e95c5784d --- /dev/null +++ b/src/infinicore/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids_infiniop.cc @@ -0,0 +1,41 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_mask_topk_ids.hpp" +#include "infiniop/ops/dsv4_mask_topk_ids.h" + +namespace infinicore::op::dsv4_mask_topk_ids_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4MaskTopkIds, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, topk_ids, num_token_non_padded; +}; + +void *plan(Tensor topk_ids, const Tensor &num_token_non_padded) { + size_t seed = hash_combine(topk_ids, num_token_non_padded); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, Dsv4MaskTopkIds, seed, topk_ids->desc(), num_token_non_padded->desc()); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4MaskTopkIds, descriptor); + return new PlannedMeta{ + descriptor, graph::GraphTensor(workspace), graph::GraphTensor(topk_ids), graph::GraphTensor(num_token_non_padded)}; +} + +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4MaskTopkIds( + p->descriptor->desc, + p->workspace->data(), + p->workspace->numel(), + p->topk_ids->data(), + p->num_token_non_padded->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4MaskTopkIds, &plan, &run, &cleanup); + +} // namespace infinicore::op::dsv4_mask_topk_ids_impl::infiniop diff --git a/src/infinicore/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8.cc b/src/infinicore/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8.cc new file mode 100644 index 000000000..a0de958f2 --- /dev/null +++ b/src/infinicore/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8.cc @@ -0,0 +1,39 @@ +#include "infinicore/ops/dsv4_per_token_quant_int8.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4PerTokenQuantInt8); + +namespace { +Size leading_rows(const Tensor &x) { + Size rows = 1; + const auto &shape = x->shape(); + for (Size i = 0; i + 1 < shape.size(); ++i) { + rows *= shape[i]; + } + return rows; +} +} // namespace + +Dsv4PerTokenQuantInt8::Dsv4PerTokenQuantInt8(Tensor q, Tensor scale, const Tensor &x) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q, scale, x); + INFINICORE_GRAPH_OP_DISPATCH(q->device().getType(), q, scale, x); +} + +void Dsv4PerTokenQuantInt8::execute(Tensor q, Tensor scale, const Tensor &x) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4PerTokenQuantInt8, q, scale, x); +} + +std::pair dsv4_per_token_quant_int8(const Tensor &x) { + auto q = Tensor::empty(x->shape(), DataType::I8, x->device()); + auto scale = Tensor::empty({leading_rows(x), 1}, DataType::F32, x->device()); + dsv4_per_token_quant_int8_(q, scale, x); + return {q, scale}; +} + +void dsv4_per_token_quant_int8_(Tensor q, Tensor scale, const Tensor &x) { + Dsv4PerTokenQuantInt8::execute(q, scale, x); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8_infiniop.cc b/src/infinicore/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8_infiniop.cc new file mode 100644 index 000000000..f8434e297 --- /dev/null +++ b/src/infinicore/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8_infiniop.cc @@ -0,0 +1,42 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_per_token_quant_int8.hpp" +#include "infiniop/ops/dsv4_per_token_quant_int8.h" + +namespace infinicore::op::dsv4_per_token_quant_int8_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4PerTokenQuantInt8, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, q, scale, x; +}; + +void *plan(Tensor q, Tensor scale, const Tensor &x) { + size_t seed = hash_combine(q, scale, x); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, Dsv4PerTokenQuantInt8, + seed, q->desc(), scale->desc(), x->desc()); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4PerTokenQuantInt8, descriptor); + return new PlannedMeta{descriptor, graph::GraphTensor(workspace), graph::GraphTensor(q), graph::GraphTensor(scale), graph::GraphTensor(x)}; +} + +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4PerTokenQuantInt8( + p->descriptor->desc, + p->workspace->data(), + p->workspace->numel(), + p->q->data(), + p->scale->data(), + p->x->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4PerTokenQuantInt8, &plan, &run, &cleanup); + +} // namespace infinicore::op::dsv4_per_token_quant_int8_impl::infiniop diff --git a/src/infinicore/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self.cc b/src/infinicore/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self.cc new file mode 100644 index 000000000..83bcfee8a --- /dev/null +++ b/src/infinicore/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/dsv4_rmsnorm_self.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4RMSNormSelf); + +Dsv4RMSNormSelf::Dsv4RMSNormSelf(Tensor y, const Tensor &x, float epsilon) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); + INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x, epsilon); +} + +void Dsv4RMSNormSelf::execute(Tensor y, const Tensor &x, float epsilon) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4RMSNormSelf, y, x, epsilon); +} + +Tensor dsv4_rmsnorm_self(const Tensor &x, float epsilon) { + auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); + dsv4_rmsnorm_self_(y, x, epsilon); + return y; +} + +void dsv4_rmsnorm_self_(Tensor y, const Tensor &x, float epsilon) { + Dsv4RMSNormSelf::execute(y, x, epsilon); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self_infiniop.cc b/src/infinicore/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self_infiniop.cc new file mode 100644 index 000000000..5636d09d9 --- /dev/null +++ b/src/infinicore/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self_infiniop.cc @@ -0,0 +1,33 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_rmsnorm_self.hpp" +#include "infiniop/ops/dsv4_rmsnorm_self.h" + +namespace infinicore::op::dsv4_rmsnorm_self_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4RMSNormSelf, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, y, x; +}; + +void *plan(Tensor y, const Tensor &x, float epsilon) { + size_t seed = hash_combine(y, x, epsilon); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(Descriptor, descriptor, Dsv4RMSNormSelf, seed, y->desc(), x->desc(), epsilon); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4RMSNormSelf, descriptor); + return new PlannedMeta{descriptor, graph::GraphTensor(workspace), graph::GraphTensor(y), graph::GraphTensor(x)}; +} + +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4RMSNormSelf(p->descriptor->desc, p->workspace->data(), p->workspace->numel(), p->y->data(), p->x->data(), context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4RMSNormSelf, &plan, &run, &cleanup); + +} // namespace infinicore::op::dsv4_rmsnorm_self_impl::infiniop diff --git a/src/infinicore/ops/dsv4_silu_and_mul/dsv4_silu_and_mul.cc b/src/infinicore/ops/dsv4_silu_and_mul/dsv4_silu_and_mul.cc new file mode 100644 index 000000000..b9d22833a --- /dev/null +++ b/src/infinicore/ops/dsv4_silu_and_mul/dsv4_silu_and_mul.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/dsv4_silu_and_mul.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4SiluAndMul); + +Dsv4SiluAndMul::Dsv4SiluAndMul(Tensor y, const Tensor &gate, const Tensor &up) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, gate, up); + INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, gate, up); +} + +void Dsv4SiluAndMul::execute(Tensor y, const Tensor &gate, const Tensor &up) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4SiluAndMul, y, gate, up); +} + +Tensor dsv4_silu_and_mul(const Tensor &gate, const Tensor &up) { + auto y = Tensor::empty(gate->shape(), gate->dtype(), gate->device()); + dsv4_silu_and_mul_(y, gate, up); + return y; +} + +void dsv4_silu_and_mul_(Tensor y, const Tensor &gate, const Tensor &up) { + Dsv4SiluAndMul::execute(y, gate, up); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_silu_and_mul/dsv4_silu_and_mul_infiniop.cc b/src/infinicore/ops/dsv4_silu_and_mul/dsv4_silu_and_mul_infiniop.cc new file mode 100644 index 000000000..547ed40cb --- /dev/null +++ b/src/infinicore/ops/dsv4_silu_and_mul/dsv4_silu_and_mul_infiniop.cc @@ -0,0 +1,33 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_silu_and_mul.hpp" +#include "infiniop/ops/dsv4_silu_and_mul.h" + +namespace infinicore::op::dsv4_silu_and_mul_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4SiluAndMul, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, y, gate, up; +}; + +void *plan(Tensor y, const Tensor &gate, const Tensor &up) { + size_t seed = hash_combine(y, gate, up); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(Descriptor, descriptor, Dsv4SiluAndMul, seed, y->desc(), gate->desc(), up->desc()); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4SiluAndMul, descriptor); + return new PlannedMeta{descriptor, graph::GraphTensor(workspace), graph::GraphTensor(y), graph::GraphTensor(gate), graph::GraphTensor(up)}; +} + +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4SiluAndMul(p->descriptor->desc, p->workspace->data(), p->workspace->numel(), p->y->data(), p->gate->data(), p->up->data(), context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4SiluAndMul, &plan, &run, &cleanup); + +} // namespace infinicore::op::dsv4_silu_and_mul_impl::infiniop diff --git a/src/infinicore/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant.cc b/src/infinicore/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant.cc new file mode 100644 index 000000000..cf356842f --- /dev/null +++ b/src/infinicore/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant.cc @@ -0,0 +1,13 @@ +#include "infinicore/ops/dsv4_silu_mul_quant.hpp" +#include "../../utils.hpp" +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4SiluMulQuant); +Dsv4SiluMulQuant::Dsv4SiluMulQuant(Tensor q, Tensor scale, const Tensor &gate, const Tensor &up) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q, scale, gate, up); + INFINICORE_GRAPH_OP_DISPATCH(q->device().getType(), q, scale, gate, up); +} +void Dsv4SiluMulQuant::execute(Tensor q, Tensor scale, const Tensor &gate, const Tensor &up) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4SiluMulQuant, q, scale, gate, up); +} +void dsv4_silu_mul_quant_(Tensor q, Tensor scale, const Tensor &gate, const Tensor &up) { Dsv4SiluMulQuant::execute(q, scale, gate, up); } +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant_infiniop.cc b/src/infinicore/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant_infiniop.cc new file mode 100644 index 000000000..18bb3c3d1 --- /dev/null +++ b/src/infinicore/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant_infiniop.cc @@ -0,0 +1,25 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_silu_mul_quant.hpp" +#include "infiniop/ops/dsv4_silu_mul_quant.h" +namespace infinicore::op::dsv4_silu_mul_quant_impl::infiniop { +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4SiluMulQuant, 100); +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, q, scale, gate, up; +}; +void *plan(Tensor q, Tensor scale, const Tensor &gate, const Tensor &up) { + size_t seed = hash_combine(q, scale, gate, up); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(Descriptor, descriptor, Dsv4SiluMulQuant, seed, q->desc(), scale->desc(), gate->desc(), up->desc()); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4SiluMulQuant, descriptor); + return new PlannedMeta{descriptor, graph::GraphTensor(workspace), graph::GraphTensor(q), graph::GraphTensor(scale), graph::GraphTensor(gate), graph::GraphTensor(up)}; +} +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4SiluMulQuant(p->descriptor->desc, p->workspace->data(), p->workspace->numel(), p->q->data(), p->scale->data(), p->gate->data(), p->up->data(), context::getStream())); +} +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4SiluMulQuant, &plan, &run, &cleanup); +} // namespace infinicore::op::dsv4_silu_mul_quant_impl::infiniop diff --git a/src/infinicore/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices.cc b/src/infinicore/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices.cc new file mode 100644 index 000000000..7d14038fb --- /dev/null +++ b/src/infinicore/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices.cc @@ -0,0 +1,12 @@ +#include "infinicore/ops/dsv4_swa_prefill_indices.hpp" +#include "../../utils.hpp" +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4SwaPrefillIndices); +Dsv4SwaPrefillIndices::Dsv4SwaPrefillIndices(Tensor indices, int window_size) { + INFINICORE_GRAPH_OP_DISPATCH(indices->device().getType(), indices, window_size); +} +void Dsv4SwaPrefillIndices::execute(Tensor indices, int window_size) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4SwaPrefillIndices, indices, window_size); +} +void dsv4_swa_prefill_indices_(Tensor indices, int window_size) { Dsv4SwaPrefillIndices::execute(indices, window_size); } +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices_infiniop.cc b/src/infinicore/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices_infiniop.cc new file mode 100644 index 000000000..99ef7e55f --- /dev/null +++ b/src/infinicore/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices_infiniop.cc @@ -0,0 +1,25 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_swa_prefill_indices.hpp" +#include "infiniop/ops/dsv4_swa_prefill_indices.h" +namespace infinicore::op::dsv4_swa_prefill_indices_impl::infiniop { +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4SwaPrefillIndices, 100); +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, indices; +}; +void *plan(Tensor indices, int window_size) { + size_t seed = hash_combine(indices, window_size); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(Descriptor, descriptor, Dsv4SwaPrefillIndices, seed, indices->desc(), window_size); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4SwaPrefillIndices, descriptor); + return new PlannedMeta{descriptor, graph::GraphTensor(workspace), graph::GraphTensor(indices)}; +} +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4SwaPrefillIndices(p->descriptor->desc, p->workspace->data(), p->workspace->numel(), p->indices->data(), context::getStream())); +} +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4SwaPrefillIndices, &plan, &run, &cleanup); +} // namespace infinicore::op::dsv4_swa_prefill_indices_impl::infiniop diff --git a/src/infinicore/ops/dsv4_topk_transform/dsv4_topk_transform.cc b/src/infinicore/ops/dsv4_topk_transform/dsv4_topk_transform.cc new file mode 100644 index 000000000..539a133e2 --- /dev/null +++ b/src/infinicore/ops/dsv4_topk_transform/dsv4_topk_transform.cc @@ -0,0 +1,29 @@ +#include "infinicore/ops/dsv4_topk_transform.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dsv4TopkTransform); + +Dsv4TopkTransform::Dsv4TopkTransform(Tensor out, const Tensor &scores, const Tensor &seq_lens, const Tensor &page_tables, int page_size) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, scores, seq_lens, page_tables); + INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, scores, seq_lens, page_tables, page_size); +} + +void Dsv4TopkTransform::execute(Tensor out, const Tensor &scores, const Tensor &seq_lens, const Tensor &page_tables, int page_size) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dsv4TopkTransform, out, scores, seq_lens, page_tables, page_size); +} + +Tensor dsv4_topk_transform(const Tensor &scores, const Tensor &seq_lens, const Tensor &page_tables, int page_size) { + auto shape = scores->shape(); + INFINICORE_ASSERT(shape.size() == 2 && shape[1] % 64 == 0); + auto out = Tensor::empty({shape[0], shape[1] / 64}, DataType::I32, scores->device()); + dsv4_topk_transform_(out, scores, seq_lens, page_tables, page_size); + return out; +} + +void dsv4_topk_transform_(Tensor out, const Tensor &scores, const Tensor &seq_lens, const Tensor &page_tables, int page_size) { + Dsv4TopkTransform::execute(out, scores, seq_lens, page_tables, page_size); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/dsv4_topk_transform/dsv4_topk_transform_infiniop.cc b/src/infinicore/ops/dsv4_topk_transform/dsv4_topk_transform_infiniop.cc new file mode 100644 index 000000000..dd2778031 --- /dev/null +++ b/src/infinicore/ops/dsv4_topk_transform/dsv4_topk_transform_infiniop.cc @@ -0,0 +1,33 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/dsv4_topk_transform.hpp" +#include "infiniop/ops/dsv4_topk_transform.h" + +namespace infinicore::op::dsv4_topk_transform_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dsv4TopkTransform, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, out, scores, seq_lens, page_tables; +}; + +void *plan(Tensor out, const Tensor &scores, const Tensor &seq_lens, const Tensor &page_tables, int page_size) { + size_t seed = hash_combine(out, scores, seq_lens, page_tables, page_size); + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(Descriptor, descriptor, Dsv4TopkTransform, seed, out->desc(), scores->desc(), seq_lens->desc(), page_tables->desc(), page_size); + INFINIOP_WORKSPACE_TENSOR(workspace, Dsv4TopkTransform, descriptor); + return new PlannedMeta{descriptor, graph::GraphTensor(workspace), graph::GraphTensor(out), graph::GraphTensor(scores), graph::GraphTensor(seq_lens), graph::GraphTensor(page_tables)}; +} + +void run(void *planned_meta) { + auto p = reinterpret_cast(planned_meta); + INFINICORE_CHECK_ERROR(infiniopDsv4TopkTransform(p->descriptor->desc, p->workspace->data(), p->workspace->numel(), p->out->data(), p->scores->data(), p->seq_lens->data(), p->page_tables->data(), context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dsv4TopkTransform, &plan, &run, &cleanup); + +} // namespace infinicore::op::dsv4_topk_transform_impl::infiniop diff --git a/src/infiniop/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8.h b/src/infiniop/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8.h new file mode 100644 index 000000000..e26e3462d --- /dev/null +++ b/src/infiniop/ops/dsv4_act_quant_fp8/dsv4_act_quant_fp8.h @@ -0,0 +1,19 @@ +#ifndef DSV4_ACT_QUANT_FP8_H +#define DSV4_ACT_QUANT_FP8_H +#include "../../operator.h" +#include "info.h" +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_act_quant_fp8::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + size_t _workspace_size; \ + Descriptor(Info info, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id}, _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t xq_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc, float fp8_max); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *xq, void *scale, const void *x, void *stream) const; \ + }; \ + } +#endif diff --git a/src/infiniop/ops/dsv4_act_quant_fp8/info.h b/src/infiniop/ops/dsv4_act_quant_fp8/info.h new file mode 100644 index 000000000..fc6d8de57 --- /dev/null +++ b/src/infiniop/ops/dsv4_act_quant_fp8/info.h @@ -0,0 +1,24 @@ +#ifndef DSV4_ACT_QUANT_FP8_INFO_H +#define DSV4_ACT_QUANT_FP8_INFO_H +#include "../../../utils.h" +#include "../../tensor.h" +namespace op::dsv4_act_quant_fp8 { +struct Info { + size_t rows, cols; + infiniDtype_t dtype; + float fp8_max; +}; +inline infiniStatus_t createInfo(Info *info, infiniopTensorDescriptor_t xq, infiniopTensorDescriptor_t scale, infiniopTensorDescriptor_t x, float fp8_max) { + CHECK_OR_RETURN(info && xq && scale && x, INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(xq->dtype() == INFINI_DTYPE_F8 && scale->dtype() == INFINI_DTYPE_F32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(x->dtype() == INFINI_DTYPE_BF16 || x->dtype() == INFINI_DTYPE_F16, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(xq->ndim() == x->ndim() && x->ndim() >= 2 && xq->shape() == x->shape(), INFINI_STATUS_BAD_TENSOR_SHAPE); + size_t rows = x->numel() / x->dim(x->ndim() - 1); + size_t cols = x->dim(x->ndim() - 1); + CHECK_OR_RETURN(scale->numel() == rows, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(xq->isContiguous() && scale->isContiguous() && x->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + *info = Info{rows, cols, x->dtype(), fp8_max}; + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_act_quant_fp8 +#endif diff --git a/src/infiniop/ops/dsv4_act_quant_fp8/nvidia/dsv4_act_quant_fp8_nvidia.cu b/src/infiniop/ops/dsv4_act_quant_fp8/nvidia/dsv4_act_quant_fp8_nvidia.cu new file mode 100644 index 000000000..a8433fe10 --- /dev/null +++ b/src/infiniop/ops/dsv4_act_quant_fp8/nvidia/dsv4_act_quant_fp8_nvidia.cu @@ -0,0 +1,107 @@ + +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_act_quant_fp8_nvidia.cuh" +namespace { +__device__ __forceinline__ float warpMax(float v) { +#if defined(ENABLE_HYGON_API) + for (int o = 32; o > 0; o >>= 1) { + v = fmaxf(v, __shfl_xor(v, o, 64)); + } +#else + for (int o = 16; o > 0; o >>= 1) { + v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, o, 32)); + } +#endif + return v; +} +template +__device__ float toFloat(T v) { return static_cast(v); } +#if defined(ENABLE_HYGON_API) +using fp8_out_t = uint8_t; +__device__ __forceinline__ uint8_t f32ToE4m3fn(float v) { + if (isnan(v)) { + return 0x7f; + } + int sign = v < 0.0f; + float a = fminf(fabsf(v), 448.0f); + if (a == 0.0f) { + return static_cast(sign << 7); + } + int exp_field; + int mant; + if (a < 0.015625f) { + exp_field = 0; + mant = static_cast(nearbyintf(a * 512.0f)); + mant = max(0, min(7, mant)); + } else { + int exp_unbiased = static_cast(floorf(log2f(a))); + float base = exp2f(static_cast(exp_unbiased)); + exp_field = exp_unbiased + 7; + mant = static_cast(nearbyintf((a / base - 1.0f) * 8.0f)); + if (mant == 8) { + mant = 0; + exp_field += 1; + } + if (exp_field >= 15) { + exp_field = 15; + mant = min(6, mant); + } + } + return static_cast((sign << 7) | (exp_field << 3) | mant); +} +__device__ __forceinline__ fp8_out_t toFp8(float v) { return f32ToE4m3fn(v); } +#else +using fp8_out_t = cuda_fp8_e4m3; +__device__ __forceinline__ fp8_out_t toFp8(float v) { return cuda_fp8_e4m3(v); } +#endif + +template +__global__ void kernel(const T *x, fp8_out_t *xq, float *scale, int rows, int cols, float fp8_max) { + int row = blockIdx.x; + if (row >= rows) { + return; + } + int tid = threadIdx.x, off = row * cols; + float lmax = 0.f; + for (int i = tid; i < cols; i += BS) { + lmax = fmaxf(lmax, fabsf(toFloat(x[off + i]))); + } + lmax = warpMax(lmax); + __shared__ float inv; + if (tid == 0) { + float am = fmaxf(lmax, 1e-12f); + inv = fp8_max / am; + scale[row] = am / fp8_max; + } + __syncthreads(); + float s = inv; + for (int i = tid; i < cols; i += BS) { + xq[off + i] = toFp8(toFloat(x[off + i]) * s); + } +} +template +infiniStatus_t launch(const op::dsv4_act_quant_fp8::Info &info, void *xq, void *scale, const void *x, cudaStream_t st) { + kernel<128, T><<>>(static_cast(x), static_cast(xq), static_cast(scale), (int)info.rows, (int)info.cols, info.fp8_max); + return INFINI_STATUS_SUCCESS; +} +} // namespace +namespace op::dsv4_act_quant_fp8::nvidia { +infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t xq_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc, float fp8_max) { + Info info; + CHECK_STATUS(createInfo(&info, xq_desc, scale_desc, x_desc, fp8_max)); + *desc_ptr = new Descriptor(info, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} +infiniStatus_t Descriptor::calculate(void *, size_t, void *xq, void *scale, const void *x, void *stream) const { + cudaStream_t st = (cudaStream_t)stream; + switch (_info.dtype) { + case INFINI_DTYPE_BF16: + return launch<__nv_bfloat16>(_info, xq, scale, x, st); + case INFINI_DTYPE_F16: + return launch(_info, xq, scale, x, st); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} +} // namespace op::dsv4_act_quant_fp8::nvidia diff --git a/src/infiniop/ops/dsv4_act_quant_fp8/nvidia/dsv4_act_quant_fp8_nvidia.cuh b/src/infiniop/ops/dsv4_act_quant_fp8/nvidia/dsv4_act_quant_fp8_nvidia.cuh new file mode 100644 index 000000000..d8c160ba7 --- /dev/null +++ b/src/infiniop/ops/dsv4_act_quant_fp8/nvidia/dsv4_act_quant_fp8_nvidia.cuh @@ -0,0 +1,5 @@ +#ifndef DSV4_ACT_QUANT_FP8_NVIDIA_CUH +#define DSV4_ACT_QUANT_FP8_NVIDIA_CUH +#include "../dsv4_act_quant_fp8.h" +DESCRIPTOR(nvidia) +#endif diff --git a/src/infiniop/ops/dsv4_act_quant_fp8/operator.cc b/src/infiniop/ops/dsv4_act_quant_fp8/operator.cc new file mode 100644 index 000000000..9c359c1e8 --- /dev/null +++ b/src/infiniop/ops/dsv4_act_quant_fp8/operator.cc @@ -0,0 +1,108 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_act_quant_fp8.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_act_quant_fp8_nvidia.cuh" +#endif +__INFINI_C infiniStatus_t infiniopCreateDsv4ActQuantFp8Descriptor(infiniopHandle_t handle, infiniopDsv4ActQuantFp8Descriptor_t *desc_ptr, infiniopTensorDescriptor_t xq_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc, float fp8_max) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_act_quant_fp8::NAMESPACE::Descriptor::create(handle, reinterpret_cast(desc_ptr), xq_desc, scale_desc, x_desc, fp8_max) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} +__INFINI_C infiniStatus_t infiniopGetDsv4ActQuantFp8WorkspaceSize(infiniopDsv4ActQuantFp8Descriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} +__INFINI_C infiniStatus_t infiniopDsv4ActQuantFp8(infiniopDsv4ActQuantFp8Descriptor_t desc, void *workspace, size_t workspace_size, void *xq, void *scale, const void *x, void *stream) { +#define CALC(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate(workspace, workspace_size, xq, scale, x, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALC(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALC(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALC(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALC(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALC(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALC +} +__INFINI_C infiniStatus_t infiniopDestroyDsv4ActQuantFp8Descriptor(infiniopDsv4ActQuantFp8Descriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant.h b/src/infiniop/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant.h new file mode 100644 index 000000000..a9db33ccb --- /dev/null +++ b/src/infiniop/ops/dsv4_add_rmsnorm_quant/dsv4_add_rmsnorm_quant.h @@ -0,0 +1,21 @@ +#ifndef DSV4_ADD_RMSNORM_QUANT_H +#define DSV4_ADD_RMSNORM_QUANT_H +#include "../../operator.h" +#include "info.h" +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_add_rmsnorm_quant::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + float _epsilon; \ + size_t _workspace_size; \ + Descriptor(Info info, float epsilon, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id}, _info(info), _epsilon(epsilon), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + float epsilon() const { return _epsilon; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t res_desc, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc, infiniopTensorDescriptor_t weight_desc, float epsilon); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *res, void *q, void *scale, const void *x, const void *weight, void *stream) const; \ + }; \ + } +#endif diff --git a/src/infiniop/ops/dsv4_add_rmsnorm_quant/info.h b/src/infiniop/ops/dsv4_add_rmsnorm_quant/info.h new file mode 100644 index 000000000..b01b0d614 --- /dev/null +++ b/src/infiniop/ops/dsv4_add_rmsnorm_quant/info.h @@ -0,0 +1,24 @@ +#ifndef DSV4_ADD_RMSNORM_QUANT_INFO_H +#define DSV4_ADD_RMSNORM_QUANT_INFO_H +#include "../../../utils.h" +#include "../../tensor.h" +namespace op::dsv4_add_rmsnorm_quant { +struct Info { + size_t rows, cols; + infiniDtype_t dtype; +}; +inline infiniStatus_t createInfo(Info *info, infiniopTensorDescriptor_t res, infiniopTensorDescriptor_t q, infiniopTensorDescriptor_t scale, infiniopTensorDescriptor_t x, infiniopTensorDescriptor_t weight) { + CHECK_OR_RETURN(info && res && q && scale && x && weight, INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(q->dtype() == INFINI_DTYPE_I8 && scale->dtype() == INFINI_DTYPE_F32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(res->dtype() == x->dtype() && res->dtype() == weight->dtype() && (res->dtype() == INFINI_DTYPE_BF16 || res->dtype() == INFINI_DTYPE_F16), INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(res->ndim() == x->ndim() && q->ndim() == res->ndim() && res->ndim() >= 2 && weight->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(res->shape() == x->shape() && q->shape() == res->shape(), INFINI_STATUS_BAD_TENSOR_SHAPE); + size_t rows = res->numel() / res->dim(res->ndim() - 1); + size_t cols = res->dim(res->ndim() - 1); + CHECK_OR_RETURN(weight->dim(0) == cols && scale->numel() == rows, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(res->isContiguous() && q->isContiguous() && scale->isContiguous() && x->isContiguous() && weight->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + *info = Info{rows, cols, res->dtype()}; + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_add_rmsnorm_quant +#endif diff --git a/src/infiniop/ops/dsv4_add_rmsnorm_quant/nvidia/dsv4_add_rmsnorm_quant_nvidia.cu b/src/infiniop/ops/dsv4_add_rmsnorm_quant/nvidia/dsv4_add_rmsnorm_quant_nvidia.cu new file mode 100644 index 000000000..30e5ea58a --- /dev/null +++ b/src/infiniop/ops/dsv4_add_rmsnorm_quant/nvidia/dsv4_add_rmsnorm_quant_nvidia.cu @@ -0,0 +1,134 @@ + +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_add_rmsnorm_quant_nvidia.cuh" +namespace { +__device__ __forceinline__ float warpSum(float v) { +#if defined(ENABLE_HYGON_API) + for (int o = 32; o > 0; o >>= 1) { + v += __shfl_xor(v, o, 64); + } +#else + for (int o = 16; o > 0; o >>= 1) { + v += __shfl_xor_sync(0xffffffff, v, o, 32); + } +#endif + return v; +} +__device__ __forceinline__ float warpMax(float v) { +#if defined(ENABLE_HYGON_API) + for (int o = 32; o > 0; o >>= 1) { + v = fmaxf(v, __shfl_xor(v, o, 64)); + } +#else + for (int o = 16; o > 0; o >>= 1) { + v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, o, 32)); + } +#endif + return v; +} +__device__ __forceinline__ float roundHe(float v) { +#if defined(ENABLE_HYGON_API) + return __builtin_rintf(v); +#else + return nearbyintf(v); +#endif +} +template +__device__ float toFloat(T v) { return static_cast(v); } +template +__device__ T fromFloat(float v) { return static_cast(v); } +template +__global__ void kernel(T *res, const T *x, const T *w, int8_t *q, float *s, int rows, int cols, float eps) { + int row = blockIdx.x; + if (row >= rows) { + return; + } + int tid = threadIdx.x, off = row * cols; + __shared__ float s_rrms, s_inv, s_red[BS / 64 + 1]; + float added[EPT], ss = 0.f; +#pragma unroll + for (int i = 0; i < EPT; i++) { + int idx = tid + i * BS; + if (idx < cols) { + float a = toFloat(res[off + idx]) + toFloat(x[off + idx]); + res[off + idx] = fromFloat(a); + added[i] = a; + ss += a * a; + } + } + int lane = tid & 63, wid = tid >> 6; + ss = warpSum(ss); + if (lane == 0) { + s_red[wid] = ss; + } + __syncthreads(); + if (wid == 0) { + ss = (lane < (BS >> 6)) ? s_red[lane] : 0.f; + ss = warpSum(ss); + if (lane == 0) { + s_rrms = rsqrtf(ss / (float)cols + eps); + } + } + __syncthreads(); + float rrms = s_rrms, lmax = 0.f; +#pragma unroll + for (int i = 0; i < EPT; i++) { + int idx = tid + i * BS; + if (idx < cols) { + float nm = added[i] * rrms * toFloat(w[idx]); + added[i] = nm; + lmax = fmaxf(lmax, fabsf(nm)); + } + } + lmax = warpMax(lmax); + if (lane == 0) { + s_red[wid] = lmax; + } + __syncthreads(); + if (wid == 0) { + lmax = (lane < (BS >> 6)) ? s_red[lane] : 0.f; + lmax = warpMax(lmax); + if (lane == 0) { + float am = fmaxf(lmax, 1e-10f); + s_inv = 127.f / am; + s[row] = am / 127.f; + } + } + __syncthreads(); + float inv = s_inv; +#pragma unroll + for (int i = 0; i < EPT; i++) { + int idx = tid + i * BS; + if (idx < cols) { + int qi = (int)roundHe(added[i] * inv); + qi = max(-128, min(127, qi)); + q[off + idx] = (int8_t)qi; + } + } +} +template +infiniStatus_t launch(const op::dsv4_add_rmsnorm_quant::Info &info, float eps, void *res, void *q, void *scale, const void *x, const void *weight, cudaStream_t st) { + kernel<256, 16, T><<>>(static_cast(res), static_cast(x), static_cast(weight), static_cast(q), static_cast(scale), (int)info.rows, (int)info.cols, eps); + return INFINI_STATUS_SUCCESS; +} +} // namespace +namespace op::dsv4_add_rmsnorm_quant::nvidia { +infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t res_desc, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc, infiniopTensorDescriptor_t weight_desc, float epsilon) { + Info info; + CHECK_STATUS(createInfo(&info, res_desc, q_desc, scale_desc, x_desc, weight_desc)); + *desc_ptr = new Descriptor(info, epsilon, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} +infiniStatus_t Descriptor::calculate(void *, size_t, void *res, void *q, void *scale, const void *x, const void *weight, void *stream) const { + cudaStream_t st = (cudaStream_t)stream; + switch (_info.dtype) { + case INFINI_DTYPE_BF16: + return launch<__nv_bfloat16>(_info, _epsilon, res, q, scale, x, weight, st); + case INFINI_DTYPE_F16: + return launch(_info, _epsilon, res, q, scale, x, weight, st); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} +} // namespace op::dsv4_add_rmsnorm_quant::nvidia diff --git a/src/infiniop/ops/dsv4_add_rmsnorm_quant/nvidia/dsv4_add_rmsnorm_quant_nvidia.cuh b/src/infiniop/ops/dsv4_add_rmsnorm_quant/nvidia/dsv4_add_rmsnorm_quant_nvidia.cuh new file mode 100644 index 000000000..c6eaffd01 --- /dev/null +++ b/src/infiniop/ops/dsv4_add_rmsnorm_quant/nvidia/dsv4_add_rmsnorm_quant_nvidia.cuh @@ -0,0 +1,5 @@ +#ifndef DSV4_ADD_RMSNORM_QUANT_NVIDIA_CUH +#define DSV4_ADD_RMSNORM_QUANT_NVIDIA_CUH +#include "../dsv4_add_rmsnorm_quant.h" +DESCRIPTOR(nvidia) +#endif diff --git a/src/infiniop/ops/dsv4_add_rmsnorm_quant/operator.cc b/src/infiniop/ops/dsv4_add_rmsnorm_quant/operator.cc new file mode 100644 index 000000000..5c4679989 --- /dev/null +++ b/src/infiniop/ops/dsv4_add_rmsnorm_quant/operator.cc @@ -0,0 +1,108 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_add_rmsnorm_quant.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_add_rmsnorm_quant_nvidia.cuh" +#endif +__INFINI_C infiniStatus_t infiniopCreateDsv4AddRMSNormQuantDescriptor(infiniopHandle_t handle, infiniopDsv4AddRMSNormQuantDescriptor_t *desc_ptr, infiniopTensorDescriptor_t res_desc, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc, infiniopTensorDescriptor_t weight_desc, float epsilon) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_add_rmsnorm_quant::NAMESPACE::Descriptor::create(handle, reinterpret_cast(desc_ptr), res_desc, q_desc, scale_desc, x_desc, weight_desc, epsilon) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} +__INFINI_C infiniStatus_t infiniopGetDsv4AddRMSNormQuantWorkspaceSize(infiniopDsv4AddRMSNormQuantDescriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} +__INFINI_C infiniStatus_t infiniopDsv4AddRMSNormQuant(infiniopDsv4AddRMSNormQuantDescriptor_t desc, void *workspace, size_t workspace_size, void *res, void *q, void *scale, const void *x, const void *weight, void *stream) { +#define CALC(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate(workspace, workspace_size, res, q, scale, x, weight, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALC(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALC(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALC(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALC(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALC(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALC +} +__INFINI_C infiniStatus_t infiniopDestroyDsv4AddRMSNormQuantDescriptor(infiniopDsv4AddRMSNormQuantDescriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_fused_rope/dsv4_fused_rope.h b/src/infiniop/ops/dsv4_fused_rope/dsv4_fused_rope.h new file mode 100644 index 000000000..3d29e7dbb --- /dev/null +++ b/src/infiniop/ops/dsv4_fused_rope/dsv4_fused_rope.h @@ -0,0 +1,19 @@ +#ifndef DSV4_FUSED_ROPE_H +#define DSV4_FUSED_ROPE_H +#include "../../operator.h" +#include "info.h" +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_fused_rope::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + size_t _workspace_size; \ + Descriptor(Info info, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id}, _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t k_desc, infiniopTensorDescriptor_t freq_real_desc, infiniopTensorDescriptor_t freq_imag_desc, int has_k); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *q, void *k, const void *freq_real, const void *freq_imag, void *stream) const; \ + }; \ + } +#endif diff --git a/src/infiniop/ops/dsv4_fused_rope/info.h b/src/infiniop/ops/dsv4_fused_rope/info.h new file mode 100644 index 000000000..53621a717 --- /dev/null +++ b/src/infiniop/ops/dsv4_fused_rope/info.h @@ -0,0 +1,33 @@ +#ifndef DSV4_FUSED_ROPE_INFO_H +#define DSV4_FUSED_ROPE_INFO_H +#include "../../../utils.h" +#include "../../tensor.h" +namespace op::dsv4_fused_rope { +struct Info { + infiniDtype_t dtype; + size_t seq_len; + size_t q_heads; + size_t k_heads; + size_t rope_dim; + bool has_k; +}; +inline infiniStatus_t createInfo(Info *info, infiniopTensorDescriptor_t q, infiniopTensorDescriptor_t k, infiniopTensorDescriptor_t fr, infiniopTensorDescriptor_t fi, int has_k) { + CHECK_OR_RETURN(info && q && fr && fi, INFINI_STATUS_NULL_POINTER); + CHECK_DTYPE(q->dtype(), INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + CHECK_OR_RETURN(fr->dtype() == INFINI_DTYPE_F32 && fi->dtype() == INFINI_DTYPE_F32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(q->ndim() == 3 && fr->ndim() == 2 && fi->ndim() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(q->dim(2) % 2 == 0, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(fr->dim(0) == q->dim(0) && fi->dim(0) == q->dim(0) && fr->dim(1) == q->dim(2) / 2 && fi->dim(1) == q->dim(2) / 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + size_t k_heads = 0; + if (has_k) { + CHECK_OR_RETURN(k != nullptr, INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(k->dtype() == q->dtype() && k->ndim() == 3 && k->dim(0) == q->dim(0) && k->dim(2) == q->dim(2), INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(k->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + k_heads = k->dim(1); + } + CHECK_OR_RETURN(q->isContiguous() && fr->isContiguous() && fi->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + *info = Info{q->dtype(), q->dim(0), q->dim(1), k_heads, q->dim(2), has_k != 0}; + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_fused_rope +#endif diff --git a/src/infiniop/ops/dsv4_fused_rope/nvidia/dsv4_fused_rope_nvidia.cu b/src/infiniop/ops/dsv4_fused_rope/nvidia/dsv4_fused_rope_nvidia.cu new file mode 100644 index 000000000..dd1f49047 --- /dev/null +++ b/src/infiniop/ops/dsv4_fused_rope/nvidia/dsv4_fused_rope_nvidia.cu @@ -0,0 +1,68 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_fused_rope_nvidia.cuh" +namespace { +template +__device__ float toFloat(T v) { return static_cast(v); } +template +__device__ T fromFloat(float v) { return static_cast(v); } +template +__global__ void kernel(T *q, T *k, const float *fr, const float *fi, size_t seq_len, size_t q_heads, size_t k_heads, size_t rope_dim, bool has_k) { + size_t s = blockIdx.x; + size_t h = blockIdx.y; + size_t half = rope_dim / 2; + if (s >= seq_len || h >= q_heads) { + return; + } + for (size_t i = threadIdx.x; i < half; i += blockDim.x) { + float real = fr[s * half + i]; + float imag = fi[s * half + i]; + size_t even = s * q_heads * rope_dim + h * rope_dim + 2 * i; + size_t odd = even + 1; + float x0 = toFloat(q[even]); + float x1 = toFloat(q[odd]); + q[even] = fromFloat(x0 * real - x1 * imag); + q[odd] = fromFloat(x0 * imag + x1 * real); + } + if (has_k && h < k_heads) { + for (size_t i = threadIdx.x; i < half; i += blockDim.x) { + float real = fr[s * half + i]; + float imag = fi[s * half + i]; + size_t even = s * k_heads * rope_dim + h * rope_dim + 2 * i; + size_t odd = even + 1; + float x0 = toFloat(k[even]); + float x1 = toFloat(k[odd]); + k[even] = fromFloat(x0 * real - x1 * imag); + k[odd] = fromFloat(x0 * imag + x1 * real); + } + } +} +template +infiniStatus_t launch(const op::dsv4_fused_rope::Info &info, void *q, void *k, const void *fr, const void *fi, cudaStream_t stream) { + dim3 grid(info.seq_len, info.q_heads); + int threads = static_cast((info.rope_dim / 2) < 256 ? (info.rope_dim / 2) : 256); + kernel<<>>(static_cast(q), static_cast(k), static_cast(fr), static_cast(fi), info.seq_len, info.q_heads, info.k_heads, info.rope_dim, info.has_k); + return INFINI_STATUS_SUCCESS; +} +} // namespace +namespace op::dsv4_fused_rope::nvidia { +infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t k_desc, infiniopTensorDescriptor_t freq_real_desc, infiniopTensorDescriptor_t freq_imag_desc, int has_k) { + Info info; + CHECK_STATUS(createInfo(&info, q_desc, k_desc, freq_real_desc, freq_imag_desc, has_k)); + *desc_ptr = new Descriptor(info, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} +infiniStatus_t Descriptor::calculate(void *, size_t, void *q, void *k, const void *fr, const void *fi, void *stream) const { + cudaStream_t s = reinterpret_cast(stream); + switch (_info.dtype) { + case INFINI_DTYPE_F16: + return launch(_info, q, k, fr, fi, s); + case INFINI_DTYPE_BF16: + return launch<__nv_bfloat16>(_info, q, k, fr, fi, s); + case INFINI_DTYPE_F32: + return launch(_info, q, k, fr, fi, s); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} +} // namespace op::dsv4_fused_rope::nvidia diff --git a/src/infiniop/ops/dsv4_fused_rope/nvidia/dsv4_fused_rope_nvidia.cuh b/src/infiniop/ops/dsv4_fused_rope/nvidia/dsv4_fused_rope_nvidia.cuh new file mode 100644 index 000000000..842926669 --- /dev/null +++ b/src/infiniop/ops/dsv4_fused_rope/nvidia/dsv4_fused_rope_nvidia.cuh @@ -0,0 +1,5 @@ +#ifndef DSV4_FUSED_ROPE_NVIDIA_CUH +#define DSV4_FUSED_ROPE_NVIDIA_CUH +#include "../dsv4_fused_rope.h" +DESCRIPTOR(nvidia) +#endif diff --git a/src/infiniop/ops/dsv4_fused_rope/operator.cc b/src/infiniop/ops/dsv4_fused_rope/operator.cc new file mode 100644 index 000000000..e0b7ef182 --- /dev/null +++ b/src/infiniop/ops/dsv4_fused_rope/operator.cc @@ -0,0 +1,112 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_fused_rope.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_fused_rope_nvidia.cuh" +#endif +__INFINI_C infiniStatus_t infiniopCreateDsv4FusedRopeDescriptor(infiniopHandle_t handle, infiniopDsv4FusedRopeDescriptor_t *desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t k_desc, infiniopTensorDescriptor_t freq_real_desc, infiniopTensorDescriptor_t freq_imag_desc, int has_k) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_fused_rope::NAMESPACE::Descriptor::create(handle, reinterpret_cast(desc_ptr), q_desc, k_desc, freq_real_desc, freq_imag_desc, has_k) + switch (handle->device) { + +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} +__INFINI_C infiniStatus_t infiniopGetDsv4FusedRopeWorkspaceSize(infiniopDsv4FusedRopeDescriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} +__INFINI_C infiniStatus_t infiniopDsv4FusedRope(infiniopDsv4FusedRopeDescriptor_t desc, void *workspace, size_t workspace_size, void *q, void *k, const void *freq_real, const void *freq_imag, void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate(workspace, workspace_size, q, k, freq_real, freq_imag, stream) + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} +__INFINI_C infiniStatus_t infiniopDestroyDsv4FusedRopeDescriptor(infiniopDsv4FusedRopeDescriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32.h b/src/infiniop/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32.h new file mode 100644 index 000000000..afcafacf5 --- /dev/null +++ b/src/infiniop/ops/dsv4_linear_bf16_fp32/dsv4_linear_bf16_fp32.h @@ -0,0 +1,26 @@ +#ifndef DSV4_LINEAR_BF16_FP32_H +#define DSV4_LINEAR_BF16_FP32_H + +#include "../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_linear_bf16_fp32::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Info _info; \ + size_t _workspace_size; \ + Opaque *_opaque; \ + Descriptor(Info info, size_t workspace_size, Opaque *opaque, infiniDevice_t device_type, int device_id) \ + : InfiniopDescriptor{device_type, device_id}, _info(info), _workspace_size(workspace_size), _opaque(opaque) {} \ + \ + public: \ + ~Descriptor(); \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t x_desc, infiniopTensorDescriptor_t w_desc); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *y, const void *x, const void *w, void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/dsv4_linear_bf16_fp32/info.h b/src/infiniop/ops/dsv4_linear_bf16_fp32/info.h new file mode 100644 index 000000000..7407453cb --- /dev/null +++ b/src/infiniop/ops/dsv4_linear_bf16_fp32/info.h @@ -0,0 +1,36 @@ +#ifndef DSV4_LINEAR_BF16_FP32_INFO_H +#define DSV4_LINEAR_BF16_FP32_INFO_H + +#include "../../../utils.h" +#include "../../tensor.h" + +namespace op::dsv4_linear_bf16_fp32 { + +struct Info { + size_t m; + size_t n; + size_t k; + infiniDtype_t x_dtype; +}; + +inline infiniStatus_t createInfo( + Info *info, + infiniopTensorDescriptor_t y, + infiniopTensorDescriptor_t x, + infiniopTensorDescriptor_t w) { + CHECK_OR_RETURN(info && y && x && w, INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(y->dtype() == INFINI_DTYPE_F32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(w->dtype() == INFINI_DTYPE_F32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(x->dtype() == INFINI_DTYPE_BF16 || x->dtype() == INFINI_DTYPE_F16, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(y->ndim() == 2 && x->ndim() == 2 && w->ndim() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(y->dim(0) == x->dim(0), INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(y->dim(1) == w->dim(0), INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(x->dim(1) == w->dim(1), INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(y->isContiguous() && x->isContiguous() && w->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + *info = Info{y->dim(0), y->dim(1), x->dim(1), x->dtype()}; + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::dsv4_linear_bf16_fp32 + +#endif diff --git a/src/infiniop/ops/dsv4_linear_bf16_fp32/nvidia/dsv4_linear_bf16_fp32_nvidia.cu b/src/infiniop/ops/dsv4_linear_bf16_fp32/nvidia/dsv4_linear_bf16_fp32_nvidia.cu new file mode 100644 index 000000000..722046bad --- /dev/null +++ b/src/infiniop/ops/dsv4_linear_bf16_fp32/nvidia/dsv4_linear_bf16_fp32_nvidia.cu @@ -0,0 +1,112 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_handle.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_linear_bf16_fp32_nvidia.cuh" + +namespace { + +template +__global__ void castToFloatKernel(const T *x, float *x_fp32, size_t total) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < total; idx += stride) { + x_fp32[idx] = static_cast(x[idx]); + } +} + +template +infiniStatus_t castToFloat(const void *x, void *workspace, size_t total, cudaStream_t stream) { + constexpr int block = 256; + size_t blocks = (total + block - 1) / block; + int grid = static_cast(blocks < 4096 ? blocks : 4096); + castToFloatKernel<<>>(static_cast(x), static_cast(workspace), total); + return INFINI_STATUS_SUCCESS; +} + +} // namespace + +namespace op::dsv4_linear_bf16_fp32::nvidia { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t w_desc) { + auto handle = reinterpret_cast(handle_); + Info info; + CHECK_STATUS(createInfo(&info, y_desc, x_desc, w_desc)); + size_t workspace_size = info.m * info.k * sizeof(float); + *desc_ptr = new Descriptor(info, workspace_size, new Opaque{handle->internal()}, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, void *y, const void *x, const void *w, void *stream) const { + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + cudaStream_t cuda_stream = reinterpret_cast(stream); + size_t x_numel = _info.m * _info.k; + switch (_info.x_dtype) { + case INFINI_DTYPE_F16: + CHECK_STATUS(castToFloat(x, workspace, x_numel, cuda_stream)); + break; + case INFINI_DTYPE_BF16: + CHECK_STATUS(castToFloat<__nv_bfloat16>(x, workspace, x_numel, cuda_stream)); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + +#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) + cudaDataType compute_type = CUDA_R_32F; +#else + cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; +#endif + + const float alpha = 1.0f; + const float beta = 0.0f; + const int m = static_cast(_info.m); + const int n = static_cast(_info.n); + const int k = static_cast(_info.k); + const float *x_fp32 = static_cast(workspace); + + CHECK_STATUS(_opaque->internal->useCublas( + cuda_stream, + [&](cublasHandle_t handle) { + CHECK_CUBLAS(cublasGemmEx( + handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + n, + m, + k, + &alpha, + w, + CUDA_R_32F, + k, + x_fp32, + CUDA_R_32F, + k, + &beta, + y, + CUDA_R_32F, + n, + compute_type, + CUBLAS_GEMM_DEFAULT)); + return INFINI_STATUS_SUCCESS; + })); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::dsv4_linear_bf16_fp32::nvidia diff --git a/src/infiniop/ops/dsv4_linear_bf16_fp32/nvidia/dsv4_linear_bf16_fp32_nvidia.cuh b/src/infiniop/ops/dsv4_linear_bf16_fp32/nvidia/dsv4_linear_bf16_fp32_nvidia.cuh new file mode 100644 index 000000000..07d523cf6 --- /dev/null +++ b/src/infiniop/ops/dsv4_linear_bf16_fp32/nvidia/dsv4_linear_bf16_fp32_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef DSV4_LINEAR_BF16_FP32_NVIDIA_CUH +#define DSV4_LINEAR_BF16_FP32_NVIDIA_CUH + +#include "../dsv4_linear_bf16_fp32.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/dsv4_linear_bf16_fp32/operator.cc b/src/infiniop/ops/dsv4_linear_bf16_fp32/operator.cc new file mode 100644 index 000000000..fdcaf9375 --- /dev/null +++ b/src/infiniop/ops/dsv4_linear_bf16_fp32/operator.cc @@ -0,0 +1,134 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_linear_bf16_fp32.h" + +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_linear_bf16_fp32_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateDsv4LinearBf16Fp32Descriptor( + infiniopHandle_t handle, + infiniopDsv4LinearBf16Fp32Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t y_desc, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t w_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_linear_bf16_fp32::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + y_desc, \ + x_desc, \ + w_desc) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetDsv4LinearBf16Fp32WorkspaceSize( + infiniopDsv4LinearBf16Fp32Descriptor_t desc, + size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C infiniStatus_t infiniopDsv4LinearBf16Fp32( + infiniopDsv4LinearBf16Fp32Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *y, + const void *x, + const void *w, + void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate( \ + workspace, workspace_size, y, x, w, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyDsv4LinearBf16Fp32Descriptor( + infiniopDsv4LinearBf16Fp32Descriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids.h b/src/infiniop/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids.h new file mode 100644 index 000000000..a14bca403 --- /dev/null +++ b/src/infiniop/ops/dsv4_mask_topk_ids/dsv4_mask_topk_ids.h @@ -0,0 +1,23 @@ +#ifndef DSV4_MASK_TOPK_IDS_H +#define DSV4_MASK_TOPK_IDS_H + +#include "../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_mask_topk_ids::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + size_t _workspace_size; \ + Descriptor(Info info, size_t workspace_size, infiniDevice_t device_type, int device_id) \ + : InfiniopDescriptor{device_type, device_id}, _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t topk_ids_desc, infiniopTensorDescriptor_t num_token_non_padded_desc); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *topk_ids, const void *num_token_non_padded, void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/dsv4_mask_topk_ids/info.h b/src/infiniop/ops/dsv4_mask_topk_ids/info.h new file mode 100644 index 000000000..2d0cd417f --- /dev/null +++ b/src/infiniop/ops/dsv4_mask_topk_ids/info.h @@ -0,0 +1,30 @@ +#ifndef DSV4_MASK_TOPK_IDS_INFO_H +#define DSV4_MASK_TOPK_IDS_INFO_H + +#include "../../../utils.h" +#include "../../tensor.h" + +namespace op::dsv4_mask_topk_ids { + +struct Info { + size_t batch; + size_t topk; +}; + +inline infiniStatus_t createInfo( + Info *info, + infiniopTensorDescriptor_t topk_ids, + infiniopTensorDescriptor_t num_token_non_padded) { + CHECK_OR_RETURN(info && topk_ids && num_token_non_padded, INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(topk_ids->dtype() == INFINI_DTYPE_I32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(num_token_non_padded->dtype() == INFINI_DTYPE_I32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(topk_ids->ndim() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(num_token_non_padded->numel() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(topk_ids->isContiguous() && num_token_non_padded->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + *info = Info{topk_ids->dim(0), topk_ids->dim(1)}; + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::dsv4_mask_topk_ids + +#endif diff --git a/src/infiniop/ops/dsv4_mask_topk_ids/nvidia/dsv4_mask_topk_ids_nvidia.cu b/src/infiniop/ops/dsv4_mask_topk_ids/nvidia/dsv4_mask_topk_ids_nvidia.cu new file mode 100644 index 000000000..7f7237ce5 --- /dev/null +++ b/src/infiniop/ops/dsv4_mask_topk_ids/nvidia/dsv4_mask_topk_ids_nvidia.cu @@ -0,0 +1,55 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_mask_topk_ids_nvidia.cuh" + +namespace { + +__global__ void kernel(int32_t *topk_ids, const int32_t *num_token_non_padded, size_t batch, size_t topk) { + int keep = num_token_non_padded[0]; + if (keep < 0) { + keep = 0; + } + if (static_cast(keep) > batch) { + keep = static_cast(batch); + } + size_t total = batch * topk; + for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) { + size_t row = idx / topk; + if (row >= static_cast(keep)) { + topk_ids[idx] = -1; + } + } +} + +} // namespace + +namespace op::dsv4_mask_topk_ids::nvidia { + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t topk_ids_desc, + infiniopTensorDescriptor_t num_token_non_padded_desc) { + Info info; + CHECK_STATUS(createInfo(&info, topk_ids_desc, num_token_non_padded_desc)); + *desc_ptr = new Descriptor(info, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *, size_t, void *topk_ids, const void *num_token_non_padded, void *stream) const { + cudaStream_t s = reinterpret_cast(stream); + size_t total = _info.batch * _info.topk; + int blocks = static_cast((total + 255) / 256); + if (blocks > 1024) { + blocks = 1024; + } + kernel<<>>( + static_cast(topk_ids), + static_cast(num_token_non_padded), + _info.batch, + _info.topk); + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::dsv4_mask_topk_ids::nvidia diff --git a/src/infiniop/ops/dsv4_mask_topk_ids/nvidia/dsv4_mask_topk_ids_nvidia.cuh b/src/infiniop/ops/dsv4_mask_topk_ids/nvidia/dsv4_mask_topk_ids_nvidia.cuh new file mode 100644 index 000000000..9d2165ce4 --- /dev/null +++ b/src/infiniop/ops/dsv4_mask_topk_ids/nvidia/dsv4_mask_topk_ids_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef DSV4_MASK_TOPK_IDS_NVIDIA_CUH +#define DSV4_MASK_TOPK_IDS_NVIDIA_CUH + +#include "../dsv4_mask_topk_ids.h" + +DESCRIPTOR(nvidia) + +#endif diff --git a/src/infiniop/ops/dsv4_mask_topk_ids/operator.cc b/src/infiniop/ops/dsv4_mask_topk_ids/operator.cc new file mode 100644 index 000000000..41dc82c7d --- /dev/null +++ b/src/infiniop/ops/dsv4_mask_topk_ids/operator.cc @@ -0,0 +1,131 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_mask_topk_ids.h" + +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_mask_topk_ids_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateDsv4MaskTopkIdsDescriptor( + infiniopHandle_t handle, + infiniopDsv4MaskTopkIdsDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t topk_ids_desc, + infiniopTensorDescriptor_t num_token_non_padded_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_mask_topk_ids::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + topk_ids_desc, \ + num_token_non_padded_desc) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetDsv4MaskTopkIdsWorkspaceSize( + infiniopDsv4MaskTopkIdsDescriptor_t desc, + size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C infiniStatus_t infiniopDsv4MaskTopkIds( + infiniopDsv4MaskTopkIdsDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *topk_ids, + const void *num_token_non_padded, + void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate( \ + workspace, workspace_size, topk_ids, num_token_non_padded, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyDsv4MaskTopkIdsDescriptor( + infiniopDsv4MaskTopkIdsDescriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8.h b/src/infiniop/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8.h new file mode 100644 index 000000000..e84c5fa8e --- /dev/null +++ b/src/infiniop/ops/dsv4_per_token_quant_int8/dsv4_per_token_quant_int8.h @@ -0,0 +1,22 @@ +#ifndef DSV4_PER_TOKEN_QUANT_INT8_H +#define DSV4_PER_TOKEN_QUANT_INT8_H + +#include "../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_per_token_quant_int8::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + size_t _workspace_size; \ + Descriptor(Info info, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id}, _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *q, void *scale, const void *x, void *stream) const; \ + }; \ + } + +#endif diff --git a/src/infiniop/ops/dsv4_per_token_quant_int8/info.h b/src/infiniop/ops/dsv4_per_token_quant_int8/info.h new file mode 100644 index 000000000..41e007193 --- /dev/null +++ b/src/infiniop/ops/dsv4_per_token_quant_int8/info.h @@ -0,0 +1,40 @@ +#ifndef DSV4_PER_TOKEN_QUANT_INT8_INFO_H +#define DSV4_PER_TOKEN_QUANT_INT8_INFO_H + +#include "../../../utils.h" +#include "../../tensor.h" + +namespace op::dsv4_per_token_quant_int8 { + +struct Info { + infiniDtype_t dtype; + size_t rows; + size_t cols; +}; + +inline size_t leadingRows(infiniopTensorDescriptor_t desc) { + size_t rows = 1; + for (size_t i = 0; i + 1 < desc->ndim(); ++i) { + rows *= desc->dim(i); + } + return rows; +} + +inline infiniStatus_t createInfo(Info *info, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc) { + CHECK_OR_RETURN(info != nullptr && q_desc != nullptr && scale_desc != nullptr && x_desc != nullptr, INFINI_STATUS_NULL_POINTER); + CHECK_DTYPE(x_desc->dtype(), INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + CHECK_OR_RETURN(q_desc->dtype() == INFINI_DTYPE_I8 && scale_desc->dtype() == INFINI_DTYPE_F32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(x_desc->ndim() >= 2 && q_desc->ndim() == x_desc->ndim(), INFINI_STATUS_BAD_TENSOR_SHAPE); + for (size_t i = 0; i < x_desc->ndim(); ++i) { + CHECK_OR_RETURN(q_desc->dim(i) == x_desc->dim(i), INFINI_STATUS_BAD_TENSOR_SHAPE); + } + CHECK_OR_RETURN(x_desc->isContiguous() && q_desc->isContiguous() && scale_desc->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + size_t rows = leadingRows(x_desc); + CHECK_OR_RETURN(scale_desc->numel() == rows, INFINI_STATUS_BAD_TENSOR_SHAPE); + *info = Info{x_desc->dtype(), rows, x_desc->dim(x_desc->ndim() - 1)}; + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::dsv4_per_token_quant_int8 + +#endif diff --git a/src/infiniop/ops/dsv4_per_token_quant_int8/nvidia/dsv4_per_token_quant_int8_nvidia.cu b/src/infiniop/ops/dsv4_per_token_quant_int8/nvidia/dsv4_per_token_quant_int8_nvidia.cu new file mode 100644 index 000000000..013b91454 --- /dev/null +++ b/src/infiniop/ops/dsv4_per_token_quant_int8/nvidia/dsv4_per_token_quant_int8_nvidia.cu @@ -0,0 +1,153 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_per_token_quant_int8_nvidia.cuh" + +namespace { +constexpr int BLOCK = 256; + +template +__device__ float toFloat(T v) { + return static_cast(v); +} + +__device__ __forceinline__ float warpMax(float v) { +#if defined(ENABLE_HYGON_API) + for (int o = 32; o > 0; o >>= 1) { + v = fmaxf(v, __shfl_xor(v, o, 64)); + } +#else + for (int o = 16; o > 0; o >>= 1) { + v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, o, 32)); + } +#endif + return v; +} + +__device__ __forceinline__ float roundHe(float v) { +#if defined(ENABLE_HYGON_API) + return __builtin_rintf(v); +#else + return nearbyintf(v); +#endif +} + +template +__global__ void fallbackKernel(const T *__restrict__ x, int8_t *__restrict__ q, float *__restrict__ scale, size_t rows, size_t cols) { + size_t row = blockIdx.x; + if (row >= rows) { + return; + } + __shared__ float red[BLOCK]; + float local = 0.0f; + size_t base = row * cols; + for (size_t i = threadIdx.x; i < cols; i += blockDim.x) { + local = fmaxf(local, fabsf(toFloat(x[base + i]))); + } + red[threadIdx.x] = local; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + red[threadIdx.x] = fmaxf(red[threadIdx.x], red[threadIdx.x + stride]); + } + __syncthreads(); + } + float amax = fmaxf(red[0], 1e-10f); + float inv = 127.0f / amax; + if (threadIdx.x == 0) { + scale[row] = amax / 127.0f; + } + for (size_t i = threadIdx.x; i < cols; i += blockDim.x) { + int qi = static_cast(nearbyintf(toFloat(x[base + i]) * inv)); + qi = max(-128, min(127, qi)); + q[base + i] = static_cast(qi); + } +} + +#if defined(ENABLE_HYGON_API) +template +__global__ void hygonBf16Kernel(const __nv_bfloat16 *__restrict__ x, int8_t *__restrict__ q, float *__restrict__ scale, int rows, int cols) { + int row = blockIdx.x; + if (row >= rows) { + return; + } + int tid = threadIdx.x; + int off = row * cols; + __shared__ float s_inv; + __shared__ float s_red[BS / 64 + 1]; + float vals[EPT]; + float local_max = 0.0f; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + int idx = tid + i * BS; + if (idx < cols) { + float v = static_cast(x[off + idx]); + vals[i] = v; + local_max = fmaxf(local_max, fabsf(v)); + } + } + int lane = tid & 63; + int wid = tid >> 6; + local_max = warpMax(local_max); + if (lane == 0) { + s_red[wid] = local_max; + } + __syncthreads(); + if (wid == 0) { + local_max = lane < (BS >> 6) ? s_red[lane] : 0.0f; + local_max = warpMax(local_max); + if (lane == 0) { + float amax = fmaxf(local_max, 1e-10f); + s_inv = 127.0f / amax; + scale[row] = amax / 127.0f; + } + } + __syncthreads(); + float inv = s_inv; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + int idx = tid + i * BS; + if (idx < cols) { + int qi = static_cast(roundHe(vals[i] * inv)); + qi = max(-128, min(127, qi)); + q[off + idx] = static_cast(qi); + } + } +} +#endif + +template +infiniStatus_t launchFallback(const op::dsv4_per_token_quant_int8::Info &info, void *q, void *scale, const void *x, cudaStream_t stream) { + fallbackKernel<<>>(static_cast(x), static_cast(q), static_cast(scale), info.rows, info.cols); + return INFINI_STATUS_SUCCESS; +} +} // namespace + +namespace op::dsv4_per_token_quant_int8::nvidia { + +infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc) { + Info info; + CHECK_STATUS(createInfo(&info, q_desc, scale_desc, x_desc)); + *desc_ptr = new Descriptor(info, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate(void *, size_t, void *q, void *scale, const void *x, void *stream) const { + cudaStream_t s = reinterpret_cast(stream); + switch (_info.dtype) { + case INFINI_DTYPE_F16: + return launchFallback(_info, q, scale, x, s); + case INFINI_DTYPE_BF16: +#if defined(ENABLE_HYGON_API) + hygonBf16Kernel<256, 16><<<_info.rows, 256, 0, s>>>(static_cast(x), static_cast(q), static_cast(scale), static_cast(_info.rows), static_cast(_info.cols)); + return INFINI_STATUS_SUCCESS; +#else + return launchFallback<__nv_bfloat16>(_info, q, scale, x, s); +#endif + case INFINI_DTYPE_F32: + return launchFallback(_info, q, scale, x, s); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::dsv4_per_token_quant_int8::nvidia diff --git a/src/infiniop/ops/dsv4_per_token_quant_int8/nvidia/dsv4_per_token_quant_int8_nvidia.cuh b/src/infiniop/ops/dsv4_per_token_quant_int8/nvidia/dsv4_per_token_quant_int8_nvidia.cuh new file mode 100644 index 000000000..35e525fa1 --- /dev/null +++ b/src/infiniop/ops/dsv4_per_token_quant_int8/nvidia/dsv4_per_token_quant_int8_nvidia.cuh @@ -0,0 +1,5 @@ +#ifndef DSV4_PER_TOKEN_QUANT_INT8_NVIDIA_CUH +#define DSV4_PER_TOKEN_QUANT_INT8_NVIDIA_CUH +#include "../dsv4_per_token_quant_int8.h" +DESCRIPTOR(nvidia) +#endif diff --git a/src/infiniop/ops/dsv4_per_token_quant_int8/operator.cc b/src/infiniop/ops/dsv4_per_token_quant_int8/operator.cc new file mode 100644 index 000000000..4df941726 --- /dev/null +++ b/src/infiniop/ops/dsv4_per_token_quant_int8/operator.cc @@ -0,0 +1,117 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_per_token_quant_int8.h" + +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_per_token_quant_int8_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateDsv4PerTokenQuantInt8Descriptor(infiniopHandle_t handle, infiniopDsv4PerTokenQuantInt8Descriptor_t *desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t x_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_per_token_quant_int8::NAMESPACE::Descriptor::create(handle, reinterpret_cast(desc_ptr), q_desc, scale_desc, x_desc) + switch (handle->device) { + +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetDsv4PerTokenQuantInt8WorkspaceSize(infiniopDsv4PerTokenQuantInt8Descriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} + +__INFINI_C infiniStatus_t infiniopDsv4PerTokenQuantInt8(infiniopDsv4PerTokenQuantInt8Descriptor_t desc, void *workspace, size_t workspace_size, void *q, void *scale, const void *x, void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate(workspace, workspace_size, q, scale, x, stream) + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyDsv4PerTokenQuantInt8Descriptor(infiniopDsv4PerTokenQuantInt8Descriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self.h b/src/infiniop/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self.h new file mode 100644 index 000000000..209a31529 --- /dev/null +++ b/src/infiniop/ops/dsv4_rmsnorm_self/dsv4_rmsnorm_self.h @@ -0,0 +1,21 @@ +#ifndef DSV4_RMSNORM_SELF_H +#define DSV4_RMSNORM_SELF_H +#include "../../operator.h" +#include "info.h" +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_rmsnorm_self::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + float _epsilon; \ + size_t _workspace_size; \ + Descriptor(Info info, float epsilon, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id}, _info(info), _epsilon(epsilon), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + float epsilon() const { return _epsilon; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t x_desc, float epsilon); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *y, const void *x, void *stream) const; \ + }; \ + } +#endif diff --git a/src/infiniop/ops/dsv4_rmsnorm_self/info.h b/src/infiniop/ops/dsv4_rmsnorm_self/info.h new file mode 100644 index 000000000..74e36ba01 --- /dev/null +++ b/src/infiniop/ops/dsv4_rmsnorm_self/info.h @@ -0,0 +1,31 @@ +#ifndef DSV4_RMSNORM_SELF_INFO_H +#define DSV4_RMSNORM_SELF_INFO_H +#include "../../../utils.h" +#include "../../tensor.h" +namespace op::dsv4_rmsnorm_self { +struct Info { + infiniDtype_t dtype; + size_t rows; + size_t cols; +}; +inline size_t leadingRows(infiniopTensorDescriptor_t desc) { + size_t rows = 1; + for (size_t i = 0; i + 1 < desc->ndim(); ++i) { + rows *= desc->dim(i); + } + return rows; +} +inline infiniStatus_t createInfo(Info *info, infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t x_desc) { + CHECK_OR_RETURN(info && y_desc && x_desc, INFINI_STATUS_NULL_POINTER); + CHECK_DTYPE(x_desc->dtype(), INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + CHECK_OR_RETURN(y_desc->dtype() == x_desc->dtype(), INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(x_desc->ndim() >= 2 && y_desc->ndim() == x_desc->ndim(), INFINI_STATUS_BAD_TENSOR_SHAPE); + for (size_t i = 0; i < x_desc->ndim(); ++i) { + CHECK_OR_RETURN(y_desc->dim(i) == x_desc->dim(i), INFINI_STATUS_BAD_TENSOR_SHAPE); + } + CHECK_OR_RETURN(x_desc->isContiguous() && y_desc->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + *info = Info{x_desc->dtype(), leadingRows(x_desc), x_desc->dim(x_desc->ndim() - 1)}; + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_rmsnorm_self +#endif diff --git a/src/infiniop/ops/dsv4_rmsnorm_self/nvidia/dsv4_rmsnorm_self_nvidia.cu b/src/infiniop/ops/dsv4_rmsnorm_self/nvidia/dsv4_rmsnorm_self_nvidia.cu new file mode 100644 index 000000000..33537433a --- /dev/null +++ b/src/infiniop/ops/dsv4_rmsnorm_self/nvidia/dsv4_rmsnorm_self_nvidia.cu @@ -0,0 +1,141 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_rmsnorm_self_nvidia.cuh" + +namespace { +constexpr int BLOCK = 256; + +template +__device__ float toFloat(T v) { + return static_cast(v); +} + +template +__device__ T fromFloat(float v) { + return static_cast(v); +} + +__device__ __forceinline__ float warpSum(float v) { +#if defined(ENABLE_HYGON_API) + for (int o = 32; o > 0; o >>= 1) { + v += __shfl_xor(v, o, 64); + } +#else + for (int o = 16; o > 0; o >>= 1) { + v += __shfl_xor_sync(0xffffffff, v, o, 32); + } +#endif + return v; +} + +template +__global__ void fallbackKernel(const T *x, T *y, size_t rows, size_t cols, float eps) { + size_t row = blockIdx.x; + if (row >= rows) { + return; + } + __shared__ float red[BLOCK]; + float local = 0.0f; + size_t base = row * cols; + for (size_t i = threadIdx.x; i < cols; i += blockDim.x) { + float v = toFloat(x[base + i]); + local += v * v; + } + red[threadIdx.x] = local; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + red[threadIdx.x] += red[threadIdx.x + stride]; + } + __syncthreads(); + } + float rrms = rsqrtf(red[0] / static_cast(cols) + eps); + for (size_t i = threadIdx.x; i < cols; i += blockDim.x) { + y[base + i] = fromFloat(toFloat(x[base + i]) * rrms); + } +} + +#if defined(ENABLE_HYGON_API) +template +__global__ void hygonBf16Kernel(const __nv_bfloat16 *x, __nv_bfloat16 *y, int rows, int cols, float eps) { + int row = blockIdx.x; + if (row >= rows) { + return; + } + int tid = threadIdx.x; + int off = row * cols; + __shared__ float s_rrms; + __shared__ float s_red[BS / 64 + 1]; + float vals[EPT]; + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + int idx = tid + i * BS; + if (idx < cols) { + float v = static_cast(x[off + idx]); + vals[i] = v; + ss += v * v; + } + } + int lane = tid & 63; + int wid = tid >> 6; + ss = warpSum(ss); + if (lane == 0) { + s_red[wid] = ss; + } + __syncthreads(); + if (wid == 0) { + ss = lane < (BS >> 6) ? s_red[lane] : 0.0f; + ss = warpSum(ss); + if (lane == 0) { + s_rrms = rsqrtf(ss / static_cast(cols) + eps); + } + } + __syncthreads(); + float rrms = s_rrms; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + int idx = tid + i * BS; + if (idx < cols) { + y[off + idx] = static_cast<__nv_bfloat16>(vals[i] * rrms); + } + } +} +#endif + +template +infiniStatus_t launchFallback(const op::dsv4_rmsnorm_self::Info &info, void *y, const void *x, float eps, cudaStream_t stream) { + fallbackKernel<<>>(static_cast(x), static_cast(y), info.rows, info.cols, eps); + return INFINI_STATUS_SUCCESS; +} +} // namespace + +namespace op::dsv4_rmsnorm_self::nvidia { + +infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t x_desc, float epsilon) { + Info info; + CHECK_STATUS(createInfo(&info, y_desc, x_desc)); + *desc_ptr = new Descriptor(info, epsilon, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate(void *, size_t, void *y, const void *x, void *stream) const { + cudaStream_t s = reinterpret_cast(stream); + switch (_info.dtype) { + case INFINI_DTYPE_F16: + return launchFallback(_info, y, x, _epsilon, s); + case INFINI_DTYPE_BF16: +#if defined(ENABLE_HYGON_API) + hygonBf16Kernel<256, 16><<<_info.rows, 256, 0, s>>>(static_cast(x), static_cast<__nv_bfloat16 *>(y), static_cast(_info.rows), static_cast(_info.cols), _epsilon); + return INFINI_STATUS_SUCCESS; +#else + return launchFallback<__nv_bfloat16>(_info, y, x, _epsilon, s); +#endif + case INFINI_DTYPE_F32: + return launchFallback(_info, y, x, _epsilon, s); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::dsv4_rmsnorm_self::nvidia diff --git a/src/infiniop/ops/dsv4_rmsnorm_self/nvidia/dsv4_rmsnorm_self_nvidia.cuh b/src/infiniop/ops/dsv4_rmsnorm_self/nvidia/dsv4_rmsnorm_self_nvidia.cuh new file mode 100644 index 000000000..f0fbf1053 --- /dev/null +++ b/src/infiniop/ops/dsv4_rmsnorm_self/nvidia/dsv4_rmsnorm_self_nvidia.cuh @@ -0,0 +1,5 @@ +#ifndef DSV4_RMSNORM_SELF_NVIDIA_CUH +#define DSV4_RMSNORM_SELF_NVIDIA_CUH +#include "../dsv4_rmsnorm_self.h" +DESCRIPTOR(nvidia) +#endif diff --git a/src/infiniop/ops/dsv4_rmsnorm_self/operator.cc b/src/infiniop/ops/dsv4_rmsnorm_self/operator.cc new file mode 100644 index 000000000..543446c05 --- /dev/null +++ b/src/infiniop/ops/dsv4_rmsnorm_self/operator.cc @@ -0,0 +1,112 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_rmsnorm_self.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_rmsnorm_self_nvidia.cuh" +#endif +__INFINI_C infiniStatus_t infiniopCreateDsv4RMSNormSelfDescriptor(infiniopHandle_t handle, infiniopDsv4RMSNormSelfDescriptor_t *desc_ptr, infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t x_desc, float epsilon) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_rmsnorm_self::NAMESPACE::Descriptor::create(handle, reinterpret_cast(desc_ptr), y_desc, x_desc, epsilon) + switch (handle->device) { + +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} +__INFINI_C infiniStatus_t infiniopGetDsv4RMSNormSelfWorkspaceSize(infiniopDsv4RMSNormSelfDescriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} +__INFINI_C infiniStatus_t infiniopDsv4RMSNormSelf(infiniopDsv4RMSNormSelfDescriptor_t desc, void *workspace, size_t workspace_size, void *y, const void *x, void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate(workspace, workspace_size, y, x, stream) + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} +__INFINI_C infiniStatus_t infiniopDestroyDsv4RMSNormSelfDescriptor(infiniopDsv4RMSNormSelfDescriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_silu_and_mul/dsv4_silu_and_mul.h b/src/infiniop/ops/dsv4_silu_and_mul/dsv4_silu_and_mul.h new file mode 100644 index 000000000..aa6045e63 --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_and_mul/dsv4_silu_and_mul.h @@ -0,0 +1,19 @@ +#ifndef DSV4_SILU_AND_MUL_H +#define DSV4_SILU_AND_MUL_H +#include "../../operator.h" +#include "info.h" +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_silu_and_mul::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + size_t _workspace_size; \ + Descriptor(Info info, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id}, _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t gate_desc, infiniopTensorDescriptor_t up_desc); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *y, const void *gate, const void *up, void *stream) const; \ + }; \ + } +#endif diff --git a/src/infiniop/ops/dsv4_silu_and_mul/info.h b/src/infiniop/ops/dsv4_silu_and_mul/info.h new file mode 100644 index 000000000..2d0a86ab1 --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_and_mul/info.h @@ -0,0 +1,31 @@ +#ifndef DSV4_SILU_AND_MUL_INFO_H +#define DSV4_SILU_AND_MUL_INFO_H +#include "../../../utils.h" +#include "../../tensor.h" +namespace op::dsv4_silu_and_mul { +struct Info { + infiniDtype_t dtype; + size_t rows; + size_t cols; +}; +inline size_t leadingRows(infiniopTensorDescriptor_t desc) { + size_t rows = 1; + for (size_t i = 0; i + 1 < desc->ndim(); ++i) { + rows *= desc->dim(i); + } + return rows; +} +inline infiniStatus_t createInfo(Info *info, infiniopTensorDescriptor_t y, infiniopTensorDescriptor_t gate, infiniopTensorDescriptor_t up) { + CHECK_OR_RETURN(info && y && gate && up, INFINI_STATUS_NULL_POINTER); + CHECK_DTYPE(gate->dtype(), INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + CHECK_OR_RETURN(y->dtype() == gate->dtype() && up->dtype() == gate->dtype(), INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(gate->ndim() >= 1 && y->ndim() == gate->ndim() && up->ndim() == gate->ndim(), INFINI_STATUS_BAD_TENSOR_SHAPE); + for (size_t i = 0; i < gate->ndim(); ++i) { + CHECK_OR_RETURN(y->dim(i) == gate->dim(i) && up->dim(i) == gate->dim(i), INFINI_STATUS_BAD_TENSOR_SHAPE); + } + CHECK_OR_RETURN(y->isContiguous() && gate->isContiguous() && up->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + *info = Info{gate->dtype(), leadingRows(gate), gate->dim(gate->ndim() - 1)}; + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_silu_and_mul +#endif diff --git a/src/infiniop/ops/dsv4_silu_and_mul/nvidia/dsv4_silu_and_mul_nvidia.cu b/src/infiniop/ops/dsv4_silu_and_mul/nvidia/dsv4_silu_and_mul_nvidia.cu new file mode 100644 index 000000000..d3124f461 --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_and_mul/nvidia/dsv4_silu_and_mul_nvidia.cu @@ -0,0 +1,88 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_silu_and_mul_nvidia.cuh" + +namespace { +constexpr int BLOCK = 256; + +template +__device__ float toFloat(T v) { + return static_cast(v); +} + +template +__device__ T fromFloat(float v) { + return static_cast(v); +} + +template +__global__ void fallbackKernel(const T *gate, const T *up, T *y, size_t total) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + for (; idx < total; idx += stride) { + float g = toFloat(gate[idx]); + float u = toFloat(up[idx]); + y[idx] = fromFloat((1.0f / (1.0f + expf(-g))) * g * u); + } +} + +#if defined(ENABLE_HYGON_API) +template +__global__ void hygonBf16Kernel(const __nv_bfloat16 *gate, const __nv_bfloat16 *up, __nv_bfloat16 *y, int rows, int cols) { + int row = blockIdx.x; + if (row >= rows) { + return; + } + int tid = threadIdx.x; + int off = row * cols; +#pragma unroll + for (int i = 0; i < EPT; ++i) { + int idx = tid + i * BS; + if (idx < cols) { + float g = static_cast(gate[off + idx]); + float u = static_cast(up[off + idx]); + y[off + idx] = static_cast<__nv_bfloat16>((1.0f / (1.0f + expf(-g))) * g * u); + } + } +} +#endif + +template +infiniStatus_t launchFallback(const op::dsv4_silu_and_mul::Info &info, void *y, const void *gate, const void *up, cudaStream_t stream) { + size_t total = info.rows * info.cols; + size_t blocks = (total + BLOCK - 1) / BLOCK; + int grid = static_cast(blocks < 4096 ? blocks : 4096); + fallbackKernel<<>>(static_cast(gate), static_cast(up), static_cast(y), total); + return INFINI_STATUS_SUCCESS; +} +} // namespace + +namespace op::dsv4_silu_and_mul::nvidia { + +infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t gate_desc, infiniopTensorDescriptor_t up_desc) { + Info info; + CHECK_STATUS(createInfo(&info, y_desc, gate_desc, up_desc)); + *desc_ptr = new Descriptor(info, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate(void *, size_t, void *y, const void *gate, const void *up, void *stream) const { + cudaStream_t s = reinterpret_cast(stream); + switch (_info.dtype) { + case INFINI_DTYPE_F16: + return launchFallback(_info, y, gate, up, s); + case INFINI_DTYPE_BF16: +#if defined(ENABLE_HYGON_API) + hygonBf16Kernel<256, 8><<<_info.rows, 256, 0, s>>>(static_cast(gate), static_cast(up), static_cast<__nv_bfloat16 *>(y), static_cast(_info.rows), static_cast(_info.cols)); + return INFINI_STATUS_SUCCESS; +#else + return launchFallback<__nv_bfloat16>(_info, y, gate, up, s); +#endif + case INFINI_DTYPE_F32: + return launchFallback(_info, y, gate, up, s); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::dsv4_silu_and_mul::nvidia diff --git a/src/infiniop/ops/dsv4_silu_and_mul/nvidia/dsv4_silu_and_mul_nvidia.cuh b/src/infiniop/ops/dsv4_silu_and_mul/nvidia/dsv4_silu_and_mul_nvidia.cuh new file mode 100644 index 000000000..2555771f4 --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_and_mul/nvidia/dsv4_silu_and_mul_nvidia.cuh @@ -0,0 +1,5 @@ +#ifndef DSV4_SILU_AND_MUL_NVIDIA_CUH +#define DSV4_SILU_AND_MUL_NVIDIA_CUH +#include "../dsv4_silu_and_mul.h" +DESCRIPTOR(nvidia) +#endif diff --git a/src/infiniop/ops/dsv4_silu_and_mul/operator.cc b/src/infiniop/ops/dsv4_silu_and_mul/operator.cc new file mode 100644 index 000000000..4ab80f8f0 --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_and_mul/operator.cc @@ -0,0 +1,112 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_silu_and_mul.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_silu_and_mul_nvidia.cuh" +#endif +__INFINI_C infiniStatus_t infiniopCreateDsv4SiluAndMulDescriptor(infiniopHandle_t handle, infiniopDsv4SiluAndMulDescriptor_t *desc_ptr, infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t gate_desc, infiniopTensorDescriptor_t up_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_silu_and_mul::NAMESPACE::Descriptor::create(handle, reinterpret_cast(desc_ptr), y_desc, gate_desc, up_desc) + switch (handle->device) { + +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} +__INFINI_C infiniStatus_t infiniopGetDsv4SiluAndMulWorkspaceSize(infiniopDsv4SiluAndMulDescriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} +__INFINI_C infiniStatus_t infiniopDsv4SiluAndMul(infiniopDsv4SiluAndMulDescriptor_t desc, void *workspace, size_t workspace_size, void *y, const void *gate, const void *up, void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate(workspace, workspace_size, y, gate, up, stream) + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} +__INFINI_C infiniStatus_t infiniopDestroyDsv4SiluAndMulDescriptor(infiniopDsv4SiluAndMulDescriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant.h b/src/infiniop/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant.h new file mode 100644 index 000000000..d5c58828b --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_mul_quant/dsv4_silu_mul_quant.h @@ -0,0 +1,19 @@ +#ifndef DSV4_SILU_MUL_QUANT_H +#define DSV4_SILU_MUL_QUANT_H +#include "../../operator.h" +#include "info.h" +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_silu_mul_quant::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + size_t _workspace_size; \ + Descriptor(Info info, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id}, _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t gate_desc, infiniopTensorDescriptor_t up_desc); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *q, void *scale, const void *gate, const void *up, void *stream) const; \ + }; \ + } +#endif diff --git a/src/infiniop/ops/dsv4_silu_mul_quant/info.h b/src/infiniop/ops/dsv4_silu_mul_quant/info.h new file mode 100644 index 000000000..c1f1521a3 --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_mul_quant/info.h @@ -0,0 +1,24 @@ +#ifndef DSV4_SILU_MUL_QUANT_INFO_H +#define DSV4_SILU_MUL_QUANT_INFO_H +#include "../../../utils.h" +#include "../../tensor.h" +namespace op::dsv4_silu_mul_quant { +struct Info { + size_t rows, cols; + infiniDtype_t dtype; +}; +inline infiniStatus_t createInfo(Info *info, infiniopTensorDescriptor_t q, infiniopTensorDescriptor_t scale, infiniopTensorDescriptor_t gate, infiniopTensorDescriptor_t up) { + CHECK_OR_RETURN(info && q && scale && gate && up, INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(q->dtype() == INFINI_DTYPE_I8 && scale->dtype() == INFINI_DTYPE_F32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(gate->dtype() == up->dtype() && (gate->dtype() == INFINI_DTYPE_BF16 || gate->dtype() == INFINI_DTYPE_F16), INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(q->ndim() == gate->ndim() && gate->ndim() == up->ndim() && gate->ndim() >= 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(q->shape() == gate->shape() && up->shape() == gate->shape(), INFINI_STATUS_BAD_TENSOR_SHAPE); + size_t rows = gate->numel() / gate->dim(gate->ndim() - 1); + size_t cols = gate->dim(gate->ndim() - 1); + CHECK_OR_RETURN(scale->numel() == rows, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(q->isContiguous() && scale->isContiguous() && gate->isContiguous() && up->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + *info = Info{rows, cols, gate->dtype()}; + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_silu_mul_quant +#endif diff --git a/src/infiniop/ops/dsv4_silu_mul_quant/nvidia/dsv4_silu_mul_quant_nvidia.cu b/src/infiniop/ops/dsv4_silu_mul_quant/nvidia/dsv4_silu_mul_quant_nvidia.cu new file mode 100644 index 000000000..e2fd6359c --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_mul_quant/nvidia/dsv4_silu_mul_quant_nvidia.cu @@ -0,0 +1,97 @@ + +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_silu_mul_quant_nvidia.cuh" +namespace { +__device__ __forceinline__ float warpMax(float v) { +#if defined(ENABLE_HYGON_API) + for (int o = 32; o > 0; o >>= 1) { + v = fmaxf(v, __shfl_xor(v, o, 64)); + } +#else + for (int o = 16; o > 0; o >>= 1) { + v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, o, 32)); + } +#endif + return v; +} +__device__ __forceinline__ float roundHe(float v) { +#if defined(ENABLE_HYGON_API) + return __builtin_rintf(v); +#else + return nearbyintf(v); +#endif +} +template +__device__ float toFloat(T v) { return static_cast(v); } +template +__global__ void kernel(const T *g, const T *u, int8_t *q, float *s, int rows, int cols) { + int row = blockIdx.x; + if (row >= rows) { + return; + } + int tid = threadIdx.x, off = row * cols; + __shared__ float s_inv, s_red[BS / 64 + 1]; + float vals[EPT], lmax = 0.f; +#pragma unroll + for (int i = 0; i < EPT; i++) { + int idx = tid + i * BS; + if (idx < cols) { + float gv = toFloat(g[off + idx]), uv = toFloat(u[off + idx]); + float h = (1.f / (1.f + expf(-gv))) * gv * uv; + vals[i] = h; + lmax = fmaxf(lmax, fabsf(h)); + } + } + int lane = tid & 63, wid = tid >> 6; + lmax = warpMax(lmax); + if (lane == 0) { + s_red[wid] = lmax; + } + __syncthreads(); + if (wid == 0) { + lmax = (lane < (BS >> 6)) ? s_red[lane] : 0.f; + lmax = warpMax(lmax); + if (lane == 0) { + float am = fmaxf(lmax, 1e-10f); + s_inv = 127.f / am; + s[row] = am / 127.f; + } + } + __syncthreads(); + float inv = s_inv; +#pragma unroll + for (int i = 0; i < EPT; i++) { + int idx = tid + i * BS; + if (idx < cols) { + int qi = (int)roundHe(vals[i] * inv); + qi = max(-128, min(127, qi)); + q[off + idx] = (int8_t)qi; + } + } +} +template +infiniStatus_t launch(const op::dsv4_silu_mul_quant::Info &info, void *q, void *scale, const void *gate, const void *up, cudaStream_t st) { + kernel<256, 8, T><<>>(static_cast(gate), static_cast(up), static_cast(q), static_cast(scale), (int)info.rows, (int)info.cols); + return INFINI_STATUS_SUCCESS; +} +} // namespace +namespace op::dsv4_silu_mul_quant::nvidia { +infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t gate_desc, infiniopTensorDescriptor_t up_desc) { + Info info; + CHECK_STATUS(createInfo(&info, q_desc, scale_desc, gate_desc, up_desc)); + *desc_ptr = new Descriptor(info, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} +infiniStatus_t Descriptor::calculate(void *, size_t, void *q, void *scale, const void *gate, const void *up, void *stream) const { + cudaStream_t st = (cudaStream_t)stream; + switch (_info.dtype) { + case INFINI_DTYPE_BF16: + return launch<__nv_bfloat16>(_info, q, scale, gate, up, st); + case INFINI_DTYPE_F16: + return launch(_info, q, scale, gate, up, st); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} +} // namespace op::dsv4_silu_mul_quant::nvidia diff --git a/src/infiniop/ops/dsv4_silu_mul_quant/nvidia/dsv4_silu_mul_quant_nvidia.cuh b/src/infiniop/ops/dsv4_silu_mul_quant/nvidia/dsv4_silu_mul_quant_nvidia.cuh new file mode 100644 index 000000000..2d7feb35d --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_mul_quant/nvidia/dsv4_silu_mul_quant_nvidia.cuh @@ -0,0 +1,5 @@ +#ifndef DSV4_SILU_MUL_QUANT_NVIDIA_CUH +#define DSV4_SILU_MUL_QUANT_NVIDIA_CUH +#include "../dsv4_silu_mul_quant.h" +DESCRIPTOR(nvidia) +#endif diff --git a/src/infiniop/ops/dsv4_silu_mul_quant/operator.cc b/src/infiniop/ops/dsv4_silu_mul_quant/operator.cc new file mode 100644 index 000000000..485c58285 --- /dev/null +++ b/src/infiniop/ops/dsv4_silu_mul_quant/operator.cc @@ -0,0 +1,108 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_silu_mul_quant.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_silu_mul_quant_nvidia.cuh" +#endif +__INFINI_C infiniStatus_t infiniopCreateDsv4SiluMulQuantDescriptor(infiniopHandle_t handle, infiniopDsv4SiluMulQuantDescriptor_t *desc_ptr, infiniopTensorDescriptor_t q_desc, infiniopTensorDescriptor_t scale_desc, infiniopTensorDescriptor_t gate_desc, infiniopTensorDescriptor_t up_desc) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_silu_mul_quant::NAMESPACE::Descriptor::create(handle, reinterpret_cast(desc_ptr), q_desc, scale_desc, gate_desc, up_desc) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} +__INFINI_C infiniStatus_t infiniopGetDsv4SiluMulQuantWorkspaceSize(infiniopDsv4SiluMulQuantDescriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} +__INFINI_C infiniStatus_t infiniopDsv4SiluMulQuant(infiniopDsv4SiluMulQuantDescriptor_t desc, void *workspace, size_t workspace_size, void *q, void *scale, const void *gate, const void *up, void *stream) { +#define CALC(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate(workspace, workspace_size, q, scale, gate, up, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALC(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALC(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALC(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALC(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALC(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALC +} +__INFINI_C infiniStatus_t infiniopDestroyDsv4SiluMulQuantDescriptor(infiniopDsv4SiluMulQuantDescriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices.h b/src/infiniop/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices.h new file mode 100644 index 000000000..537990c48 --- /dev/null +++ b/src/infiniop/ops/dsv4_swa_prefill_indices/dsv4_swa_prefill_indices.h @@ -0,0 +1,19 @@ +#ifndef DSV4_SWA_PREFILL_INDICES_H +#define DSV4_SWA_PREFILL_INDICES_H +#include "../../operator.h" +#include "info.h" +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_swa_prefill_indices::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + size_t _workspace_size; \ + Descriptor(Info info, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id}, _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t indices_desc, int window_size); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *indices, void *stream) const; \ + }; \ + } +#endif diff --git a/src/infiniop/ops/dsv4_swa_prefill_indices/info.h b/src/infiniop/ops/dsv4_swa_prefill_indices/info.h new file mode 100644 index 000000000..c8a1b1653 --- /dev/null +++ b/src/infiniop/ops/dsv4_swa_prefill_indices/info.h @@ -0,0 +1,19 @@ +#ifndef DSV4_SWA_PREFILL_INDICES_INFO_H +#define DSV4_SWA_PREFILL_INDICES_INFO_H +#include "../../../utils.h" +#include "../../tensor.h" +namespace op::dsv4_swa_prefill_indices { +struct Info { + size_t batch, seq_len; + int window_size; +}; +inline infiniStatus_t createInfo(Info *info, infiniopTensorDescriptor_t indices, int window_size) { + CHECK_OR_RETURN(info && indices, INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(indices->dtype() == INFINI_DTYPE_I32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(indices->ndim() == 2 && indices->isContiguous(), INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(window_size > 0, INFINI_STATUS_BAD_PARAM); + *info = Info{indices->dim(0), indices->dim(1), window_size}; + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_swa_prefill_indices +#endif diff --git a/src/infiniop/ops/dsv4_swa_prefill_indices/nvidia/dsv4_swa_prefill_indices_nvidia.cu b/src/infiniop/ops/dsv4_swa_prefill_indices/nvidia/dsv4_swa_prefill_indices_nvidia.cu new file mode 100644 index 000000000..df3dcfb43 --- /dev/null +++ b/src/infiniop/ops/dsv4_swa_prefill_indices/nvidia/dsv4_swa_prefill_indices_nvidia.cu @@ -0,0 +1,29 @@ + +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_swa_prefill_indices_nvidia.cuh" +namespace { +__global__ void kernel(int32_t *indices, int seq_len, int window_size, int batch) { + int b = blockIdx.x; + if (b >= batch) { + return; + } + int tid = threadIdx.x; + for (int i = tid; i < seq_len; i += blockDim.x) { + int start = max(0, i - window_size + 1); + indices[b * seq_len + i] = start; + } +} +} // namespace +namespace op::dsv4_swa_prefill_indices::nvidia { +infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t indices_desc, int window_size) { + Info info; + CHECK_STATUS(createInfo(&info, indices_desc, window_size)); + *desc_ptr = new Descriptor(info, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} +infiniStatus_t Descriptor::calculate(void *, size_t, void *indices, void *stream) const { + kernel<<<_info.batch, 128, 0, (cudaStream_t)stream>>>(static_cast(indices), (int)_info.seq_len, _info.window_size, (int)_info.batch); + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_swa_prefill_indices::nvidia diff --git a/src/infiniop/ops/dsv4_swa_prefill_indices/nvidia/dsv4_swa_prefill_indices_nvidia.cuh b/src/infiniop/ops/dsv4_swa_prefill_indices/nvidia/dsv4_swa_prefill_indices_nvidia.cuh new file mode 100644 index 000000000..5f76fc8a0 --- /dev/null +++ b/src/infiniop/ops/dsv4_swa_prefill_indices/nvidia/dsv4_swa_prefill_indices_nvidia.cuh @@ -0,0 +1,5 @@ +#ifndef DSV4_SWA_PREFILL_INDICES_NVIDIA_CUH +#define DSV4_SWA_PREFILL_INDICES_NVIDIA_CUH +#include "../dsv4_swa_prefill_indices.h" +DESCRIPTOR(nvidia) +#endif diff --git a/src/infiniop/ops/dsv4_swa_prefill_indices/operator.cc b/src/infiniop/ops/dsv4_swa_prefill_indices/operator.cc new file mode 100644 index 000000000..4f27cf36b --- /dev/null +++ b/src/infiniop/ops/dsv4_swa_prefill_indices/operator.cc @@ -0,0 +1,108 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_swa_prefill_indices.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_swa_prefill_indices_nvidia.cuh" +#endif +__INFINI_C infiniStatus_t infiniopCreateDsv4SwaPrefillIndicesDescriptor(infiniopHandle_t handle, infiniopDsv4SwaPrefillIndicesDescriptor_t *desc_ptr, infiniopTensorDescriptor_t indices_desc, int window_size) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_swa_prefill_indices::NAMESPACE::Descriptor::create(handle, reinterpret_cast(desc_ptr), indices_desc, window_size) + switch (handle->device) { +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} +__INFINI_C infiniStatus_t infiniopGetDsv4SwaPrefillIndicesWorkspaceSize(infiniopDsv4SwaPrefillIndicesDescriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} +__INFINI_C infiniStatus_t infiniopDsv4SwaPrefillIndices(infiniopDsv4SwaPrefillIndicesDescriptor_t desc, void *workspace, size_t workspace_size, void *indices, void *stream) { +#define CALC(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate(workspace, workspace_size, indices, stream) + switch (desc->device_type) { +#ifdef ENABLE_NVIDIA_API + CALC(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALC(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALC(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALC(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALC(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALC +} +__INFINI_C infiniStatus_t infiniopDestroyDsv4SwaPrefillIndicesDescriptor(infiniopDsv4SwaPrefillIndicesDescriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/dsv4_topk_transform/dsv4_topk_transform.h b/src/infiniop/ops/dsv4_topk_transform/dsv4_topk_transform.h new file mode 100644 index 000000000..663ef15c3 --- /dev/null +++ b/src/infiniop/ops/dsv4_topk_transform/dsv4_topk_transform.h @@ -0,0 +1,19 @@ +#ifndef DSV4_TOPK_TRANSFORM_H +#define DSV4_TOPK_TRANSFORM_H +#include "../../operator.h" +#include "info.h" +#define DESCRIPTOR(NAMESPACE) \ + namespace op::dsv4_topk_transform::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + Info _info; \ + size_t _workspace_size; \ + Descriptor(Info info, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id}, _info(info), _workspace_size(workspace_size) {} \ + \ + public: \ + size_t workspaceSize() const { return _workspace_size; } \ + const Info &info() const { return _info; } \ + static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t scores_desc, infiniopTensorDescriptor_t seq_lens_desc, infiniopTensorDescriptor_t page_tables_desc, int page_size); \ + infiniStatus_t calculate(void *workspace, size_t workspace_size, void *out, const void *scores, const void *seq_lens, const void *page_tables, void *stream) const; \ + }; \ + } +#endif diff --git a/src/infiniop/ops/dsv4_topk_transform/info.h b/src/infiniop/ops/dsv4_topk_transform/info.h new file mode 100644 index 000000000..5e0bce5c6 --- /dev/null +++ b/src/infiniop/ops/dsv4_topk_transform/info.h @@ -0,0 +1,23 @@ +#ifndef DSV4_TOPK_TRANSFORM_INFO_H +#define DSV4_TOPK_TRANSFORM_INFO_H +#include "../../../utils.h" +#include "../../tensor.h" +namespace op::dsv4_topk_transform { +struct Info { + size_t batch; + size_t index_topk; + int page_size; +}; +inline infiniStatus_t createInfo(Info *info, infiniopTensorDescriptor_t out, infiniopTensorDescriptor_t scores, infiniopTensorDescriptor_t seq_lens, infiniopTensorDescriptor_t page_tables, int page_size) { + CHECK_OR_RETURN(info && out && scores && seq_lens && page_tables, INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(out->dtype() == INFINI_DTYPE_I32 && seq_lens->dtype() == INFINI_DTYPE_I32 && page_tables->dtype() == INFINI_DTYPE_I32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(scores->dtype() == INFINI_DTYPE_F32, INFINI_STATUS_BAD_TENSOR_DTYPE); + CHECK_OR_RETURN(out->ndim() == 2 && scores->ndim() == 2 && seq_lens->ndim() == 1 && page_tables->ndim() >= 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(out->dim(0) == scores->dim(0) && out->dim(0) == seq_lens->dim(0) && out->dim(0) == page_tables->dim(0), INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(scores->dim(1) == out->dim(1) * 64, INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(out->isContiguous() && scores->isContiguous() && seq_lens->isContiguous() && page_tables->isContiguous(), INFINI_STATUS_BAD_TENSOR_STRIDES); + *info = Info{out->dim(0), out->dim(1), page_size}; + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_topk_transform +#endif diff --git a/src/infiniop/ops/dsv4_topk_transform/nvidia/dsv4_topk_transform_nvidia.cu b/src/infiniop/ops/dsv4_topk_transform/nvidia/dsv4_topk_transform_nvidia.cu new file mode 100644 index 000000000..e1287916c --- /dev/null +++ b/src/infiniop/ops/dsv4_topk_transform/nvidia/dsv4_topk_transform_nvidia.cu @@ -0,0 +1,32 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "dsv4_topk_transform_nvidia.cuh" +namespace { +__global__ void kernel(int32_t *out, const int32_t *seq_lens, size_t batch, size_t index_topk) { + size_t b = blockIdx.x; + if (b >= batch) { + return; + } + int valid = seq_lens[b]; + if (valid < 0) { + valid = 0; + } + size_t n_valid = static_cast(valid) < index_topk ? static_cast(valid) : index_topk; + for (size_t i = threadIdx.x; i < index_topk; i += blockDim.x) { + out[b * index_topk + i] = i < n_valid ? static_cast(i) : -1; + } +} +} // namespace +namespace op::dsv4_topk_transform::nvidia { +infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr, infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t scores_desc, infiniopTensorDescriptor_t seq_lens_desc, infiniopTensorDescriptor_t page_tables_desc, int page_size) { + Info info; + CHECK_STATUS(createInfo(&info, out_desc, scores_desc, seq_lens_desc, page_tables_desc, page_size)); + *desc_ptr = new Descriptor(info, 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} +infiniStatus_t Descriptor::calculate(void *, size_t, void *out, const void *, const void *seq_lens, const void *, void *stream) const { + cudaStream_t s = reinterpret_cast(stream); + kernel<<<_info.batch, 128, 0, s>>>(static_cast(out), static_cast(seq_lens), _info.batch, _info.index_topk); + return INFINI_STATUS_SUCCESS; +} +} // namespace op::dsv4_topk_transform::nvidia diff --git a/src/infiniop/ops/dsv4_topk_transform/nvidia/dsv4_topk_transform_nvidia.cuh b/src/infiniop/ops/dsv4_topk_transform/nvidia/dsv4_topk_transform_nvidia.cuh new file mode 100644 index 000000000..f25f626d3 --- /dev/null +++ b/src/infiniop/ops/dsv4_topk_transform/nvidia/dsv4_topk_transform_nvidia.cuh @@ -0,0 +1,5 @@ +#ifndef DSV4_TOPK_TRANSFORM_NVIDIA_CUH +#define DSV4_TOPK_TRANSFORM_NVIDIA_CUH +#include "../dsv4_topk_transform.h" +DESCRIPTOR(nvidia) +#endif diff --git a/src/infiniop/ops/dsv4_topk_transform/operator.cc b/src/infiniop/ops/dsv4_topk_transform/operator.cc new file mode 100644 index 000000000..6ffbdc7d2 --- /dev/null +++ b/src/infiniop/ops/dsv4_topk_transform/operator.cc @@ -0,0 +1,112 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/dsv4_topk_transform.h" +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/dsv4_topk_transform_nvidia.cuh" +#endif +__INFINI_C infiniStatus_t infiniopCreateDsv4TopkTransformDescriptor(infiniopHandle_t handle, infiniopDsv4TopkTransformDescriptor_t *desc_ptr, infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t scores_desc, infiniopTensorDescriptor_t seq_lens_desc, infiniopTensorDescriptor_t page_tables_desc, int page_size) { +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::dsv4_topk_transform::NAMESPACE::Descriptor::create(handle, reinterpret_cast(desc_ptr), out_desc, scores_desc, seq_lens_desc, page_tables_desc, page_size) + switch (handle->device) { + +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CREATE +} +__INFINI_C infiniStatus_t infiniopGetDsv4TopkTransformWorkspaceSize(infiniopDsv4TopkTransformDescriptor_t desc, size_t *size) { +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET +} +__INFINI_C infiniStatus_t infiniopDsv4TopkTransform(infiniopDsv4TopkTransformDescriptor_t desc, void *workspace, size_t workspace_size, void *out, const void *scores, const void *seq_lens, const void *page_tables, void *stream) { +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc)->calculate(workspace, workspace_size, out, scores, seq_lens, page_tables, stream) + switch (desc->device_type) { + +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} +__INFINI_C infiniStatus_t infiniopDestroyDsv4TopkTransformDescriptor(infiniopDsv4TopkTransformDescriptor_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_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/test/infiniop/_dsv4_common.py b/test/infiniop/_dsv4_common.py new file mode 100644 index 000000000..64276b214 --- /dev/null +++ b/test/infiniop/_dsv4_common.py @@ -0,0 +1,247 @@ +import ctypes +from ctypes import c_float, c_int, c_uint64 + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + InfiniDtypeNames, + TestTensor, + TestWorkspace, + check_error, + infiniopOperatorDescriptor_t, +) + +_TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.BF16] +_QUANT_CASES = [((4, 512),), ((64, 4096),), ((8, 256, 4096),)] +_RMS_CASES = [((1, 1, 512),), ((1, 64, 512),), ((1, 128, 512),)] +_SILU_CASES = [((16, 2048),), ((64, 2048),)] +_ROPE_CASES = [((128, 64, 64),), ((1, 64, 64),), ((64, 8, 64),)] +_TOPK_CASES = [((2, 64),), ((4, 128),)] + + +def _workspace(descriptor, getter, device): + size = c_uint64(0) + check_error(getter(descriptor, ctypes.byref(size))) + return TestWorkspace(size.value, device), size.value + + +def test_quant(handle, device, shape, dtype=InfiniDtype.BF16, sync=None): + print( + f"Testing DSV4 per_token_quant_int8 on {InfiniDeviceNames[device]} shape:{shape} dtype:{InfiniDtypeNames[dtype]}" + ) + x = TestTensor(shape, None, dtype, device) + q = TestTensor(shape, None, InfiniDtype.I8, device, mode="zeros") + rows = x.torch_tensor().numel() // x.torch_tensor().shape[-1] + scale = TestTensor((rows, 1), None, InfiniDtype.F32, device, mode="zeros") + + ref_absmax = ( + x.torch_tensor().float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-10) + ) + ref_q = ( + torch.round(x.torch_tensor().float() * (127.0 / ref_absmax)) + .clamp(-128, 127) + .to(torch.int8) + ) + ref_s = ref_absmax / 127.0 + if sync: + sync() + + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4PerTokenQuantInt8Descriptor( + handle, ctypes.byref(desc), q.descriptor, scale.descriptor, x.descriptor + ) + ) + for tensor in [x, q, scale]: + tensor.destroy_desc() + workspace, workspace_size = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4PerTokenQuantInt8WorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4PerTokenQuantInt8( + desc, + workspace.data(), + workspace_size, + q.data(), + scale.data(), + x.data(), + None, + ) + ) + assert (q.actual_tensor().float() - ref_q.float()).abs().max().item() <= 1.0 + assert torch.allclose( + scale.actual_tensor().reshape(ref_s.shape), ref_s, atol=1e-5, rtol=1e-5 + ) + check_error(LIBINFINIOP.infiniopDestroyDsv4PerTokenQuantInt8Descriptor(desc)) + + +def test_rms(handle, device, shape, dtype=InfiniDtype.BF16, sync=None): + print( + f"Testing DSV4 rmsnorm_self on {InfiniDeviceNames[device]} shape:{shape} dtype:{InfiniDtypeNames[dtype]}" + ) + x = TestTensor(shape, None, dtype, device) + y = TestTensor(shape, None, dtype, device, mode="zeros") + eps = 1e-6 + ref = x.torch_tensor() * torch.rsqrt( + x.torch_tensor().pow(2).mean(-1, keepdim=True) + eps + ) + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4RMSNormSelfDescriptor( + handle, ctypes.byref(desc), y.descriptor, x.descriptor, c_float(eps) + ) + ) + for tensor in [x, y]: + tensor.destroy_desc() + workspace, workspace_size = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4RMSNormSelfWorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4RMSNormSelf( + desc, workspace.data(), workspace_size, y.data(), x.data(), None + ) + ) + assert torch.allclose(y.actual_tensor(), ref, atol=2e-2, rtol=2e-2) + check_error(LIBINFINIOP.infiniopDestroyDsv4RMSNormSelfDescriptor(desc)) + + +def test_silu(handle, device, shape, dtype=InfiniDtype.BF16, sync=None): + print( + f"Testing DSV4 silu_and_mul on {InfiniDeviceNames[device]} shape:{shape} dtype:{InfiniDtypeNames[dtype]}" + ) + gate = TestTensor(shape, None, dtype, device) + up = TestTensor(shape, None, dtype, device) + y = TestTensor(shape, None, dtype, device, mode="zeros") + ref = ( + torch.sigmoid(gate.torch_tensor().float()) + * gate.torch_tensor() + * up.torch_tensor() + ).to(gate.torch_tensor().dtype) + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4SiluAndMulDescriptor( + handle, ctypes.byref(desc), y.descriptor, gate.descriptor, up.descriptor + ) + ) + for tensor in [gate, up, y]: + tensor.destroy_desc() + workspace, workspace_size = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4SiluAndMulWorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4SiluAndMul( + desc, + workspace.data(), + workspace_size, + y.data(), + gate.data(), + up.data(), + None, + ) + ) + assert torch.allclose(y.actual_tensor(), ref, atol=1e-2, rtol=1e-2) + check_error(LIBINFINIOP.infiniopDestroyDsv4SiluAndMulDescriptor(desc)) + + +def test_rope(handle, device, shape, dtype=InfiniDtype.BF16, sync=None): + print( + f"Testing DSV4 fused_rope on {InfiniDeviceNames[device]} shape:{shape} dtype:{InfiniDtypeNames[dtype]}" + ) + q = TestTensor(shape, None, dtype, device) + rope_dim = shape[-1] + freq_real = TestTensor((shape[0], rope_dim // 2), None, InfiniDtype.F32, device) + freq_imag = TestTensor((shape[0], rope_dim // 2), None, InfiniDtype.F32, device) + q_ref = q.torch_tensor().clone() + even = q_ref[..., 0::2].float() + odd = q_ref[..., 1::2].float() + real = freq_real.torch_tensor().unsqueeze(-2) + imag = freq_imag.torch_tensor().unsqueeze(-2) + q_ref[..., 0::2] = (even * real - odd * imag).to(q_ref.dtype) + q_ref[..., 1::2] = (even * imag + odd * real).to(q_ref.dtype) + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4FusedRopeDescriptor( + handle, + ctypes.byref(desc), + q.descriptor, + None, + freq_real.descriptor, + freq_imag.descriptor, + c_int(0), + ) + ) + for tensor in [q, freq_real, freq_imag]: + tensor.destroy_desc() + workspace, workspace_size = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4FusedRopeWorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4FusedRope( + desc, + workspace.data(), + workspace_size, + q.data(), + None, + freq_real.data(), + freq_imag.data(), + None, + ) + ) + assert torch.allclose(q.actual_tensor(), q_ref, atol=1e-2, rtol=1e-2) + check_error(LIBINFINIOP.infiniopDestroyDsv4FusedRopeDescriptor(desc)) + + +def test_topk(handle, device, case, dtype=InfiniDtype.F32, sync=None): + batch, n_valid = case + index_topk = 512 + print( + f"Testing DSV4 topk_transform on {InfiniDeviceNames[device]} batch:{batch} n_valid:{n_valid}" + ) + scores = TestTensor((batch, 64 * index_topk), None, InfiniDtype.F32, device) + seq_lens = TestTensor((batch,), None, InfiniDtype.I32, device, mode="zeros") + seq_lens.set_tensor(torch.full((batch,), n_valid, dtype=torch.int32)) + page_tables = TestTensor((batch, 16), None, InfiniDtype.I32, device, mode="zeros") + out = TestTensor((batch, index_topk), None, InfiniDtype.I32, device, mode="zeros") + ref = torch.full((batch, index_topk), -1, dtype=torch.int32) + ref[:, :n_valid] = torch.arange(n_valid, dtype=torch.int32) + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4TopkTransformDescriptor( + handle, + ctypes.byref(desc), + out.descriptor, + scores.descriptor, + seq_lens.descriptor, + page_tables.descriptor, + c_int(64), + ) + ) + for tensor in [scores, seq_lens, page_tables, out]: + tensor.destroy_desc() + workspace, workspace_size = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4TopkTransformWorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4TopkTransform( + desc, + workspace.data(), + workspace_size, + out.data(), + scores.data(), + seq_lens.data(), + page_tables.data(), + None, + ) + ) + assert torch.equal(out.actual_tensor().cpu(), ref) + check_error(LIBINFINIOP.infiniopDestroyDsv4TopkTransformDescriptor(desc)) diff --git a/test/infiniop/dsv4_act_quant_fp8.py b/test/infiniop/dsv4_act_quant_fp8.py new file mode 100644 index 000000000..fc61a7b9c --- /dev/null +++ b/test/infiniop/dsv4_act_quant_fp8.py @@ -0,0 +1,72 @@ +import ctypes +from ctypes import c_float, c_uint64 + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopOperatorDescriptor_t, + test_operator, +) + + +def _workspace(desc, getter, device): + size = c_uint64(0) + check_error(getter(desc, ctypes.byref(size))) + return TestWorkspace(size.value, device), size.value + + +_CASES = [((4, 512),), ((16, 4096),)] + + +def test_op(handle, device, shape, dtype=InfiniDtype.BF16, sync=None): + print(f"Testing DSV4 act_quant_fp8 on {InfiniDeviceNames[device]} shape:{shape}") + x = TestTensor(shape, None, dtype, device) + xq = TestTensor(shape, None, InfiniDtype.F8, device, mode="zeros") + rows = x.torch_tensor().numel() // shape[-1] + scale = TestTensor((rows, 1), None, InfiniDtype.F32, device, mode="zeros") + fp8_max = 448.0 + absmax = x.torch_tensor().float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) + ref_s = absmax / fp8_max + ref_q = (x.torch_tensor().float() / ref_s).to(torch.float8_e4m3fn) + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4ActQuantFp8Descriptor( + handle, + ctypes.byref(desc), + xq.descriptor, + scale.descriptor, + x.descriptor, + c_float(fp8_max), + ) + ) + for t in [x, xq, scale]: + t.destroy_desc() + ws, wsz = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4ActQuantFp8WorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4ActQuantFp8( + desc, ws.data(), wsz, xq.data(), scale.data(), x.data(), None + ) + ) + assert torch.allclose( + scale.actual_tensor().reshape(ref_s.shape), ref_s, atol=1e-5, rtol=1e-5 + ) + assert torch.allclose(xq.actual_tensor().float(), ref_q.float(), atol=4.0, rtol=0.0) + check_error(LIBINFINIOP.infiniopDestroyDsv4ActQuantFp8Descriptor(desc)) + + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_op, _CASES, [InfiniDtype.BF16]) + print("\033[92mDSV4 act_quant_fp8 Test passed!\033[0m") diff --git a/test/infiniop/dsv4_add_rmsnorm_quant.py b/test/infiniop/dsv4_add_rmsnorm_quant.py new file mode 100644 index 000000000..ec9d665c2 --- /dev/null +++ b/test/infiniop/dsv4_add_rmsnorm_quant.py @@ -0,0 +1,94 @@ +import ctypes +from ctypes import c_float, c_uint64 + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopOperatorDescriptor_t, + test_operator, +) + + +def _workspace(desc, getter, device): + size = c_uint64(0) + check_error(getter(desc, ctypes.byref(size))) + return TestWorkspace(size.value, device), size.value + + +_CASES = [((4, 512),), ((16, 4096),)] + + +def test_op(handle, device, shape, dtype=InfiniDtype.BF16, sync=None): + print( + f"Testing DSV4 add_rmsnorm_quant on {InfiniDeviceNames[device]} shape:{shape}" + ) + res = TestTensor(shape, None, dtype, device) + x = TestTensor(shape, None, dtype, device) + weight = TestTensor((shape[-1],), None, dtype, device) + q = TestTensor(shape, None, InfiniDtype.I8, device, mode="zeros") + rows = res.torch_tensor().numel() // shape[-1] + scale = TestTensor((rows, 1), None, InfiniDtype.F32, device, mode="zeros") + eps = 1e-6 + added = res.torch_tensor().float() + x.torch_tensor().float() + norm = ( + added + * torch.rsqrt(added.pow(2).mean(-1, keepdim=True) + eps) + * weight.torch_tensor().float() + ) + absmax = norm.abs().amax(dim=-1, keepdim=True).clamp(min=1e-10) + ref_q = torch.round(norm * (127.0 / absmax)).clamp(-128, 127).to(torch.int8) + ref_s = absmax / 127.0 + ref_res = added.to(res.torch_tensor().dtype) + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4AddRMSNormQuantDescriptor( + handle, + ctypes.byref(desc), + res.descriptor, + q.descriptor, + scale.descriptor, + x.descriptor, + weight.descriptor, + c_float(eps), + ) + ) + for t in [res, x, weight, q, scale]: + t.destroy_desc() + ws, wsz = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4AddRMSNormQuantWorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4AddRMSNormQuant( + desc, + ws.data(), + wsz, + res.data(), + q.data(), + scale.data(), + x.data(), + weight.data(), + None, + ) + ) + assert torch.allclose(res.actual_tensor(), ref_res, atol=1e-2, rtol=1e-2) + assert (q.actual_tensor().float() - ref_q.float()).abs().max().item() <= 1.0 + assert torch.allclose( + scale.actual_tensor().reshape(ref_s.shape), ref_s, atol=2e-4, rtol=2e-4 + ) + check_error(LIBINFINIOP.infiniopDestroyDsv4AddRMSNormQuantDescriptor(desc)) + + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_op, _CASES, [InfiniDtype.BF16]) + print("\033[92mDSV4 add_rmsnorm_quant Test passed!\033[0m") diff --git a/test/infiniop/dsv4_fused_rope.py b/test/infiniop/dsv4_fused_rope.py new file mode 100644 index 000000000..dbbe07e9e --- /dev/null +++ b/test/infiniop/dsv4_fused_rope.py @@ -0,0 +1,14 @@ +from _dsv4_common import ( + _ROPE_CASES, + _TENSOR_DTYPES, + get_args, + get_test_devices, + test_operator, + test_rope, +) + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_rope, _ROPE_CASES, _TENSOR_DTYPES) + print("\033[92mDSV4 fused_rope Test passed!\033[0m") diff --git a/test/infiniop/dsv4_linear_bf16_fp32.py b/test/infiniop/dsv4_linear_bf16_fp32.py new file mode 100644 index 000000000..6e2d074f9 --- /dev/null +++ b/test/infiniop/dsv4_linear_bf16_fp32.py @@ -0,0 +1,64 @@ +import ctypes +from ctypes import c_uint64 + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopOperatorDescriptor_t, + test_operator, +) + +_TEST_CASES = [((1, 4096, 1024),), ((4, 1024, 2048),), ((16, 4096, 1024),)] + + +def _workspace(descriptor, getter, device): + size = c_uint64(0) + check_error(getter(descriptor, ctypes.byref(size))) + return TestWorkspace(size.value, device), size.value + + +def test_linear(handle, device, case, dtype=InfiniDtype.BF16, sync=None): + m, k, n = case + print( + f"Testing DSV4 linear_bf16_fp32 on {InfiniDeviceNames[device]} m:{m} k:{k} n:{n}" + ) + x = TestTensor((m, k), None, dtype, device) + w = TestTensor((n, k), None, InfiniDtype.F32, device) + y = TestTensor((m, n), None, InfiniDtype.F32, device, mode="zeros") + ref = x.torch_tensor().float() @ w.torch_tensor().T + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4LinearBf16Fp32Descriptor( + handle, ctypes.byref(desc), y.descriptor, x.descriptor, w.descriptor + ) + ) + for tensor in [x, w, y]: + tensor.destroy_desc() + workspace, workspace_size = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4LinearBf16Fp32WorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4LinearBf16Fp32( + desc, workspace.data(), workspace_size, y.data(), x.data(), w.data(), None + ) + ) + assert torch.allclose(y.actual_tensor(), ref, atol=1e-2, rtol=1e-2), ( + (y.actual_tensor() - ref).abs().max().item() + ) + check_error(LIBINFINIOP.infiniopDestroyDsv4LinearBf16Fp32Descriptor(desc)) + + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_linear, _TEST_CASES, [InfiniDtype.BF16]) + print("\033[92mDSV4 linear_bf16_fp32 Test passed!\033[0m") diff --git a/test/infiniop/dsv4_mask_topk_ids.py b/test/infiniop/dsv4_mask_topk_ids.py new file mode 100644 index 000000000..dfe341454 --- /dev/null +++ b/test/infiniop/dsv4_mask_topk_ids.py @@ -0,0 +1,66 @@ +import ctypes +from ctypes import c_uint64 + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopOperatorDescriptor_t, + test_operator, +) + +_TEST_CASES = [((4, 32, 0),), ((4, 32, 2),), ((4, 32, 4),), ((16, 8, 7),)] + + +def _workspace(descriptor, getter, device): + size = c_uint64(0) + check_error(getter(descriptor, ctypes.byref(size))) + return TestWorkspace(size.value, device), size.value + + +def test_mask_topk_ids(handle, device, case, dtype=InfiniDtype.I32, sync=None): + batch, topk, num_non_padded = case + print( + f"Testing DSV4 mask_topk_ids on {InfiniDeviceNames[device]} " + f"batch:{batch} topk:{topk} num_non_padded:{num_non_padded}" + ) + topk_ids = TestTensor( + (batch, topk), None, InfiniDtype.I32, device, randint_low=0, randint_high=256 + ) + num = TestTensor((), None, InfiniDtype.I32, device, mode="zeros") + num.set_tensor(torch.tensor(num_non_padded, dtype=torch.int32)) + ref = topk_ids.torch_tensor().clone() + ref[num_non_padded:] = -1 + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4MaskTopkIdsDescriptor( + handle, ctypes.byref(desc), topk_ids.descriptor, num.descriptor + ) + ) + for tensor in [topk_ids, num]: + tensor.destroy_desc() + workspace, workspace_size = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4MaskTopkIdsWorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4MaskTopkIds( + desc, workspace.data(), workspace_size, topk_ids.data(), num.data(), None + ) + ) + assert torch.equal(topk_ids.actual_tensor().cpu(), ref.cpu()) + check_error(LIBINFINIOP.infiniopDestroyDsv4MaskTopkIdsDescriptor(desc)) + + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_mask_topk_ids, _TEST_CASES, [InfiniDtype.I32]) + print("\033[92mDSV4 mask_topk_ids Test passed!\033[0m") diff --git a/test/infiniop/dsv4_per_token_quant_int8.py b/test/infiniop/dsv4_per_token_quant_int8.py new file mode 100644 index 000000000..4f4b48c8c --- /dev/null +++ b/test/infiniop/dsv4_per_token_quant_int8.py @@ -0,0 +1,14 @@ +from _dsv4_common import ( + _QUANT_CASES, + _TENSOR_DTYPES, + get_args, + get_test_devices, + test_operator, + test_quant, +) + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_quant, _QUANT_CASES, _TENSOR_DTYPES) + print("\033[92mDSV4 per_token_quant_int8 Test passed!\033[0m") diff --git a/test/infiniop/dsv4_rmsnorm_self.py b/test/infiniop/dsv4_rmsnorm_self.py new file mode 100644 index 000000000..1d07fc8fb --- /dev/null +++ b/test/infiniop/dsv4_rmsnorm_self.py @@ -0,0 +1,14 @@ +from _dsv4_common import ( + _RMS_CASES, + _TENSOR_DTYPES, + get_args, + get_test_devices, + test_operator, + test_rms, +) + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_rms, _RMS_CASES, _TENSOR_DTYPES) + print("\033[92mDSV4 rmsnorm_self Test passed!\033[0m") diff --git a/test/infiniop/dsv4_silu_and_mul.py b/test/infiniop/dsv4_silu_and_mul.py new file mode 100644 index 000000000..c99084f3d --- /dev/null +++ b/test/infiniop/dsv4_silu_and_mul.py @@ -0,0 +1,14 @@ +from _dsv4_common import ( + _SILU_CASES, + _TENSOR_DTYPES, + get_args, + get_test_devices, + test_operator, + test_silu, +) + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_silu, _SILU_CASES, _TENSOR_DTYPES) + print("\033[92mDSV4 silu_and_mul Test passed!\033[0m") diff --git a/test/infiniop/dsv4_silu_mul_quant.py b/test/infiniop/dsv4_silu_mul_quant.py new file mode 100644 index 000000000..30b7600b8 --- /dev/null +++ b/test/infiniop/dsv4_silu_mul_quant.py @@ -0,0 +1,77 @@ +import ctypes +from ctypes import c_uint64 + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopOperatorDescriptor_t, + test_operator, +) + + +def _workspace(desc, getter, device): + size = c_uint64(0) + check_error(getter(desc, ctypes.byref(size))) + return TestWorkspace(size.value, device), size.value + + +_CASES = [((16, 2048),), ((64, 2048),)] + + +def test_op(handle, device, shape, dtype=InfiniDtype.BF16, sync=None): + print(f"Testing DSV4 silu_mul_quant on {InfiniDeviceNames[device]} shape:{shape}") + gate = TestTensor(shape, None, dtype, device) + up = TestTensor(shape, None, dtype, device) + q = TestTensor(shape, None, InfiniDtype.I8, device, mode="zeros") + rows = gate.torch_tensor().numel() // shape[-1] + scale = TestTensor((rows, 1), None, InfiniDtype.F32, device, mode="zeros") + h = ( + torch.sigmoid(gate.torch_tensor().float()) + * gate.torch_tensor().float() + * up.torch_tensor().float() + ) + absmax = h.abs().amax(dim=-1, keepdim=True).clamp(min=1e-10) + ref_q = torch.round(h * (127.0 / absmax)).clamp(-128, 127).to(torch.int8) + ref_s = absmax / 127.0 + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4SiluMulQuantDescriptor( + handle, + ctypes.byref(desc), + q.descriptor, + scale.descriptor, + gate.descriptor, + up.descriptor, + ) + ) + for t in [gate, up, q, scale]: + t.destroy_desc() + ws, wsz = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4SiluMulQuantWorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4SiluMulQuant( + desc, ws.data(), wsz, q.data(), scale.data(), gate.data(), up.data(), None + ) + ) + assert (q.actual_tensor().float() - ref_q.float()).abs().max().item() <= 1.0 + assert torch.allclose( + scale.actual_tensor().reshape(ref_s.shape), ref_s, atol=1e-5, rtol=1e-5 + ) + check_error(LIBINFINIOP.infiniopDestroyDsv4SiluMulQuantDescriptor(desc)) + + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_op, _CASES, [InfiniDtype.BF16]) + print("\033[92mDSV4 silu_mul_quant Test passed!\033[0m") diff --git a/test/infiniop/dsv4_swa_prefill_indices.py b/test/infiniop/dsv4_swa_prefill_indices.py new file mode 100644 index 000000000..4245095c0 --- /dev/null +++ b/test/infiniop/dsv4_swa_prefill_indices.py @@ -0,0 +1,62 @@ +import ctypes +from ctypes import c_int32, c_uint64 + +import torch +from libinfiniop import ( + LIBINFINIOP, + InfiniDeviceNames, + InfiniDtype, + TestTensor, + TestWorkspace, + check_error, + get_args, + get_test_devices, + infiniopOperatorDescriptor_t, + test_operator, +) + + +def _workspace(desc, getter, device): + size = c_uint64(0) + check_error(getter(desc, ctypes.byref(size))) + return TestWorkspace(size.value, device), size.value + + +_CASES = [((2, 128, 32),), ((4, 256, 128),)] + + +def test_op(handle, device, case, dtype=InfiniDtype.I32, sync=None): + batch, seq_len, window = case + print( + f"Testing DSV4 swa_prefill_indices on {InfiniDeviceNames[device]} batch:{batch} seq:{seq_len} window:{window}" + ) + indices = TestTensor((batch, seq_len), None, InfiniDtype.I32, device, mode="zeros") + ref = torch.empty((batch, seq_len), dtype=torch.int32) + for i in range(seq_len): + ref[:, i] = max(0, i - window + 1) + if sync: + sync() + desc = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateDsv4SwaPrefillIndicesDescriptor( + handle, ctypes.byref(desc), indices.descriptor, c_int32(window) + ) + ) + indices.destroy_desc() + ws, wsz = _workspace( + desc, LIBINFINIOP.infiniopGetDsv4SwaPrefillIndicesWorkspaceSize, device + ) + check_error( + LIBINFINIOP.infiniopDsv4SwaPrefillIndices( + desc, ws.data(), wsz, indices.data(), None + ) + ) + assert torch.equal(indices.actual_tensor().cpu(), ref) + check_error(LIBINFINIOP.infiniopDestroyDsv4SwaPrefillIndicesDescriptor(desc)) + + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_op, _CASES, [InfiniDtype.I32]) + print("\033[92mDSV4 swa_prefill_indices Test passed!\033[0m") diff --git a/test/infiniop/dsv4_topk_transform.py b/test/infiniop/dsv4_topk_transform.py new file mode 100644 index 000000000..4dda87d28 --- /dev/null +++ b/test/infiniop/dsv4_topk_transform.py @@ -0,0 +1,14 @@ +from _dsv4_common import ( + _TOPK_CASES, + InfiniDtype, + get_args, + get_test_devices, + test_operator, + test_topk, +) + +if __name__ == "__main__": + args = get_args() + for device in get_test_devices(args): + test_operator(device, test_topk, _TOPK_CASES, [InfiniDtype.F32]) + print("\033[92mDSV4 topk_transform Test passed!\033[0m") diff --git a/test/infiniop/libinfiniop/op_register.py b/test/infiniop/libinfiniop/op_register.py index a45241b56..c53f17920 100644 --- a/test/infiniop/libinfiniop/op_register.py +++ b/test/infiniop/libinfiniop/op_register.py @@ -3125,3 +3125,326 @@ def nsa_paged_attention_(lib): lib.infiniopDestroyNsaPagedAttentionDescriptor.argtypes = [ infiniopOperatorDescriptor_t, ] + + +@OpRegister.operator +def dsv4_ops_(lib): + lib.infiniopCreateDsv4PerTokenQuantInt8Descriptor.restype = c_int32 + lib.infiniopCreateDsv4PerTokenQuantInt8Descriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + ] + lib.infiniopGetDsv4PerTokenQuantInt8WorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4PerTokenQuantInt8WorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4PerTokenQuantInt8.restype = c_int32 + lib.infiniopDsv4PerTokenQuantInt8.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4PerTokenQuantInt8Descriptor.restype = c_int32 + lib.infiniopDestroyDsv4PerTokenQuantInt8Descriptor.argtypes = [ + infiniopOperatorDescriptor_t + ] + + lib.infiniopCreateDsv4RMSNormSelfDescriptor.restype = c_int32 + lib.infiniopCreateDsv4RMSNormSelfDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + c_float, + ] + lib.infiniopGetDsv4RMSNormSelfWorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4RMSNormSelfWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4RMSNormSelf.restype = c_int32 + lib.infiniopDsv4RMSNormSelf.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4RMSNormSelfDescriptor.restype = c_int32 + lib.infiniopDestroyDsv4RMSNormSelfDescriptor.argtypes = [ + infiniopOperatorDescriptor_t + ] + + lib.infiniopCreateDsv4SiluAndMulDescriptor.restype = c_int32 + lib.infiniopCreateDsv4SiluAndMulDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + ] + lib.infiniopGetDsv4SiluAndMulWorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4SiluAndMulWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4SiluAndMul.restype = c_int32 + lib.infiniopDsv4SiluAndMul.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4SiluAndMulDescriptor.restype = c_int32 + lib.infiniopDestroyDsv4SiluAndMulDescriptor.argtypes = [ + infiniopOperatorDescriptor_t + ] + + lib.infiniopCreateDsv4TopkTransformDescriptor.restype = c_int32 + lib.infiniopCreateDsv4TopkTransformDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + c_int32, + ] + lib.infiniopGetDsv4TopkTransformWorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4TopkTransformWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4TopkTransform.restype = c_int32 + lib.infiniopDsv4TopkTransform.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4TopkTransformDescriptor.restype = c_int32 + lib.infiniopDestroyDsv4TopkTransformDescriptor.argtypes = [ + infiniopOperatorDescriptor_t + ] + + lib.infiniopCreateDsv4FusedRopeDescriptor.restype = c_int32 + lib.infiniopCreateDsv4FusedRopeDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + c_int32, + ] + lib.infiniopGetDsv4FusedRopeWorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4FusedRopeWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4FusedRope.restype = c_int32 + lib.infiniopDsv4FusedRope.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4FusedRopeDescriptor.restype = c_int32 + lib.infiniopDestroyDsv4FusedRopeDescriptor.argtypes = [infiniopOperatorDescriptor_t] + + +@OpRegister.operator +def dsv4_more_ops_(lib): + lib.infiniopCreateDsv4MaskTopkIdsDescriptor.restype = c_int32 + lib.infiniopCreateDsv4MaskTopkIdsDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + ] + lib.infiniopGetDsv4MaskTopkIdsWorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4MaskTopkIdsWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4MaskTopkIds.restype = c_int32 + lib.infiniopDsv4MaskTopkIds.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4MaskTopkIdsDescriptor.restype = c_int32 + lib.infiniopDestroyDsv4MaskTopkIdsDescriptor.argtypes = [ + infiniopOperatorDescriptor_t, + ] + + lib.infiniopCreateDsv4LinearBf16Fp32Descriptor.restype = c_int32 + lib.infiniopCreateDsv4LinearBf16Fp32Descriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + ] + lib.infiniopGetDsv4LinearBf16Fp32WorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4LinearBf16Fp32WorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4LinearBf16Fp32.restype = c_int32 + lib.infiniopDsv4LinearBf16Fp32.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4LinearBf16Fp32Descriptor.restype = c_int32 + lib.infiniopDestroyDsv4LinearBf16Fp32Descriptor.argtypes = [ + infiniopOperatorDescriptor_t, + ] + + +@OpRegister.operator +def dsv4_fused_quant_ops_(lib): + lib.infiniopCreateDsv4SiluMulQuantDescriptor.restype = c_int32 + lib.infiniopCreateDsv4SiluMulQuantDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + ] + lib.infiniopGetDsv4SiluMulQuantWorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4SiluMulQuantWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4SiluMulQuant.restype = c_int32 + lib.infiniopDsv4SiluMulQuant.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4SiluMulQuantDescriptor.restype = c_int32 + lib.infiniopDestroyDsv4SiluMulQuantDescriptor.argtypes = [ + infiniopOperatorDescriptor_t + ] + + lib.infiniopCreateDsv4AddRMSNormQuantDescriptor.restype = c_int32 + lib.infiniopCreateDsv4AddRMSNormQuantDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + c_float, + ] + lib.infiniopGetDsv4AddRMSNormQuantWorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4AddRMSNormQuantWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4AddRMSNormQuant.restype = c_int32 + lib.infiniopDsv4AddRMSNormQuant.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4AddRMSNormQuantDescriptor.restype = c_int32 + lib.infiniopDestroyDsv4AddRMSNormQuantDescriptor.argtypes = [ + infiniopOperatorDescriptor_t + ] + + lib.infiniopCreateDsv4ActQuantFp8Descriptor.restype = c_int32 + lib.infiniopCreateDsv4ActQuantFp8Descriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + infiniopTensorDescriptor_t, + c_float, + ] + lib.infiniopGetDsv4ActQuantFp8WorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4ActQuantFp8WorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4ActQuantFp8.restype = c_int32 + lib.infiniopDsv4ActQuantFp8.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4ActQuantFp8Descriptor.restype = c_int32 + lib.infiniopDestroyDsv4ActQuantFp8Descriptor.argtypes = [ + infiniopOperatorDescriptor_t + ] + + lib.infiniopCreateDsv4SwaPrefillIndicesDescriptor.restype = c_int32 + lib.infiniopCreateDsv4SwaPrefillIndicesDescriptor.argtypes = [ + infiniopHandle_t, + POINTER(infiniopOperatorDescriptor_t), + infiniopTensorDescriptor_t, + c_int32, + ] + lib.infiniopGetDsv4SwaPrefillIndicesWorkspaceSize.restype = c_int32 + lib.infiniopGetDsv4SwaPrefillIndicesWorkspaceSize.argtypes = [ + infiniopOperatorDescriptor_t, + POINTER(c_size_t), + ] + lib.infiniopDsv4SwaPrefillIndices.restype = c_int32 + lib.infiniopDsv4SwaPrefillIndices.argtypes = [ + infiniopOperatorDescriptor_t, + c_void_p, + c_size_t, + c_void_p, + c_void_p, + ] + lib.infiniopDestroyDsv4SwaPrefillIndicesDescriptor.restype = c_int32 + lib.infiniopDestroyDsv4SwaPrefillIndicesDescriptor.argtypes = [ + infiniopOperatorDescriptor_t + ]