From 50a3ca101f63307e130f868c7409985be99b2e21 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 3 Jul 2026 10:33:21 +0800 Subject: [PATCH 01/16] Add Hygon MoE operator support --- .../hygon/moe_fused_dense_hygon.cc | 104 ++++++++++++++++++ .../ops/moe_align/nvidia/moe_align_nvidia.cu | 6 +- src/infiniop/ops/moe_align/operator.cc | 14 ++- .../nvidia/moe_fused_dense_nvidia.cu | 4 +- src/infiniop/ops/moe_fused_dense/operator.cc | 14 ++- .../nvidia/moe_fused_gate_nvidia.cu | 18 ++- src/infiniop/ops/moe_fused_gate/operator.cc | 14 ++- .../ops/moe_sum/nvidia/moe_sum_nvidia.cu | 4 +- src/infiniop/ops/moe_sum/operator.cc | 14 ++- .../nvidia/moe_topk_sigmoid_nvidia.cu | 16 ++- src/infiniop/ops/moe_topk_sigmoid/operator.cc | 14 ++- .../nvidia/moe_topk_softmax_nvidia.cu | 32 +++++- src/infiniop/ops/moe_topk_softmax/operator.cc | 14 ++- .../nvidia/prepare_moe_input_nvidia.cu | 4 +- .../ops/prepare_moe_input/operator.cc | 14 ++- 15 files changed, 263 insertions(+), 23 deletions(-) create mode 100644 src/infinicore/ops/moe_fused_dense/hygon/moe_fused_dense_hygon.cc diff --git a/src/infinicore/ops/moe_fused_dense/hygon/moe_fused_dense_hygon.cc b/src/infinicore/ops/moe_fused_dense/hygon/moe_fused_dense_hygon.cc new file mode 100644 index 000000000..592edec60 --- /dev/null +++ b/src/infinicore/ops/moe_fused_dense/hygon/moe_fused_dense_hygon.cc @@ -0,0 +1,104 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/ops/moe_fused_dense.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include +#include +#include + +namespace infinicore::op::moe_fused_dense_impl::hygon { + +struct PlannedMeta { + graph::GraphTensor output; + graph::GraphTensor hidden_states; + graph::GraphTensor w13; + graph::GraphTensor w2; + graph::GraphTensor topk_weights; + graph::GraphTensor topk_ids; +}; + +void *plan(Tensor output, + const Tensor &hidden_states, + const Tensor &w13, + const Tensor &w2, + const Tensor &topk_weights, + const Tensor &topk_ids, + const Tensor &, + const Tensor &, + const Tensor &) { + return new PlannedMeta{ + graph::GraphTensor(output), + graph::GraphTensor(hidden_states), + graph::GraphTensor(w13), + graph::GraphTensor(w2), + graph::GraphTensor(topk_weights), + graph::GraphTensor(topk_ids)}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + const bool output_need_copy_back = !p->output->is_contiguous(); + Tensor output_work_ic = output_need_copy_back ? p->output->contiguous() : Tensor(p->output); + + auto output = infinicore::adaptor::to_aten_tensor(output_work_ic); + auto hidden_states = infinicore::adaptor::to_aten_tensor(p->hidden_states); + auto w13 = infinicore::adaptor::to_aten_tensor(p->w13); + auto w2 = infinicore::adaptor::to_aten_tensor(p->w2); + auto topk_weights = infinicore::adaptor::to_aten_tensor(p->topk_weights); + auto topk_ids = infinicore::adaptor::to_aten_tensor(p->topk_ids); + + const int64_t num_experts = w13.size(0); + const int64_t intermediate_size = w2.size(2); + const int64_t topk = topk_ids.size(1); + + auto result = at::zeros_like(output); + auto topk_ids_i64 = topk_ids.to(at::kLong); + + for (int64_t k = 0; k < topk; ++k) { + auto ids_k = topk_ids_i64.select(1, k); + for (int64_t expert = 0; expert < num_experts; ++expert) { + auto token_indices = at::nonzero(ids_k == expert).flatten(); + if (token_indices.numel() == 0) { + continue; + } + + auto hidden = hidden_states.index_select(0, token_indices); + auto w13_e = w13.select(0, expert); + auto gate_up = at::matmul(hidden, w13_e.transpose(0, 1)); + auto gate = gate_up.narrow(1, 0, intermediate_size); + auto up = gate_up.narrow(1, intermediate_size, intermediate_size); + auto activated = (gate / (1 + at::exp(-gate))) * up; + + auto w2_e = w2.select(0, expert); + auto expert_out = at::matmul(activated, w2_e.transpose(0, 1)); + auto weights = topk_weights.select(1, k) + .index_select(0, token_indices) + .to(expert_out.scalar_type()) + .unsqueeze(1); + result.index_add_(0, token_indices, expert_out * weights); + } + } + + output.copy_(result); + if (output_need_copy_back) { + p->output->copy_from(output_work_ic); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + MoeFusedDense::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); + MoeFusedDense::run_dispatcher().registerDevice(Device::Type::HYGON, &run); + MoeFusedDense::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); + return true; +}(); + +} // namespace infinicore::op::moe_fused_dense_impl::hygon +#endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/src/infiniop/ops/moe_align/nvidia/moe_align_nvidia.cu b/src/infiniop/ops/moe_align/nvidia/moe_align_nvidia.cu index 0b96a4d43..8a0222a4f 100644 --- a/src/infiniop/ops/moe_align/nvidia/moe_align_nvidia.cu +++ b/src/infiniop/ops/moe_align/nvidia/moe_align_nvidia.cu @@ -7,7 +7,7 @@ * Licensed under the Apache License, Version 2.0. */ -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) #include "moe_align_nvidia.cuh" @@ -71,7 +71,7 @@ size_t next_pow2(size_t value) { } template -constexpr T ceil_div(T a, T b) { +__host__ __device__ constexpr T ceil_div(T a, T b) { return (a + b - 1) / b; } @@ -415,4 +415,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::moe_align::nvidia -#endif // ENABLE_NVIDIA_API +#endif // ENABLE_NVIDIA_API || ENABLE_HYGON_API diff --git a/src/infiniop/ops/moe_align/operator.cc b/src/infiniop/ops/moe_align/operator.cc index 908aa9de6..13b30486d 100644 --- a/src/infiniop/ops/moe_align/operator.cc +++ b/src/infiniop/ops/moe_align/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/moe_align.h" -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) #include "nvidia/moe_align_nvidia.cuh" #endif @@ -31,6 +31,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeAlignDescriptor( switch (handle->device) { #ifdef ENABLE_NVIDIA_API CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -51,6 +54,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeAlignWorkspaceSize( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -80,6 +86,9 @@ __INFINI_C infiniStatus_t infiniopMoeAlign( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -99,6 +108,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeAlignDescriptor( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu b/src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu index 0b6357f9f..3c2978bb5 100644 --- a/src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu +++ b/src/infiniop/ops/moe_fused_dense/nvidia/moe_fused_dense_nvidia.cu @@ -1,4 +1,4 @@ -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) #include "moe_fused_dense_nvidia.cuh" @@ -700,4 +700,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::moe_fused_dense::nvidia -#endif // ENABLE_NVIDIA_API +#endif // ENABLE_NVIDIA_API || ENABLE_HYGON_API diff --git a/src/infiniop/ops/moe_fused_dense/operator.cc b/src/infiniop/ops/moe_fused_dense/operator.cc index a4f295212..0e416fe76 100644 --- a/src/infiniop/ops/moe_fused_dense/operator.cc +++ b/src/infiniop/ops/moe_fused_dense/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/moe_fused_dense.h" -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) #include "nvidia/moe_fused_dense_nvidia.cuh" #endif @@ -27,6 +27,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeFusedDenseDescriptor( switch (handle->device) { #ifdef ENABLE_NVIDIA_API CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -44,6 +47,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeFusedDenseWorkspaceSize( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -73,6 +79,9 @@ __INFINI_C infiniStatus_t infiniopMoeFusedDense( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -89,6 +98,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeFusedDenseDescriptor( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/moe_fused_gate/nvidia/moe_fused_gate_nvidia.cu b/src/infiniop/ops/moe_fused_gate/nvidia/moe_fused_gate_nvidia.cu index c12561919..eb9bc8785 100644 --- a/src/infiniop/ops/moe_fused_gate/nvidia/moe_fused_gate_nvidia.cu +++ b/src/infiniop/ops/moe_fused_gate/nvidia/moe_fused_gate_nvidia.cu @@ -7,7 +7,7 @@ * Licensed under the Apache License, Version 2.0. */ -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) #include "moe_fused_gate_nvidia.cuh" @@ -94,7 +94,9 @@ __device__ void moe_fused_gate_impl( float row_chunk[MAX_VPT]; float bias_chunk[MAX_VPT]; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < MAX_VPT; ++ii) { if (ii < params.VPT) { const int expert = first_elt_read_by_thread + ii; @@ -110,7 +112,9 @@ __device__ void moe_fused_gate_impl( int expert = first_elt_read_by_thread; float max_val = -FLT_MAX; float max_val_second = -FLT_MAX; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < MAX_VPT; ++ii) { if (ii < params.VPT) { const float val = bias_chunk[ii]; @@ -124,7 +128,9 @@ __device__ void moe_fused_gate_impl( } float max_sum = max_val + max_val_second; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int mask = params.THREADS_PER_ROW / 2; mask > 0; mask /= 2) { const float other_max_sum = __shfl_xor_sync(0xffffffff, max_sum, mask, params.THREADS_PER_ROW); const int other_expert = __shfl_xor_sync(0xffffffff, expert, mask, params.THREADS_PER_ROW); @@ -136,7 +142,9 @@ __device__ void moe_fused_gate_impl( const int thread_to_clear_in_group = expert / params.VPT; if (thread_group_idx == thread_to_clear_in_group) { +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < MAX_VPT; ++ii) { if (ii < params.VPT) { bias_chunk[ii] = FLT_MAX; @@ -152,7 +160,9 @@ __device__ void moe_fused_gate_impl( float max_val = bias_chunk[0]; int expert = first_elt_read_by_thread; if (max_val != FLT_MAX) { +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 1; ii < MAX_VPT; ++ii) { if (ii < params.VPT) { const float val = bias_chunk[ii]; @@ -166,7 +176,9 @@ __device__ void moe_fused_gate_impl( max_val = -FLT_MAX; } +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int mask = params.THREADS_PER_ROW / 2; mask > 0; mask /= 2) { const float other_max = __shfl_xor_sync(0xffffffff, max_val, mask, params.THREADS_PER_ROW); const int other_expert = __shfl_xor_sync(0xffffffff, expert, mask, params.THREADS_PER_ROW); @@ -201,7 +213,9 @@ __device__ void moe_fused_gate_impl( __syncthreads(); if (thread_group_idx == 0) { +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < topk; ++ii) { const int64_t idx = topk * thread_row + ii; output[idx] = output[idx] / output_sum; @@ -411,4 +425,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::moe_fused_gate::nvidia -#endif // ENABLE_NVIDIA_API +#endif // ENABLE_NVIDIA_API || ENABLE_HYGON_API diff --git a/src/infiniop/ops/moe_fused_gate/operator.cc b/src/infiniop/ops/moe_fused_gate/operator.cc index af2089bbe..afe06749f 100644 --- a/src/infiniop/ops/moe_fused_gate/operator.cc +++ b/src/infiniop/ops/moe_fused_gate/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/moe_fused_gate.h" -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) #include "nvidia/moe_fused_gate_nvidia.cuh" #endif @@ -28,6 +28,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeFusedGateDescriptor( switch (handle->device) { #ifdef ENABLE_NVIDIA_API CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -45,6 +48,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeFusedGateWorkspaceSize( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -68,6 +74,9 @@ __INFINI_C infiniStatus_t infiniopMoeFusedGate( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -84,6 +93,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeFusedGateDescriptor( switch (desc->device_type) { #ifdef ENABLE_NVIDIA_API DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/moe_sum/nvidia/moe_sum_nvidia.cu b/src/infiniop/ops/moe_sum/nvidia/moe_sum_nvidia.cu index 2e24cb358..d88852394 100644 --- a/src/infiniop/ops/moe_sum/nvidia/moe_sum_nvidia.cu +++ b/src/infiniop/ops/moe_sum/nvidia/moe_sum_nvidia.cu @@ -1,4 +1,4 @@ -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "moe_sum_nvidia.cuh" @@ -117,4 +117,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::moe_sum::nvidia -#endif // ENABLE_NVIDIA_API || ENABLE_ILUVATAR_API +#endif // ENABLE_NVIDIA_API || ENABLE_ILUVATAR_API || ENABLE_HYGON_API diff --git a/src/infiniop/ops/moe_sum/operator.cc b/src/infiniop/ops/moe_sum/operator.cc index 175117587..7d8e57a82 100644 --- a/src/infiniop/ops/moe_sum/operator.cc +++ b/src/infiniop/ops/moe_sum/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/moe_sum.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/moe_sum_nvidia.cuh" #endif @@ -33,6 +33,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeSumDescriptor( #endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -59,6 +62,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeSumWorkspaceSize( #endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -89,6 +95,9 @@ __INFINI_C infiniStatus_t infiniopMoeSum( #endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -114,6 +123,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeSumDescriptor( #endif #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu b/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu index a2bbf298c..1a6cd52d2 100644 --- a/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu +++ b/src/infiniop/ops/moe_topk_sigmoid/nvidia/moe_topk_sigmoid_nvidia.cu @@ -7,7 +7,7 @@ * Licensed under the Apache License, Version 2.0. */ -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "moe_topk_sigmoid_nvidia.cuh" @@ -192,13 +192,17 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSigmoid( T row_chunk_temp[VPT]; auto *row_chunk_vec_ptr = reinterpret_cast(&row_chunk_temp); const auto *vec_thread_read_ptr = reinterpret_cast(thread_read_ptr); +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; } float row_chunk[VPT]; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < VPT; ++ii) { float val = convert_to_float(row_chunk_temp[ii]); val = 1.0f / (1.0f + expf(-val)); @@ -216,9 +220,13 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSigmoid( for (int k_idx = 0; k_idx < k; ++k_idx) { float max_val = row_chunk[0]; int expert = start_col; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG) { +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { float val = row_chunk[ldg * ELTS_PER_LDG + ii]; if (val > max_val) { @@ -228,7 +236,9 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSigmoid( } } +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { float other_max = __shfl_xor_sync(0xffffffff, max_val, mask, THREADS_PER_ROW); int other_expert = __shfl_xor_sync(0xffffffff, expert, mask, THREADS_PER_ROW); @@ -262,7 +272,9 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSigmoid( if (renormalize && thread_group_idx == 0) { const float inv = 1.0f / row_sum_for_renormalize; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int k_idx = 0; k_idx < k; ++k_idx) { const int idx = k * thread_row + k_idx; output[idx] *= inv; @@ -464,4 +476,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::moe_topk_sigmoid::nvidia -#endif // ENABLE_NVIDIA_API || ENABLE_ILUVATAR_API +#endif // ENABLE_NVIDIA_API || ENABLE_ILUVATAR_API || ENABLE_HYGON_API diff --git a/src/infiniop/ops/moe_topk_sigmoid/operator.cc b/src/infiniop/ops/moe_topk_sigmoid/operator.cc index 97f255c70..a8744ee40 100644 --- a/src/infiniop/ops/moe_topk_sigmoid/operator.cc +++ b/src/infiniop/ops/moe_topk_sigmoid/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/moe_topk_sigmoid.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/moe_topk_sigmoid_nvidia.cuh" #endif @@ -25,6 +25,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeTopkSigmoidDescriptor( #endif #ifdef ENABLE_ILUVATAR_API CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -45,6 +48,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeTopkSigmoidWorkspaceSize( #endif #ifdef ENABLE_ILUVATAR_API GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -71,6 +77,9 @@ __INFINI_C infiniStatus_t infiniopMoeTopkSigmoid( #endif #ifdef ENABLE_ILUVATAR_API CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -90,6 +99,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeTopkSigmoidDescriptor( #endif #ifdef ENABLE_ILUVATAR_API DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu b/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu index 952b3ca5d..402dac7f6 100644 --- a/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu +++ b/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu @@ -7,7 +7,7 @@ * Licensed under the Apache License, Version 2.0. */ -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "moe_topk_softmax_nvidia.cuh" @@ -181,7 +181,9 @@ __launch_bounds__(TPB) __global__ void moeTopKFast( TopKPairArgMax reducer; const TopKPair result_pair = BlockReduce(tmp_storage).Reduce(thread_pair, reducer); if (threadIdx.x == 0) { +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int i = 0; i < TopKPair::PAIR; ++i) { if (k_idx * 2 + i >= k) { break; @@ -308,19 +310,25 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( T row_chunk_temp[VPT]; auto *row_chunk_vec_ptr = reinterpret_cast(&row_chunk_temp); const auto *vec_thread_read_ptr = reinterpret_cast(thread_read_ptr); +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; } float row_chunk[VPT]; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < VPT; ++ii) { row_chunk[ii] = convert_to_float(row_chunk_temp[ii]); } - if (moe_softcapping != 0.0f) { + if (moe_softcapping != 0.0f || correction_bias != nullptr) { +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < VPT; ++ii) { float val = row_chunk[ii]; if (moe_softcapping != 0.0f) { @@ -331,27 +339,37 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( } float thread_max = row_chunk[0]; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 1; ii < VPT; ++ii) { thread_max = fmaxf(thread_max, row_chunk[ii]); } +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { thread_max = fmaxf(thread_max, __shfl_xor_sync(0xffffffff, thread_max, mask, THREADS_PER_ROW)); } float row_sum = 0.0f; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < VPT; ++ii) { row_chunk[ii] = expf(row_chunk[ii] - thread_max); row_sum += row_chunk[ii]; } +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { row_sum += __shfl_xor_sync(0xffffffff, row_sum, mask, THREADS_PER_ROW); } const float reciprocal_row_sum = 1.0f / row_sum; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < VPT; ++ii) { row_chunk[ii] *= reciprocal_row_sum; } @@ -362,9 +380,13 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( float max_prob = row_chunk[0]; float max_choice = correction_bias == nullptr ? max_prob : max_prob + correction_bias[start_col]; int expert = start_col; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG) { +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { const int expert_idx = col + ii; float prob = row_chunk[ldg * ELTS_PER_LDG + ii]; @@ -377,7 +399,9 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( } } +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { float other_choice = __shfl_xor_sync(0xffffffff, max_choice, mask, THREADS_PER_ROW); float other_prob = __shfl_xor_sync(0xffffffff, max_prob, mask, THREADS_PER_ROW); @@ -408,7 +432,9 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( if (renormalize && thread_group_idx == 0) { const float inv = 1.0f / row_sum_for_renormalize; +#ifndef ENABLE_HYGON_API #pragma unroll +#endif for (int k_idx = 0; k_idx < k; ++k_idx) { const int idx = k * thread_row + k_idx; output[idx] *= inv; @@ -618,4 +644,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::moe_topk_softmax::nvidia -#endif // ENABLE_NVIDIA_API || ENABLE_ILUVATAR_API +#endif // ENABLE_NVIDIA_API || ENABLE_ILUVATAR_API || ENABLE_HYGON_API diff --git a/src/infiniop/ops/moe_topk_softmax/operator.cc b/src/infiniop/ops/moe_topk_softmax/operator.cc index 213ad93ce..8cf30478d 100644 --- a/src/infiniop/ops/moe_topk_softmax/operator.cc +++ b/src/infiniop/ops/moe_topk_softmax/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/moe_topk_softmax.h" -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) #include "nvidia/moe_topk_softmax_nvidia.cuh" #endif @@ -27,6 +27,9 @@ __INFINI_C infiniStatus_t infiniopCreateMoeTopkSoftmaxDescriptor( #endif #ifdef ENABLE_ILUVATAR_API CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -47,6 +50,9 @@ __INFINI_C infiniStatus_t infiniopGetMoeTopkSoftmaxWorkspaceSize( #endif #ifdef ENABLE_ILUVATAR_API GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -73,6 +79,9 @@ __INFINI_C infiniStatus_t infiniopMoeTopkSoftmax( #endif #ifdef ENABLE_ILUVATAR_API CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -92,6 +101,9 @@ __INFINI_C infiniStatus_t infiniopDestroyMoeTopkSoftmaxDescriptor( #endif #ifdef ENABLE_ILUVATAR_API DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; diff --git a/src/infiniop/ops/prepare_moe_input/nvidia/prepare_moe_input_nvidia.cu b/src/infiniop/ops/prepare_moe_input/nvidia/prepare_moe_input_nvidia.cu index 2ecc737f5..03ee446f6 100644 --- a/src/infiniop/ops/prepare_moe_input/nvidia/prepare_moe_input_nvidia.cu +++ b/src/infiniop/ops/prepare_moe_input/nvidia/prepare_moe_input_nvidia.cu @@ -1,4 +1,4 @@ -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) #include "prepare_moe_input_nvidia.cuh" @@ -207,4 +207,4 @@ infiniStatus_t Descriptor::calculate( } // namespace op::prepare_moe_input::nvidia -#endif // ENABLE_NVIDIA_API +#endif // ENABLE_NVIDIA_API || ENABLE_HYGON_API diff --git a/src/infiniop/ops/prepare_moe_input/operator.cc b/src/infiniop/ops/prepare_moe_input/operator.cc index aba7cd0a3..ac9a9d901 100644 --- a/src/infiniop/ops/prepare_moe_input/operator.cc +++ b/src/infiniop/ops/prepare_moe_input/operator.cc @@ -2,7 +2,7 @@ #include "../../handle.h" #include "infiniop/ops/prepare_moe_input.h" -#ifdef ENABLE_NVIDIA_API +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) #include "nvidia/prepare_moe_input_nvidia.cuh" #endif @@ -46,6 +46,9 @@ __INFINI_C infiniStatus_t infiniopCreatePrepareMoeInputDescriptor( #endif #ifdef ENABLE_METAX_API CREATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -69,6 +72,9 @@ __INFINI_C infiniStatus_t infiniopGetPrepareMoeInputWorkspaceSize( #endif #ifdef ENABLE_METAX_API GET(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -102,6 +108,9 @@ __INFINI_C infiniStatus_t infiniopPrepareMoeInput( #endif #ifdef ENABLE_METAX_API CALCULATE(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; @@ -124,6 +133,9 @@ __INFINI_C infiniStatus_t infiniopDestroyPrepareMoeInputDescriptor( #endif #ifdef ENABLE_METAX_API DESTROY(INFINI_DEVICE_METAX, metax); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); #endif default: return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; From 4265f845deb42f7f3774f05a77a0b992e33b9671 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Thu, 9 Jul 2026 09:26:08 +0800 Subject: [PATCH 02/16] feat(hygon): add graph-safe MoE inference ops --- .../infinicore/adaptor/lightop_adaptor.hpp | 61 ++++ include/infinicore/ops.hpp | 1 + include/infinicore/ops/moe_w16a16_marlin.hpp | 48 +++ src/infinicore/adaptor/lightop_adaptor.cc | 255 ++++++++++++++++ .../hygon/mha_kvcache_flashattn_hygon.cc | 12 + .../hygon/moe_w16a16_marlin_hygon.cc | 280 ++++++++++++++++++ .../moe_w16a16_marlin/moe_w16a16_marlin.cc | 90 ++++++ .../hygon/silu_and_mul_lightop_hygon.cc | 59 ++++ .../ops/moe_align/nvidia/moe_align_nvidia.cu | 14 +- xmake/hygon.lua | 4 +- 10 files changed, 819 insertions(+), 5 deletions(-) create mode 100644 include/infinicore/adaptor/lightop_adaptor.hpp create mode 100644 include/infinicore/ops/moe_w16a16_marlin.hpp create mode 100644 src/infinicore/adaptor/lightop_adaptor.cc create mode 100644 src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc create mode 100644 src/infinicore/ops/moe_w16a16_marlin/moe_w16a16_marlin.cc create mode 100644 src/infinicore/ops/silu_and_mul/hygon/silu_and_mul_lightop_hygon.cc diff --git a/include/infinicore/adaptor/lightop_adaptor.hpp b/include/infinicore/adaptor/lightop_adaptor.hpp new file mode 100644 index 000000000..46256c4bc --- /dev/null +++ b/include/infinicore/adaptor/lightop_adaptor.hpp @@ -0,0 +1,61 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#pragma once + +#include + +#include +#include + +namespace infinicore::adaptor::lightop { + +bool available(); + +bool enabled_by_env(); + +void preload_basic_ops(); + +void preload_silu_and_mul(); + +void fused_rms_norm_contiguous( + at::Tensor &out, + at::Tensor &input, + at::Tensor &weight, + double epsilon); + +void fuse_silu_and_mul( + at::Tensor &input, + at::Tensor &output); + +void moe_sum( + at::Tensor &input, + at::Tensor &output, + const std::optional &bias = std::nullopt, + const std::optional &expert_mask = std::nullopt, + const std::optional &local_num_tokens = std::nullopt, + float factor = 1.0f, + int expect_m = -1); + +std::vector moe_fused_gate( + at::Tensor &input, + at::Tensor &bias, + int64_t num_expert_group, + int64_t topk_group, + int64_t topk, + int64_t num_fused_shared_experts, + double routed_scaling_factor); + +void moe_gemm_marlin_w16a16( + at::Tensor input, + at::Tensor b_qweight, + at::Tensor output, + const std::optional &topk_weights, + at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded, + int64_t top_k, + int mode, + int delta); + +} // namespace infinicore::adaptor::lightop + +#endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index b960e4a49..8a5269c70 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -57,6 +57,7 @@ #include "ops/moe_sum.hpp" #include "ops/moe_topk_sigmoid.hpp" #include "ops/moe_topk_softmax.hpp" +#include "ops/moe_w16a16_marlin.hpp" #include "ops/nrm2.hpp" #include "ops/ones.hpp" #include "ops/paged_attention.hpp" diff --git a/include/infinicore/ops/moe_w16a16_marlin.hpp b/include/infinicore/ops/moe_w16a16_marlin.hpp new file mode 100644 index 000000000..58325c3e1 --- /dev/null +++ b/include/infinicore/ops/moe_w16a16_marlin.hpp @@ -0,0 +1,48 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MoeW16A16MarlinFusedDense, + Tensor, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + size_t, + int, + int, + int, + int); + +Tensor moe_w16a16_marlin_pack(const Tensor &weight); + +void moe_w16a16_marlin_fused_dense_( + Tensor output, + Tensor cache13, + Tensor cache2, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + int delta0, + int mode1, + int delta1); + +} // namespace infinicore::op diff --git a/src/infinicore/adaptor/lightop_adaptor.cc b/src/infinicore/adaptor/lightop_adaptor.cc new file mode 100644 index 000000000..b21639be1 --- /dev/null +++ b/src/infinicore/adaptor/lightop_adaptor.cc @@ -0,0 +1,255 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/adaptor/lightop_adaptor.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +namespace infinicore::adaptor::lightop { +namespace { + +constexpr const char *kDefaultLightopSo = + "/usr/local/lib/python3.10/dist-packages/lightop/op.cpython-310-x86_64-linux-gnu.so"; + +constexpr const char *kFusedRmsNormSymbol = + "_ZN2at6native25fused_rms_norm_contiguousERNS_6TensorES2_S2_d"; +constexpr const char *kFuseSiluAndMulSymbol = + "_ZN2at6native17fuse_silu_and_mulERNS_6TensorES2_"; +constexpr const char *kMoeSumSymbol = + "_ZN2at6native7moe_sumERNS_6TensorES2_RKSt8optionalIS1_ES6_S6_fi"; +constexpr const char *kMoeFusedGateSymbol = + "_ZN2at6native14moe_fused_gateERNS_6TensorES2_lllld"; +constexpr const char *kMoeGemmW16A16Symbol = + "_ZN2at6native15moe_gemm_w16a16ENS_6TensorES1_S1_St8optionalIS1_ES1_S1_S1_lii"; +constexpr const char *kMoeMarlinW16A16AsmSymbol = + "_ZN2at6native21moe_marlin_w16a16_asmENS_6TensorES1_S1_St8optionalIS1_ES1_S1_S1_iii"; + +class LightopLibrary { +public: + bool available() { + std::lock_guard lock(mutex_); + return ensure_open_locked(false); + } + + void *symbol(const char *name) { + std::lock_guard lock(mutex_); + if (!ensure_open_locked(true)) { + throw std::runtime_error(error_); + } + + dlerror(); + void *fn = dlsym(handle_, name); + const char *err = dlerror(); + if (err != nullptr || fn == nullptr) { + std::ostringstream oss; + oss << "failed to resolve lightop symbol " << name; + if (err != nullptr) { + oss << ": " << err; + } + throw std::runtime_error(oss.str()); + } + return fn; + } + +private: + bool ensure_open_locked(bool update_error) { + if (handle_ != nullptr) { + return true; + } + + const char *path_env = std::getenv("INFINICORE_LIGHTOP_SO"); + const char *path = (path_env != nullptr && path_env[0] != '\0') ? path_env : kDefaultLightopSo; + handle_ = dlopen(path, RTLD_LAZY | RTLD_LOCAL); + if (handle_ != nullptr) { + error_.clear(); + return true; + } + + if (update_error || error_.empty()) { + const char *err = dlerror(); + std::ostringstream oss; + oss << "failed to load lightop shared library " << path; + if (err != nullptr) { + oss << ": " << err; + } + error_ = oss.str(); + } + return false; + } + + std::mutex mutex_; + void *handle_ = nullptr; + std::string error_; +}; + +LightopLibrary &library() { + static LightopLibrary lib; + return lib; +} + +template +Fn resolve(const char *symbol) { + return reinterpret_cast(library().symbol(symbol)); +} + +using FusedRmsNormFn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, double); +using FuseSiluAndMulFn = void (*)(at::Tensor &, at::Tensor &); +using MoeSumFn = void (*)( + at::Tensor &, + at::Tensor &, + const std::optional &, + const std::optional &, + const std::optional &, + float, + int); +using MoeFusedGateFn = std::vector (*)( + at::Tensor &, + at::Tensor &, + int64_t, + int64_t, + int64_t, + int64_t, + double); +using MoeGemmW16A16Fn = at::Tensor (*)( + at::Tensor, + at::Tensor, + at::Tensor, + std::optional, + at::Tensor, + at::Tensor, + at::Tensor, + int64_t, + int, + int); +using MoeMarlinW16A16AsmFn = at::Tensor (*)( + at::Tensor, + at::Tensor, + at::Tensor, + std::optional, + at::Tensor, + at::Tensor, + at::Tensor, + int, + int, + int); + +FusedRmsNormFn fused_rms_norm_fn() { + static auto fn = resolve(kFusedRmsNormSymbol); + return fn; +} + +FuseSiluAndMulFn fuse_silu_and_mul_fn() { + static auto fn = resolve(kFuseSiluAndMulSymbol); + return fn; +} + +MoeSumFn moe_sum_fn() { + static auto fn = resolve(kMoeSumSymbol); + return fn; +} + +MoeFusedGateFn moe_fused_gate_fn() { + static auto fn = resolve(kMoeFusedGateSymbol); + return fn; +} + +MoeGemmW16A16Fn moe_gemm_w16a16_fn() { + static auto fn = resolve(kMoeGemmW16A16Symbol); + return fn; +} + +MoeMarlinW16A16AsmFn moe_marlin_w16a16_asm_fn() { + static auto fn = resolve(kMoeMarlinW16A16AsmSymbol); + return fn; +} + +} // namespace + +bool available() { + return library().available(); +} + +bool enabled_by_env() { + const char *value = std::getenv("INFINICORE_ENABLE_HYGON_LIGHTOP"); + if (value == nullptr) { + return false; + } + std::string normalized(value); + for (auto &ch : normalized) { + ch = static_cast(std::tolower(static_cast(ch))); + } + return normalized == "1" || normalized == "true" || normalized == "on" || normalized == "yes"; +} + +void preload_basic_ops() { + (void)fused_rms_norm_fn(); + (void)fuse_silu_and_mul_fn(); + (void)moe_sum_fn(); + (void)moe_fused_gate_fn(); + (void)moe_gemm_w16a16_fn(); + (void)moe_marlin_w16a16_asm_fn(); +} + +void preload_silu_and_mul() { + (void)fuse_silu_and_mul_fn(); +} + +void fused_rms_norm_contiguous(at::Tensor &out, at::Tensor &input, at::Tensor &weight, double epsilon) { + fused_rms_norm_fn()(out, input, weight, epsilon); +} + +void fuse_silu_and_mul(at::Tensor &input, at::Tensor &output) { + fuse_silu_and_mul_fn()(input, output); +} + +void moe_sum(at::Tensor &input, + at::Tensor &output, + const std::optional &bias, + const std::optional &expert_mask, + const std::optional &local_num_tokens, + float factor, + int expect_m) { + moe_sum_fn()(input, output, bias, expert_mask, local_num_tokens, factor, expect_m); +} + +std::vector moe_fused_gate(at::Tensor &input, + at::Tensor &bias, + int64_t num_expert_group, + int64_t topk_group, + int64_t topk, + int64_t num_fused_shared_experts, + double routed_scaling_factor) { + return moe_fused_gate_fn()(input, bias, num_expert_group, topk_group, topk, num_fused_shared_experts, routed_scaling_factor); +} + +void moe_gemm_marlin_w16a16(at::Tensor input, + at::Tensor b_qweight, + at::Tensor output, + const std::optional &topk_weights, + at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded, + int64_t top_k, + int mode, + int delta) { + if (mode < 1000) { + moe_gemm_w16a16_fn()( + input, b_qweight, output, topk_weights, + sorted_token_ids, expert_ids, num_tokens_post_padded, + top_k, mode, delta); + } else { + moe_marlin_w16a16_asm_fn()( + input, b_qweight, output, topk_weights, + sorted_token_ids, expert_ids, num_tokens_post_padded, + static_cast(top_k), mode, delta); + } +} + +} // namespace infinicore::adaptor::lightop + +#endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc b/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc index 9d2cdc2b5..160147436 100644 --- a/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc +++ b/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc @@ -11,9 +11,18 @@ #include #include +#include namespace infinicore::op::mha_kvcache_impl::flashattn { +namespace { +bool isCurrentHipStreamCapturing() { + hipStreamCaptureStatus capture_status = hipStreamCaptureStatusNone; + hipError_t status = hipStreamIsCapturing(infinicore::adaptor::get_hip_stream().stream(), &capture_status); + return status == hipSuccess && capture_status != hipStreamCaptureStatusNone; +} +} // namespace + struct PlannedMeta { graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; std::optional alibi_slopes; @@ -143,6 +152,9 @@ void run(void *planned_meta) { false, 0); + if (!isCurrentHipStreamCapturing()) { + c10::hip::device_synchronize(); + } if (!result.empty() && result[0].defined()) { out_tensor.copy_(result[0]); } diff --git a/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc b/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc new file mode 100644 index 000000000..001e01875 --- /dev/null +++ b/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc @@ -0,0 +1,280 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/ops/moe_w16a16_marlin.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/lightop_adaptor.hpp" +#include "infinicore/ops/common/dispatcher.hpp" +#include "infinicore/ops/moe_sum.hpp" +#include "infinicore/ops/silu_and_mul.hpp" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace infinicore::op { +namespace moe_w16a16_marlin_pack_impl { +using schema = Tensor (*)(const Tensor &); +common::OpDispatcher &dispatcher(); +} // namespace moe_w16a16_marlin_pack_impl +} // namespace infinicore::op + +namespace infinicore::op::moe_w16a16_marlin_impl::hygon { + +namespace { + +std::vector weight_perm_data() { + std::vector perm; + perm.reserve(2048); + for (int i = 0; i < 64; ++i) { + for (int col = 0; col < 2; ++col) { + const int cur_col = (i % 16) * 2 + col; + for (int row = 0; row < 4; ++row) { + const int cur_row = (i / 16) * 4 + row; + perm.push_back(static_cast(cur_row * 32 + cur_col)); + } + } + } + return perm; +} + +at::Tensor pack_one_expert(const at::Tensor &weight) { + if (weight.dim() != 2) { + throw std::runtime_error("w16a16 marlin pack expects each expert weight to be 2D"); + } + auto q_w = weight.transpose(0, 1).contiguous(); + const int64_t size_k = q_w.size(0); + const int64_t size_n = q_w.size(1); + if (size_k % 16 != 0 || size_n % 32 != 0) { + throw std::runtime_error("w16a16 marlin pack requires K % 16 == 0 and N % 32 == 0"); + } + const auto perm_vec = weight_perm_data(); + auto weight_perm = at::tensor( + perm_vec, + at::TensorOptions().dtype(at::kLong).device(q_w.device())); + auto packed = q_w.reshape({size_k / 16, 16, size_n / 32, 32}) + .permute({0, 2, 1, 3}) + .reshape({size_k / 16, size_n * 16}); + packed = packed.reshape({-1, static_cast(perm_vec.size())}) + .index_select(1, weight_perm) + .reshape({size_k / 16, size_n * 16}) + .contiguous(); + return packed; +} + + +bool debug_enabled() { + const char *value = std::getenv("INFINICORE_DEBUG_HYGON_MARLIN"); + return value != nullptr && value[0] != '\0' && std::string(value) != "0"; +} + +std::string tensor_desc(const at::Tensor &tensor) { + std::ostringstream oss; + oss << "shape=["; + for (int64_t i = 0; i < tensor.dim(); ++i) { + if (i != 0) { + oss << ","; + } + oss << tensor.size(i); + } + oss << "] stride=["; + for (int64_t i = 0; i < tensor.dim(); ++i) { + if (i != 0) { + oss << ","; + } + oss << tensor.stride(i); + } + oss << "] dtype=" << tensor.scalar_type() + << " device=" << tensor.device() + << " contiguous=" << (tensor.is_contiguous() ? "true" : "false"); + return oss.str(); +} + +void debug_tensor(const char *name, const at::Tensor &tensor) { + if (debug_enabled()) { + std::cerr << "[hygon-marlin] " << name << " " << tensor_desc(tensor) << std::endl; + } +} + +Tensor pack(const Tensor &weight) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto weight_at = infinicore::adaptor::to_aten_tensor(weight); + if (weight_at.dim() != 3) { + throw std::runtime_error("w16a16 marlin pack expects weight shape [E, N, K]"); + } + const int64_t num_experts = weight_at.size(0); + auto packed0 = pack_one_expert(weight_at.select(0, 0)); + auto output_ic = Tensor::empty( + {static_cast(num_experts), + static_cast(packed0.size(0)), + static_cast(packed0.size(1))}, + weight->dtype(), + weight->device()); + auto output_at = infinicore::adaptor::to_aten_tensor(output_ic); + output_at.select(0, 0).copy_(packed0); + for (int64_t expert = 1; expert < num_experts; ++expert) { + auto packed = pack_one_expert(weight_at.select(0, expert)); + output_at.select(0, expert).copy_(packed); + } + return output_ic; +} + +} // namespace + +struct PlannedMeta { + graph::GraphTensor output; + graph::GraphTensor cache13; + graph::GraphTensor cache2; + graph::GraphTensor hidden_states; + graph::GraphTensor w13_marlin; + graph::GraphTensor w2_marlin; + graph::GraphTensor topk_weights; + graph::GraphTensor sorted_token_ids; + graph::GraphTensor expert_ids; + graph::GraphTensor num_tokens_post_padded; + size_t top_k; + int mode0; + int delta0; + int mode1; + int delta1; +}; + +void *plan(Tensor output, + Tensor cache13, + Tensor cache2, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + int delta0, + int mode1, + int delta1) { + infinicore::adaptor::lightop::preload_basic_ops(); + return new PlannedMeta{ + graph::GraphTensor(output), graph::GraphTensor(cache13), graph::GraphTensor(cache2), + graph::GraphTensor(hidden_states), graph::GraphTensor(w13_marlin), graph::GraphTensor(w2_marlin), + graph::GraphTensor(topk_weights), graph::GraphTensor(sorted_token_ids), graph::GraphTensor(expert_ids), + graph::GraphTensor(num_tokens_post_padded), top_k, mode0, delta0, mode1, delta1}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + const auto hidden_shape = p->hidden_states->shape(); + const auto w13_shape = p->w13_marlin->shape(); + const auto w2_shape = p->w2_marlin->shape(); + if (hidden_shape.size() != 2 || w13_shape.size() != 3 || w2_shape.size() != 3) { + throw std::runtime_error("w16a16 marlin fused dense expects hidden [M,K], w13/w2 [E,*,*]"); + } + const size_t m = hidden_shape[0]; + const size_t k = hidden_shape[1]; + const size_t n2 = w13_shape[2] / 16; + const size_t n = n2 / 2; + if (w13_shape[1] * 16 != k || w2_shape[1] * 16 != n || w2_shape[2] / 16 != k) { + throw std::runtime_error("w16a16 marlin fused dense weight shape mismatch"); + } + + const bool output_need_copy_back = !p->output->is_contiguous(); + Tensor output_work_ic = output_need_copy_back ? p->output->contiguous() : Tensor(p->output); + Tensor hidden_work_ic = p->hidden_states->is_contiguous() ? Tensor(p->hidden_states) : p->hidden_states->contiguous(); + + const size_t top_k = p->top_k; + const size_t cache1_numel = m * top_k * n2; + const size_t cache3_numel = m * top_k * k; + const size_t cache2_numel = m * top_k * n; + auto cache1_ic = p->cache13->narrow({{0, 0, cache1_numel}})->view({m * top_k, n2}); + auto cache3_ic = p->cache13->narrow({{0, 0, cache3_numel}})->view({m * top_k, k}); + auto cache2_ic = p->cache2->narrow({{0, 0, cache2_numel}})->view({m * top_k, n}); + + auto hidden = infinicore::adaptor::to_aten_tensor(hidden_work_ic); + auto w13 = infinicore::adaptor::to_aten_tensor(p->w13_marlin); + auto w2 = infinicore::adaptor::to_aten_tensor(p->w2_marlin); + auto cache1 = infinicore::adaptor::to_aten_tensor(cache1_ic); + auto cache2 = infinicore::adaptor::to_aten_tensor(cache2_ic); + auto cache3 = infinicore::adaptor::to_aten_tensor(cache3_ic); + auto topk_weights = infinicore::adaptor::to_aten_tensor(p->topk_weights); + auto sorted_token_ids = infinicore::adaptor::to_aten_tensor(p->sorted_token_ids); + auto expert_ids = infinicore::adaptor::to_aten_tensor(p->expert_ids); + auto num_tokens_post_padded = infinicore::adaptor::to_aten_tensor(p->num_tokens_post_padded); + + if (debug_enabled()) { + std::cerr << "[hygon-marlin] top_k=" << top_k << " mode0=" << p->mode0 << " delta0=" << p->delta0 + << " mode1=" << p->mode1 << " delta1=" << p->delta1 << std::endl; + debug_tensor("hidden", hidden); + debug_tensor("w13", w13); + debug_tensor("cache1", cache1); + debug_tensor("cache2", cache2); + debug_tensor("cache3", cache3); + debug_tensor("topk_weights", topk_weights); + debug_tensor("sorted_token_ids", sorted_token_ids); + debug_tensor("expert_ids", expert_ids); + debug_tensor("num_tokens_post_padded", num_tokens_post_padded); + } + + try { + infinicore::adaptor::lightop::moe_gemm_marlin_w16a16( + hidden, w13, cache1, std::nullopt, sorted_token_ids, expert_ids, + num_tokens_post_padded, static_cast(top_k), p->mode0, p->delta0); + } catch (const std::exception &e) { + throw std::runtime_error(std::string("Hygon W16A16 Marlin GEMM1 failed: ") + e.what()); + } + + infinicore::op::silu_and_mul_(cache2_ic, cache1_ic); + + std::optional topk_weights_opt(topk_weights); + try { + infinicore::adaptor::lightop::moe_gemm_marlin_w16a16( + cache2, w2, cache3, topk_weights_opt, sorted_token_ids, expert_ids, + num_tokens_post_padded, 1, p->mode1, p->delta1); + } catch (const std::exception &e) { + throw std::runtime_error(std::string("Hygon W16A16 Marlin GEMM2 failed: ") + e.what()); + } + + auto cache3_reduce_ic = cache3_ic->view({m, top_k, k}); + auto cache3_reduce = infinicore::adaptor::to_aten_tensor(cache3_reduce_ic); + auto output_work = infinicore::adaptor::to_aten_tensor(output_work_ic); + infinicore::adaptor::lightop::moe_sum( + cache3_reduce, + output_work, + std::nullopt, + std::nullopt, + std::nullopt, + 1.0f, + -1); + + if (output_need_copy_back) { + p->output->copy_from(output_work_ic); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + if (!infinicore::adaptor::lightop::available()) { + return false; + } + infinicore::op::moe_w16a16_marlin_pack_impl::dispatcher().registerDevice(Device::Type::HYGON, &pack); + MoeW16A16MarlinFusedDense::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); + MoeW16A16MarlinFusedDense::run_dispatcher().registerDevice(Device::Type::HYGON, &run); + MoeW16A16MarlinFusedDense::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); + return true; +}(); + +} // namespace infinicore::op::moe_w16a16_marlin_impl::hygon + +#endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/src/infinicore/ops/moe_w16a16_marlin/moe_w16a16_marlin.cc b/src/infinicore/ops/moe_w16a16_marlin/moe_w16a16_marlin.cc new file mode 100644 index 000000000..d8a1f5ffb --- /dev/null +++ b/src/infinicore/ops/moe_w16a16_marlin/moe_w16a16_marlin.cc @@ -0,0 +1,90 @@ +#include "infinicore/ops/moe_w16a16_marlin.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +namespace moe_w16a16_marlin_pack_impl { +using schema = Tensor (*)(const Tensor &); +common::OpDispatcher &dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} +} // namespace moe_w16a16_marlin_pack_impl + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeW16A16MarlinFusedDense); + +MoeW16A16MarlinFusedDense::MoeW16A16MarlinFusedDense( + Tensor output, + Tensor cache13, + Tensor cache2, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + int delta0, + int mode1, + int delta1) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + output, cache13, cache2, hidden_states, w13_marlin, w2_marlin, + topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded); + INFINICORE_GRAPH_OP_DISPATCH( + output->device().getType(), output, cache13, cache2, hidden_states, + w13_marlin, w2_marlin, topk_weights, sorted_token_ids, expert_ids, + num_tokens_post_padded, top_k, mode0, delta0, mode1, delta1); +} + +void MoeW16A16MarlinFusedDense::execute( + Tensor output, + Tensor cache13, + Tensor cache2, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + int delta0, + int mode1, + int delta1) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MoeW16A16MarlinFusedDense, output, cache13, cache2, hidden_states, + w13_marlin, w2_marlin, topk_weights, sorted_token_ids, expert_ids, + num_tokens_post_padded, top_k, mode0, delta0, mode1, delta1); +} + +Tensor moe_w16a16_marlin_pack(const Tensor &weight) { + return moe_w16a16_marlin_pack_impl::dispatcher().lookup(weight->device().getType())(weight); +} + +void moe_w16a16_marlin_fused_dense_( + Tensor output, + Tensor cache13, + Tensor cache2, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + int delta0, + int mode1, + int delta1) { + MoeW16A16MarlinFusedDense::execute( + output, cache13, cache2, hidden_states, w13_marlin, w2_marlin, + topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, + top_k, mode0, delta0, mode1, delta1); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/silu_and_mul/hygon/silu_and_mul_lightop_hygon.cc b/src/infinicore/ops/silu_and_mul/hygon/silu_and_mul_lightop_hygon.cc new file mode 100644 index 000000000..0472eb68c --- /dev/null +++ b/src/infinicore/ops/silu_and_mul/hygon/silu_and_mul_lightop_hygon.cc @@ -0,0 +1,59 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/ops/silu_and_mul.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/lightop_adaptor.hpp" + +#include + +namespace infinicore::op::silu_and_mul_impl::lightop_hygon { + +struct PlannedMeta { + graph::GraphTensor out; + graph::GraphTensor x; +}; + +void *plan(Tensor out, const Tensor &x) { + infinicore::adaptor::lightop::preload_silu_and_mul(); + return new PlannedMeta{ + graph::GraphTensor(out), + graph::GraphTensor(x)}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + const bool out_need_copy_back = !p->out->is_contiguous(); + Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); + Tensor x_work_ic = p->x->is_contiguous() ? Tensor(p->x) : p->x->contiguous(); + + auto out = infinicore::adaptor::to_aten_tensor(out_work_ic); + auto x = infinicore::adaptor::to_aten_tensor(x_work_ic); + + infinicore::adaptor::lightop::fuse_silu_and_mul(x, out); + + if (out_need_copy_back) { + p->out->copy_from(out_work_ic); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + if (!infinicore::adaptor::lightop::enabled_by_env() + || !infinicore::adaptor::lightop::available()) { + return false; + } + SiluAndMul::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); + SiluAndMul::run_dispatcher().registerDevice(Device::Type::HYGON, &run); + SiluAndMul::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); + return true; +}(); + +} // namespace infinicore::op::silu_and_mul_impl::lightop_hygon + +#endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/src/infiniop/ops/moe_align/nvidia/moe_align_nvidia.cu b/src/infiniop/ops/moe_align/nvidia/moe_align_nvidia.cu index 8a0222a4f..774de29f5 100644 --- a/src/infiniop/ops/moe_align/nvidia/moe_align_nvidia.cu +++ b/src/infiniop/ops/moe_align/nvidia/moe_align_nvidia.cu @@ -356,13 +356,23 @@ infiniStatus_t Descriptor::calculate( auto *cumsum_buffer = static_cast(workspace); constexpr int warp_size = 32; - int threads = 1024; +#if defined(ENABLE_HYGON_API) && !defined(ENABLE_NVIDIA_API) + constexpr int max_threads_per_block = 256; +#else + constexpr int max_threads_per_block = 1024; +#endif + int threads = max_threads_per_block; threads = ((threads + warp_size - 1) / warp_size) * warp_size; const int32_t num_experts = static_cast(_info.num_experts + 1); const int32_t block_size = static_cast(_info.block_size); const int32_t max_num_tokens_padded = static_cast(_info.max_num_tokens_padded); - const bool small_batch_expert_mode = (_info.numel < 1024) && (num_experts <= 64); + const bool small_batch_expert_mode = +#if defined(ENABLE_HYGON_API) && !defined(ENABLE_NVIDIA_API) + false; +#else + (_info.numel < 1024) && (num_experts <= 64); +#endif if (small_batch_expert_mode) { const int32_t expert_threads = std::max(num_experts, warp_size); diff --git a/xmake/hygon.lua b/xmake/hygon.lua index ed7da7b85..9e5175ce9 100644 --- a/xmake/hygon.lua +++ b/xmake/hygon.lua @@ -220,9 +220,7 @@ target("infiniccl-hygon") add_cxxflags("-fPIC") -- 添加海光DCU特定的编译标志 - -- 检测实际GPU架构，如果未指定则默认使用gfx906 - local hygon_arch = os.getenv("HYGON_ARCH") or "gfx906" - add_cuflags("-arch=" .. hygon_arch) + add_cuflags("-arch=" .. HYGON_ARCH) -- 使用NCCL (NVIDIA Collective Communications Library) add_links("nccl") From e6a89ae68815d3b0aa0e7d2744a698f106dfef16 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 10 Jul 2026 20:11:19 +0800 Subject: [PATCH 03/16] fix: honor singleton strides in contiguity checks --- src/infinicore-test/test_tensor_destructor.cc | 21 +++++++++++++++++++ src/infinicore/tensor/tensor.cc | 8 +++++++ 2 files changed, 29 insertions(+) diff --git a/src/infinicore-test/test_tensor_destructor.cc b/src/infinicore-test/test_tensor_destructor.cc index c75dd546b..4288ec088 100644 --- a/src/infinicore-test/test_tensor_destructor.cc +++ b/src/infinicore-test/test_tensor_destructor.cc @@ -202,6 +202,27 @@ TestResult TensorDestructorTest::testStridedTensor() { std::cout << dim << " "; } std::cout << std::endl; + + auto singleton_strided = Tensor::strided_empty( + {1, 1, 2048}, {2560, 2560, 1}, DataType::F32, Device::Type::CPU); + if (!singleton_strided->is_contiguous()) { + std::cerr << "Size-one dimensions must not constrain contiguity" << std::endl; + return false; + } + + auto genuinely_strided = Tensor::strided_empty( + {2, 1, 3}, {4, 100, 1}, DataType::F32, Device::Type::CPU); + if (genuinely_strided->is_contiguous()) { + std::cerr << "A gap between non-singleton dimensions must be non-contiguous" << std::endl; + return false; + } + + auto empty_strided = Tensor::strided_empty( + {0, 3}, {99, 1}, DataType::F32, Device::Type::CPU); + if (!empty_strided->is_contiguous()) { + std::cerr << "Empty tensors must be contiguous" << std::endl; + return false; + } } std::cout << "Destroyed strided tensor successfully" << std::endl; diff --git a/src/infinicore/tensor/tensor.cc b/src/infinicore/tensor/tensor.cc index 3e6f2ea3c..9e685f8dd 100644 --- a/src/infinicore/tensor/tensor.cc +++ b/src/infinicore/tensor/tensor.cc @@ -106,8 +106,16 @@ Size TensorImpl::ndim() const { } bool TensorImpl::is_contiguous() const { + if (numel() == 0) { + return true; + } + Stride expected_stride = 1; for (int i = meta_.shape.size() - 1; i >= 0; --i) { + // Size-one dimensions do not constrain the physical layout. + if (meta_.shape[i] == 1) { + continue; + } if (meta_.strides[i] != expected_stride) { return false; } From ff0dbd65c946ac83caad5cba8833632580cce91f Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Sat, 11 Jul 2026 00:43:42 +0800 Subject: [PATCH 04/16] Optimize-Hygon-TP2-graph-allreduce --- src/infiniccl/cuda/infiniccl_cuda.cu | 394 +++++++++++++++++++++++++++ 1 file changed, 394 insertions(+) diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu index e9a177dda..4de94db0c 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ b/src/infiniccl/cuda/infiniccl_cuda.cu @@ -6,6 +6,19 @@ #include #include +#if defined(ENABLE_HYGON_API) +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#endif + #include "../../utils.h" #define CHECK_NCCL(API__) CHECK_INTERNAL(API__, ncclSuccess) @@ -63,6 +76,372 @@ inline ncclComm_t getNcclComm(infinicclComm_t comm) { namespace infiniccl::cuda { +#if defined(ENABLE_HYGON_API) +namespace { + +constexpr int kHygonTp2MaxBlocks = 80; +constexpr size_t kHygonTp2StageCapacityElements = 1u << 22; + +struct HygonTp2Signal { + alignas(128) uint32_t start[kHygonTp2MaxBlocks][8]; + alignas(128) uint32_t end[kHygonTp2MaxBlocks][8]; + alignas(128) uint32_t flag[kHygonTp2MaxBlocks]; +}; + +struct alignas(16) HygonTp2RankData { + const void *ptrs[2]; +}; + +struct alignas(16) HygonTp2RankSignals { + HygonTp2Signal *signals[2]; +}; + +struct alignas(16) HygonBf16Pack { + __nv_bfloat16 values[8]; +}; + +struct HygonCudaDriverApi { + void *library = nullptr; + decltype(&cuMemGetAllocationGranularity) mem_get_allocation_granularity = nullptr; + decltype(&cuMemAddressReserve) mem_address_reserve = nullptr; + decltype(&cuMemCreate) mem_create = nullptr; + decltype(&cuMemMap) mem_map = nullptr; + decltype(&cuMemSetAccess) mem_set_access = nullptr; + decltype(&cuMemUnmap) mem_unmap = nullptr; + decltype(&cuMemRelease) mem_release = nullptr; + decltype(&cuMemAddressFree) mem_address_free = nullptr; + bool available = false; + + HygonCudaDriverApi() { + library = dlopen("libcuda.so.1", RTLD_NOW | RTLD_LOCAL); + if (library == nullptr) { + library = dlopen("/opt/dtk/cuda/cuda/lib64/libcuda.so.1", RTLD_NOW | RTLD_LOCAL); + } + available = library != nullptr && + load(mem_get_allocation_granularity, "cuMemGetAllocationGranularity") && + load(mem_address_reserve, "cuMemAddressReserve") && + load(mem_create, "cuMemCreate") && + load(mem_map, "cuMemMap") && + load(mem_set_access, "cuMemSetAccess") && + load(mem_unmap, "cuMemUnmap") && + load(mem_release, "cuMemRelease") && + load(mem_address_free, "cuMemAddressFree"); + } + +private: + template + bool load(T &symbol, const char *name) { + symbol = reinterpret_cast(dlsym(library, name)); + return symbol != nullptr; + } +}; + +HygonCudaDriverApi &hygon_cuda_driver_api() { + static HygonCudaDriverApi api; + return api; +} + +struct HygonVmmAllocation { + void *ptr = nullptr; + size_t size = 0; + CUmemGenericAllocationHandle handle = 0; +}; + +struct HygonTp2AllReduceState { + int device_ids[2] = {0, 1}; + HygonVmmAllocation stages[2]; + HygonTp2Signal *signal_hosts[2] = {nullptr, nullptr}; + HygonTp2Signal *signals[2] = {nullptr, nullptr}; + HygonTp2RankData *rank_data[2] = {nullptr, nullptr}; + HygonTp2RankSignals rank_signals{}; + struct CaptureCursor { + unsigned long long id = 0; + size_t next_element = 0; + bool initialized = false; + }; + std::mutex capture_mutex; + CaptureCursor capture_cursors[2]; + + ~HygonTp2AllReduceState() { + int previous_device = 0; + const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess; + for (int rank = 0; rank < 2; ++rank) { + cudaSetDevice(device_ids[rank]); + if (rank_data[rank] != nullptr) cudaFree(rank_data[rank]); + if (signal_hosts[rank] != nullptr) cudaFreeHost(signal_hosts[rank]); + auto &driver = hygon_cuda_driver_api(); + if (driver.available && stages[rank].ptr != nullptr) { + const auto address = reinterpret_cast(stages[rank].ptr); + driver.mem_unmap(address, stages[rank].size); + driver.mem_address_free(address, stages[rank].size); + } + if (driver.available && stages[rank].handle != 0) { + driver.mem_release(stages[rank].handle); + } + } + if (restore_device) cudaSetDevice(previous_device); + } +}; + +std::mutex hygon_tp2_states_mutex; +std::unordered_map> hygon_tp2_states; + +bool allocate_hygon_vmm(HygonVmmAllocation &allocation, + int owner_device, + const int device_ids[2], + size_t requested_size) { + auto &driver = hygon_cuda_driver_api(); + if (!driver.available) return false; + CUmemAllocationProp properties{}; + properties.type = CU_MEM_ALLOCATION_TYPE_PINNED; + properties.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + properties.location.id = owner_device; + size_t granularity = 0; + if (driver.mem_get_allocation_granularity( + &granularity, &properties, CU_MEM_ALLOC_GRANULARITY_MINIMUM) != CUDA_SUCCESS || + granularity == 0) return false; + allocation.size = (requested_size + granularity - 1) / granularity * granularity; + CUdeviceptr address = 0; + if (driver.mem_address_reserve( + &address, allocation.size, granularity, 0, 0) != CUDA_SUCCESS) { + allocation = {}; + return false; + } + allocation.ptr = reinterpret_cast(address); + if (driver.mem_create( + &allocation.handle, allocation.size, &properties, 0) != CUDA_SUCCESS) { + driver.mem_address_free(address, allocation.size); + allocation = {}; + return false; + } + if (driver.mem_map( + address, allocation.size, 0, allocation.handle, 0) != CUDA_SUCCESS) { + driver.mem_release(allocation.handle); + driver.mem_address_free(address, allocation.size); + allocation = {}; + return false; + } + CUmemAccessDesc access[2]{}; + for (int rank = 0; rank < 2; ++rank) { + access[rank].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access[rank].location.id = device_ids[rank]; + access[rank].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + } + if (driver.mem_set_access(address, allocation.size, access, 2) != CUDA_SUCCESS) { + driver.mem_unmap(address, allocation.size); + driver.mem_release(allocation.handle); + driver.mem_address_free(address, allocation.size); + allocation = {}; + return false; + } + return true; +} + +template +__device__ __forceinline__ uint32_t hygon_tp2_start_sync( + const HygonTp2RankSignals &rank_signals, + HygonTp2Signal *self_signal, + int rank) { + const uint32_t next_flag = self_signal->flag[blockIdx.x] + 1; + if (threadIdx.x < NumRanks) { + __scoped_atomic_store_n( + &rank_signals.signals[threadIdx.x]->start[blockIdx.x][rank], + next_flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); + while (__scoped_atomic_load_n( + &self_signal->start[blockIdx.x][threadIdx.x], + __ATOMIC_RELAXED, __MEMORY_SCOPE_DEVICE) < next_flag) { + } + } + __syncthreads(); + if (threadIdx.x == 0) self_signal->flag[blockIdx.x] = next_flag; + return next_flag; +} + +template +__device__ __forceinline__ void hygon_tp2_end_sync( + const HygonTp2RankSignals &rank_signals, + HygonTp2Signal *self_signal, + int rank, + uint32_t flag) { + __syncthreads(); + if (threadIdx.x < NumRanks) { + __scoped_atomic_store_n( + &rank_signals.signals[threadIdx.x]->end[blockIdx.x][rank], + flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM); + while (__scoped_atomic_load_n( + &self_signal->end[blockIdx.x][threadIdx.x], + __ATOMIC_RELAXED, __MEMORY_SCOPE_DEVICE) < flag) { + } + } + __syncthreads(); +} + +__global__ __launch_bounds__(512, 1) void hygon_tp2_bf16_allreduce_kernel( + const HygonTp2RankData *rank_data, + HygonTp2RankSignals rank_signals, + HygonTp2Signal *self_signal, + const __nv_bfloat16 *input, + __nv_bfloat16 *output, + int rank, + size_t pack_count, + size_t pack_offset) { + constexpr int num_ranks = 2; + constexpr int threads_per_rank = 512 / num_ranks; + constexpr int pack_size = 8; + __shared__ __nv_bfloat16 shared[threads_per_rank * num_ranks * pack_size]; + const HygonTp2RankData data = *rank_data; + const int source_rank = threadIdx.x / threads_per_rank; + const int lane = threadIdx.x % threads_per_rank; + if (threadIdx.x < threads_per_rank) { + auto *local_stage = reinterpret_cast( + const_cast(data.ptrs[rank])); + const auto *local_input = reinterpret_cast(input); + for (size_t index = blockIdx.x * threads_per_rank + threadIdx.x; + index < pack_count; + index += gridDim.x * threads_per_rank) { + local_stage[pack_offset + index] = local_input[index]; + } + __threadfence_system(); + } + __syncthreads(); + const uint32_t sync_flag = + hygon_tp2_start_sync(rank_signals, self_signal, rank); + for (size_t index = blockIdx.x * threads_per_rank + lane; + index < pack_count; + index += gridDim.x * threads_per_rank) { + auto *shared_packs = reinterpret_cast(shared); + const auto *source = reinterpret_cast(data.ptrs[source_rank]); + shared_packs[threadIdx.x] = source[pack_offset + index]; + __syncthreads(); + if (source_rank == 0) { + HygonBf16Pack reduced; +#pragma unroll + for (int element = 0; element < pack_size; ++element) { + const float value = + __bfloat162float(shared[threadIdx.x * pack_size + element]) + + __bfloat162float(shared[(threads_per_rank + threadIdx.x) * pack_size + element]); + reduced.values[element] = __float2bfloat16(value); + } + reinterpret_cast(output)[index] = reduced; + } + __syncthreads(); + } + if (pack_offset == 0) { + hygon_tp2_end_sync( + rank_signals, self_signal, rank, sync_flag); + } +} + +std::shared_ptr create_hygon_tp2_state( + int ndevice, const int *device_ids) { + if (ndevice != 2 || device_ids == nullptr) return nullptr; + auto state = std::make_shared(); + state->device_ids[0] = device_ids[0]; + state->device_ids[1] = device_ids[1]; + int previous_device = 0; + const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess; + auto fail = [&]() -> std::shared_ptr { + if (restore_device) cudaSetDevice(previous_device); + return nullptr; + }; + const size_t stage_bytes = kHygonTp2StageCapacityElements * sizeof(__nv_bfloat16); + for (int rank = 0; rank < 2; ++rank) { + if (cudaSetDevice(device_ids[rank]) != cudaSuccess || + !allocate_hygon_vmm(state->stages[rank], device_ids[rank], + state->device_ids, stage_bytes)) return fail(); + void *signal_host = nullptr; + if (cudaHostAlloc(&signal_host, sizeof(HygonTp2Signal), + cudaHostAllocMapped) != cudaSuccess) return fail(); + state->signal_hosts[rank] = static_cast(signal_host); + std::memset(state->signal_hosts[rank], 0, sizeof(HygonTp2Signal)); + void *signal_device = nullptr; + if (cudaHostGetDevicePointer(&signal_device, signal_host, 0) != cudaSuccess) return fail(); + state->signals[rank] = static_cast(signal_device); + if (cudaMalloc(reinterpret_cast(&state->rank_data[rank]), + sizeof(HygonTp2RankData)) != cudaSuccess) return fail(); + } + HygonTp2RankData host_rank_data{{state->stages[0].ptr, state->stages[1].ptr}}; + state->rank_signals.signals[0] = state->signals[0]; + state->rank_signals.signals[1] = state->signals[1]; + for (int rank = 0; rank < 2; ++rank) { + cudaSetDevice(device_ids[rank]); + if (cudaMemcpy(state->rank_data[rank], &host_rank_data, + sizeof(host_rank_data), cudaMemcpyHostToDevice) != cudaSuccess) return fail(); + } + if (restore_device) cudaSetDevice(previous_device); + return state; +} + +void register_hygon_tp2_state(infinicclComm_t *comms, + int ndevice, + const int *device_ids) { + auto state = create_hygon_tp2_state(ndevice, device_ids); + if (state == nullptr) return; + std::lock_guard lock(hygon_tp2_states_mutex); + for (int rank = 0; rank < 2; ++rank) hygon_tp2_states.emplace(comms[rank], state); +} + +void erase_hygon_tp2_state(infinicclComm_t comm) { + std::lock_guard lock(hygon_tp2_states_mutex); + hygon_tp2_states.erase(comm); +} + +std::shared_ptr get_hygon_tp2_state(infinicclComm_t comm) { + std::lock_guard lock(hygon_tp2_states_mutex); + auto found = hygon_tp2_states.find(comm); + return found == hygon_tp2_states.end() ? nullptr : found->second; +} + +bool reserve_hygon_tp2_stage( + const std::shared_ptr &state, + int rank, + unsigned long long capture_id, + size_t count, + size_t *element_offset) { + std::lock_guard lock(state->capture_mutex); + auto &cursor = state->capture_cursors[rank]; + if (!cursor.initialized || cursor.id != capture_id) { + cursor.id = capture_id; + cursor.next_element = 0; + cursor.initialized = true; + } + const size_t aligned_offset = (cursor.next_element + 7) & ~size_t{7}; + if (aligned_offset > kHygonTp2StageCapacityElements - count) return false; + *element_offset = aligned_offset; + cursor.next_element = aligned_offset + count; + return true; +} + +bool try_hygon_tp2_graph_allreduce( + void *sendbuf, void *recvbuf, size_t count, + infiniDtype_t datatype, infinicclReduceOp_t op, + infinicclComm_t comm, cudaStream_t stream) { + if (comm == nullptr || comm->world_size != 2 || + datatype != INFINI_DTYPE_BF16 || op != INFINICCL_SUM || + count == 0 || count > kHygonTp2StageCapacityElements || (count % 8) != 0) return false; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + unsigned long long capture_id = 0; + if (cudaStreamGetCaptureInfo(stream, &capture_status, &capture_id) != cudaSuccess || + capture_status != cudaStreamCaptureStatusActive) return false; + auto state = get_hygon_tp2_state(comm); + if (state == nullptr || comm->rank < 0 || comm->rank >= 2) return false; + const int rank = comm->rank; + size_t element_offset = 0; + if (!reserve_hygon_tp2_stage(state, rank, capture_id, count, &element_offset)) return false; + const size_t pack_count = count / 8; + int blocks = static_cast( + std::min(kHygonTp2MaxBlocks, (pack_count + 255) / 256)); + blocks = std::max(blocks, 1); + hygon_tp2_bf16_allreduce_kernel<<>>( + state->rank_data[rank], state->rank_signals, state->signals[rank], + static_cast(sendbuf), + static_cast<__nv_bfloat16 *>(recvbuf), rank, pack_count, element_offset / 8); + return cudaGetLastError() == cudaSuccess; +} + +} // namespace +#endif + infiniStatus_t commInitAll( infinicclComm_t *comms, int ndevice, @@ -75,6 +454,10 @@ infiniStatus_t commInitAll( comms[i] = new InfinicclComm{INFINI_DEVICE_NVIDIA, device_ids[i], (void *)(nccl_comms[i]), i, ndevice}; } +#if defined(ENABLE_HYGON_API) + register_hygon_tp2_state(comms, ndevice, device_ids); +#endif + return INFINI_STATUS_SUCCESS; } @@ -112,6 +495,9 @@ infiniStatus_t commInitRank( } infiniStatus_t commDestroy(infinicclComm_t comm) { +#if defined(ENABLE_HYGON_API) + erase_hygon_tp2_state(comm); +#endif CHECK_NCCL(ncclCommDestroy(getNcclComm(comm))); delete comm; return INFINI_STATUS_SUCCESS; @@ -140,6 +526,14 @@ infiniStatus_t allReduce( INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U32, INFINI_DTYPE_U64); +#if defined(ENABLE_HYGON_API) + if (try_hygon_tp2_graph_allreduce( + sendbuf, recvbuf, count, datatype, op, comm, + getCudaStream(stream))) { + return INFINI_STATUS_SUCCESS; + } +#endif + CHECK_NCCL(ncclAllReduce(sendbuf, recvbuf, count, getNcclDtype(datatype), getNcclRedOp(op), getNcclComm(comm), getCudaStream(stream))); From dcd3c0e81d0c2656f1aff8c574479be16513a071 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Sun, 12 Jul 2026 14:20:28 +0800 Subject: [PATCH 05/16] feat(hygon): optimize Qwen3 MoE graph inference on TP8 --- include/infinicore/adaptor/aten_adaptor.hpp | 2 +- .../infinicore/adaptor/lightop_adaptor.hpp | 89 ++- include/infinicore/graph/graph.hpp | 6 +- include/infinicore/nn/rope.hpp | 4 + include/infinicore/ops.hpp | 2 + include/infinicore/ops/moe_w8a8_marlin.hpp | 60 ++ .../infinicore/ops/rms_rotary_embedding.hpp | 35 + src/infiniccl/cuda/infiniccl_cuda.cu | 566 +++++++++++++++ src/infinicore/adaptor/lightop_adaptor.cc | 650 ++++++++++++++++-- .../allocators/pinnable_block_allocator.cc | 26 +- src/infinicore/nn/rope.cc | 18 + .../hygon/mha_kvcache_flashattn_hygon.cc | 92 +-- .../hygon/moe_align_lightop_hygon.cc | 129 ++++ .../hygon/moe_w16a16_marlin_hygon.cc | 53 +- .../hygon/moe_w8a8_marlin_hygon.cc | 304 ++++++++ .../ops/moe_w8a8_marlin/moe_w8a8_marlin.cc | 150 ++++ .../hygon/mha_varlen_flashattn_hygon.cc | 75 +- .../per_channel_quant_i8_lightop_hygon.cc | 92 +++ .../rms_rotary_embedding_lightop_hygon.cc | 89 +++ .../rms_rotary_embedding.cc | 76 ++ .../hygon/scaled_mm_i8_lightop_hygon.cc | 102 +++ .../hygon/silu_and_mul_lightop_hygon.cc | 3 +- .../ops/embedding/nvidia/embedding_nvidia.cu | 50 +- 23 files changed, 2354 insertions(+), 319 deletions(-) create mode 100644 include/infinicore/ops/moe_w8a8_marlin.hpp create mode 100644 include/infinicore/ops/rms_rotary_embedding.hpp create mode 100644 src/infinicore/ops/moe_align/hygon/moe_align_lightop_hygon.cc create mode 100644 src/infinicore/ops/moe_w8a8_marlin/hygon/moe_w8a8_marlin_hygon.cc create mode 100644 src/infinicore/ops/moe_w8a8_marlin/moe_w8a8_marlin.cc create mode 100644 src/infinicore/ops/per_channel_quant_i8/hygon/per_channel_quant_i8_lightop_hygon.cc create mode 100644 src/infinicore/ops/rms_rotary_embedding/hygon/rms_rotary_embedding_lightop_hygon.cc create mode 100644 src/infinicore/ops/rms_rotary_embedding/rms_rotary_embedding.cc create mode 100644 src/infinicore/ops/scaled_mm_i8/hygon/scaled_mm_i8_lightop_hygon.cc diff --git a/include/infinicore/adaptor/aten_adaptor.hpp b/include/infinicore/adaptor/aten_adaptor.hpp index daf428f50..c65eee830 100644 --- a/include/infinicore/adaptor/aten_adaptor.hpp +++ b/include/infinicore/adaptor/aten_adaptor.hpp @@ -44,7 +44,7 @@ inline at::ScalarType to_at_dtype(DataType dtype) { case DataType::I64: return at::kLong; default: - throw std::runtime_error("Unsupported dtype for ATen"); + throw std::runtime_error("Unsupported dtype for ATen: " + infinicore::toString(dtype)); } } diff --git a/include/infinicore/adaptor/lightop_adaptor.hpp b/include/infinicore/adaptor/lightop_adaptor.hpp index 46256c4bc..0d3d9ed21 100644 --- a/include/infinicore/adaptor/lightop_adaptor.hpp +++ b/include/infinicore/adaptor/lightop_adaptor.hpp @@ -3,29 +3,42 @@ #include +#include #include -#include namespace infinicore::adaptor::lightop { bool available(); -bool enabled_by_env(); +void preload_moe_w16a16_ops(); -void preload_basic_ops(); +void preload_moe_w8a8_ops(); + +void preload_moe_align(); + +void preload_moe_w8a8_marlin_asm(); void preload_silu_and_mul(); -void fused_rms_norm_contiguous( - at::Tensor &out, - at::Tensor &input, - at::Tensor &weight, - double epsilon); +void preload_rms_rotary_embedding(); void fuse_silu_and_mul( at::Tensor &input, at::Tensor &output); +void rms_rotary_embedding_fuse( + at::Tensor &positions, + at::Tensor &query, + at::Tensor &key, + int64_t head_size, + at::Tensor &cos_sin_cache, + bool is_neox, + at::Tensor q_weight, + at::Tensor k_weight, + const std::optional &q_bias = std::nullopt, + const std::optional &k_bias = std::nullopt, + double epsilon = 1e-6); + void moe_sum( at::Tensor &input, at::Tensor &output, @@ -35,15 +48,18 @@ void moe_sum( float factor = 1.0f, int expect_m = -1); -std::vector moe_fused_gate( - at::Tensor &input, - at::Tensor &bias, - int64_t num_expert_group, - int64_t topk_group, - int64_t topk, - int64_t num_fused_shared_experts, - double routed_scaling_factor); - +void moe_align_block_size( + at::Tensor topk_ids, + int64_t num_experts, + int64_t block_size, + at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded, + const std::optional &expert_map = std::nullopt, + const std::optional &expert_mask = std::nullopt, + const std::optional &num_local_tokens = std::nullopt, + bool is_ep = false, + bool fuse_fill = true); void moe_gemm_marlin_w16a16( at::Tensor input, at::Tensor b_qweight, @@ -56,6 +72,45 @@ void moe_gemm_marlin_w16a16( int mode, int delta); +void moe_gemm_marlin_w8a8( + at::Tensor input, + at::Tensor b_qweight, + at::Tensor output, + at::Tensor a_scale, + at::Tensor b_scale, + const std::optional &topk_weights, + at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded, + int64_t top_k, + int mode, + int delta); + +void fuse_silu_mul_quant( + at::Tensor &input, + at::Tensor &output, + at::Tensor &scales, + std::optional &num_local_tokens, + int topk, + int expect_m, + std::optional &expert_ids); + +void preload_w8a8_linear_ops(); + +void per_token_dynamic_quant_int8( + at::Tensor &output, + const at::Tensor &input, + at::Tensor &scales, + const at::Tensor &smooth); + +void blaslt_w8a8_gemm( + at::Tensor &output, + const at::Tensor &a, + const at::Tensor &b, + const at::Tensor &scale_a, + const at::Tensor &scale_b, + const std::optional &bias); + } // namespace infinicore::adaptor::lightop #endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/include/infinicore/graph/graph.hpp b/include/infinicore/graph/graph.hpp index 2cc97e023..9106ab3cd 100644 --- a/include/infinicore/graph/graph.hpp +++ b/include/infinicore/graph/graph.hpp @@ -31,9 +31,9 @@ class DispatchableGraphOperator : public GraphOperator { protected: using run_schema = void (*)(void *); using cleanup_schema = void (*)(void **); - void *planned_meta_; - run_schema runner_; - cleanup_schema deleter_; + void *planned_meta_ = nullptr; + run_schema runner_ = nullptr; + cleanup_schema deleter_ = nullptr; }; class Graph { diff --git a/include/infinicore/nn/rope.hpp b/include/infinicore/nn/rope.hpp index eaeba8712..e9ddc6dc4 100644 --- a/include/infinicore/nn/rope.hpp +++ b/include/infinicore/nn/rope.hpp @@ -85,6 +85,9 @@ class RoPE : public Module { double theta() const { return theta_; } Algo algo() const { return algo_; } DataType dtype() const { return dtype_; } + Tensor sin_cache() const { return sin_cache_; } + Tensor cos_cache() const { return cos_cache_; } + Tensor cos_sin_cache() const { return cos_sin_cache_; } const std::optional> &mrope_section() const { return mrope_section_; } bool mrope_interleaved() const { return mrope_interleaved_; } const Tensor &sin_cache() const { return sin_cache_; } @@ -97,6 +100,7 @@ class RoPE : public Module { // Buffers (sin and cos cache tables) - not exposed in state_dict INFINICORE_NN_BUFFER(sin_cache); INFINICORE_NN_BUFFER(cos_cache); + INFINICORE_NN_BUFFER(cos_sin_cache); private: void initialize_cache(); diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index 8a5269c70..a95cc17e7 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -58,6 +58,7 @@ #include "ops/moe_topk_sigmoid.hpp" #include "ops/moe_topk_softmax.hpp" #include "ops/moe_w16a16_marlin.hpp" +#include "ops/moe_w8a8_marlin.hpp" #include "ops/nrm2.hpp" #include "ops/ones.hpp" #include "ops/paged_attention.hpp" @@ -73,6 +74,7 @@ #include "ops/recurrent_gated_delta_rule.hpp" #include "ops/relu.hpp" #include "ops/rms_norm.hpp" +#include "ops/rms_rotary_embedding.hpp" #include "ops/rope.hpp" #include "ops/rot.hpp" #include "ops/rotg.hpp" diff --git a/include/infinicore/ops/moe_w8a8_marlin.hpp b/include/infinicore/ops/moe_w8a8_marlin.hpp new file mode 100644 index 000000000..f2a014860 --- /dev/null +++ b/include/infinicore/ops/moe_w8a8_marlin.hpp @@ -0,0 +1,60 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(MoeW8A8MarlinFusedDense, + Tensor, + Tensor, + Tensor, + Tensor, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + size_t, + int, + size_t, + int, + int, + int); + +Tensor moe_w8a8_marlin_pack(const Tensor &weight); + +void moe_w8a8_marlin_fused_dense_( + Tensor output, + Tensor cache13, + Tensor cache2_i8, + Tensor input_i8, + Tensor input_scale, + Tensor cache2_scale, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &w13_scale, + const Tensor &w2_scale, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + size_t block_size_m, + int delta0, + int mode1, + int delta1); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/rms_rotary_embedding.hpp b/include/infinicore/ops/rms_rotary_embedding.hpp new file mode 100644 index 000000000..540f80871 --- /dev/null +++ b/include/infinicore/ops/rms_rotary_embedding.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +#include + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(RMSRotaryEmbedding, + Tensor, + Tensor, + const Tensor &, + int64_t, + const Tensor &, + bool, + const Tensor &, + const Tensor &, + float); + +bool rms_rotary_embedding_fuse_available(const Device &device); + +void rms_rotary_embedding_fuse_(Tensor query, + Tensor key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache, + bool is_neox, + const Tensor &q_weight, + const Tensor &k_weight, + float epsilon = 1e-6f); + +} // namespace infinicore::op diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu index 4de94db0c..9fdcd6749 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ b/src/infiniccl/cuda/infiniccl_cuda.cu @@ -8,6 +8,7 @@ #if defined(ENABLE_HYGON_API) #include +#include #include #include #include @@ -81,6 +82,15 @@ namespace { constexpr int kHygonTp2MaxBlocks = 80; constexpr size_t kHygonTp2StageCapacityElements = 1u << 22; +constexpr int kHygonTp8WorldSize = 8; +constexpr int kHygonTp8Threads = 512; +constexpr int kHygonTp8MaxBlocks = 80; +constexpr size_t kHygonTp8OneStageMaxBytes = 80u * 1024u; +constexpr size_t kHygonTp8TwoStageMaxBytes = 512u * 1024u; +constexpr int kHygonHipSuccess = 0; +constexpr unsigned int kHygonHipDeviceMallocUncached = 0x3; +constexpr unsigned int kHygonHipEventDisableTiming = 0x2; +constexpr unsigned int kHygonHipEventReleaseToSystem = 0x80000000u; struct HygonTp2Signal { alignas(128) uint32_t start[kHygonTp2MaxBlocks][8]; @@ -100,6 +110,31 @@ struct alignas(16) HygonBf16Pack { __nv_bfloat16 values[8]; }; +static_assert(sizeof(HygonBf16Pack) == 16); + +struct HygonTp8Signal { + alignas(128) uint32_t start[kHygonTp8MaxBlocks][kHygonTp8WorldSize]; + alignas(128) uint32_t end[kHygonTp8MaxBlocks][kHygonTp8WorldSize]; + alignas(128) uint32_t flag[kHygonTp8MaxBlocks]; +}; + +constexpr size_t kHygonTp8TwoStageScratchBytes = + kHygonTp8TwoStageMaxBytes / kHygonTp8WorldSize + + (kHygonTp8WorldSize - 1) * sizeof(HygonBf16Pack); +constexpr size_t kHygonTp8SignalAllocationBytes = + sizeof(HygonTp8Signal) + kHygonTp8TwoStageScratchBytes; + +struct alignas(16) HygonTp8RankData { + const void *ptrs[kHygonTp8WorldSize]; +}; + +struct alignas(16) HygonTp8RankSignals { + HygonTp8Signal *signals[kHygonTp8WorldSize]; +}; + +static_assert(sizeof(HygonTp8RankData) == 64 && alignof(HygonTp8RankData) == 16); +static_assert(sizeof(HygonTp8RankSignals) == 64 && alignof(HygonTp8RankSignals) == 16); + struct HygonCudaDriverApi { void *library = nullptr; decltype(&cuMemGetAllocationGranularity) mem_get_allocation_granularity = nullptr; @@ -141,6 +176,51 @@ HygonCudaDriverApi &hygon_cuda_driver_api() { return api; } +struct HygonHipExtApi { + void *library = nullptr; + int (*ext_malloc_with_flags)(void **, size_t, unsigned int) = nullptr; + int (*memset)(void *, int, size_t) = nullptr; + int (*free)(void *) = nullptr; + int (*event_create_with_flags)(void **, unsigned int) = nullptr; + int (*event_destroy)(void *) = nullptr; + int (*ext_launch_kernel)( + const void *, dim3, dim3, void **, size_t, + void *, void *, void *, int) = nullptr; + bool available = false; + + HygonHipExtApi() { + constexpr const char *candidates[] = { + "libgalaxyhip.so.5", + "libgalaxyhip.so", + "/opt/dtk/hip/lib/libgalaxyhip.so.5", + "/opt/dtk/lib/libgalaxyhip.so.5", + }; + for (const char *candidate : candidates) { + library = dlopen(candidate, RTLD_NOW | RTLD_LOCAL); + if (library != nullptr) break; + } + available = library != nullptr && + load(ext_malloc_with_flags, "hipExtMallocWithFlags") && + load(memset, "hipMemset") && + load(free, "hipFree") && + load(event_create_with_flags, "hipEventCreateWithFlags") && + load(event_destroy, "hipEventDestroy") && + load(ext_launch_kernel, "hipExtLaunchKernel"); + } + +private: + template + bool load(T &symbol, const char *name) { + symbol = reinterpret_cast(dlsym(library, name)); + return symbol != nullptr; + } +}; + +HygonHipExtApi &hygon_hip_ext_api() { + static HygonHipExtApi *api = new HygonHipExtApi(); + return *api; +} + struct HygonVmmAllocation { void *ptr = nullptr; size_t size = 0; @@ -186,6 +266,49 @@ struct HygonTp2AllReduceState { std::mutex hygon_tp2_states_mutex; std::unordered_map> hygon_tp2_states; +struct HygonTp8AllReduceState { + int device_ids[kHygonTp8WorldSize]{}; + HygonTp8Signal *signals[kHygonTp8WorldSize]{}; + void *release_events[kHygonTp8WorldSize]{}; + HygonTp8RankSignals rank_signals{}; + + struct CaptureRendezvous { + std::mutex mutex; + std::condition_variable condition; + uint64_t generation = 0; + uint32_t arrived_mask = 0; + int departed = 0; + bool ready = false; + bool use_custom = false; + bool metadata_error = false; + const void *sendbufs[kHygonTp8WorldSize]{}; + void *recvbufs[kHygonTp8WorldSize]{}; + size_t counts[kHygonTp8WorldSize]{}; + infiniDtype_t datatypes[kHygonTp8WorldSize]{}; + infinicclReduceOp_t ops[kHygonTp8WorldSize]{}; + bool eligible[kHygonTp8WorldSize]{}; + } rendezvous; + + ~HygonTp8AllReduceState() { + int previous_device = 0; + const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess; + auto &hip = hygon_hip_ext_api(); + for (int rank = 0; rank < kHygonTp8WorldSize; ++rank) { + cudaSetDevice(device_ids[rank]); + if (hip.available && release_events[rank] != nullptr) { + hip.event_destroy(release_events[rank]); + } + if (hip.available && signals[rank] != nullptr) { + hip.free(signals[rank]); + } + } + if (restore_device) cudaSetDevice(previous_device); + } +}; + +std::mutex hygon_tp8_states_mutex; +std::unordered_map> hygon_tp8_states; + bool allocate_hygon_vmm(HygonVmmAllocation &allocation, int owner_device, const int device_ids[2], @@ -439,6 +562,438 @@ bool try_hygon_tp2_graph_allreduce( return cudaGetLastError() == cudaSuccess; } +template +__device__ __forceinline__ uint32_t hygon_tp8_start_sync( + const HygonTp8RankSignals &rank_signals, + HygonTp8Signal *self_signal, + int rank, + uint32_t *block_flag) { + if (threadIdx.x == 0) { + *block_flag = __scoped_atomic_load_n( + &self_signal->flag[blockIdx.x], + __ATOMIC_ACQUIRE, __MEMORY_SCOPE_SYSTEM) + + 1; + } + __syncthreads(); + const uint32_t next_flag = *block_flag; + if (threadIdx.x < NumRanks) { + __scoped_atomic_store_n( + &rank_signals.signals[threadIdx.x]->start[blockIdx.x][rank], + next_flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM); + while (__scoped_atomic_load_n( + &self_signal->start[blockIdx.x][threadIdx.x], + __ATOMIC_ACQUIRE, __MEMORY_SCOPE_SYSTEM) != next_flag) { + } + } + __syncthreads(); + if (threadIdx.x == 0) { + __scoped_atomic_store_n( + &self_signal->flag[blockIdx.x], next_flag, + __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM); + } + return next_flag; +} + +template +__device__ __forceinline__ void hygon_tp8_end_sync( + const HygonTp8RankSignals &rank_signals, + HygonTp8Signal *self_signal, + int rank, + uint32_t flag) { + __syncthreads(); + if (threadIdx.x < NumRanks) { + __scoped_atomic_store_n( + &rank_signals.signals[threadIdx.x]->end[blockIdx.x][rank], + flag, __ATOMIC_RELEASE, __MEMORY_SCOPE_SYSTEM); + while (__scoped_atomic_load_n( + &self_signal->end[blockIdx.x][threadIdx.x], + __ATOMIC_ACQUIRE, __MEMORY_SCOPE_SYSTEM) != flag) { + } + } + __syncthreads(); +} + +__global__ __launch_bounds__(kHygonTp8Threads, 1) void hygon_tp8_bf16_allreduce_kernel( + HygonTp8RankData rank_data, + HygonTp8RankSignals rank_signals, + HygonTp8Signal *self_signal, + __nv_bfloat16 *output, + int rank, + size_t pack_count) { + constexpr int num_ranks = kHygonTp8WorldSize; + constexpr int threads_per_rank = kHygonTp8Threads / num_ranks; + __shared__ HygonBf16Pack shared_packs[kHygonTp8Threads]; + __shared__ uint32_t block_flag; + + const int source_rank = threadIdx.x / threads_per_rank; + const int lane = threadIdx.x % threads_per_rank; + const uint32_t sync_flag = hygon_tp8_start_sync( + rank_signals, self_signal, rank, &block_flag); + + for (size_t base = blockIdx.x * threads_per_rank; + base < pack_count; + base += gridDim.x * threads_per_rank) { + const size_t index = base + lane; + HygonBf16Pack source_pack{}; + if (index < pack_count) { + const auto *source = reinterpret_cast( + rank_data.ptrs[source_rank]); + source_pack = source[index]; + } + shared_packs[threadIdx.x] = source_pack; + __syncthreads(); + + if (source_rank == 0 && index < pack_count) { + float reduced[8]; +#pragma unroll + for (int element = 0; element < 8; ++element) { + reduced[element] = __bfloat162float( + shared_packs[threadIdx.x].values[element]); + } +#pragma unroll + for (int peer = 1; peer < num_ranks; ++peer) { +#pragma unroll + for (int element = 0; element < 8; ++element) { + reduced[element] += __bfloat162float( + shared_packs[peer * threads_per_rank + threadIdx.x] + .values[element]); + } + } + HygonBf16Pack result; +#pragma unroll + for (int element = 0; element < 8; ++element) { + result.values[element] = __float2bfloat16(reduced[element]); + } + reinterpret_cast(output)[index] = result; + } + __syncthreads(); + } + + hygon_tp8_end_sync( + rank_signals, self_signal, rank, sync_flag); +} + +__device__ __forceinline__ HygonBf16Pack *hygon_tp8_two_stage_scratch( + HygonTp8Signal *signal) { + return reinterpret_cast(signal + 1); +} + +__global__ __launch_bounds__(kHygonTp8Threads, 1) void hygon_tp8_bf16_allreduce_2stage_kernel( + HygonTp8RankData rank_data, + HygonTp8RankSignals rank_signals, + HygonTp8Signal *self_signal, + __nv_bfloat16 *output, + int rank, + size_t pack_count) { + constexpr int num_ranks = kHygonTp8WorldSize; + constexpr int threads_per_rank = kHygonTp8Threads / num_ranks; + __shared__ HygonBf16Pack shared_packs[kHygonTp8Threads]; + __shared__ uint32_t block_flag; + + const int source_rank = threadIdx.x / threads_per_rank; + const int lane = threadIdx.x % threads_per_rank; + const size_t thread_index = blockIdx.x * threads_per_rank + lane; + const size_t thread_stride = gridDim.x * threads_per_rank; + const size_t part = pack_count / num_ranks; + const size_t remainder = pack_count % num_ranks; + const size_t slice_begin = static_cast(rank) * part; + const size_t slice_end = rank == num_ranks - 1 + ? pack_count + : slice_begin + part; + const size_t largest_part = part + remainder; + HygonBf16Pack *local_scratch = + hygon_tp8_two_stage_scratch(self_signal); + + const uint32_t sync_flag = hygon_tp8_start_sync( + rank_signals, self_signal, rank, &block_flag); + + // Stage 1: each rank reduces one disjoint slice into its peer-visible + // uncached scratch buffer. + for (size_t base = slice_begin + blockIdx.x * threads_per_rank; + base < slice_end; + base += gridDim.x * threads_per_rank) { + const size_t index = base + lane; + HygonBf16Pack source_pack{}; + if (index < slice_end) { + const auto *source = reinterpret_cast( + rank_data.ptrs[source_rank]); + source_pack = source[index]; + } + shared_packs[threadIdx.x] = source_pack; + __syncthreads(); + + if (source_rank == 0 && index < slice_end) { + float reduced[8]; +#pragma unroll + for (int element = 0; element < 8; ++element) { + reduced[element] = __bfloat162float( + shared_packs[threadIdx.x].values[element]); + } +#pragma unroll + for (int peer = 1; peer < num_ranks; ++peer) { +#pragma unroll + for (int element = 0; element < 8; ++element) { + reduced[element] += __bfloat162float( + shared_packs[peer * threads_per_rank + threadIdx.x] + .values[element]); + } + } + HygonBf16Pack result; +#pragma unroll + for (int element = 0; element < 8; ++element) { + result.values[element] = __float2bfloat16(reduced[element]); + } + local_scratch[index - slice_begin] = result; + } + __syncthreads(); + } + + // Release makes the reduced slices visible before peers gather them. + hygon_tp8_end_sync( + rank_signals, self_signal, rank, sync_flag); + + // Stage 2: the eight thread groups gather one source rank each. + const HygonBf16Pack *source_scratch = + hygon_tp8_two_stage_scratch(rank_signals.signals[source_rank]); + auto *packed_output = reinterpret_cast(output); + for (size_t offset = thread_index; + offset < largest_part; + offset += thread_stride) { + if (source_rank == num_ranks - 1 || offset < part) { + packed_output[static_cast(source_rank) * part + offset] = + source_scratch[offset]; + } + } + return; +} + +std::shared_ptr create_hygon_tp8_state( + int ndevice, const int *device_ids) { + if (ndevice != kHygonTp8WorldSize || device_ids == nullptr) return nullptr; + auto &hip = hygon_hip_ext_api(); + if (!hip.available) return nullptr; + auto state = std::make_shared(); + std::memcpy(state->device_ids, device_ids, sizeof(state->device_ids)); + + int previous_device = 0; + const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess; + auto fail = [&]() -> std::shared_ptr { + if (restore_device) cudaSetDevice(previous_device); + return nullptr; + }; + + for (int rank = 0; rank < kHygonTp8WorldSize; ++rank) { + if (cudaSetDevice(device_ids[rank]) != cudaSuccess) return fail(); + for (int peer = 0; peer < kHygonTp8WorldSize; ++peer) { + if (peer == rank) continue; + int can_access = 0; + if (cudaDeviceCanAccessPeer( + &can_access, device_ids[rank], device_ids[peer]) != cudaSuccess || + can_access == 0) return fail(); + const cudaError_t enable_status = + cudaDeviceEnablePeerAccess(device_ids[peer], 0); + if (enable_status == cudaErrorPeerAccessAlreadyEnabled) { + (void)cudaGetLastError(); + } else if (enable_status != cudaSuccess) { + return fail(); + } + } + if (hip.ext_malloc_with_flags( + reinterpret_cast(&state->signals[rank]), + kHygonTp8SignalAllocationBytes, + kHygonHipDeviceMallocUncached) != kHygonHipSuccess || + hip.memset(state->signals[rank], 0, + kHygonTp8SignalAllocationBytes) != + kHygonHipSuccess) { + return fail(); + } + if (hip.event_create_with_flags( + &state->release_events[rank], + kHygonHipEventReleaseToSystem | kHygonHipEventDisableTiming) != + kHygonHipSuccess) { + return fail(); + } + state->rank_signals.signals[rank] = state->signals[rank]; + } + if (restore_device) cudaSetDevice(previous_device); + return state; +} + +void register_hygon_tp8_state(infinicclComm_t *comms, + int ndevice, + const int *device_ids) { + auto state = create_hygon_tp8_state(ndevice, device_ids); + if (state == nullptr) return; + std::lock_guard lock(hygon_tp8_states_mutex); + for (int rank = 0; rank < kHygonTp8WorldSize; ++rank) { + hygon_tp8_states.emplace(comms[rank], state); + } +} + +void erase_hygon_tp8_state(infinicclComm_t comm) { + std::lock_guard lock(hygon_tp8_states_mutex); + hygon_tp8_states.erase(comm); +} + +std::shared_ptr get_hygon_tp8_state(infinicclComm_t comm) { + std::lock_guard lock(hygon_tp8_states_mutex); + auto found = hygon_tp8_states.find(comm); + return found == hygon_tp8_states.end() ? nullptr : found->second; +} + +bool rendezvous_hygon_tp8_graph_inputs( + const std::shared_ptr &state, + int rank, + void *sendbuf, + void *recvbuf, + size_t count, + infiniDtype_t datatype, + infinicclReduceOp_t op, + HygonTp8RankData *rank_data, + bool *metadata_error) { + const bool local_eligible = + sendbuf != nullptr && recvbuf != nullptr && sendbuf != recvbuf && + datatype == INFINI_DTYPE_BF16 && op == INFINICCL_SUM && + count != 0 && (count % 8) == 0 && + count <= kHygonTp8TwoStageMaxBytes / sizeof(__nv_bfloat16) && + (reinterpret_cast(sendbuf) % alignof(HygonBf16Pack)) == 0 && + (reinterpret_cast(recvbuf) % alignof(HygonBf16Pack)) == 0; + + auto &rendezvous = state->rendezvous; + std::unique_lock lock(rendezvous.mutex); + const uint64_t generation = rendezvous.generation; + const uint32_t rank_bit = uint32_t{1} << rank; + if ((rendezvous.arrived_mask & rank_bit) != 0) std::abort(); + + rendezvous.arrived_mask |= rank_bit; + rendezvous.sendbufs[rank] = sendbuf; + rendezvous.recvbufs[rank] = recvbuf; + rendezvous.counts[rank] = count; + rendezvous.datatypes[rank] = datatype; + rendezvous.ops[rank] = op; + rendezvous.eligible[rank] = local_eligible; + + constexpr uint32_t all_ranks_mask = + (uint32_t{1} << kHygonTp8WorldSize) - 1; + if (rendezvous.arrived_mask == all_ranks_mask) { + bool signatures_match = true; + bool use_custom = true; + for (int peer = 0; peer < kHygonTp8WorldSize; ++peer) { + signatures_match = + signatures_match && + rendezvous.counts[peer] == rendezvous.counts[0] && + rendezvous.datatypes[peer] == rendezvous.datatypes[0] && + rendezvous.ops[peer] == rendezvous.ops[0]; + use_custom = use_custom && rendezvous.eligible[peer]; + } + rendezvous.metadata_error = !signatures_match; + rendezvous.use_custom = signatures_match && use_custom; + rendezvous.ready = true; + rendezvous.condition.notify_all(); + } else { + rendezvous.condition.wait(lock, [&] { + return rendezvous.ready && rendezvous.generation == generation; + }); + } + + const bool use_custom = rendezvous.use_custom; + *metadata_error = rendezvous.metadata_error; + if (use_custom) { + for (int peer = 0; peer < kHygonTp8WorldSize; ++peer) { + rank_data->ptrs[peer] = rendezvous.sendbufs[peer]; + } + } + + if (++rendezvous.departed == kHygonTp8WorldSize) { + rendezvous.arrived_mask = 0; + rendezvous.departed = 0; + rendezvous.ready = false; + rendezvous.use_custom = false; + rendezvous.metadata_error = false; + ++rendezvous.generation; + rendezvous.condition.notify_all(); + } else { + rendezvous.condition.wait(lock, [&] { + return rendezvous.generation != generation; + }); + } + return use_custom; +} + +enum class HygonTp8AllReduceResult { + Fallback, + Success, + Error, +}; + +HygonTp8AllReduceResult try_hygon_tp8_graph_allreduce( + void *sendbuf, void *recvbuf, size_t count, + infiniDtype_t datatype, infinicclReduceOp_t op, + infinicclComm_t comm, cudaStream_t stream) { + if (comm == nullptr || comm->world_size != kHygonTp8WorldSize || + comm->rank < 0 || comm->rank >= kHygonTp8WorldSize) { + return HygonTp8AllReduceResult::Fallback; + } + + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + unsigned long long capture_id = 0; + if (cudaStreamGetCaptureInfo(stream, &capture_status, &capture_id) != cudaSuccess) { + return HygonTp8AllReduceResult::Error; + } + (void)capture_id; + if (capture_status != cudaStreamCaptureStatusActive) { + return HygonTp8AllReduceResult::Fallback; + } + + auto state = get_hygon_tp8_state(comm); + if (state == nullptr) return HygonTp8AllReduceResult::Fallback; + const int rank = comm->rank; + HygonTp8RankData rank_data{}; + bool metadata_error = false; + if (!rendezvous_hygon_tp8_graph_inputs( + state, rank, sendbuf, recvbuf, count, datatype, op, + &rank_data, &metadata_error)) { + if (metadata_error) return HygonTp8AllReduceResult::Error; + return HygonTp8AllReduceResult::Fallback; + } + + constexpr int threads_per_rank = + kHygonTp8Threads / kHygonTp8WorldSize; + size_t pack_count = count / 8; + const bool use_two_stage = + count * sizeof(__nv_bfloat16) >= kHygonTp8OneStageMaxBytes; + const size_t work_pack_count = use_two_stage + ? pack_count / kHygonTp8WorldSize + + pack_count % kHygonTp8WorldSize + : pack_count; + int blocks = static_cast(std::min( + kHygonTp8MaxBlocks, + (work_pack_count + threads_per_rank - 1) / threads_per_rank)); + blocks = std::max(blocks, 1); + HygonTp8RankSignals rank_signals = state->rank_signals; + HygonTp8Signal *self_signal = state->signals[rank]; + auto *output = static_cast<__nv_bfloat16 *>(recvbuf); + int kernel_rank = rank; + void *args[] = { + &rank_data, &rank_signals, &self_signal, + &output, &kernel_rank, &pack_count, + }; + auto &hip = hygon_hip_ext_api(); + const void *kernel = use_two_stage + ? reinterpret_cast( + hygon_tp8_bf16_allreduce_2stage_kernel) + : reinterpret_cast( + hygon_tp8_bf16_allreduce_kernel); + const int launch_status = hip.ext_launch_kernel( + kernel, + dim3(blocks), dim3(kHygonTp8Threads), args, 0, + reinterpret_cast(stream), nullptr, + state->release_events[rank], 0); + return launch_status == kHygonHipSuccess + ? HygonTp8AllReduceResult::Success + : HygonTp8AllReduceResult::Error; +} + } // namespace #endif @@ -456,6 +1011,7 @@ infiniStatus_t commInitAll( #if defined(ENABLE_HYGON_API) register_hygon_tp2_state(comms, ndevice, device_ids); + register_hygon_tp8_state(comms, ndevice, device_ids); #endif return INFINI_STATUS_SUCCESS; @@ -497,6 +1053,7 @@ infiniStatus_t commInitRank( infiniStatus_t commDestroy(infinicclComm_t comm) { #if defined(ENABLE_HYGON_API) erase_hygon_tp2_state(comm); + erase_hygon_tp8_state(comm); #endif CHECK_NCCL(ncclCommDestroy(getNcclComm(comm))); delete comm; @@ -527,6 +1084,15 @@ infiniStatus_t allReduce( INFINI_DTYPE_U32, INFINI_DTYPE_U64); #if defined(ENABLE_HYGON_API) + const auto tp8_result = try_hygon_tp8_graph_allreduce( + sendbuf, recvbuf, count, datatype, op, comm, + getCudaStream(stream)); + if (tp8_result == HygonTp8AllReduceResult::Success) { + return INFINI_STATUS_SUCCESS; + } + if (tp8_result == HygonTp8AllReduceResult::Error) { + return INFINI_STATUS_INTERNAL_ERROR; + } if (try_hygon_tp2_graph_allreduce( sendbuf, recvbuf, count, datatype, op, comm, getCudaStream(stream))) { diff --git a/src/infinicore/adaptor/lightop_adaptor.cc b/src/infinicore/adaptor/lightop_adaptor.cc index b21639be1..37f50477e 100644 --- a/src/infinicore/adaptor/lightop_adaptor.cc +++ b/src/infinicore/adaptor/lightop_adaptor.cc @@ -1,33 +1,55 @@ #if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) #include "infinicore/adaptor/lightop_adaptor.hpp" +#include "infinicore/adaptor/aten_adaptor.hpp" #include +#include -#include +#include #include +#include #include #include #include #include +#include namespace infinicore::adaptor::lightop { namespace { constexpr const char *kDefaultLightopSo = "/usr/local/lib/python3.10/dist-packages/lightop/op.cpython-310-x86_64-linux-gnu.so"; - -constexpr const char *kFusedRmsNormSymbol = - "_ZN2at6native25fused_rms_norm_contiguousERNS_6TensorES2_S2_d"; +constexpr const char *kDefaultLmslimQuantSo = + "/usr/local/lib/python3.10/dist-packages/lmslimquant.cpython-310-x86_64-linux-gnu.so"; +constexpr const char *kDefaultLightopGpuTarget = "gfx936"; +constexpr const char *kDefaultLightopAsmDir = + "/usr/local/lib/python3.10/dist-packages/lightop/hsa/gfx936/"; constexpr const char *kFuseSiluAndMulSymbol = "_ZN2at6native17fuse_silu_and_mulERNS_6TensorES2_"; +constexpr const char *kRmsRotaryEmbeddingFuseSymbol = + "_ZN2at6native25rms_rotary_embedding_fuseERNS_6TensorES2_S2_lS2_bS1_S1_St8optionalIS1_ES4_d"; constexpr const char *kMoeSumSymbol = "_ZN2at6native7moe_sumERNS_6TensorES2_RKSt8optionalIS1_ES6_S6_fi"; -constexpr const char *kMoeFusedGateSymbol = - "_ZN2at6native14moe_fused_gateERNS_6TensorES2_lllld"; +constexpr const char *kMoeAlignBlockSizeSymbol = + "_ZN2at6native20moe_align_block_sizeENS_6TensorEllS1_S1_S1_RKSt8optionalIS1_ES5_S5_bb"; constexpr const char *kMoeGemmW16A16Symbol = "_ZN2at6native15moe_gemm_w16a16ENS_6TensorES1_S1_St8optionalIS1_ES1_S1_S1_lii"; constexpr const char *kMoeMarlinW16A16AsmSymbol = "_ZN2at6native21moe_marlin_w16a16_asmENS_6TensorES1_S1_St8optionalIS1_ES1_S1_S1_iii"; +constexpr const char *kMoeGemmW8A8Symbol = + "_ZN2at6native20moe_gemm_marlin_w8a8ENS_6TensorES1_S1_S1_S1_St8optionalIS1_ES1_S1_S1_lii"; +constexpr const char *kMoeMarlinW8A8AsmSymbol = + "_ZN2at6native19moe_marlin_w8a8_asmENS_6TensorES1_S1_S1_S1_St8optionalIS1_ES1_S1_S1_jii"; +constexpr const char *kFuseSiluMulQuantSymbol = + "_ZN2at6native19fuse_silu_mul_quantERNS_6TensorES2_S2_RSt8optionalIS1_EiiS5_"; +constexpr const char *kPerTokenDynamicQuantInt8Symbol = + "_ZN2at6native28per_token_dynamic_quant_int8ERNS_6TensorERKS1_S2_S4_"; +constexpr const char *kBlasltW8A8Bf16Symbol = + "_ZN14hipblaslt_gemm14w8a8_bf16_gemmERKN2at6TensorES3_S3_S3_RS1_llllRKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEES3_S3_RKSt8optionalIS1_E"; +constexpr const char *kBlasltW8A8Fp16Symbol = + "_ZN14hipblaslt_gemm14w8a8_fp16_gemmERKN2at6TensorES3_S3_S3_RS1_llllRKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEES3_S3_RKSt8optionalIS1_E"; + +void ensure_default_lightop_env(); class LightopLibrary { public: @@ -62,9 +84,11 @@ class LightopLibrary { return true; } + ensure_default_lightop_env(); + const char *path_env = std::getenv("INFINICORE_LIGHTOP_SO"); const char *path = (path_env != nullptr && path_env[0] != '\0') ? path_env : kDefaultLightopSo; - handle_ = dlopen(path, RTLD_LAZY | RTLD_LOCAL); + handle_ = dlopen(path, RTLD_LAZY | RTLD_GLOBAL); if (handle_ != nullptr) { error_.clear(); return true; @@ -92,13 +116,325 @@ LightopLibrary &library() { return lib; } +constexpr const char *kMoeW8A8MarlinMode1001Co = + "moe_w8a8_channel/moe_w8a8_i8_marlin_64x256x128_TN_BF16_UP.co"; +constexpr const char *kMoeW8A8MarlinMode1001Kernel = + "MOE_W8A8_I8_PERCHANNEL_MARLIN_ASM_TN_MT64x256x128_WGM1_UP"; +constexpr uint32_t kMoeW8A8MarlinNBlock = 256; +constexpr uint32_t kMoeW8A8MarlinWorkgroupSize = 768; + +struct MoeW8A8MarlinMode1001Args { + uint32_t n_block_count; + uint32_t max_m_block_count; + void *output; + void *weight; + void *input; + void *weight_scale; + void *input_scale; + void *topk_weights; + void *sorted_token_ids; + void *expert_ids; + void *num_tokens_post_padded; + uint32_t num_experts; + uint32_t m; + uint32_t n; + uint32_t k; + uint32_t flag0; + uint32_t flag1; + uint32_t output_stride; + uint32_t flag2; + uint32_t flag3; + uint32_t max_tokens_padded; + uint32_t top_k; + float inverse_top_k; + float output_scale; + uint32_t reserved0; + uint32_t reserved1; + uint32_t reserved2; +}; +static_assert(sizeof(MoeW8A8MarlinMode1001Args) == 144); + +struct MoeW8A8MarlinDeviceKernel { + hipModule_t module = nullptr; + hipFunction_t function = nullptr; +}; + +std::mutex &moe_w8a8_marlin_kernel_mutex() { + static std::mutex mutex; + return mutex; +} + +std::unordered_map &moe_w8a8_marlin_kernels() { + static std::unordered_map kernels; + return kernels; +} + +std::string hip_error_message(const std::string &operation, hipError_t status) { + std::ostringstream oss; + oss << operation << " failed with HIP status " << static_cast(status); + const char *message = hipGetErrorString(status); + if (message != nullptr) { + oss << " (" << message << ")"; + } + return oss.str(); +} + +MoeW8A8MarlinDeviceKernel get_moe_w8a8_marlin_mode1001_kernel() { + int device = -1; + auto status = hipGetDevice(&device); + if (status != hipSuccess) { + throw std::runtime_error(hip_error_message("hipGetDevice", status)); + } + + std::lock_guard lock(moe_w8a8_marlin_kernel_mutex()); + auto &kernels = moe_w8a8_marlin_kernels(); + auto found = kernels.find(device); + if (found != kernels.end()) { + return found->second; + } + + ensure_default_lightop_env(); + const char *asm_dir_env = std::getenv("LIGHTOP_ASM_DIR"); + std::string asm_dir = + asm_dir_env != nullptr && asm_dir_env[0] != '\0' + ? asm_dir_env + : kDefaultLightopAsmDir; + if (!asm_dir.empty() && asm_dir.back() != '/') { + asm_dir.push_back('/'); + } + const std::string co_path = asm_dir + kMoeW8A8MarlinMode1001Co; + + MoeW8A8MarlinDeviceKernel kernel; + status = hipModuleLoad(&kernel.module, co_path.c_str()); + if (status != hipSuccess) { + throw std::runtime_error(hip_error_message("hipModuleLoad(" + co_path + ")", status)); + } + status = hipModuleGetFunction( + &kernel.function, kernel.module, kMoeW8A8MarlinMode1001Kernel); + if (status != hipSuccess) { + (void)hipModuleUnload(kernel.module); + throw std::runtime_error(hip_error_message("hipModuleGetFunction", status)); + } + + kernels.emplace(device, kernel); + return kernel; +} + +uint32_t checked_u32(int64_t value, const char *name) { + if (value < 0 || + static_cast(value) > std::numeric_limits::max()) { + throw std::runtime_error(std::string("Hygon W8A8 Marlin ") + name + " exceeds uint32"); + } + return static_cast(value); +} + +bool can_launch_moe_w8a8_marlin_mode1001( + const at::Tensor &input, + const at::Tensor &weight, + const at::Tensor &output, + const at::Tensor &input_scale, + const at::Tensor &weight_scale, + const std::optional &topk_weights, + const at::Tensor &sorted_token_ids, + const at::Tensor &expert_ids, + const at::Tensor &num_tokens_post_padded, + int64_t top_k, + int mode, + int delta) { + return mode == 1001 && delta == 1 && !topk_weights.has_value() && + top_k > 0 && + input.dim() == 2 && weight.dim() == 3 && output.dim() == 3 && + input_scale.dim() == 2 && weight_scale.dim() == 3 && + input.scalar_type() == at::kChar && + weight.scalar_type() == at::kChar && + output.scalar_type() == at::kBFloat16 && + input_scale.scalar_type() == at::kFloat && + weight_scale.scalar_type() == at::kFloat && + sorted_token_ids.scalar_type() == at::kInt && + expert_ids.scalar_type() == at::kInt && + num_tokens_post_padded.scalar_type() == at::kInt && + input.is_contiguous() && weight.is_contiguous() && + output.is_contiguous() && input_scale.is_contiguous() && + weight_scale.is_contiguous() && sorted_token_ids.is_contiguous() && + expert_ids.is_contiguous() && num_tokens_post_padded.is_contiguous(); +} + +void launch_moe_w8a8_marlin_mode1001( + at::Tensor &input, + at::Tensor &weight, + at::Tensor &output, + at::Tensor &input_scale, + at::Tensor &weight_scale, + at::Tensor &sorted_token_ids, + at::Tensor &expert_ids, + at::Tensor &num_tokens_post_padded, + int64_t top_k) { + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t num_experts = weight.size(0); + const int64_t n = output.size(2); + if (output.size(0) != m || output.size(1) != top_k || + weight.size(1) * 64 != k || weight.size(2) != n * 64 || + input_scale.size(0) != m || input_scale.size(1) != 1 || + weight_scale.size(0) != num_experts || + weight_scale.size(1) != n || weight_scale.size(2) != 1 || + num_tokens_post_padded.numel() != 1) { + throw std::runtime_error("Hygon W8A8 Marlin mode 1001 tensor shape mismatch"); + } + + const uint32_t n_u32 = checked_u32(n, "N"); + const uint32_t top_k_u32 = checked_u32(top_k, "top_k"); + const uint32_t n_block_count = + (n_u32 + kMoeW8A8MarlinNBlock - 1) / kMoeW8A8MarlinNBlock; + const uint32_t max_m_block_count = + checked_u32(expert_ids.numel(), "max_m_block_count"); + + MoeW8A8MarlinMode1001Args args{ + n_block_count, + max_m_block_count, + output.data_ptr(), + weight.data_ptr(), + input.data_ptr(), + weight_scale.data_ptr(), + input_scale.data_ptr(), + nullptr, + sorted_token_ids.data_ptr(), + expert_ids.data_ptr(), + num_tokens_post_padded.data_ptr(), + checked_u32(num_experts, "num_experts"), + checked_u32(m, "M"), + n_u32, + checked_u32(k, "K"), + 1, + 1, + n_u32, + 1, + 1, + checked_u32(sorted_token_ids.numel(), "max_tokens_padded"), + top_k_u32, + 1.0f / static_cast(top_k_u32), + 1.0f, + 0, + 0, + 0}; + + size_t args_size = sizeof(args); + void *launch_config[] = { + HIP_LAUNCH_PARAM_BUFFER_POINTER, + &args, + HIP_LAUNCH_PARAM_BUFFER_SIZE, + &args_size, + HIP_LAUNCH_PARAM_END}; + + auto kernel = get_moe_w8a8_marlin_mode1001_kernel(); + auto status = hipModuleLaunchKernel( + kernel.function, + n_block_count, + 1, + max_m_block_count, + kMoeW8A8MarlinWorkgroupSize, + 1, + 1, + 0, + infinicore::adaptor::get_hip_stream().stream(), + nullptr, + launch_config); + if (status != hipSuccess) { + throw std::runtime_error( + hip_error_message("hipModuleLaunchKernel(W8A8 Marlin mode 1001)", status)); + } +} + +class LmslimQuantLibrary { +public: + void *symbol(const char *name) { + std::lock_guard lock(mutex_); + if (!ensure_open_locked(true)) { + throw std::runtime_error(error_); + } + + dlerror(); + void *fn = dlsym(handle_, name); + const char *err = dlerror(); + if (err != nullptr || fn == nullptr) { + std::ostringstream oss; + oss << "failed to resolve lmslimquant symbol " << name; + if (err != nullptr) { + oss << ": " << err; + } + throw std::runtime_error(oss.str()); + } + return fn; + } + +private: + bool ensure_open_locked(bool update_error) { + if (handle_ != nullptr) { + return true; + } + + const char *path_env = std::getenv("INFINICORE_LMSLIMQUANT_SO"); + const char *path = (path_env != nullptr && path_env[0] != '\0') ? path_env : kDefaultLmslimQuantSo; + handle_ = dlopen(path, RTLD_LAZY | RTLD_GLOBAL); + if (handle_ != nullptr) { + error_.clear(); + return true; + } + + if (update_error || error_.empty()) { + const char *err = dlerror(); + std::ostringstream oss; + oss << "failed to load lmslimquant shared library " << path; + if (err != nullptr) { + oss << ": " << err; + } + error_ = oss.str(); + } + return false; + } + + std::mutex mutex_; + void *handle_ = nullptr; + std::string error_; +}; + +LmslimQuantLibrary &lmslimquant_library() { + static LmslimQuantLibrary lib; + return lib; +} + +void ensure_default_lightop_env() { + if (std::getenv("LIGHTOP_GPU_TARGET") == nullptr) { + setenv("LIGHTOP_GPU_TARGET", kDefaultLightopGpuTarget, 0); + } + if (std::getenv("LIGHTOP_ASM_DIR") == nullptr) { + setenv("LIGHTOP_ASM_DIR", kDefaultLightopAsmDir, 0); + } +} + template Fn resolve(const char *symbol) { return reinterpret_cast(library().symbol(symbol)); } -using FusedRmsNormFn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, double); +template +Fn resolve_lmslimquant(const char *symbol) { + return reinterpret_cast(lmslimquant_library().symbol(symbol)); +} + using FuseSiluAndMulFn = void (*)(at::Tensor &, at::Tensor &); +using RmsRotaryEmbeddingFuseFn = void (*)( + at::Tensor &, + at::Tensor &, + at::Tensor &, + long, + at::Tensor &, + bool, + at::Tensor, + at::Tensor, + std::optional, + std::optional, + double); using MoeSumFn = void (*)( at::Tensor &, at::Tensor &, @@ -107,14 +443,18 @@ using MoeSumFn = void (*)( const std::optional &, float, int); -using MoeFusedGateFn = std::vector (*)( - at::Tensor &, - at::Tensor &, - int64_t, - int64_t, +using MoeAlignBlockSizeFn = void (*)( + at::Tensor, int64_t, int64_t, - double); + at::Tensor, + at::Tensor, + at::Tensor, + const std::optional &, + const std::optional &, + const std::optional &, + bool, + bool); using MoeGemmW16A16Fn = at::Tensor (*)( at::Tensor, at::Tensor, @@ -138,23 +478,66 @@ using MoeMarlinW16A16AsmFn = at::Tensor (*)( int, int); -FusedRmsNormFn fused_rms_norm_fn() { - static auto fn = resolve(kFusedRmsNormSymbol); - return fn; -} +using MoeGemmW8A8Fn = at::Tensor (*)( + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + std::optional, + at::Tensor, + at::Tensor, + at::Tensor, + int64_t, + int, + int); +using MoeMarlinW8A8AsmFn = at::Tensor (*)( + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + std::optional, + at::Tensor, + at::Tensor, + at::Tensor, + unsigned int, + int, + int); +using FuseSiluMulQuantFn = void (*)(at::Tensor &, at::Tensor &, at::Tensor &, std::optional &, int, int, std::optional &); +using PerTokenDynamicQuantInt8Fn = void (*)(at::Tensor &, const at::Tensor &, at::Tensor &, const at::Tensor &); +using BlasltW8A8GemmFn = void (*)( + const at::Tensor &, + const at::Tensor &, + const at::Tensor &, + const at::Tensor &, + at::Tensor &, + long, + long, + long, + long, + const std::string &, + const at::Tensor &, + const at::Tensor &, + const std::optional &); FuseSiluAndMulFn fuse_silu_and_mul_fn() { static auto fn = resolve(kFuseSiluAndMulSymbol); return fn; } +RmsRotaryEmbeddingFuseFn rms_rotary_embedding_fuse_fn() { + static auto fn = resolve(kRmsRotaryEmbeddingFuseSymbol); + return fn; +} + MoeSumFn moe_sum_fn() { static auto fn = resolve(kMoeSumSymbol); return fn; } -MoeFusedGateFn moe_fused_gate_fn() { - static auto fn = resolve(kMoeFusedGateSymbol); +MoeAlignBlockSizeFn moe_align_block_size_fn() { + static auto fn = resolve(kMoeAlignBlockSizeSymbol); return fn; } @@ -168,45 +551,106 @@ MoeMarlinW16A16AsmFn moe_marlin_w16a16_asm_fn() { return fn; } +MoeGemmW8A8Fn moe_gemm_w8a8_fn() { + static auto fn = resolve(kMoeGemmW8A8Symbol); + return fn; +} + +MoeMarlinW8A8AsmFn moe_marlin_w8a8_asm_fn() { + static auto fn = resolve(kMoeMarlinW8A8AsmSymbol); + return fn; +} + +FuseSiluMulQuantFn fuse_silu_mul_quant_fn() { + static auto fn = resolve(kFuseSiluMulQuantSymbol); + return fn; +} + +PerTokenDynamicQuantInt8Fn per_token_dynamic_quant_int8_fn() { + static auto fn = resolve(kPerTokenDynamicQuantInt8Symbol); + return fn; +} + +BlasltW8A8GemmFn blaslt_w8a8_bf16_fn() { + static auto fn = resolve_lmslimquant(kBlasltW8A8Bf16Symbol); + return fn; +} + +BlasltW8A8GemmFn blaslt_w8a8_fp16_fn() { + static auto fn = resolve_lmslimquant(kBlasltW8A8Fp16Symbol); + return fn; +} + } // namespace bool available() { return library().available(); } -bool enabled_by_env() { - const char *value = std::getenv("INFINICORE_ENABLE_HYGON_LIGHTOP"); - if (value == nullptr) { - return false; - } - std::string normalized(value); - for (auto &ch : normalized) { - ch = static_cast(std::tolower(static_cast(ch))); - } - return normalized == "1" || normalized == "true" || normalized == "on" || normalized == "yes"; -} - -void preload_basic_ops() { - (void)fused_rms_norm_fn(); - (void)fuse_silu_and_mul_fn(); +void preload_moe_w16a16_ops() { (void)moe_sum_fn(); - (void)moe_fused_gate_fn(); (void)moe_gemm_w16a16_fn(); (void)moe_marlin_w16a16_asm_fn(); } +void preload_moe_w8a8_ops() { + (void)moe_sum_fn(); + (void)moe_gemm_w8a8_fn(); + (void)moe_marlin_w8a8_asm_fn(); + (void)fuse_silu_mul_quant_fn(); +} + +void preload_moe_align() { + (void)moe_align_block_size_fn(); +} + void preload_silu_and_mul() { (void)fuse_silu_and_mul_fn(); } -void fused_rms_norm_contiguous(at::Tensor &out, at::Tensor &input, at::Tensor &weight, double epsilon) { - fused_rms_norm_fn()(out, input, weight, epsilon); +void preload_moe_w8a8_marlin_asm() { + (void)get_moe_w8a8_marlin_mode1001_kernel(); +} + +void preload_rms_rotary_embedding() { + (void)rms_rotary_embedding_fuse_fn(); +} + +void preload_w8a8_linear_ops() { + (void)per_token_dynamic_quant_int8_fn(); + (void)blaslt_w8a8_bf16_fn(); + (void)blaslt_w8a8_fp16_fn(); } void fuse_silu_and_mul(at::Tensor &input, at::Tensor &output) { fuse_silu_and_mul_fn()(input, output); } +void rms_rotary_embedding_fuse(at::Tensor &positions, + at::Tensor &query, + at::Tensor &key, + int64_t head_size, + at::Tensor &cos_sin_cache, + bool is_neox, + at::Tensor q_weight, + at::Tensor k_weight, + const std::optional &q_bias, + const std::optional &k_bias, + double epsilon) { + rms_rotary_embedding_fuse_fn()( + positions, + query, + key, + static_cast(head_size), + cos_sin_cache, + is_neox, + q_weight, + k_weight, + q_bias, + k_bias, + epsilon); +} + void moe_sum(at::Tensor &input, at::Tensor &output, const std::optional &bias, @@ -217,14 +661,30 @@ void moe_sum(at::Tensor &input, moe_sum_fn()(input, output, bias, expert_mask, local_num_tokens, factor, expect_m); } -std::vector moe_fused_gate(at::Tensor &input, - at::Tensor &bias, - int64_t num_expert_group, - int64_t topk_group, - int64_t topk, - int64_t num_fused_shared_experts, - double routed_scaling_factor) { - return moe_fused_gate_fn()(input, bias, num_expert_group, topk_group, topk, num_fused_shared_experts, routed_scaling_factor); +void moe_align_block_size( + at::Tensor topk_ids, + int64_t num_experts, + int64_t block_size, + at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded, + const std::optional &expert_map, + const std::optional &expert_mask, + const std::optional &num_local_tokens, + bool is_ep, + bool fuse_fill) { + moe_align_block_size_fn()( + topk_ids, + static_cast(num_experts), + static_cast(block_size), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + expert_map, + expert_mask, + num_local_tokens, + is_ep, + fuse_fill); } void moe_gemm_marlin_w16a16(at::Tensor input, @@ -250,6 +710,104 @@ void moe_gemm_marlin_w16a16(at::Tensor input, } } +void moe_gemm_marlin_w8a8(at::Tensor input, + at::Tensor b_qweight, + at::Tensor output, + at::Tensor a_scale, + at::Tensor b_scale, + const std::optional &topk_weights, + at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded, + int64_t top_k, + int mode, + int delta) { + if (mode < 1000) { + moe_gemm_w8a8_fn()( + input, b_qweight, output, a_scale, b_scale, topk_weights, + sorted_token_ids, expert_ids, num_tokens_post_padded, + top_k, mode, delta); + } else { + if (can_launch_moe_w8a8_marlin_mode1001( + input, b_qweight, output, a_scale, b_scale, topk_weights, + sorted_token_ids, expert_ids, num_tokens_post_padded, + top_k, mode, delta)) { + launch_moe_w8a8_marlin_mode1001( + input, b_qweight, output, a_scale, b_scale, + sorted_token_ids, expert_ids, num_tokens_post_padded, top_k); + return; + } + moe_marlin_w8a8_asm_fn()( + input, b_qweight, output, a_scale, b_scale, topk_weights, + sorted_token_ids, expert_ids, num_tokens_post_padded, + static_cast(top_k), mode, delta); + } +} + +void fuse_silu_mul_quant(at::Tensor &input, + at::Tensor &output, + at::Tensor &scales, + std::optional &num_local_tokens, + int topk, + int expect_m, + std::optional &expert_ids) { + fuse_silu_mul_quant_fn()( + input, + output, + scales, + num_local_tokens, + topk, + expect_m, + expert_ids); +} + +void per_token_dynamic_quant_int8(at::Tensor &output, + const at::Tensor &input, + at::Tensor &scales, + const at::Tensor &smooth) { + per_token_dynamic_quant_int8_fn()(output, input, scales, smooth); +} + +void blaslt_w8a8_gemm(at::Tensor &output, + const at::Tensor &a, + const at::Tensor &b, + const at::Tensor &scale_a, + const at::Tensor &scale_b, + const std::optional &bias) { + if (a.dim() != 2 || b.dim() != 2 || output.dim() != 2) { + throw std::runtime_error("lmslimquant W8A8 GEMM expects 2D tensors"); + } + if (!a.is_contiguous() || !b.is_contiguous() || !output.is_contiguous()) { + throw std::runtime_error("lmslimquant W8A8 GEMM expects contiguous a, b, and output"); + } + if (a.scalar_type() != at::kChar || b.scalar_type() != at::kChar || + scale_a.scalar_type() != at::kFloat || scale_b.scalar_type() != at::kFloat) { + throw std::runtime_error("lmslimquant W8A8 GEMM expects int8 inputs and float32 scales"); + } + + const long m = static_cast(output.size(0)); + const long n = static_cast(output.size(1)); + const long k = static_cast(a.size(1)); + if (a.size(0) != m || b.size(0) != n || b.size(1) != k) { + throw std::runtime_error("lmslimquant W8A8 GEMM shape mismatch"); + } + + static const at::Tensor alpha = at::tensor(1, at::TensorOptions().dtype(at::kInt)); + static const at::Tensor beta = at::tensor(0, at::TensorOptions().dtype(at::kInt)); + static const std::string transpose = "TN"; + constexpr long batch = 1; + + if (output.scalar_type() == at::kBFloat16) { + blaslt_w8a8_bf16_fn()(b, a, scale_b, scale_a, output, m, n, k, batch, transpose, alpha, beta, bias); + return; + } + if (output.scalar_type() == at::kHalf) { + blaslt_w8a8_fp16_fn()(b, a, scale_b, scale_a, output, m, n, k, batch, transpose, alpha, beta, bias); + return; + } + throw std::runtime_error("lmslimquant W8A8 GEMM only supports FP16/BF16 output"); +} + } // namespace infinicore::adaptor::lightop #endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.cc b/src/infinicore/context/allocators/pinnable_block_allocator.cc index 32e5c5e9b..07350b201 100644 --- a/src/infinicore/context/allocators/pinnable_block_allocator.cc +++ b/src/infinicore/context/allocators/pinnable_block_allocator.cc @@ -29,9 +29,6 @@ PinnableBlockAllocator::PinnableBlockAllocator(Device device) {8 * 1024 * 1024, {}}, // 8 MB {16 * 1024 * 1024, {}}, // 16 MB {32 * 1024 * 1024, {}}, // 32 MB - {64 * 1024 * 1024, {}}, // 64 MB - {128 * 1024 * 1024, {}}, // 128 MB - {256 * 1024 * 1024, {}}, // 256 MB }; } @@ -83,8 +80,13 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) { // 2. Large block allocation // Try to reuse a frozen or free large block - auto it = std::find_if(large_blocks_.begin(), large_blocks_.end(), - [size](const std::shared_ptr &b) { return b->size >= size && !b->in_use; }); + auto it = large_blocks_.end(); + for (auto candidate = large_blocks_.begin(); candidate != large_blocks_.end(); ++candidate) { + if (!(*candidate)->in_use && (*candidate)->size >= size + && (it == large_blocks_.end() || (*candidate)->size < (*it)->size)) { + it = candidate; + } + } if (it != large_blocks_.end()) { block = *it; @@ -94,6 +96,20 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) { return reinterpret_cast(block->ptr); } + // A growing sequence of exact-size allocations should not leave every + // smaller eager-only block cached indefinitely. Graph-frozen and live + // blocks must stay resident because captured graphs may still reference + // their addresses. + for (auto stale = large_blocks_.begin(); stale != large_blocks_.end();) { + if (!(*stale)->in_use && !(*stale)->frozen && (*stale)->size < size) { + INFINICORE_CHECK_ERROR(infinirtFree((*stale)->ptr)); + all_blocks_.erase((*stale)->ptr); + stale = large_blocks_.erase(stale); + } else { + ++stale; + } + } + // Allocate new large block block = std::make_shared(); block->size = size; diff --git a/src/infinicore/nn/rope.cc b/src/infinicore/nn/rope.cc index 115e162e2..f971f0cbf 100644 --- a/src/infinicore/nn/rope.cc +++ b/src/infinicore/nn/rope.cc @@ -57,6 +57,7 @@ void RoPE::initialize_cache() { // Create sin and cos cache tables: [max_seq_len, cache_dim] INFINICORE_NN_BUFFER_INIT(sin_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); INFINICORE_NN_BUFFER_INIT(cos_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); + INFINICORE_NN_BUFFER_INIT(cos_sin_cache, ({max_seq_len_, rotary_dim_}, dtype_, device_)); // Pre-compute sin and cos values // Frequency generation always uses GPT-J style (theta^(-2j/rotary_dim)). @@ -68,6 +69,7 @@ void RoPE::initialize_cache() { // Allocate CPU buffers std::vector sin_data(max_seq_len_ * cache_dim); std::vector cos_data(max_seq_len_ * cache_dim); + std::vector cos_sin_data(max_seq_len_ * rotary_dim_); for (size_t pos = 0; pos < max_seq_len_; pos++) { for (size_t dim_idx = 0; dim_idx < cache_dim; dim_idx++) { @@ -84,6 +86,8 @@ void RoPE::initialize_cache() { sin_data[pos * cache_dim + dim_idx] = std::sin(angle) * mag_scale; cos_data[pos * cache_dim + dim_idx] = std::cos(angle) * mag_scale; + cos_sin_data[pos * rotary_dim_ + dim_idx] = cos_data[pos * cache_dim + dim_idx]; + cos_sin_data[pos * rotary_dim_ + cache_dim + dim_idx] = sin_data[pos * cache_dim + dim_idx]; } } @@ -93,40 +97,54 @@ void RoPE::initialize_cache() { // Direct use of F32 data auto sin_f32_cpu = Tensor::from_blob(sin_data.data(), {max_seq_len_, cache_dim}, DataType::F32, cpu_device); auto cos_f32_cpu = Tensor::from_blob(cos_data.data(), {max_seq_len_, cache_dim}, DataType::F32, cpu_device); + auto cos_sin_f32_cpu = Tensor::from_blob(cos_sin_data.data(), {max_seq_len_, rotary_dim_}, DataType::F32, cpu_device); sin_cache_->copy_from(sin_f32_cpu); cos_cache_->copy_from(cos_f32_cpu); + cos_sin_cache_->copy_from(cos_sin_f32_cpu); } else if (dtype_ == DataType::BF16) { // Convert F32 to BF16 using the same conversion as Python's ml_dtypes.bfloat16 // This uses round-to-nearest-even (matching _f32_to_bf16 implementation) std::vector sin_bf16_data(max_seq_len_ * cache_dim); std::vector cos_bf16_data(max_seq_len_ * cache_dim); + std::vector cos_sin_bf16_data(max_seq_len_ * rotary_dim_); for (size_t i = 0; i < sin_data.size(); i++) { sin_bf16_data[i] = utils::cast(sin_data[i]); cos_bf16_data[i] = utils::cast(cos_data[i]); } + for (size_t i = 0; i < cos_sin_data.size(); i++) { + cos_sin_bf16_data[i] = utils::cast(cos_sin_data[i]); + } auto sin_bf16_cpu = Tensor::from_blob(sin_bf16_data.data(), {max_seq_len_, cache_dim}, DataType::BF16, cpu_device); auto cos_bf16_cpu = Tensor::from_blob(cos_bf16_data.data(), {max_seq_len_, cache_dim}, DataType::BF16, cpu_device); + auto cos_sin_bf16_cpu = Tensor::from_blob(cos_sin_bf16_data.data(), {max_seq_len_, rotary_dim_}, DataType::BF16, cpu_device); // copy_from handles cross-device copying to target device sin_cache_->copy_from(sin_bf16_cpu); cos_cache_->copy_from(cos_bf16_cpu); + cos_sin_cache_->copy_from(cos_sin_bf16_cpu); } else if (dtype_ == DataType::F16) { // Convert F32 to F16 std::vector sin_f16_data(max_seq_len_ * cache_dim); std::vector cos_f16_data(max_seq_len_ * cache_dim); + std::vector cos_sin_f16_data(max_seq_len_ * rotary_dim_); for (size_t i = 0; i < sin_data.size(); i++) { sin_f16_data[i] = utils::cast(sin_data[i]); cos_f16_data[i] = utils::cast(cos_data[i]); } + for (size_t i = 0; i < cos_sin_data.size(); i++) { + cos_sin_f16_data[i] = utils::cast(cos_sin_data[i]); + } auto sin_f16_cpu = Tensor::from_blob(sin_f16_data.data(), {max_seq_len_, cache_dim}, DataType::F16, cpu_device); auto cos_f16_cpu = Tensor::from_blob(cos_f16_data.data(), {max_seq_len_, cache_dim}, DataType::F16, cpu_device); + auto cos_sin_f16_cpu = Tensor::from_blob(cos_sin_f16_data.data(), {max_seq_len_, rotary_dim_}, DataType::F16, cpu_device); sin_cache_->copy_from(sin_f16_cpu); cos_cache_->copy_from(cos_f16_cpu); + cos_sin_cache_->copy_from(cos_sin_f16_cpu); } else { throw std::runtime_error( "RoPE cache dtype conversion not yet supported for dtype: " diff --git a/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc b/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc index 160147436..095f06e74 100644 --- a/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc +++ b/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc @@ -4,25 +4,10 @@ #include "../../../adaptor/flash_attn/hygon/flash_attn_hygon.hpp" #include "infinicore/adaptor/aten_adaptor.hpp" -#include -#include -#include -#include - -#include #include -#include namespace infinicore::op::mha_kvcache_impl::flashattn { -namespace { -bool isCurrentHipStreamCapturing() { - hipStreamCaptureStatus capture_status = hipStreamCaptureStatusNone; - hipError_t status = hipStreamIsCapturing(infinicore::adaptor::get_hip_stream().stream(), &capture_status); - return status == hipSuccess && capture_status != hipStreamCaptureStatusNone; -} -} // namespace - struct PlannedMeta { graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; std::optional alibi_slopes; @@ -65,70 +50,25 @@ void run(void *planned_meta) { ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; - if (std::getenv("INFINICORE_HYGON_ATEN_FALLBACK")) { - namespace idx = at::indexing; - auto seqlens_t = infinicore::adaptor::to_aten_tensor(p->seqlens_k); - auto block_table_t = infinicore::adaptor::to_aten_tensor(p->block_table); - auto seqlens_cpu = seqlens_t.to(at::kCPU); - auto block_table_cpu = block_table_t.to(at::kCPU); - - auto result = at::empty_like(out_tensor); - const int64_t batch_size = q.size(0); - const int64_t seqlen_q = q.size(1); - const int64_t num_heads = q.size(2); - const int64_t block_size = k_cache.size(1); - const int64_t num_kv_heads = k_cache.size(2); - const int64_t group_size = num_heads / num_kv_heads; - - for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - const int64_t seq_len = seqlens_cpu.index({batch_idx}).item(); - std::vector keys; - std::vector values; - keys.reserve(seq_len); - values.reserve(seq_len); - for (int64_t logical_pos = 0; logical_pos < seq_len; ++logical_pos) { - const int64_t block_id = block_table_cpu.index({batch_idx, logical_pos / block_size}).item(); - const int64_t off = logical_pos % block_size; - keys.push_back(k_cache.index({block_id, off, idx::Slice(), idx::Slice()})); - values.push_back(v_cache.index({block_id, off, idx::Slice(), idx::Slice()})); - } - auto K = at::stack(keys, 0); - auto V = at::stack(values, 0); - if (group_size > 1) { - K = K.repeat_interleave(group_size, 1); - V = V.repeat_interleave(group_size, 1); - } - auto cur_q = q.index({batch_idx}); - auto scores = at::matmul(cur_q.permute({1, 0, 2}).to(at::kFloat), K.permute({1, 2, 0}).to(at::kFloat)) * p->scale; - auto mask = at::full({seqlen_q, seq_len}, -std::numeric_limits::infinity(), q.options().dtype(at::kFloat)); - const int64_t prefix_len = seq_len - seqlen_q; - for (int64_t query_pos = 0; query_pos < seqlen_q; ++query_pos) { - mask.index_put_({query_pos, idx::Slice(0, prefix_len + query_pos + 1)}, 0.0); - } - auto attn = at::softmax(scores + mask.unsqueeze(0), -1).to(q.dtype()); - auto cur_out = at::matmul(attn, V.permute({1, 0, 2})).permute({1, 0, 2}); - result.index_put_({batch_idx}, cur_out); - } - - out_tensor.copy_(result); - if (out_need_copy_back) { - p->out->copy_from(out_work); - } - return; - } - std::optional k_new = std::nullopt; std::optional v_new = std::nullopt; std::optional rotary_cos = std::nullopt; std::optional rotary_sin = std::nullopt; std::optional cache_batch_idx = std::nullopt; std::optional leftpad_k = std::nullopt; - const bool use_dynamic_out = q.dim() == 4 && k_cache.dim() == 4 - && q.size(1) == 1 && q.size(2) > k_cache.size(2) - && q.size(3) % 8 == 0 && !alibi_slopes.has_value(); - - auto out = use_dynamic_out ? std::optional(std::nullopt) - : std::optional(out_tensor); + const bool needs_grouped_out_alias = q.dim() == 4 + && k_cache.dim() == 4 && v_cache.dim() == 4 + && q.size(1) == 1 && k_cache.size(2) > 0 + && q.size(2) > k_cache.size(2) + && q.size(2) % k_cache.size(2) == 0 + && q.size(3) == v_cache.size(3) + && q.size(3) % 8 == 0 + && v_cache.sizes() == k_cache.sizes() + && !alibi_slopes.has_value(); + auto direct_out = needs_grouped_out_alias + ? out_tensor.view({q.size(0), q.size(2) / k_cache.size(2), k_cache.size(2), v_cache.size(3)}) + : out_tensor; + auto out = std::optional(direct_out); auto result = flash::mha_fwd_kvcache( q, @@ -152,10 +92,8 @@ void run(void *planned_meta) { false, 0); - if (!isCurrentHipStreamCapturing()) { - c10::hip::device_synchronize(); - } - if (!result.empty() && result[0].defined()) { + if (!result.empty() && result[0].defined() + && result[0].data_ptr() != out_tensor.data_ptr()) { out_tensor.copy_(result[0]); } if (out_need_copy_back) { diff --git a/src/infinicore/ops/moe_align/hygon/moe_align_lightop_hygon.cc b/src/infinicore/ops/moe_align/hygon/moe_align_lightop_hygon.cc new file mode 100644 index 000000000..6748cdd93 --- /dev/null +++ b/src/infinicore/ops/moe_align/hygon/moe_align_lightop_hygon.cc @@ -0,0 +1,129 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/ops/moe_align.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/lightop_adaptor.hpp" + +#include +#include + +#include + +namespace infinicore::op::moe_align_impl::infiniop { +void *plan(Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + size_t num_experts, + size_t block_size, + bool pad_sorted_token_ids); +void run(void *planned_meta); +void cleanup(void **planned_meta_ptr); +} // namespace infinicore::op::moe_align_impl::infiniop +namespace infinicore::op::moe_align_impl::lightop_hygon { + +struct PlannedMeta { + graph::GraphTensor sorted_token_ids; + graph::GraphTensor expert_ids; + graph::GraphTensor num_tokens_post_padded; + graph::GraphTensor topk_ids; + size_t num_experts; + size_t block_size; + bool pad_sorted_token_ids; + bool use_lightop; + void *fallback; +}; + +void *plan(Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + const size_t num_experts, + const size_t block_size, + const bool pad_sorted_token_ids) { + const auto shape = topk_ids->shape(); + const bool use_lightop = + shape.size() == 2 && shape[0] == 1 && shape[1] == 8 && num_experts == 128; + if (use_lightop) { + infinicore::adaptor::lightop::preload_moe_align(); + } + + void *fallback = nullptr; + if (!use_lightop) { + fallback = infinicore::op::moe_align_impl::infiniop::plan( + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + num_experts, + block_size, + pad_sorted_token_ids); + } + + return new PlannedMeta{ + graph::GraphTensor(sorted_token_ids), + graph::GraphTensor(expert_ids), + graph::GraphTensor(num_tokens_post_padded), + graph::GraphTensor(topk_ids), + num_experts, + block_size, + pad_sorted_token_ids, + use_lightop, + fallback}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + if (!p->use_lightop) { + infinicore::op::moe_align_impl::infiniop::run(p->fallback); + return; + } + + auto topk_ids = infinicore::adaptor::to_aten_tensor(p->topk_ids); + auto sorted_token_ids = infinicore::adaptor::to_aten_tensor(p->sorted_token_ids); + auto expert_ids = infinicore::adaptor::to_aten_tensor(p->expert_ids); + auto num_tokens_post_padded = infinicore::adaptor::to_aten_tensor(p->num_tokens_post_padded); + + if (p->pad_sorted_token_ids) { + sorted_token_ids.fill_(topk_ids.numel()); + } + + const std::optional none = std::nullopt; + infinicore::adaptor::lightop::moe_align_block_size( + topk_ids, + static_cast(p->num_experts), + static_cast(p->block_size), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + none, + none, + none, + false, + false); +} + +void cleanup(void **planned_meta_ptr) { + auto *p = *reinterpret_cast(planned_meta_ptr); + if (p->fallback != nullptr) { + infinicore::op::moe_align_impl::infiniop::cleanup(&p->fallback); + } + delete p; + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + if (!infinicore::adaptor::lightop::available()) { + return false; + } + MoeAlign::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); + MoeAlign::run_dispatcher().registerDevice(Device::Type::HYGON, &run); + MoeAlign::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); + return true; +}(); + +} // namespace infinicore::op::moe_align_impl::lightop_hygon + +#endif // ENABLE_HYGON_API && ENABLE_ATEN \ No newline at end of file diff --git a/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc b/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc index 001e01875..18d46e409 100644 --- a/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc +++ b/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc @@ -10,10 +10,7 @@ #include #include -#include -#include #include -#include #include #include #include @@ -68,40 +65,6 @@ at::Tensor pack_one_expert(const at::Tensor &weight) { return packed; } - -bool debug_enabled() { - const char *value = std::getenv("INFINICORE_DEBUG_HYGON_MARLIN"); - return value != nullptr && value[0] != '\0' && std::string(value) != "0"; -} - -std::string tensor_desc(const at::Tensor &tensor) { - std::ostringstream oss; - oss << "shape=["; - for (int64_t i = 0; i < tensor.dim(); ++i) { - if (i != 0) { - oss << ","; - } - oss << tensor.size(i); - } - oss << "] stride=["; - for (int64_t i = 0; i < tensor.dim(); ++i) { - if (i != 0) { - oss << ","; - } - oss << tensor.stride(i); - } - oss << "] dtype=" << tensor.scalar_type() - << " device=" << tensor.device() - << " contiguous=" << (tensor.is_contiguous() ? "true" : "false"); - return oss.str(); -} - -void debug_tensor(const char *name, const at::Tensor &tensor) { - if (debug_enabled()) { - std::cerr << "[hygon-marlin] " << name << " " << tensor_desc(tensor) << std::endl; - } -} - Tensor pack(const Tensor &weight) { c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); auto weight_at = infinicore::adaptor::to_aten_tensor(weight); @@ -160,7 +123,7 @@ void *plan(Tensor output, int delta0, int mode1, int delta1) { - infinicore::adaptor::lightop::preload_basic_ops(); + infinicore::adaptor::lightop::preload_moe_w16a16_ops(); return new PlannedMeta{ graph::GraphTensor(output), graph::GraphTensor(cache13), graph::GraphTensor(cache2), graph::GraphTensor(hidden_states), graph::GraphTensor(w13_marlin), graph::GraphTensor(w2_marlin), @@ -209,20 +172,6 @@ void run(void *planned_meta) { auto expert_ids = infinicore::adaptor::to_aten_tensor(p->expert_ids); auto num_tokens_post_padded = infinicore::adaptor::to_aten_tensor(p->num_tokens_post_padded); - if (debug_enabled()) { - std::cerr << "[hygon-marlin] top_k=" << top_k << " mode0=" << p->mode0 << " delta0=" << p->delta0 - << " mode1=" << p->mode1 << " delta1=" << p->delta1 << std::endl; - debug_tensor("hidden", hidden); - debug_tensor("w13", w13); - debug_tensor("cache1", cache1); - debug_tensor("cache2", cache2); - debug_tensor("cache3", cache3); - debug_tensor("topk_weights", topk_weights); - debug_tensor("sorted_token_ids", sorted_token_ids); - debug_tensor("expert_ids", expert_ids); - debug_tensor("num_tokens_post_padded", num_tokens_post_padded); - } - try { infinicore::adaptor::lightop::moe_gemm_marlin_w16a16( hidden, w13, cache1, std::nullopt, sorted_token_ids, expert_ids, diff --git a/src/infinicore/ops/moe_w8a8_marlin/hygon/moe_w8a8_marlin_hygon.cc b/src/infinicore/ops/moe_w8a8_marlin/hygon/moe_w8a8_marlin_hygon.cc new file mode 100644 index 000000000..6e18bf0b9 --- /dev/null +++ b/src/infinicore/ops/moe_w8a8_marlin/hygon/moe_w8a8_marlin_hygon.cc @@ -0,0 +1,304 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/ops/moe_w8a8_marlin.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/lightop_adaptor.hpp" +#include "infinicore/ops/common/dispatcher.hpp" +#include "infinicore/ops/moe_sum.hpp" +#include "infinicore/ops/per_channel_quant_i8.hpp" + +#include +#include + +#include +#include +#include +#include + +namespace infinicore::op { +namespace moe_w8a8_marlin_pack_impl { +using schema = Tensor (*)(const Tensor &); +common::OpDispatcher &dispatcher(); +} // namespace moe_w8a8_marlin_pack_impl +} // namespace infinicore::op + +namespace infinicore::op::moe_w8a8_marlin_impl::hygon { + +namespace { + +at::Tensor pack_one_expert(const at::Tensor &weight) { + if (weight.dim() != 2) { + throw std::runtime_error("w8a8 marlin pack expects each expert weight to be 2D"); + } + if (weight.scalar_type() != at::kChar) { + throw std::runtime_error("w8a8 marlin pack expects int8 weights"); + } + auto q_w = weight.transpose(0, 1).contiguous(); + const int64_t size_k = q_w.size(0); + const int64_t size_n = q_w.size(1); + constexpr int64_t k_tile = 64; + if (size_k % k_tile != 0) { + throw std::runtime_error("w8a8 marlin pack requires K % 64 == 0"); + } + auto packed = q_w.reshape({size_k / k_tile, k_tile, size_n}) + .transpose(1, 2) + .reshape({size_k / k_tile, size_n * k_tile}) + .contiguous(); + return packed; +} + +Tensor pack(const Tensor &weight) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto weight_at = infinicore::adaptor::to_aten_tensor(weight); + if (weight_at.dim() != 3) { + throw std::runtime_error("w8a8 marlin pack expects weight shape [E, N, K]"); + } + const int64_t num_experts = weight_at.size(0); + auto packed0 = pack_one_expert(weight_at.select(0, 0)); + auto output_ic = Tensor::empty( + {static_cast(num_experts), + static_cast(packed0.size(0)), + static_cast(packed0.size(1))}, + weight->dtype(), + weight->device()); + auto output_at = infinicore::adaptor::to_aten_tensor(output_ic); + output_at.select(0, 0).copy_(packed0); + for (int64_t expert = 1; expert < num_experts; ++expert) { + auto packed = pack_one_expert(weight_at.select(0, expert)); + output_at.select(0, expert).copy_(packed); + } + return output_ic; +} + +} // namespace + +struct PlannedMeta { + graph::GraphTensor output; + graph::GraphTensor cache13; + graph::GraphTensor cache2_i8; + graph::GraphTensor input_i8; + graph::GraphTensor input_scale; + graph::GraphTensor cache2_scale; + graph::GraphTensor hidden_states; + graph::GraphTensor w13_marlin; + graph::GraphTensor w2_marlin; + graph::GraphTensor w13_scale; + graph::GraphTensor w2_scale; + graph::GraphTensor topk_weights; + graph::GraphTensor sorted_token_ids; + graph::GraphTensor expert_ids; + graph::GraphTensor num_tokens_post_padded; + size_t top_k; + int mode0; + size_t block_size_m; + int delta0; + int mode1; + int delta1; +}; + +void *plan(Tensor output, + Tensor cache13, + Tensor cache2_i8, + Tensor input_i8, + Tensor input_scale, + Tensor cache2_scale, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &w13_scale, + const Tensor &w2_scale, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + size_t block_size_m, + int delta0, + int mode1, + int delta1) { + infinicore::adaptor::lightop::preload_moe_w8a8_ops(); + if (mode0 == 1001 && delta0 == 1) { + infinicore::adaptor::lightop::preload_moe_w8a8_marlin_asm(); + } + return new PlannedMeta{ + graph::GraphTensor(output), + graph::GraphTensor(cache13), + graph::GraphTensor(cache2_i8), + graph::GraphTensor(input_i8), + graph::GraphTensor(input_scale), + graph::GraphTensor(cache2_scale), + graph::GraphTensor(hidden_states), + graph::GraphTensor(w13_marlin), + graph::GraphTensor(w2_marlin), + graph::GraphTensor(w13_scale), + graph::GraphTensor(w2_scale), + graph::GraphTensor(topk_weights), + graph::GraphTensor(sorted_token_ids), + graph::GraphTensor(expert_ids), + graph::GraphTensor(num_tokens_post_padded), + top_k, + mode0, + block_size_m, + delta0, + mode1, + delta1}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + const auto hidden_shape = p->hidden_states->shape(); + const auto w13_shape = p->w13_marlin->shape(); + const auto w2_shape = p->w2_marlin->shape(); + if (hidden_shape.size() != 2 || w13_shape.size() != 3 || w2_shape.size() != 3) { + throw std::runtime_error("w8a8 marlin fused dense expects hidden [M,K], w13/w2 [E,*,*]"); + } + const size_t m = hidden_shape[0]; + const size_t k = hidden_shape[1]; + const size_t top_k = p->top_k; + const size_t n = w2_shape[1] * 64; + const size_t n2 = n * 2; + if (w13_shape[1] * 64 != k || + w13_shape[2] != n2 * 64 || + w2_shape[2] != k * 64) { + throw std::runtime_error("w8a8 marlin fused dense weight shape mismatch"); + } + const infinicore::Shape input_i8_shape{m, k}; + const infinicore::Shape input_scale_shape{m, 1}; + const infinicore::Shape cache2_i8_shape{m * top_k, n}; + const infinicore::Shape cache2_scale_shape{m * top_k, 1}; + if (p->input_i8->shape() != input_i8_shape || + p->input_scale->shape() != input_scale_shape || + p->cache2_i8->shape() != cache2_i8_shape || + p->cache2_scale->shape() != cache2_scale_shape) { + throw std::runtime_error("w8a8 marlin fused dense workspace shape mismatch"); + } + if (p->input_i8->dtype() != infinicore::DataType::I8 || + p->cache2_i8->dtype() != infinicore::DataType::I8 || + p->input_scale->dtype() != infinicore::DataType::F32 || + p->cache2_scale->dtype() != infinicore::DataType::F32) { + throw std::runtime_error("w8a8 marlin fused dense workspace dtype mismatch"); + } + + const bool output_need_copy_back = !p->output->is_contiguous(); + Tensor output_work_ic = output_need_copy_back ? p->output->contiguous() : Tensor(p->output); + Tensor hidden_work_ic = p->hidden_states->is_contiguous() ? Tensor(p->hidden_states) : p->hidden_states->contiguous(); + + const size_t cache1_numel = m * top_k * n2; + const size_t cache3_numel = m * top_k * k; + auto cache1_ic = p->cache13->narrow({{0, 0, cache1_numel}})->view({m, top_k, n2}); + auto cache1_2d_ic = cache1_ic->view({m * top_k, n2}); + auto cache3_ic = p->cache13->narrow({{0, 0, cache3_numel}})->view({m, top_k, k}); + + infinicore::op::per_channel_quant_i8_( + hidden_work_ic->view({m, k}), + Tensor(p->input_i8), + Tensor(p->input_scale)); + + auto qhidden = infinicore::adaptor::to_aten_tensor(p->input_i8); + auto hidden_scale = infinicore::adaptor::to_aten_tensor(p->input_scale); + auto w13 = infinicore::adaptor::to_aten_tensor(p->w13_marlin); + auto w2 = infinicore::adaptor::to_aten_tensor(p->w2_marlin); + auto w13_scale = infinicore::adaptor::to_aten_tensor(p->w13_scale); + auto w2_scale = infinicore::adaptor::to_aten_tensor(p->w2_scale); + auto cache1 = infinicore::adaptor::to_aten_tensor(cache1_ic); + auto cache1_2d = infinicore::adaptor::to_aten_tensor(cache1_2d_ic); + auto qcache2 = infinicore::adaptor::to_aten_tensor(p->cache2_i8); + auto cache2_scale = infinicore::adaptor::to_aten_tensor(p->cache2_scale); + auto cache3 = infinicore::adaptor::to_aten_tensor(cache3_ic); + auto topk_weights = infinicore::adaptor::to_aten_tensor(p->topk_weights); + auto sorted_token_ids = infinicore::adaptor::to_aten_tensor(p->sorted_token_ids); + auto expert_ids = infinicore::adaptor::to_aten_tensor(p->expert_ids); + auto num_tokens_post_padded = infinicore::adaptor::to_aten_tensor(p->num_tokens_post_padded); + if (p->block_size_m == 0) { + throw std::runtime_error("w8a8 marlin fused dense requires nonzero block_size_m"); + } + + const size_t num_pairs = m * top_k; + const size_t num_experts = w13_shape[0]; + const size_t vllm_max_tokens_padded = + num_pairs < num_experts + ? std::min(num_pairs * p->block_size_m, + num_pairs + num_experts * (p->block_size_m - 1)) + : num_pairs + num_experts * (p->block_size_m - 1); + const size_t vllm_max_blocks = + (vllm_max_tokens_padded + p->block_size_m - 1) / p->block_size_m; + auto sorted_token_ids_lightop = + vllm_max_tokens_padded < static_cast(sorted_token_ids.size(0)) + ? sorted_token_ids.narrow(0, 0, static_cast(vllm_max_tokens_padded)) + : sorted_token_ids; + auto expert_ids_lightop = + vllm_max_blocks < static_cast(expert_ids.size(0)) + ? expert_ids.narrow(0, 0, static_cast(vllm_max_blocks)) + : expert_ids; + + try { + infinicore::adaptor::lightop::moe_gemm_marlin_w8a8( + qhidden, w13, cache1, hidden_scale, w13_scale, std::nullopt, + sorted_token_ids_lightop, expert_ids_lightop, num_tokens_post_padded, + static_cast(top_k), p->mode0, p->delta0); + } catch (const std::exception &e) { + throw std::runtime_error(std::string("Hygon W8A8 Marlin GEMM1 failed: ") + e.what()); + } + + std::optional num_local_tokens = std::nullopt; + std::optional silu_expert_ids = std::nullopt; + try { + infinicore::adaptor::lightop::fuse_silu_mul_quant( + cache1_2d, + qcache2, + cache2_scale, + num_local_tokens, + 1, + -1, + silu_expert_ids); + } catch (const std::exception &e) { + throw std::runtime_error(std::string("Hygon W8A8 Marlin fuse_silu_mul_quant failed: ") + e.what()); + } + + std::optional topk_weights_opt(topk_weights); + try { + infinicore::adaptor::lightop::moe_gemm_marlin_w8a8( + qcache2, w2, cache3, cache2_scale, w2_scale, topk_weights_opt, + sorted_token_ids_lightop, expert_ids_lightop, num_tokens_post_padded, + 1, p->mode1, p->delta1); + } catch (const std::exception &e) { + throw std::runtime_error(std::string("Hygon W8A8 Marlin GEMM2 failed: ") + e.what()); + } + + auto output_work = infinicore::adaptor::to_aten_tensor(output_work_ic); + infinicore::adaptor::lightop::moe_sum( + cache3, + output_work, + std::nullopt, + std::nullopt, + std::nullopt, + 1.0f, + -1); + + if (output_need_copy_back) { + p->output->copy_from(output_work_ic); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + if (!infinicore::adaptor::lightop::available()) { + return false; + } + infinicore::op::moe_w8a8_marlin_pack_impl::dispatcher().registerDevice(Device::Type::HYGON, &pack); + MoeW8A8MarlinFusedDense::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); + MoeW8A8MarlinFusedDense::run_dispatcher().registerDevice(Device::Type::HYGON, &run); + MoeW8A8MarlinFusedDense::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); + return true; +}(); + +} // namespace infinicore::op::moe_w8a8_marlin_impl::hygon + +#endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/src/infinicore/ops/moe_w8a8_marlin/moe_w8a8_marlin.cc b/src/infinicore/ops/moe_w8a8_marlin/moe_w8a8_marlin.cc new file mode 100644 index 000000000..7c24be3b3 --- /dev/null +++ b/src/infinicore/ops/moe_w8a8_marlin/moe_w8a8_marlin.cc @@ -0,0 +1,150 @@ +#include "infinicore/ops/moe_w8a8_marlin.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +namespace moe_w8a8_marlin_pack_impl { +using schema = Tensor (*)(const Tensor &); +common::OpDispatcher &dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +} +} // namespace moe_w8a8_marlin_pack_impl + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeW8A8MarlinFusedDense); + +MoeW8A8MarlinFusedDense::MoeW8A8MarlinFusedDense( + Tensor output, + Tensor cache13, + Tensor cache2_i8, + Tensor input_i8, + Tensor input_scale, + Tensor cache2_scale, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &w13_scale, + const Tensor &w2_scale, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + size_t block_size_m, + int delta0, + int mode1, + int delta1) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE( + output, cache13, cache2_i8, input_i8, input_scale, cache2_scale, + hidden_states, w13_marlin, w2_marlin, w13_scale, w2_scale, + topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded); + INFINICORE_GRAPH_OP_DISPATCH( + output->device().getType(), + output, + cache13, + cache2_i8, + input_i8, + input_scale, + cache2_scale, + hidden_states, + w13_marlin, + w2_marlin, + w13_scale, + w2_scale, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + top_k, + mode0, + block_size_m, + delta0, + mode1, + delta1); +} + +void MoeW8A8MarlinFusedDense::execute( + Tensor output, + Tensor cache13, + Tensor cache2_i8, + Tensor input_i8, + Tensor input_scale, + Tensor cache2_scale, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &w13_scale, + const Tensor &w2_scale, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + size_t block_size_m, + int delta0, + int mode1, + int delta1) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN( + MoeW8A8MarlinFusedDense, + output, + cache13, + cache2_i8, + input_i8, + input_scale, + cache2_scale, + hidden_states, + w13_marlin, + w2_marlin, + w13_scale, + w2_scale, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + top_k, + mode0, + block_size_m, + delta0, + mode1, + delta1); +} + + +Tensor moe_w8a8_marlin_pack(const Tensor &weight) { + return moe_w8a8_marlin_pack_impl::dispatcher().lookup(weight->device().getType())(weight); +} + +void moe_w8a8_marlin_fused_dense_( + Tensor output, + Tensor cache13, + Tensor cache2_i8, + Tensor input_i8, + Tensor input_scale, + Tensor cache2_scale, + const Tensor &hidden_states, + const Tensor &w13_marlin, + const Tensor &w2_marlin, + const Tensor &w13_scale, + const Tensor &w2_scale, + const Tensor &topk_weights, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded, + size_t top_k, + int mode0, + size_t block_size_m, + int delta0, + int mode1, + int delta1) { + MoeW8A8MarlinFusedDense::execute( + output, cache13, cache2_i8, input_i8, input_scale, cache2_scale, + hidden_states, w13_marlin, w2_marlin, w13_scale, w2_scale, + topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, + top_k, mode0, block_size_m, delta0, mode1, delta1); +} + + +} // namespace infinicore::op diff --git a/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc b/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc index b9ffe1570..b55046ce7 100644 --- a/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc +++ b/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc @@ -4,15 +4,11 @@ #ifdef ENABLE_ATEN #include "infinicore/adaptor/aten_adaptor.hpp" #include -#include #include #endif #include "../../../adaptor/flash_attn/hygon/flash_attn_hygon.hpp" -#include -#include -#include #include namespace infinicore::op::mha_varlen_impl::flashattn { @@ -113,74 +109,6 @@ void run(void *planned_meta) { return; } - if (std::getenv("INFINICORE_HYGON_ATEN_FALLBACK")) { - namespace idx = at::indexing; - auto cu_q_cpu = cu_seqlens_q.to(at::kCPU); - auto cu_k_cpu = cu_seqlens_kv.to(at::kCPU); - const int64_t num_seqs = cu_q_cpu.size(0) - 1; - auto result = at::zeros_like(out_work); - - if (!p->block_table.has_value()) { - for (int64_t i = 0; i < num_seqs; ++i) { - const int64_t q_start = cu_q_cpu.index({i}).item(); - const int64_t q_end = cu_q_cpu.index({i + 1}).item(); - const int64_t k_start = cu_k_cpu.index({i}).item(); - const int64_t k_end = cu_k_cpu.index({i + 1}).item(); - auto cur_q = q.index({idx::Slice(q_start, q_end)}).unsqueeze(0).transpose(1, 2); - auto cur_k = k.index({idx::Slice(k_start, k_end)}).unsqueeze(0).transpose(1, 2); - auto cur_v = v.index({idx::Slice(k_start, k_end)}).unsqueeze(0).transpose(1, 2); - auto cur_out = at::scaled_dot_product_attention( - cur_q, cur_k, cur_v, std::nullopt, 0.0, true, std::optional(static_cast(p->scale))); - result.index_put_({idx::Slice(q_start, q_end)}, cur_out.transpose(1, 2).squeeze(0)); - } - } else { - auto block_table_t = infinicore::adaptor::to_aten_tensor(*p->block_table); - auto block_table_cpu = block_table_t.to(at::kCPU); - const int64_t block_size = k.size(1); - for (int64_t i = 0; i < num_seqs; ++i) { - const int64_t q_start = cu_q_cpu.index({i}).item(); - const int64_t q_end = cu_q_cpu.index({i + 1}).item(); - const int64_t q_len = q_end - q_start; - const int64_t h_len = (cu_k_cpu.index({i + 1}).item() - cu_k_cpu.index({i}).item()) - q_len; - const int64_t total_len = h_len + q_len; - auto cur_q = q.index({idx::Slice(q_start, q_end)}); - std::vector keys; - std::vector values; - keys.reserve(total_len); - values.reserve(total_len); - for (int64_t j = 0; j < total_len; ++j) { - const int64_t b_id = block_table_cpu.index({i, j / block_size}).item(); - const int64_t off = j % block_size; - keys.push_back(k.index({b_id, off, idx::Slice(), idx::Slice()})); - values.push_back(v.index({b_id, off, idx::Slice(), idx::Slice()})); - } - auto K = at::stack(keys, 0); - auto V = at::stack(values, 0); - const int64_t q_heads = cur_q.size(1); - const int64_t kv_heads = K.size(1); - if (q_heads != kv_heads) { - const int64_t repeat = q_heads / kv_heads; - K = K.repeat_interleave(repeat, 1); - V = V.repeat_interleave(repeat, 1); - } - auto scores = at::matmul(cur_q.permute({1, 0, 2}).to(at::kFloat), K.permute({1, 2, 0}).to(at::kFloat)) * p->scale; - auto mask = at::full({q_len, total_len}, -std::numeric_limits::infinity(), q.options().dtype(at::kFloat)); - for (int64_t t = 0; t < q_len; ++t) { - mask.index_put_({t, idx::Slice(0, h_len + t + 1)}, 0.0); - } - auto attn = at::softmax(scores + mask.unsqueeze(0), -1).to(q.dtype()); - auto cur_out = at::matmul(attn, V.permute({1, 0, 2})).permute({1, 0, 2}); - result.index_put_({idx::Slice(q_start, q_end)}, cur_out); - } - } - - out_work.copy_(result); - if (out_need_copy_back) { - p->out->copy_from(out_work_ic); - } - return; - } - auto out = std::optional(out_work); std::optional seqused_k = std::nullopt; std::optional leftpad_k = std::nullopt; @@ -240,7 +168,8 @@ void run(void *planned_meta) { 0.0, false, std::nullopt); - if (!result.empty() && result[0].defined()) { + if (!result.empty() && result[0].defined() + && result[0].data_ptr() != out_work.data_ptr()) { out_work.copy_(result[0]); } if (out_need_copy_back) { diff --git a/src/infinicore/ops/per_channel_quant_i8/hygon/per_channel_quant_i8_lightop_hygon.cc b/src/infinicore/ops/per_channel_quant_i8/hygon/per_channel_quant_i8_lightop_hygon.cc new file mode 100644 index 000000000..1bb31ec0d --- /dev/null +++ b/src/infinicore/ops/per_channel_quant_i8/hygon/per_channel_quant_i8_lightop_hygon.cc @@ -0,0 +1,92 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/ops/per_channel_quant_i8.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/lightop_adaptor.hpp" + +#include +#include + +#include + +namespace infinicore::op::per_channel_quant_i8_impl::lightop_hygon { + +struct PlannedMeta { + graph::GraphTensor x; + graph::GraphTensor x_packed; + graph::GraphTensor x_scale; + at::Tensor smooth; +}; + +void *plan(const Tensor &x, Tensor x_packed, Tensor x_scale) { + infinicore::adaptor::lightop::preload_w8a8_linear_ops(); + if (x->ndim() != 2 || x_packed->ndim() != 2 || x_scale->ndim() != 2) { + throw std::runtime_error("Hygon per_channel_quant_i8 expects 2D tensors"); + } + const auto m = x->shape()[0]; + const auto k = x->shape()[1]; + if (x_packed->shape() != x->shape() || x_scale->shape() != std::vector{m, 1}) { + throw std::runtime_error("Hygon per_channel_quant_i8 shape mismatch"); + } + if (x_packed->dtype() != DataType::I8 || x_scale->dtype() != DataType::F32) { + throw std::runtime_error("Hygon per_channel_quant_i8 output dtype mismatch"); + } + + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto options = at::TensorOptions() + .dtype(at::kFloat) + .device(infinicore::adaptor::to_at_device(x->device())) + .requires_grad(false); + auto smooth = at::ones({static_cast(k)}, options); + + return new PlannedMeta{ + graph::GraphTensor(x), + graph::GraphTensor(x_packed), + graph::GraphTensor(x_scale), + smooth}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + Tensor x_work = p->x->is_contiguous() ? Tensor(p->x) : p->x->contiguous(); + + const bool packed_need_copy_back = !p->x_packed->is_contiguous(); + const bool scale_need_copy_back = !p->x_scale->is_contiguous(); + Tensor packed_work = packed_need_copy_back ? p->x_packed->contiguous() : Tensor(p->x_packed); + Tensor scale_work = scale_need_copy_back ? p->x_scale->contiguous() : Tensor(p->x_scale); + + auto x = infinicore::adaptor::to_aten_tensor(x_work); + auto x_packed = infinicore::adaptor::to_aten_tensor(packed_work); + auto x_scale = infinicore::adaptor::to_aten_tensor(scale_work); + + infinicore::adaptor::lightop::per_token_dynamic_quant_int8( + x_packed, x, x_scale, p->smooth); + + if (packed_need_copy_back) { + p->x_packed->copy_from(packed_work); + } + if (scale_need_copy_back) { + p->x_scale->copy_from(scale_work); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + if (!infinicore::adaptor::lightop::available()) { + return false; + } + PerChannelQuantI8::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); + PerChannelQuantI8::run_dispatcher().registerDevice(Device::Type::HYGON, &run); + PerChannelQuantI8::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); + return true; +}(); + +} // namespace infinicore::op::per_channel_quant_i8_impl::lightop_hygon + +#endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/src/infinicore/ops/rms_rotary_embedding/hygon/rms_rotary_embedding_lightop_hygon.cc b/src/infinicore/ops/rms_rotary_embedding/hygon/rms_rotary_embedding_lightop_hygon.cc new file mode 100644 index 000000000..a988d70b8 --- /dev/null +++ b/src/infinicore/ops/rms_rotary_embedding/hygon/rms_rotary_embedding_lightop_hygon.cc @@ -0,0 +1,89 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/ops/rms_rotary_embedding.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/lightop_adaptor.hpp" + +#include + +#include + +namespace infinicore::op::rms_rotary_embedding_impl::lightop_hygon { + +struct PlannedMeta { + graph::GraphTensor query; + graph::GraphTensor key; + graph::GraphTensor positions; + graph::GraphTensor cos_sin_cache; + graph::GraphTensor q_weight; + graph::GraphTensor k_weight; + int64_t head_size; + bool is_neox; + float epsilon; +}; + +void *plan(Tensor query, + Tensor key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache, + bool is_neox, + const Tensor &q_weight, + const Tensor &k_weight, + float epsilon) { + infinicore::adaptor::lightop::preload_rms_rotary_embedding(); + return new PlannedMeta{ + graph::GraphTensor(query), + graph::GraphTensor(key), + graph::GraphTensor(positions), + graph::GraphTensor(cos_sin_cache), + graph::GraphTensor(q_weight), + graph::GraphTensor(k_weight), + head_size, + is_neox, + epsilon}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + auto query = infinicore::adaptor::to_aten_tensor(p->query); + auto key = infinicore::adaptor::to_aten_tensor(p->key); + auto positions = infinicore::adaptor::to_aten_tensor(p->positions); + auto cos_sin_cache = infinicore::adaptor::to_aten_tensor(p->cos_sin_cache); + auto q_weight = infinicore::adaptor::to_aten_tensor(p->q_weight); + auto k_weight = infinicore::adaptor::to_aten_tensor(p->k_weight); + + infinicore::adaptor::lightop::rms_rotary_embedding_fuse( + positions, + query, + key, + p->head_size, + cos_sin_cache, + p->is_neox, + q_weight, + k_weight, + std::nullopt, + std::nullopt, + static_cast(p->epsilon)); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + if (!infinicore::adaptor::lightop::available()) { + return false; + } + RMSRotaryEmbedding::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); + RMSRotaryEmbedding::run_dispatcher().registerDevice(Device::Type::HYGON, &run); + RMSRotaryEmbedding::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); + return true; +}(); + +} // namespace infinicore::op::rms_rotary_embedding_impl::lightop_hygon + +#endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/src/infinicore/ops/rms_rotary_embedding/rms_rotary_embedding.cc b/src/infinicore/ops/rms_rotary_embedding/rms_rotary_embedding.cc new file mode 100644 index 000000000..f686b9049 --- /dev/null +++ b/src/infinicore/ops/rms_rotary_embedding/rms_rotary_embedding.cc @@ -0,0 +1,76 @@ +#include "infinicore/ops/rms_rotary_embedding.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RMSRotaryEmbedding); + +RMSRotaryEmbedding::RMSRotaryEmbedding(Tensor query, + Tensor key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache, + bool is_neox, + const Tensor &q_weight, + const Tensor &k_weight, + float epsilon) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(query, key, positions, cos_sin_cache, q_weight, k_weight); + INFINICORE_GRAPH_OP_DISPATCH(query->device().getType(), + query, + key, + positions, + head_size, + cos_sin_cache, + is_neox, + q_weight, + k_weight, + epsilon); +} + +void RMSRotaryEmbedding::execute(Tensor query, + Tensor key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache, + bool is_neox, + const Tensor &q_weight, + const Tensor &k_weight, + float epsilon) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(RMSRotaryEmbedding, + query, + key, + positions, + head_size, + cos_sin_cache, + is_neox, + q_weight, + k_weight, + epsilon); +} + +bool rms_rotary_embedding_fuse_available(const Device &device) { + return RMSRotaryEmbedding::plan_dispatcher().lookup(device.getType()) != nullptr; +} + +void rms_rotary_embedding_fuse_(Tensor query, + Tensor key, + const Tensor &positions, + int64_t head_size, + const Tensor &cos_sin_cache, + bool is_neox, + const Tensor &q_weight, + const Tensor &k_weight, + float epsilon) { + RMSRotaryEmbedding::execute(query, + key, + positions, + head_size, + cos_sin_cache, + is_neox, + q_weight, + k_weight, + epsilon); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/scaled_mm_i8/hygon/scaled_mm_i8_lightop_hygon.cc b/src/infinicore/ops/scaled_mm_i8/hygon/scaled_mm_i8_lightop_hygon.cc new file mode 100644 index 000000000..a22701f5a --- /dev/null +++ b/src/infinicore/ops/scaled_mm_i8/hygon/scaled_mm_i8_lightop_hygon.cc @@ -0,0 +1,102 @@ +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/ops/scaled_mm_i8.hpp" + +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/lightop_adaptor.hpp" + +#include +#include + +#include + +namespace infinicore::op::scaled_mm_i8_impl::lightop_hygon { + +struct PlannedMeta { + graph::GraphTensor c; + graph::GraphTensor a_p; + graph::GraphTensor a_s; + graph::GraphTensor b_p; + graph::GraphTensor b_s; + std::optional bias; +}; + +void *plan(Tensor c, + const Tensor &a_p, + const Tensor &a_s, + const Tensor &b_p, + const Tensor &b_s, + std::optional bias) { + infinicore::adaptor::lightop::preload_w8a8_linear_ops(); + if (c->ndim() != 2 || a_p->ndim() != 2 || b_p->ndim() != 2) { + throw std::runtime_error("Hygon scaled_mm_i8 expects 2D tensors"); + } + if (a_p->dtype() != DataType::I8 || b_p->dtype() != DataType::I8 || + a_s->dtype() != DataType::F32 || b_s->dtype() != DataType::F32) { + throw std::runtime_error("Hygon scaled_mm_i8 expects int8 inputs and float32 scales"); + } + if (a_p->shape()[0] != c->shape()[0] || b_p->shape()[1] != c->shape()[1] || + a_p->shape()[1] != b_p->shape()[0]) { + throw std::runtime_error("Hygon scaled_mm_i8 matrix shape mismatch"); + } + + return new PlannedMeta{ + graph::GraphTensor(c), + graph::GraphTensor(a_p), + graph::GraphTensor(a_s), + graph::GraphTensor(b_p), + graph::GraphTensor(b_s), + bias ? std::optional(graph::GraphTensor(*bias)) : std::nullopt}; +} + +void run(void *planned_meta) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto *p = reinterpret_cast(planned_meta); + + const bool c_need_copy_back = !p->c->is_contiguous(); + Tensor c_work = c_need_copy_back ? p->c->contiguous() : Tensor(p->c); + Tensor a_work = p->a_p->is_contiguous() ? Tensor(p->a_p) : p->a_p->contiguous(); + Tensor a_scale_work = p->a_s->is_contiguous() ? Tensor(p->a_s) : p->a_s->contiguous(); + Tensor b_scale_work = p->b_s->is_contiguous() ? Tensor(p->b_s) : p->b_s->contiguous(); + + Tensor b_work = Tensor(p->b_p); + + auto c = infinicore::adaptor::to_aten_tensor(c_work); + auto a = infinicore::adaptor::to_aten_tensor(a_work); + auto a_scale = infinicore::adaptor::to_aten_tensor(a_scale_work); + auto b_scale = infinicore::adaptor::to_aten_tensor(b_scale_work); + std::optional bias = std::nullopt; + if (p->bias.has_value()) { + bias = infinicore::adaptor::to_aten_tensor(Tensor(p->bias.value())); + } + + Tensor b_nk_work = b_work->permute({1, 0}); + if (!b_nk_work->is_contiguous()) { + b_nk_work = b_nk_work->contiguous(); + } + auto b_nk = infinicore::adaptor::to_aten_tensor(b_nk_work); + infinicore::adaptor::lightop::blaslt_w8a8_gemm( + c, a, b_nk, a_scale, b_scale, bias); + + if (c_need_copy_back) { + p->c->copy_from(c_work); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + if (!infinicore::adaptor::lightop::available()) { + return false; + } + I8Gemm::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); + I8Gemm::run_dispatcher().registerDevice(Device::Type::HYGON, &run); + I8Gemm::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); + return true; +}(); + +} // namespace infinicore::op::scaled_mm_i8_impl::lightop_hygon + +#endif // ENABLE_HYGON_API && ENABLE_ATEN diff --git a/src/infinicore/ops/silu_and_mul/hygon/silu_and_mul_lightop_hygon.cc b/src/infinicore/ops/silu_and_mul/hygon/silu_and_mul_lightop_hygon.cc index 0472eb68c..dcd4c8df5 100644 --- a/src/infinicore/ops/silu_and_mul/hygon/silu_and_mul_lightop_hygon.cc +++ b/src/infinicore/ops/silu_and_mul/hygon/silu_and_mul_lightop_hygon.cc @@ -44,8 +44,7 @@ void cleanup(void **planned_meta_ptr) { } static bool registered = []() { - if (!infinicore::adaptor::lightop::enabled_by_env() - || !infinicore::adaptor::lightop::available()) { + if (!infinicore::adaptor::lightop::available()) { return false; } SiluAndMul::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); diff --git a/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cu b/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cu index 8414e187e..7db3d6532 100644 --- a/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cu +++ b/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cu @@ -2,7 +2,6 @@ #include "../../../devices/nvidia/nvidia_common.cuh" #include "../../../devices/nvidia/nvidia_kernel_common.cuh" #include "../../../tensor.h" -#include "../cuda/embedding_kernel.cuh" #include "embedding_nvidia.cuh" #include @@ -14,8 +13,9 @@ INFINIOP_CUDA_KERNEL embeddingKernel( size_t num_indices, size_t embedding_dim, size_t vocab_size) { - // Calculate global thread index - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + // Assign one block to each index so threads copy the embedding row + // cooperatively. Decode commonly has a single index and a wide row. + size_t idx = blockIdx.x; if (idx < num_indices) { // Get the index value @@ -27,35 +27,8 @@ INFINIOP_CUDA_KERNEL embeddingKernel( const T *src = weight + static_cast(index_val) * embedding_dim; T *dst = output + idx * embedding_dim; - // Choose optimal copy strategy based on type and alignment - if constexpr (std::is_same_v) { - // Check alignment for float4 (16 bytes) - bool aligned_16 = is_aligned(src, 16) && is_aligned(dst, 16); - if (aligned_16 && embedding_dim >= 4 && embedding_dim % 4 == 0) { - copyVectorizedFloat4(dst, src, embedding_dim); - } else if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - // Try float2 if not aligned to 16 bytes - copyVectorizedFloat2(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else if constexpr (std::is_same_v) { - // Use half2 for vectorized access - if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - copyVectorizedHalf2(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else if constexpr (std::is_same_v) { - // Use bfloat162 for vectorized access - if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - copyVectorizedBFloat162(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else { - // Fallback to scalar copy with __ldg - copyScalar(dst, src, embedding_dim); + for (size_t col = threadIdx.x; col < embedding_dim; col += blockDim.x) { + dst[col] = __ldg(&src[col]); } } } @@ -135,17 +108,8 @@ infiniStatus_t Descriptor::calculate( auto cuda_stream = reinterpret_cast(stream); - // Dynamic block size optimization based on embedding_dim - // Smaller embedding_dim benefits from larger block size (better occupancy) - // Larger embedding_dim benefits from smaller block size (more registers per thread) - size_t block_size = 256; // Default - if (_embedding_dim <= 64) { - block_size = 512; // Small embedding_dim: use larger block for better occupancy - } else if (_embedding_dim >= 1024) { - block_size = 128; // Large embedding_dim: use smaller block to reduce register pressure - } - - size_t grid_size = (_num_indices + block_size - 1) / block_size; + constexpr size_t block_size = 256; + size_t grid_size = _num_indices; // Launch kernel based on dtypes if (_input_dtype == INFINI_DTYPE_I32) { From 60cc8c561722a84aed818cedd8c417e70106f42a Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 13 Jul 2026 15:11:55 +0800 Subject: [PATCH 06/16] fix-hygon-w16a16-marlin-graph-safe --- .../infinicore/adaptor/lightop_adaptor.hpp | 4 + src/infinicore/adaptor/lightop_adaptor.cc | 276 +++++++++++++++++- .../hygon/moe_w16a16_marlin_hygon.cc | 55 +++- 3 files changed, 328 insertions(+), 7 deletions(-) diff --git a/include/infinicore/adaptor/lightop_adaptor.hpp b/include/infinicore/adaptor/lightop_adaptor.hpp index 0d3d9ed21..e73662768 100644 --- a/include/infinicore/adaptor/lightop_adaptor.hpp +++ b/include/infinicore/adaptor/lightop_adaptor.hpp @@ -12,6 +12,10 @@ bool available(); void preload_moe_w16a16_ops(); +void preload_moe_w16a16_ops(bool preload_legacy_gemm, bool preload_legacy_asm); + +void preload_moe_w16a16_marlin_asm(bool down_stage); + void preload_moe_w8a8_ops(); void preload_moe_align(); diff --git a/src/infinicore/adaptor/lightop_adaptor.cc b/src/infinicore/adaptor/lightop_adaptor.cc index 37f50477e..a694d5cfc 100644 --- a/src/infinicore/adaptor/lightop_adaptor.cc +++ b/src/infinicore/adaptor/lightop_adaptor.cc @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -116,6 +117,74 @@ LightopLibrary &library() { return lib; } +constexpr const char *kMoeW16A16MarlinMode1000UpCo = + "moe_w16a16_channel/moe_w16a16_marlin_128x256x64_TN_BF16_UP.co"; +constexpr const char *kMoeW16A16MarlinMode1000DownCo = + "moe_w16a16_channel/moe_w16a16_marlin_128x256x64_TN_BF16_DOWN.co"; +constexpr const char *kMoeW16A16MarlinMode1000UpKernel = + "MOE_W16A16_BF16_PERCHANNEL_MARLIN_ASM_TN_MT128x256x64_WGM1_UP"; +constexpr const char *kMoeW16A16MarlinMode1000DownKernel = + "MOE_W16A16_BF16_PERCHANNEL_MARLIN_ASM_TN_MT128x256x64_WGM1_DOWN"; +constexpr uint32_t kMoeW16A16MarlinNBlock = 256; +constexpr uint32_t kMoeW16A16MarlinMBlock = 128; +constexpr uint32_t kMoeW16A16MarlinKBlock = 64; +constexpr uint32_t kMoeW16A16MarlinWorkgroupSize = 768; + +struct MoeW16A16MarlinMode1000Args { + uint32_t n_block_count; + uint32_t m_block_count; + void *output; + void *weight; + void *input; + void *scale_a; + void *scale_b; + void *topk_weights; + void *sorted_token_ids; + void *expert_ids; + void *num_tokens_post_padded; + uint32_t num_experts; + uint32_t m; + uint32_t n; + uint32_t k; + uint32_t stride_asm; + uint32_t stride_ask; + uint32_t stride_bse; + uint32_t stride_bsn; + uint32_t stride_bsk; + uint32_t sorted_token_lens; + uint32_t top_k; + float inverse_top_k; + float inverse_delta; + uint32_t reserved0; + uint32_t reserved1; + uint32_t reserved2; +}; +static_assert(sizeof(MoeW16A16MarlinMode1000Args) == 144); +static_assert(offsetof(MoeW16A16MarlinMode1000Args, output) == 8); +static_assert(offsetof(MoeW16A16MarlinMode1000Args, topk_weights) == 48); +static_assert(offsetof(MoeW16A16MarlinMode1000Args, num_experts) == 80); +static_assert(offsetof(MoeW16A16MarlinMode1000Args, sorted_token_lens) == 116); +static_assert(offsetof(MoeW16A16MarlinMode1000Args, inverse_top_k) == 124); +static_assert(offsetof(MoeW16A16MarlinMode1000Args, inverse_delta) == 128); +static_assert(offsetof(MoeW16A16MarlinMode1000Args, reserved0) == 132); + +struct MoeW16A16MarlinDeviceKernel { + hipModule_t module = nullptr; + hipFunction_t function = nullptr; +}; + +std::mutex &moe_w16a16_marlin_kernel_mutex() { + static std::mutex mutex; + return mutex; +} + +std::unordered_map & +moe_w16a16_marlin_kernels(bool down_stage) { + static std::unordered_map up_kernels; + static std::unordered_map down_kernels; + return down_stage ? down_kernels : up_kernels; +} + constexpr const char *kMoeW8A8MarlinMode1001Co = "moe_w8a8_channel/moe_w8a8_i8_marlin_64x256x128_TN_BF16_UP.co"; constexpr const char *kMoeW8A8MarlinMode1001Kernel = @@ -179,6 +248,53 @@ std::string hip_error_message(const std::string &operation, hipError_t status) { return oss.str(); } +MoeW16A16MarlinDeviceKernel get_moe_w16a16_marlin_mode1000_kernel( + bool down_stage) { + int device = -1; + auto status = hipGetDevice(&device); + if (status != hipSuccess) { + throw std::runtime_error(hip_error_message("hipGetDevice", status)); + } + + std::lock_guard lock(moe_w16a16_marlin_kernel_mutex()); + auto &kernels = moe_w16a16_marlin_kernels(down_stage); + auto found = kernels.find(device); + if (found != kernels.end()) { + return found->second; + } + + ensure_default_lightop_env(); + const char *asm_dir_env = std::getenv("LIGHTOP_ASM_DIR"); + std::string asm_dir = + asm_dir_env != nullptr && asm_dir_env[0] != '\0' + ? asm_dir_env + : kDefaultLightopAsmDir; + if (!asm_dir.empty() && asm_dir.back() != '/') { + asm_dir.push_back('/'); + } + const char *co = down_stage + ? kMoeW16A16MarlinMode1000DownCo + : kMoeW16A16MarlinMode1000UpCo; + const char *function = down_stage + ? kMoeW16A16MarlinMode1000DownKernel + : kMoeW16A16MarlinMode1000UpKernel; + const std::string co_path = asm_dir + co; + + MoeW16A16MarlinDeviceKernel kernel; + status = hipModuleLoad(&kernel.module, co_path.c_str()); + if (status != hipSuccess) { + throw std::runtime_error(hip_error_message("hipModuleLoad(" + co_path + ")", status)); + } + status = hipModuleGetFunction(&kernel.function, kernel.module, function); + if (status != hipSuccess) { + (void)hipModuleUnload(kernel.module); + throw std::runtime_error(hip_error_message("hipModuleGetFunction", status)); + } + + kernels.emplace(device, kernel); + return kernel; +} + MoeW8A8MarlinDeviceKernel get_moe_w8a8_marlin_mode1001_kernel() { int device = -1; auto status = hipGetDevice(&device); @@ -228,6 +344,143 @@ uint32_t checked_u32(int64_t value, const char *name) { return static_cast(value); } +uint32_t checked_w16_u32(int64_t value, const char *name) { + if (value < 0 || + static_cast(value) > std::numeric_limits::max()) { + throw std::runtime_error(std::string("Hygon W16A16 Marlin ") + name + " exceeds uint32"); + } + return static_cast(value); +} + +void launch_moe_w16a16_marlin_mode1000( + at::Tensor &input, + at::Tensor &weight, + at::Tensor &output, + const std::optional &topk_weights, + at::Tensor &sorted_token_ids, + at::Tensor &expert_ids, + at::Tensor &num_tokens_post_padded, + int64_t top_k, + int delta) { + const auto device = input.device(); + if (delta <= 0 || top_k <= 0 || + input.dim() != 2 || weight.dim() != 3 || output.dim() != 2 || + sorted_token_ids.dim() != 1 || expert_ids.dim() != 1 || + input.scalar_type() != at::kBFloat16 || + weight.scalar_type() != at::kBFloat16 || + output.scalar_type() != at::kBFloat16 || + sorted_token_ids.scalar_type() != at::kInt || + expert_ids.scalar_type() != at::kInt || + num_tokens_post_padded.scalar_type() != at::kInt || + !input.is_contiguous() || !weight.is_contiguous() || + !output.is_contiguous() || !sorted_token_ids.is_contiguous() || + !expert_ids.is_contiguous() || !num_tokens_post_padded.is_contiguous() || + weight.device() != device || output.device() != device || + sorted_token_ids.device() != device || expert_ids.device() != device || + num_tokens_post_padded.device() != device || + (topk_weights.has_value() && + (topk_weights->scalar_type() != at::kFloat || + !topk_weights->is_contiguous() || topk_weights->device() != device))) { + throw std::runtime_error("Hygon W16A16 Marlin mode 1000 tensor contract mismatch"); + } + + const int64_t m = input.size(0); + const int64_t k = input.size(1); + const int64_t num_experts = weight.size(0); + const int64_t n = weight.size(2) / 16; + const int64_t sorted_token_lens = sorted_token_ids.numel(); + if (m <= 0 || k <= 0 || n <= 0 || num_experts <= 0 || + m > std::numeric_limits::max() / top_k || + k % kMoeW16A16MarlinKBlock != 0 || + weight.size(1) != k / 16 || weight.size(2) != n * 16 || + output.size(0) != m * top_k || output.size(1) != n || + num_tokens_post_padded.numel() != 1 || sorted_token_lens <= 0 || + (topk_weights.has_value() && topk_weights->numel() != m)) { + throw std::runtime_error("Hygon W16A16 Marlin mode 1000 tensor shape mismatch"); + } + + const uint32_t n_u32 = checked_w16_u32(n, "N"); + const uint32_t sorted_token_lens_u32 = + checked_w16_u32(sorted_token_lens, "sorted_token_lens"); + (void)checked_w16_u32(expert_ids.numel(), "expert_ids capacity"); + const uint32_t top_k_u32 = checked_w16_u32(top_k, "top_k"); + const uint32_t n_block_count = + 1 + (n_u32 - 1) / kMoeW16A16MarlinNBlock; + const uint32_t m_block_count = + 1 + (sorted_token_lens_u32 - 1) / kMoeW16A16MarlinMBlock; + if (static_cast(expert_ids.numel()) < m_block_count) { + throw std::runtime_error("Hygon W16A16 Marlin mode 1000 expert_ids capacity mismatch"); + } + + MoeW16A16MarlinMode1000Args args{ + n_block_count, + m_block_count, + output.data_ptr(), + weight.data_ptr(), + input.data_ptr(), + nullptr, + nullptr, + topk_weights.has_value() ? topk_weights->data_ptr() : nullptr, + sorted_token_ids.data_ptr(), + expert_ids.data_ptr(), + num_tokens_post_padded.data_ptr(), + checked_w16_u32(num_experts, "num_experts"), + checked_w16_u32(m, "M"), + n_u32, + checked_w16_u32(k, "K"), + 0, + 0, + 0, + 0, + 0, + sorted_token_lens_u32, + top_k_u32, + 1.0f / static_cast(top_k_u32), + 1.0f / static_cast(delta), + 0, + 0, + 0}; + + int current_device = -1; + auto status = hipGetDevice(¤t_device); + if (status != hipSuccess) { + throw std::runtime_error(hip_error_message("hipGetDevice", status)); + } + if (input.get_device() != current_device) { + throw std::runtime_error("Hygon W16A16 Marlin mode 1000 current device mismatch"); + } + + size_t args_size = sizeof(args); + void *launch_config[] = { + HIP_LAUNCH_PARAM_BUFFER_POINTER, + &args, + HIP_LAUNCH_PARAM_BUFFER_SIZE, + &args_size, + HIP_LAUNCH_PARAM_END}; + + const bool down_stage = topk_weights.has_value(); + auto kernel = get_moe_w16a16_marlin_mode1000_kernel(down_stage); + status = hipModuleLaunchKernel( + kernel.function, + n_block_count, + 1, + m_block_count, + kMoeW16A16MarlinWorkgroupSize, + 1, + 1, + 0, + infinicore::adaptor::get_hip_stream().stream(), + nullptr, + launch_config); + if (status != hipSuccess) { + const char *stage = down_stage ? "DOWN" : "UP"; + throw std::runtime_error( + hip_error_message( + std::string("hipModuleLaunchKernel(W16A16 Marlin mode 1000 ") + stage + ")", + status)); + } +} + bool can_launch_moe_w8a8_marlin_mode1001( const at::Tensor &input, const at::Tensor &weight, @@ -587,10 +840,22 @@ bool available() { return library().available(); } -void preload_moe_w16a16_ops() { +void preload_moe_w16a16_ops(bool preload_legacy_gemm, bool preload_legacy_asm) { (void)moe_sum_fn(); - (void)moe_gemm_w16a16_fn(); - (void)moe_marlin_w16a16_asm_fn(); + if (preload_legacy_gemm) { + (void)moe_gemm_w16a16_fn(); + } + if (preload_legacy_asm) { + (void)moe_marlin_w16a16_asm_fn(); + } +} + +void preload_moe_w16a16_ops() { + preload_moe_w16a16_ops(true, true); +} + +void preload_moe_w16a16_marlin_asm(bool down_stage) { + (void)get_moe_w16a16_marlin_mode1000_kernel(down_stage); } void preload_moe_w8a8_ops() { @@ -702,6 +967,11 @@ void moe_gemm_marlin_w16a16(at::Tensor input, input, b_qweight, output, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded, top_k, mode, delta); + } else if (mode == 1000 && input.scalar_type() == at::kBFloat16) { + launch_moe_w16a16_marlin_mode1000( + input, b_qweight, output, topk_weights, + sorted_token_ids, expert_ids, num_tokens_post_padded, + top_k, delta); } else { moe_marlin_w16a16_asm_fn()( input, b_qweight, output, topk_weights, diff --git a/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc b/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc index 18d46e409..523d176e1 100644 --- a/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc +++ b/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc @@ -3,6 +3,7 @@ #include "infinicore/adaptor/aten_adaptor.hpp" #include "infinicore/adaptor/lightop_adaptor.hpp" +#include "infinicore/context/context.hpp" #include "infinicore/ops/common/dispatcher.hpp" #include "infinicore/ops/moe_sum.hpp" #include "infinicore/ops/silu_and_mul.hpp" @@ -123,7 +124,41 @@ void *plan(Tensor output, int delta0, int mode1, int delta1) { - infinicore::adaptor::lightop::preload_moe_w16a16_ops(); + infinicore::context::setDevice(hidden_states->device()); + const auto device = hidden_states->device(); + const auto same_device = [&](const Tensor &tensor) { + return tensor && tensor->device() == device; + }; + if (!same_device(output) || !same_device(cache13) || !same_device(cache2) || + !same_device(w13_marlin) || !same_device(w2_marlin) || + !same_device(topk_weights) || !same_device(sorted_token_ids) || + !same_device(expert_ids) || !same_device(num_tokens_post_padded)) { + throw std::runtime_error("w16a16 marlin fused dense tensors must be on one device"); + } + + const bool bf16 = hidden_states->dtype() == infinicore::DataType::BF16; + const bool direct_mode0 = bf16 && mode0 == 1000; + const bool direct_mode1 = bf16 && mode1 == 1000; + if ((direct_mode0 || direct_mode1) && + (!output->is_contiguous() || !cache13->is_contiguous() || + !cache2->is_contiguous() || !hidden_states->is_contiguous() || + !w13_marlin->is_contiguous() || !w2_marlin->is_contiguous() || + !topk_weights->is_contiguous() || !sorted_token_ids->is_contiguous() || + !expert_ids->is_contiguous() || !num_tokens_post_padded->is_contiguous())) { + throw std::runtime_error("Hygon W16A16 Marlin mode 1000 requires contiguous tensors"); + } + + const bool preload_legacy_gemm = mode0 < 1000 || mode1 < 1000; + const bool preload_legacy_asm = + (mode0 >= 1000 && !direct_mode0) || (mode1 >= 1000 && !direct_mode1); + infinicore::adaptor::lightop::preload_moe_w16a16_ops( + preload_legacy_gemm, preload_legacy_asm); + if (direct_mode0) { + infinicore::adaptor::lightop::preload_moe_w16a16_marlin_asm(false); + } + if (direct_mode1) { + infinicore::adaptor::lightop::preload_moe_w16a16_marlin_asm(true); + } return new PlannedMeta{ graph::GraphTensor(output), graph::GraphTensor(cache13), graph::GraphTensor(cache2), graph::GraphTensor(hidden_states), graph::GraphTensor(w13_marlin), graph::GraphTensor(w2_marlin), @@ -132,8 +167,9 @@ void *plan(Tensor output, } void run(void *planned_meta) { - c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); auto *p = reinterpret_cast(planned_meta); + infinicore::context::setDevice(p->hidden_states->device()); + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); const auto hidden_shape = p->hidden_states->shape(); const auto w13_shape = p->w13_marlin->shape(); @@ -149,9 +185,20 @@ void run(void *planned_meta) { throw std::runtime_error("w16a16 marlin fused dense weight shape mismatch"); } - const bool output_need_copy_back = !p->output->is_contiguous(); + const bool uses_direct_mode1000 = + p->hidden_states->dtype() == infinicore::DataType::BF16 && + (p->mode0 == 1000 || p->mode1 == 1000); + if (uses_direct_mode1000 && + (!p->output->is_contiguous() || !p->hidden_states->is_contiguous())) { + throw std::runtime_error("Hygon W16A16 Marlin mode 1000 requires contiguous input and output"); + } + const bool output_need_copy_back = + !uses_direct_mode1000 && !p->output->is_contiguous(); Tensor output_work_ic = output_need_copy_back ? p->output->contiguous() : Tensor(p->output); - Tensor hidden_work_ic = p->hidden_states->is_contiguous() ? Tensor(p->hidden_states) : p->hidden_states->contiguous(); + Tensor hidden_work_ic = + uses_direct_mode1000 || p->hidden_states->is_contiguous() + ? Tensor(p->hidden_states) + : p->hidden_states->contiguous(); const size_t top_k = p->top_k; const size_t cache1_numel = m * top_k * n2; From 6267fb89c55cd0dfa12f09c2d1888bb114febbfe Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 13 Jul 2026 15:16:18 +0800 Subject: [PATCH 07/16] perf-hygon-tp8-bf16-graph-allreduce --- src/infiniccl/cuda/infiniccl_cuda.cu | 84 ++++++++++++++-------------- 1 file changed, 41 insertions(+), 43 deletions(-) diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu index 9fdcd6749..ba8534e66 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ b/src/infiniccl/cuda/infiniccl_cuda.cu @@ -85,6 +85,7 @@ constexpr size_t kHygonTp2StageCapacityElements = 1u << 22; constexpr int kHygonTp8WorldSize = 8; constexpr int kHygonTp8Threads = 512; constexpr int kHygonTp8MaxBlocks = 80; +constexpr int kHygonTp8OneStageMaxBlocks = 16; constexpr size_t kHygonTp8OneStageMaxBytes = 80u * 1024u; constexpr size_t kHygonTp8TwoStageMaxBytes = 512u * 1024u; constexpr int kHygonHipSuccess = 0; @@ -621,52 +622,44 @@ __global__ __launch_bounds__(kHygonTp8Threads, 1) void hygon_tp8_bf16_allreduce_ int rank, size_t pack_count) { constexpr int num_ranks = kHygonTp8WorldSize; - constexpr int threads_per_rank = kHygonTp8Threads / num_ranks; - __shared__ HygonBf16Pack shared_packs[kHygonTp8Threads]; __shared__ uint32_t block_flag; - const int source_rank = threadIdx.x / threads_per_rank; - const int lane = threadIdx.x % threads_per_rank; const uint32_t sync_flag = hygon_tp8_start_sync( rank_signals, self_signal, rank, &block_flag); - for (size_t base = blockIdx.x * threads_per_rank; - base < pack_count; - base += gridDim.x * threads_per_rank) { - const size_t index = base + lane; - HygonBf16Pack source_pack{}; - if (index < pack_count) { - const auto *source = reinterpret_cast( - rank_data.ptrs[source_rank]); - source_pack = source[index]; - } - shared_packs[threadIdx.x] = source_pack; - __syncthreads(); - - if (source_rank == 0 && index < pack_count) { - float reduced[8]; -#pragma unroll - for (int element = 0; element < 8; ++element) { - reduced[element] = __bfloat162float( - shared_packs[threadIdx.x].values[element]); - } + const size_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t thread_stride = + static_cast(gridDim.x) * blockDim.x; + auto *packed_output = reinterpret_cast(output); + for (size_t index = thread_index; + index < pack_count; + index += thread_stride) { + const auto *source = reinterpret_cast( + rank_data.ptrs[0]); + HygonBf16Pack source_pack = source[index]; + float reduced[8]; #pragma unroll - for (int peer = 1; peer < num_ranks; ++peer) { + for (int element = 0; element < 8; ++element) { + reduced[element] = __bfloat162float(source_pack.values[element]); + } #pragma unroll - for (int element = 0; element < 8; ++element) { - reduced[element] += __bfloat162float( - shared_packs[peer * threads_per_rank + threadIdx.x] - .values[element]); - } - } - HygonBf16Pack result; + for (int peer = 1; peer < num_ranks; ++peer) { + source = reinterpret_cast( + rank_data.ptrs[peer]); + source_pack = source[index]; #pragma unroll for (int element = 0; element < 8; ++element) { - result.values[element] = __float2bfloat16(reduced[element]); + reduced[element] += + __bfloat162float(source_pack.values[element]); } - reinterpret_cast(output)[index] = result; } - __syncthreads(); + HygonBf16Pack result; +#pragma unroll + for (int element = 0; element < 8; ++element) { + result.values[element] = __float2bfloat16(reduced[element]); + } + packed_output[index] = result; } hygon_tp8_end_sync( @@ -957,18 +950,23 @@ HygonTp8AllReduceResult try_hygon_tp8_graph_allreduce( return HygonTp8AllReduceResult::Fallback; } - constexpr int threads_per_rank = - kHygonTp8Threads / kHygonTp8WorldSize; size_t pack_count = count / 8; const bool use_two_stage = count * sizeof(__nv_bfloat16) >= kHygonTp8OneStageMaxBytes; - const size_t work_pack_count = use_two_stage - ? pack_count / kHygonTp8WorldSize + - pack_count % kHygonTp8WorldSize - : pack_count; + const size_t work_pack_count = + use_two_stage + ? pack_count / kHygonTp8WorldSize + + pack_count % kHygonTp8WorldSize + : pack_count; + const size_t work_threads = + use_two_stage + ? kHygonTp8Threads / kHygonTp8WorldSize + : kHygonTp8Threads; + const int max_blocks = + use_two_stage ? kHygonTp8MaxBlocks : kHygonTp8OneStageMaxBlocks; int blocks = static_cast(std::min( - kHygonTp8MaxBlocks, - (work_pack_count + threads_per_rank - 1) / threads_per_rank)); + max_blocks, + (work_pack_count + work_threads - 1) / work_threads)); blocks = std::max(blocks, 1); HygonTp8RankSignals rank_signals = state->rank_signals; HygonTp8Signal *self_signal = state->signals[rank]; From 8ebab1bf17741e8210898f123ca11a27319491d2 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 13 Jul 2026 13:00:56 +0000 Subject: [PATCH 08/16] feat(hygon): auto-detect LightOP device config --- .../infinicore/adaptor/lightop_adaptor.hpp | 17 +- src/infinicore/adaptor/lightop_adaptor.cc | 167 ++++++++++++++---- 2 files changed, 152 insertions(+), 32 deletions(-) diff --git a/include/infinicore/adaptor/lightop_adaptor.hpp b/include/infinicore/adaptor/lightop_adaptor.hpp index e73662768..37b158d60 100644 --- a/include/infinicore/adaptor/lightop_adaptor.hpp +++ b/include/infinicore/adaptor/lightop_adaptor.hpp @@ -1,6 +1,21 @@ -#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) #pragma once +#include +#include + +namespace infinicore::adaptor::lightop { + +struct DeviceInfo { + std::string gpu_target; + int compute_units = 0; +}; + +DeviceInfo device_info(std::size_t device_index); + +} // namespace infinicore::adaptor::lightop + +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) + #include #include diff --git a/src/infinicore/adaptor/lightop_adaptor.cc b/src/infinicore/adaptor/lightop_adaptor.cc index a694d5cfc..f2178cbb0 100644 --- a/src/infinicore/adaptor/lightop_adaptor.cc +++ b/src/infinicore/adaptor/lightop_adaptor.cc @@ -1,19 +1,25 @@ -#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) #include "infinicore/adaptor/lightop_adaptor.hpp" + +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) #include "infinicore/adaptor/aten_adaptor.hpp" #include #include +#include +#include +#include #include #include #include +#include #include #include #include #include #include #include +#include namespace infinicore::adaptor::lightop { namespace { @@ -22,9 +28,8 @@ constexpr const char *kDefaultLightopSo = "/usr/local/lib/python3.10/dist-packages/lightop/op.cpython-310-x86_64-linux-gnu.so"; constexpr const char *kDefaultLmslimQuantSo = "/usr/local/lib/python3.10/dist-packages/lmslimquant.cpython-310-x86_64-linux-gnu.so"; -constexpr const char *kDefaultLightopGpuTarget = "gfx936"; -constexpr const char *kDefaultLightopAsmDir = - "/usr/local/lib/python3.10/dist-packages/lightop/hsa/gfx936/"; +constexpr const char *kLightopAsmRoot = + "/usr/local/lib/python3.10/dist-packages/lightop/hsa/"; constexpr const char *kFuseSiluAndMulSymbol = "_ZN2at6native17fuse_silu_and_mulERNS_6TensorES2_"; constexpr const char *kRmsRotaryEmbeddingFuseSymbol = @@ -50,7 +55,98 @@ constexpr const char *kBlasltW8A8Bf16Symbol = constexpr const char *kBlasltW8A8Fp16Symbol = "_ZN14hipblaslt_gemm14w8a8_fp16_gemmERKN2at6TensorES3_S3_S3_RS1_llllRKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEES3_S3_RKSt8optionalIS1_E"; -void ensure_default_lightop_env(); +std::string hip_error_message(const std::string &operation, hipError_t status); + +struct LightopRuntimeConfig { + std::string gpu_target; + std::string asm_dir; + int compute_units = 0; +}; + +std::string normalize_gpu_target(std::string target) { + const auto feature_pos = target.find(':'); + if (feature_pos != std::string::npos) { + target.resize(feature_pos); + } + std::transform(target.begin(), target.end(), target.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + if (target.size() <= 3 || target.compare(0, 3, "gfx") != 0 || + !std::all_of(target.begin() + 3, target.end(), [](unsigned char ch) { + return std::isalnum(ch) != 0; + })) { + throw std::runtime_error("invalid Hygon GPU target: " + target); + } + return target; +} + +void set_lightop_env(const char *name, const std::string &value) { + if (setenv(name, value.c_str(), 1) != 0) { + throw std::runtime_error( + std::string("failed to set ") + name + ": " + std::strerror(errno)); + } +} + +DeviceInfo query_device_info(int device) { + hipDeviceProp_t properties{}; + const auto status = hipGetDeviceProperties(&properties, device); + if (status != hipSuccess) { + throw std::runtime_error(hip_error_message("hipGetDeviceProperties", status)); + } + + const char *arch_end = std::find( + properties.gcnArchName, + properties.gcnArchName + sizeof(properties.gcnArchName), + '\0'); + if (arch_end == properties.gcnArchName + sizeof(properties.gcnArchName)) { + throw std::runtime_error("hipGetDeviceProperties returned an unterminated gcnArchName"); + } + + if (properties.multiProcessorCount <= 0) { + throw std::runtime_error("hipGetDeviceProperties returned an invalid compute unit count"); + } + + return DeviceInfo{ + normalize_gpu_target(std::string( + properties.gcnArchName, + static_cast(arch_end - properties.gcnArchName))), + properties.multiProcessorCount, + }; +} + +const DeviceInfo &cached_device_info(int device) { + thread_local std::unordered_map cache; + const auto found = cache.find(device); + if (found != cache.end()) { + return found->second; + } + return cache.emplace(device, query_device_info(device)).first->second; +} + +LightopRuntimeConfig make_lightop_runtime_config() { + int device = -1; + auto status = hipGetDevice(&device); + if (status != hipSuccess) { + throw std::runtime_error(hip_error_message("hipGetDevice", status)); + } + + const auto &detected = cached_device_info(device); + std::string gpu_target = detected.gpu_target; + std::string asm_dir = std::string(kLightopAsmRoot) + gpu_target + "/"; + + set_lightop_env("LIGHTOP_GPU_TARGET", gpu_target); + set_lightop_env("LIGHTOP_ASM_DIR", asm_dir); + return LightopRuntimeConfig{ + std::move(gpu_target), + std::move(asm_dir), + detected.compute_units, + }; +} + +const LightopRuntimeConfig &lightop_runtime_config() { + static const LightopRuntimeConfig config = make_lightop_runtime_config(); + return config; +} class LightopLibrary { public: @@ -85,7 +181,14 @@ class LightopLibrary { return true; } - ensure_default_lightop_env(); + try { + (void)lightop_runtime_config(); + } catch (const std::exception &exception) { + if (update_error || error_.empty()) { + error_ = exception.what(); + } + return false; + } const char *path_env = std::getenv("INFINICORE_LIGHTOP_SO"); const char *path = (path_env != nullptr && path_env[0] != '\0') ? path_env : kDefaultLightopSo; @@ -263,15 +366,12 @@ MoeW16A16MarlinDeviceKernel get_moe_w16a16_marlin_mode1000_kernel( return found->second; } - ensure_default_lightop_env(); - const char *asm_dir_env = std::getenv("LIGHTOP_ASM_DIR"); - std::string asm_dir = - asm_dir_env != nullptr && asm_dir_env[0] != '\0' - ? asm_dir_env - : kDefaultLightopAsmDir; - if (!asm_dir.empty() && asm_dir.back() != '/') { - asm_dir.push_back('/'); + const auto &runtime_config = lightop_runtime_config(); + const auto ¤t_device = cached_device_info(device); + if (current_device.gpu_target != runtime_config.gpu_target) { + throw std::runtime_error("LightOP does not support mixed Hygon GPU architectures in one process"); } + const auto &asm_dir = runtime_config.asm_dir; const char *co = down_stage ? kMoeW16A16MarlinMode1000DownCo : kMoeW16A16MarlinMode1000UpCo; @@ -309,15 +409,12 @@ MoeW8A8MarlinDeviceKernel get_moe_w8a8_marlin_mode1001_kernel() { return found->second; } - ensure_default_lightop_env(); - const char *asm_dir_env = std::getenv("LIGHTOP_ASM_DIR"); - std::string asm_dir = - asm_dir_env != nullptr && asm_dir_env[0] != '\0' - ? asm_dir_env - : kDefaultLightopAsmDir; - if (!asm_dir.empty() && asm_dir.back() != '/') { - asm_dir.push_back('/'); + const auto &runtime_config = lightop_runtime_config(); + const auto ¤t_device = cached_device_info(device); + if (current_device.gpu_target != runtime_config.gpu_target) { + throw std::runtime_error("LightOP does not support mixed Hygon GPU architectures in one process"); } + const auto &asm_dir = runtime_config.asm_dir; const std::string co_path = asm_dir + kMoeW8A8MarlinMode1001Co; MoeW8A8MarlinDeviceKernel kernel; @@ -656,15 +753,6 @@ LmslimQuantLibrary &lmslimquant_library() { return lib; } -void ensure_default_lightop_env() { - if (std::getenv("LIGHTOP_GPU_TARGET") == nullptr) { - setenv("LIGHTOP_GPU_TARGET", kDefaultLightopGpuTarget, 0); - } - if (std::getenv("LIGHTOP_ASM_DIR") == nullptr) { - setenv("LIGHTOP_ASM_DIR", kDefaultLightopAsmDir, 0); - } -} - template Fn resolve(const char *symbol) { return reinterpret_cast(library().symbol(symbol)); @@ -836,6 +924,13 @@ BlasltW8A8GemmFn blaslt_w8a8_fp16_fn() { } // namespace +DeviceInfo device_info(std::size_t device_index) { + if (device_index > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("Hygon device index exceeds int range"); + } + return cached_device_info(static_cast(device_index)); +} + bool available() { return library().available(); } @@ -1080,4 +1175,14 @@ void blaslt_w8a8_gemm(at::Tensor &output, } // namespace infinicore::adaptor::lightop +#else + +namespace infinicore::adaptor::lightop { + +DeviceInfo device_info(std::size_t) { + return {}; +} + +} // namespace infinicore::adaptor::lightop + #endif // ENABLE_HYGON_API && ENABLE_ATEN From f81dec804954a8704505ae02b71e1927b8fe1d08 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Thu, 23 Jul 2026 18:31:08 +0800 Subject: [PATCH 09/16] feat(hygon): support fp32 cross-entropy output --- .../ops/cross_entropy/cuda/kernel.cuh | 10 +- src/infiniop/ops/cross_entropy/info.h | 1 + .../nvidia/cross_entropy_nvidia.cu | 95 +++++++++++++------ test/infiniop/cross_entropy.py | 70 ++++++++++++-- 4 files changed, 137 insertions(+), 39 deletions(-) diff --git a/src/infiniop/ops/cross_entropy/cuda/kernel.cuh b/src/infiniop/ops/cross_entropy/cuda/kernel.cuh index c048c1233..2e7fe558d 100644 --- a/src/infiniop/ops/cross_entropy/cuda/kernel.cuh +++ b/src/infiniop/ops/cross_entropy/cuda/kernel.cuh @@ -4,9 +4,13 @@ #include "../../../devices/nvidia/nvidia_common.cuh" #include "../../../reduce/cuda/reduce.cuh" -template +template __device__ void crossEntropyKernel( - Tdata *y_, + Tout *y_, const Tdata *x_, const void *target_, size_t outer_size, @@ -73,7 +77,7 @@ __device__ void crossEntropyKernel( log_term = 0.0f; } - y_[row_idx] = static_cast(log_term - target_logit); + y_[row_idx] = static_cast(log_term - target_logit); } } diff --git a/src/infiniop/ops/cross_entropy/info.h b/src/infiniop/ops/cross_entropy/info.h index a915a4fe4..f79953cde 100644 --- a/src/infiniop/ops/cross_entropy/info.h +++ b/src/infiniop/ops/cross_entropy/info.h @@ -8,6 +8,7 @@ struct CrossEntropyInfo { int dtype; + int output_dtype; int target_dtype; size_t outer_size; size_t vocab_size; diff --git a/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cu b/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cu index 77e3d2d58..518b46058 100644 --- a/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cu +++ b/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cu @@ -3,12 +3,16 @@ #include "../cuda/kernel.cuh" #include "cross_entropy_nvidia.cuh" -template +template INFINIOP_CUDA_KERNEL crossEntropy( - Tdata *y, const Tdata *x, const void *target, + Tout *y, const Tdata *x, const void *target, size_t outer_size, size_t vocab_size, ptrdiff_t x_stride) { - crossEntropyKernel( + crossEntropyKernel( y, x, target, outer_size, vocab_size, x_stride); } @@ -29,11 +33,21 @@ infiniStatus_t Descriptor::create( infiniopTensorDescriptor_t x_desc, infiniopTensorDescriptor_t target_desc) { + auto y_dtype = y_desc->dtype(); auto x_dtype = x_desc->dtype(); auto t_dtype = target_desc->dtype(); - CrossEntropyInfo info; + CHECK_DTYPE(x_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + CHECK_DTYPE(y_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); + CHECK_DTYPE(t_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); + + if (y_dtype != x_dtype && y_dtype != INFINI_DTYPE_F32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + CrossEntropyInfo info{}; info.dtype = x_dtype; + info.output_dtype = y_dtype; info.target_dtype = t_dtype; info.vocab_size = x_desc->shape().back(); @@ -48,40 +62,65 @@ infiniStatus_t Descriptor::create( return INFINI_STATUS_SUCCESS; } +template +infiniStatus_t launchTypedKernel(void *y, const void *x, const void *target, + const CrossEntropyInfo &info, cudaStream_t stream) { + dim3 grid(static_cast(info.outer_size), 1, 1); + if (info.output_dtype == INFINI_DTYPE_F32) { + crossEntropy + <<>>( + (float *)y, + (const Tdata *)x, + target, + info.outer_size, + info.vocab_size, + info.x_stride); + } else if (info.output_dtype == info.dtype) { + crossEntropy + <<>>( + (Tdata *)y, + (const Tdata *)x, + target, + info.outer_size, + info.vocab_size, + info.x_stride); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + return INFINI_STATUS_SUCCESS; +} + template infiniStatus_t launchKernel(void *y, const void *x, const void *target, const CrossEntropyInfo &info, cudaStream_t stream) { - - dim3 grid(static_cast(info.outer_size), 1, 1); - if (info.target_dtype == INFINI_DTYPE_I64) { if (info.dtype == INFINI_DTYPE_F16) { - crossEntropy - <<>>((half *)y, (const half *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_BF16) { - crossEntropy - <<>>((__nv_bfloat16 *)y, (const __nv_bfloat16 *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_F32) { - crossEntropy - <<>>((float *)y, (const float *)x, target, info.outer_size, info.vocab_size, info.x_stride); + return launchTypedKernel( + y, x, target, info, stream); + } + if (info.dtype == INFINI_DTYPE_BF16) { + return launchTypedKernel( + y, x, target, info, stream); + } + if (info.dtype == INFINI_DTYPE_F32) { + return launchTypedKernel( + y, x, target, info, stream); } } else if (info.target_dtype == INFINI_DTYPE_I32) { - if (info.dtype == INFINI_DTYPE_F16) { - crossEntropy - <<>>((half *)y, (const half *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_BF16) { - crossEntropy - <<>>((__nv_bfloat16 *)y, (const __nv_bfloat16 *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_F32) { - crossEntropy - <<>>((float *)y, (const float *)x, target, info.outer_size, info.vocab_size, info.x_stride); + return launchTypedKernel( + y, x, target, info, stream); + } + if (info.dtype == INFINI_DTYPE_BF16) { + return launchTypedKernel( + y, x, target, info, stream); + } + if (info.dtype == INFINI_DTYPE_F32) { + return launchTypedKernel( + y, x, target, info, stream); } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; } - - return INFINI_STATUS_SUCCESS; + return INFINI_STATUS_BAD_TENSOR_DTYPE; } infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, diff --git a/test/infiniop/cross_entropy.py b/test/infiniop/cross_entropy.py index 987f2d11a..44ea84aee 100644 --- a/test/infiniop/cross_entropy.py +++ b/test/infiniop/cross_entropy.py @@ -11,6 +11,7 @@ get_tolerance, profile_operation, TestWorkspace, + InfiniDeviceEnum, InfiniDtype, InfiniDtypeNames, InfiniDeviceNames, @@ -27,29 +28,52 @@ ] _TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.BF16, InfiniDtype.F32] +_MIXED_OUTPUT_DTYPES = [InfiniDtype.F16, InfiniDtype.BF16] +# Hygon dispatches this operator through the same CUDA implementation as NVIDIA. +_MIXED_OUTPUT_DEVICES = {InfiniDeviceEnum.NVIDIA, InfiniDeviceEnum.HYGON} +_MIXED_OUTPUT_TEST_CASES_ = [((2, 4, 10), None, None)] _TOLERANCE_MAP = { InfiniDtype.F16: {"atol": 1e-3, "rtol": 1e-2}, InfiniDtype.BF16: {"atol": 1e-2, "rtol": 2e-2}, InfiniDtype.F32: {"atol": 1e-5, "rtol": 1e-5}, } +_TORCH_DTYPE_MAP = { + InfiniDtype.F16: torch.float16, + InfiniDtype.BF16: torch.bfloat16, + InfiniDtype.F32: torch.float32, +} # ------------------------------------------------------------ # PyTorch 参考实现 # ------------------------------------------------------------ -def cross_entropy_ref(logits, target): +def cross_entropy_ref(logits, target, output_dtype): vocab = logits.shape[-1] logits_flat = logits.reshape(-1, vocab).float() target_flat = target.reshape(-1).long() loss = torch.nn.functional.cross_entropy(logits_flat, target_flat, reduction="none") - return loss.view(target.shape).to(logits.dtype) - - -def test(handle, device, shape, x_stride=None, y_stride=None, dtype=InfiniDtype.F16, sync=None): + return loss.view(target.shape).to(_TORCH_DTYPE_MAP[output_dtype]) + + +def test( + handle, + device, + shape, + x_stride=None, + y_stride=None, + dtype=InfiniDtype.F16, + sync=None, + output_dtype=None, +): logits_shape = shape label_shape = shape[:-1] vocab = shape[-1] + output_dtype = dtype if output_dtype is None else output_dtype - print(f"Testing CrossEntropy on {InfiniDeviceNames[device]} logits:{logits_shape} dtype:{InfiniDtypeNames[dtype]}") + print( + f"Testing CrossEntropy on {InfiniDeviceNames[device]} " + f"logits:{logits_shape} input_dtype:{InfiniDtypeNames[dtype]} " + f"output_dtype:{InfiniDtypeNames[output_dtype]}" + ) x = TestTensor(logits_shape, x_stride, dtype, device) target = TestTensor(label_shape, None, InfiniDtype.I64, device) @@ -59,8 +83,10 @@ def test(handle, device, shape, x_stride=None, y_stride=None, dtype=InfiniDtype. tgt.copy_(torch.randint(0, vocab, label_shape, dtype=torch.int64, device=tgt.device)) target.actual_tensor().copy_(tgt) - reference = cross_entropy_ref(x.torch_tensor(), target.torch_tensor()) - y = TestTensor(label_shape, y_stride, dtype, device) + reference = cross_entropy_ref( + x.torch_tensor(), target.torch_tensor(), output_dtype + ) + y = TestTensor(label_shape, y_stride, output_dtype, device) descriptor = infiniopOperatorDescriptor_t() check_error( @@ -99,8 +125,36 @@ def run(): check_error(LIBINFINIOP.infiniopDestroyCrossEntropyDescriptor(descriptor)) +def test_mixed_output( + handle, + device, + shape, + x_stride=None, + y_stride=None, + dtype=InfiniDtype.F16, + sync=None, +): + test( + handle, + device, + shape, + x_stride, + y_stride, + dtype, + sync, + output_dtype=InfiniDtype.F32, + ) + + if __name__ == "__main__": args = get_args() for device in get_test_devices(args): test_operator(device, test, _TEST_CASES_, _TENSOR_DTYPES) + if device in _MIXED_OUTPUT_DEVICES: + test_operator( + device, + test_mixed_output, + _MIXED_OUTPUT_TEST_CASES_, + _MIXED_OUTPUT_DTYPES, + ) print("\033[92mTest passed!\033[0m") From e3aeb11e8d342f3ae51bd89832dd2158216d14ef Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 24 Jul 2026 10:24:55 +0800 Subject: [PATCH 10/16] perf(hygon): fuse MoE align padding fill --- .../ops/moe_align/hygon/moe_align_lightop_hygon.cc | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/infinicore/ops/moe_align/hygon/moe_align_lightop_hygon.cc b/src/infinicore/ops/moe_align/hygon/moe_align_lightop_hygon.cc index 6748cdd93..67c2a63ea 100644 --- a/src/infinicore/ops/moe_align/hygon/moe_align_lightop_hygon.cc +++ b/src/infinicore/ops/moe_align/hygon/moe_align_lightop_hygon.cc @@ -86,10 +86,6 @@ void run(void *planned_meta) { auto expert_ids = infinicore::adaptor::to_aten_tensor(p->expert_ids); auto num_tokens_post_padded = infinicore::adaptor::to_aten_tensor(p->num_tokens_post_padded); - if (p->pad_sorted_token_ids) { - sorted_token_ids.fill_(topk_ids.numel()); - } - const std::optional none = std::nullopt; infinicore::adaptor::lightop::moe_align_block_size( topk_ids, @@ -102,7 +98,7 @@ void run(void *planned_meta) { none, none, false, - false); + p->pad_sorted_token_ids); } void cleanup(void **planned_meta_ptr) { From 26648e49138c161927a8f172b52e7ffe08e1eca5 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 24 Jul 2026 11:33:04 +0800 Subject: [PATCH 11/16] perf(hygon): add LightOP paged KV attention path --- .../infinicore/adaptor/lightop_adaptor.hpp | 12 +++ .../flash_attn/hygon/flash_attn_hygon.cc | 38 ++++++++ .../flash_attn/hygon/flash_attn_hygon.hpp | 17 ++++ src/infinicore/adaptor/lightop_adaptor.cc | 39 ++++++++ .../hygon/mha_kvcache_flashattn_hygon.cc | 52 ++++++++++- .../hygon/mha_varlen_flashattn_hygon.cc | 53 +++++++---- .../paged_caching/paged_caching_infiniop.cc | 89 +++++++++++++++++-- 7 files changed, 273 insertions(+), 27 deletions(-) diff --git a/include/infinicore/adaptor/lightop_adaptor.hpp b/include/infinicore/adaptor/lightop_adaptor.hpp index 37b158d60..9c7f7ba4d 100644 --- a/include/infinicore/adaptor/lightop_adaptor.hpp +++ b/include/infinicore/adaptor/lightop_adaptor.hpp @@ -41,6 +41,8 @@ void preload_silu_and_mul(); void preload_rms_rotary_embedding(); +void preload_reshape_and_cache_cuda(); + void fuse_silu_and_mul( at::Tensor &input, at::Tensor &output); @@ -58,6 +60,16 @@ void rms_rotary_embedding_fuse( const std::optional &k_bias = std::nullopt, double epsilon = 1e-6); +void reshape_and_cache_cuda( + at::Tensor &key, + at::Tensor &value, + at::Tensor &key_cache, + at::Tensor &value_cache, + at::Tensor &slot_mapping, + const std::string &kv_cache_dtype, + at::Tensor &k_scale, + at::Tensor &v_scale); + void moe_sum( at::Tensor &input, at::Tensor &output, diff --git a/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.cc b/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.cc index a26b400ba..64bda76b1 100644 --- a/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.cc +++ b/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.cc @@ -33,6 +33,22 @@ using mha_fwd_kvcache_fn_t = std::vector (*)( int num_splits, const std::optional &s_aux_); +using paged_attention_fn_t = void (*)( + at::Tensor &out, + at::Tensor &q, + at::Tensor &k_cache, + at::Tensor &v_cache, + float scale, + at::Tensor &block_table, + at::Tensor &cache_lens, + const std::optional &alibi_slopes, + const std::string &kv_cache_dtype, + const std::optional &q_descale, + const std::optional &k_descale, + const std::optional &v_descale, + int max_context_len, + const std::optional &s_aux); + using mha_varlen_fwd_fn_t = std::vector (*)( at::Tensor &q, const at::Tensor &k, @@ -134,6 +150,28 @@ mha_fwd_kvcache(at::Tensor &q, softcap, is_rotary_interleaved, num_splits, s_aux); } +void paged_attention( + at::Tensor &out, + at::Tensor &q, + at::Tensor &k_cache, + at::Tensor &v_cache, + float scale, + at::Tensor &block_table, + at::Tensor &cache_lens, + const std::optional &alibi_slopes, + const std::string &kv_cache_dtype, + const std::optional &q_descale, + const std::optional &k_descale, + const std::optional &v_descale, + int max_context_len, + const std::optional &s_aux) { + static auto fn = reinterpret_cast( + resolve_flash_extension_symbol("paged_attention")); + fn(out, q, k_cache, v_cache, scale, block_table, cache_lens, + alibi_slopes, kv_cache_dtype, q_descale, k_descale, v_descale, + max_context_len, s_aux); +} + std::vector vllm_mha_varlen_fwd(at::Tensor &q, const at::Tensor &k, diff --git a/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.hpp b/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.hpp index 8600c8711..1bc8985eb 100644 --- a/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.hpp +++ b/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.hpp @@ -5,6 +5,7 @@ #include #include +#include #include namespace flash { @@ -31,6 +32,22 @@ mha_fwd_kvcache(at::Tensor &q, bool is_rotary_interleaved, int num_splits); +void paged_attention( + at::Tensor &out, + at::Tensor &q, + at::Tensor &k_cache, + at::Tensor &v_cache, + float scale, + at::Tensor &block_table, + at::Tensor &cache_lens, + const std::optional &alibi_slopes, + const std::string &kv_cache_dtype, + const std::optional &q_descale, + const std::optional &k_descale, + const std::optional &v_descale, + int max_context_len, + const std::optional &s_aux); + std::vector vllm_mha_varlen_fwd(at::Tensor &q, const at::Tensor &k, diff --git a/src/infinicore/adaptor/lightop_adaptor.cc b/src/infinicore/adaptor/lightop_adaptor.cc index f2178cbb0..72a7c8b23 100644 --- a/src/infinicore/adaptor/lightop_adaptor.cc +++ b/src/infinicore/adaptor/lightop_adaptor.cc @@ -34,6 +34,8 @@ constexpr const char *kFuseSiluAndMulSymbol = "_ZN2at6native17fuse_silu_and_mulERNS_6TensorES2_"; constexpr const char *kRmsRotaryEmbeddingFuseSymbol = "_ZN2at6native25rms_rotary_embedding_fuseERNS_6TensorES2_S2_lS2_bS1_S1_St8optionalIS1_ES4_d"; +constexpr const char *kReshapeAndCacheCudaSymbol = + "_ZN2at6native22reshape_and_cache_cudaERNS_6TensorES2_S2_S2_S2_RKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEES2_S2_"; constexpr const char *kMoeSumSymbol = "_ZN2at6native7moe_sumERNS_6TensorES2_RKSt8optionalIS1_ES6_S6_fi"; constexpr const char *kMoeAlignBlockSizeSymbol = @@ -776,6 +778,15 @@ using RmsRotaryEmbeddingFuseFn = void (*)( std::optional, std::optional, double); +using ReshapeAndCacheCudaFn = void (*)( + at::Tensor &, + at::Tensor &, + at::Tensor &, + at::Tensor &, + at::Tensor &, + const std::string &, + at::Tensor &, + at::Tensor &); using MoeSumFn = void (*)( at::Tensor &, at::Tensor &, @@ -872,6 +883,11 @@ RmsRotaryEmbeddingFuseFn rms_rotary_embedding_fuse_fn() { return fn; } +ReshapeAndCacheCudaFn reshape_and_cache_cuda_fn() { + static auto fn = resolve(kReshapeAndCacheCudaSymbol); + return fn; +} + MoeSumFn moe_sum_fn() { static auto fn = resolve(kMoeSumSymbol); return fn; @@ -976,6 +992,10 @@ void preload_rms_rotary_embedding() { (void)rms_rotary_embedding_fuse_fn(); } +void preload_reshape_and_cache_cuda() { + (void)reshape_and_cache_cuda_fn(); +} + void preload_w8a8_linear_ops() { (void)per_token_dynamic_quant_int8_fn(); (void)blaslt_w8a8_bf16_fn(); @@ -1011,6 +1031,25 @@ void rms_rotary_embedding_fuse(at::Tensor &positions, epsilon); } +void reshape_and_cache_cuda(at::Tensor &key, + at::Tensor &value, + at::Tensor &key_cache, + at::Tensor &value_cache, + at::Tensor &slot_mapping, + const std::string &kv_cache_dtype, + at::Tensor &k_scale, + at::Tensor &v_scale) { + reshape_and_cache_cuda_fn()( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale); +} + void moe_sum(at::Tensor &input, at::Tensor &output, const std::optional &bias, diff --git a/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc b/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc index 095f06e74..621a8bffa 100644 --- a/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc +++ b/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc @@ -6,6 +6,10 @@ #include +#include +#include +#include + namespace infinicore::op::mha_kvcache_impl::flashattn { struct PlannedMeta { @@ -44,12 +48,56 @@ void run(void *planned_meta) { auto q = infinicore::adaptor::to_aten_tensor(p->q); auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache); auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache); - auto seqlens_k = std::optional(infinicore::adaptor::to_aten_tensor(p->seqlens_k)); - auto block_table = std::optional(infinicore::adaptor::to_aten_tensor(p->block_table)); + auto seqlens_k_tensor = infinicore::adaptor::to_aten_tensor(p->seqlens_k); + auto block_table_tensor = infinicore::adaptor::to_aten_tensor(p->block_table); + auto seqlens_k = std::optional(seqlens_k_tensor); + auto block_table = std::optional(block_table_tensor); auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; + const bool use_paged_attention = + q.dim() == 4 && q.size(1) == 1 + && k_cache.dim() == 4 && v_cache.dim() == 4 + && k_cache.size(0) == v_cache.size(0) + && k_cache.size(1) == v_cache.size(1) + && k_cache.size(2) == v_cache.size(3) + && k_cache.size(3) == v_cache.size(2) + && k_cache.size(2) == 64 + && q.size(2) % k_cache.size(1) == 0 + && q.size(3) == k_cache.size(3) + && q.is_contiguous() && k_cache.is_contiguous() && v_cache.is_contiguous() + && seqlens_k_tensor.dim() == 1 && block_table_tensor.dim() == 2 + && !alibi_slopes.has_value(); + if (use_paged_attention) { + const auto max_context_len_64 = block_table_tensor.size(1) * k_cache.size(2); + if (max_context_len_64 > std::numeric_limits::max()) { + throw std::runtime_error("paged_attention max context length exceeds int range"); + } + auto paged_out = out_tensor.view({q.size(0), q.size(2), q.size(3)}); + const std::optional none = std::nullopt; + static const std::string kv_cache_dtype = "auto"; + flash::paged_attention( + paged_out, + q, + k_cache, + v_cache, + p->scale, + block_table_tensor, + seqlens_k_tensor, + none, + kv_cache_dtype, + none, + none, + none, + static_cast(max_context_len_64), + none); + if (out_need_copy_back) { + p->out->copy_from(out_work); + } + return; + } + std::optional k_new = std::nullopt; std::optional v_new = std::nullopt; std::optional rotary_cos = std::nullopt; diff --git a/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc b/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc index b55046ce7..331f99e21 100644 --- a/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc +++ b/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc @@ -125,25 +125,40 @@ void run(void *planned_meta) { auto k_work = k.contiguous(); auto v_work = v.contiguous(); if (block_table.has_value() && k.dim() == 4 && v.dim() == 4) { - const int64_t num_blocks = k.size(0); - const int64_t block_size = k.size(1); - const int64_t num_kv_heads = k.size(2); - const int64_t head_dim = k.size(3); - if (block_size % 64 != 0) { - throw std::runtime_error("[mha_varlen/hygon] flash-attn requires paged KV block size to be divisible by 64"); - } - const int64_t pages_per_block = block_size / 64; - k_work = k_work.reshape({num_blocks, pages_per_block, 64, num_kv_heads, head_dim}) - .reshape({num_blocks * pages_per_block, 64, num_kv_heads, head_dim}) - .contiguous(); - v_work = v_work.reshape({num_blocks, pages_per_block, 64, num_kv_heads, head_dim}) - .reshape({num_blocks * pages_per_block, 64, num_kv_heads, head_dim}) - .contiguous(); - if (pages_per_block != 1) { - auto offsets = at::arange(pages_per_block, block_table->options()).view({1, 1, pages_per_block}); - block_table = ((*block_table).unsqueeze(-1) * pages_per_block + offsets) - .reshape({block_table->size(0), block_table->size(1) * pages_per_block}) - .contiguous(); + const bool vllm_cache_layout = + k.size(0) == v.size(0) + && k.size(1) == v.size(1) + && k.size(2) == v.size(3) + && k.size(3) == v.size(2); + if (vllm_cache_layout) { + if (k.size(2) != 64) { + throw std::runtime_error("[mha_varlen/hygon] vLLM cache layout requires block size 64"); + } + // LightOP paged attention stores K/V as BHSD/BHDS, while the + // flash-attn varlen prefill ABI consumes BSHD for both caches. + k_work = k.permute({0, 2, 1, 3}).contiguous(); + v_work = v.permute({0, 3, 1, 2}).contiguous(); + } else { + const int64_t num_blocks = k.size(0); + const int64_t block_size = k.size(1); + const int64_t num_kv_heads = k.size(2); + const int64_t head_dim = k.size(3); + if (block_size % 64 != 0) { + throw std::runtime_error("[mha_varlen/hygon] flash-attn requires paged KV block size to be divisible by 64"); + } + const int64_t pages_per_block = block_size / 64; + k_work = k_work.reshape({num_blocks, pages_per_block, 64, num_kv_heads, head_dim}) + .reshape({num_blocks * pages_per_block, 64, num_kv_heads, head_dim}) + .contiguous(); + v_work = v_work.reshape({num_blocks, pages_per_block, 64, num_kv_heads, head_dim}) + .reshape({num_blocks * pages_per_block, 64, num_kv_heads, head_dim}) + .contiguous(); + if (pages_per_block != 1) { + auto offsets = at::arange(pages_per_block, block_table->options()).view({1, 1, pages_per_block}); + block_table = ((*block_table).unsqueeze(-1) * pages_per_block + offsets) + .reshape({block_table->size(0), block_table->size(1) * pages_per_block}) + .contiguous(); + } } } auto result = flash::vllm_mha_varlen_fwd( diff --git a/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc b/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc index 5e8be049a..b87e65dc0 100644 --- a/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc +++ b/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc @@ -2,17 +2,66 @@ #include "../infiniop_impl.hpp" +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/lightop_adaptor.hpp" + +#include + +#include +#include +#endif + namespace infinicore::op::paged_caching_impl::infiniop { INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PagedCaching, 100); struct PlannedMeta { + bool use_hygon_lightop; std::shared_ptr descriptor; - - graph::GraphTensor workspace, k_cache, v_cache, k, v, slot_mapping; + std::optional workspace; + graph::GraphTensor k_cache, v_cache, k, v, slot_mapping; +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) + at::Tensor k_scale, v_scale; +#endif }; +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +bool is_hygon_vllm_cache_layout(const Tensor &k_cache, const Tensor &v_cache) { + if (k_cache->device().getType() != Device::Type::HYGON) { + return false; + } + const auto &k_shape = k_cache->shape(); + const auto &v_shape = v_cache->shape(); + return k_shape.size() == 4 && v_shape.size() == 4 + && k_shape[0] == v_shape[0] + && k_shape[1] == v_shape[1] + && k_shape[2] == v_shape[3] + && k_shape[3] == v_shape[2] + && k_shape[2] == 64; +} +#endif + void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) + if (is_hygon_vllm_cache_layout(k_cache, v_cache)) { + infinicore::adaptor::lightop::preload_reshape_and_cache_cuda(); + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto options = infinicore::adaptor::to_aten_tensor(k).options().dtype(at::kFloat); + return new PlannedMeta{ + true, + nullptr, + std::nullopt, + graph::GraphTensor(k_cache), + graph::GraphTensor(v_cache), + graph::GraphTensor(k), + graph::GraphTensor(v), + graph::GraphTensor(slot_mapping), + at::ones({1}, options), + at::ones({1}, options)}; + } +#endif + size_t key = hash_combine(k_cache, v_cache, k, v, slot_mapping); INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( @@ -22,23 +71,51 @@ void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, con INFINIOP_WORKSPACE_TENSOR(workspace, PagedCaching, descriptor); return new PlannedMeta{ + false, descriptor, graph::GraphTensor(workspace), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), - graph::GraphTensor(slot_mapping)}; + graph::GraphTensor(slot_mapping) +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) + , at::Tensor{}, at::Tensor{} +#endif + }; } void run(void *planned_meta) { auto *p = reinterpret_cast(planned_meta); +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) + if (p->use_hygon_lightop) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + auto k = infinicore::adaptor::to_aten_tensor(p->k); + auto v = infinicore::adaptor::to_aten_tensor(p->v); + auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache); + auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache); + auto slot_mapping = infinicore::adaptor::to_aten_tensor(p->slot_mapping); + static const std::string kv_cache_dtype = "auto"; + infinicore::adaptor::lightop::reshape_and_cache_cuda( + k, + v, + k_cache, + v_cache, + slot_mapping, + kv_cache_dtype, + p->k_scale, + p->v_scale); + return; + } +#endif + + auto &workspace = p->workspace.value(); INFINICORE_CHECK_ERROR( infiniopPagedCaching( p->descriptor->desc, - p->workspace->data(), - p->workspace->numel(), + workspace->data(), + workspace->numel(), p->k_cache->data(), p->v_cache->data(), p->k->data(), @@ -54,4 +131,4 @@ void cleanup(void **planned_meta_ptr) { INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(PagedCaching, &plan, &run, &cleanup); -} // namespace infinicore::op::paged_caching_impl::infiniop +} // namespace infinicore::op::paged_caching_impl::infiniop \ No newline at end of file From 40b11dbc3db1655321c05c8eb46d149571c51636 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 27 Jul 2026 14:39:40 +0800 Subject: [PATCH 12/16] fix(hygon): stabilize threaded graph inference --- .../infinicore/ops/distributed/allreduce.hpp | 1 + src/infiniccl/cuda/infiniccl_cuda.cu | 53 +++++++++++++++---- src/infinicore/graph/graph.cc | 50 +++++++++++++++-- src/infinicore/ops/distributed/allreduce.cc | 5 ++ 4 files changed, 94 insertions(+), 15 deletions(-) diff --git a/include/infinicore/ops/distributed/allreduce.hpp b/include/infinicore/ops/distributed/allreduce.hpp index 39f74243a..d2fd25600 100644 --- a/include/infinicore/ops/distributed/allreduce.hpp +++ b/include/infinicore/ops/distributed/allreduce.hpp @@ -12,6 +12,7 @@ class AllReduce : public graph::GraphOperator { AllReduce(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); ~AllReduce(); void run() const override; + bool supports_device_graph_capture() const override; static void execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); private: diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu index ba8534e66..24a3b78d2 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ b/src/infiniccl/cuda/infiniccl_cuda.cu @@ -226,6 +226,7 @@ struct HygonVmmAllocation { void *ptr = nullptr; size_t size = 0; CUmemGenericAllocationHandle handle = 0; + bool hip_uncached = false; }; struct HygonTp2AllReduceState { @@ -251,13 +252,19 @@ struct HygonTp2AllReduceState { if (rank_data[rank] != nullptr) cudaFree(rank_data[rank]); if (signal_hosts[rank] != nullptr) cudaFreeHost(signal_hosts[rank]); auto &driver = hygon_cuda_driver_api(); - if (driver.available && stages[rank].ptr != nullptr) { + if (driver.available && stages[rank].handle != 0) { const auto address = reinterpret_cast(stages[rank].ptr); driver.mem_unmap(address, stages[rank].size); driver.mem_address_free(address, stages[rank].size); - } - if (driver.available && stages[rank].handle != 0) { driver.mem_release(stages[rank].handle); + } else if (stages[rank].hip_uncached && + stages[rank].ptr != nullptr) { + auto &hip = hygon_hip_ext_api(); + if (hip.available) { + hip.free(stages[rank].ptr); + } + } else if (stages[rank].ptr != nullptr) { + cudaFree(stages[rank].ptr); } } if (restore_device) cudaSetDevice(previous_device); @@ -450,10 +457,8 @@ __global__ __launch_bounds__(512, 1) void hygon_tp2_bf16_allreduce_kernel( } __syncthreads(); } - if (pack_offset == 0) { - hygon_tp2_end_sync( - rank_signals, self_signal, rank, sync_flag); - } + hygon_tp2_end_sync( + rank_signals, self_signal, rank, sync_flag); } std::shared_ptr create_hygon_tp2_state( @@ -470,9 +475,37 @@ std::shared_ptr create_hygon_tp2_state( }; const size_t stage_bytes = kHygonTp2StageCapacityElements * sizeof(__nv_bfloat16); for (int rank = 0; rank < 2; ++rank) { - if (cudaSetDevice(device_ids[rank]) != cudaSuccess || - !allocate_hygon_vmm(state->stages[rank], device_ids[rank], - state->device_ids, stage_bytes)) return fail(); + if (cudaSetDevice(device_ids[rank]) != cudaSuccess) return fail(); + const int peer = 1 - rank; + int can_access_peer = 0; + if (cudaDeviceCanAccessPeer( + &can_access_peer, + device_ids[rank], + device_ids[peer]) != cudaSuccess || + can_access_peer == 0) { + return fail(); + } + const cudaError_t enable_status = + cudaDeviceEnablePeerAccess(device_ids[peer], 0); + if (enable_status == cudaErrorPeerAccessAlreadyEnabled) { + (void)cudaGetLastError(); + } else if (enable_status != cudaSuccess) { + return fail(); + } + auto &hip = hygon_hip_ext_api(); + if (!hip.available) return fail(); + state->stages[rank].size = stage_bytes; + if (hip.ext_malloc_with_flags( + &state->stages[rank].ptr, + stage_bytes, + kHygonHipDeviceMallocUncached) != kHygonHipSuccess) { + return fail(); + } + state->stages[rank].hip_uncached = true; + if (hip.memset( + state->stages[rank].ptr, + 0, + stage_bytes) != kHygonHipSuccess) return fail(); void *signal_host = nullptr; if (cudaHostAlloc(&signal_host, sizeof(HygonTp2Signal), cudaHostAllocMapped) != cudaSuccess) return fail(); diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc index 72f422adc..e39d87705 100644 --- a/src/infinicore/graph/graph.cc +++ b/src/infinicore/graph/graph.cc @@ -5,6 +5,17 @@ #include #include +#include + +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include +#include + +#include +#endif + namespace infinicore::graph { /* ========================= @@ -33,10 +44,13 @@ DispatchableGraphOperator::~DispatchableGraphOperator() { * ========================= */ struct Graph::DeviceGraph { - infinirtGraph_t graph; - infinirtGraphExec_t exec; - infinirtGraphNode_t node; + infinirtGraph_t graph = nullptr; + infinirtGraphExec_t exec = nullptr; + infinirtGraphNode_t node = nullptr; std::vector log_buffer; +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) + std::unique_ptr hygon_graph; +#endif DeviceGraph() : graph(nullptr), exec(nullptr), node(nullptr) { log_buffer.resize(4 * 1024); @@ -51,7 +65,23 @@ struct Graph::DeviceGraph { } } + bool is_instantiated() const { +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) + if (hygon_graph) { + return true; + } +#endif + return exec != nullptr; + } + void launch() { +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) + if (hygon_graph) { + c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); + hygon_graph->replay(); + return; + } +#endif INFINICORE_CHECK_ERROR(infinirtGraphLuanch(exec, context::getStream())); } }; @@ -95,11 +125,20 @@ void Graph::add_operator(std::shared_ptr op) { } void Graph::instantiate() { +#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN) + const bool is_hygon = + context::getDevice().getType() == Device::Type::HYGON; +#else + constexpr bool is_hygon = false; +#endif + segments_.clear(); // Warm the complete op list before splitting it into replay segments. for (size_t iter = 0; iter < 5; ++iter) { - this->run(); + for (const auto &op : op_list_) { + op->run(); + } } infinicore::context::syncStream(); @@ -109,9 +148,10 @@ void Graph::instantiate() { spdlog::info("device graph segments disabled; replaying recorded operators"); return; } +#endif for (const auto &op : op_list_) { - const bool capture_safe = op->is_device_graph_capture_safe(); + const bool capture_safe = is_hygon || op->is_device_graph_capture_safe(); if (segments_.empty() || segments_.back()->capture_safe != capture_safe) { segments_.push_back(std::make_unique(capture_safe)); } diff --git a/src/infinicore/ops/distributed/allreduce.cc b/src/infinicore/ops/distributed/allreduce.cc index ddfc238c9..7d140ac69 100644 --- a/src/infinicore/ops/distributed/allreduce.cc +++ b/src/infinicore/ops/distributed/allreduce.cc @@ -34,6 +34,11 @@ void AllReduce::run() const { infinicore::context::getStream())); } +bool AllReduce::supports_device_graph_capture() const { + const auto *meta = reinterpret_cast(planned_meta_); + return meta->input->device().getType() != Device::Type::HYGON; +} + void AllReduce::execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator) { INFINICORE_GRAPH_OP_RECORD_OR_RUN(AllReduce, output, input, op, communicator); } From a06372a7d368c81a0c087f4efe47d497a82fe4c9 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Tue, 28 Jul 2026 09:47:39 +0800 Subject: [PATCH 13/16] feat(hygon): expose optimized inference operators --- include/infinicore/ops.hpp | 4 + include/infinicore/ops/hygon_moe_marlin.hpp | 61 ++ include/infinicore/ops/moe_marlin_config.hpp | 42 ++ .../infinicore/ops/paged_flash_attention.hpp | 30 + .../ops/select_last_token_hidden_states.hpp | 16 + .../ops/hygon_moe_marlin/hygon_moe_marlin.cc | 648 ++++++++++++++++++ .../moe_marlin_config/moe_marlin_config.cc | 215 ++++++ .../paged_flash_attention.cc | 139 ++++ .../select_last_token_hidden_states.cc | 77 +++ 9 files changed, 1232 insertions(+) create mode 100644 include/infinicore/ops/hygon_moe_marlin.hpp create mode 100644 include/infinicore/ops/moe_marlin_config.hpp create mode 100644 include/infinicore/ops/paged_flash_attention.hpp create mode 100644 include/infinicore/ops/select_last_token_hidden_states.hpp create mode 100644 src/infinicore/ops/hygon_moe_marlin/hygon_moe_marlin.cc create mode 100644 src/infinicore/ops/moe_marlin_config/moe_marlin_config.cc create mode 100644 src/infinicore/ops/paged_flash_attention/paged_flash_attention.cc create mode 100644 src/infinicore/ops/select_last_token_hidden_states/select_last_token_hidden_states.cc diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index a95cc17e7..461856a73 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -44,6 +44,7 @@ #include "ops/gelutanh.hpp" #include "ops/hardswish.hpp" #include "ops/hardtanh.hpp" +#include "ops/hygon_moe_marlin.hpp" #include "ops/kv_caching.hpp" #include "ops/kimi_delta_attention.hpp" #include "ops/layer_norm.hpp" @@ -54,6 +55,7 @@ #include "ops/moe_align.hpp" #include "ops/moe_fused_dense.hpp" #include "ops/moe_fused_gate.hpp" +#include "ops/moe_marlin_config.hpp" #include "ops/moe_sum.hpp" #include "ops/moe_topk_sigmoid.hpp" #include "ops/moe_topk_softmax.hpp" @@ -64,6 +66,7 @@ #include "ops/paged_attention.hpp" #include "ops/paged_attention_prefill.hpp" #include "ops/paged_caching.hpp" +#include "ops/paged_flash_attention.hpp" #include "ops/per_tensor_dequant_i8.hpp" #include "ops/per_tensor_quant_i8.hpp" #include "ops/prepare_moe_input.hpp" @@ -83,6 +86,7 @@ #include "ops/rwkv5_wkv.hpp" #include "ops/scal.hpp" #include "ops/select_last_token_hidden.hpp" +#include "ops/select_last_token_hidden_states.hpp" #include "ops/sigmoid.hpp" #include "ops/silu.hpp" #include "ops/silu_and_mul.hpp" diff --git a/include/infinicore/ops/hygon_moe_marlin.hpp b/include/infinicore/ops/hygon_moe_marlin.hpp new file mode 100644 index 000000000..540c9fde7 --- /dev/null +++ b/include/infinicore/ops/hygon_moe_marlin.hpp @@ -0,0 +1,61 @@ +#pragma once + +#include "../tensor.hpp" + +#include + +namespace infinicore::op { + +enum class HygonMoeMarlinWeightFormat { + W16A16, + W8A8, +}; + +struct HygonMoeMarlinWeights { + Tensor packed_w13; + Tensor packed_w2; + Tensor packed_w13_scale; + Tensor packed_w2_scale; + HygonMoeMarlinWeightFormat format = + HygonMoeMarlinWeightFormat::W16A16; +}; + +struct HygonMoeMarlinWorkspace { + Tensor output; + Tensor cache13; + Tensor cache2; + Tensor input_i8; + Tensor input_scale; + Tensor cache2_i8; + Tensor cache2_scale; + Tensor sorted_token_ids; + Tensor expert_ids; + Tensor num_tokens_post_padded; + + size_t cache13_capacity = 0; + size_t cache2_capacity = 0; + size_t sorted_token_ids_capacity = 0; + size_t expert_ids_capacity = 0; +}; + +struct HygonMoeMarlinOutput { + Tensor hidden_states; + Tensor sorted_token_ids; + Tensor expert_ids; + Tensor num_tokens_post_padded; + bool has_routing_metadata = false; +}; + +HygonMoeMarlinOutput hygon_moe_marlin_fused( + const Tensor &hidden_states, + const Tensor &topk_weights, + const Tensor &topk_ids, + const Tensor &expert_map, + const HygonMoeMarlinWeights &weights, + HygonMoeMarlinWorkspace &workspace, + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size, + size_t fallback_align_block_size); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/moe_marlin_config.hpp b/include/infinicore/ops/moe_marlin_config.hpp new file mode 100644 index 000000000..e3bd3b40f --- /dev/null +++ b/include/infinicore/ops/moe_marlin_config.hpp @@ -0,0 +1,42 @@ +#pragma once + +#include "../device.hpp" +#include "../dtype.hpp" + +#include + +namespace infinicore::op { + +struct HygonMarlinGemmConfig { + int mode = 103; + int delta = 1; + size_t block_size_m = 16; + bool found = false; +}; + +struct HygonW16A16MarlinRuntimeConfig { + HygonMarlinGemmConfig gemm1; + HygonMarlinGemmConfig gemm2; + bool supported = false; +}; + +struct HygonW8A8MarlinRuntimeConfig { + HygonMarlinGemmConfig gemm1; + HygonMarlinGemmConfig gemm2; + bool supported = false; +}; + +HygonW16A16MarlinRuntimeConfig select_hygon_w16a16_marlin_config( + size_t num_tokens, + size_t hidden_size, + size_t intermediate_size, + DataType hidden_dtype, + size_t device_index); + +HygonW8A8MarlinRuntimeConfig select_hygon_w8a8_marlin_config( + size_t num_tokens, + size_t hidden_size, + size_t intermediate_size, + size_t device_index); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/paged_flash_attention.hpp b/include/infinicore/ops/paged_flash_attention.hpp new file mode 100644 index 000000000..535f4e95c --- /dev/null +++ b/include/infinicore/ops/paged_flash_attention.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include "../tensor.hpp" + +#include +#include + +namespace infinicore::op { + +// Update a paged KV cache and run prefill or decode FlashAttention. +// +// On Hygon, this function also owns the LightOP cache-layout policy and +// serializes graph capture across TP threads because the vendor extension +// keeps process-global launch state. +Tensor paged_flash_attention( + const Tensor &query, + const Tensor &key, + const Tensor &value, + const Tensor &kv_cache, + const Tensor &total_sequence_lengths, + const std::optional &input_offsets, + const std::optional &cu_seqlens, + const Tensor &block_tables, + const Tensor &slot_mapping, + size_t num_heads, + size_t num_kv_heads, + size_t head_dim, + float scale); + +} // namespace infinicore::op diff --git a/include/infinicore/ops/select_last_token_hidden_states.hpp b/include/infinicore/ops/select_last_token_hidden_states.hpp new file mode 100644 index 000000000..675e3191d --- /dev/null +++ b/include/infinicore/ops/select_last_token_hidden_states.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "../tensor.hpp" + +namespace infinicore::op { + +// Select the final token of every packed request. +// +// hidden_states: [batch, tokens, hidden_size] +// input_offsets: [num_requests + 1], I32 +// result: [1, num_requests, hidden_size] +Tensor select_last_token_hidden_states( + const Tensor &hidden_states, + const Tensor &input_offsets); + +} // namespace infinicore::op diff --git a/src/infinicore/ops/hygon_moe_marlin/hygon_moe_marlin.cc b/src/infinicore/ops/hygon_moe_marlin/hygon_moe_marlin.cc new file mode 100644 index 000000000..974f8910b --- /dev/null +++ b/src/infinicore/ops/hygon_moe_marlin/hygon_moe_marlin.cc @@ -0,0 +1,648 @@ +#include "infinicore/ops/hygon_moe_marlin.hpp" + +#include "infinicore/context/context.hpp" +#include "infinicore/ops/moe_align.hpp" +#include "infinicore/ops/moe_marlin_config.hpp" +#include "infinicore/ops/moe_w16a16_marlin.hpp" +#include "infinicore/ops/moe_w8a8_marlin.hpp" + +#include +#include +#include +#include + +namespace infinicore::op { +namespace { + +constexpr size_t kHygonMoeSliceTokens = 16384; + +struct RoutingMetadata { + Tensor sorted_token_ids; + Tensor expert_ids; + Tensor num_tokens_post_padded; +}; + +bool same_device(const Tensor &tensor, const Device &device) { + return tensor + && tensor->device().getType() == device.getType() + && tensor->device().getIndex() == device.getIndex(); +} + +void ensure_tensor( + Tensor &tensor, + const Shape &shape, + DataType dtype, + const Device &device, + const char *name) { + if (!same_device(tensor, device) + || tensor->dtype() != dtype + || tensor->shape() != shape) { + if (context::isGraphRecording()) { + throw std::runtime_error( + std::string("Hygon MoE Marlin ") + name + + " workspace was not initialized before graph capture"); + } + tensor = Tensor::empty(shape, dtype, device); + } +} + +std::string shape_to_string(const Shape &shape) { + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < shape.size(); ++i) { + if (i != 0) { + oss << ", "; + } + oss << shape[i]; + } + oss << "]"; + return oss.str(); +} + +void check_packed_weight_tensor( + const Tensor &tensor, + const std::string &name, + const Device &device, + DataType dtype, + const Shape &shape) { + if (!tensor) { + throw std::runtime_error( + "Hygon MoE Marlin requires " + name); + } + if (tensor->device().getType() != device.getType() + || tensor->device().getIndex() != device.getIndex()) { + throw std::runtime_error( + "Hygon MoE Marlin requires packed weights on the hidden_states device"); + } + if (tensor->dtype() != dtype) { + throw std::runtime_error( + "Hygon MoE Marlin packed tensor dtype mismatch for " + name); + } + if (tensor->shape() != shape) { + throw std::runtime_error( + "Hygon MoE Marlin packed weight shape mismatch for " + name + + ": expected " + shape_to_string(shape) + + ", got " + shape_to_string(tensor->shape())); + } +} + +RoutingMetadata prepare_routing( + const Tensor &topk_ids, + const Tensor &expert_map, + HygonMoeMarlinWorkspace &workspace, + size_t num_local_experts, + size_t block_size) { + const auto &topk_shape = topk_ids->shape(); + if (topk_shape.size() != 2) { + throw std::runtime_error( + "Hygon MoE Marlin requires topk_ids [M, top_k]"); + } + const size_t num_pairs = topk_shape[0] * topk_shape[1]; + const size_t align_num_experts = num_local_experts + 1; + const size_t max_num_tokens_padded = + num_pairs < align_num_experts + ? num_pairs * block_size + : num_pairs + align_num_experts * (block_size - 1); + const size_t sorted_token_ids_capacity = + ((max_num_tokens_padded + 3) / 4) * 4; + const size_t max_num_blocks = + (max_num_tokens_padded + block_size - 1) / block_size; + const auto device = topk_ids->device(); + + if (!same_device(workspace.sorted_token_ids, device) + || workspace.sorted_token_ids_capacity + < sorted_token_ids_capacity) { + if (context::isGraphRecording()) { + throw std::runtime_error( + "Hygon MoE Marlin sorted_token_ids workspace was not initialized before graph capture"); + } + workspace.sorted_token_ids = Tensor::empty( + {sorted_token_ids_capacity}, + DataType::I32, + device); + workspace.sorted_token_ids_capacity = + sorted_token_ids_capacity; + } + if (!same_device(workspace.expert_ids, device) + || workspace.expert_ids_capacity < max_num_blocks) { + if (context::isGraphRecording()) { + throw std::runtime_error( + "Hygon MoE Marlin expert_ids workspace was not initialized before graph capture"); + } + workspace.expert_ids = Tensor::empty( + {max_num_blocks}, + DataType::I32, + device); + workspace.expert_ids_capacity = max_num_blocks; + } + if (!same_device(workspace.num_tokens_post_padded, device)) { + if (context::isGraphRecording()) { + throw std::runtime_error( + "Hygon MoE Marlin num_tokens_post_padded workspace was not initialized before graph capture"); + } + workspace.num_tokens_post_padded = Tensor::empty( + {1}, + DataType::I32, + device); + } + + auto sorted_token_ids = workspace.sorted_token_ids->narrow( + {{0, 0, sorted_token_ids_capacity}}); + auto expert_ids = workspace.expert_ids->narrow( + {{0, 0, max_num_blocks}}); + + if (expert_map) { + moe_align_with_expert_map_( + sorted_token_ids, + expert_ids, + workspace.num_tokens_post_padded, + topk_ids, + expert_map, + num_local_experts, + block_size, + true); + } else { + moe_align_( + sorted_token_ids, + expert_ids, + workspace.num_tokens_post_padded, + topk_ids, + num_local_experts, + block_size, + true); + } + + return RoutingMetadata{ + sorted_token_ids, + expert_ids, + workspace.num_tokens_post_padded, + }; +} + +Tensor run_w16a16( + const Tensor &hidden_states, + const Tensor &topk_weights, + const Tensor &topk_ids, + const RoutingMetadata &routing, + const HygonMoeMarlinWeights &weights, + HygonMoeMarlinWorkspace &workspace, + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size, + const HygonW16A16MarlinRuntimeConfig &config) { + const auto activation_dtype = hidden_states->dtype(); + if (activation_dtype != DataType::BF16 + && activation_dtype != DataType::F16) { + throw std::runtime_error( + "Hygon W16A16 Marlin MoE requires BF16 or FP16 activations"); + } + check_packed_weight_tensor( + weights.packed_w13, + "w13", + hidden_states->device(), + activation_dtype, + {num_local_experts, + hidden_size / 16, + intermediate_size * 2 * 16}); + check_packed_weight_tensor( + weights.packed_w2, + "w2", + hidden_states->device(), + activation_dtype, + {num_local_experts, + intermediate_size / 16, + hidden_size * 16}); + + const size_t top_k = topk_ids->shape()[1]; + const size_t num_tokens = hidden_states->shape()[0]; + if (num_tokens > kHygonMoeSliceTokens) { + throw std::runtime_error( + "Hygon W16A16 Marlin MoE inputs above 16384 tokens must be sliced"); + } + const size_t cache13_required = + num_tokens * top_k + * std::max(intermediate_size * 2, hidden_size); + const size_t cache2_required = + num_tokens * top_k * intermediate_size; + + ensure_tensor( + workspace.output, + hidden_states->shape(), + hidden_states->dtype(), + hidden_states->device(), + "output"); + if (!same_device(workspace.cache13, hidden_states->device()) + || workspace.cache13->dtype() != hidden_states->dtype() + || workspace.cache13_capacity < cache13_required) { + if (context::isGraphRecording()) { + throw std::runtime_error( + "Hygon W16A16 Marlin cache13 workspace was not initialized before graph capture"); + } + workspace.cache13 = Tensor::empty( + {cache13_required}, + hidden_states->dtype(), + hidden_states->device()); + workspace.cache13_capacity = cache13_required; + } + if (!same_device(workspace.cache2, hidden_states->device()) + || workspace.cache2->dtype() != hidden_states->dtype() + || workspace.cache2_capacity < cache2_required) { + if (context::isGraphRecording()) { + throw std::runtime_error( + "Hygon W16A16 Marlin cache2 workspace was not initialized before graph capture"); + } + workspace.cache2 = Tensor::empty( + {cache2_required}, + hidden_states->dtype(), + hidden_states->device()); + workspace.cache2_capacity = cache2_required; + } + + moe_w16a16_marlin_fused_dense_( + workspace.output, + workspace.cache13, + workspace.cache2, + hidden_states, + weights.packed_w13, + weights.packed_w2, + topk_weights, + routing.sorted_token_ids, + routing.expert_ids, + routing.num_tokens_post_padded, + top_k, + config.gemm1.mode, + config.gemm1.delta, + config.gemm2.mode, + config.gemm2.delta); + return workspace.output; +} + +Tensor run_w8a8( + const Tensor &hidden_states, + const Tensor &topk_weights, + const Tensor &topk_ids, + const RoutingMetadata &routing, + const HygonMoeMarlinWeights &weights, + HygonMoeMarlinWorkspace &workspace, + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size, + const HygonW8A8MarlinRuntimeConfig &config) { + check_packed_weight_tensor( + weights.packed_w13, + "w13", + hidden_states->device(), + DataType::I8, + {num_local_experts, + hidden_size / 64, + intermediate_size * 2 * 64}); + check_packed_weight_tensor( + weights.packed_w2, + "w2", + hidden_states->device(), + DataType::I8, + {num_local_experts, + intermediate_size / 64, + hidden_size * 64}); + check_packed_weight_tensor( + weights.packed_w13_scale, + "w13_scale", + hidden_states->device(), + DataType::F32, + {num_local_experts, intermediate_size * 2, 1}); + check_packed_weight_tensor( + weights.packed_w2_scale, + "w2_scale", + hidden_states->device(), + DataType::F32, + {num_local_experts, hidden_size, 1}); + + const size_t top_k = topk_ids->shape()[1]; + const size_t num_tokens = hidden_states->shape()[0]; + const size_t cache13_required = + num_tokens * top_k + * std::max(intermediate_size * 2, hidden_size); + + ensure_tensor( + workspace.output, + hidden_states->shape(), + hidden_states->dtype(), + hidden_states->device(), + "output"); + if (!same_device(workspace.cache13, hidden_states->device()) + || workspace.cache13->dtype() != hidden_states->dtype() + || workspace.cache13_capacity < cache13_required) { + if (context::isGraphRecording()) { + throw std::runtime_error( + "Hygon W8A8 Marlin cache13 workspace was not initialized before graph capture"); + } + workspace.cache13 = Tensor::empty( + {cache13_required}, + hidden_states->dtype(), + hidden_states->device()); + workspace.cache13_capacity = cache13_required; + } + ensure_tensor( + workspace.input_i8, + {num_tokens, hidden_size}, + DataType::I8, + hidden_states->device(), + "input_i8"); + ensure_tensor( + workspace.input_scale, + {num_tokens, 1}, + DataType::F32, + hidden_states->device(), + "input_scale"); + ensure_tensor( + workspace.cache2_i8, + {num_tokens * top_k, intermediate_size}, + DataType::I8, + hidden_states->device(), + "cache2_i8"); + ensure_tensor( + workspace.cache2_scale, + {num_tokens * top_k, 1}, + DataType::F32, + hidden_states->device(), + "cache2_scale"); + + moe_w8a8_marlin_fused_dense_( + workspace.output, + workspace.cache13, + workspace.cache2_i8, + workspace.input_i8, + workspace.input_scale, + workspace.cache2_scale, + hidden_states, + weights.packed_w13, + weights.packed_w2, + weights.packed_w13_scale, + weights.packed_w2_scale, + topk_weights, + routing.sorted_token_ids, + routing.expert_ids, + routing.num_tokens_post_padded, + top_k, + config.gemm1.mode, + config.gemm1.block_size_m, + config.gemm1.delta, + config.gemm2.mode, + config.gemm2.delta); + return workspace.output; +} + +HygonMoeMarlinOutput run_sliced( + const Tensor &hidden_states, + const Tensor &topk_weights, + const Tensor &topk_ids, + const HygonMoeMarlinWeights &weights, + HygonMoeMarlinWorkspace &workspace, + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size) { + if (context::isGraphRecording()) { + throw std::runtime_error( + "Hygon sliced MoE Marlin cannot allocate or copy outputs during graph capture"); + } + + ensure_tensor( + workspace.output, + hidden_states->shape(), + hidden_states->dtype(), + hidden_states->device(), + "output"); + HygonMoeMarlinWorkspace slice_workspace; + const auto activation_dtype = hidden_states->dtype(); + const auto device_index = hidden_states->device().getIndex(); + + HygonW16A16MarlinRuntimeConfig full_w16_config; + HygonW8A8MarlinRuntimeConfig full_w8_config; + if (weights.format == HygonMoeMarlinWeightFormat::W16A16) { + full_w16_config = select_hygon_w16a16_marlin_config( + kHygonMoeSliceTokens, + hidden_size, + intermediate_size, + activation_dtype, + device_index); + if (!full_w16_config.supported) { + throw std::runtime_error( + "No LightOP W16A16 Marlin config found for full Hygon slice"); + } + } else { + full_w8_config = select_hygon_w8a8_marlin_config( + kHygonMoeSliceTokens, + hidden_size, + intermediate_size, + device_index); + if (!full_w8_config.supported) { + throw std::runtime_error( + "No LightOP W8A8 Marlin config found for full Hygon slice"); + } + } + + const size_t num_tokens = hidden_states->shape()[0]; + size_t offset = 0; + while (offset < num_tokens) { + const size_t slice_tokens = std::min( + kHygonMoeSliceTokens, + num_tokens - offset); + auto hidden_slice = + hidden_states->narrow({{0, offset, slice_tokens}}); + auto topk_weights_slice = + topk_weights->narrow({{0, offset, slice_tokens}}); + auto topk_ids_slice = + topk_ids->narrow({{0, offset, slice_tokens}}); + + if (weights.format == HygonMoeMarlinWeightFormat::W16A16) { + const auto config = + slice_tokens == kHygonMoeSliceTokens + ? full_w16_config + : select_hygon_w16a16_marlin_config( + slice_tokens, + hidden_size, + intermediate_size, + activation_dtype, + device_index); + if (!config.supported) { + throw std::runtime_error( + "No LightOP W16A16 Marlin config found for sliced Hygon shape"); + } + const auto routing = prepare_routing( + topk_ids_slice, + Tensor(), + slice_workspace, + num_local_experts, + config.gemm1.block_size_m); + const auto slice_output = run_w16a16( + hidden_slice, + topk_weights_slice, + topk_ids_slice, + routing, + weights, + slice_workspace, + num_local_experts, + hidden_size, + intermediate_size, + config); + workspace.output + ->narrow({{0, offset, slice_tokens}}) + ->copy_from(slice_output); + } else { + const auto config = + slice_tokens == kHygonMoeSliceTokens + ? full_w8_config + : select_hygon_w8a8_marlin_config( + slice_tokens, + hidden_size, + intermediate_size, + device_index); + if (!config.supported) { + throw std::runtime_error( + "No LightOP W8A8 Marlin config found for sliced Hygon shape"); + } + const auto routing = prepare_routing( + topk_ids_slice, + Tensor(), + slice_workspace, + num_local_experts, + config.gemm1.block_size_m); + const auto slice_output = run_w8a8( + hidden_slice, + topk_weights_slice, + topk_ids_slice, + routing, + weights, + slice_workspace, + num_local_experts, + hidden_size, + intermediate_size, + config); + workspace.output + ->narrow({{0, offset, slice_tokens}}) + ->copy_from(slice_output); + } + offset += slice_tokens; + } + + return HygonMoeMarlinOutput{ + workspace.output, + Tensor(), + Tensor(), + Tensor(), + false, + }; +} + +} // namespace + +HygonMoeMarlinOutput hygon_moe_marlin_fused( + const Tensor &hidden_states, + const Tensor &topk_weights, + const Tensor &topk_ids, + const Tensor &expert_map, + const HygonMoeMarlinWeights &weights, + HygonMoeMarlinWorkspace &workspace, + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size, + size_t fallback_align_block_size) { + const auto &hidden_shape = hidden_states->shape(); + if (hidden_shape.size() != 2) { + throw std::runtime_error( + "Hygon MoE Marlin requires hidden_states [M, K]"); + } + if (hidden_shape[1] != hidden_size) { + throw std::runtime_error( + "Hygon MoE Marlin hidden size mismatch"); + } + if (topk_weights->shape() != topk_ids->shape() + || topk_ids->shape().size() != 2 + || topk_ids->shape()[0] != hidden_shape[0]) { + throw std::runtime_error( + "Hygon MoE Marlin topk tensors must have shape [M, top_k]"); + } + if (hidden_shape[0] > kHygonMoeSliceTokens) { + return run_sliced( + hidden_states, + topk_weights, + topk_ids, + weights, + workspace, + num_local_experts, + hidden_size, + intermediate_size); + } + + size_t block_size = fallback_align_block_size; + Tensor output; + RoutingMetadata routing; + if (weights.format == HygonMoeMarlinWeightFormat::W16A16) { + const auto config = select_hygon_w16a16_marlin_config( + hidden_shape[0], + hidden_size, + intermediate_size, + hidden_states->dtype(), + hidden_states->device().getIndex()); + if (!config.supported) { + throw std::runtime_error( + "No LightOP W16A16 Marlin config found for this Hygon shape"); + } + block_size = config.gemm1.block_size_m; + routing = prepare_routing( + topk_ids, + expert_map, + workspace, + num_local_experts, + block_size); + output = run_w16a16( + hidden_states, + topk_weights, + topk_ids, + routing, + weights, + workspace, + num_local_experts, + hidden_size, + intermediate_size, + config); + } else { + const auto config = select_hygon_w8a8_marlin_config( + hidden_shape[0], + hidden_size, + intermediate_size, + hidden_states->device().getIndex()); + if (!config.supported) { + throw std::runtime_error( + "No LightOP W8A8 Marlin config found for this Hygon shape"); + } + block_size = config.gemm1.block_size_m; + routing = prepare_routing( + topk_ids, + expert_map, + workspace, + num_local_experts, + block_size); + output = run_w8a8( + hidden_states, + topk_weights, + topk_ids, + routing, + weights, + workspace, + num_local_experts, + hidden_size, + intermediate_size, + config); + } + + return HygonMoeMarlinOutput{ + output, + routing.sorted_token_ids, + routing.expert_ids, + routing.num_tokens_post_padded, + true, + }; +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_marlin_config/moe_marlin_config.cc b/src/infinicore/ops/moe_marlin_config/moe_marlin_config.cc new file mode 100644 index 000000000..d207770e7 --- /dev/null +++ b/src/infinicore/ops/moe_marlin_config/moe_marlin_config.cc @@ -0,0 +1,215 @@ +#include "infinicore/ops/moe_marlin_config.hpp" + +#include "infinicore/adaptor/lightop_adaptor.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace infinicore::op { +namespace { + +enum class HygonMarlinModePolicy { + LegacyOnly, + LegacyAndBf16Mode1000, + All, +}; + +std::string lightop_config_dir() { + const char *value = std::getenv("INFINICORE_LIGHTOP_CONFIG_DIR"); + if (value != nullptr && value[0] != '\0') { + return value; + } + + // Keep the old override working while ownership moves from InfiniLM. + value = std::getenv("INFINILM_LIGHTOP_CONFIG_DIR"); + if (value != nullptr && value[0] != '\0') { + return value; + } + return "/usr/local/lib/python3.10/dist-packages/lightop/configs"; +} + +std::string normalize_hygon_gpu_target(std::string target, bool uppercase) { + const auto feature_pos = target.find(':'); + if (feature_pos != std::string::npos) { + target.resize(feature_pos); + } + std::transform(target.begin(), target.end(), target.begin(), [uppercase](unsigned char ch) { + return static_cast(uppercase ? std::toupper(ch) : std::tolower(ch)); + }); + + std::string lowercase = target; + std::transform(lowercase.begin(), lowercase.end(), lowercase.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + if (lowercase.size() <= 3 + || lowercase.compare(0, 3, "gfx") != 0 + || !std::all_of(lowercase.begin() + 3, lowercase.end(), [](unsigned char ch) { + return std::isalnum(ch) != 0; + })) { + throw std::runtime_error("Invalid Hygon GPU target for LightOP config: " + target); + } + return target; +} + +HygonMarlinGemmConfig load_lightop_marlin_config( + size_t n, + size_t k, + size_t m, + const std::string &file_prefix, + const adaptor::lightop::DeviceInfo &device_info, + HygonMarlinModePolicy mode_policy, + bool uppercase_device_name, + bool num_cus_with_cu_prefix) { + HygonMarlinGemmConfig result; + if (device_info.gpu_target.empty() || device_info.compute_units <= 0) { + throw std::runtime_error("Unable to query Hygon device properties for LightOP config"); + } + + const std::string device_name = normalize_hygon_gpu_target( + device_info.gpu_target, + uppercase_device_name); + const std::string num_cus = std::to_string(device_info.compute_units); + const std::string num_cus_suffix = + num_cus_with_cu_prefix ? ("_CU" + num_cus) : ("_" + num_cus); + const std::string file_name = + lightop_config_dir() + "/" + file_prefix + "_" + + std::to_string(n) + "_" + std::to_string(k) + "_" + + device_name + num_cus_suffix + ".json"; + + std::ifstream file(file_name); + if (!file.is_open()) { + return result; + } + + nlohmann::json config_json; + file >> config_json; + const std::string shape_key = std::to_string(n) + "_" + std::to_string(k); + if (!config_json.contains(shape_key) || !config_json.at(shape_key).is_object()) { + return result; + } + const auto &configs = config_json.at(shape_key); + + auto usable = [&](size_t token) -> bool { + const auto key = std::to_string(token); + if (!configs.contains(key) || !configs.at(key).is_object()) { + return false; + } + const int mode = configs.at(key).value("MODE", result.mode); + return mode < 1000 + || mode_policy == HygonMarlinModePolicy::All + || (mode_policy == HygonMarlinModePolicy::LegacyAndBf16Mode1000 + && mode == 1000); + }; + + size_t chosen = 0; + bool has_choice = false; + size_t chosen_ge = std::numeric_limits::max(); + size_t closest_diff = std::numeric_limits::max(); + for (auto it = configs.begin(); it != configs.end(); ++it) { + size_t token = 0; + try { + token = static_cast(std::stoull(it.key())); + } catch (const std::exception &) { + continue; + } + if (!usable(token)) { + continue; + } + if (token >= m && token < chosen_ge) { + chosen_ge = token; + chosen = token; + has_choice = true; + } + const size_t diff = token > m ? token - m : m - token; + if (diff < closest_diff) { + closest_diff = diff; + if (chosen_ge == std::numeric_limits::max()) { + chosen = token; + has_choice = true; + } + } + } + if (!has_choice) { + return result; + } + + const auto &config = configs.at(std::to_string(chosen)); + result.mode = config.value("MODE", result.mode); + result.delta = config.value("DELTA", result.delta); + result.block_size_m = config.value("BLOCK_SIZE_M", result.block_size_m); + result.found = config.contains("MODE"); + return result; +} + +} // namespace + +HygonW16A16MarlinRuntimeConfig select_hygon_w16a16_marlin_config( + size_t num_tokens, + size_t hidden_size, + size_t intermediate_size, + DataType hidden_dtype, + size_t device_index) { + HygonW16A16MarlinRuntimeConfig config; + const auto device_info = adaptor::lightop::device_info(device_index); + const auto mode_policy = hidden_dtype == DataType::BF16 + ? HygonMarlinModePolicy::LegacyAndBf16Mode1000 + : HygonMarlinModePolicy::LegacyOnly; + config.gemm1 = load_lightop_marlin_config( + intermediate_size * 2, + hidden_size, + num_tokens, + "MOE_W16A16_CUDA_MARLIN", + device_info, + mode_policy, + false, + false); + config.gemm2 = load_lightop_marlin_config( + hidden_size, + intermediate_size, + num_tokens, + "MOE_W16A16_CUDA_MARLIN", + device_info, + mode_policy, + false, + false); + config.supported = config.gemm1.found && config.gemm2.found; + return config; +} + +HygonW8A8MarlinRuntimeConfig select_hygon_w8a8_marlin_config( + size_t num_tokens, + size_t hidden_size, + size_t intermediate_size, + size_t device_index) { + HygonW8A8MarlinRuntimeConfig config; + const auto device_info = adaptor::lightop::device_info(device_index); + config.gemm1 = load_lightop_marlin_config( + intermediate_size * 2, + hidden_size, + num_tokens, + "MOE_BLOCKINT8_CUDA_MARLIN", + device_info, + HygonMarlinModePolicy::All, + true, + true); + config.gemm2 = load_lightop_marlin_config( + hidden_size, + intermediate_size, + num_tokens, + "MOE_BLOCKINT8_CUDA_MARLIN", + device_info, + HygonMarlinModePolicy::All, + true, + true); + config.supported = config.gemm1.found && config.gemm2.found; + return config; +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/paged_flash_attention/paged_flash_attention.cc b/src/infinicore/ops/paged_flash_attention/paged_flash_attention.cc new file mode 100644 index 000000000..3a8745ac3 --- /dev/null +++ b/src/infinicore/ops/paged_flash_attention/paged_flash_attention.cc @@ -0,0 +1,139 @@ +#include "infinicore/ops/paged_flash_attention.hpp" + +#include "infinicore/ops/mha_kvcache.hpp" +#include "infinicore/ops/mha_varlen.hpp" +#include "infinicore/ops/paged_caching.hpp" + +#include +#include +#include +#include + +namespace infinicore::op { +namespace { + +std::pair update_paged_kv_cache( + const Tensor &key, + const Tensor &value, + const Tensor &kv_cache, + const Tensor &slot_mapping, + size_t num_heads, + size_t num_kv_heads, + size_t head_dim) { + auto k_cache_layer = kv_cache->narrow({{0, 0, 1}})->squeeze(0); + auto v_cache_layer = kv_cache->narrow({{0, 1, 1}})->squeeze(0); + const auto &cache_shape = k_cache_layer->shape(); + const bool use_hygon_lightop_paged_attention = + key->device().getType() == Device::Type::HYGON + && cache_shape.size() == 4 + && cache_shape[1] == 64 + && cache_shape[2] == num_kv_heads + && cache_shape[3] == head_dim + && num_heads == 8 + && num_kv_heads == 1 + && head_dim == 128; + if (use_hygon_lightop_paged_attention) { + const auto num_blocks = cache_shape[0]; + const auto block_size = cache_shape[1]; + auto k_cache_lightop = k_cache_layer->view( + {num_blocks, num_kv_heads, block_size, head_dim}); + auto v_cache_lightop = v_cache_layer->view( + {num_blocks, num_kv_heads, head_dim, block_size}); + paged_caching_( + k_cache_lightop, + v_cache_lightop, + key, + value, + slot_mapping); + return {k_cache_lightop, v_cache_lightop}; + } + + paged_caching_( + k_cache_layer->permute({0, 2, 1, 3}), + v_cache_layer->permute({0, 2, 1, 3}), + key, + value, + slot_mapping); + return {k_cache_layer, v_cache_layer}; +} + +} // namespace + +Tensor paged_flash_attention( + const Tensor &query, + const Tensor &key, + const Tensor &value, + const Tensor &kv_cache, + const Tensor &total_sequence_lengths, + const std::optional &input_offsets, + const std::optional &cu_seqlens, + const Tensor &block_tables, + const Tensor &slot_mapping, + size_t num_heads, + size_t num_kv_heads, + size_t head_dim, + float scale) { + static std::mutex hygon_paged_flash_attention_mutex; + std::unique_lock hygon_lock( + hygon_paged_flash_attention_mutex, + std::defer_lock); + if (query->device().getType() == Device::Type::HYGON) { + hygon_lock.lock(); + } + + auto [k_total, v_total] = update_paged_kv_cache( + key, + value, + kv_cache, + slot_mapping, + num_heads, + num_kv_heads, + head_dim); + + const size_t seq_len = query->shape()[0]; + const bool is_prefill = seq_len != total_sequence_lengths->shape()[0]; + auto attn_output = Tensor::empty( + {seq_len, num_heads, head_dim}, + query->dtype(), + query->device()); + + if (is_prefill) { + const auto cache_block_size = kv_cache->shape()[2]; + const auto max_cache_seqlen = + block_tables->shape()[1] * cache_block_size; + if (seq_len > static_cast(std::numeric_limits::max()) + || max_cache_seqlen + > static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "FlashAttention sequence length exceeds int range"); + } + mha_varlen_( + attn_output, + query, + k_total, + v_total, + input_offsets.value(), + cu_seqlens.value(), + block_tables, + static_cast(seq_len), + static_cast(max_cache_seqlen), + std::nullopt, + scale); + } else { + auto q_for_fa = query->view({seq_len, 1, num_heads, head_dim}); + auto attn_out_4d = mha_kvcache( + q_for_fa, + k_total, + v_total, + total_sequence_lengths, + block_tables, + std::nullopt, + scale); + attn_output = + attn_out_4d->view({seq_len, num_heads, head_dim}); + } + + return attn_output->view({1, seq_len, num_heads * head_dim}); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/select_last_token_hidden_states/select_last_token_hidden_states.cc b/src/infinicore/ops/select_last_token_hidden_states/select_last_token_hidden_states.cc new file mode 100644 index 000000000..b1d4a3a4f --- /dev/null +++ b/src/infinicore/ops/select_last_token_hidden_states/select_last_token_hidden_states.cc @@ -0,0 +1,77 @@ +#include "infinicore/ops/select_last_token_hidden_states.hpp" + +#include +#include + +namespace infinicore::op { + +Tensor select_last_token_hidden_states( + const Tensor &hidden_states, + const Tensor &input_offsets) { + if (input_offsets->dtype() != DataType::I32) { + throw std::runtime_error( + "select_last_token_hidden_states: input_offsets must have I32 dtype"); + } + if (input_offsets->ndim() != 1 + || input_offsets->size(0) < 2) { + throw std::runtime_error( + "select_last_token_hidden_states: input_offsets must be a 1D tensor with at least two elements"); + } + if (hidden_states->ndim() != 3) { + throw std::runtime_error( + "select_last_token_hidden_states: expected rank-3 hidden_states"); + } + + const auto num_requests = input_offsets->size(0) - 1; + const auto hidden_size = hidden_states->size(2); + const auto total_tokens = + hidden_states->size(0) * hidden_states->size(1); + if (total_tokens < num_requests) { + throw std::runtime_error( + "select_last_token_hidden_states: more requests than input tokens"); + } + if (total_tokens == num_requests) { + return hidden_states; + } + + auto input_offsets_cpu = input_offsets->to(Device::cpu()); + const auto *offsets = reinterpret_cast( + input_offsets_cpu->data()); + if (offsets[0] != 0 + || offsets[num_requests] < 0 + || static_cast(offsets[num_requests]) + != total_tokens) { + throw std::runtime_error( + "select_last_token_hidden_states: input_offsets must cover all input tokens"); + } + for (size_t i = 0; i < num_requests; ++i) { + const auto begin = offsets[i]; + const auto end = offsets[i + 1]; + if (begin < 0 + || end <= begin + || static_cast(end) > total_tokens) { + throw std::runtime_error( + "select_last_token_hidden_states: input_offsets must be strictly increasing and in range"); + } + } + + auto flat_hidden_states = + hidden_states->view({total_tokens, hidden_size}); + auto selected_hidden_states = Tensor::empty( + {1, num_requests, hidden_size}, + hidden_states->dtype(), + hidden_states->device()); + for (size_t i = 0; i < num_requests; ++i) { + const auto token_index = + static_cast(offsets[i + 1] - 1); + selected_hidden_states + ->narrow({{1, i, 1}}) + ->view({1, hidden_size}) + ->copy_from( + flat_hidden_states->narrow( + {{0, token_index, 1}})); + } + return selected_hidden_states; +} + +} // namespace infinicore::op From c5ad69d0f07adf3212411b76a87c93242dbcc9cb Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Tue, 4 Aug 2026 16:35:09 +0800 Subject: [PATCH 14/16] fix(core): resolve RoPE cache getter rebase conflict --- include/infinicore/nn/rope.hpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/include/infinicore/nn/rope.hpp b/include/infinicore/nn/rope.hpp index e9ddc6dc4..877c48e94 100644 --- a/include/infinicore/nn/rope.hpp +++ b/include/infinicore/nn/rope.hpp @@ -85,8 +85,6 @@ class RoPE : public Module { double theta() const { return theta_; } Algo algo() const { return algo_; } DataType dtype() const { return dtype_; } - Tensor sin_cache() const { return sin_cache_; } - Tensor cos_cache() const { return cos_cache_; } Tensor cos_sin_cache() const { return cos_sin_cache_; } const std::optional> &mrope_section() const { return mrope_section_; } bool mrope_interleaved() const { return mrope_interleaved_; } From d40aed57f2927b468023e84e8209d6624556149d Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Tue, 4 Aug 2026 16:40:14 +0800 Subject: [PATCH 15/16] fix(graph): remove stale rebase preprocessor guard --- src/infinicore/graph/graph.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc index e39d87705..5064a0f52 100644 --- a/src/infinicore/graph/graph.cc +++ b/src/infinicore/graph/graph.cc @@ -148,7 +148,6 @@ void Graph::instantiate() { spdlog::info("device graph segments disabled; replaying recorded operators"); return; } -#endif for (const auto &op : op_list_) { const bool capture_safe = is_hygon || op->is_device_graph_capture_safe(); From 4638ed26aa9a401521612de79ec07f76020f67c9 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Wed, 5 Aug 2026 15:57:10 +0800 Subject: [PATCH 16/16] fix-hygon-align-allreduce-graph-capture-hook --- include/infinicore/ops/distributed/allreduce.hpp | 2 +- src/infinicore/ops/distributed/allreduce.cc | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/include/infinicore/ops/distributed/allreduce.hpp b/include/infinicore/ops/distributed/allreduce.hpp index d2fd25600..f2099028c 100644 --- a/include/infinicore/ops/distributed/allreduce.hpp +++ b/include/infinicore/ops/distributed/allreduce.hpp @@ -12,7 +12,7 @@ class AllReduce : public graph::GraphOperator { AllReduce(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); ~AllReduce(); void run() const override; - bool supports_device_graph_capture() const override; + bool is_device_graph_capture_safe() const override; static void execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); private: diff --git a/src/infinicore/ops/distributed/allreduce.cc b/src/infinicore/ops/distributed/allreduce.cc index 7d140ac69..124ed8989 100644 --- a/src/infinicore/ops/distributed/allreduce.cc +++ b/src/infinicore/ops/distributed/allreduce.cc @@ -34,7 +34,7 @@ void AllReduce::run() const { infinicore::context::getStream())); } -bool AllReduce::supports_device_graph_capture() const { +bool AllReduce::is_device_graph_capture_safe() const { const auto *meta = reinterpret_cast(planned_meta_); return meta->input->device().getType() != Device::Type::HYGON; }