diff --git a/include/infinicore/adaptor/lightop_adaptor.hpp b/include/infinicore/adaptor/lightop_adaptor.hpp new file mode 100644 index 000000000..9c7f7ba4d --- /dev/null +++ b/include/infinicore/adaptor/lightop_adaptor.hpp @@ -0,0 +1,147 @@ +#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 +#include + +namespace infinicore::adaptor::lightop { + +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(); + +void preload_moe_w8a8_marlin_asm(); + +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); + +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 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, + 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); + +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, + 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); + +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/nn/rope.hpp b/include/infinicore/nn/rope.hpp index 37f350678..622b0c2be 100644 --- a/include/infinicore/nn/rope.hpp +++ b/include/infinicore/nn/rope.hpp @@ -85,6 +85,7 @@ class RoPE : public Module { double theta() const { return theta_; } Algo algo() const { return algo_; } DataType dtype() const { return dtype_; } + const 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_; } diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index 5e93e1457..5593f7a1d 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/kimi_delta_attention.hpp" #include "ops/kv_caching.hpp" #include "ops/layer_norm.hpp" @@ -55,14 +56,18 @@ #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" +#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" #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" @@ -73,6 +78,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/rotary_embedding.hpp" @@ -82,6 +88,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..a2e1f8888 --- /dev/null +++ b/include/infinicore/ops/hygon_moe_marlin.hpp @@ -0,0 +1,60 @@ +#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/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/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/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/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/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/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu index 9b305afd0..5cc076ec5 100644 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ b/src/infiniccl/cuda/infiniccl_cuda.cu @@ -6,6 +6,20 @@ #include #include +#if defined(ENABLE_HYGON_API) +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#endif + #include "../../utils.h" #define CHECK_NCCL(API__) CHECK_INTERNAL(API__, ncclSuccess) @@ -70,6 +84,1001 @@ infiniStatus_t getCommName( return getCommNameFromHandle(comm, comm_name, comm_name_size); } +#if defined(ENABLE_HYGON_API) +namespace { + +constexpr int kHygonTp2MaxBlocks = 80; +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; +constexpr unsigned int kHygonHipDeviceMallocUncached = 0x3; +constexpr unsigned int kHygonHipEventDisableTiming = 0x2; +constexpr unsigned int kHygonHipEventReleaseToSystem = 0x80000000u; + +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]; +}; + +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; + 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 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; + CUmemGenericAllocationHandle handle = 0; + bool hip_uncached = false; +}; + +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].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); + 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); + } + } +}; + +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], + 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(); + } + 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) { + 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(); + } + 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; +} + +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; + __shared__ uint32_t block_flag; + + const uint32_t sync_flag = hygon_tp8_start_sync( + rank_signals, self_signal, rank, &block_flag); + + 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 element = 0; element < 8; ++element) { + reduced[element] = __bfloat162float(source_pack.values[element]); + } +#pragma unroll + 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) { + reduced[element] += __bfloat162float(source_pack.values[element]); + } + } + 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( + 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; + } + + 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_threads = use_two_stage + ? kHygonTp8Threads / kHygonTp8WorldSize + : kHygonTp8Threads; + const int max_blocks = use_two_stage ? kHygonTp8MaxBlocks : kHygonTp8OneStageMaxBlocks; + int blocks = static_cast(std::min( + 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]; + 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 + infiniStatus_t commInitAll( infinicclComm_t *comms, int ndevice, @@ -82,6 +1091,11 @@ 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); + register_hygon_tp8_state(comms, ndevice, device_ids); +#endif + return INFINI_STATUS_SUCCESS; } @@ -119,6 +1133,10 @@ 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; return INFINI_STATUS_SUCCESS; @@ -147,6 +1165,23 @@ infiniStatus_t allReduce( INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, 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))) { + return INFINI_STATUS_SUCCESS; + } +#endif + CHECK_NCCL(ncclAllReduce(sendbuf, recvbuf, count, getNcclDtype(datatype), getNcclRedOp(op), getNcclComm(comm), getCudaStream(stream))); 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/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 new file mode 100644 index 000000000..f7ac0b561 --- /dev/null +++ b/src/infinicore/adaptor/lightop_adaptor.cc @@ -0,0 +1,1153 @@ +#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 { + +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 *kLightopAsmRoot = "/usr/local/lib/python3.10/dist-packages/lightop/hsa/"; +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 = "_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"; + +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: + 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; + } + + 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; + 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 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; +} + +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 = "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(); +} + +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; + } + + 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; + 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); + 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; + } + + 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; + 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); +} + +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, + 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; +} + +template +Fn resolve(const char *symbol) { + return reinterpret_cast(library().symbol(symbol)); +} + +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 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 &, + const std::optional &, + const std::optional &, + const std::optional &, + float, + int); +using MoeAlignBlockSizeFn = void (*)( + at::Tensor, + int64_t, + int64_t, + 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, + 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); + +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; +} + +ReshapeAndCacheCudaFn reshape_and_cache_cuda_fn() { + static auto fn = resolve(kReshapeAndCacheCudaSymbol); + return fn; +} + +MoeSumFn moe_sum_fn() { + static auto fn = resolve(kMoeSumSymbol); + return fn; +} + +MoeAlignBlockSizeFn moe_align_block_size_fn() { + static auto fn = resolve(kMoeAlignBlockSizeSymbol); + 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; +} + +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 + +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(); +} + +void preload_moe_w16a16_ops(bool preload_legacy_gemm, bool preload_legacy_asm) { + (void)moe_sum_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() { + (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 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_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(); + (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 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, + 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); +} + +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, + 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 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, + sorted_token_ids, expert_ids, num_tokens_post_padded, + static_cast(top_k), mode, 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) { + 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 + +#else + +namespace infinicore::adaptor::lightop { + +DeviceInfo device_info(std::size_t) { + return {}; +} + +} // 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..9f286a782 100644 --- a/src/infinicore/context/allocators/pinnable_block_allocator.cc +++ b/src/infinicore/context/allocators/pinnable_block_allocator.cc @@ -33,6 +33,10 @@ PinnableBlockAllocator::PinnableBlockAllocator(Device device) {128 * 1024 * 1024, {}}, // 128 MB {256 * 1024 * 1024, {}}, // 256 MB }; + if (device_.getType() == Device::Type::HYGON) { + // Large model shapes benefit from exact-size reuse instead of coarse classes. + size_classes_.resize(8); + } } // ------------------- allocate ------------------- @@ -85,6 +89,15 @@ std::byte *PinnableBlockAllocator::allocate(size_t size) { // 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; }); + if (device_.getType() == Device::Type::HYGON) { + 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 +107,22 @@ 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. + if (device_.getType() == Device::Type::HYGON) { + 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 6c7e52ad4..83d461384 100644 --- a/src/infinicore/nn/rope.cc +++ b/src/infinicore/nn/rope.cc @@ -77,8 +77,12 @@ void RoPE::initialize_cache() { 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_)); -#ifdef ENABLE_INFINIOPS_API - if ((device_.getType() == Device::Type::NVIDIA || device_.getType() == Device::Type::METAX) && !mrope_section_) { +#if defined(ENABLE_INFINIOPS_API) || defined(ENABLE_HYGON_API) + const auto device_type = device_.getType(); + if ((device_type == Device::Type::NVIDIA + || device_type == Device::Type::METAX + || device_type == Device::Type::HYGON) + && !mrope_section_) { INFINICORE_NN_BUFFER_INIT(cos_sin_cache, ({max_seq_len_, rotary_dim_}, dtype_, device_)); } #endif 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..3f44ac5a0 --- /dev/null +++ b/src/infinicore/ops/hygon_moe_marlin/hygon_moe_marlin.cc @@ -0,0 +1,636 @@ +#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/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc b/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc index 9d2cdc2b5..8aae170f2 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,13 +4,11 @@ #include "../../../adaptor/flash_attn/hygon/flash_attn_hygon.hpp" #include "infinicore/adaptor/aten_adaptor.hpp" -#include -#include +#include + #include #include - -#include -#include +#include namespace infinicore::op::mha_kvcache_impl::flashattn { @@ -50,58 +48,49 @@ 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; - 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); + 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"); } - - out_tensor.copy_(result); + 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); } @@ -114,12 +103,19 @@ void run(void *planned_meta) { 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, @@ -143,7 +139,8 @@ void run(void *planned_meta) { false, 0); - 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..fc7527252 --- /dev/null +++ b/src/infinicore/ops/moe_align/hygon/moe_align_lightop_hygon.cc @@ -0,0 +1,124 @@ +#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); + + 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, + p->pad_sorted_token_ids); +} + +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 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/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..58cf12121 --- /dev/null +++ b/src/infinicore/ops/moe_marlin_config/moe_marlin_config.cc @@ -0,0 +1,400 @@ +#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, +}; + +// LightOP config files have a fixed object -> object -> integer object schema. +// Parsing only that schema keeps this path dependency-free on every platform. +class LightopConfigReader { +public: + explicit LightopConfigReader(const std::string &text) : _text(text) {} + + bool consume(char expected) { + skip_whitespace(); + if (_pos < _text.size() && _text[_pos] == expected) { + ++_pos; + return true; + } + return false; + } + + void expect(char expected) { + if (!consume(expected)) { + fail(std::string("expected '") + expected + "'"); + } + } + + std::string key() { + expect('"'); + const size_t begin = _pos; + while (_pos < _text.size() && _text[_pos] != '"') { + const unsigned char ch = static_cast(_text[_pos]); + if (ch == '\\' || ch < 0x20) { + fail("LightOP config keys must be unescaped strings"); + } + ++_pos; + } + if (_pos == _text.size()) { + fail("unterminated key"); + } + const std::string result = _text.substr(begin, _pos - begin); + ++_pos; + return result; + } + + long long integer() { + skip_whitespace(); + const size_t begin = _pos; + if (_pos < _text.size() && _text[_pos] == '-') { + ++_pos; + } + const size_t digits_begin = _pos; + while (_pos < _text.size() && _text[_pos] >= '0' && _text[_pos] <= '9') { + ++_pos; + } + if (_pos == digits_begin) { + fail("expected integer"); + } + if (_pos - digits_begin > 1 && _text[digits_begin] == '0') { + fail("integer has a leading zero"); + } + if (_pos < _text.size() + && (_text[_pos] == '.' || _text[_pos] == 'e' || _text[_pos] == 'E')) { + fail("expected integer, found non-integral number"); + } + try { + return std::stoll(_text.substr(begin, _pos - begin)); + } catch (const std::exception &) { + fail("integer is out of range"); + } + } + + void finish() { + skip_whitespace(); + if (_pos != _text.size()) { + fail("unexpected trailing content"); + } + } + +private: + void skip_whitespace() { + while (_pos < _text.size() + && std::isspace(static_cast(_text[_pos])) != 0) { + ++_pos; + } + } + + [[noreturn]] void fail(const std::string &message) const { + throw std::runtime_error( + "invalid LightOP JSON at byte " + std::to_string(_pos) + ": " + message); + } + + const std::string &_text; + size_t _pos = 0; +}; + +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; +} + +bool parse_size(const std::string &text, size_t &value) { + if (text.empty()) { + return false; + } + size_t result = 0; + for (const char ch : text) { + if (ch < '0' || ch > '9') { + return false; + } + const size_t digit = static_cast(ch - '0'); + if (result > (std::numeric_limits::max() - digit) / 10) { + return false; + } + result = result * 10 + digit; + } + value = result; + return true; +} + +int checked_int(long long value, const std::string &field) { + if (value < std::numeric_limits::min() + || value > std::numeric_limits::max()) { + throw std::runtime_error("LightOP field " + field + " is out of int range"); + } + return static_cast(value); +} + +HygonMarlinGemmConfig parse_gemm_config(LightopConfigReader &json) { + HygonMarlinGemmConfig config; + bool has_mode = false; + json.expect('{'); + if (json.consume('}')) { + return config; + } + while (true) { + const std::string field = json.key(); + json.expect(':'); + const long long value = json.integer(); + if (field == "MODE") { + config.mode = checked_int(value, field); + has_mode = true; + } else if (field == "DELTA") { + config.delta = checked_int(value, field); + } else if (field == "BLOCK_SIZE_M") { + if (value < 0 + || static_cast(value) + > std::numeric_limits::max()) { + throw std::runtime_error("LightOP field BLOCK_SIZE_M is out of size_t range"); + } + config.block_size_m = static_cast(value); + } + if (json.consume('}')) { + break; + } + json.expect(','); + } + config.found = has_mode; + return config; +} + +bool mode_is_usable(int mode, HygonMarlinModePolicy mode_policy) { + return mode < 1000 + || mode_policy == HygonMarlinModePolicy::All + || (mode_policy == HygonMarlinModePolicy::LegacyAndBf16Mode1000 && mode == 1000); +} + +struct ConfigChoice { + HygonMarlinGemmConfig config; + size_t chosen_ge = std::numeric_limits::max(); + size_t closest_diff = std::numeric_limits::max(); + + void consider( + size_t token, + const HygonMarlinGemmConfig &candidate, + size_t requested_tokens, + HygonMarlinModePolicy mode_policy) { + if (!candidate.found || !mode_is_usable(candidate.mode, mode_policy)) { + return; + } + if (token >= requested_tokens && token < chosen_ge) { + chosen_ge = token; + config = candidate; + } + const size_t diff = token > requested_tokens + ? token - requested_tokens + : requested_tokens - token; + if (chosen_ge == std::numeric_limits::max() && diff < closest_diff) { + closest_diff = diff; + config = candidate; + } + } +}; + +void parse_token_configs( + LightopConfigReader &json, + bool select, + size_t requested_tokens, + HygonMarlinModePolicy mode_policy, + ConfigChoice &choice) { + json.expect('{'); + if (json.consume('}')) { + return; + } + while (true) { + const std::string token_key = json.key(); + json.expect(':'); + const HygonMarlinGemmConfig candidate = parse_gemm_config(json); + size_t token = 0; + if (select && parse_size(token_key, token)) { + choice.consider(token, candidate, requested_tokens, mode_policy); + } + if (json.consume('}')) { + break; + } + json.expect(','); + } +} + +HygonMarlinGemmConfig parse_lightop_marlin_config( + const std::string &contents, + const std::string &shape_key, + size_t requested_tokens, + HygonMarlinModePolicy mode_policy) { + LightopConfigReader json(contents); + ConfigChoice choice; + json.expect('{'); + if (!json.consume('}')) { + while (true) { + const std::string current_shape = json.key(); + json.expect(':'); + parse_token_configs( + json, + current_shape == shape_key, + requested_tokens, + mode_policy, + choice); + if (json.consume('}')) { + break; + } + json.expect(','); + } + } + json.finish(); + return choice.config; +} + +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; + } + + const std::string shape_key = std::to_string(n) + "_" + std::to_string(k); + const std::string contents{ + std::istreambuf_iterator(file), + std::istreambuf_iterator()}; + try { + return parse_lightop_marlin_config(contents, shape_key, m, mode_policy); + } catch (const std::exception &error) { + throw std::runtime_error( + "Failed to parse LightOP config " + file_name + ": " + error.what()); + } +} + +} // 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/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..9a6ee75e3 --- /dev/null +++ b/src/infinicore/ops/moe_w16a16_marlin/hygon/moe_w16a16_marlin_hygon.cc @@ -0,0 +1,262 @@ +#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/context/context.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 + +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; +} + +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::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), + 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) { + 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(); + 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 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 = 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; + 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); + + 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/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..0c5a407dc --- /dev/null +++ b/src/infinicore/ops/moe_w8a8_marlin/hygon/moe_w8a8_marlin_hygon.cc @@ -0,0 +1,292 @@ +#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..3d8a6ca55 --- /dev/null +++ b/src/infinicore/ops/moe_w8a8_marlin/moe_w8a8_marlin.cc @@ -0,0 +1,148 @@ +#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..a90d6a9d1 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; @@ -197,25 +125,39 @@ 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( @@ -240,7 +182,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/paged_caching/hygon/paged_caching_lightop_hygon.cc b/src/infinicore/ops/paged_caching/hygon/paged_caching_lightop_hygon.cc new file mode 100644 index 000000000..d2dc9f6a4 --- /dev/null +++ b/src/infinicore/ops/paged_caching/hygon/paged_caching_lightop_hygon.cc @@ -0,0 +1,141 @@ +#include "infinicore/ops/paged_caching.hpp" + +#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::lightop_hygon { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PagedCaching, 100); + +struct PlannedMeta { + bool use_hygon_lightop; + std::shared_ptr descriptor; + 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( + Descriptor, descriptor, PagedCaching, + key, k_cache->desc(), v_cache->desc(), k->desc(), v->desc(), slot_mapping->desc()); + + 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) +#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, + workspace->data(), + workspace->numel(), + p->k_cache->data(), + p->v_cache->data(), + p->k->data(), + p->v->data(), + p->slot_mapping->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + PagedCaching::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); + PagedCaching::run_dispatcher().registerDevice(Device::Type::HYGON, &run); + PagedCaching::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); + return true; +}(); + +} // namespace infinicore::op::paged_caching_impl::lightop_hygon 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..efe58fd49 --- /dev/null +++ b/src/infinicore/ops/paged_flash_attention/paged_flash_attention.cc @@ -0,0 +1,136 @@ +#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/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..0c075389d --- /dev/null +++ b/src/infinicore/ops/scaled_mm_i8/hygon/scaled_mm_i8_lightop_hygon.cc @@ -0,0 +1,100 @@ +#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/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..75e26d037 --- /dev/null +++ b/src/infinicore/ops/select_last_token_hidden_states/select_last_token_hidden_states.cc @@ -0,0 +1,74 @@ +#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 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..dcd4c8df5 --- /dev/null +++ b/src/infinicore/ops/silu_and_mul/hygon/silu_and_mul_lightop_hygon.cc @@ -0,0 +1,58 @@ +#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::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/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; } 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..774de29f5 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; } @@ -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); @@ -415,4 +425,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 993f44445..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 @@ -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_softmax/nvidia/moe_topk_softmax_nvidia.cu b/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu index 2fc0e992b..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 @@ -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/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; diff --git a/xmake/hygon.lua b/xmake/hygon.lua index ba06ae972..933a4d04f 100644 --- a/xmake/hygon.lua +++ b/xmake/hygon.lua @@ -221,9 +221,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")