feat(qy): MiniCPM5 FA, longctx MoE/RoPE, CUDA-graph capture safety

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<const cache::PagedKVCacheConfig *>(cache_config);
         if (paged_cache_config == nullptr) {

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

csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp

@@ -8,7 +8,7 @@
 #if defined(ENABLE_HYGON_API)
 #include <c10/hip/HIPGuard.h>
 #include <c10/hip/HIPStream.h>
-#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API)
+#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API)
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAStream.h>
@@ -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,14 +62,17 @@ 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
 
 #if defined(ENABLE_MOORE_API)
 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

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 <ATen/ATen.h>
+#include <cstdint>
+#include <optional>
+#include <vector>
+
+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<at::Tensor> &b_bias,
+    const std::optional<at::Tensor> &b_scale,
+    const std::optional<at::Tensor> &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<int64_t> &block_shape,
+    int64_t m);
+
+} // namespace infinicore::adaptor::denglin_fused_moe
+#endif

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_schema> &plan_dispatcher();       \
+        static common::OpDispatcher<run_schema> &run_dispatcher();         \
+        static common::OpDispatcher<cleanup_schema> &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__

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<Tensor, Tensor, Tensor> moe_align(
     const Tensor &topk_ids,

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,

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<Tensor, Tensor> moe_topk_sigmoid(
     const Tensor &gating_output,

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);

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,

csrc/infinicore/src/adaptor/aten_adaptor.cc

@@ -1,6 +1,8 @@
 #ifdef ENABLE_ATEN
 #include "infinicore/adaptor/aten_adaptor.hpp"
 
+#include <stdexcept>
+
 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

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 <cstdlib>
+#include <dlfcn.h>
+#include <mutex>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+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<long> const&, long)
+using invoke_fused_moe_opt_fn = void (*)(
+    at::Tensor &,
+    at::Tensor &,
+    at::Tensor &,
+    const std::optional<at::Tensor> &,
+    const std::optional<at::Tensor> &,
+    const std::optional<at::Tensor> &,
+    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> &,
+    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<invoke_fused_moe_opt_fn>(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<at::Tensor> &b_bias,
+    const std::optional<at::Tensor> &b_scale,
+    const std::optional<at::Tensor> &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<int64_t> &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