diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index b8eeecc3d..40e60fa0c 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -104,9 +104,10 @@ InferEngine::InferEngine( && device_type != infinicore::Device::Type::kCambricon && device_type != infinicore::Device::Type::kIluvatar && device_type != infinicore::Device::Type::kHygon - && device_type != infinicore::Device::Type::kAscend) { + && device_type != infinicore::Device::Type::kAscend + && device_type != infinicore::Device::Type::kQy) { throw std::invalid_argument( - "flash-attn is only available on NVIDIA, MetaX, Moore, Cambricon, Iluvatar, Hygon, and Ascend devices"); + "flash-attn is only available on NVIDIA, MetaX, Moore, Cambricon, Iluvatar, Hygon, Ascend, and QY devices"); } const auto *paged_cache_config = dynamic_cast(cache_config); if (paged_cache_config == nullptr) { diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index ddfe1ee4e..c0af07fb7 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -493,7 +493,8 @@ void RankWorker::thread_loop() { const bool parameter_greedy = top_p == 0.0f || top_k == 1 || temperature == 0.0f; const auto logits_dtype = logits->dtype(); - const bool batch_greedy = rank_info_.device.type() == infinicore::Device::Type::kNvidia + const bool batch_greedy = (rank_info_.device.type() == infinicore::Device::Type::kNvidia + || rank_info_.device.type() == infinicore::Device::Type::kQy) && parameter_greedy && n_out > 0 && logits_positions == n_out diff --git a/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp b/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp index f48f640cd..107d38e26 100644 --- a/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp +++ b/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp @@ -8,7 +8,7 @@ #if defined(ENABLE_HYGON_API) #include #include -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) #include #include #include @@ -41,7 +41,9 @@ inline at::ScalarType to_at_dtype(DataType dtype) { inline at::Device to_at_device(const Device &device) { // PyTorch ATen only exposes standard device types (e.g. kCPU/kCUDA). // Treat CUDA-compatible devices as CUDA devices for ATen interoperability. - if (device.type() == Device::Type::kNvidia || device.type() == Device::Type::kMetax || device.type() == Device::Type::kHygon) { + if (device.type() == Device::Type::kNvidia || device.type() == Device::Type::kMetax + || device.type() == Device::Type::kHygon || device.type() == Device::Type::kQy + || device.type() == Device::Type::kThead || device.type() == Device::Type::kIluvatar) { return at::Device(at::kCUDA, device.index()); } else if (device.type() == Device::Type::kCpu) { return at::Device(at::kCPU); @@ -60,7 +62,7 @@ at::Tensor to_aten_tensor(const infinicore::Tensor &t); #if defined(ENABLE_HYGON_API) c10::hip::HIPStream get_hip_stream(); -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) c10::cuda::CUDAStream get_cuda_stream(); #endif @@ -68,6 +70,9 @@ c10::cuda::CUDAStream get_cuda_stream(); c10::musa::MUSAStream get_musa_stream(); #endif +/** Align ATen/CUDA current stream with InfiniCore's active stream. */ +void set_aten_stream_to_infinicore(); + } // namespace infinicore::adaptor #endif // ENABLE_ATEN diff --git a/csrc/infinicore/include/infinicore/adaptor/denglin_fused_moe_adaptor.hpp b/csrc/infinicore/include/infinicore/adaptor/denglin_fused_moe_adaptor.hpp new file mode 100644 index 000000000..a6cb99da8 --- /dev/null +++ b/csrc/infinicore/include/infinicore/adaptor/denglin_fused_moe_adaptor.hpp @@ -0,0 +1,43 @@ +#pragma once + +#if defined(ENABLE_ATEN) && defined(ENABLE_QY_API) +#include +#include +#include +#include + +namespace infinicore::adaptor::denglin_fused_moe { + +/** True when vllm._dl_C.invoke_fused_moe_opt is loadable (same backend as dl-vllm). */ +bool available(); + +/** + * Thin wrapper around Denglin/vLLM ``invoke_fused_moe_opt`` (torch.ops._dl_C). + * Matches vLLM ``dl_fused_moe.dl_invoke_fused_moe`` calling convention. + */ +void invoke_fused_moe_opt( + at::Tensor &a, + at::Tensor &b, + at::Tensor &c, + const std::optional &b_bias, + const std::optional &b_scale, + const std::optional &b_zp, + at::Tensor &topk_weights, + at::Tensor &topk_ids, + at::Tensor &sorted_token_ids, + at::Tensor &expert_ids, + at::Tensor &num_tokens_post_padded, + bool mul_routed_weight, + int64_t top_k, + int64_t block_size_m, + int64_t block_size_n, + int64_t block_size_k, + bool use_fp8_w8a8, + bool use_int8_w8a16, + bool use_int4_w4a16, + bool use_mxfp4_w4a16, + const std::vector &block_shape, + int64_t m); + +} // namespace infinicore::adaptor::denglin_fused_moe +#endif diff --git a/csrc/infinicore/include/infinicore/graph/graph.hpp b/csrc/infinicore/include/infinicore/graph/graph.hpp index b06249ed2..4a70181e6 100644 --- a/csrc/infinicore/include/infinicore/graph/graph.hpp +++ b/csrc/infinicore/include/infinicore/graph/graph.hpp @@ -102,6 +102,23 @@ class Graph { static void execute(__VA_ARGS__); \ }; +// Same as INFINICORE_GRAPH_OP_CLASS, but keeps the op in eager segments during +// CUDA-graph capture (host D2H/H2D, ATen temps, or non-capturable kernels). +#define INFINICORE_GRAPH_OP_CLASS_CAPTURE_UNSAFE(__OP_NAME__, ...) \ + class __OP_NAME__ : public graph::DispatchableGraphOperator { \ + public: \ + using schema = void (*)(__VA_ARGS__); \ + using plan_schema = void *(*)(__VA_ARGS__); \ + static common::OpDispatcher &plan_dispatcher(); \ + static common::OpDispatcher &run_dispatcher(); \ + static common::OpDispatcher &cleanup_dispatcher(); \ + __OP_NAME__(__VA_ARGS__); \ + static void execute(__VA_ARGS__); \ + bool is_device_graph_capture_safe() const override { \ + return false; \ + } \ + }; + #define INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(__OP_NAME__) \ common::OpDispatcher<__OP_NAME__::plan_schema> &__OP_NAME__::plan_dispatcher() { \ static common::OpDispatcher<__OP_NAME__::plan_schema> dispatcher_; \ @@ -116,15 +133,21 @@ class Graph { return dispatcher_; \ } -#define INFINICORE_GRAPH_OP_DISPATCH(__DEVICE_TYPE__, ...) \ - do { \ - const auto device_type = (__DEVICE_TYPE__); \ - const auto plan = plan_dispatcher().lookup(device_type); \ - const auto runner = run_dispatcher().lookup(device_type); \ - const auto deleter = cleanup_dispatcher().lookup(device_type); \ - planned_meta_ = plan(__VA_ARGS__); \ - runner_ = runner; \ - deleter_ = deleter; \ +#define INFINICORE_GRAPH_OP_DISPATCH(__DEVICE_TYPE__, ...) \ + do { \ + const auto device_type = (__DEVICE_TYPE__); \ + try { \ + const auto plan = plan_dispatcher().lookup(device_type); \ + const auto runner = run_dispatcher().lookup(device_type); \ + const auto deleter = cleanup_dispatcher().lookup(device_type); \ + planned_meta_ = plan(__VA_ARGS__); \ + runner_ = runner; \ + deleter_ = deleter; \ + } catch (const std::exception &e) { \ + throw std::runtime_error(std::string(e.what()) + " [" \ + + std::string(__PRETTY_FUNCTION__) \ + + "]"); \ + } \ } while (0) #define INFINICORE_DETAIL_FIRST_ARG(__FIRST__, ...) __FIRST__ diff --git a/csrc/infinicore/include/infinicore/ops/moe_align.hpp b/csrc/infinicore/include/infinicore/ops/moe_align.hpp index 2ce7ca8f4..7a7031397 100644 --- a/csrc/infinicore/include/infinicore/ops/moe_align.hpp +++ b/csrc/infinicore/include/infinicore/ops/moe_align.hpp @@ -9,8 +9,9 @@ namespace infinicore::op { -INFINICORE_GRAPH_OP_CLASS(MoeAlign, Tensor, Tensor, Tensor, const Tensor &, const size_t, const size_t, const bool); -INFINICORE_GRAPH_OP_CLASS(MoeAlignWithExpertMap, Tensor, Tensor, Tensor, const Tensor &, const Tensor &, const size_t, const size_t, const bool); +// QY uses the ATen host-path implementation (CPU round-trips); not capture-safe. +INFINICORE_GRAPH_OP_CLASS_CAPTURE_UNSAFE(MoeAlign, Tensor, Tensor, Tensor, const Tensor &, const size_t, const size_t, const bool); +INFINICORE_GRAPH_OP_CLASS_CAPTURE_UNSAFE(MoeAlignWithExpertMap, Tensor, Tensor, Tensor, const Tensor &, const Tensor &, const size_t, const size_t, const bool); std::tuple moe_align( const Tensor &topk_ids, diff --git a/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp b/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp index 0a34c734c..8c7a6b677 100644 --- a/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp +++ b/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp @@ -7,16 +7,18 @@ namespace infinicore::op { -INFINICORE_GRAPH_OP_CLASS(MoeFusedDense, - Tensor, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &); +// QY MoE runs through ATen + _dl_C and is not CUDA-graph capture safe. +// Keep it in eager graph segments so FA/gemm can still capture around it. +INFINICORE_GRAPH_OP_CLASS_CAPTURE_UNSAFE(MoeFusedDense, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &, + const Tensor &); Tensor moe_fused_dense( const Tensor &hidden_states, diff --git a/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp b/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp index e796c1dc8..27cba570b 100644 --- a/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp +++ b/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp @@ -9,12 +9,13 @@ namespace infinicore::op { -INFINICORE_GRAPH_OP_CLASS(MoeTopkSigmoid, - Tensor, - Tensor, - const Tensor &, - const Tensor &, - const bool); +// QY uses the ATen implementation; not CUDA-graph capture-safe. +INFINICORE_GRAPH_OP_CLASS_CAPTURE_UNSAFE(MoeTopkSigmoid, + Tensor, + Tensor, + const Tensor &, + const Tensor &, + const bool); std::tuple moe_topk_sigmoid( const Tensor &gating_output, diff --git a/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp b/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp index 5ee5e0d0b..a850903ca 100644 --- a/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp +++ b/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp @@ -6,7 +6,8 @@ namespace infinicore::op { -INFINICORE_GRAPH_OP_CLASS(MulScalar, Tensor, const Tensor &, double); +// Used in MiniCPM5 MoE routing-weight scale; keep eager with MoE segments. +INFINICORE_GRAPH_OP_CLASS_CAPTURE_UNSAFE(MulScalar, Tensor, const Tensor &, double); Tensor mul_scalar(const Tensor &a, double alpha); void mul_scalar_(Tensor c, const Tensor &a, double alpha); diff --git a/csrc/infinicore/include/infinicore/ops/rope.hpp b/csrc/infinicore/include/infinicore/ops/rope.hpp index 8fd630ce1..3f4253e84 100644 --- a/csrc/infinicore/include/infinicore/ops/rope.hpp +++ b/csrc/infinicore/include/infinicore/ops/rope.hpp @@ -8,7 +8,9 @@ namespace infinicore::op { -INFINICORE_GRAPH_OP_CLASS(RoPE, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, infinicore::nn::RoPE::Algo); +// QY longctx RoPE packs via cudaMemcpy*/host sync; keep eager during graph capture +// so FA/gemm can still capture around it (same pattern as MoE ATen/_dl_C). +INFINICORE_GRAPH_OP_CLASS_CAPTURE_UNSAFE(RoPE, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, infinicore::nn::RoPE::Algo); // Internal void rope_(Tensor x_out, diff --git a/csrc/infinicore/src/adaptor/aten_adaptor.cc b/csrc/infinicore/src/adaptor/aten_adaptor.cc index 87eb033f4..55b3a52d2 100644 --- a/csrc/infinicore/src/adaptor/aten_adaptor.cc +++ b/csrc/infinicore/src/adaptor/aten_adaptor.cc @@ -1,6 +1,8 @@ #ifdef ENABLE_ATEN #include "infinicore/adaptor/aten_adaptor.hpp" +#include + namespace infinicore::adaptor { at::Tensor to_aten_tensor(const infinicore::Tensor &t) { @@ -51,6 +53,14 @@ c10::musa::MUSAStream get_musa_stream() { } #endif +void set_aten_stream_to_infinicore() { +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) + c10::cuda::setCurrentCUDAStream(get_cuda_stream()); +#else + throw std::runtime_error("ATen stream synchronization is unsupported on this device"); +#endif +} + } // namespace infinicore::adaptor #endif // ENABLE_ATEN diff --git a/csrc/infinicore/src/adaptor/denglin_fused_moe_adaptor.cc b/csrc/infinicore/src/adaptor/denglin_fused_moe_adaptor.cc new file mode 100644 index 000000000..d7798072f --- /dev/null +++ b/csrc/infinicore/src/adaptor/denglin_fused_moe_adaptor.cc @@ -0,0 +1,132 @@ +#if defined(ENABLE_ATEN) && defined(ENABLE_QY_API) +#include "infinicore/adaptor/denglin_fused_moe_adaptor.hpp" +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include +#include +#include +#include +#include +#include + +namespace infinicore::adaptor::denglin_fused_moe { +namespace { + +// Mangled name of free function in vllm._dl_C: +// invoke_fused_moe_opt(at::Tensor&, at::Tensor&, at::Tensor&, +// optional const& x3, at::Tensor& x5, bool, long x4, bool x4, +// vector const&, long) +using invoke_fused_moe_opt_fn = void (*)( + at::Tensor &, + at::Tensor &, + at::Tensor &, + const std::optional &, + const std::optional &, + const std::optional &, + at::Tensor &, + at::Tensor &, + at::Tensor &, + at::Tensor &, + at::Tensor &, + bool, + int64_t, + int64_t, + int64_t, + int64_t, + bool, + bool, + bool, + bool, + const std::vector &, + int64_t); + +struct Symbols { + void *handle = nullptr; + invoke_fused_moe_opt_fn invoke = nullptr; + std::string error; +}; + +Symbols &symbols() { + static Symbols syms; + static std::once_flag once; + std::call_once(once, []() { + const char *env_path = std::getenv("INFINICORE_DENGLIN_DL_C_SO"); + // Prefer exact override, then common dl-vLLM image layouts (3.12/3.10). + const char *paths[] = { + env_path, + "/usr/local/lib/python3.12/dist-packages/vllm/_dl_C.cpython-312-x86_64-linux-gnu.so", + "/usr/local/lib/python3.10/dist-packages/vllm/_dl_C.cpython-310-x86_64-linux-gnu.so", + "/usr/lib/python3/dist-packages/vllm/_dl_C.cpython-312-x86_64-linux-gnu.so", + }; + constexpr const char *kSym = "_Z20invoke_fused_moe_optRN2at6TensorES1_S1_RKSt8optionalIS0_ES5_S5_" + "S1_S1_S1_S1_S1_bllllbbbbRKSt6vectorIlSaIlEEl"; + + for (const char *path : paths) { + if (path == nullptr || path[0] == '\0') { + continue; + } + void *handle = dlopen(path, RTLD_LAZY | RTLD_LOCAL); + if (!handle) { + syms.error = dlerror() ? dlerror() : "dlopen failed"; + continue; + } + auto fn = reinterpret_cast(dlsym(handle, kSym)); + if (!fn) { + syms.error = dlerror() ? dlerror() : "dlsym invoke_fused_moe_opt failed"; + dlclose(handle); + continue; + } + syms.handle = handle; + syms.invoke = fn; + syms.error.clear(); + return; + } + }); + return syms; +} + +} // namespace + +bool available() { + return symbols().invoke != nullptr; +} + +void invoke_fused_moe_opt( + at::Tensor &a, + at::Tensor &b, + at::Tensor &c, + const std::optional &b_bias, + const std::optional &b_scale, + const std::optional &b_zp, + at::Tensor &topk_weights, + at::Tensor &topk_ids, + at::Tensor &sorted_token_ids, + at::Tensor &expert_ids, + at::Tensor &num_tokens_post_padded, + bool mul_routed_weight, + int64_t top_k, + int64_t block_size_m, + int64_t block_size_n, + int64_t block_size_k, + bool use_fp8_w8a8, + bool use_int8_w8a16, + bool use_int4_w4a16, + bool use_mxfp4_w4a16, + const std::vector &block_shape, + int64_t m) { + auto &syms = symbols(); + if (!syms.invoke) { + throw std::runtime_error( + "Denglin fused MoE unavailable (vllm._dl_C.invoke_fused_moe_opt): " + syms.error); + } + infinicore::adaptor::set_aten_stream_to_infinicore(); + syms.invoke( + a, b, c, b_bias, b_scale, b_zp, + topk_weights, topk_ids, sorted_token_ids, expert_ids, num_tokens_post_padded, + mul_routed_weight, top_k, block_size_m, block_size_n, block_size_k, + use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16, use_mxfp4_w4a16, + block_shape, m); +} + +} // namespace infinicore::adaptor::denglin_fused_moe +#endif diff --git a/csrc/infinicore/src/context/context_impl.cc b/csrc/infinicore/src/context/context_impl.cc index a243eb58f..b5622f366 100644 --- a/csrc/infinicore/src/context/context_impl.cc +++ b/csrc/infinicore/src/context/context_impl.cc @@ -20,6 +20,7 @@ constexpr std::array(Device::Type::kCount)> kD Device::Type::kIluvatar, Device::Type::kHygon, Device::Type::kThead, + Device::Type::kQy, Device::Type::kCpu, }; @@ -237,6 +238,7 @@ ContextImpl::ContextImpl() { initializeDeviceType(); initializeDeviceType(); initializeDeviceType(); + initializeDeviceType(); const bool has_runtime = std::any_of( runtime_table_.begin(), runtime_table_.end(), diff --git a/csrc/infinicore/src/nn/rope.cc b/csrc/infinicore/src/nn/rope.cc index aced91856..9b5785a91 100644 --- a/csrc/infinicore/src/nn/rope.cc +++ b/csrc/infinicore/src/nn/rope.cc @@ -78,6 +78,8 @@ void RoPE::initialize_cache() { INFINICORE_NN_BUFFER_INIT(cos_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); #ifdef ENABLE_INFINIOPS_API + // QY: prefer sin/cos + ATen RoPE. InfiniOps RotaryEmbedding no-ops on 4D + // [B,S,H,D] (MiniCPM5 static) and 3D reshape still mismatches HF GPT-NeoX. if ((device_.type() == Device::Type::kNvidia || device_.type() == Device::Type::kMetax || device_.type() == Device::Type::kIluvatar diff --git a/csrc/infinicore/src/ops/cat/cat.cc b/csrc/infinicore/src/ops/cat/cat.cc index c19e0896b..dfe54d859 100644 --- a/csrc/infinicore/src/ops/cat/cat.cc +++ b/csrc/infinicore/src/ops/cat/cat.cc @@ -14,7 +14,8 @@ bool use_slice_copy_cat(Device::Type device_type, int dim, int ndim) { && (device_type == Device::Type::kNvidia || device_type == Device::Type::kHygon || device_type == Device::Type::kIluvatar - || device_type == Device::Type::kThead); + || device_type == Device::Type::kThead + || device_type == Device::Type::kQy); } class CatInfo { diff --git a/csrc/infinicore/src/ops/causal_softmax/causal_softmax_aten.cc b/csrc/infinicore/src/ops/causal_softmax/causal_softmax_aten.cc new file mode 100644 index 000000000..6bdc3af03 --- /dev/null +++ b/csrc/infinicore/src/ops/causal_softmax/causal_softmax_aten.cc @@ -0,0 +1,77 @@ +#include "infinicore/ops/causal_softmax.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include + +#include +#include + +namespace infinicore::op::causal_softmax_impl::aten { +namespace { + +struct PlannedMeta { + Tensor output; + Tensor input; +}; + +void *plan(Tensor output, const Tensor &input) { + INFINICORE_ASSERT(output->device().type() == Device::Type::kQy); + return new PlannedMeta{output, input}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + using namespace ::infinicore::adaptor; + set_aten_stream_to_infinicore(); + + auto input = to_aten_tensor(p->input); + auto output = to_aten_tensor(p->output); + INFINICORE_ASSERT(input.dim() == 2 || input.dim() == 3); + + const auto seq_len = input.size(-2); + const auto total_seq_len = input.size(-1); + INFINICORE_ASSERT(seq_len <= total_seq_len); + + const bool reuses_input = input.is_same(output); + const bool needs_mask = seq_len != 1 || !reuses_input; + + at::Tensor scores; + if (needs_mask) { + const int64_t score_diagonal = total_seq_len - seq_len; + const int64_t mask_diagonal = total_seq_len - seq_len + 1; + // Match InfiniOps causal_softmax: tril(scores) + triu(-inf mask). + auto tril_scores = at::tril(input, score_diagonal); + auto mask = at::full({seq_len, total_seq_len}, + -std::numeric_limits::infinity(), + input.options()); + mask = at::triu(mask, mask_diagonal); + // Broadcast mask over leading dims if present. + scores = tril_scores + mask; + } else { + scores = input; + } + + auto probs = at::softmax(scores, /*dim=*/-1); + output.copy_(probs); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + CausalSoftmax::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + CausalSoftmax::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + CausalSoftmax::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::causal_softmax_impl::aten +#endif diff --git a/csrc/infinicore/src/ops/infiniops_impl.hpp b/csrc/infinicore/src/ops/infiniops_impl.hpp index a9051f1dc..3b8483689 100644 --- a/csrc/infinicore/src/ops/infiniops_impl.hpp +++ b/csrc/infinicore/src/ops/infiniops_impl.hpp @@ -65,6 +65,8 @@ inline infini::ops::Device toInfiniOpsDevice(const Device &device) { return infini::ops::Device{infini::ops::Device::Type::kHygon, static_cast(device.index())}; case Device::Type::kThead: return infini::ops::Device{infini::ops::Device::Type::kThead, static_cast(device.index())}; + case Device::Type::kQy: + return infini::ops::Device{infini::ops::Device::Type::kQy, static_cast(device.index())}; default: throw std::runtime_error("InfiniOps backend does not support this device type."); } @@ -80,6 +82,7 @@ inline bool isSupportedDevice(Device::Type device_type) { case Device::Type::kAscend: case Device::Type::kHygon: case Device::Type::kThead: + case Device::Type::kQy: return true; default: return false; @@ -148,6 +151,7 @@ void registerSupportedDevices(Dispatcher &dispatcher, Function function) { dispatcher.registerDevice(Device::Type::kAscend, function); dispatcher.registerDevice(Device::Type::kHygon, function); dispatcher.registerDevice(Device::Type::kThead, function); + dispatcher.registerDevice(Device::Type::kQy, function); } struct TensorMeta { diff --git a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc index 12fdf2b74..6514124d8 100644 --- a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc +++ b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc @@ -1,6 +1,8 @@ #include "infinicore/ops/mha_kvcache.hpp" #include "../../utils.hpp" +#include + namespace infinicore::op { INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MhaKVCache); @@ -70,6 +72,26 @@ bool MhaKVCache::supports_device_graph_capture( && k_cache->size(3) == 128 && seqlens_k->shape() == Shape({16}) && block_table->shape() == Shape({16, 512}); + // MiniCPM5-2B TP1 decode: fused QKV Q view, paged KV [NUM_BLOCKS,256,2,128]. + const bool minicpm5_2b_shape = q->shape() == Shape({1, 1, 16, 128}) + && out->shape() == Shape({1, 1, 16, 128}) + && k_cache->size(0) > 0 + && k_cache->size(1) == 256 + && k_cache->size(2) == 2 + && k_cache->size(3) == 128 + && seqlens_k->shape() == Shape({1}) + && block_table->size(0) == 1 + && block_table->size(1) > 0; + // MiniCPM5-16A3 TP1 decode: gated attention Q|gate interleaved head stride. + const bool minicpm5_16a3_shape = q->shape() == Shape({1, 1, 32, 128}) + && out->shape() == Shape({1, 1, 32, 128}) + && k_cache->size(0) > 0 + && k_cache->size(1) == 256 + && k_cache->size(2) == 2 + && k_cache->size(3) == 128 + && seqlens_k->shape() == Shape({1}) + && block_table->size(0) == 1 + && block_table->size(1) > 0; const bool p13_layout = q->is_contiguous() && out->is_contiguous(); // P09 reads Q directly from the fused QKV projection while output is dense. const bool p09_layout = q->strides() == Strides({5120, 5120, 128, 1}) @@ -83,33 +105,75 @@ bool MhaKVCache::supports_device_graph_capture( // P12 has the same fused QKV projection layout at batch size one. const bool p12_layout = q->strides() == Strides({4608, 4608, 128, 1}) && out->strides() == Strides({4096, 4096, 128, 1}); + // MiniCPM5-2B: QKV width (16+2+2)*128=2560; dense out width 16*128=2048. + const bool minicpm5_2b_layout + = (q->strides() == Strides({2560, 2560, 128, 1}) + && out->strides() == Strides({2048, 2048, 128, 1})) + || (q->is_contiguous() && out->is_contiguous()); + // MiniCPM5-16A3 gated: QKV width (64+2+2)*128=8704 with per-head Q|gate + // stride 256; dense out width 32*128=4096. + const bool minicpm5_16a3_layout + = (q->strides() == Strides({8704, 8704, 256, 1}) + && out->strides() == Strides({4096, 4096, 128, 1})) + || (q->is_contiguous() && out->is_contiguous()); const bool reviewed_shape_and_layout = (p12_shape && p12_layout) || (p13_shape && p13_layout) || (p09_shape && p09_layout) || (p11_shape && p11_layout) - || (p14_shape && p14_layout); + || (p14_shape && p14_layout) + || (minicpm5_2b_shape && minicpm5_2b_layout) + || (minicpm5_16a3_shape && minicpm5_16a3_layout); const auto dtype = q->dtype(); - return device.type() == Device::Type::kNvidia - && same_device(q) - && same_device(k_cache) - && same_device(v_cache) - && same_device(seqlens_k) - && same_device(block_table) - && reviewed_shape_and_layout - && out->shape() == q->shape() - && k_cache->size(0) > 0 - && v_cache->shape() == k_cache->shape() - && (dtype == DataType::kFloat16 - || dtype == DataType::kBFloat16) - && out->dtype() == dtype - && k_cache->dtype() == dtype - && v_cache->dtype() == dtype - && seqlens_k->dtype() == DataType::kInt32 - && block_table->dtype() == DataType::kInt32 - && k_cache->is_contiguous() - && v_cache->is_contiguous() - && seqlens_k->is_contiguous() - && block_table->is_contiguous(); + const bool supported_device = device.type() == Device::Type::kNvidia + || device.type() == Device::Type::kQy; + const bool ok = supported_device + && same_device(q) + && same_device(k_cache) + && same_device(v_cache) + && same_device(seqlens_k) + && same_device(block_table) + && reviewed_shape_and_layout + && out->shape() == q->shape() + && k_cache->size(0) > 0 + && v_cache->shape() == k_cache->shape() + && (dtype == DataType::kFloat16 + || dtype == DataType::kBFloat16) + && out->dtype() == dtype + && k_cache->dtype() == dtype + && v_cache->dtype() == dtype + && seqlens_k->dtype() == DataType::kInt32 + && block_table->dtype() == DataType::kInt32 + && k_cache->is_contiguous() + && v_cache->is_contiguous() + && seqlens_k->is_contiguous() + && block_table->is_contiguous(); + if (!ok + && device.type() == Device::Type::kQy + && (minicpm5_2b_shape || minicpm5_16a3_shape)) { + static bool logged = false; + if (!logged) { + logged = true; + const auto &qs = q->shape(); + const auto &qstr = q->strides(); + const auto &ostr = out->strides(); + const auto &ks = k_cache->shape(); + const auto &bts = block_table->shape(); + std::fprintf( + stderr, + "[mha_kvcache] QY graph-capture allowlist miss: " + "q=[%ld,%ld,%ld,%ld] q_strides=[%ld,%ld,%ld,%ld] " + "out_strides=[%ld,%ld,%ld,%ld] k_cache=[%ld,%ld,%ld,%ld] " + "block_table=[%ld,%ld] contiguous(q/out)=%d/%d\n", + (long)qs[0], (long)qs[1], (long)qs[2], (long)qs[3], + (long)qstr[0], (long)qstr[1], (long)qstr[2], (long)qstr[3], + (long)ostr[0], (long)ostr[1], (long)ostr[2], (long)ostr[3], + (long)ks[0], (long)ks[1], (long)ks[2], (long)ks[3], + (long)bts[0], (long)bts[1], + q->is_contiguous() ? 1 : 0, + out->is_contiguous() ? 1 : 0); + } + } + return ok; } MhaKVCache::MhaKVCache(Tensor out, diff --git a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc index 7a03a736d..15268fada 100644 --- a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc +++ b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc @@ -45,7 +45,8 @@ bool is_supported(const Tensor &out, && device_type != Device::Type::kCambricon && device_type != Device::Type::kIluvatar && device_type != Device::Type::kHygon - && device_type != Device::Type::kAscend) + && device_type != Device::Type::kAscend + && device_type != Device::Type::kQy) || q->ndim() != 4 || out->ndim() != 4 || k_cache->ndim() != 4 @@ -119,9 +120,11 @@ struct PlannedMeta { std::unique_ptr make_graph_safe_provider( PlannedMeta &planned) { + const auto device_type = planned.q.device.type(); + const auto implementation_index = implementation_index_for_device(device_type); const auto config = ::infinicore::op::infiniops::configForImplementation< infini::ops::FlashAttnWithKvcache>( - infini::ops::Device::Type::kNvidia, 16); + device_type, implementation_index); const auto q = planned.q.tensor(planned.q_tensor); const auto k_cache = planned.k_cache.tensor(planned.k_cache_tensor); @@ -323,6 +326,9 @@ static bool registered = []() { MhaKVCache::plan_dispatcher().registerDevice(Device::Type::kAscend, &plan); MhaKVCache::run_dispatcher().registerDevice(Device::Type::kAscend, &run); MhaKVCache::cleanup_dispatcher().registerDevice(Device::Type::kAscend, &cleanup); + MhaKVCache::plan_dispatcher().registerDevice(Device::Type::kQy, &plan); + MhaKVCache::run_dispatcher().registerDevice(Device::Type::kQy, &run); + MhaKVCache::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup); return true; }(); diff --git a/csrc/infinicore/src/ops/moe_align/moe_align_aten.cc b/csrc/infinicore/src/ops/moe_align/moe_align_aten.cc new file mode 100644 index 000000000..5db3e7af6 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_align/moe_align_aten.cc @@ -0,0 +1,235 @@ +#include "infinicore/ops/moe_align.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace infinicore::op::moe_align_impl::aten { +namespace { + +bool stream_is_capturing(cudaStream_t stream) { + cudaStreamCaptureStatus status = cudaStreamCaptureStatusNone; + auto st = cudaStreamIsCapturing(stream, &status); + return st == cudaSuccess && status != cudaStreamCaptureStatusNone; +} + +cudaError_t sync_unless_capturing(cudaStream_t stream) { + if (stream_is_capturing(stream)) { + return cudaGetLastError(); + } + return cudaStreamSynchronize(stream); +} + +bool host_timer_enabled() { + static int cached = -1; + if (cached < 0) { + const char *env = std::getenv("INFINI_MOE_HOST_TIMER"); + cached = (env != nullptr && env[0] == '1' && env[1] == '\0') ? 1 : 0; + } + return cached == 1; +} + +struct HostTimer { + const char *op; + int64_t m; + std::chrono::steady_clock::time_point t0; + bool active; + + HostTimer(const char *op_name, int64_t m_val) + : op(op_name), m(m_val), active(host_timer_enabled()) { + if (active) { + t0 = std::chrono::steady_clock::now(); + } + } + + ~HostTimer() { + if (!active) { + return; + } + auto t1 = std::chrono::steady_clock::now(); + double ms = std::chrono::duration(t1 - t0).count(); + std::fprintf(stderr, "host_timer op=%s ms=%.3f m=%lld\n", op, ms, + static_cast(m)); + } +}; + +at::Tensor ensure_i32_contig(const at::Tensor &t) { + INFINICORE_ASSERT(t.is_contiguous()); + if (t.scalar_type() == at::kInt) { + return t; + } + INFINICORE_ASSERT(t.scalar_type() == at::kLong); + // Small [M*topk] cast — avoid ATen .to() OOR on Denglin. + auto out = at::empty(t.sizes(), t.options().dtype(at::kInt)); + const int64_t n = t.numel(); + std::vector host(static_cast(n)); + cudaMemcpy(host.data(), t.data_ptr(), + static_cast(n * static_cast(sizeof(int64_t))), + cudaMemcpyDeviceToHost); + std::vector host32(static_cast(n)); + for (int64_t i = 0; i < n; ++i) { + host32[static_cast(i)] = static_cast(host[static_cast(i)]); + } + cudaMemcpy(out.data_ptr(), host32.data(), + static_cast(n * static_cast(sizeof(int32_t))), + cudaMemcpyHostToDevice); + return out; +} + +void run_moe_align_gpu(at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded, + at::Tensor topk_ids, + const at::Tensor *expert_map, + int64_t num_experts, + int64_t block_size, + bool pad_sorted_token_ids) { + INFINICORE_ASSERT(sorted_token_ids.is_contiguous()); + INFINICORE_ASSERT(expert_ids.is_contiguous()); + INFINICORE_ASSERT(num_tokens_post_padded.is_contiguous()); + INFINICORE_ASSERT(topk_ids.is_contiguous()); + + auto topk_i32 = ensure_i32_contig(topk_ids); + at::Tensor map_i32; + const void *map_ptr = nullptr; + if (expert_map != nullptr) { + INFINICORE_ASSERT(expert_map->is_contiguous()); + map_i32 = ensure_i32_contig(*expert_map); + map_ptr = map_i32.data_ptr(); + } + + auto workspace = at::empty( + {static_cast( + ::infini::ops::qy::moe_align_workspace_bytes( + static_cast(num_experts)) + / sizeof(int32_t))}, + sorted_token_ids.options().dtype(at::kInt)); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + ::infini::ops::qy::moe_align_launch( + sorted_token_ids.data_ptr(), expert_ids.data_ptr(), + num_tokens_post_padded.data_ptr(), topk_i32.data_ptr(), map_ptr, + static_cast(topk_i32.numel()), + static_cast(num_experts), static_cast(block_size), + static_cast(sorted_token_ids.numel()), + static_cast(expert_ids.numel()), pad_sorted_token_ids, + workspace.data_ptr(), reinterpret_cast(stream)); + auto err = sync_unless_capturing(stream); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("QY moe_align sync: ") + cudaGetErrorString(err)); + } +} + +struct PlannedMeta { + Tensor sorted_token_ids; + Tensor expert_ids; + Tensor num_tokens_post_padded; + Tensor topk_ids; + Tensor expert_map; + bool has_expert_map; + size_t num_experts; + size_t block_size; + bool pad_sorted_token_ids; +}; + +void *plan_no_map(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) { + INFINICORE_ASSERT(sorted_token_ids->device().type() == Device::Type::kQy); + return new PlannedMeta{ + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + Tensor{}, + false, + num_experts, + block_size, + pad_sorted_token_ids}; +} + +void *plan_with_map(Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + const Tensor &expert_map, + const size_t num_experts, + const size_t block_size, + const bool pad_sorted_token_ids) { + INFINICORE_ASSERT(sorted_token_ids->device().type() == Device::Type::kQy); + return new PlannedMeta{ + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + expert_map, + true, + num_experts, + block_size, + pad_sorted_token_ids}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + using namespace ::infinicore::adaptor; + set_aten_stream_to_infinicore(); + + HostTimer timer("moe_align", static_cast(p->topk_ids->numel())); + + at::Tensor expert_map_aten; + const at::Tensor *expert_map_ptr = nullptr; + if (p->has_expert_map) { + expert_map_aten = to_aten_tensor(p->expert_map); + expert_map_ptr = &expert_map_aten; + } + + run_moe_align_gpu( + to_aten_tensor(p->sorted_token_ids), + to_aten_tensor(p->expert_ids), + to_aten_tensor(p->num_tokens_post_padded), + to_aten_tensor(p->topk_ids), + expert_map_ptr, + static_cast(p->num_experts), + static_cast(p->block_size), + p->pad_sorted_token_ids); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + MoeAlign::plan_dispatcher().registerDevice(Device::Type::kQy, &plan_no_map, true); + MoeAlign::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MoeAlign::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + + MoeAlignWithExpertMap::plan_dispatcher().registerDevice(Device::Type::kQy, &plan_with_map, true); + MoeAlignWithExpertMap::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MoeAlignWithExpertMap::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::moe_align_impl::aten +#endif diff --git a/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc index dea77265e..b0f4f9111 100644 --- a/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc +++ b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc @@ -18,9 +18,9 @@ MoeFusedDense::MoeFusedDense(Tensor output, INFINICORE_ASSERT_TENSORS_SAME_DEVICE( output, hidden_states, w13, w2, topk_weights, topk_ids, sorted_token_ids, expert_ids, num_tokens_post_padded); - throw std::runtime_error( - "`moe_fused_dense` requires an InfiniOps dense expert provider; " - "quantized WNA16/Marlin providers do not accept floating-point expert weights"); + INFINICORE_GRAPH_OP_DISPATCH( + output->device().type(), output, hidden_states, w13, w2, topk_weights, topk_ids, + sorted_token_ids, expert_ids, num_tokens_post_padded); } void MoeFusedDense::execute(Tensor output, diff --git a/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense_qy_dl.cc b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense_qy_dl.cc new file mode 100644 index 000000000..5ce1a2514 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense_qy_dl.cc @@ -0,0 +1,243 @@ +#if defined(ENABLE_ATEN) && defined(ENABLE_QY_API) + +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/denglin_fused_moe_adaptor.hpp" +#include "infinicore/graph/graph.hpp" +#include "infinicore/ops/moe_fused_dense.hpp" + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace infinicore::op::moe_fused_dense_impl::qy_dl { +namespace { + +bool stream_is_capturing(cudaStream_t stream) { + cudaStreamCaptureStatus status = cudaStreamCaptureStatusNone; + auto st = cudaStreamIsCapturing(stream, &status); + return st == cudaSuccess && status != cudaStreamCaptureStatusNone; +} + +cudaError_t sync_unless_capturing(cudaStream_t stream) { + if (stream_is_capturing(stream)) { + return cudaGetLastError(); + } + return cudaStreamSynchronize(stream); +} + +cudaError_t sync_device_unless_capturing(cudaStream_t stream) { + if (stream_is_capturing(stream)) { + return cudaGetLastError(); + } + return cudaDeviceSynchronize(); +} + +bool host_timer_enabled() { + static int cached = -1; + if (cached < 0) { + const char *env = std::getenv("INFINI_MOE_HOST_TIMER"); + cached = (env != nullptr && env[0] == '1' && env[1] == '\0') ? 1 : 0; + } + return cached == 1; +} + +struct HostTimer { + const char *op; + int64_t m; + std::chrono::steady_clock::time_point t0; + bool active; + + HostTimer(const char *op_name, int64_t m_val) + : op(op_name), m(m_val), active(host_timer_enabled()) { + if (active) { + t0 = std::chrono::steady_clock::now(); + } + } + + ~HostTimer() { + if (!active) { + return; + } + auto t1 = std::chrono::steady_clock::now(); + double ms = std::chrono::duration(t1 - t0).count(); + std::fprintf(stderr, "host_timer op=%s ms=%.3f m=%lld\n", op, ms, + static_cast(m)); + } +}; + +} // namespace + +struct PlannedMeta { + graph::GraphTensor output; + graph::GraphTensor hidden_states; + graph::GraphTensor w13; + graph::GraphTensor w2; + graph::GraphTensor topk_weights; + graph::GraphTensor topk_ids; + graph::GraphTensor sorted_token_ids; + graph::GraphTensor expert_ids; + graph::GraphTensor num_tokens_post_padded; +}; + +void *plan(Tensor output, + const Tensor &hidden_states, + const Tensor &w13, + const Tensor &w2, + const Tensor &topk_weights, + const Tensor &topk_ids, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded) { + if (!adaptor::denglin_fused_moe::available()) { + throw std::runtime_error( + "QY moe_fused_dense requires vllm._dl_C (set INFINICORE_DENGLIN_DL_C_SO if needed)"); + } + return new PlannedMeta{ + graph::GraphTensor(output), + graph::GraphTensor(hidden_states), + graph::GraphTensor(w13), + graph::GraphTensor(w2), + graph::GraphTensor(topk_weights), + graph::GraphTensor(topk_ids), + graph::GraphTensor(sorted_token_ids), + graph::GraphTensor(expert_ids), + graph::GraphTensor(num_tokens_post_padded)}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + adaptor::set_aten_stream_to_infinicore(); + + auto hidden = adaptor::to_aten_tensor(p->hidden_states); + auto w13 = adaptor::to_aten_tensor(p->w13); + auto w2 = adaptor::to_aten_tensor(p->w2); + auto topk_weights = adaptor::to_aten_tensor(p->topk_weights); + auto topk_ids = adaptor::to_aten_tensor(p->topk_ids); + auto sorted_token_ids = adaptor::to_aten_tensor(p->sorted_token_ids); + auto expert_ids = adaptor::to_aten_tensor(p->expert_ids); + auto num_tokens_post_padded = adaptor::to_aten_tensor(p->num_tokens_post_padded); + auto output = adaptor::to_aten_tensor(p->output); + + const int64_t m = hidden.size(0); + const int64_t h = hidden.size(1); + const int64_t topk = topk_ids.size(1); + const int64_t inter2 = w13.size(1); + const int64_t inter = inter2 / 2; + + // Workspace capacity rounding makes sorted_cap/expert_cap recovery inexact + // (e.g. 17 instead of 16 for typical prefill). MiniCPM5 / CudaFusedMoeRunner + // default moe_align_block_size is 16 — use that for _dl_C. + constexpr int64_t block_size_m = 16; + constexpr int64_t block_size_n = 64; + constexpr int64_t block_size_k = 64; + const std::vector block_shape; // empty == vLLM default + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + // Stage 1: hidden @ w13 -> [M, topk, 2*inter] (vLLM fused_experts_impl) + auto cache1 = at::empty({m, topk, inter2}, hidden.options()); + auto a1 = hidden.view({-1, h}); + auto c1 = cache1.view({-1, topk, inter2}); + { + HostTimer timer("stage1", m); + adaptor::denglin_fused_moe::invoke_fused_moe_opt( + a1, w13, c1, + std::nullopt, std::nullopt, std::nullopt, + topk_weights, topk_ids, sorted_token_ids, expert_ids, num_tokens_post_padded, + /*mul_routed_weight=*/false, + topk, block_size_m, block_size_n, block_size_k, + false, false, false, false, + block_shape, m); + auto err = sync_device_unless_capturing(stream); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_fused_dense stage1 (w13) OOR/sync: ") + + cudaGetErrorString(err)); + } + } + + // SiLU/GLU activation via thin QY kernel (avoids D2H/H2D and ATen OOR). + auto flat_gate_up = cache1.reshape({m * topk, inter2}).contiguous(); + auto activated = at::empty({m * topk, inter}, hidden.options()); + { + HostTimer timer("silu", m); + ::infini::ops::qy::silu_and_mul_launch( + activated.data_ptr(), flat_gate_up.data_ptr(), + static_cast(m * topk), static_cast(inter), + static_cast(p->hidden_states->dtype()), + reinterpret_cast(stream)); + auto err = sync_unless_capturing(stream); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_fused_dense silu sync: ") + cudaGetErrorString(err)); + } + } + + // Stage 2: activated @ w2 -> [M, topk, H], top_k=1, mul routed weights + auto cache3 = at::empty({m, topk, h}, hidden.options()); + auto a2 = activated.contiguous(); + auto c2 = cache3.view({-1, /*top_k=*/1, h}); + { + HostTimer timer("stage2", m); + adaptor::denglin_fused_moe::invoke_fused_moe_opt( + a2, w2, c2, + std::nullopt, std::nullopt, std::nullopt, + topk_weights, topk_ids, sorted_token_ids, expert_ids, num_tokens_post_padded, + /*mul_routed_weight=*/true, + /*top_k=*/1, block_size_m, block_size_n, block_size_k, + false, false, false, false, + block_shape, m); + auto err = sync_device_unless_capturing(stream); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_fused_dense stage2 (w2) OOR/sync: ") + + cudaGetErrorString(err)); + } + } + + // Expert combine via thin QY moe_sum (arbitrary topk). + auto cache3_contig = cache3.is_contiguous() ? cache3 : cache3.contiguous(); + if (!output.is_contiguous()) { + throw std::runtime_error( + "QY moe_fused_dense combine requires contiguous output"); + } + { + HostTimer timer("combine", m); + ::infini::ops::qy::moe_sum_launch( + output.data_ptr(), cache3_contig.data_ptr(), + static_cast(m), static_cast(topk), + static_cast(h), static_cast(p->hidden_states->dtype()), + reinterpret_cast(stream)); + auto err = sync_unless_capturing(stream); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_fused_dense combine sync: ") + cudaGetErrorString(err)); + } + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + // Override the ALLDEVICE infiniop registration for Denglin QY. + MoeFusedDense::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + MoeFusedDense::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MoeFusedDense::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::moe_fused_dense_impl::qy_dl + +#endif diff --git a/csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid_aten.cc b/csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid_aten.cc new file mode 100644 index 000000000..3d8697a60 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid_aten.cc @@ -0,0 +1,161 @@ +#include "infinicore/ops/moe_topk_sigmoid.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace infinicore::op::moe_topk_sigmoid_impl::aten { +namespace { + +bool stream_is_capturing(cudaStream_t stream) { + cudaStreamCaptureStatus status = cudaStreamCaptureStatusNone; + auto st = cudaStreamIsCapturing(stream, &status); + return st == cudaSuccess && status != cudaStreamCaptureStatusNone; +} + +cudaError_t sync_unless_capturing(cudaStream_t stream) { + if (stream_is_capturing(stream)) { + return cudaGetLastError(); + } + return cudaStreamSynchronize(stream); +} + +bool host_timer_enabled() { + static int cached = -1; + if (cached < 0) { + const char *env = std::getenv("INFINI_MOE_HOST_TIMER"); + cached = (env != nullptr && env[0] == '1' && env[1] == '\0') ? 1 : 0; + } + return cached == 1; +} + +struct HostTimer { + const char *op; + int64_t m; + std::chrono::steady_clock::time_point t0; + bool active; + + HostTimer(const char *op_name, int64_t m_val) + : op(op_name), m(m_val), active(host_timer_enabled()) { + if (active) { + t0 = std::chrono::steady_clock::now(); + } + } + + ~HostTimer() { + if (!active) { + return; + } + auto t1 = std::chrono::steady_clock::now(); + double ms = std::chrono::duration(t1 - t0).count(); + std::fprintf(stderr, "host_timer op=%s ms=%.3f m=%lld\n", op, ms, + static_cast(m)); + } +}; + +struct PlannedMeta { + Tensor topk_weights; + Tensor topk_indices; + Tensor gating_output; + Tensor correction_bias; + bool renormalize; +}; + +void *plan(Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias, + const bool renormalize) { + INFINICORE_ASSERT(topk_weights->device().type() == Device::Type::kQy); + return new PlannedMeta{ + topk_weights, + topk_indices, + gating_output, + correction_bias, + renormalize}; +} + +void run(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + using namespace ::infinicore::adaptor; + set_aten_stream_to_infinicore(); + + auto gating_dev = to_aten_tensor(planned->gating_output); + auto weights_dev = to_aten_tensor(planned->topk_weights); + auto indices_dev = to_aten_tensor(planned->topk_indices); + + INFINICORE_ASSERT(gating_dev.is_contiguous()); + INFINICORE_ASSERT(weights_dev.is_contiguous()); + INFINICORE_ASSERT(indices_dev.is_contiguous()); + INFINICORE_ASSERT(gating_dev.dim() == 2); + INFINICORE_ASSERT(weights_dev.dim() == 2); + + const auto tokens = gating_dev.size(0); + const auto experts = gating_dev.size(1); + const auto topk = weights_dev.size(1); + INFINICORE_ASSERT(indices_dev.size(0) == tokens); + INFINICORE_ASSERT(indices_dev.size(1) == topk); + + HostTimer timer("moe_topk_sigmoid", tokens); + + const void *bias_ptr = nullptr; + int bias_dtype_tag = 0; + if (planned->correction_bias) { + auto bias = to_aten_tensor(planned->correction_bias); + INFINICORE_ASSERT(bias.is_contiguous()); + INFINICORE_ASSERT(bias.numel() == experts); + // f32/bf16/fp16 bias consumed on device (kernel converts to float). + if (bias.scalar_type() != at::kFloat && bias.scalar_type() != at::kHalf && bias.scalar_type() != at::kBFloat16) { + throw std::runtime_error( + "QY moe_topk_sigmoid: unsupported correction_bias dtype"); + } + bias_ptr = bias.data_ptr(); + bias_dtype_tag = static_cast(planned->correction_bias->dtype()); + } + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + // Kernel writes the actual weight dtype (f16/bf16/f32) and index dtype + // (i32/i64) — no D2H/H2D post-cast. + ::infini::ops::qy::moe_topk_sigmoid_launch( + weights_dev.data_ptr(), indices_dev.data_ptr(), gating_dev.data_ptr(), + bias_ptr, static_cast(tokens), static_cast(experts), + static_cast(topk), planned->renormalize, + static_cast(planned->gating_output->dtype()), + static_cast(planned->topk_weights->dtype()), + static_cast(planned->topk_indices->dtype()), bias_dtype_tag, + reinterpret_cast(stream)); + auto err = sync_unless_capturing(stream); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("QY moe_topk_sigmoid sync: ") + cudaGetErrorString(err)); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + MoeTopkSigmoid::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + MoeTopkSigmoid::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MoeTopkSigmoid::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::moe_topk_sigmoid_impl::aten +#endif diff --git a/csrc/infinicore/src/ops/mul_scalar/mul_scalar_aten.cc b/csrc/infinicore/src/ops/mul_scalar/mul_scalar_aten.cc new file mode 100644 index 000000000..4f407c471 --- /dev/null +++ b/csrc/infinicore/src/ops/mul_scalar/mul_scalar_aten.cc @@ -0,0 +1,58 @@ +#include "infinicore/ops/mul_scalar.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/context/context.hpp" + +#include + +namespace infinicore::op::mul_scalar_impl::aten { +namespace { + +struct PlannedMeta { + Tensor c; + Tensor a; + double alpha; +}; + +void *plan(Tensor c, const Tensor &a, double alpha) { + INFINICORE_ASSERT(c->device().type() == Device::Type::kQy); + return new PlannedMeta{c, a, alpha}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + + INFINICORE_ASSERT(p->c->device() == p->a->device()); + INFINICORE_ASSERT(p->c->dtype() == p->a->dtype()); + INFINICORE_ASSERT(p->c->numel() == p->a->numel()); + INFINICORE_ASSERT(p->c->is_contiguous()); + + Tensor a = p->a->is_contiguous() ? p->a : p->a->contiguous(); + INFINICORE_ASSERT(a->is_contiguous()); + INFINICORE_ASSERT(a->numel() == p->c->numel()); + INFINICORE_ASSERT(a->dtype() == p->c->dtype()); + + ::infini::ops::qy::mul_scalar_launch( + p->c->data(), a->data(), static_cast(a->numel()), + static_cast(a->dtype()), p->alpha, + reinterpret_cast(context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + MulScalar::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + MulScalar::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MulScalar::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::mul_scalar_impl::aten +#endif diff --git a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc index 043ce070a..a3389041f 100644 --- a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc +++ b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc @@ -47,7 +47,8 @@ bool is_supported(const Tensor &out, && device_type != Device::Type::kCambricon && device_type != Device::Type::kIluvatar && device_type != Device::Type::kHygon - && device_type != Device::Type::kAscend) + && device_type != Device::Type::kAscend + && device_type != Device::Type::kQy) || q->ndim() != 3 || out->ndim() != 3 || ((paged && (k->ndim() != 4 || v->ndim() != 4)) @@ -266,6 +267,12 @@ static bool registered = []() { Device::Type::kAscend, &run); MultiheadAttentionVarlen::cleanup_dispatcher().registerDevice( Device::Type::kAscend, &cleanup); + MultiheadAttentionVarlen::plan_dispatcher().registerDevice( + Device::Type::kQy, &plan); + MultiheadAttentionVarlen::run_dispatcher().registerDevice( + Device::Type::kQy, &run); + MultiheadAttentionVarlen::cleanup_dispatcher().registerDevice( + Device::Type::kQy, &cleanup); return true; }(); diff --git a/csrc/infinicore/src/ops/random_sample/random_sample.cc b/csrc/infinicore/src/ops/random_sample/random_sample.cc index 6d93af3c1..55bbb654f 100644 --- a/csrc/infinicore/src/ops/random_sample/random_sample.cc +++ b/csrc/infinicore/src/ops/random_sample/random_sample.cc @@ -25,13 +25,15 @@ bool tryGreedyWithInfiniOps( && device_type != Device::Type::kMoore && device_type != Device::Type::kCambricon && device_type != Device::Type::kAscend - && device_type != Device::Type::kHygon) + && device_type != Device::Type::kHygon + && device_type != Device::Type::kQy) || (random_value != 0.0f && top_p != 0.0f && top_k != 1 && temperature != 0.0f) || (logits->ndim() != 1 - && !(device_type == Device::Type::kNvidia && batched)) + && !(device_type == Device::Type::kNvidia && batched) + && !(device_type == Device::Type::kQy && batched)) || logits->numel() == 0 || !logits->is_contiguous() || (dtype != DataType::kFloat16 && dtype != DataType::kBFloat16 && dtype != DataType::kFloat32) diff --git a/csrc/infinicore/src/ops/rope/rope_aten.cc b/csrc/infinicore/src/ops/rope/rope_aten.cc new file mode 100644 index 000000000..cf2e37b1b --- /dev/null +++ b/csrc/infinicore/src/ops/rope/rope_aten.cc @@ -0,0 +1,354 @@ +#include "infinicore/ops/rope.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/nn/rope.hpp" + +#include +#include + +#include +#include + +#include +#include +#include +#include + +namespace infinicore::op::rope_impl::aten { +namespace { + +bool stream_is_capturing(cudaStream_t stream) { + cudaStreamCaptureStatus status = cudaStreamCaptureStatusNone; + auto st = cudaStreamIsCapturing(stream, &status); + return st == cudaSuccess && status != cudaStreamCaptureStatusNone; +} + +cudaError_t sync_unless_capturing(cudaStream_t stream) { + if (stream_is_capturing(stream)) { + return cudaGetLastError(); + } + return cudaStreamSynchronize(stream); +} + +// QY FA+paged long-prefill RoPE: Denglin ATen index_select/gather and strided +// copy_/mul on gated as_strided Q/K hit LaunchOutOfResources. Pack/unpack with +// cudaMemcpyAsync (contiguous head vectors), rotate on a compact workspace, +// gather cos/sin via cache-table slice. +constexpr int64_t kQyRopeChunkTokens = 128; + +bool host_timer_enabled() { + static int cached = -1; + if (cached < 0) { + const char *env = std::getenv("INFINI_MOE_HOST_TIMER"); + const char *env2 = std::getenv("INFINI_ROPE_HOST_TIMER"); + cached = ((env != nullptr && env[0] == '1' && env[1] == '\0') || (env2 != nullptr && env2[0] == '1' && env2[1] == '\0')) + ? 1 + : 0; + } + return cached == 1; +} + +struct HostTimer { + const char *op; + int64_t m; + std::chrono::steady_clock::time_point t0; + bool active; + + HostTimer(const char *op_name, int64_t m_val) + : op(op_name), m(m_val), active(host_timer_enabled()) { + if (active) { + t0 = std::chrono::steady_clock::now(); + } + } + + ~HostTimer() { + if (!active) { + return; + } + auto t1 = std::chrono::steady_clock::now(); + double ms = std::chrono::duration(t1 - t0).count(); + std::fprintf(stderr, "host_timer op=%s ms=%.3f m=%lld\n", op, ms, + static_cast(m)); + } +}; + +struct PlannedMeta { + Tensor x_out; + Tensor x; + Tensor pos; + Tensor sin; + Tensor cos; + infinicore::nn::RoPE::Algo algo; +}; + +void *plan(Tensor x_out, + const Tensor &x, + const Tensor &pos, + const Tensor &sin, + const Tensor &cos, + infinicore::nn::RoPE::Algo algo) { + INFINICORE_ASSERT(x_out->device().type() == Device::Type::kQy); + return new PlannedMeta{x_out, x, pos, sin, cos, algo}; +} + +int64_t seq_dim_of(const at::Tensor &x) { + if (x.dim() == 2) { + return 0; + } + INFINICORE_ASSERT(x.dim() >= 3); + return x.dim() - 3; +} + +at::Tensor d2h_contig(const at::Tensor &dev) { + INFINICORE_ASSERT(dev.is_contiguous()); + auto host = at::empty(dev.sizes(), dev.options().device(at::kCPU)); + cudaMemcpy(host.data_ptr(), + dev.data_ptr(), + static_cast(dev.numel() * static_cast(dev.element_size())), + cudaMemcpyDeviceToHost); + return host; +} + +at::Tensor gather_cache_rows(const at::Tensor &table, const at::Tensor &pos_chunk) { + const int64_t rows = pos_chunk.numel(); + INFINICORE_ASSERT(pos_chunk.dim() == 1); + INFINICORE_ASSERT(table.dim() == 2); + if (rows == 0) { + return at::empty({0, table.size(1)}, table.options()); + } + + // Host indices once — avoid ATen index_select / .contiguous() on device. + INFINICORE_ASSERT(pos_chunk.is_contiguous()); + auto pos_host = d2h_contig(pos_chunk); + const auto *pos_ptr = pos_host.data_ptr(); + + auto out = at::empty({rows, table.size(1)}, table.options()); + const auto elem = static_cast(table.element_size()); + auto *dst_base = static_cast(out.data_ptr()); + const auto *src_base = static_cast(table.data_ptr()); + const int64_t row_bytes = table.size(1) * elem; + const int64_t table_rows = table.size(0); + + const int64_t p0 = pos_ptr[0]; + bool consecutive = p0 >= 0 && (p0 + rows) <= table_rows; + for (int64_t i = 0; consecutive && i < rows; ++i) { + consecutive = (pos_ptr[i] == p0 + i); + } + if (consecutive && table.stride(0) == table.size(1)) { + cudaMemcpy(dst_base, + src_base + p0 * table.stride(0) * elem, + static_cast(rows * row_bytes), + cudaMemcpyDeviceToDevice); + return out; + } + + for (int64_t i = 0; i < rows; ++i) { + const int64_t p = pos_ptr[i]; + INFINICORE_ASSERT(p >= 0 && p < table_rows); + cudaMemcpy(dst_base + i * row_bytes, + src_base + p * table.stride(0) * elem, + static_cast(row_bytes), + cudaMemcpyDeviceToDevice); + } + return out; +} + +void log_rope_layout_once(const at::Tensor &x) { + static bool logged = false; + if (logged || !host_timer_enabled()) { + return; + } + logged = true; + std::fprintf(stderr, + "rope_layout dim=%ld sizes=[", static_cast(x.dim())); + for (int64_t i = 0; i < x.dim(); ++i) { + std::fprintf(stderr, "%ld%s", static_cast(x.size(i)), + (i + 1 < x.dim()) ? "," : ""); + } + std::fprintf(stderr, "] strides=["); + for (int64_t i = 0; i < x.dim(); ++i) { + std::fprintf(stderr, "%ld%s", static_cast(x.stride(i)), + (i + 1 < x.dim()) ? "," : ""); + } + std::fprintf(stderr, "] contig=%d\n", x.is_contiguous() ? 1 : 0); +} + +// Pack strided [seq, heads, dim] (or [seq, dim]) into contiguous via QY kernel. +void pack_head_vectors(const at::Tensor &src, at::Tensor &dst, cudaStream_t stream) { + INFINICORE_ASSERT(src.sizes() == dst.sizes()); + INFINICORE_ASSERT(dst.is_contiguous()); + if (src.is_contiguous()) { + cudaMemcpyAsync(dst.data_ptr(), + src.data_ptr(), + static_cast(src.numel() * static_cast(src.element_size())), + cudaMemcpyDeviceToDevice, + stream); + return; + } + log_rope_layout_once(src); + const int elem = static_cast(src.element_size()); + const int64_t dim = src.size(-1); + INFINICORE_ASSERT(src.stride(-1) == 1); + + if (src.dim() == 2) { + ::infini::ops::qy::rope_pack_launch( + dst.data_ptr(), src.data_ptr(), static_cast(src.size(0)), + /*heads=*/1, static_cast(dim), src.stride(0), /*s1=*/0, + /*s2=*/1, elem, reinterpret_cast(stream)); + return; + } + INFINICORE_ASSERT(src.dim() == 3); + ::infini::ops::qy::rope_pack_launch( + dst.data_ptr(), src.data_ptr(), static_cast(src.size(0)), + static_cast(src.size(1)), static_cast(dim), + src.stride(0), src.stride(1), src.stride(2), elem, + reinterpret_cast(stream)); +} + +void unpack_head_vectors(const at::Tensor &src, at::Tensor &dst, cudaStream_t stream) { + INFINICORE_ASSERT(src.sizes() == dst.sizes()); + INFINICORE_ASSERT(src.is_contiguous()); + if (dst.is_contiguous()) { + cudaMemcpyAsync(dst.data_ptr(), + src.data_ptr(), + static_cast(src.numel() * static_cast(src.element_size())), + cudaMemcpyDeviceToDevice, + stream); + return; + } + const int elem = static_cast(src.element_size()); + const int64_t dim = src.size(-1); + INFINICORE_ASSERT(dst.stride(-1) == 1); + + if (dst.dim() == 2) { + ::infini::ops::qy::rope_unpack_launch( + dst.data_ptr(), src.data_ptr(), static_cast(dst.size(0)), + /*heads=*/1, static_cast(dim), dst.stride(0), /*s1=*/0, + /*s2=*/1, elem, reinterpret_cast(stream)); + return; + } + INFINICORE_ASSERT(dst.dim() == 3); + ::infini::ops::qy::rope_unpack_launch( + dst.data_ptr(), src.data_ptr(), static_cast(dst.size(0)), + static_cast(dst.size(1)), static_cast(dim), + dst.stride(0), dst.stride(1), dst.stride(2), elem, + reinterpret_cast(stream)); +} + +at::Tensor rotate_contig_gpu(const at::Tensor &x_contig, + const at::Tensor &pos_chunk, + const at::Tensor &sin_table, + const at::Tensor &cos_table, + infinicore::nn::RoPE::Algo algo, + int dtype_tag, + cudaStream_t stream) { + INFINICORE_ASSERT(x_contig.is_contiguous()); + INFINICORE_ASSERT(pos_chunk.is_contiguous()); + INFINICORE_ASSERT(sin_table.is_contiguous()); + INFINICORE_ASSERT(cos_table.is_contiguous()); + + const int64_t seqlen = x_contig.size(0); + const int64_t nhead = (x_contig.dim() == 2) ? 1 : x_contig.size(1); + const int64_t head_dim = x_contig.size(-1); + INFINICORE_ASSERT(pos_chunk.numel() == seqlen); + INFINICORE_ASSERT(sin_table.size(1) == head_dim / 2); + INFINICORE_ASSERT(cos_table.size(1) == head_dim / 2); + + const int algo_tag = (algo == infinicore::nn::RoPE::Algo::GPT_J) ? 0 : 1; + if (algo != infinicore::nn::RoPE::Algo::GPT_NEOX && algo != infinicore::nn::RoPE::Algo::GPT_J) { + throw std::runtime_error("Unsupported RoPE algorithm for QY ATen path"); + } + + auto out = at::empty(x_contig.sizes(), x_contig.options()); + ::infini::ops::qy::rope_launch( + out.data_ptr(), x_contig.data_ptr(), + pos_chunk.data_ptr(), sin_table.data_ptr(), + cos_table.data_ptr(), static_cast(seqlen), + static_cast(nhead), static_cast(head_dim), algo_tag, + dtype_tag, reinterpret_cast(stream)); + return out; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + using namespace ::infinicore::adaptor; + set_aten_stream_to_infinicore(); + + HostTimer timer("rope", static_cast(p->x->size(0))); + + auto x = to_aten_tensor(p->x); + auto pos = to_aten_tensor(p->pos).to(at::kLong).reshape({-1}).contiguous(); + auto sin_table = to_aten_tensor(p->sin); + auto cos_table = to_aten_tensor(p->cos); + auto out = to_aten_tensor(p->x_out); + + const int64_t seq_len = pos.numel(); + const int64_t seq_dim = seq_dim_of(x); + INFINICORE_ASSERT(seq_dim == 0); + INFINICORE_ASSERT(x.size(0) == seq_len); + INFINICORE_ASSERT(out.sizes() == x.sizes()); + INFINICORE_ASSERT(x.stride(-1) == 1); + INFINICORE_ASSERT(out.stride(-1) == 1); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + at::Tensor x_contig = x.is_contiguous() ? x : at::empty(x.sizes(), x.options()); + if (!x.is_contiguous()) { + HostTimer pack_timer("rope_pack", seq_len); + pack_head_vectors(x, x_contig, stream); + sync_unless_capturing(stream); + } + + // Prefer full-table + pos lookup in the kernel to avoid per-chunk gather. + // Fall back to gathered rows only if tables are not contiguous. + at::Tensor sin_use = sin_table; + at::Tensor cos_use = cos_table; + at::Tensor pos_use = pos; + if (!sin_table.is_contiguous() || !cos_table.is_contiguous()) { + HostTimer gather_timer("rope_gather", seq_len); + sin_use = gather_cache_rows(sin_table, pos); + cos_use = gather_cache_rows(cos_table, pos); + // Remap positions to [0..seq) against gathered rows. + auto pos_remap = at::empty({seq_len}, pos.options().device(at::kCPU)); + auto *ptr = pos_remap.data_ptr(); + for (int64_t i = 0; i < seq_len; ++i) { + ptr[i] = i; + } + pos_use = at::empty({seq_len}, pos.options()); + cudaMemcpy(pos_use.data_ptr(), ptr, + static_cast(seq_len * static_cast(sizeof(int64_t))), + cudaMemcpyHostToDevice); + } + + { + HostTimer rot_timer("rope_rot", seq_len); + auto rotated = rotate_contig_gpu(x_contig, pos_use, sin_use, cos_use, p->algo, + static_cast(p->x->dtype()), stream); + { + HostTimer unpack_timer("rope_unpack", seq_len); + unpack_head_vectors(rotated, out, stream); + sync_unless_capturing(stream); + } + } + + (void)kQyRopeChunkTokens; +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + RoPE::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + RoPE::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + RoPE::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::rope_impl::aten +#endif diff --git a/csrc/infinicore/src/pybind11/device.hpp b/csrc/infinicore/src/pybind11/device.hpp index b8c8e76a6..482968c84 100644 --- a/csrc/infinicore/src/pybind11/device.hpp +++ b/csrc/infinicore/src/pybind11/device.hpp @@ -20,7 +20,8 @@ inline void bind(py::module &m) { .value("MOORE", Device::Type::kMoore) .value("ILUVATAR", Device::Type::kIluvatar) .value("HYGON", Device::Type::kHygon) - .value("THEAD", Device::Type::kThead); + .value("THEAD", Device::Type::kThead) + .value("QY", Device::Type::kQy); device .def(py::init(), diff --git a/csrc/layers/attention/backends/static_attn.cpp b/csrc/layers/attention/backends/static_attn.cpp index 69af850f8..a7395f226 100644 --- a/csrc/layers/attention/backends/static_attn.cpp +++ b/csrc/layers/attention/backends/static_attn.cpp @@ -5,6 +5,8 @@ #include "infinicore/ops/per_tensor_dequant_i8.hpp" #include "infinicore/ops/per_tensor_quant_i8.hpp" +#include + namespace infinilm::layers::attention::backends { namespace { @@ -100,14 +102,42 @@ infinicore::Tensor StaticAttentionImpl::forward(const AttentionLayer &layer, auto K = k_total->view({batch_size * num_kv_heads_, total_seq_len, head_dim_}); auto V = v_total->view({batch_size * num_kv_heads_, total_seq_len, value_head_dim}); - auto K_transposed = K->permute({0, 2, 1}); // [bs * n_kv_head, head_dim, total_seq_len] - - auto attn_weight = infinicore::op::matmul(Q, K_transposed, scale_); // [bs * n_kv_head, ng * seq_len, total_seq_len] - - auto attn_weight_softmax = attn_weight->view({batch_size * num_heads_, seq_len, total_seq_len}); - infinicore::op::causal_softmax_(attn_weight_softmax, attn_weight_softmax); - - auto out = infinicore::op::matmul(attn_weight, V); // [bs * n_kv_head, ng * seq_len, value_head_dim] + // Denglin static-attn GEMM/softmax launch fails around ~600+ query tokens + // ("too many resources requested for launch"). Chunk Q along the sequence + // so each causal matmul stays small; FA/paged remains M2. + constexpr size_t kQyAttnChunkTokens = 128; + const bool chunk_qy = q_reshaped->device().type() == infinicore::Device::Type::kQy + && (seq_len > kQyAttnChunkTokens || total_seq_len > 512); + + auto run_chunk = [&](size_t q_begin, size_t q_rows) -> infinicore::Tensor { + // Queries at absolute positions [cache_pos + q_begin, ...); keys needed + // are only [0, cache_pos + q_begin + q_rows) under causal masking. + const size_t abs_begin = cache_pos + q_begin; + const size_t kv_len = abs_begin + q_rows; + auto Q_chunk = Q->narrow({{1, ngroup * q_begin, ngroup * q_rows}}); + auto K_chunk = K->narrow({{1, 0, kv_len}}); + auto V_chunk = V->narrow({{1, 0, kv_len}}); + auto K_T = K_chunk->permute({0, 2, 1}); + auto attn_weight = infinicore::op::matmul(Q_chunk, K_T, scale_); + auto attn_weight_softmax = attn_weight->view({batch_size * num_heads_, q_rows, kv_len}); + infinicore::op::causal_softmax_(attn_weight_softmax, attn_weight_softmax); + return infinicore::op::matmul(attn_weight, V_chunk); + }; + + infinicore::Tensor out; + if (!chunk_qy) { + out = run_chunk(/*q_begin=*/0, seq_len); + } else { + out = infinicore::Tensor::empty( + {batch_size * num_kv_heads_, ngroup * seq_len, value_head_dim}, + Q->dtype(), + Q->device()); + for (size_t begin = 0; begin < seq_len; begin += kQyAttnChunkTokens) { + const size_t rows = std::min(kQyAttnChunkTokens, seq_len - begin); + auto chunk_out = run_chunk(begin, rows); + out->narrow({{1, ngroup * begin, ngroup * rows}})->copy_from(chunk_out); + } + } return out->view({batch_size, num_heads_, seq_len, value_head_dim}) ->permute({0, 2, 1, 3}) diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index e33e16495..0348f94d5 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -4,6 +4,8 @@ #include "infinicore/ops/moe_align.hpp" #include "infinicore/ops/moe_fused_dense.hpp" +#include +#include #include #include #include @@ -76,17 +78,58 @@ void check_packed_weight_tensor(const infinicore::Tensor &tensor, CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, const MoeWeights &weights, MoeWorkspace &workspace) const { - auto runner_input = prepare_runner_input( - dispatch_output, - workspace); + // Denglin fused MoE materializes [M, topk, H] intermediates. A 7k-token + // prefill OOMs a 32GB KS38; vLLM chunks, but its 16k chunk is still too + // big here. Run ≤1024-token chunks on QY (graph capture stays one shot). + constexpr size_t kQyChunkTokens = 512; + const auto &hidden = dispatch_output.hidden_states; + const size_t num_tokens = hidden ? hidden->shape()[0] : 0; + const bool chunk_qy = hidden + && hidden->device().type() == infinicore::Device::Type::kQy + && !infinicore::context::isGraphRecording() + && num_tokens > kQyChunkTokens + && hidden->shape().size() == 2; + std::cerr << "moe_run device=" << static_cast(hidden ? hidden->device().type() : infinicore::Device::Type::kCpu) + << " m=" << num_tokens + << " ndim=" << (hidden ? hidden->shape().size() : 0) + << " chunk=" << chunk_qy << std::endl; + + if (!chunk_qy) { + auto runner_input = prepare_runner_input(dispatch_output, workspace); + auto runner_output = run_fused_core(runner_input, weights, workspace); + return CombineInput{ + CombineInputFormat::Standard, + runner_output.hidden_states, + runner_input.topk_output, + runner_input.routing_metadata, + }; + } - auto runner_output = run_fused_core(runner_input, weights, workspace); + ensure_tensor( + workspace.fused_moe_output, + hidden->shape(), + hidden->dtype(), + hidden->device()); + workspace.fused_moe_output_tokens_capacity = num_tokens; + const auto &topk = dispatch_output.topk_output; + for (size_t begin = 0; begin < num_tokens; begin += kQyChunkTokens) { + const size_t rows = std::min(kQyChunkTokens, num_tokens - begin); + DispatchOutput slice = dispatch_output; + slice.hidden_states = hidden->narrow({{0, begin, rows}})->contiguous(); + slice.topk_output.topk_weights = topk.topk_weights->narrow({{0, begin, rows}})->contiguous(); + slice.topk_output.topk_ids = topk.topk_ids->narrow({{0, begin, rows}})->contiguous(); + // Fresh align buffers so block_size recovered from numel matches this chunk. + MoeWorkspace chunk_workspace; + auto runner_input = prepare_runner_input(slice, chunk_workspace); + auto runner_output = run_fused_core(runner_input, weights, chunk_workspace); + workspace.fused_moe_output->narrow({{0, begin, rows}})->copy_from(runner_output.hidden_states); + } return CombineInput{ CombineInputFormat::Standard, - runner_output.hidden_states, - runner_input.topk_output, - runner_input.routing_metadata, + workspace.fused_moe_output, + dispatch_output.topk_output, + MoeRoutingMetadata{}, }; } diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_attention.cpp b/csrc/models/minicpm5_moe/minicpm5_moe_attention.cpp new file mode 100644 index 000000000..ff4b44f3e --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_attention.cpp @@ -0,0 +1,181 @@ +#include "minicpm5_moe_attention.hpp" + +#include "../../global_state/global_state.hpp" +#include "../../layers/attention/attention.hpp" +#include "../../utils.hpp" + +#include +#include +#include +#include + +namespace infinilm::models::minicpm5_moe { + +MiniCPM5MoEAttention::MiniCPM5MoEAttention( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + layer_idx_ = layer_idx; + const auto &dtype{model_config->get_dtype()}; + hidden_size_ = model_config->get("hidden_size"); + head_dim_ = model_config->get("head_dim"); + use_gated_attention_ = model_config->get_or("use_gated_attention", true); + + size_t total_num_heads = model_config->get("num_attention_heads"); + size_t total_num_kv_heads = model_config->get("num_key_value_heads"); + bool use_bias = model_config->get_or("attention_bias", false); + bool use_output_bias = model_config->get_or("attention_output_bias", false); + + attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend; + const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank(); + int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size(); + + num_attention_heads_ = total_num_heads / tp_size; + num_key_value_heads_ = total_num_kv_heads < static_cast(tp_size) + ? 1 + : total_num_kv_heads / tp_size; + + // Gated path doubles q_proj (per-head [q | gate]), same packing as qwen3_next. + const size_t q_heads_for_proj = total_num_heads * (use_gated_attention_ ? 2 : 1); + + auto quantization_method = model_config->get_quantization_method(); + auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { + this->register_parameter(n, std::move(p)); + }; + qkv_proj_ = std::make_shared( + hidden_size_, head_dim_, q_heads_for_proj, total_num_kv_heads, + "q_proj", "k_proj", "v_proj", register_fn, + quantization_method, use_bias, dtype, device, rank_info); + o_proj_ = this->register_module( + "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method, + use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm); + + rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); + + float scaling = 1.0f / std::sqrt(static_cast(head_dim_)); + attn_ = std::make_shared( + num_attention_heads_, head_dim_, scaling, num_key_value_heads_, layer_idx_, + kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_); + + infinilm::layers::attention::init_kv_cache_quant_params( + register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_); +} + +infinicore::Tensor MiniCPM5MoEAttention::forward( + const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const { + if (::infinilm::backends::AttentionBackend::STATIC_ATTN == attention_backend_) { + return forward_static_(positions, hidden_states); + } + return forward_paged_(positions, hidden_states); +} + +infinicore::Tensor MiniCPM5MoEAttention::forward_static_( + const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + auto shape = hidden_states->shape(); + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + + auto [q_proj_out, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + + infinicore::Tensor q_reshaped; + infinicore::Tensor gate; + if (use_gated_attention_) { + auto q_proj_heads = q_proj_out->as_strided( + {batch_size * seq_len, num_attention_heads_, head_dim_ * 2}, + {q_proj_out->stride(1), static_cast(head_dim_ * 2), 1}); + auto q = q_proj_heads->narrow({{2, 0, head_dim_}}); + gate = q_proj_heads->narrow({{2, head_dim_, head_dim_}}) + ->contiguous() + ->view({batch_size, seq_len, num_attention_heads_ * head_dim_}); + q_reshaped = q->as_strided( + {batch_size, seq_len, num_attention_heads_, head_dim_}, + {static_cast(seq_len * num_attention_heads_ * head_dim_), + static_cast(num_attention_heads_ * head_dim_), + static_cast(head_dim_), + 1}); + } else { + q_reshaped = q_proj_out->view({batch_size, seq_len, num_attention_heads_, head_dim_}); + } + + auto k_reshaped = k->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + auto v_reshaped = v->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + + auto pos_shape = position_ids->shape(); + infinicore::Tensor pos_ids_for_rope = position_ids; + if (pos_shape.size() == 2) { + auto pos_narrowed = position_ids->narrow({{0, 0, 1}}); + pos_ids_for_rope = pos_narrowed->contiguous()->view({pos_shape[1]}); + } else if (pos_shape.size() == 1) { + pos_ids_for_rope = position_ids->contiguous(); + } else { + throw std::runtime_error( + "MiniCPM5MoEAttention: Unexpected position_ids shape"); + } + + rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + + auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); + if (use_gated_attention_) { + attn_output = infinicore::op::mul(attn_output, infinicore::op::sigmoid(gate)); + } + return o_proj_->forward(attn_output); +} + +infinicore::Tensor MiniCPM5MoEAttention::forward_paged_( + const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + auto shape = hidden_states->shape(); + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + + ASSERT_EQ(batch_size, 1); + + auto [q_proj_out, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + + infinicore::Tensor q_reshaped; + infinicore::Tensor gate; + if (use_gated_attention_) { + auto q_proj_heads = q_proj_out->as_strided( + {seq_len, num_attention_heads_, head_dim_ * 2}, + {q_proj_out->stride(1), static_cast(head_dim_ * 2), 1}); + auto q = q_proj_heads->narrow({{2, 0, head_dim_}}); + gate = q_proj_heads->narrow({{2, head_dim_, head_dim_}}) + ->contiguous() + ->view({seq_len, num_attention_heads_ * head_dim_}); + q_reshaped = q; + } else { + q_reshaped = q_proj_out->view({seq_len, num_attention_heads_, head_dim_}); + } + + auto k_reshaped = k->view({seq_len, num_key_value_heads_, head_dim_}); + auto v_reshaped = v->view({seq_len, num_key_value_heads_, head_dim_}); + + auto pos_shape = position_ids->shape(); + infinicore::Tensor pos_ids_for_rope = position_ids; + if (pos_shape.size() == 2) { + auto pos_narrowed = position_ids->narrow({{0, 0, 1}}); + pos_ids_for_rope = pos_narrowed->view({pos_shape[1]}); + } else if (pos_shape.size() == 1) { + pos_ids_for_rope = position_ids; + } else { + throw std::runtime_error("MiniCPM5MoEAttention: Unexpected position_ids shape"); + } + + rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + + auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); + if (use_gated_attention_) { + attn_output = infinicore::op::mul( + attn_output, infinicore::op::sigmoid(gate)->view(attn_output->shape())); + } + return o_proj_->forward(attn_output); +} + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_attention.hpp b/csrc/models/minicpm5_moe/minicpm5_moe_attention.hpp new file mode 100644 index 000000000..5e17d1c90 --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_attention.hpp @@ -0,0 +1,57 @@ +#pragma once + +#include "../../layers/common_modules.hpp" + +namespace infinilm::models::minicpm5_moe { + +/** + * GQA attention with optional gated output (doubled q_proj, per-head [q|gate]). + * Matches HF MiniCPM5MoEAttention / vllm_minicpm5 gated path (no q/k RMSNorm). + */ +class MiniCPM5MoEAttention : public infinicore::nn::Module { +public: + MiniCPM5MoEAttention(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const; + + void process_weights_after_loading() override { + qkv_proj_->process_weights_after_loading(); + } + + void reset_runtime_state() const override { + qkv_proj_->reset_runtime_state(); + } + + size_t layer_idx() const { return layer_idx_; } + size_t num_heads() const { return num_attention_heads_; } + size_t num_kv_heads() const { return num_key_value_heads_; } + size_t head_dim() const { return head_dim_; } + size_t hidden_size() const { return hidden_size_; } + +protected: + infinicore::Tensor forward_static_(const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const; + infinicore::Tensor forward_paged_(const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const; + + std::shared_ptr qkv_proj_; + std::shared_ptr o_proj_; + std::shared_ptr rotary_emb_; + + std::shared_ptr attn_; + ::infinilm::backends::AttentionBackend attention_backend_; + size_t layer_idx_; + size_t num_attention_heads_; + size_t num_key_value_heads_; + size_t hidden_size_; + size_t head_dim_; + bool use_gated_attention_{true}; + + INFINICORE_NN_PARAMETER(kv_cache_k_scale); + INFINICORE_NN_PARAMETER(kv_cache_v_scale); +}; + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.cpp b/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.cpp new file mode 100644 index 000000000..e1f8ec9cf --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.cpp @@ -0,0 +1,42 @@ +#include "minicpm5_moe_decoder_layer.hpp" + +namespace infinilm::models::minicpm5_moe { + +MiniCPM5MoEDecoderLayer::MiniCPM5MoEDecoderLayer( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const double rms_norm_eps = model_config->get("rms_norm_eps"); + INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); + INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device); + + const size_t first_dense_layers = model_config->get_or("first_k_dense_replace", 0); + use_moe_ = model_config->get_or("n_routed_experts", 0) > 0 + && layer_idx >= first_dense_layers; + if (use_moe_) { + moe_mlp_ = this->register_module( + "mlp", model_config, layer_idx, device); + } else { + auto dense_config = make_minicpm5_mlp_config( + model_config, model_config->get("intermediate_size")); + dense_mlp_ = this->register_module( + "mlp", dense_config, device); + } +} + +std::tuple +MiniCPM5MoEDecoderLayer::forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + input_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = self_attn_->forward(positions, hidden_states); + post_attention_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = use_moe_ ? moe_mlp_->forward(hidden_states) + : dense_mlp_->forward(hidden_states); + return {hidden_states, residual}; +} + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.hpp b/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.hpp new file mode 100644 index 000000000..fe0b1b85d --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.hpp @@ -0,0 +1,38 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/mlp/mlp.hpp" +#include "minicpm5_moe_attention.hpp" +#include "minicpm5_moe_sparse_moe_block.hpp" + +#include +#include +#include + +#include +#include +#include + +namespace infinilm::models::minicpm5_moe { + +class MiniCPM5MoEDecoderLayer : public infinicore::nn::Module { +public: + MiniCPM5MoEDecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + std::tuple + forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + +protected: + INFINICORE_NN_MODULE(MiniCPM5MoEAttention, self_attn); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); + INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, dense_mlp); + INFINICORE_NN_MODULE(MiniCPM5MoESparseMoeBlock, moe_mlp); + bool use_moe_{false}; +}; + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.cpp b/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.cpp new file mode 100644 index 000000000..5be340a72 --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.cpp @@ -0,0 +1,71 @@ +#include "minicpm5_moe_for_causal_lm.hpp" + +#include "../models_registry.hpp" + +#include +#include + +namespace infinilm::models::minicpm5_moe { + +std::shared_ptr +create_minicpm5_moe_model_config(std::shared_ptr model_config) { + const std::string model_type = model_config->get("model_type"); + if (model_type != "minicpm5_moe") { + throw std::runtime_error( + "create_minicpm5_moe_model_config: model_type is not minicpm5_moe"); + } + + auto &config_json = model_config->get_config_json(); + + if (!config_json.contains("head_dim")) { + config_json["head_dim"] = model_config->get("hidden_size") + / model_config->get("num_attention_heads"); + } + if (!config_json.contains("attention_bias")) { + config_json["attention_bias"] = false; + } + if (!config_json.contains("attention_output_bias")) { + config_json["attention_output_bias"] = config_json.value("attention_bias", false); + } + if (!config_json.contains("mlp_bias")) { + config_json["mlp_bias"] = false; + } + if (!config_json.contains("rope_theta")) { + config_json["rope_theta"] = 10000.0; + } + if (!config_json.contains("n_group")) { + config_json["n_group"] = 1; + } + if (!config_json.contains("topk_group")) { + config_json["topk_group"] = 1; + } + if (!config_json.contains("norm_topk_prob")) { + config_json["norm_topk_prob"] = true; + } + if (!config_json.contains("dtype") && config_json.contains("torch_dtype")) { + config_json["dtype"] = config_json["torch_dtype"]; + } + + config_json["num_experts"] = config_json.value("n_routed_experts", 0); + config_json["e_score_correction_bias"] = true; + config_json["apply_routed_scaling_factor_on_output"] = false; + // n_group/topk_group default to 1 → grouped selection is a no-op; use sigmoid + // (same escape hatch as kimi_k25 until fused_gate is graph-safe on Denglin). + config_json["moe_router_backend"] = "sigmoid"; + if (!config_json.contains("use_gated_attention")) { + config_json["use_gated_attention"] = true; + } + + // Llama / MiniCPM rotate_half = GPT_NEOX. + model_config->set_rope_algo(infinicore::nn::RoPE::Algo::GPT_NEOX); + return model_config; +} + +} // namespace infinilm::models::minicpm5_moe + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + minicpm5_moe, + infinilm::models::minicpm5_moe::MiniCPM5MoEForCausalLM, + infinilm::models::minicpm5_moe::create_minicpm5_moe_model_config); +} // namespace diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.hpp b/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.hpp new file mode 100644 index 000000000..06e0a316a --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include "../../layers/causal_lm_templates/text_causal_lm.hpp" +#include "../../layers/causal_lm_templates/text_model.hpp" +#include "minicpm5_moe_decoder_layer.hpp" + +#include + +namespace infinilm::models::minicpm5_moe { + +using MiniCPM5MoEModel = infinilm::layers::causal_lm_templates::TextModel; +using MiniCPM5MoEForCausalLM = infinilm::layers::causal_lm_templates::TextCausalLM; + +std::shared_ptr +create_minicpm5_moe_model_config(std::shared_ptr model_config); + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.cpp b/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.cpp new file mode 100644 index 000000000..f1b2626ea --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.cpp @@ -0,0 +1,58 @@ +#include "minicpm5_moe_sparse_moe_block.hpp" + +#include +#include + +namespace infinilm::models::minicpm5_moe { + +std::shared_ptr +make_minicpm5_mlp_config(const std::shared_ptr &model_config, + size_t intermediate_size) { + auto config_json = model_config->get_config_json(); + config_json["intermediate_size"] = intermediate_size; + return std::make_shared(std::move(config_json)); +} + +MiniCPM5MoESparseMoeBlock::MiniCPM5MoESparseMoeBlock( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + INFINICORE_NN_MODULE_INIT(gate, model_config, device); + INFINICORE_NN_MODULE_INIT(experts, model_config, device); + INFINICORE_NN_MODULE_INIT(fused_moe, model_config, device, layer_idx); + + routed_scaling_factor_ = model_config->get_or("routed_scaling_factor", 1.0f); + const size_t n_shared_experts = model_config->get_or("n_shared_experts", 0); + has_shared_experts_ = n_shared_experts > 0; + if (has_shared_experts_) { + const size_t shared_intermediate_size = model_config->get("moe_intermediate_size") * n_shared_experts; + auto shared_config = make_minicpm5_mlp_config(model_config, shared_intermediate_size); + INFINICORE_NN_MODULE_INIT(shared_experts, shared_config, device); + } +} + +infinicore::Tensor MiniCPM5MoESparseMoeBlock::forward( + const infinicore::Tensor &hidden_states) const { + ASSERT(hidden_states->ndim() == 3); + const auto shape = hidden_states->shape(); + auto flattened = hidden_states->view({shape[0] * shape[1], shape[2]}); + + auto [routing_weights, selected_experts] = gate_->forward(flattened); + // Fold routed_scaling into topk weights (M×topk) instead of mul_scalar on + // the full hidden [B,S,H] — Denglin ATen mul_Scalar OORs on long prefill. + if (routed_scaling_factor_ != 1.0f) { + routing_weights = infinicore::op::mul_scalar(routing_weights, routed_scaling_factor_); + } + infinilm::layers::moe::TopKOutput topk_output{ + routing_weights, + selected_experts, + infinicore::Tensor(), + }; + auto routed = fused_moe_->forward(flattened, topk_output, experts_->moe_weights())->view(shape); + if (has_shared_experts_) { + return infinicore::op::add(routed, shared_experts_->forward(hidden_states)); + } + return routed; +} + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.hpp b/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.hpp new file mode 100644 index 000000000..38a117f40 --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/mlp/mlp.hpp" +#include "../../layers/moe/experts/fused_moe_experts.hpp" +#include "../../layers/moe/fused_moe.hpp" +#include "../../layers/moe/router/topk_router.hpp" + +#include +#include + +namespace infinilm::models::minicpm5_moe { + +std::shared_ptr +make_minicpm5_mlp_config(const std::shared_ptr &model_config, + size_t intermediate_size); + +/** + * Sparse MoE + shared expert (DeepSeek-V3 / MiniCPM5 style). + * Router: sigmoid + e_score_correction_bias; routed_scaling applied on routed output. + */ +class MiniCPM5MoESparseMoeBlock final : public infinicore::nn::Module { +public: + MiniCPM5MoESparseMoeBlock(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +protected: + INFINICORE_NN_MODULE(infinilm::layers::moe::TopKRouter, gate); + INFINICORE_NN_MODULE(infinilm::layers::moe::FusedMoeExperts, experts); + INFINICORE_NN_MODULE(infinilm::layers::moe::FusedMoE, fused_moe); + INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, shared_experts); + bool has_shared_experts_{false}; + float routed_scaling_factor_{1.0f}; +}; + +} // namespace infinilm::models::minicpm5_moe diff --git a/python/infinicore/device.py b/python/infinicore/device.py index 3e9461569..775e9999e 100644 --- a/python/infinicore/device.py +++ b/python/infinicore/device.py @@ -16,6 +16,7 @@ "iluvatar": (_infinicore.Device.Type.ILUVATAR, "iluvatar"), "hygon": (_infinicore.Device.Type.HYGON, "hygon"), "thead": (_infinicore.Device.Type.THEAD, "thead"), + "qy": (_infinicore.Device.Type.QY, "qy"), } _NATIVE_TO_NAME = {native: name for native, name in set(_TYPE_ALIASES.values())} diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py index 6459dda6b..249c35dea 100644 --- a/python/infinicore/tensor.py +++ b/python/infinicore/tensor.py @@ -181,16 +181,20 @@ def from_torch(torch_tensor, *, device=None): ) index = 0 if owner.device.index is None else owner.device.index + _cuda_compat = ("cuda", "qy", "metax", "iluvatar", "hygon", "thead") if device is None: if owner.device.type == "cuda": current = get_device() - if current.type not in ("cpu", "cuda"): + if current.type in ("qy", "thead", "hygon", "iluvatar", "metax"): + target = Device(current.type, index) + elif current.type in ("cpu", "cuda"): + target = Device("nvidia", index) + else: raise ValueError( "Torch reports a CUDA-compatible device without its vendor; " "pass device= explicitly" ) - target = Device("nvidia", index) else: target = Device(owner.device.type, index) elif isinstance(device, str) and ":" not in device: @@ -201,7 +205,7 @@ def from_torch(torch_tensor, *, device=None): if target.index != index: raise ValueError("device index must match the source Torch tensor") if owner.device.type == "cuda": - if target.type not in ("cuda", "metax", "iluvatar", "hygon", "thead"): + if target.type not in _cuda_compat: raise ValueError("device is not compatible with a Torch CUDA tensor") elif target.type != owner.device.type: raise ValueError("device type must match the source Torch tensor") @@ -218,11 +222,11 @@ def from_torch(torch_tensor, *, device=None): ) result = empty(list(owner.shape), dtype=dtype, device=target) - if owner.device.type == "cuda": + if owner.device.type == "cuda" or target.type == "qy": torch.cuda.synchronize(owner.device) result.copy_(borrowed) - if owner.device.type == "cuda": + if owner.device.type == "cuda" or target.type == "qy": sync_stream() return result diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 857b71c01..ab1dd6c6d 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -3,6 +3,7 @@ import json import os import shutil +import subprocess import warnings from infinilm.config.attention import resolve_attention_backend @@ -201,7 +202,7 @@ def _add_common_args(self): type=str, default="auto", help=( - "device platform: auto, cpu, nvidia, metax, moore, iluvatar, " + "device platform: auto, cpu, nvidia, qy, metax, moore, iluvatar, " "cambricon, ascend, hygon, thead, or backend name " "(cuda/mlu/musa/npu)" ), @@ -524,11 +525,28 @@ def detect_device(self): env_checks = [ ("metax", ["MACA_PATH", "MACA_HOME", "MACA_ROOT"]), ("hygon", ["DTK_HOME", "DTK_PATH"]), + ("qy", ["QY_ROOT", "QY_HOME", "QY_PATH"]), ] for device_name, env_names in env_checks: if any(os.getenv(env_name) for env_name in env_names): return device_name + # QY exposes inventory via dlsmi; require a successful probe. + dlsmi = shutil.which("dlsmi") + if dlsmi: + try: + subprocess.run( + [dlsmi], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=True, + timeout=2, + ) + except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired): + pass + else: + return "qy" + command_checks = [ ("cambricon", ["cnmon"]), ("ascend", ["npu-smi"]), @@ -557,6 +575,8 @@ def get_device_str(self, device): "musa": "musa", "npu": "npu", "nvidia": "cuda", + # Keep InfiniCore device as `qy` (kQy); torch tensors still use CUDA. + "qy": "qy", "cambricon": "mlu", "ascend": "npu", "metax": "metax", diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..fed7a60be 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -784,6 +784,7 @@ def add_request( request_id: Optional[str] = None, # For server use request_data: Optional[dict] = None, + chat_template_kwargs: Optional[dict] = None, ) -> InferenceRequest: """Add a request to the engine. @@ -838,7 +839,9 @@ def add_request( ) prompt = self.engine.apply_chat_template( - messages, add_generation_prompt=add_generation_prompt + messages, + add_generation_prompt=add_generation_prompt, + chat_template_kwargs=chat_template_kwargs or {}, ) mm_inputs = resolve_multimodal_inputs(messages) @@ -918,6 +921,7 @@ def add_chat_request( sampling_params=sampling_params, request_id=request_id, request_data=request_data, + chat_template_kwargs=kwargs.get("chat_template_kwargs") or {}, ) async def stream_request( diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index 5a178a08b..85c061ca8 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -166,6 +166,7 @@ def _init_device(self): supported_devices = [ "cpu", "cuda", + "qy", "hygon", "thead", "iluvatar", diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index f4b7c50fc..bc454413f 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -342,7 +342,10 @@ def load_model_state_dict_by_file( ) already_loaded_keys.append(key) - model.load_state_dict(model_param_infini, strict=True) + # MiniCPM5 MoE registers packed w13/w2 views absent from HF `.bin` + # checkpoints; check_parameters filters those expected missing keys. + bin_strict = model_type not in ("minicpm5_moe",) + model.load_state_dict(model_param_infini, strict=bin_strict) infinicore.sync_device() del model_param_infini del model_params diff --git a/python/infinilm/processors/sentencepiece_processor.py b/python/infinilm/processors/sentencepiece_processor.py index c8c3c1b33..857752d55 100644 --- a/python/infinilm/processors/sentencepiece_processor.py +++ b/python/infinilm/processors/sentencepiece_processor.py @@ -16,10 +16,29 @@ def __init__(self, model_dir_path: str): super().__init__(model_dir_path) self._fix_tokenizer_decode(self.tokenizer) + @staticmethod + def _uses_sentencepiece_metaspace(tokenizer) -> bool: + """True for SentencePiece Metaspace (▁ pieces), false for ByteLevel (Ġ).""" + vocab = tokenizer.get_vocab() + return any(token.startswith("▁") for token in vocab) + @staticmethod def _fix_tokenizer_decode(tokenizer): """Fix leading space loss when tokenizer decodes incrementally. + ByteLevel tokenizers (MiniCPM5 sidecar ``tokenizer_bytelevel``) must + keep the native ``decode()`` so GPT-2 byte-unicode pieces become UTF-8. + Joining ``convert_ids_to_tokens`` leaves Chinese as mojibake (e.g. 二). + """ + if not SentencePieceProcessor._uses_sentencepiece_metaspace(tokenizer): + return + + SentencePieceProcessor._patch_metaspace_decode(tokenizer) + + @staticmethod + def _patch_metaspace_decode(tokenizer): + """Fix leading space loss when tokenizer decodes incrementally. + Problem: Fast tokenizer's Rust backend trims leading spaces derived from ▁ (U+2581) during single-token decoding, causing words to concatenate. @@ -78,6 +97,7 @@ class MistralProcessor(SentencePieceProcessor): @register_processor("minicpm") @register_processor("minicpm4") +@register_processor("minicpm5_moe") class MiniCPMProcessor(SentencePieceProcessor): pass diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 462c31084..26b733612 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -15,6 +15,7 @@ import uvicorn from fastapi import FastAPI, Request from fastapi.responses import JSONResponse, StreamingResponse + from infinilm.base_config import BaseConfig from infinilm.config import KVTransferConfig from infinilm.llm import AsyncLLMEngine, FinishReason, SamplingParams @@ -23,7 +24,8 @@ logger = logging.getLogger(__name__) DEFAULT_STREAM_TIMEOUT = 100.0 -DEFAULT_REQUEST_TIMEOUT = 1000.0 +# Long-prefill adapting (CPU RoPE + Q-chunked attn) can exceed 1000s at ~4–7k. +DEFAULT_REQUEST_TIMEOUT = 7200.0 def chunk_json( diff --git a/scripts/build_infini_stack.py b/scripts/build_infini_stack.py index 35a4e80b7..515740935 100644 --- a/scripts/build_infini_stack.py +++ b/scripts/build_infini_stack.py @@ -9,9 +9,12 @@ from typing import Dict, List, Mapping, Optional, Sequence PROJECT_ROOT = Path(__file__).resolve().parents[1] -SUPPORTED_BACKENDS = ("nvidia", "iluvatar", "hygon", "moore") +SUPPORTED_BACKENDS = ("nvidia", "iluvatar", "hygon", "moore", "qy") DEFAULT_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops.json" MOORE_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops_moore.json" +QY_M0_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops_qy_m0.json" +QY_M1_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops_qy_m1.json" +QY_M2_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops_qy_m2.json" SUBMODULES = { "InfiniRT": Path("submodules/InfiniRT"), "InfiniOps": Path("submodules/InfiniOps"), @@ -80,7 +83,12 @@ def git_capture(arguments: Sequence[str], cwd: Path) -> str: ) -def validate_submodule(project_root: Path, relative_path: Path) -> str: +def validate_submodule( + project_root: Path, + relative_path: Path, + *, + allow_dirty: bool = False, +) -> str: submodule_root = project_root / relative_path if not (submodule_root / "CMakeLists.txt").is_file(): raise RuntimeError( @@ -93,14 +101,20 @@ def validate_submodule(project_root: Path, relative_path: Path) -> str: relative_path, ) actual = git_capture(["rev-parse", "HEAD"], submodule_root) - if actual != expected: + if actual != expected and not allow_dirty: raise RuntimeError( f"Submodule revision mismatch for {relative_path}: " f"expected {expected}, found {actual}" ) status = git_capture(["status", "--porcelain"], submodule_root) - if status: + if status and not allow_dirty: raise RuntimeError(f"Submodule worktree is dirty: {relative_path}") + if allow_dirty and (actual != expected or status): + print( + f"warning: allowing dirty/mismatched submodule {relative_path} " + f"(expected {expected}, actual {actual})", + flush=True, + ) return actual @@ -219,6 +233,15 @@ def build_infiniops_commands( iluvatar_arch: Optional[str] = None, hygon_arch: Optional[str] = None, ) -> List[List[str]]: + # QY: enable linked FA when the ops config selects slot-16 flash providers (M2+). + ops_text = operator_config.read_text(encoding="utf-8") + qy_wants_linked = ( + backend == "qy" + and "flash_attn_varlen_func" in ops_text + and "flash_attn_with_kvcache" in ops_text + ) + linked = "ON" if (backend != "qy" or qy_wants_linked) else "OFF" + generate_bindings = "ON" if backend == "qy" else "OFF" configure = [ "cmake", "-S", @@ -227,11 +250,11 @@ def build_infiniops_commands( str(build), "-DWITH_CPU=ON", backend_cmake_option(backend), - "-DWITH_LINKED=ON", + f"-DWITH_LINKED={linked}", "-DWITH_TORCH=ON", "-DAUTO_DETECT_DEVICES=OFF", "-DAUTO_DETECT_BACKENDS=OFF", - "-DGENERATE_PYTHON_BINDINGS=OFF", + f"-DGENERATE_PYTHON_BINDINGS={generate_bindings}", f"-DPython_EXECUTABLE={sys.executable}", f"-DINFINI_RT_ROOT={prefix}", f"-DINFINI_OPS_OPS={operator_config}", @@ -239,20 +262,46 @@ def build_infiniops_commands( f"-DCMAKE_INSTALL_PREFIX={prefix}", ] append_gpu_arch(configure, backend, cuda_arch, iluvatar_arch, hygon_arch) + if backend == "qy": + qy_root = Path( + os.environ.get("QY_ROOT") + or os.environ.get("QY_HOME") + or os.environ.get("QY_PATH") + or "/usr/local/denglin" + ) + configure.append(f"-DQY_ROOT={qy_root}") + try: + import pybind11 - return [ - configure, - [ - "cmake", - "--build", - str(build), - "--target", - "infiniops", - "--parallel", - str(jobs), - ], - ["cmake", "--install", str(build)], + configure.append(f"-Dpybind11_DIR={pybind11.get_cmake_dir()}") + except Exception: + pass + + build_cmd = [ + "cmake", + "--build", + str(build), + "--target", + "infiniops", + "--parallel", + str(jobs), ] + commands = [configure, build_cmd] + # QY M0 enables GENERATE_PYTHON_BINDINGS; install expects the ops module. + if backend == "qy": + commands.append( + [ + "cmake", + "--build", + str(build), + "--target", + "ops", + "--parallel", + str(jobs), + ] + ) + commands.append(["cmake", "--install", str(build)]) + return commands def build_infiniccl_commands( @@ -267,23 +316,31 @@ def build_infiniccl_commands( iluvatar_arch: Optional[str] = None, hygon_arch: Optional[str] = None, ) -> List[List[str]]: + # QY InfiniCCL/NCCL is M3; TP1 keeps InfiniCCL CPU-only (no WITH_QY yet). + ccl_backend = backend + ccl_test = test + if backend == "qy": + ccl_backend = None + ccl_test = False + configure = [ "cmake", "-S", str(source), "-B", str(build), - backend_cmake_option(backend), - f"-DWITH_NCCL={'OFF' if backend == 'moore' else 'ON'}", + "-DWITH_NCCL=OFF" if backend in ("moore", "qy") else "-DWITH_NCCL=ON", f"-DWITH_MCCL={'ON' if backend == 'moore' else 'OFF'}", "-DWITH_OMPI=OFF", "-DWITH_MPICH=OFF", "-DAUTO_DETECT_DEVICES=OFF", "-DAUTO_DETECT_BACKENDS=OFF", - f"-DBUILD_EXAMPLES={'ON' if test else 'OFF'}", + f"-DBUILD_EXAMPLES={'ON' if ccl_test else 'OFF'}", f"-DCMAKE_BUILD_TYPE={build_type}", f"-DCMAKE_INSTALL_PREFIX={prefix}", ] + if ccl_backend is not None: + configure.insert(5, backend_cmake_option(ccl_backend)) append_gpu_arch(configure, backend, cuda_arch, iluvatar_arch, hygon_arch) if backend == "iluvatar": configure.extend( @@ -310,7 +367,7 @@ def build_infiniccl_commands( ["cmake", "--build", str(build), "--parallel", str(jobs)], ["cmake", "--install", str(build)], ] - if test: + if ccl_test: commands.append( [ str(build / "examples/ccl/all_reduce"), @@ -465,6 +522,14 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: action="store_true", help="Print commands without executing them.", ) + parser.add_argument( + "--allow-dirty-submodules", + action="store_true", + help=( + "Allow dirty or gitlink-mismatched InfiniCore submodules " + "(for local M0 feature-branch builds)." + ), + ) args = parser.parse_args(argv) if args.jobs < 1: parser.error("--jobs must be positive") @@ -479,13 +544,16 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: if args.backend == "hygon" and args.hygon_arch is None: args.hygon_arch = "gfx936" if args.operator_config is None: - if args.backend not in ("nvidia", "moore"): + if args.backend == "qy": + args.operator_config = QY_M1_OPERATOR_CONFIG + elif args.backend not in ("nvidia", "moore"): parser.error(f"--operator-config is required with --backend={args.backend}") - args.operator_config = ( - MOORE_OPERATOR_CONFIG - if args.backend == "moore" - else DEFAULT_OPERATOR_CONFIG - ) + else: + args.operator_config = ( + MOORE_OPERATOR_CONFIG + if args.backend == "moore" + else DEFAULT_OPERATOR_CONFIG + ) if args.build_root is None: args.build_root = Path("build/integration") / args.backend return args @@ -509,7 +577,11 @@ def main(argv: Optional[Sequence[str]] = None) -> int: operator_config_path = resolve_from_project(args.operator_config) operator_config = read_operator_config(operator_config_path) revisions = { - name: validate_submodule(infinicore_root, relative_path) + name: validate_submodule( + infinicore_root, + relative_path, + allow_dirty=args.allow_dirty_submodules, + ) for name, relative_path in SUBMODULES.items() } infinilm_revision = git_capture(["rev-parse", "HEAD"], PROJECT_ROOT) @@ -527,6 +599,29 @@ def main(argv: Optional[Sequence[str]] = None) -> int: ], ) ) + elif args.backend == "qy": + qy_root = Path( + build_env.setdefault( + "QY_ROOT", + build_env.get("QY_HOME") + or build_env.get("QY_PATH") + or "/usr/local/denglin", + ) + ) + build_env["CUDA_HOME"] = str(qy_root / "sdk") + build_env["CUDA_PATH"] = str(qy_root / "sdk") + build_env["LD_LIBRARY_PATH"] = os.pathsep.join( + filter( + None, + [ + str(qy_root / "sdk/lib"), + build_env.get("LD_LIBRARY_PATH", ""), + ], + ) + ) + build_env["PATH"] = os.pathsep.join( + filter(None, [str(qy_root / "sdk/bin"), build_env.get("PATH", "")]) + ) print(f"InfiniLM: {infinilm_revision}") print(f"InfiniCore: {infinicore_revision}") diff --git a/scripts/configs/infiniops_ops_qy_m0.json b/scripts/configs/infiniops_ops_qy_m0.json new file mode 100644 index 000000000..ca05294d1 --- /dev/null +++ b/scripts/configs/infiniops_ops_qy_m0.json @@ -0,0 +1,8 @@ +{ + "add": { + "implementations": [0] + }, + "gemm": { + "implementations": [0] + } +} diff --git a/scripts/configs/infiniops_ops_qy_m1.json b/scripts/configs/infiniops_ops_qy_m1.json new file mode 100644 index 000000000..83c615156 --- /dev/null +++ b/scripts/configs/infiniops_ops_qy_m1.json @@ -0,0 +1,22 @@ +{ + "add": {"implementations": [0]}, + "argmax": {"implementations": [0]}, + "convolution": {"implementations": [0]}, + "copy": {"implementations": [0]}, + "embedding": {"implementations": [0]}, + "fill": {"implementations": [0]}, + "fused_add_rms_norm": {"implementations": [0]}, + "gelu": {"implementations": [0]}, + "gemm": {"implementations": [0]}, + "mul": {"implementations": [0]}, + "relu": {"implementations": [0]}, + "rms_norm": {"implementations": [0]}, + "rotary_embedding": {"implementations": [0]}, + "sigmoid": {"implementations": [0]}, + "silu": {"implementations": [0]}, + "silu_and_mul": {"implementations": [0]}, + "softmax": {"implementations": [0]}, + "topk_softmax": {"implementations": [0]}, + "tril": {"implementations": [8]}, + "triu": {"implementations": [8]} +} diff --git a/scripts/configs/infiniops_ops_qy_m2.json b/scripts/configs/infiniops_ops_qy_m2.json new file mode 100644 index 000000000..54d0df8f4 --- /dev/null +++ b/scripts/configs/infiniops_ops_qy_m2.json @@ -0,0 +1,25 @@ +{ + "add": {"implementations": [0]}, + "argmax": {"implementations": [0]}, + "convolution": {"implementations": [0]}, + "copy": {"implementations": [0]}, + "embedding": {"implementations": [0]}, + "fill": {"implementations": [0]}, + "flash_attn_varlen_func": {"implementations": [16]}, + "flash_attn_with_kvcache": {"implementations": [16]}, + "fused_add_rms_norm": {"implementations": [0]}, + "gelu": {"implementations": [0]}, + "gemm": {"implementations": [0]}, + "mul": {"implementations": [0]}, + "relu": {"implementations": [0]}, + "reshape_and_cache_flash": {"implementations": [0]}, + "rms_norm": {"implementations": [0]}, + "rotary_embedding": {"implementations": [0]}, + "sigmoid": {"implementations": [0]}, + "silu": {"implementations": [0]}, + "silu_and_mul": {"implementations": [0]}, + "softmax": {"implementations": [0]}, + "topk_softmax": {"implementations": [0]}, + "tril": {"implementations": [8]}, + "triu": {"implementations": [8]} +} diff --git a/scripts/debug_first_token.py b/scripts/debug_first_token.py new file mode 100644 index 000000000..9ab965ac6 --- /dev/null +++ b/scripts/debug_first_token.py @@ -0,0 +1,316 @@ +#!/usr/bin/env python3 +"""Dump first-token argmax / top-5 logits for MiniCPM5-2B on QY. + +Runs one prefill forward (InferEngine) with static attn (graph off) and +optionally flash-attn + paged KV. Writes JSON under --out-dir. + +Tensor→numpy patterns follow test/models/llama/test_llama_inference.py +(via InfiniCore Tensor.to_numpy). Does not wire HookRegistry. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +import traceback +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import numpy as np + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--model", + default=os.environ.get("MODEL_DIR", "/models/minicpm5-2b"), + help="HF model directory", + ) + p.add_argument("--device", default="qy", help="InfiniCore device type") + p.add_argument("--device-index", type=int, default=0) + p.add_argument( + "--prompt", + default="你好", + help="User prompt (chat-templated like the OpenAI server)", + ) + p.add_argument( + "--modes", + default="static,fa", + help="Comma-separated modes: static and/or fa", + ) + p.add_argument( + "--out-dir", + default=os.environ.get( + "OUT_DIR", + str(Path(__file__).resolve().parents[2] / "logs"), + ), + help="Directory for m1-2b-first-token-{mode}.json", + ) + p.add_argument("--max-cache-len", type=int, default=4096) + p.add_argument("--num-blocks", type=int, default=64) + p.add_argument("--block-size", type=int, default=256) + p.add_argument("--top-k", type=int, default=5) + return p.parse_args() + + +def _logits_hash(arr: np.ndarray) -> str: + return hashlib.sha256(np.ascontiguousarray(arr).tobytes()).hexdigest()[:16] + + +def _top_k(logits_1d: np.ndarray, k: int) -> List[Dict[str, Any]]: + k = min(k, int(logits_1d.shape[0])) + idx = np.argpartition(logits_1d, -k)[-k:] + idx = idx[np.argsort(-logits_1d[idx])] + return [{"token_id": int(i), "logit": float(logits_1d[i])} for i in idx] + + +def _tokenize_chat(model_dir: str, user_prompt: str) -> Tuple[Any, List[int], str]: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) + messages = [{"role": "user", "content": user_prompt}] + rendered = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + input_ids = tokenizer.encode(rendered, add_special_tokens=False) + return tokenizer, input_ids, rendered + + +def _build_engine( + model_path: str, + device, + *, + mode: str, + max_cache_len: int, + num_blocks: int, + block_size: int, +): + from infinilm.cache.cache import PagedKVCacheConfig, StaticKVCacheConfig + from infinilm.distributed import DistConfig + from infinilm.infer_engine import InferEngine + from infinilm.modeling_utils import load_model_state_dict_by_file + + if mode == "static": + cache_config = StaticKVCacheConfig( + max_batch_size=1, max_cache_len=max_cache_len + ) + attention_backend = "default" + enable_graph = False + elif mode == "fa": + cache_config = PagedKVCacheConfig( + num_blocks=num_blocks, + block_size=block_size, + max_batch_size=1, + ) + attention_backend = "flash-attn" + enable_graph = False + else: + raise ValueError(f"unknown mode: {mode}") + + engine = InferEngine( + model_path=model_path, + device=device, + distributed_config=DistConfig(1), + cache_config=cache_config, + enable_graph_compiling=enable_graph, + attention_backend=attention_backend, + weight_load_mode="sync", + ) + load_model_state_dict_by_file(engine, model_path, dtype=engine.dtype) + return engine + + +def _tensor_to_f32_numpy(tensor) -> np.ndarray: + """InfiniCore→numpy; bf16 has no numpy dtype so go through torch.""" + import ctypes + + import infinicore + import torch + + source = tensor + if source.device.type != "cpu": + source = source.to(infinicore.device("cpu", 0)) + try: + infinicore.sync_device() + except Exception: + pass + if not source.is_contiguous(): + source = source.contiguous() + + shape = list(source.shape) + dtype_name = str(source.dtype) + if "bfloat16" in dtype_name: + torch_t = torch.empty(shape, dtype=torch.bfloat16) + if torch_t.numel(): + ctypes.memmove(torch_t.data_ptr(), source.data_ptr(), torch_t.numel() * 2) + return torch_t.float().numpy() + if "float16" in dtype_name or dtype_name.endswith("f16"): + torch_t = torch.empty(shape, dtype=torch.float16) + if torch_t.numel(): + ctypes.memmove(torch_t.data_ptr(), source.data_ptr(), torch_t.numel() * 2) + return torch_t.float().numpy() + return np.asarray(source.to_numpy(), dtype=np.float32) + + +def _prefill_logits(engine, input_ids: List[int], mode: str, block_size: int): + import infinicore + + seq_len = len(input_ids) + batch_size = 1 + past_seq_len = 0 + + ids = infinicore.from_list([input_ids], dtype=infinicore.int64) + position_ids = infinicore.from_list( + [list(range(past_seq_len, past_seq_len + seq_len))], + dtype=infinicore.int64, + ) + past_kv_lengths = infinicore.from_list( + [past_seq_len] * batch_size, dtype=infinicore.int32 + ) + total_kv_lengths = infinicore.from_list( + [past_seq_len + seq_len] * batch_size, dtype=infinicore.int32 + ) + cu_seqlens = infinicore.from_list( + [(past_seq_len + seq_len) * i for i in range(batch_size + 1)], + dtype=infinicore.int32, + ) + input_offsets = infinicore.from_list( + [seq_len * i for i in range(batch_size + 1)], + dtype=infinicore.int32, + ) + + block_tables = None + slot_mapping = None + if mode == "fa": + # Match InferEngine.generate paged prefill bookkeeping (batch=1). + max_blocks = (seq_len + block_size - 1) // block_size + max_blocks = max(max_blocks, 1) + block_tables = infinicore.from_list( + [list(range(max_blocks))], dtype=infinicore.int32 + ) + slot_mapping = infinicore.from_list( + list(range(seq_len)), dtype=infinicore.int64 + ) + ids = ids.view([1, seq_len]) + + out = engine.forward_raw( + input_ids=ids, + position_ids=position_ids, + past_kv_lengths=past_kv_lengths, + total_kv_lengths=total_kv_lengths, + input_offsets=input_offsets, + cu_seqlens=cu_seqlens, + block_tables=block_tables, + slot_mapping=slot_mapping, + sample_all_positions=True, + ) + arr = _tensor_to_f32_numpy(out["logits"]) + # Expected [B, T, V] or [B, 1, V] (last token only). + if arr.ndim == 3: + last = arr[0, -1, :] + elif arr.ndim == 2: + last = arr[-1, :] + else: + raise RuntimeError(f"unexpected logits shape: {arr.shape}") + return arr, last.astype(np.float32, copy=False) + + +def _run_mode(args: argparse.Namespace, mode: str) -> Dict[str, Any]: + import infinicore + + record: Dict[str, Any] = { + "mode": mode, + "model": args.model, + "device": args.device, + "prompt": args.prompt, + "status": "ok", + } + try: + tokenizer, input_ids, rendered = _tokenize_chat(args.model, args.prompt) + record["rendered_prompt"] = rendered + record["input_ids"] = input_ids + record["num_prompt_tokens"] = len(input_ids) + + device = infinicore.device(args.device, args.device_index) + engine = _build_engine( + args.model, + device, + mode=mode, + max_cache_len=args.max_cache_len, + num_blocks=args.num_blocks, + block_size=args.block_size, + ) + full_logits, last = _prefill_logits( + engine, input_ids, mode=mode, block_size=args.block_size + ) + argmax_id = int(np.argmax(last)) + top = _top_k(last, args.top_k) + # Attach decoded text for readability. + for item in top: + item["token"] = tokenizer.decode( + [item["token_id"]], skip_special_tokens=False + ) + record.update( + { + "argmax": argmax_id, + "argmax_token": tokenizer.decode( + [argmax_id], skip_special_tokens=False + ), + "top5": top, + "logits_shape": list(full_logits.shape), + "logits_dtype": str(full_logits.dtype), + "logits_hash": _logits_hash(last), + "logits_l2": float(np.linalg.norm(last)), + "logits_stats": { + "min": float(last.min()), + "max": float(last.max()), + "mean": float(last.mean()), + }, + } + ) + except Exception as e: + record["status"] = "error" + record["error"] = f"{type(e).__name__}: {e}" + record["traceback"] = traceback.format_exc() + return record + + +def main() -> int: + args = _parse_args() + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + modes = [m.strip() for m in args.modes.split(",") if m.strip()] + rc = 0 + for mode in modes: + print(f"=== mode={mode} ===", flush=True) + record = _run_mode(args, mode) + out_path = out_dir / f"m1-2b-first-token-{mode}.json" + with open(out_path, "w", encoding="utf-8") as f: + json.dump(record, f, ensure_ascii=False, indent=2) + f.write("\n") + print(f"wrote {out_path} status={record.get('status')}", flush=True) + if record.get("status") == "ok": + print( + f" argmax={record['argmax']} ({record['argmax_token']!r}) " + f"hash={record['logits_hash']} l2={record['logits_l2']:.4f}", + flush=True, + ) + for i, item in enumerate(record["top5"]): + print( + f" top{i + 1}: id={item['token_id']} " + f"logit={item['logit']:.6f} token={item['token']!r}", + flush=True, + ) + else: + print(f" ERROR: {record.get('error')}", flush=True) + rc = 1 + return rc + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/static/test_infinicore_python_contracts.py b/test/static/test_infinicore_python_contracts.py index ae6e501ec..e6b4cc667 100644 --- a/test/static/test_infinicore_python_contracts.py +++ b/test/static/test_infinicore_python_contracts.py @@ -66,8 +66,7 @@ def test_ppu_detection_selects_thead(self) -> None: config._torch_device_available = lambda name: False self.assertEqual(config.get_device_str("thead"), "thead") self.assertEqual(config.detect_device(), "thead") - with self.assertRaisesRegex(ValueError, "unsupported device platform 'qy'"): - config.get_device_str("qy") + self.assertEqual(config.get_device_str("qy"), "qy") def test_rope_wrapper_preserves_arguments_and_output_identity(self) -> None: calls = [] diff --git a/test/static/test_modern_infinicore_compatibility.py b/test/static/test_modern_infinicore_compatibility.py index 29988fbb5..b49179465 100644 --- a/test/static/test_modern_infinicore_compatibility.py +++ b/test/static/test_modern_infinicore_compatibility.py @@ -267,7 +267,11 @@ def test_current_infiniops_backend_bridges_are_registered(self) -> None: engine, ) self.assertIn( - "flash-attn is only available on NVIDIA, MetaX, Moore, Cambricon, Iluvatar, Hygon, and Ascend devices", + "device_type != infinicore::Device::Type::kQy", + engine, + ) + self.assertIn( + "flash-attn is only available on NVIDIA, MetaX, Moore, Cambricon, Iluvatar, Hygon, Ascend, and QY devices", engine, ) self.assertIn( diff --git a/xmake.lua b/xmake.lua index 42e93b626..ec8a5e51b 100644 --- a/xmake.lua +++ b/xmake.lua @@ -1,6 +1,20 @@ -add_requires("pybind11") -if is_plat("linux") then - add_requireconfs("pybind11.python", { configs = { headeronly = true } }) +-- Prefer pip/system pybind11 to avoid xmake-repo network fetches in offline containers. +local PYBIND11_INCLUDE = os.getenv("PYBIND11_INCLUDE") +if not PYBIND11_INCLUDE or PYBIND11_INCLUDE == "" then + local ok, out = pcall(function() + return os.iorunv(os.getenv("PYTHON") or "python3", { + "-c", "import pybind11; print(pybind11.get_include())" + }) + end) + if ok and out and out:trim() ~= "" then + PYBIND11_INCLUDE = out:trim() + end +end +if not PYBIND11_INCLUDE or PYBIND11_INCLUDE == "" then + add_requires("pybind11") + if is_plat("linux") then + add_requireconfs("pybind11.python", { configs = { headeronly = true } }) + end end local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") @@ -16,6 +30,24 @@ local OPERATOR_CALLS_HEADER = INFINI_ROOT .. "/include/infini/operator_call_inst set_toolchains("gcc") +-- Fallback when xmake python plugin / network repos are unavailable. +rule("python.module") + on_config(function (target) + target:set("kind", "shared") + target:set("prefixname", "") + local suffix = try { + function () + return os.iorunv(os.getenv("PYTHON") or "python3", { + "-c", "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX') or '')" + }):trim() + end + } + if suffix and suffix ~= "" and suffix ~= "None" then + target:set("extension", suffix) + end + end) +rule_end() + option("cxx11-abi") set_default(nil) set_showmenu(true) @@ -50,11 +82,18 @@ if os.isdir(MUSA_ROOT .. "/include") then add_linkdirs(MUSA_ROOT .. "/lib", MUSA_ROOT .. "/lib64") end +local QY_SDK_PRESENT = os.getenv("QY_ROOT") or os.getenv("QY_HOME") or os.getenv("QY_PATH") + or os.isdir("/usr/local/denglin/sdk") or os.isdir("/usr/local/dlgpu/sdk") + target("infinicore_runtime") set_default(false) set_kind("shared") set_languages("cxx17") add_defines("ENABLE_INFINIOPS_API", "USE_INFINIRT_GRAPH") + -- Denglin QY: CUDA-compatible streams + ATen bridge for `_dl_C` MoE. + if QY_SDK_PRESENT then + add_defines("ENABLE_QY_API", "ENABLE_ATEN") + end if not is_plat("windows") then add_cxflags("-fPIC") @@ -75,6 +114,7 @@ target("infinicore_runtime") add_links("infiniops", "infiniccl", "infinirt", { public = true }) add_files("csrc/infinicore/src/*.cc") + add_files("csrc/infinicore/src/adaptor/*.cc") add_files("csrc/infinicore/src/context/*.cc") add_files("csrc/infinicore/src/context/*/*.cc") add_files("csrc/infinicore/src/graph/*.cc") @@ -91,6 +131,39 @@ target("infinicore_runtime") remove_files("csrc/infinicore/src/ops/*/*_moore.cc") on_load(function (target) + if QY_SDK_PRESENT then + local python = os.getenv("PYTHON") or "python3" + local torch_dir = try { + function () + return os.iorunv(python, {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim() + end + } + if torch_dir and torch_dir ~= "" then + target:add("includedirs", + path.join(torch_dir, "include"), + path.join(torch_dir, "include/torch/csrc/api/include"), + { public = true }) + target:add("linkdirs", path.join(torch_dir, "lib"), { public = true }) + target:add("rpathdirs", path.join(torch_dir, "lib")) + target:add("links", + "torch", "torch_cpu", "torch_cuda", "torch_python", + "c10", "c10_cuda", + { public = true }) + end + local qy_root = os.getenv("QY_ROOT") or os.getenv("QY_HOME") or os.getenv("QY_PATH") + if (not qy_root or qy_root == "") and os.isdir("/usr/local/dlgpu/sdk") then + qy_root = "/usr/local/dlgpu" + end + if (not qy_root or qy_root == "") and os.isdir("/usr/local/denglin/sdk") then + qy_root = "/usr/local/denglin" + end + if qy_root and os.isdir(qy_root .. "/sdk/include") then + target:add("includedirs", qy_root .. "/sdk/include", { public = true }) + target:add("linkdirs", qy_root .. "/sdk/lib", { public = true }) + target:add("links", "curt", { public = true }) + end + end + if not os.isfile(OPERATOR_CALLS_HEADER) then return end @@ -124,7 +197,11 @@ target("infinicore_runtime") target_end() target("_infinicore") - add_packages("pybind11") + if PYBIND11_INCLUDE and PYBIND11_INCLUDE ~= "" then + add_includedirs(PYBIND11_INCLUDE) + else + add_packages("pybind11") + end set_default(false) add_rules("python.module", { soabi = true }) set_languages("cxx17") @@ -150,7 +227,11 @@ target("_infinicore") target_end() target("_infinilm") - add_packages("pybind11") + if PYBIND11_INCLUDE and PYBIND11_INCLUDE ~= "" then + add_includedirs(PYBIND11_INCLUDE) + else + add_packages("pybind11") + end set_default(false) add_rules("python.module", { soabi = true }) set_languages("cxx17") @@ -178,6 +259,7 @@ target("_infinilm") add_files("csrc/models/internlm3/*.cpp") add_files("csrc/models/llama/*.cpp") add_files("csrc/models/minicpm4/*.cpp") + add_files("csrc/models/minicpm5_moe/*.cpp") add_files("csrc/models/minicpm_eagle/*.cpp") add_files("csrc/models/qwen2/*.cpp") add_files("csrc/models/qwen3/*.cpp")